feat(nemo-gym): serve a written resources server on its own

pyproject.toml

@@ -62,10 +62,12 @@ trl = [
     "transformers>=5.2",  # renders the tools as TRL shows them to a model; 5.2 is what TRL's environment_factory needs
 ]
 # agent_env.training.nemo_gym: a bundle served to NeMo Gym as a resources server. rc0 is the version NeMo RL's Gym
-# checkout reports; Gym needs Python 3.13.14 or later.
+# checkout reports; Gym needs Python 3.13.14 or later. omegaconf builds the config Gym's server client takes when the
+# server runs on its own.
 nemo-gym = [
     "agentenv-framework[openenv]",
     "nemo-gym>=0.7.0rc0; python_full_version >= '3.13.14'",
+    "omegaconf>=2.3; python_full_version >= '3.13.14'",
 ]
 # The AWS store backends (S3, Secrets Manager, DynamoDB, ECR), selected by impl pointer in config.toml.
 aws = [

src/agent_env/training/nemo_gym/__init__.py

@@ -1,4 +1,5 @@
 """NeMo Gym, and the trainers that take its environments (NeMo RL among them), on a bundle's tasks served as an OpenEnv
 env: ``python -m agent_env.training.nemo_gym`` writes the bundle's resources server into a Gym root, with its config and
 a JSONL file of rows per split, and ``server.AgentEnvResourcesServer`` is that server, which runs in the venv NeMo Gym
-builds for it."""
+builds for it, or on its own with ``python -m agent_env.training.nemo_gym serve`` for a trainer that takes its URL, as
+verifiers' NeMo Gym taskset does."""

src/agent_env/training/nemo_gym/__main__.py

@@ -1,11 +1,33 @@
-"""``python -m agent_env.training.nemo_gym SOURCE ROOT``: writes NeMo Gym's resources server for a bundle's tasks."""
+"""``python -m agent_env.training.nemo_gym SOURCE ROOT``: writes NeMo Gym's resources server for a bundle's tasks.
+``python -m agent_env.training.nemo_gym serve FOLDER``: runs a written one on its own, for a trainer that takes a
+resources server's URL, as verifiers' NeMo Gym taskset does."""
 
 import click
 
 from agent_env.training.nemo_gym.prepare import write_server
 
+try:
+    from agent_env.training.nemo_gym.server import serve as run_server
+except ImportError as exc:
+    run_server, UNSERVABLE = None, str(exc)
 
-@click.command()
+
+class _WriteByDefault(click.Group):
+    """``serve`` and ``write`` by name; anything else is ``write``'s arguments, as the command took them before it had
+    subcommands."""
+
+    def parse_args(self, ctx: click.Context, args: list[str]) -> list[str]:
+        if args and args[0] not in self.commands and args[0] != "--help":
+            args = ["write", *args]
+        return super().parse_args(ctx, args)
+
+
+@click.group(cls=_WriteByDefault)
+def main() -> None:
+    """Writes NeMo Gym's resources server for a bundle's tasks, or serves a written one on its own."""
+
+
+@main.command()
 @click.argument("source")
 @click.argument("root", type=click.Path(file_okay=False))
 @click.option("--name", help="The server's name; the bundle's env name by default.")
@@ -16,8 +38,8 @@
                    "reward is then their sum.")
 @click.option("--max-concurrent-envs", type=int, default=8, show_default=True,
               help="The episodes the server runs at once, which should cover a step's rollouts: the rest wait their turn.")
-def main(source: str, root: str, name: str | None, sandbox: str | None, reward: str | None, reward_components: bool,
-         max_concurrent_envs: int) -> None:
+def write(source: str, root: str, name: str | None, sandbox: str | None, reward: str | None, reward_components: bool,
+          max_concurrent_envs: int) -> None:
     """Writes the resources server for SOURCE, a bundle or the URL of an OpenEnv server serving one, into the Gym root
     ROOT, with its config and a JSONL file of rows per split. Playing one task of each toolset to find its tools
     deploys its envs."""
