fix(training): pickle AgentEnv.factory's classes so TRL's AsyncGRPO can use them

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."""
 

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.")