diff --git a/src/agent_env/training/trl.py b/src/agent_env/training/trl.py index 9f4123a63..83efa5302 100644 --- a/src/agent_env/training/trl.py +++ b/src/agent_env/training/trl.py @@ -9,11 +9,13 @@ import asyncio import contextlib +import copyreg import json import logging import re import threading import time +import uuid import weakref from collections.abc import Callable, Coroutine, Mapping from dataclasses import dataclass, field @@ -186,14 +188,10 @@ async def _environment(self, toolset: _Toolset) -> type[_Environment]: if first is None: raise RuntimeError(f"task {index} of split {split!r} ended before its tools could be listed: " f"{_error_of(turn.observation)}") - methods = {name: tool_method(tool) for name, tool in tools.items()} shared = _Shared(url=self.url, offered=first, tool_timeout=self._tool_timeout, reset_timeout=self._reset_timeout) self._shared.append(shared) - environment = type(toolset.name, (_Environment,), {**methods, "_shared": shared}) - for method in methods.values(): - method.__qualname__ = f"{toolset.name}.{method.__name__}" - return environment + return _environment_class(toolset.name, tuple(tools.values()), shared) def _split(self, split: str) -> list[dict[str, Any]]: if split not in self._tasks: @@ -207,15 +205,17 @@ class _Toolset: first: tuple[str, int, int] # the split, index and turn count of its first task -@dataclass(frozen=True, eq=False) +@dataclass(frozen=True) class _Shared: """What an environment class's instances share: the server, the tools every episode's first turn must offer, and - the deadlines.""" + the deadlines. The key makes it equal to its pickled copies only, so ``AgentEnv.close`` also closes what a class + rebuilt from one opened.""" url: str offered: frozenset[str] tool_timeout: float reset_timeout: float + key: str = field(default_factory=lambda: uuid.uuid4().hex) class _Failed(Exception): @@ -251,6 +251,7 @@ class _Environment: turn, as AsyncGRPO can call two tools of a completion at once.""" _shared: ClassVar[_Shared] + _tools: ClassVar[tuple[dict[str, Any], ...]] def __init__(self) -> None: self._client: GenericEnvClient | None = None @@ -400,6 +401,29 @@ async def _close(self) -> None: await client.close() +class _EnvironmentClass(type): + """The type of ``AgentEnv.factory``'s classes, which pickle as the toolset and server they're built from, since + AsyncGRPO builds its environments in a process of its own.""" + + +def _environment_class(name: str, tools: tuple[dict[str, Any], ...], shared: _Shared) -> type[_Environment]: + methods = {tool["name"]: tool_method(tool) for tool in tools} + for method in methods.values(): + method.__qualname__ = f"{name}.{method.__name__}" + return _EnvironmentClass(name, (_Environment,), {**methods, "_shared": shared, "_tools": tools}) + + +def _reduce_environment_class(environment: _EnvironmentClass) -> tuple[Any, ...] | str: + """A class ``_environment_class`` built pickles as what it's built from, and a subclass of one by its name, as any + class does, so it keeps what it overrides.""" + if "_tools" not in vars(environment): + return environment.__qualname__ + return _environment_class, (environment.__name__, environment._tools, environment._shared) + + +copyreg.pickle(_EnvironmentClass, _reduce_environment_class) + + class _Server: """An app under uvicorn on a thread of its own, its lifespan included, on a free loopback port.""" diff --git a/tst/unit/training/trl_test.py b/tst/unit/training/trl_test.py index 09290aad7..0821abe8d 100644 --- a/tst/unit/training/trl_test.py +++ b/tst/unit/training/trl_test.py @@ -7,6 +7,8 @@ import inspect import json import logging +import multiprocessing +import pickle import pytest from agentenv_protocol import EnvironmentTool, ToolResult @@ -85,6 +87,27 @@ def test_a_rollout_plays_its_rows_task_and_an_environment_plays_one_after_anothe assert len(teardowns.finished) == 3 +def test_the_factory_pickles_and_plays_in_a_spawned_process_as_asyncgrpos_rollout_process_does(served): + env = served({"a": PLAYED}) + + with multiprocessing.get_context("spawn").Pool(1) as pool: + assert pool.apply(_play, (env.factory, env.dataset[0])) == 1.0 + rebuilt = pickle.loads(pickle.dumps(env.factory))() + rebuilt.reset(**env.dataset[0]) + assert (_tools(type(rebuilt)), rebuilt.add_item(item="crate")) == (_tools(env.factory), "add_item {'item': 'crate'}") + env.close() + assert rebuilt._client is None + with pytest.raises(pickle.PicklingError, match="recorded"): + pickle.dumps(type("recorded", (env.factory,), {"get_reward": lambda self: 0.5})) + + +def _play(factory, row): + environment = factory() + environment.reset(**row) + environment.end_turn(response="Added a crate.") + return environment.get_reward() + + def test_end_turn_answers_with_the_next_turns_prompt_and_a_turn_left_open_ends_empty(served): env = served({"a": [ITEMS, ALICE, TURN, AGAIN, REPLIES]}) environment = env.factory() @@ -280,7 +303,7 @@ def test_tasks_whose_agent_gets_other_tools_take_trls_multi_environment_mode(ser "items": ["add_item", "end_turn", "get_reward", "reset"], "pantry": ["end_turn", "get_reward", "reset", "take_item"]} assert [(row["environment"], row["index"]) for row in env.dataset] == [("items", 0), ("pantry", 1)] - pantry = env.factory["pantry"]() + pantry = pickle.loads(pickle.dumps(env.factory))["pantry"]() pantry.reset(**{key: value for key, value in env.dataset[1].items() if key != "environment"}) assert pantry.take_item() == f"take_item on {PANTRY_MCP}" pantry.end_turn(response="Took one.")