From 5daf3e5cbebc26a483b01cb2bd9e1c6207942f36 Mon Sep 17 00:00:00 2001 From: Felipe Vieira Frujeri Date: Fri, 30 Oct 2026 06:05:09 +0000 Subject: [PATCH] feat(environment): route BaseEnvironmentServerWithTools tool calls to the worker that runs the episode With several uvicorn workers, a tool call for an episode can reach a worker that does not hold the episode, so BaseEnvironmentServerWithTools refused num_workers > 1. It now routes sessions to their owner, as Resources and Agent Servers do: - serve_tools() claims the episode, by its capture key, on the worker that runs it; the episode's cleanup releases the claim. A second serve of the same attempt on another worker is refused. - The router reads a tool call's episode from its episode token header, without the nonce, which only the episode's worker checks. The MCP grant carries the same header and /mcp is stateless, so every MCP request (initialize, tools/list, tools/call) is routed the same way. - session_routing gains a generic header source for the session ID (session_header with the server's reader) and SessionHolder, for a server that claims and releases each session itself. SessionKind's cookie key, body field and seed and close routes become optional. Resources and Agent Servers are unchanged. - An end-to-end test runs a test-only BaseEnvironmentServerWithTools with four workers, with episodes that call their tools over direct HTTP and over MCP; the routing e2e job also runs when the base class changes. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Felipe Vieira Frujeri --- .github/workflows/unit-tests.yml | 1 + .../server_runtime/base_environment_server.py | 56 +++- .../request_handling/session_routing.py | 99 +++++++- .../environment_servers/tool_counter/app.py | 149 +++++++++++ .../test_episode_tools_routing_e2e.py | 220 ++++++++++++++++ ...test_base_environment_server_with_tools.py | 11 +- .../test_episode_tools_session_routing.py | 240 ++++++++++++++++++ 7 files changed, 757 insertions(+), 19 deletions(-) create mode 100644 tests/e2e/session_routing/episode_tools_servers/environment_servers/tool_counter/app.py create mode 100644 tests/e2e/session_routing/test_episode_tools_routing_e2e.py create mode 100644 tests/unit_tests/server_runtime/test_episode_tools_session_routing.py diff --git a/.github/workflows/unit-tests.yml b/.github/workflows/unit-tests.yml index 75cc074b930..2d52488e15c 100644 --- a/.github/workflows/unit-tests.yml +++ b/.github/workflows/unit-tests.yml @@ -69,6 +69,7 @@ jobs: nemo_gym/client/** nemo_gym/agents/base_responses_api_agent.py nemo_gym/server_runtime/base_resources_server.py + nemo_gym/server_runtime/base_environment_server.py nemo_gym/server_runtime/mcp_auto_exposure.py tests/e2e/session_routing/** diff --git a/nemo_gym/server_runtime/base_environment_server.py b/nemo_gym/server_runtime/base_environment_server.py index b1565aad902..90603352254 100644 --- a/nemo_gym/server_runtime/base_environment_server.py +++ b/nemo_gym/server_runtime/base_environment_server.py @@ -17,7 +17,16 @@ from anyio import CancelScope from fastapi import Body, FastAPI, Request from fastapi.responses import JSONResponse -from pydantic import ConfigDict, Field, JsonValue, PositiveFloat, PositiveInt, field_validator, model_validator +from pydantic import ( + ConfigDict, + Field, + JsonValue, + PositiveFloat, + PositiveInt, + ValidationError, + field_validator, + model_validator, +) from typing_extensions import Self from nemo_gym.client.http_client import format_exception_with_response_body @@ -38,6 +47,7 @@ from nemo_gym.server_runtime.episode_tokens import EPISODE_TOKEN_HEADER, EpisodeTokenClaims, EpisodeTokens from nemo_gym.server_runtime.episode_tools import EpisodeToolset, ToolArgumentsError, ToolCallRecord, ToolSpec from nemo_gym.server_runtime.request_handling.rollout_context import RolloutContextMiddleware +from nemo_gym.server_runtime.request_handling.session_routing import SessionKind, SessionOwners from nemo_gym.server_runtime.tool_access import ( DirectHTTPToolAccess, MCPStreamableHTTPConnection, @@ -540,6 +550,15 @@ async def aggregate_metrics(self, body: AggregateMetricsRequest = Body()) -> Agg """Aggregate per-rollout scores into task-level metrics.""" +# An episode's tool calls name it only by the episode token header; no cookie, seed or close names it. +EPISODE_TOOL_SESSIONS = SessionKind( + cookie_key=None, + id_field=None, + seed_path=None, + close_path=None, + cookie_session=False, +) + ToolTransport = Literal["direct_http", "mcp"] # The path ``nemo_gym.server_runtime.mcp_auto_exposure.mount_mcp_endpoint`` serves. Kept here so that a server without MCP never # imports the MCP SDK. @@ -615,17 +634,23 @@ class BaseEnvironmentServerWithTools(BaseEnvironmentServer[EpisodeRequestT, Epis Both transports check the token the same way. Calls to one episode run one at a time, whatever their transport, and each is recorded in ``ServedTools.calls``. + With several workers, the worker that runs the episode owns it (see ``session_routing``), so a tool call that + reaches another worker, over either transport, is forwarded there, by the episode its token names. + When the server shuts down, it finishes the cleanups of every episode still running, such as stopping its sandbox. """ config: BaseEnvironmentServerWithToolsConfig + routes_sessions_to_owner = True _episode_tokens: Optional[EpisodeTokens] = None _open_episodes: Optional[dict[EpisodeId, _OpenEpisode]] = None # Keyed by id(): two attempts of one episode can run at once without checkpointing. _running_cleanups: Optional[dict[int, CleanupContext]] = None _reserved_tool_names: frozenset[str] = frozenset() + # With several workers, which worker runs each episode (see nemo_gym.server_runtime.request_handling.session_routing). + _session_owners: Optional[SessionOwners] = None def model_post_init(self, context: Any, /) -> None: super().model_post_init(context) @@ -682,8 +707,27 @@ async def finish_running_cleanups(self) -> None: """ await asyncio.gather(*(cleanup.aclose() for cleanup in list(self._running_cleanups.values()))) - def multiple_workers_unsupported_reason(self) -> Optional[str]: - return "a tool call reaches whichever worker accepts it, and only the worker that runs the episode holds it" + def owned_session_routing(self) -> dict[str, Any]: + return { + "kind": EPISODE_TOOL_SESSIONS, + "session_header": EPISODE_TOKEN_HEADER, + "session_from_header": self._episode_session_id, + "session_holder": self, + } + + def _episode_session_id(self, token: str) -> Optional[str]: + # Routing needs only the episode the token names; the worker that runs the episode checks its nonce. + claims = self._episode_tokens.read(token) + return claims.episode_id.capture_key if claims is not None else None + + def holds_session(self, session_id: str) -> bool: + try: + return EpisodeId.from_capture_key(session_id) in self._open_episodes + except (ValueError, ValidationError): + return False + + def use_session_owners(self, owners: SessionOwners) -> None: + self._session_owners = owners async def serve_tools( self, request: EpisodeRequestT, toolset: EpisodeToolset, cleanup: CleanupContext @@ -709,6 +753,10 @@ async def serve_tools( if unnamable: raise ValueError(f"MCP tool names may use only letters, digits, _ and -: {', '.join(unnamable)}") + session_owners = self._session_owners + if session_owners is not None and session_owners.acquire(episode_id.capture_key) != session_owners.routing_id: + raise ValueError(f"Episode {episode_id.capture_key} already serves its tools on another worker") + token, claims = self._episode_tokens.mint(episode_id) episode = _OpenEpisode(claims=claims, toolset=toolset, specs=specs, calls=[]) self._open_episodes[episode_id] = episode @@ -716,6 +764,8 @@ async def serve_tools( async def stop_serving() -> None: if self._open_episodes.get(episode_id) is episode: del self._open_episodes[episode_id] + if session_owners is not None: + session_owners.release(episode_id.capture_key) cleanup.register_cleanup("episode tools", stop_serving) base_url = self.server_client._resolve_base_url(self.config.name).rstrip("/") diff --git a/nemo_gym/server_runtime/request_handling/session_routing.py b/nemo_gym/server_runtime/request_handling/session_routing.py index 0bc21f3fd68..4e4deda5740 100644 --- a/nemo_gym/server_runtime/request_handling/session_routing.py +++ b/nemo_gym/server_runtime/request_handling/session_routing.py @@ -39,6 +39,7 @@ | Any request to an agent server | The session cookie's ``agent_session_id`` | | An agent session seed or close | The body's ``agent_session_id`` | | An agent rollout stop without a cookie | The path's ``/rollouts/``, which the session's worker also owns | +| A tool call to a ``BaseEnvironmentServerWithTools`` | The episode named in its episode token header | The body is read only on the seed and close routes, where it names the session before any cookie exists. If the cookie and the body name different sessions, the reply is 409. @@ -64,7 +65,9 @@ - a resources session's ownership lasts while the session is open, from a successful seed until its close or ``/verify``, and expires ``session_close_retry_window_seconds`` after the last request for it otherwise. - A ``/verify`` that says a close follows, as an environment server's does, leaves the session open for that close. + A ``/verify`` that says a close follows, as an environment server's does, leaves the session open for that close; +- an episode of a ``BaseEnvironmentServerWithTools`` is owned by the worker that runs it, from when it starts serving + the episode's tools until it stops (see ``SessionHolder``). An agent worker also owns ``/rollouts/`` from the start of a session's seed until the session closes, so a rollout stop, which names only the rollout, reaches the worker that runs it. @@ -163,15 +166,19 @@ def _paths(paths: tuple[str, ...]) -> re.Pattern[str]: @dataclass(frozen=True) class SessionKind: - """Where one kind of server's requests name their session.""" + """Where one kind of server's requests name their session. + + A kind whose requests name their session only in a header, such as the episode tool calls of a + ``BaseEnvironmentServerWithTools``, has no cookie key, body field or seed and close routes. + """ #: The key of the session ID in the session cookie. - cookie_key: str + cookie_key: Optional[str] #: The body field that names the session on the seed and close routes. - id_field: str + id_field: Optional[str] #: The seed and close routes, whose body names the session. - seed_path: str - close_path: str + seed_path: Optional[str] + close_path: Optional[str] #: Whether the session ID is the cookie session's own ID, as on a resources server: #: an episode seed's ID becomes it, and a request without one gets a new ID, #: owned by this worker if the request starts a session or its handler uses the session. @@ -182,8 +189,9 @@ class SessionKind: path_key: Optional[re.Pattern[str]] = None @property - def body_paths(self) -> re.Pattern[str]: - return _paths((self.seed_path, self.close_path)) + def body_paths(self) -> Optional[re.Pattern[str]]: + paths = tuple(path for path in (self.seed_path, self.close_path) if path is not None) + return _paths(paths) if paths else None def create_session_socket_dir() -> str: @@ -545,6 +553,42 @@ def __len__(self) -> int: return len(self._until.keys() | self._retained) +class SessionHolder(Protocol): + """A server that acquires each of its sessions itself, and owns it for as long as it holds the session. + + A ``BaseEnvironmentServerWithTools`` is one: it acquires an episode when it starts serving the episode's tools, and + releases it when it stops. + """ + + def holds_session(self, session_id: str) -> bool: + """Whether this worker still holds the session.""" + ... + + def use_session_owners(self, owners: SessionOwners) -> None: + """Keep the session owners, to acquire a session when it starts and release it when it ends.""" + ... + + +class _HeldSessionLifetimes: + """A holder's own state decides what it owns: it acquires and releases each session itself.""" + + def __init__(self, holder: SessionHolder) -> None: + self.holder = holder + + def holds(self, session_id: str) -> bool: + return self.holder.holds_session(session_id) + + def served(self, session_id: str, scope: Scope, status: Optional[int]) -> None: + pass + + def starts(self, path: str) -> bool: + # The holder acquires a session when it starts it, not on a request. + return False + + def sweep(self) -> list[str]: + return [] + + class AgentSessionHolder(Protocol): """An agent server, whose sessions are owned by ``agent_session_id`` across workers.""" @@ -610,6 +654,8 @@ def __init__( mcp_token_header: Optional[str] = None, mcp_token_serializer: Optional[itsdangerous.URLSafeSerializer] = None, mcp_path: str = "/mcp", + session_header: Optional[str] = None, + session_from_header: Optional[Callable[[str], Optional[str]]] = None, ) -> None: self.app = app self.kind = kind @@ -626,13 +672,20 @@ def __init__( self.mcp_token_header = mcp_token_header self.mcp_token_serializer = mcp_token_serializer self.mcp_path = mcp_path + # A header that names the session, and the server's reader of its value; ahead of the cookie and MCP token. + self.session_header = session_header + self.session_from_header = session_from_header async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: if scope["type"] != "http": await self.app(scope, receive, send) return owners = self.owners - names_by_body = scope["method"] == "POST" and self.body_paths.fullmatch(scope["path"]) is not None + names_by_body = ( + scope["method"] == "POST" + and self.body_paths is not None + and self.body_paths.fullmatch(scope["path"]) is not None + ) if owners is None and not names_by_body: await self.app(scope, receive, send) return @@ -722,7 +775,8 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: elif forwarded: # The owner released the ID while the request was forwarded here, and another worker took ownership since. await _conflict( - f"This {self.kind.id_field} moved to another server worker while the request was forwarded. Retry it.", + f"This {self.kind.id_field or 'session'} moved to another server worker while the request was " + "forwarded. Retry it.", scope, receive, send, @@ -749,7 +803,13 @@ async def send_and_record(message: Message) -> None: self.lifetimes.served(session_id, scope, status[0] if status else None) def _session_id_from_headers(self, connection: HTTPConnection) -> Optional[str]: - cookie = connection.cookies.get(self.session_cookie) + if self.session_header is not None: + value = connection.headers.get(self.session_header) + if value is not None: + # A value the reader rejects names no session; the handler refuses the request. + return self.session_from_header(value) + # A kind without a cookie key never names its session in the cookie. + cookie = connection.cookies.get(self.session_cookie) if self.kind.cookie_key is not None else None token = connection.headers.get(self.mcp_token_header) if self.mcp_token_serializer is not None else None path = connection.scope["path"] is_mcp = path == self.mcp_path or path.startswith(self.mcp_path + "/") @@ -964,7 +1024,10 @@ def install_session_routing( secret_key: str, mcp_token_header: Optional[str] = None, mcp_token_serializer: Optional[itsdangerous.URLSafeSerializer] = None, + session_header: Optional[str] = None, + session_from_header: Optional[Callable[[str], Optional[str]]] = None, agent_sessions: Optional[AgentSessionHolder] = None, + session_holder: Optional[SessionHolder] = None, close_retry_window_seconds: Optional[float] = None, session_starts: tuple[str, ...] = (), session_ends: tuple[str, ...] = (), @@ -978,8 +1041,10 @@ def install_session_routing( With one, it records session owners in it and returns this worker's ID, also stored in ``app.state.nemo_gym_routing_id``. Pass the agent for an agent server, whose own records decide how long it owns its sessions, + the holder for a server that acquires and releases each session itself (see ``SessionHolder``), or the close retry window and the session start and end routes for a resources server (see ``ResourcesSessionLifetimes``). + A server whose requests name their session in a header passes the header and its reader of the header's value. """ route_by = dict( kind=kind, @@ -987,7 +1052,11 @@ def install_session_routing( secret_key=secret_key, mcp_token_header=mcp_token_header, mcp_token_serializer=mcp_token_serializer, + session_header=session_header, + session_from_header=session_from_header, ) + if session_header is not None and session_from_header is None: + raise ValueError("install_session_routing needs session_from_header with a session_header") if socket_dir is None: # An agent's own seed handler already refuses a cookie of another session, so it needs nothing here. if kind.cookie_session: @@ -1000,6 +1069,9 @@ def install_session_routing( if agent_sessions is not None: lifetimes: SessionLifetimes = _AgentSessionLifetimes(agent_sessions) window = None + elif session_holder is not None: + lifetimes = _HeldSessionLifetimes(session_holder) + window = None elif close_retry_window_seconds is not None: resources_lifetimes = ResourcesSessionLifetimes( close_retry_window_seconds, @@ -1012,13 +1084,16 @@ def install_session_routing( window = close_retry_window_seconds else: raise ValueError( - "install_session_routing needs agent_sessions or close_retry_window_seconds with a socket_dir" + "install_session_routing needs agent_sessions, session_holder or close_retry_window_seconds " + "with a socket_dir" ) owners = SessionOwners(socket_dir, routing_id, lifetimes.holds) app.state.nemo_gym_session_owners = owners if agent_sessions is not None: agent_sessions.use_agent_session_owners(owners) window = getattr(getattr(agent_sessions, "config", None), "session_close_retry_window_seconds", None) + if session_holder is not None: + session_holder.use_session_owners(owners) sweep_seconds = min(_MAX_SWEEP_SECONDS, max(_MIN_SWEEP_SECONDS, (window or _MAX_SWEEP_SECONDS) / 4)) app.add_middleware( SessionRoutingMiddleware, diff --git a/tests/e2e/session_routing/episode_tools_servers/environment_servers/tool_counter/app.py b/tests/e2e/session_routing/episode_tools_servers/environment_servers/tool_counter/app.py new file mode 100644 index 00000000000..405cce94f8e --- /dev/null +++ b/tests/e2e/session_routing/episode_tools_servers/environment_servers/tool_counter/app.py @@ -0,0 +1,149 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A ``BaseEnvironmentServerWithTools`` whose episodes call their own tools through the server's public port. + +Only ``test_episode_tools_routing_e2e.py`` runs it. +Its directory is on that test's PYTHONPATH, under the ``environment_servers`` namespace, +so uvicorn can import it as ``environment_servers.tool_counter.app`` in each worker. +""" + +import asyncio +import json +import os +from pathlib import Path +from typing import Literal, Optional + +import aiohttp +import httpx +from fastapi import Body +from pydantic import BaseModel + +from nemo_gym.config.config_types import AggregateMetrics, AggregateMetricsRequest +from nemo_gym.core.episode_types import BaseEpisodeRequest, BaseEpisodeResponse +from nemo_gym.server_runtime.base_environment_server import BaseEnvironmentServerWithTools, CleanupContext +from nemo_gym.server_runtime.episode_tools import TypedToolset, episode_tool +from nemo_gym.server_runtime.tool_access import DirectHTTPToolAccess, MCPToolAccess +from nemo_gym.utils.server_process import is_nemo_gym_fastapi_entrypoint + + +class CounterTask(BaseModel): + calls: int = 0 + # Where to write the episode's grant and worker, for a test that needs them while the episode runs. + grant_file: Optional[str] = None + # How long the episode keeps serving its tools after its calls. + hold_seconds: float = 0 + # How the episode calls its tools: as simple_agent does, or as an MCP-only harness such as Codex does. + transport: Literal["direct_http", "mcp"] = "direct_http" + + +class CounterRequest(BaseEpisodeRequest[CounterTask]): + pass + + +class CounterResult(BaseModel): + owner_pid: int + statuses: list[int] + # The worker that ran each call, as the tool reported it; None for a call that did not run. + ran_on: list[Optional[int]] + total: int + # The tools the episode listed over MCP; empty over direct HTTP. + listed: list[str] = [] + + +class CounterResponse(BaseEpisodeResponse[CounterResult]): + pass + + +class Amount(BaseModel): + amount: int + + +class Counter: + def __init__(self) -> None: + self.total = 0 + + @episode_tool + async def add(self, arguments: Amount) -> dict[str, int]: + """Add an amount to the episode's total, and report the worker that ran the call.""" + self.total += arguments.amount + return {"total": self.total, "pid": os.getpid()} + + +class ToolCounterServer(BaseEnvironmentServerWithTools[CounterRequest, CounterResponse]): + request_model = CounterRequest + response_model = CounterResponse + + async def run_episode(self, request: CounterRequest, cleanup: CleanupContext) -> CounterResponse: + task = request.task.task_input + counter = Counter() + served = await self.serve_tools(request, TypedToolset(counter), cleanup) + grant = next(access for access in served.tool_accesses if isinstance(access, DirectHTTPToolAccess)) + base_url = str(grant.base_url).rstrip("/") + headers = dict(grant.headers) + if task.grant_file is not None: + Path(task.grant_file).write_text( + json.dumps({"headers": headers, "base_url": base_url, "pid": os.getpid()}) + ) + statuses: list[int] = [] + ran_on: list[Optional[int]] = [] + listed: list[str] = [] + if task.transport == "mcp": + mcp_grant = next(access for access in served.tool_accesses if isinstance(access, MCPToolAccess)) + listed = await _call_over_mcp(mcp_grant, task.calls, statuses, ran_on) + else: + # Not Gym's shared client: a new connection for each call, so the kernel hands each one to any worker. + async with aiohttp.ClientSession(connector=aiohttp.TCPConnector(force_close=True)) as client: + for _ in range(task.calls): + async with client.post(f"{base_url}/add", json={"amount": 1}, headers=headers) as reply: + statuses.append(reply.status) + ran_on.append((await reply.json()).get("pid")) + await asyncio.sleep(task.hold_seconds) + return CounterResponse( + episode_id=request.episode_id, + task_id=request.task.task_id, + result=CounterResult( + owner_pid=os.getpid(), statuses=statuses, ran_on=ran_on, total=counter.total, listed=listed + ), + ) + + async def aggregate_metrics(self, body: AggregateMetricsRequest = Body()) -> AggregateMetrics: + return AggregateMetrics() + + +def _new_connection_per_request( + headers: Optional[dict[str, str]] = None, + timeout: Optional[httpx.Timeout] = None, + auth: Optional[httpx.Auth] = None, +) -> httpx.AsyncClient: + """The MCP SDK's client without keep-alive, so the kernel hands each MCP request to any worker.""" + return httpx.AsyncClient( + headers=headers, timeout=timeout, auth=auth, limits=httpx.Limits(max_keepalive_connections=0) + ) + + +async def _call_over_mcp( + grant: MCPToolAccess, calls: int, statuses: list[int], ran_on: list[Optional[int]] +) -> list[str]: + """List the episode's tools and call ``add`` over MCP, as an MCP-only harness does; return the listed names.""" + from mcp import ClientSession + from mcp.client.streamable_http import streamablehttp_client + + async with streamablehttp_client( + str(grant.connection.url), + headers=dict(grant.connection.headers), + httpx_client_factory=_new_connection_per_request, + ) as (read, write, _): + async with ClientSession(read, write) as session: + await session.initialize() + listed = [tool.name for tool in (await session.list_tools()).tools] + for _ in range(calls): + result = await session.call_tool("add", {"amount": 1}) + statuses.append(500 if result.isError else 200) + ran_on.append(None if result.isError else json.loads(result.content[0].text).get("pid")) + return listed + + +if __name__ == "__main__": + ToolCounterServer.run_webserver() +elif is_nemo_gym_fastapi_entrypoint(__file__): + app = ToolCounterServer.run_webserver() # noqa: F401 diff --git a/tests/e2e/session_routing/test_episode_tools_routing_e2e.py b/tests/e2e/session_routing/test_episode_tools_routing_e2e.py new file mode 100644 index 00000000000..e43a2fa5d84 --- /dev/null +++ b/tests/e2e/session_routing/test_episode_tools_routing_e2e.py @@ -0,0 +1,220 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Real worker processes of a BaseEnvironmentServerWithTools: each tool call reaches the worker that runs its episode. + +Every episode calls its own tools through the server's public port, over direct HTTP or MCP, on a new connection +each time, so the kernel hands each request to any of the 4 workers. +Skipped unless NEMO_GYM_MULTIWORKER_E2E=1; under a minute. +""" + +import asyncio +import contextlib +import glob +import json +import os +import re +import signal +import socket +import subprocess +import sys +import time +from collections import Counter +from pathlib import Path +from typing import Any, Iterator + +import aiohttp +import psutil +import pytest +import yaml + +from nemo_gym.server_runtime.runtime_dir import RUNTIME_ROOT +from tests.e2e.session_routing.test_session_routing_e2e import _free_port, _port, _server, _wait_healthy + + +pytestmark = pytest.mark.skipif( + os.getenv("NEMO_GYM_MULTIWORKER_E2E") != "1", reason="set NEMO_GYM_MULTIWORKER_E2E=1 to run" +) + +HERE = Path(__file__).resolve().parent +REPO = HERE.parents[2] +SERVERS = HERE / "episode_tools_servers" +WORKERS = 4 +EPISODES = int(os.getenv("NEMO_GYM_MULTIWORKER_E2E_EPISODES", "200")) +CALLS = 8 +# An MCP episode also initializes and lists its tools, each a request of its own. +MCP_EPISODES = int(os.getenv("NEMO_GYM_MULTIWORKER_E2E_MCP_EPISODES", "50")) +MCP_CALLS = 4 + + +class Deployment: + """The tool counter environment server, with 4 workers.""" + + def __init__(self, work_dir: Path) -> None: + self.work_dir = work_dir + self.log = work_dir / "environment.log" + self.config: dict[str, Any] = { + "dry_run": False, + "use_absolute_ip": False, + "skip_venv_if_present": True, + "head_server": {"host": "127.0.0.1", "port": _free_port()}, + "environment": _server( + "environment_servers", + "tool_counter", + num_workers=WORKERS, + cleanup_timeout_seconds=5, + tool_transports=["direct_http", "mcp"], + ), + } + self.proc: subprocess.Popen | None = None + + @property + def url(self) -> str: + return f"http://127.0.0.1:{_port(self.config['environment'])}" + + def start(self) -> None: + env = os.environ | { + "NEMO_GYM_CONFIG_DICT": yaml.safe_dump(self.config), + "NEMO_GYM_CONFIG_PATH": "environment", + # The test server is imported as environment_servers.tool_counter. + "PYTHONPATH": os.pathsep.join([str(REPO), str(SERVERS)]), + "RAY_TMPDIR": "/tmp", + } + self.proc = subprocess.Popen( + [sys.executable, "app.py"], + cwd=SERVERS / "environment_servers" / "tool_counter", + env=env, + stdout=self.log.open("a"), + stderr=subprocess.STDOUT, + start_new_session=True, + ) + _wait_healthy(self.proc, f"{self.url}/health", self.log) + self.wait_for_workers() + + def workers(self) -> list[psutil.Process]: + # uvicorn's workers, not the multiprocessing helper it may also start. + children = psutil.Process(self.proc.pid).children() + return [child for child in children if "resource_tracker" not in " ".join(child.cmdline())] + + def wait_for_workers(self, *, exclude: int | None = None, timeout: float = 60) -> None: + """Wait until every live worker's private session socket accepts connections.""" + deadline = time.time() + timeout + while True: + live = [worker for worker in self.workers() if worker.pid != exclude] + accepting = 0 + for path in glob.glob(os.path.join(RUNTIME_ROOT, f"ng-{self.proc.pid}-*", "*.sock")): + if not re.fullmatch(r"[0-9a-f]{32}\.sock", os.path.basename(path)): + continue + with socket.socket(socket.AF_UNIX) as probe: + probe.settimeout(1) + with contextlib.suppress(OSError): + probe.connect(path) + accepting += 1 + if len(live) >= WORKERS and accepting >= WORKERS: + return + if time.time() > deadline: + raise TimeoutError(f"the workers did not start; see {self.log}") + time.sleep(0.2) + + def stop(self) -> None: + if self.proc is None: + return + with contextlib.suppress(ProcessLookupError): + os.killpg(self.proc.pid, signal.SIGTERM) + try: + self.proc.wait(timeout=10) + except subprocess.TimeoutExpired: + with contextlib.suppress(ProcessLookupError): + os.killpg(self.proc.pid, signal.SIGKILL) + self.proc.wait() + + +@pytest.fixture(scope="module") +def deployment(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Deployment]: + deployment = Deployment(tmp_path_factory.mktemp("episode_tools_routing")) + try: + deployment.start() + yield deployment + finally: + deployment.stop() + + +def _run_body(index: int, **task: Any) -> dict: + return { + "episode_id": {"rollout_id": f"{index}-0"}, + "task": {"task_id": {"taskset": "counter", "task_id": str(index)}, "task_input": task}, + } + + +def test_every_tool_call_runs_on_the_worker_that_runs_its_episode(deployment: Deployment) -> None: + async def run_all() -> list[dict]: + async with aiohttp.ClientSession(cookie_jar=aiohttp.DummyCookieJar()) as client: + + async def one(index: int) -> dict: + async with client.post(f"{deployment.url}/run", json=_run_body(index, calls=CALLS)) as response: + assert response.status == 200, await response.text() + return await response.json() + + return await asyncio.gather(*(one(index) for index in range(EPISODES))) + + results = [response["result"] for response in asyncio.run(run_all())] + + assert Counter(status for result in results for status in result["statuses"]) == Counter({200: EPISODES * CALLS}) + assert all(set(result["ran_on"]) == {result["owner_pid"]} for result in results) + assert all(result["total"] == CALLS for result in results) + # The episodes ran on several workers, so most of their calls reached another worker first. + assert len({result["owner_pid"] for result in results}) > 1 + + +def test_every_mcp_request_runs_on_the_worker_that_runs_its_episode(deployment: Deployment) -> None: + """Initialize, tools/list and tools/call each reach a random worker; each is routed by the token header.""" + + async def run_all() -> list[dict]: + async with aiohttp.ClientSession(cookie_jar=aiohttp.DummyCookieJar()) as client: + + async def one(index: int) -> dict: + body = _run_body(2 * 10**5 + index, calls=MCP_CALLS, transport="mcp") + async with client.post(f"{deployment.url}/run", json=body) as response: + assert response.status == 200, await response.text() + return await response.json() + + return await asyncio.gather(*(one(index) for index in range(MCP_EPISODES))) + + results = [response["result"] for response in asyncio.run(run_all())] + + # A tools/list that reached another worker would list nothing. + assert all(result["listed"] == ["add"] for result in results) + assert Counter(status for result in results for status in result["statuses"]) == Counter( + {200: MCP_EPISODES * MCP_CALLS} + ) + assert all(set(result["ran_on"]) == {result["owner_pid"]} for result in results) + assert len({result["owner_pid"] for result in results}) > 1 + + +def test_a_tool_call_for_an_exited_workers_episode_gets_410(deployment: Deployment, tmp_path: Path) -> None: + grant_file = tmp_path / "grant.json" + + async def scenario() -> tuple[int, dict]: + async with aiohttp.ClientSession(cookie_jar=aiohttp.DummyCookieJar()) as client: + # An episode that keeps serving its tools until its worker is killed. + episode = asyncio.create_task( + client.post( + f"{deployment.url}/run", json=_run_body(10**6, grant_file=str(grant_file), hold_seconds=120) + ) + ) + while not grant_file.exists(): + await asyncio.sleep(0.05) + grant = json.loads(grant_file.read_text()) + owner = psutil.Process(grant["pid"]) + owner.send_signal(signal.SIGKILL) + owner.wait(timeout=60) + with contextlib.suppress(aiohttp.ClientError): + await episode + await asyncio.to_thread(deployment.wait_for_workers, exclude=grant["pid"]) + # A new client, so no pooled connection to the killed worker is reused. + async with aiohttp.ClientSession(cookie_jar=aiohttp.DummyCookieJar()) as client: + async with client.post(f"{grant['base_url']}/add", json={"amount": 1}, headers=grant["headers"]) as reply: + return reply.status, await reply.json() + + status, body = asyncio.run(scenario()) + + assert status == 410 and "no longer running" in body["detail"] diff --git a/tests/unit_tests/server_runtime/test_base_environment_server_with_tools.py b/tests/unit_tests/server_runtime/test_base_environment_server_with_tools.py index 12eb3d50166..c3fa930227d 100644 --- a/tests/unit_tests/server_runtime/test_base_environment_server_with_tools.py +++ b/tests/unit_tests/server_runtime/test_base_environment_server_with_tools.py @@ -342,11 +342,14 @@ async def agent(server: _Server, counter: _Counter, served: ServedTools) -> None assert response.failure is not None and "already serves its tools" in response.failure.failure_reason -def test_several_workers_are_refused_while_calls_reach_any_worker() -> None: - server, _ = _server(num_workers=2) +def test_several_workers_route_each_episode_to_the_worker_that_runs_it() -> None: + """Routing is tested in test_episode_tools_session_routing.py; here, the server takes several workers.""" + server, _ = _server(num_workers=4) - with pytest.raises(ValueError, match="needs num_workers=1"): - refuse_unsupported_num_workers(server) + refuse_unsupported_num_workers(server) + + assert server.routes_sessions_to_owner + assert server.owned_session_routing()["session_header"] == EPISODE_TOKEN_HEADER class _Sandbox: diff --git a/tests/unit_tests/server_runtime/test_episode_tools_session_routing.py b/tests/unit_tests/server_runtime/test_episode_tools_session_routing.py new file mode 100644 index 00000000000..0fd4ebed703 --- /dev/null +++ b/tests/unit_tests/server_runtime/test_episode_tools_session_routing.py @@ -0,0 +1,240 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tool calls to a ``BaseEnvironmentServerWithTools`` reach the worker that runs the episode. + +In-process workers share a real socket directory, as the workers of one server do. +""" + +import hashlib +import os +import shutil +import tempfile +from collections.abc import Iterator +from typing import Optional +from unittest.mock import MagicMock + +import pytest +from fastapi import Body +from fastapi.testclient import TestClient +from httpx import Response +from pydantic import BaseModel + +from nemo_gym.client.server_client import ServerClient +from nemo_gym.config.config_types import AggregateMetrics, AggregateMetricsRequest +from nemo_gym.core.episode_types import BaseEpisodeRequest, BaseEpisodeResponse, EpisodeId, MaterializedTask, TaskId +from nemo_gym.server_runtime.base_environment_server import ( + BaseEnvironmentServerWithTools, + BaseEnvironmentServerWithToolsConfig, + CleanupContext, + ServedTools, +) +from nemo_gym.server_runtime.episode_tokens import EPISODE_TOKEN_HEADER, EpisodeTokens +from nemo_gym.server_runtime.episode_tools import TypedToolset, episode_tool +from nemo_gym.server_runtime.request_handling.session_routing import SESSION_OWNERS_DIR, install_session_routing + + +class _TaskInput(BaseModel): + pass + + +class _Request(BaseEpisodeRequest[_TaskInput]): + pass + + +class _Response(BaseEpisodeResponse[str]): + pass + + +class _Amount(BaseModel): + amount: int + + +class _Counter: + def __init__(self) -> None: + self.total = 0 + + @episode_tool + async def add(self, arguments: _Amount) -> dict[str, int]: + """Add an amount to the episode's total.""" + self.total += arguments.amount + return {"total": self.total} + + +class _Server(BaseEnvironmentServerWithTools[_Request, _Response]): + request_model = _Request + response_model = _Response + + async def run_episode(self, request: _Request, cleanup: CleanupContext) -> _Response: + raise NotImplementedError("the tests serve episodes directly") + + async def aggregate_metrics(self, body: AggregateMetricsRequest = Body()) -> AggregateMetrics: + return AggregateMetrics() + + +def _request(rollout_id: str, attempt: int = 0) -> _Request: + return _Request( + episode_id=EpisodeId(rollout_id=rollout_id, attempt=attempt), + task=MaterializedTask(task_id=TaskId(taskset="t", task_id="0"), task_input=_TaskInput()), + ) + + +class _Worker: + """One worker of a ``BaseEnvironmentServerWithTools``, wired as run_webserver wires it.""" + + def __init__(self, socket_dir: Optional[str]) -> None: + config = BaseEnvironmentServerWithToolsConfig( + name="counter_environment", + host="127.0.0.1", + port=1234, + entrypoint="app.py", + cleanup_timeout_seconds=1, + num_workers=2 if socket_dir is not None else 1, + ) + server_client = MagicMock(spec=ServerClient) + server_client._resolve_base_url.return_value = "http://127.0.0.1:1234" + self.server = _Server(config=config, server_client=server_client) + self.app = self.server.setup_webserver() + key = self.server.get_session_middleware_key() + self.routing_id = install_session_routing( + self.app, + socket_dir=socket_dir, + session_cookie=key, + secret_key=key, + **self.server.owned_session_routing(), + ) + self.client = TestClient(self.app, raise_server_exceptions=False) + + def serve(self, rollout_id: str, attempt: int = 0) -> tuple[ServedTools, CleanupContext, _Counter]: + """Serve one episode's tools on this worker, in its event loop, as its run() would.""" + request = _request(rollout_id, attempt) + cleanup = CleanupContext(episode_id=request.episode_id, cleanup_timeout_seconds=1) + counter = _Counter() + served = self.client.portal.call(self.server.serve_tools, request, TypedToolset(counter), cleanup) + return served, cleanup, counter + + def end(self, cleanup: CleanupContext) -> None: + self.client.portal.call(cleanup.aclose) + + def add(self, token: str, amount: int = 1) -> Response: + return self.client.post("/add", json={"amount": amount}, headers={EPISODE_TOKEN_HEADER: token}) + + +def _token(served: ServedTools) -> str: + (grant,) = served.tool_accesses + return grant.headers[EPISODE_TOKEN_HEADER] + + +def _owner(socket_dir: str, session_id: str) -> Optional[str]: + path = os.path.join(socket_dir, SESSION_OWNERS_DIR, hashlib.sha256(session_id.encode("utf-8")).hexdigest()) + try: + with open(path) as owner: + return owner.read() + except FileNotFoundError: + return None + + +@pytest.fixture +def socket_dir() -> Iterator[str]: + # Under /tmp, as in production: AF_UNIX paths are limited to about 100 bytes. + path = tempfile.mkdtemp(prefix="ng-test-", dir="/tmp") + yield path + shutil.rmtree(path, ignore_errors=True) + + +@pytest.fixture +def workers(socket_dir: str) -> Iterator[tuple[_Worker, _Worker]]: + a, b = _Worker(socket_dir), _Worker(socket_dir) + with a.client, b.client: + yield a, b + + +def test_a_tool_call_that_reaches_another_worker_runs_on_the_episodes_worker( + socket_dir: str, workers: tuple[_Worker, _Worker] +) -> None: + a, b = workers + served, _, counter = a.serve("3-0") + + first, second = b.add(_token(served), 2), b.add(_token(served), 3) + + assert (first.status_code, first.json(), second.json()) == (200, {"total": 2}, {"total": 5}) + assert counter.total == 5 and [call.name for call in served.calls] == ["add", "add"] + assert _owner(socket_dir, "3-0") == a.routing_id + + +def test_routing_reads_the_episode_and_only_the_episodes_worker_checks_the_nonce( + workers: tuple[_Worker, _Worker], +) -> None: + a, b = workers + served, _, counter = a.serve("3-0") + forged, _ = EpisodeTokens(a.server.get_session_middleware_key()).mint(EpisodeId(rollout_id="3-0")) + + reply = b.add(forged) + + # Forwarded to the episode's worker, which refuses the nonce of another grant. + assert reply.status_code == 403 and "not granted for this episode" in reply.json()["detail"] + assert counter.total == 0 and served.calls == [] + + +def test_a_token_that_names_no_episode_is_refused_where_it_lands(workers: tuple[_Worker, _Worker]) -> None: + a, b = workers + a.serve("3-0") + + assert b.add("garbage").status_code == 401 + assert b.client.post("/add", json={"amount": 1}).status_code == 401 + + +def test_each_attempt_of_an_episode_has_its_own_owner(socket_dir: str, workers: tuple[_Worker, _Worker]) -> None: + a, b = workers + first_attempt, _, first_counter = a.serve("3-0") + retry, _, retry_counter = b.serve("3-0", attempt=1) + + a.add(_token(retry), 10) + b.add(_token(first_attempt), 1) + + assert (first_counter.total, retry_counter.total) == (1, 10) + assert (_owner(socket_dir, "3-0"), _owner(socket_dir, "3-0-a1")) == (a.routing_id, b.routing_id) + + +def test_an_episode_running_on_another_worker_cannot_serve_its_tools_again(workers: tuple[_Worker, _Worker]) -> None: + a, b = workers + a.serve("3-0") + + with pytest.raises(ValueError, match="already serves its tools on another worker"): + b.serve("3-0") + + +def test_ownership_ends_with_the_episode(socket_dir: str, workers: tuple[_Worker, _Worker]) -> None: + a, b = workers + served, cleanup, _ = a.serve("3-0") + a.end(cleanup) + + reply = b.add(_token(served)) + + assert _owner(socket_dir, "3-0") is None + # Nobody owns the episode, so the worker it reached answers that it does not run it. + assert reply.status_code == 410 and "not running on this server" in reply.json()["detail"] + + +def test_a_tool_call_for_an_exited_workers_episode_gets_410(socket_dir: str) -> None: + a, b = _Worker(socket_dir), _Worker(socket_dir) + with b.client: + a.client.__enter__() + served, _, _ = a.serve("3-0") + # The worker exits: its private socket closes, and it stays the owner. + a.client.__exit__(None, None, None) + + reply = b.add(_token(served)) + + assert _owner(socket_dir, "3-0") == a.routing_id + assert reply.status_code == 410 and "no longer running" in reply.json()["detail"] + + +def test_one_worker_serves_tool_calls_without_owners() -> None: + worker = _Worker(socket_dir=None) + with worker.client: + served, _, counter = worker.serve("3-0") + + reply = worker.add(_token(served), 4) + + assert worker.routing_id is None + assert reply.json() == {"total": 4} and counter.total == 4