From cf5e855b5feeb0d16d591b2cf4c65779d69c552a Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 30 Oct 2026 12:01:50 -0700 Subject: [PATCH 1/3] feat(nemo-gym): serve a written resources server on its own `python -m agent_env.training.nemo_gym serve FOLDER --port P` runs the resources server `write` wrote into FOLDER without Gym's head server or the agent servers its config names, for a trainer that takes a resources server's URL. verifiers' NeMo Gym taskset (prime-rl's) does: its managed mode starts only servers from NeMo Gym's own package. The server is built as verifiers' own launcher builds Gym's servers. The command becomes a group with `write` and `serve`; any other first argument is write's, so `python -m agent_env.training.nemo_gym SOURCE ROOT` works as before. Writing doesn't need the nemo-gym extra, so `__main__` imports the server guarded and `serve` reports the extra's install hint when it's missing. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/training/nemo_gym/__init__.py | 3 +- src/agent_env/training/nemo_gym/__main__.py | 44 +++++++++++++++-- src/agent_env/training/nemo_gym/server.py | 28 ++++++++++- tst/nemo_gym/server_test.py | 52 +++++++++++++++++++-- 4 files changed, 117 insertions(+), 10 deletions(-) diff --git a/src/agent_env/training/nemo_gym/__init__.py b/src/agent_env/training/nemo_gym/__init__.py index b179d10f..2c41e9b2 100644 --- a/src/agent_env/training/nemo_gym/__init__.py +++ b/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.""" diff --git a/src/agent_env/training/nemo_gym/__main__.py b/src/agent_env/training/nemo_gym/__main__.py index 26962a85..416add33 100644 --- a/src/agent_env/training/nemo_gym/__main__.py +++ b/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/, 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/.jsonl.""" + if run_server is None: + raise click.ClickException(UNSERVABLE) + run_server(folder, host=host, port=port) + + if __name__ == "__main__": main() diff --git a/src/agent_env/training/nemo_gym/server.py b/src/agent_env/training/nemo_gym/server.py index a761bf9c..c23eeab6 100644 --- a/src/agent_env/training/nemo_gym/server.py +++ b/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) + 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 diff --git a/tst/nemo_gym/server_test.py b/tst/nemo_gym/server_test.py index 097098a5..c5ef1fd3 100644 --- a/tst/nemo_gym/server_test.py +++ b/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", str(folder), "--port", + str(port)], 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.""" From 62d191e0c58d88d3b5d799e63c0949d90b2f2240 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 30 Oct 2026 12:16:05 -0700 Subject: [PATCH 2/3] fix(nemo-gym): declare omegaconf, which serve's Gym client config needs serve builds Gym's ServerClient, whose config must be an omegaconf DictConfig, so the server imports omegaconf directly. Declare it in the nemo-gym extra under Gym's marker, and teach the dependency guard the nemo-gym and omegaconf module paths it can't read when the extra isn't installed. Co-Authored-By: Claude Opus 5.5 --- pyproject.toml | 4 +++- tst/unit/test_runtime_dependencies_are_imported.py | 8 +++++--- uv.lock | 2 ++ 3 files changed, 10 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 95c44355..8112115f 100644 --- a/pyproject.toml +++ b/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 = [ diff --git a/tst/unit/test_runtime_dependencies_are_imported.py b/tst/unit/test_runtime_dependencies_are_imported.py index ca82ceec..2a72f134 100644 --- a/tst/unit/test_runtime_dependencies_are_imported.py +++ b/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"}), } diff --git a/uv.lock b/uv.lock index 6bf2e58c..8e2f06a0 100644 --- a/uv.lock +++ b/uv.lock @@ -133,6 +133,7 @@ gcp = [ ] nemo-gym = [ { name = "nemo-gym", marker = "python_full_version >= '3.13.14'" }, + { name = "omegaconf", marker = "python_full_version >= '3.13.14'" }, { name = "openenv" }, ] openenv = [ @@ -196,6 +197,7 @@ requires-dist = [ { name = "modal", specifier = ">=1.5.2.dev22" }, { name = "moto", marker = "extra == 'dev'", specifier = ">=5.0.0" }, { name = "nemo-gym", marker = "python_full_version >= '3.13.14' and extra == 'nemo-gym'", specifier = ">=0.7.0rc0" }, + { name = "omegaconf", marker = "python_full_version >= '3.13.14' and extra == 'nemo-gym'", specifier = ">=2.3" }, { name = "openenv", marker = "extra == 'openenv'", specifier = ">=0.9" }, { name = "packaging", specifier = ">=22" }, { name = "pillow", marker = "extra == 'dev'", specifier = ">=10.4" }, From 270377547a3b925360a75a663cb0d0dd491d6d58 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 30 Oct 2026 12:17:00 -0700 Subject: [PATCH 3/3] fix(nemo-gym): serve resolves its folder, so serve . works The folder's name is the server's, and Path(".").name is empty. The serve test now runs serve . from inside the written folder. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/training/nemo_gym/server.py | 2 +- tst/nemo_gym/server_test.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/agent_env/training/nemo_gym/server.py b/src/agent_env/training/nemo_gym/server.py index c23eeab6..48ea6d21 100644 --- a/src/agent_env/training/nemo_gym/server.py +++ b/src/agent_env/training/nemo_gym/server.py @@ -209,7 +209,7 @@ 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) + 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( diff --git a/tst/nemo_gym/server_test.py b/tst/nemo_gym/server_test.py index c5ef1fd3..103d8ade 100644 --- a/tst/nemo_gym/server_test.py +++ b/tst/nemo_gym/server_test.py @@ -184,8 +184,8 @@ def test_serve_runs_a_written_server_on_its_own_at_the_port_it_is_given(serve, t 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", str(folder), "--port", - str(port)], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) + 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() == {}