@@ -26,5 +48,19 @@ def main(source: str, root: str, name: str | None, sandbox: str | None, reward:
     click.echo(f"wrote {server}")
 
 
+@main.command()
+@click.argument("folder", type=click.Path(exists=True, file_okay=False))
+@click.option("--host", default="127.0.0.1", show_default=True,
+              help="The address to listen on. The server checks no token, so keep it off public interfaces.")
+@click.option("--port", type=int, required=True, help="The port to listen on.")
+def serve(folder: str, host: str, port: int) -> None:
+    """Runs the resources server written into FOLDER, a Gym root's resources_servers/<name>, on its own: without Gym's
+    head server or the agent servers its config names. Point verifiers' NeMo Gym taskset at it with
+    env.taskset.task.resources_url, and its env.taskset.dataset at the folder's data/<split>.jsonl."""
+    if run_server is None:
+        raise click.ClickException(UNSERVABLE)
+    run_server(folder, host=host, port=port)
+
+
 if __name__ == "__main__":
     main()

src/agent_env/training/nemo_gym/server.py

@@ -2,15 +2,19 @@
 package ``openenv import`` writes). Gym's ``simple_agent`` seeds a session with a row, which starts the row's episode,
 calls the episode's tools at their routes, and verifies the session, which ends the turns the model left open, its last
 reply ending the first, and answers with the episode's reward; a rollout the environment failed is masked out of
-training. It runs in the venv Gym builds for it, with one worker, since a rollout's session lives in the process."""
+training. It runs in the venv Gym builds for it, or on its own through ``serve``, with one worker, since a rollout's
+session lives in the process."""
 
 import asyncio
 import logging
 import time
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
+from pathlib import Path
 from typing import Any
 
+import uvicorn
+import yaml
 from fastapi import Body, FastAPI, HTTPException, Request
 from fastapi.responses import PlainTextResponse
 from pydantic import ConfigDict, PrivateAttr, model_validator
@@ -26,7 +30,9 @@
         ResourcesCloseSessionResponse,
         SimpleResourcesServer,
     )
-    from nemo_gym.server_utils import SESSION_ID_KEY
+    from nemo_gym.config_types import BaseServerConfig
+    from nemo_gym.server_utils import SESSION_ID_KEY, ServerClient
+    from omegaconf import OmegaConf
     from openenv.core.generic_client import GenericEnvClient  # noqa: F401
 except ModuleNotFoundError as exc:
     if exc.name not in ("nemo_gym", "openenv"):
@@ -199,6 +205,24 @@ async def _shut_down(self) -> None:
                            "may be left running", served.url, STOP_SECONDS)
 
 
+def serve(folder: str | Path, *, host: str, port: int) -> None:
+    """Runs the resources server ``write_server`` wrote into ``folder`` on its own, at ``host``:``port``, without Gym's
+    head server or the agent servers its config names, for a trainer that plays the folder's rows against the server's
+    URL. As verifiers' launcher for Gym's own servers does, its client names a head server nothing calls."""
+    folder = Path(folder).resolve()
+    written = yaml.safe_load((folder / "configs" / f"{folder.name}.yaml").read_text())
+    config = written[folder.name]["resources_servers"][folder.name]
+    resources = AgentEnvResourcesServer(
+        config=AgentEnvResourcesServerConfig(**config, name=folder.name, host=host, port=port),
+        server_client=ServerClient(head_server_config=BaseServerConfig(host=host, port=port),
+                                   global_config_dict=OmegaConf.create({})),
+    )
+    app = resources.setup_webserver()
+    resources.setup_liveness(app)
+    resources.setup_exception_middleware(app)
+    uvicorn.run(app, host=host, port=port)
+
+
 class _Session(Session):
     """A rollout's session, which shows the model the tools its row lists. The server gives it the reset deadline to
     wait for room in, since Gym runs every rollout of a step at once and a server at its limit makes the rest take

tst/nemo_gym/server_test.py

@@ -1,13 +1,16 @@
 """``AgentEnvResourcesServer`` driven as NeMo Gym's ``simple_agent`` drives it, through a client that keeps the session
 cookie Gym keeps per rollout, over the bundle the unit tier serves on loopback: a row's seed starts its task's episode,
 the tool routes answer as text, verify ends the turns left open and scores the episode or masks the rollout, a rollout
-waits its turn at the server's limit, a closed or idle session ends its episode, the server takes one worker, and a
-bundle named in the config is served for the server's lifetime. CI runs it on the nemo-gym ``uv.lock`` pins and on the
-latest nemo-gym."""
+waits its turn at the server's limit, a closed or idle session ends its episode, the server takes one worker, a
+bundle named in the config is served for the server's lifetime, and ``serve`` runs a written server on its own. CI runs
+it on the nemo-gym ``uv.lock`` pins and on the latest nemo-gym."""
 
 import asyncio
 import json
 import logging
+import socket
+import subprocess
+import sys
 import time
 from unittest.mock import MagicMock
 
@@ -16,6 +19,7 @@
 import pytest
 from fastapi.testclient import TestClient
 
+from agent_env.training.nemo_gym.prepare import write_server
 from tst.unit.bundle._support import layout
 from tst.unit.openenv.environment_test import AGAIN
 from tst.unit.task.episode.run_test import ALICE, ITEMS, REPLIES, TURN
@@ -174,6 +178,28 @@ def test_a_bundle_named_in_the_config_is_served_while_the_server_runs(tmp_path,
     assert len(teardowns.finished) == 2
 
 
+def test_serve_runs_a_written_server_on_its_own_at_the_port_it_is_given(serve, tmp_path, teardowns):
+    bundle_server = serve({"a": PLAYED})
+    folder = write_server(bundle_server.url, tmp_path)
+    row = json.loads((folder / "data" / "all.jsonl").read_text().splitlines()[0])
+    port = _free_port()
+
+    served = subprocess.Popen([sys.executable, "-m", "agent_env.training.nemo_gym", "serve", ".", "--port", str(port)],
+                              cwd=folder, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
+    try:
+        with httpx.Client(base_url=_up(served, f"http://127.0.0.1:{port}")) as client:
+            assert client.post("/seed_session", json=row).json() == {}
+            assert client.post("/add_item", json={"item": "crate"}).text == "add_item {'item': 'crate'}"
+            verified = _verify(client, row, "Added a crate.")
+    finally:
+        served.terminate()
+        served.wait(timeout=30)
+    bundle_server.stop()
+
+    assert (verified["reward"], verified["mask_sample"], verified["task_id"]) == (1.0, False, "a")
+    assert len(teardowns.finished) == 2
+
+
 _EPISODE = {"rollout_id": "rollout"}
 
 
@@ -202,6 +228,26 @@ def _verified(row, reply):
     return {**row, "response": response}
 
 
+def _free_port():
+    with socket.socket() as bound:
+        bound.bind(("127.0.0.1", 0))
+        return bound.getsockname()[1]
+
+
+def _up(process, url, timeout=180):
+    """``url`` once the server ``process`` runs answers there; the first import of NeMo Gym can take a minute."""
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        if process.poll() is not None:
+            raise AssertionError(f"serve exited {process.returncode}: {process.stdout.read()}")
+        try:
+            if httpx.get(url).status_code == 200:
+                return url
+        except httpx.TransportError:
+            time.sleep(0.5)
+    raise AssertionError(f"serve didn't answer at {url} in {timeout}s")
+
+
 def _client(app):
     """A client of its own cookie jar, as each rollout of Gym's has, on the app without its lifespan, which serves
     nothing for a config that names a url."""

tst/unit/test_runtime_dependencies_are_imported.py

@@ -22,11 +22,13 @@
 PYPROJECT = SRC.parents[1] / "pyproject.toml"
 CONTAINER_ONLY = (SRC / "env" / "gateway",)
 SELF = "agentenv-framework"
-# The Smol SDK has no wheel for some development hosts, and NeMo Gym needs Python 3.13.14 or later; both are
-# deliberately absent from [dev].
+# The Smol SDK has no wheel for some development hosts, and NeMo Gym needs Python 3.13.14 or later, as the extra's
+# omegaconf does with it; all are deliberately absent from [dev].
 _UNINSTALLED_OPTIONAL_MODULES = {
     "smolmachines": frozenset({"smol"}),
-    "nemo-gym": frozenset({"nemo_gym", "nemo_gym.base_resources_server", "nemo_gym.server_utils"}),
+    "nemo-gym": frozenset({"nemo_gym", "nemo_gym.base_resources_server", "nemo_gym.config_types",
+                           "nemo_gym.server_utils"}),
+    "omegaconf": frozenset({"omegaconf"}),
 }