diff --git a/.gitignore b/.gitignore index e051977..62a8d0f 100644 --- a/.gitignore +++ b/.gitignore @@ -15,6 +15,10 @@ build/ .env .wandb.env +# Logs and experiment tracking +wandb/ +*.log + # Local slime backend config (contains ARN/bucket; commit config.yaml.example only) src/agentcore_rl_toolkit/backends/slime/examples/**/config.yaml diff --git a/AGENTS.md b/AGENTS.md index 66f187f..627195c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -80,6 +80,7 @@ agentcore-rl-toolkit/ │ │ ├── render.py # Renderer protocol; HfTemplateRenderer, TinkerRenderer │ │ ├── parsing.py # tool/reasoning output parsing (sglang optional) │ │ ├── gateway.py # RolloutGateway — assembles the serving unit +│ │ ├── server.py # ThreadedGatewayServer — serve the gateway from sync trainers │ │ ├── adapters/ # OpenAI + Anthropic wire protocol adapters │ │ └── sampling_backends/ # SamplingBackend impls (vLLM/SGLang HTTP, Tinker SDK) │ └── backends/experimental/verl/ # Experimental verl backend on the rollout gateway @@ -288,6 +289,7 @@ sample-only backends like Tinker, which cannot render themselves. | `SamplingBackend` | `sampling_backends/` | The one per-engine seam: `token_ids -> token_ids + logprobs` as a `TurnRecord`. Impls: `VllmHttpBackend`, `SglangHttpBackend`, `TinkerSdkBackend`. Placement rule: engine seams for independently reachable inference services (HTTP endpoints, hosted SDKs like Tinker) live here; seams over trainer-internal handles (e.g. `VerlSamplingBackend` over verl's Ray-based `LLMServerClient`) live with that trainer's integration under `backends/`. | | Adapters | `adapters/` | Wire-protocol translation: `OpenAIAdapter` (`/v1/chat/completions`), `AnthropicAdapter` (`/v1/messages`). An agent drives the gateway in its *native* protocol unmodified (just point `base_url` at it); both normalize to one canonical message form and share one `TrajectoryManager`. | | `RolloutGateway` | `gateway.py` | Assembles tokenizer + renderer + backend + adapters onto one aiohttp app sharing one `TrajectoryManager`. Session identity rides in the api-key / Bearer slot; `base_url` is a fixed gateway address (no per-session URLs). | +| `ThreadedGatewayServer` | `server.py` | Serves an assembled gateway on a background thread with its own event loop — the deployment shape for synchronous trainers (slime, verl). Async trainers can mount `gateway.app` into their own loop instead. | **Session model.** A session id (in the Bearer slot) keys one trajectory tree. `gateway.create_session(sid)` → agent turns are captured → `gateway.finish_session(sid)` @@ -383,6 +385,7 @@ upstream before re-syncing, diff the source file against the baseline commit bel | `rollout_gateway/trajectory.py` | `slime/agent/trajectory.py` | `90c212b5` | | `rollout_gateway/adapters/{common,openai,anthropic}.py` | `slime/agent/adapters/` | `90c212b5` | | `rollout_gateway/parsing.py` | `slime/agent/parsing.py` | `90c212b5` | +| `rollout_gateway/server.py` | `slime/agent/aiohttp_threaded.py` | `fa3c990a` | Re-sync workflow: `git -C diff 90c212b5..HEAD -- slime/agent/` shows upstream changes since the lift. Our copies are intentionally modified (torch-free; `Sample` → diff --git a/NOTICE b/NOTICE index d58e96e..7b4fcd8 100644 --- a/NOTICE +++ b/NOTICE @@ -13,3 +13,5 @@ slime, notably: - adapters/ — OpenAIAdapter, AnthropicAdapter, and the shared BaseAdapter pipeline, adapted from slime/agent/adapters/. - parsing.py — output-parsing helpers, adapted from slime/agent/parsing.py. + - server.py — background-thread serving mechanics and FilteredAccessLogger, + adapted from slime/agent/aiohttp_threaded.py. diff --git a/examples/strands_math_agent/rl_app.py b/examples/strands_math_agent/rl_app.py index d64ca03..b6a0c89 100644 --- a/examples/strands_math_agent/rl_app.py +++ b/examples/strands_math_agent/rl_app.py @@ -40,6 +40,9 @@ def invoke_agent(payload: dict, context): base_url = payload["_rollout"]["base_url"] model_id = payload["_rollout"]["model_id"] params = payload["_rollout"].get("sampling_params", {}) + # During training the rollout gateway keys the session off the api-key slot; + # "EMPTY" (the vLLM convention) is fine for plain evaluation endpoints. + api_key = payload["_rollout"].get("api_key", "EMPTY") # The ACR session id doubles as the trajectory-capture session key: rollout # gateways read it from the api-key slot. "EMPTY" for local runs and @@ -65,8 +68,13 @@ def invoke_agent(payload: dict, context): response = agent(user_input) - # Compute rewards - rewards = reward_fn(response_text=response.message["content"][0]["text"], ground_truth=answer) + # Compute rewards. Join the text blocks instead of indexing content[0]: a + # reasoning-parsing server (vllm/sglang --reasoning-parser, or the rollout + # gateway) returns reasoning_content, which Strands surfaces as a leading + # reasoningContent block before the text block. + content = response.message.get("content") or [] + response_text = "".join(block["text"] for block in content if "text" in block) + rewards = reward_fn(response_text=response_text, ground_truth=answer) return {"rewards": rewards} diff --git a/pyproject.toml b/pyproject.toml index 35092d3..22e21a1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -212,3 +212,6 @@ select = [ [tool.ruff.lint.isort] known-first-party = ["agentcore_rl_toolkit"] + +[tool.pytest.ini_options] +addopts = "--ignore=tests/rollout_gateway/integration_test" diff --git a/src/agentcore_rl_toolkit/rollout_gateway/__init__.py b/src/agentcore_rl_toolkit/rollout_gateway/__init__.py index 235cf37..8cbf23c 100644 --- a/src/agentcore_rl_toolkit/rollout_gateway/__init__.py +++ b/src/agentcore_rl_toolkit/rollout_gateway/__init__.py @@ -10,9 +10,9 @@ imports with nothing beyond the stdlib — torch-free and aiohttp-free. (``HfTemplateRenderer`` needs a HF tokenizer only when *used*, not to import; sglang parsers and tinker are lazy/optional.) -- ``RolloutGateway`` and the HTTP adapters require ``aiohttp`` (the ``[gateway]`` extra), - so ``RolloutGateway`` is exposed lazily via ``__getattr__`` — importing this package - never requires aiohttp. +- ``RolloutGateway``, ``ThreadedGatewayServer``, and the HTTP adapters require + ``aiohttp`` (the ``[gateway]`` extra), so both are exposed lazily via + ``__getattr__`` — importing this package never requires aiohttp. """ from .render import HfTemplateRenderer, ParsedOutput, Renderer @@ -32,6 +32,7 @@ "RolloutGateway", "SamplingBackend", "Status", + "ThreadedGatewayServer", "TraceRecord", "TrajectoryManager", "TurnRecord", @@ -39,9 +40,13 @@ def __getattr__(name: str): - # RolloutGateway pulls in the aiohttp adapters; keep it off the plain import path. + # These pull in aiohttp; keep them off the plain import path. if name == "RolloutGateway": from .gateway import RolloutGateway return RolloutGateway + if name == "ThreadedGatewayServer": + from .server import ThreadedGatewayServer + + return ThreadedGatewayServer raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/src/agentcore_rl_toolkit/rollout_gateway/server.py b/src/agentcore_rl_toolkit/rollout_gateway/server.py new file mode 100644 index 0000000..32d6bb8 --- /dev/null +++ b/src/agentcore_rl_toolkit/rollout_gateway/server.py @@ -0,0 +1,118 @@ +"""Run a :class:`RolloutGateway` on a background thread, for synchronous trainers. + +A sync trainer (slime, verl) blocks its thread waiting for episode results, but the +gateway must keep answering the agent's LLM calls the whole time. Served on the same +thread, the block would freeze the event loop and deadlock the episode — so this +class serves ``gateway.app`` on its own daemon thread with a long-lived loop. The +session API (``create_session`` / ``finish_session`` / ``drop_session``) stays safe +to call from the trainer thread. Async trainers don't need this: mount +``gateway.app`` on your own loop and ``await`` instead. + +The serving mechanics are adapted from slime's ``slime/agent/aiohttp_threaded.py`` +(baseline commit ``fa3c990a``; see NOTICE). + +Requires ``aiohttp`` (the ``[gateway]`` extra), like :class:`RolloutGateway` itself. +""" + +from __future__ import annotations + +import asyncio +import logging +import threading + +from aiohttp import web +from aiohttp.web_log import AccessLogger + +logger = logging.getLogger(__name__) + + +class FilteredAccessLogger(AccessLogger): + """Log only failures and slow requests; healthy fast traffic is noise.""" + + SLOW_THRESHOLD_SEC = 120.0 + + def log(self, request, response, time): + if request.method == "HEAD": + return + if response.status == 200 and time <= self.SLOW_THRESHOLD_SEC: + return + super().log(request, response, time) + + +class ThreadedGatewayServer: + """Serve an assembled :class:`RolloutGateway` on a background thread. + + ``start()`` blocks until the server is bound (or raises if binding fails), + so callers can hand out :attr:`base_url` immediately afterwards. ``port=0`` + binds an OS-assigned port, reflected in :attr:`port`/:attr:`base_url` after + ``start()``. + + Session identity rides in the api-key / Bearer slot of each request, so all + sessions share the single fixed :attr:`base_url`. + """ + + def __init__(self, gateway, *, host: str, port: int = 0, startup_timeout: float = 120.0): + self.gateway = gateway + self.host = host + self.port = port + self.startup_timeout = startup_timeout + self._loop: asyncio.AbstractEventLoop | None = None + self._thread: threading.Thread | None = None + self._runner: web.AppRunner | None = None + + @property + def base_url(self) -> str: + """Fixed OpenAI-compatible ``base_url`` shared by all sessions.""" + return f"http://{self.host}:{self.port}/v1" + + def start(self) -> None: + started = threading.Event() + startup_error: list[BaseException] = [] + + def _serve() -> None: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + # handler_cancellation=True: a client disconnect cancels the + # in-flight handler coroutine, so an agent that dies mid-request + # doesn't leave an orphaned generate call in the sampling backend. + runner = web.AppRunner( + self.gateway.app, + handler_cancellation=True, + access_log_class=FilteredAccessLogger, + ) + loop.run_until_complete(runner.setup()) + site = web.TCPSite(runner, host=self.host, port=self.port) + loop.run_until_complete(site.start()) + for sock in site._server.sockets: # resolve OS-assigned port when port=0 + self.port = sock.getsockname()[1] + break + self._loop = loop + self._runner = runner + started.set() + loop.run_forever() + except BaseException as e: # surface bind/setup failures to the caller + startup_error.append(e) + started.set() + + self._thread = threading.Thread(target=_serve, name="rollout-gateway", daemon=True) + self._thread.start() + if not started.wait(timeout=self.startup_timeout): + raise TimeoutError(f"Rollout gateway did not start within {self.startup_timeout}s") + if startup_error: + raise RuntimeError(f"Rollout gateway failed to start on {self.host}:{self.port}") from startup_error[0] + logger.info("Rollout gateway serving at %s", self.base_url) + + def shutdown(self) -> None: + if self._loop is None: + return + # Clean up on the live loop first (graceful connection shutdown + + # release of the listening socket), then stop the loop. + try: + fut = asyncio.run_coroutine_threadsafe(self._runner.cleanup(), self._loop) + fut.result(timeout=10) + except Exception: + pass + self._loop.call_soon_threadsafe(self._loop.stop) + self._thread.join(timeout=10) + self._loop = None diff --git a/tests/rollout_gateway/integration_test/sglang/test_e2e_qwen_sglang.py b/tests/rollout_gateway/integration_test/sglang/test_e2e_qwen_sglang.py index 63d2868..5201cf4 100644 --- a/tests/rollout_gateway/integration_test/sglang/test_e2e_qwen_sglang.py +++ b/tests/rollout_gateway/integration_test/sglang/test_e2e_qwen_sglang.py @@ -8,16 +8,22 @@ import time import urllib.error import urllib.request -from contextlib import asynccontextmanager +from contextlib import contextmanager import pytest -from agentcore_rl_toolkit.rollout_gateway import BaseTrace, HfTemplateRenderer, RolloutGateway +from agentcore_rl_toolkit.rollout_gateway import ( + BaseTrace, + HfTemplateRenderer, + RolloutGateway, + ThreadedGatewayServer, +) SGLANG_URL = os.environ.get("E2E_SGLANG_URL", "http://localhost:30000") MODEL = os.environ.get("E2E_MODEL", "Qwen/Qwen2.5-0.5B-Instruct") MAX_NEW_TOKENS = 512 SERVER_LOG = "/tmp/e2e_sglang_server.log" +GATEWAY_PORT = 9090 HAVE_DEPS = all(importlib.util.find_spec(m) for m in ("openai", "sglang")) @@ -105,19 +111,16 @@ def make_gateway(): return gateway, backend, renderer, _TOKENIZER -@asynccontextmanager -async def serve(gateway: RolloutGateway): - from aiohttp import web - - runner = web.AppRunner(gateway.app) - await runner.setup() - site = web.TCPSite(runner, "127.0.0.1", 0) - await site.start() - port = site._server.sockets[0].getsockname()[1] +@contextmanager +def serve(gateway: RolloutGateway): + """Serve the gateway exactly as a sync trainer would: on a + ThreadedGatewayServer, off the test's own event loop.""" + gw_server = ThreadedGatewayServer(gateway, host="127.0.0.1", port=GATEWAY_PORT) + gw_server.start() try: - yield f"http://127.0.0.1:{port}" + yield f"http://127.0.0.1:{GATEWAY_PORT}" finally: - await runner.cleanup() + gw_server.shutdown() def _trained(rec) -> list[int]: @@ -138,7 +141,7 @@ async def test_single_turn_token_exact(server): gateway.create_session(sid, sampling_defaults={"temperature": 0.0}) messages = [{"role": "user", "content": "What is 2+2? Reply with just the number. /no_think"}] - async with serve(gateway) as base_url: + with serve(gateway) as base_url: import openai client = openai.AsyncOpenAI(base_url=f"{base_url}/v1", api_key=sid) @@ -179,7 +182,7 @@ async def test_multi_turn_drift_healed_into_one_record(server): sid = "e2e:multi" gateway.create_session(sid, sampling_defaults={"temperature": 0.0}) - async with serve(gateway) as base_url: + with serve(gateway) as base_url: import openai client = openai.AsyncOpenAI(base_url=f"{base_url}/v1", api_key=sid) @@ -232,7 +235,7 @@ async def test_rewritten_echo_trains_only_final_turn(server): sid = "e2e:rewrite" gateway.create_session(sid, sampling_defaults={"temperature": 0.0}) - async with serve(gateway) as base_url: + with serve(gateway) as base_url: import openai client = openai.AsyncOpenAI(base_url=f"{base_url}/v1", api_key=sid) @@ -278,7 +281,7 @@ async def test_stateless_turns_fork_into_separate_records(server): "What is 3*3? Reply with just the number. /no_think", ] - async with serve(gateway) as base_url: + with serve(gateway) as base_url: import openai client = openai.AsyncOpenAI(base_url=f"{base_url}/v1", api_key=sid) diff --git a/tests/rollout_gateway/integration_test/vllm/test_e2e_qwen_vllm.py b/tests/rollout_gateway/integration_test/vllm/test_e2e_qwen_vllm.py index 69c7b77..9ccb092 100644 --- a/tests/rollout_gateway/integration_test/vllm/test_e2e_qwen_vllm.py +++ b/tests/rollout_gateway/integration_test/vllm/test_e2e_qwen_vllm.py @@ -8,17 +8,23 @@ import time import urllib.error import urllib.request -from contextlib import asynccontextmanager +from contextlib import contextmanager from pathlib import Path import pytest -from agentcore_rl_toolkit.rollout_gateway import BaseTrace, HfTemplateRenderer, RolloutGateway +from agentcore_rl_toolkit.rollout_gateway import ( + BaseTrace, + HfTemplateRenderer, + RolloutGateway, + ThreadedGatewayServer, +) VLLM_URL = os.environ.get("E2E_VLLM_URL", "http://localhost:8000") MODEL = os.environ.get("E2E_MODEL", "Qwen/Qwen2.5-0.5B-Instruct") MAX_NEW_TOKENS = 512 SERVER_LOG = "/tmp/e2e_vllm_server.log" +GATEWAY_PORT = 9090 HAVE_DEPS = all(importlib.util.find_spec(m) for m in ("openai", "anthropic", "vllm")) @@ -106,19 +112,16 @@ def make_gateway(): return gateway, backend, renderer, _TOKENIZER -@asynccontextmanager -async def serve(gateway: RolloutGateway): - from aiohttp import web - - runner = web.AppRunner(gateway.app) - await runner.setup() - site = web.TCPSite(runner, "127.0.0.1", 0) - await site.start() - port = site._server.sockets[0].getsockname()[1] +@contextmanager +def serve(gateway: RolloutGateway): + """Serve the gateway exactly as a sync trainer would: on a + ThreadedGatewayServer, off the test's own event loop.""" + gw_server = ThreadedGatewayServer(gateway, host="127.0.0.1", port=GATEWAY_PORT) + gw_server.start() try: - yield f"http://127.0.0.1:{port}" + yield f"http://127.0.0.1:{GATEWAY_PORT}" finally: - await runner.cleanup() + gw_server.shutdown() def _trained(rec) -> list[int]: @@ -139,7 +142,7 @@ async def test_single_turn_token_exact(server): gateway.create_session(sid, sampling_defaults={"temperature": 0.0}) messages = [{"role": "user", "content": "What is 2+2? Reply with just the number. /no_think"}] - async with serve(gateway) as base_url: + with serve(gateway) as base_url: import openai client = openai.AsyncOpenAI(base_url=f"{base_url}/v1", api_key=sid) @@ -180,7 +183,7 @@ async def test_multi_turn_drift_healed_into_one_record(server): sid = "e2e:multi" gateway.create_session(sid, sampling_defaults={"temperature": 0.0}) - async with serve(gateway) as base_url: + with serve(gateway) as base_url: import openai client = openai.AsyncOpenAI(base_url=f"{base_url}/v1", api_key=sid) @@ -233,7 +236,7 @@ async def test_rewritten_echo_trains_only_final_turn(server): sid = "e2e:rewrite" gateway.create_session(sid, sampling_defaults={"temperature": 0.0}) - async with serve(gateway) as base_url: + with serve(gateway) as base_url: import openai client = openai.AsyncOpenAI(base_url=f"{base_url}/v1", api_key=sid) @@ -279,7 +282,7 @@ async def test_stateless_turns_fork_into_separate_records(server): "What is 3*3? Reply with just the number. /no_think", ] - async with serve(gateway) as base_url: + with serve(gateway) as base_url: import openai client = openai.AsyncOpenAI(base_url=f"{base_url}/v1", api_key=sid) @@ -312,7 +315,7 @@ async def test_anthropic_client_multi_turn_token_exact(server): sid = "e2e:anthropic" gateway.create_session(sid, sampling_defaults={"temperature": 0.0}) - async with serve(gateway) as base_url: + with serve(gateway) as base_url: import anthropic # the anthropic SDK sends api_key as X-Api-Key; the adapter maps it to the sid @@ -355,7 +358,7 @@ async def test_mixed_protocol_clients_share_one_trajectory(server): sid = "e2e:mixed-sdk" gateway.create_session(sid, sampling_defaults={"temperature": 0.0}) - async with serve(gateway) as base_url: + with serve(gateway) as base_url: import anthropic import openai