diff --git a/packages/nemo_evaluator_sdk/src/nemo_evaluator_sdk/agent_eval/runtimes/gym/sandboxed.py b/packages/nemo_evaluator_sdk/src/nemo_evaluator_sdk/agent_eval/runtimes/gym/sandboxed.py index 8c01b1a593..38be467ebe 100644 --- a/packages/nemo_evaluator_sdk/src/nemo_evaluator_sdk/agent_eval/runtimes/gym/sandboxed.py +++ b/packages/nemo_evaluator_sdk/src/nemo_evaluator_sdk/agent_eval/runtimes/gym/sandboxed.py @@ -84,6 +84,8 @@ class SandboxedGymRuntimeConfig(BaseModel): "instance an environment's config defines (`mcqa_simple_agent`), not the agent component " "(`simple_agent`).", ) + num_repeats: int = Field(default=1, ge=1, description="Attempts per row; each attempt becomes one trial.") + concurrency: int = Field(default=4, ge=1, description="Maximum concurrent Gym rollout attempts.") reward_key: str = Field(default=DEFAULT_REWARD_KEY, description="Key read from each rollout record.") @@ -117,6 +119,8 @@ def runner_info(self) -> RunnerInfo: "mode": "sandboxed", "rollout_url": cfg.rollout_url, "agent_ref_name": cfg.agent_ref_name, + "num_repeats": cfg.num_repeats, + "concurrency": cfg.concurrency, "reward_key": cfg.reward_key, "timeout_s": cfg.timeout_s, }, @@ -129,26 +133,28 @@ def _request_headers(self) -> dict[str, str]: return headers async def _collect(self, examples: list[dict[str, Any]]) -> list[dict[str, Any]]: - """POST the examples and return the host's rollout records.""" + """Ask the host to repeat examples while bounding Gym's concurrent rollout attempts.""" + cfg = self._config async with httpx.AsyncClient(timeout=self._config.timeout_s) as client: response = await client.post( - self._config.rollout_url, - json={"examples": examples}, + cfg.rollout_url, + json={ + "examples": examples, + "num_repeats": cfg.num_repeats, + "concurrency": cfg.concurrency, + }, headers=self._request_headers(), ) if response.status_code >= 400: - # The body is the host's own error envelope; it names which example or server failed, - # which the status code alone does not. + # The body says which example or server failed; the status code alone does not. raise RuntimeError( - f"sandboxed Gym host returned {response.status_code} from {self._config.rollout_url}: " - f"{response.text[:2000]}" + f"sandboxed Gym host returned {response.status_code} from {cfg.rollout_url}: {response.text[:2000]}" ) body = response.json() results = body.get("results") if isinstance(body, Mapping) else None if not isinstance(results, list): raise RuntimeError( - f"sandboxed Gym host returned no `results` list from {self._config.rollout_url}; " - f"got {type(results).__name__}" + f"sandboxed Gym host returned no `results` list from {cfg.rollout_url}; got {type(results).__name__}" ) return [record for record in results if isinstance(record, dict)] diff --git a/packages/nemo_evaluator_sdk/tests/agent_eval/test_gym_sandboxed_runtime.py b/packages/nemo_evaluator_sdk/tests/agent_eval/test_gym_sandboxed_runtime.py index 7a96b58e41..ee378884dd 100644 --- a/packages/nemo_evaluator_sdk/tests/agent_eval/test_gym_sandboxed_runtime.py +++ b/packages/nemo_evaluator_sdk/tests/agent_eval/test_gym_sandboxed_runtime.py @@ -48,6 +48,7 @@ class _FakeHost: def __init__(self, *, status: int = 200, body: Any = None, rewards: dict[int, float] | None = None) -> None: self.requests: list[httpx.Request] = [] self.posted: list[dict[str, Any]] = [] + self.payloads: list[dict[str, Any]] = [] self._status = status self._body = body self._rewards = rewards @@ -55,18 +56,26 @@ def __init__(self, *, status: int = 200, body: Any = None, rewards: dict[int, fl def transport(self) -> httpx.MockTransport: def handle(request: httpx.Request) -> httpx.Response: self.requests.append(request) - self.posted = json.loads(request.content.decode())["examples"] + payload = json.loads(request.content.decode()) + self.payloads.append(payload) + examples = payload["examples"] + self.posted.extend(examples) if self._body is not None or self._status >= 400: return httpx.Response(self._status, json=self._body if self._body is not None else {"error": "boom"}) rewards = self._rewards if self._rewards is not None else {} + attempts = [ + {**example, NG_ROLLOUT_INDEX: attempt} + for example in examples + for attempt in range(payload["num_repeats"]) + ] results = [ { NG_TASK_INDEX: example[NG_TASK_INDEX], - NG_ROLLOUT_INDEX: 0, + NG_ROLLOUT_INDEX: example.get(NG_ROLLOUT_INDEX, 0), "reward": rewards.get(example[NG_TASK_INDEX], 1.0), "response": f"answer-{example[NG_TASK_INDEX]}", } - for example in self.posted + for example in attempts ] return httpx.Response(200, json={"results": results}) @@ -111,6 +120,20 @@ async def test_the_examples_posted_carry_the_index_we_stamped(tasks, tmp_path, m assert all("responses_create_params" in example for example in host.posted) +async def test_repeat_and_concurrency_contract_crosses_http_once(tasks, tmp_path, monkeypatch) -> None: + host = _FakeHost() + runner = runner_against(host, monkeypatch, num_repeats=3, concurrency=2) + + trials = await runner.run_tasks(tasks, AgentEvalRunConfig(work_dir=tmp_path)) + + assert len(host.requests) == 1 + assert host.payloads == [{"examples": host.posted, "num_repeats": 3, "concurrency": 2}] + assert len(trials) == 6 + for task in tasks: + task_trials = [trial for trial in trials if trial.task_id == task.id] + assert {trial.metadata[NG_ROLLOUT_INDEX] for trial in task_trials} == {0, 1, 2} + + async def test_the_auth_token_is_sent_as_the_proxy_header(tasks, tmp_path, monkeypatch) -> None: host = _FakeHost() runner = runner_against(host, monkeypatch, auth_token="tok-123", headers={"X-Extra": "kept"}) @@ -182,7 +205,12 @@ def handle(request: httpx.Request) -> httpx.Response: def test_runner_info_records_the_host_but_not_the_token() -> None: runner = SandboxedGymAgentTaskRunner( - config=SandboxedGymRuntimeConfig(rollout_url=ROLLOUT_URL, auth_token="sk-secret-value") + config=SandboxedGymRuntimeConfig( + rollout_url=ROLLOUT_URL, + auth_token="sk-secret-value", + num_repeats=3, + concurrency=2, + ) ) info = runner.runner_info() @@ -190,4 +218,6 @@ def test_runner_info_records_the_host_but_not_the_token() -> None: assert info.name == "gym" assert info.config["mode"] == "sandboxed" assert info.config["rollout_url"] == ROLLOUT_URL + assert info.config["num_repeats"] == 3 + assert info.config["concurrency"] == 2 assert "sk-secret-value" not in json.dumps(info.config), "the token must not reach the run bundle" diff --git a/packages/sandboxed_gym/src/sandboxed_gym/runtime/gym_host_runtime.py b/packages/sandboxed_gym/src/sandboxed_gym/runtime/gym_host_runtime.py index 4178fc76dc..1c9fd8e677 100644 --- a/packages/sandboxed_gym/src/sandboxed_gym/runtime/gym_host_runtime.py +++ b/packages/sandboxed_gym/src/sandboxed_gym/runtime/gym_host_runtime.py @@ -415,9 +415,25 @@ async def _collect_rollout_results( examples: list[dict], head_server_config: Any, rollout_helper: Any, + *, + num_repeats: int = 1, + concurrency: int = 4, ) -> list[dict]: + attempts = [ + { + **example, + **({NG_ROLLOUT_INDEX: attempt} if NG_TASK_INDEX in example else {}), + } + for example in examples + for attempt in range(num_repeats) + ] + semaphore = asyncio.Semaphore(concurrency) results: list[dict] = [] - for task in rollout_helper.run_examples(examples=examples, head_server_config=head_server_config): + for task in rollout_helper.run_examples( + examples=attempts, + head_server_config=head_server_config, + semaphore=semaphore, + ): row, nemo_gym_result = await task results.append(_with_row_identity(nemo_gym_result, row)) return results @@ -427,8 +443,26 @@ def run_rollouts_sync( examples: list[dict], head_server_config: Any, rollout_helper: Any, + *, + num_repeats: int = 1, + concurrency: int = 4, ) -> list[dict]: - return asyncio.run(_collect_rollout_results(examples, head_server_config, rollout_helper)) + return asyncio.run( + _collect_rollout_results( + examples, + head_server_config, + rollout_helper, + num_repeats=num_repeats, + concurrency=concurrency, + ) + ) + + +def _positive_int(request: dict[str, Any], field: str, default: int) -> int: + value = request.get(field, default) + if isinstance(value, bool) or not isinstance(value, int) or value < 1: + raise ValueError(f"{field} must be a positive integer") + return value class Handler(BaseHTTPRequestHandler): @@ -494,7 +528,20 @@ def do_POST(self) -> None: return try: - results = run_rollouts_sync(examples, _HEAD_SERVER_CONFIG, _ROLLOUT_HELPER) + num_repeats = _positive_int(request, "num_repeats", 1) + concurrency = _positive_int(request, "concurrency", 4) + except ValueError as exc: + self._send_json(400, _runtime_error("internal", str(exc))) + return + + try: + results = run_rollouts_sync( + examples, + _HEAD_SERVER_CONFIG, + _ROLLOUT_HELPER, + num_repeats=num_repeats, + concurrency=concurrency, + ) except Exception as exc: self._send_json( 500, diff --git a/packages/sandboxed_gym/tests/test_gym_host_runtime.py b/packages/sandboxed_gym/tests/test_gym_host_runtime.py index 361cf8096e..b0420bbb12 100644 --- a/packages/sandboxed_gym/tests/test_gym_host_runtime.py +++ b/packages/sandboxed_gym/tests/test_gym_host_runtime.py @@ -16,7 +16,14 @@ class _FakeRolloutHelper: - def run_examples(self, examples, head_server_config=None): + def __init__(self): + self.examples = [] + self.semaphore = None + + def run_examples(self, examples, head_server_config=None, semaphore=None): + self.examples = examples + self.semaphore = semaphore + async def _one(row): return row, {"response": {"output": []}, "reward": 0.0} @@ -90,6 +97,56 @@ def test_rollouts_run_returns_results(ready_server): assert body["results"][0]["reward"] == 0.0 +def test_rollouts_run_applies_repeats_and_concurrency(ready_server): + import urllib.request + + payload = json.dumps( + { + "examples": [{"agent_ref": {"name": "a"}, "_ng_task_index": 7}], + "num_repeats": 3, + "concurrency": 2, + } + ).encode() + req = urllib.request.Request( + f"{ready_server}/rollouts/run", + data=payload, + method="POST", + headers={"Content-Type": "application/json"}, + ) + + with urllib.request.urlopen(req, timeout=10) as resp: + body = json.loads(resp.read().decode()) + + assert len(body["results"]) == 3 + assert [result["_ng_rollout_index"] for result in body["results"]] == [0, 1, 2] + helper = runtime._ROLLOUT_HELPER + assert len(helper.examples) == 3 + assert helper.semaphore._value == 2 + + +@pytest.mark.parametrize( + ("field", "value"), + [("num_repeats", 0), ("num_repeats", True), ("concurrency", 0), ("concurrency", "2")], +) +def test_rollouts_run_rejects_invalid_repeat_and_concurrency(ready_server, field, value): + import urllib.error + import urllib.request + + payload = json.dumps({"examples": [], field: value}).encode() + req = urllib.request.Request( + f"{ready_server}/rollouts/run", + data=payload, + method="POST", + headers={"Content-Type": "application/json"}, + ) + + with pytest.raises(urllib.error.HTTPError) as exc: + urllib.request.urlopen(req, timeout=5) + + assert exc.value.code == 400 + assert field in json.loads(exc.value.read().decode())["error"]["message"] + + def test_rollouts_run_rejects_oversize_request(ready_server): import urllib.error import urllib.request @@ -223,7 +280,7 @@ def _boom(*a, **k): class _IdentityStrippingHelper: """A helper whose results carry no index, i.e. Gym did not copy the caller's stamp through.""" - def run_examples(self, examples, head_server_config=None): + def run_examples(self, examples, head_server_config=None, semaphore=None): async def _one(row): return row, {"response": {"output": []}, "reward": 0.5} @@ -233,7 +290,7 @@ async def _one(row): class _IdentityPreservingHelper: """A helper whose results carry Gym's own indices, which must win over the row's.""" - def run_examples(self, examples, head_server_config=None): + def run_examples(self, examples, head_server_config=None, semaphore=None): async def _one(row): return row, {"reward": 1.0, "_ng_task_index": 99, "_ng_rollout_index": 7} diff --git a/plugins/nemo-evaluator/src/nemo_evaluator/jobs/gym_sandbox.py b/plugins/nemo-evaluator/src/nemo_evaluator/jobs/gym_sandbox.py index 63fc30ec97..e8430a0d5a 100644 --- a/plugins/nemo-evaluator/src/nemo_evaluator/jobs/gym_sandbox.py +++ b/plugins/nemo-evaluator/src/nemo_evaluator/jobs/gym_sandbox.py @@ -475,6 +475,8 @@ async def run_tasks(self, tasks: Any, config: Any = None) -> Any: auth_token=descriptor.rollout_auth_token, headers=dict(descriptor.headers), agent_ref_name=self._target.agent_ref_name or self._target.agent, + num_repeats=self._target.num_repeats, + concurrency=self._target.concurrency, reward_key=self._target.reward_key, ) ) diff --git a/plugins/nemo-evaluator/tests/integration/test_sandboxed_gym_execution.py b/plugins/nemo-evaluator/tests/integration/test_sandboxed_gym_execution.py index 73d8d2e466..01d71cac7d 100644 --- a/plugins/nemo-evaluator/tests/integration/test_sandboxed_gym_execution.py +++ b/plugins/nemo-evaluator/tests/integration/test_sandboxed_gym_execution.py @@ -53,11 +53,15 @@ def _tasks(tmp_path: Path) -> list: return discover_gym_tasks(dataset) -def _target() -> GymRunnerTarget: +def _target(**overrides: Any) -> GymRunnerTarget: + fields: dict[str, Any] = { + "agent": "simple_agent", + "agent_config": "responses_api_agents/simple_agent/configs/simple_agent.yaml", + "resources_server": "mcqa", + } + fields.update(overrides) return GymRunnerTarget( - agent="simple_agent", - agent_config="responses_api_agents/simple_agent/configs/simple_agent.yaml", - resources_server="mcqa", + **fields, ) @@ -95,6 +99,7 @@ class _StubGymHostHandler(BaseHTTPRequestHandler): #: Set by the fixture; the spec the provider was asked to create a host for. received_specs: list[Any] = [] + received_payloads: list[dict[str, Any]] = [] def log_message(self, *args: Any) -> None: # keep pytest output readable return @@ -109,15 +114,21 @@ def do_GET(self) -> None: def do_POST(self) -> None: payload = json.loads(self.rfile.read(int(self.headers["Content-Length"])).decode()) + self.received_payloads.append(payload) # Echo the caller's own index back, which is what a real Gym host does with a stamped row. + attempts = [ + {**example, NG_ROLLOUT_INDEX: attempt} + for example in payload["examples"] + for attempt in range(payload["num_repeats"]) + ] results = [ { NG_TASK_INDEX: example[NG_TASK_INDEX], - NG_ROLLOUT_INDEX: 0, + NG_ROLLOUT_INDEX: example.get(NG_ROLLOUT_INDEX, 0), "reward": float(example[NG_TASK_INDEX]), "response": f"answer-{example[NG_TASK_INDEX]}", } - for example in payload["examples"] + for example in attempts ] body = json.dumps({"results": results}).encode() self.send_response(200) @@ -169,6 +180,7 @@ async def destroy_host(self, handle: Any) -> None: @pytest.fixture def stub_provider(monkeypatch: pytest.MonkeyPatch) -> Iterator[_StubHostProvider]: + _StubGymHostHandler.received_payloads = [] provider = _StubHostProvider() # Patched where the orchestrator looks it up, so the orchestrator itself stays untouched. monkeypatch.setattr("sandboxed_gym.orchestrator.get_host_provider", lambda *a, **k: provider) @@ -189,6 +201,26 @@ async def test_a_sandboxed_run_provisions_a_host_and_returns_attributed_trials( assert rewards[tasks[1].id] == 1.0, "each trial must carry its own task's reward, not a neighbour's" +async def test_sandboxed_repeat_and_concurrency_contract_reaches_the_host( + stub_provider: _StubHostProvider, tmp_path: Path +) -> None: + tasks = _tasks(tmp_path) + target = _target(num_repeats=3, concurrency=2) + runner = SessionBackedGymRunner(target=target, plan=_plan(), job_id="eval-job-repeats") + + trials = await runner.run_tasks(tasks, AgentEvalRunConfig(work_dir=tmp_path)) + + assert len(_StubGymHostHandler.received_payloads) == 1 + payload = _StubGymHostHandler.received_payloads[0] + assert payload["num_repeats"] == 3 + assert payload["concurrency"] == 2 + assert len(payload["examples"]) == 2 + assert len(trials) == 6 + for task in tasks: + task_trials = [trial for trial in trials if trial.task_id == task.id] + assert {trial.metadata[NG_ROLLOUT_INDEX] for trial in task_trials} == {0, 1, 2} + + async def test_the_host_is_created_from_the_deployment_config_and_the_targets_selection( stub_provider: _StubHostProvider, tmp_path: Path ) -> None: