fix(evaluator): honor sandboxed Gym repeat concurrency

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)]
 

packages/nemo_evaluator_sdk/tests/agent_eval/test_gym_sandboxed_runtime.py

@@ -48,25 +48,34 @@ 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
 
     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,12 +205,19 @@ 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()
 
     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"

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,

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}
 

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,
                 )
             )

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: