feat(agents): add aggregate metrics proxy helper

nemo_gym/base_responses_api_agent.py

@@ -12,6 +12,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import asyncio
 from abc import abstractmethod
 from collections.abc import Mapping
 from functools import wraps
@@ -38,10 +39,17 @@
     BaseServer,
     SimpleServer,
     apply_rollout_prefix,
+    get_response_json,
+    raise_for_status,
     rollout_path_prefix,
 )
 
 
+# Default bound on the aggregate-metrics proxy hop. ServerClient retries connection errors
+# indefinitely, so an unbounded proxy to a dead server would hang the caller forever.
+DEFAULT_AGGREGATE_METRICS_PROXY_TIMEOUT_SECS = 600.0
+
+
 class BaseResponsesAPIAgentConfig(BaseRunServerInstanceConfig):
     pass
 
@@ -150,3 +158,28 @@ async def aggregate_metrics(self, body: AggregateMetricsRequest = Body()) -> Agg
             compute_metrics_fn=self.compute_metrics,
             get_key_metrics_fn=self.get_key_metrics,
         )
+
+    async def proxy_aggregate_metrics(
+        self,
+        server_name: str,
+        body: AggregateMetricsRequest,
+        timeout_secs: Optional[float] = DEFAULT_AGGREGATE_METRICS_PROXY_TIMEOUT_SECS,
+    ) -> AggregateMetrics:
+        """Proxy aggregate metrics to another server with an optional timeout.
+
+        ServerClient retries connection errors indefinitely, so a dead resources server
+        at the end of a run could otherwise hang the collector after all rollouts are on disk.
+        """
+
+        async def _proxy() -> AggregateMetrics:
+            response = await self.server_client.post(
+                server_name=server_name,
+                url_path="/aggregate_metrics",
+                json=body,
+            )
+            await raise_for_status(response)
+            return AggregateMetrics.model_validate(await get_response_json(response))
+
+        if timeout_secs is None:
+            return await _proxy()
+        return await asyncio.wait_for(_proxy(), timeout=timeout_secs)

responses_api_agents/browsecomp_agent/app.py

@@ -864,13 +864,7 @@ async def run(self, request: Request, body: BrowsecompAgentRunRequest) -> Browse
 
     async def aggregate_metrics(self, body: AggregateMetricsRequest = Body()) -> AggregateMetrics:
         """Proxy aggregate_metrics to the resources server."""
-        response = await self.server_client.post(
-            server_name=self.config.resources_server.name,
-            url_path="/aggregate_metrics",
-            json=body,
-        )
-        await raise_for_status(response)
-        return AggregateMetrics.model_validate(await get_response_json(response))
+        return await self.proxy_aggregate_metrics(self.config.resources_server.name, body)
 
     def _compact_old_tool_messages(self, messages):
         """

responses_api_agents/conversational_tool_use/simulation/app.py

@@ -380,13 +380,7 @@ async def run(
                 await self._discard_session(cookies)
 
     async def aggregate_metrics(self, body: AggregateMetricsRequest = Body()) -> AggregateMetrics:
-        response = await self.server_client.post(
-            server_name=self.config.resources_server.name,
-            url_path="/aggregate_metrics",
-            json=body,
-        )
-        await raise_for_status(response)
-        return AggregateMetrics.model_validate(await get_response_json(response))
+        return await self.proxy_aggregate_metrics(self.config.resources_server.name, body)
 
     async def _execute_tool_call(
         self,

responses_api_agents/finance_agent/app.py

@@ -522,13 +522,7 @@ async def _run_inner(self, request: Request, body: FinanceAgentRunRequest) -> Fi
 
     async def aggregate_metrics(self, body: AggregateMetricsRequest = Body()) -> AggregateMetrics:
         """Proxy aggregate_metrics to the resources server."""
-        response = await self.server_client.post(
-            server_name=self.config.resources_server.name,
-            url_path="/aggregate_metrics",
-            json=body,
-        )
-        await raise_for_status(response)
-        return AggregateMetrics.model_validate(await get_response_json(response))
+        return await self.proxy_aggregate_metrics(self.config.resources_server.name, body)
 
 
 if __name__ == "__main__":

responses_api_agents/gymnasium_agent/app.py

@@ -164,13 +164,7 @@ async def run(self, request: Request, body: GymnasiumAgentRunRequest) -> Gymnasi
         )
 
     async def aggregate_metrics(self, body: AggregateMetricsRequest = Body()) -> AggregateMetrics:
-        response = await self.server_client.post(
-            server_name=self.config.resources_server.name,
-            url_path="/aggregate_metrics",
-            json=body,
-        )
-        await raise_for_status(response)
-        return AggregateMetrics.model_validate(await get_response_json(response))
+        return await self.proxy_aggregate_metrics(self.config.resources_server.name, body)
 
 
 if __name__ == "__main__":

responses_api_agents/non_executing_simple_agent/app.py

@@ -130,13 +130,7 @@ async def run(
 
     async def aggregate_metrics(self, body: AggregateMetricsRequest = Body()) -> AggregateMetrics:
         """Proxy aggregate_metrics to the resources server."""
-        response = await self.server_client.post(
-            server_name=self.config.resources_server.name,
-            url_path="/aggregate_metrics",
-            json=body,
-        )
-        await raise_for_status(response)
-        return AggregateMetrics.model_validate(await get_response_json(response))
+        return await self.proxy_aggregate_metrics(self.config.resources_server.name, body)
 
 
 if __name__ == "__main__":

responses_api_agents/remote_agent/app.py

@@ -79,7 +79,6 @@
 _REMOTE_RETRY_SLEEP_SECS = 0.5
 _FAILURE_PRINT_HEAD = 5
 _FAILURE_PRINT_INTERVAL = 100
-_AGGREGATE_PROXY_TIMEOUT_SECS = 600.0
 
 # Result/routing keys this server itself produces. Input rows may carry stale copies
 # (e.g. a rollouts or failures JSONL re-fed as a dataset); they must never collide with
@@ -490,22 +489,8 @@ def _empty_response(self) -> NeMoGymResponse:
         )
 
     async def aggregate_metrics(self, body: AggregateMetricsRequest = Body()) -> AggregateMetrics:
-        """Proxy aggregate_metrics to the resources server.
-
-        Bounded: the ServerClient hop otherwise retries connection errors forever, and a dead
-        resources server at end-of-run would hang the collector after all rollouts are on disk.
-        """
-
-        async def _proxy() -> AggregateMetrics:
-            response = await self.server_client.post(
-                server_name=self.config.resources_server.name,
-                url_path="/aggregate_metrics",
-                json=body,
-            )
-            await raise_for_status(response)
-            return AggregateMetrics.model_validate(await get_response_json(response))
-
-        return await asyncio.wait_for(_proxy(), timeout=_AGGREGATE_PROXY_TIMEOUT_SECS)
+        """Proxy aggregate_metrics to the resources server."""
+        return await self.proxy_aggregate_metrics(self.config.resources_server.name, body)
 
 
 if __name__ == "__main__":

responses_api_agents/remote_agent/tests/test_app.py

@@ -24,7 +24,7 @@
 from pydantic import BaseModel, ValidationError
 
 import responses_api_agents.remote_agent.app as remote_agent_app
-from nemo_gym.config_types import ResourcesServerRef
+from nemo_gym.config_types import AggregateMetrics, AggregateMetricsRequest, ResourcesServerRef
 from nemo_gym.openai_utils import NeMoGymResponseCreateParamsNonStreaming
 from nemo_gym.rollout_collection import NG_FAILURE_CLASS_KEY, NG_NO_PERSIST_KEY, NG_TERMINAL_KEY
 from nemo_gym.server_utils import ServerClient
@@ -984,25 +984,17 @@ async def _post(server_name, url_path, json=None, **kwargs):
         server_client.post = AsyncMock(side_effect=_post)
         agent = make_agent(server_client=server_client)
 
-        from nemo_gym.base_resources_server import AggregateMetricsRequest
-
         result = await agent.aggregate_metrics(AggregateMetricsRequest(verify_responses=[]))
         assert result.key_metrics == {"mean/reward": 1.0}
 
-    async def test_aggregate_metrics_bounded_when_resources_server_hangs(
-        self, monkeypatch: pytest.MonkeyPatch
-    ) -> None:
-        async def hang(*args, **kwargs):
-            await asyncio.sleep(60)
+    async def test_aggregate_metrics_delegates_to_shared_proxy(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        expected = AggregateMetrics(key_metrics={"mean/reward": 1.0})
+        proxy = AsyncMock(return_value=expected)
+        agent = make_agent()
+        monkeypatch.setattr(RemoteAgent, "proxy_aggregate_metrics", proxy)
+        body = AggregateMetricsRequest(verify_responses=[])
 
-        server_client = MagicMock(spec=ServerClient)
-        server_client.post = AsyncMock(side_effect=hang)
-        agent = make_agent(server_client=server_client)
-        monkeypatch.setattr(remote_agent_app, "_AGGREGATE_PROXY_TIMEOUT_SECS", 0.05)
+        result = await agent.aggregate_metrics(body)
 
-        with pytest.raises(asyncio.TimeoutError):
-            await agent.aggregate_metrics(
-                __import__(
-                    "nemo_gym.base_resources_server", fromlist=["AggregateMetricsRequest"]
-                ).AggregateMetricsRequest(verify_responses=[])
-            )
+        assert result is expected
+        proxy.assert_awaited_once_with("my_env", body)

responses_api_agents/simple_agent/app.py

@@ -345,13 +345,7 @@ async def run(self, request: Request, body: SimpleAgentRunRequest) -> SimpleAgen
 
     async def aggregate_metrics(self, body: AggregateMetricsRequest = Body()) -> AggregateMetrics:
         """Proxy aggregate_metrics to the resources server."""
-        response = await self.server_client.post(
-            server_name=self.config.resources_server.name,
-            url_path="/aggregate_metrics",
-            json=body,
-        )
-        await raise_for_status(response)
-        return AggregateMetrics.model_validate(await get_response_json(response))
+        return await self.proxy_aggregate_metrics(self.config.resources_server.name, body)
 
 
 if __name__ == "__main__":

responses_api_agents/speed_bench_agent/app.py

@@ -249,13 +249,7 @@ async def run(self, request: Request, body: SpeedBenchAgentRunRequest) -> SpeedB
         return SpeedBenchAgentVerifyResponse.model_validate(await get_response_json(verify_response))
 
     async def aggregate_metrics(self, body: AggregateMetricsRequest = Body()) -> AggregateMetrics:
-        api_response = await self.server_client.post(
-            server_name=self.config.resources_server.name,
-            url_path="/aggregate_metrics",
-            json=body,
-        )
-        await raise_for_status(api_response)
-        return AggregateMetrics.model_validate(await get_response_json(api_response))
+        return await self.proxy_aggregate_metrics(self.config.resources_server.name, body)
 
 
 if __name__ == "__main__":

responses_api_agents/stirrup_agent/app.py

@@ -1605,13 +1605,7 @@ async def aggregate_metrics(self, body: AggregateMetricsRequest = Body()) -> Agg
         # implementation to avoid a needless judge-server round trip.
         if self.config.execute_only:
             return await SimpleResponsesAPIAgent.aggregate_metrics(self, body)
-        response = await self.server_client.post(
-            server_name=self.config.resources_server.name,
-            url_path="/aggregate_metrics",
-            json=body,
-        )
-        await raise_for_status(response)
-        return AggregateMetrics.model_validate(await get_response_json(response))
+        return await self.proxy_aggregate_metrics(self.config.resources_server.name, body)
 
 
 if __name__ == "__main__":