diff --git a/verifiers/v1/clients/train.py b/verifiers/v1/clients/train.py index 3fad296beb..08a6ef1087 100644 --- a/verifiers/v1/clients/train.py +++ b/verifiers/v1/clients/train.py @@ -299,8 +299,9 @@ async def acquire(self) -> AsyncIterator[RendererSlot]: class TrainClient(Client): """Renders prompts to token ids and calls a vLLM `/inference/v1/generate` engine. - One client per rollout: it owns its engine connection and takes a slot on the shared - `ElasticRendererPool` for each turn.""" + Owned by the interception server and shared by the rollouts it multiplexes: they reuse + its engine connection pool, and each turn takes a slot on the shared + `ElasticRendererPool`.""" def __init__(self, config: TrainClientConfig) -> None: self.config = config diff --git a/verifiers/v1/configs/client.py b/verifiers/v1/configs/client.py index fea86c512c..a4a9b9639e 100644 --- a/verifiers/v1/configs/client.py +++ b/verifiers/v1/configs/client.py @@ -1,8 +1,9 @@ """Client configs: describe an OpenAI-compatible endpoint. A `BaseClientConfig` is an OpenAI-compatible endpoint (base_url + API-key env var -+ extra headers); `clients.resolve_client` turns one into a live `Client`, and every -rollout does so for itself. The default Prime endpoint, API key, and team fall back to ++ extra headers); `clients.resolve_client` turns one into a live `Client` — the +interception server builds one per distinct config and shares it across the rollouts +it multiplexes. The default Prime endpoint, API key, and team fall back to the active Prime CLI config, so direct `uv run eval` calls behave like `prime eval`. Both the eval entrypoint (its model client) and in-env LLM calls (e.g. a judge reward) build clients from these. `ClientConfig` is the CLI-selectable discriminated union diff --git a/verifiers/v1/gepa/runner.py b/verifiers/v1/gepa/runner.py index e0cfafba9c..10b3e031cb 100644 --- a/verifiers/v1/gepa/runner.py +++ b/verifiers/v1/gepa/runner.py @@ -73,7 +73,7 @@ async def on_complete(episode: Episode) -> None: await append_episode(run_dir, episode, write_lock) try: - # The endpoint stays config: each rollout builds and closes its own client. + # The endpoint stays config: the interception server builds the live client. ctx = ModelContext( client=config.client, model=config.model, sampling=config.sampling ) diff --git a/verifiers/v1/interception/server.py b/verifiers/v1/interception/server.py index 0c10bc190f..644f28a700 100644 --- a/verifiers/v1/interception/server.py +++ b/verifiers/v1/interception/server.py @@ -9,7 +9,9 @@ One server multiplexes many rollouts: each rollout registers separate model and state capabilities, and the server routes each to the right session. So N rollouts need one server (and, behind a remote runtime, one tunnel) per pool member rather than one each — -see `interception.pool`. +see `interception.pool`. The server also owns the model clients (one per distinct endpoint +config, assigned to each session at register and closed with the server), so its rollouts +share one bounded keepalive connection pool upstream instead of churning per-rollout TCP. The server is a pure model boundary: one request, one turn — refusal checks (limits, `@stop`s), the model call, the graph commit, retry atomicity. A run's user exchange @@ -34,6 +36,8 @@ from pydantic_core import PydanticSerializationError, from_json, to_json from verifiers.v1 import graph +from verifiers.v1.clients import Client, resolve_client +from verifiers.v1.configs.client import BaseClientConfig from verifiers.v1.dialects import DIALECTS, Dialect from verifiers.v1.dialects.base import is_sse_done_event from verifiers.v1.errors import ( @@ -130,6 +134,7 @@ def __init__( ) -> None: super().__init__() self.sessions: dict[str, RolloutSession] = {} + self.clients: dict[str, Client] = {} self.state_sessions: dict[str, RolloutSession] = {} self.state_routes: dict[str, RolloutSession] = {} self.state_service_secrets = frozenset(state_service_secrets) @@ -147,8 +152,22 @@ def load(self) -> int: """Rollouts currently registered — what the pools balance on.""" return len(self.sessions) + def _client(self, config: BaseClientConfig) -> Client: + """The server-owned client for `config` — one per distinct endpoint config, shared + by every session registered under it, so the rollouts this server multiplexes reuse + one bounded keepalive pool instead of each opening (and tearing down) their own + connections. Closed with the server.""" + key = config.model_dump_json() + client = self.clients.get(key) + if client is None: + client = self.clients[key] = resolve_client(config) + self.stack.push_async_callback(client.close) + return client + def register(self, session: RolloutSession) -> tuple[str, str]: - """Register separate capabilities for model inference and private task state.""" + """Register separate capabilities for model inference and private task state, and + assign the session its server-owned model client.""" + session.client = self._client(session.ctx.client) model_secret = secrets.token_urlsafe(16) state_secret = secrets.token_urlsafe(16) self.sessions[model_secret] = session diff --git a/verifiers/v1/rollout.py b/verifiers/v1/rollout.py index 285d3c6ca5..50a3bc5b20 100644 --- a/verifiers/v1/rollout.py +++ b/verifiers/v1/rollout.py @@ -8,7 +8,7 @@ from contextlib import AsyncExitStack from dataclasses import dataclass -from verifiers.v1.clients import ModelContext, resolve_client +from verifiers.v1.clients import ModelContext from verifiers.v1.configs.agent import AgentConfig from verifiers.v1.errors import ( HarnessError, @@ -91,9 +91,8 @@ def __init__( ) if on_trace is not None: on_trace(self.trace) - self.client = resolve_client(ctx.client) self._session = RolloutSession( - ctx, self.client, self.trace, discover_decorated(task, "stop"), limits + ctx, self.trace, discover_decorated(task, "stop"), limits ) self._stack = AsyncExitStack() self._failed = False @@ -318,8 +317,8 @@ async def step(self, messages: Messages | None = None) -> bool: return self.ok and trace.num_turns > turns_before async def abort(self) -> None: - """Free everything this run holds — the entered servers, its client, and an - owned runtime — without finalizing or scoring: the escape path when an exception + """Free everything this run holds — the entered servers and an owned + runtime — without finalizing or scoring: the escape path when an exception (a cancellation mid-setup, a lifetime bug raised to the caller) means the driver will never reach `close()`. Safe after a partial `close()`.""" self._closed = True @@ -328,8 +327,6 @@ async def abort(self) -> None: await self._harness_session.close() with contextlib.suppress(Exception): await self._stack.aclose() - with contextlib.suppress(Exception): - await self.client.close() if self.runtime is not None: with contextlib.suppress(Exception): await self.harness.cleanup(self.trace, self.runtime) @@ -424,12 +421,6 @@ async def close(self) -> Trace: logger.warning( "runtime teardown failed (rollout %s)", trace.id, exc_info=True ) - try: - await self.client.close() - except Exception: - logger.warning( - "client teardown failed (rollout %s)", trace.id, exc_info=True - ) logger.info( "rollout done: id=%s task=%s reward=%.3f turns=%d stop=%s", trace.id, diff --git a/verifiers/v1/session.py b/verifiers/v1/session.py index 322804c723..8251564963 100644 --- a/verifiers/v1/session.py +++ b/verifiers/v1/session.py @@ -2,8 +2,8 @@ One `RolloutSession` per rollout, registered on an interception server under the rollout's secret. The rollout constructs it (model ctx, trace, task `@stop`s, limits) and the server -drives it: routes each intercepted model call to it, runs `refused()` before each turn, -and stashes the real failure on `error`. `RolloutLimits` is the framework's per-rollout +drives it: assigns its model client at register, routes each intercepted model call to it, +runs `refused()` before each turn, and stashes the real failure on `error`. `RolloutLimits` is the framework's per-rollout budget (turns / tokens), checked between turns. """ @@ -64,10 +64,13 @@ def reached(self, trace: Trace) -> str | None: @dataclass class RolloutSession: ctx: ModelContext - client: Client trace: Trace stops: list[Callable[[Trace], Awaitable[bool]]] = field(default_factory=list) limits: RolloutLimits = field(default_factory=RolloutLimits) + client: Client | None = None + """The model client serving this rollout's turns. The interception server assigns it at + `register` (one server-owned client per distinct endpoint config), so every rollout it + multiplexes shares one keepalive connection pool instead of opening its own.""" error: "RolloutError | None" = None """The latest unresolved model-call failure. The harness only sees it as an HTTP error (and may swallow it, or exit non-zero), so the rollout re-raises this original error once the