diff --git a/nemo_gym/base_responses_api_agent.py b/nemo_gym/base_responses_api_agent.py index b701e38191..474b4a1f64 100644 --- a/nemo_gym/base_responses_api_agent.py +++ b/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 @@ -43,12 +44,19 @@ BaseServer, SimpleServer, apply_rollout_prefix, + get_response_json, + raise_for_status, rollout_path_prefix, ) from nemo_gym.telemetry.endpoints import traced_endpoint, traced_rollout_endpoint from nemo_gym.telemetry.span_groups import GymSpanGroup +# 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): skip_verification: bool = False skip_verification_reward: float = 0.0 @@ -209,3 +217,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) diff --git a/responses_api_agents/browsecomp_agent/app.py b/responses_api_agents/browsecomp_agent/app.py index 1895125b56..e3a257c8b4 100644 --- a/responses_api_agents/browsecomp_agent/app.py +++ b/responses_api_agents/browsecomp_agent/app.py @@ -863,13 +863,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): """ diff --git a/responses_api_agents/conversational_tool_use/simulation/app.py b/responses_api_agents/conversational_tool_use/simulation/app.py index 53ef0cfac2..f2b6a50084 100644 --- a/responses_api_agents/conversational_tool_use/simulation/app.py +++ b/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, diff --git a/responses_api_agents/finance_agent/app.py b/responses_api_agents/finance_agent/app.py index e837d5bced..4ca52a9746 100644 --- a/responses_api_agents/finance_agent/app.py +++ b/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__": diff --git a/responses_api_agents/gymnasium_agent/app.py b/responses_api_agents/gymnasium_agent/app.py index b1516448ce..2daf5fd0cb 100644 --- a/responses_api_agents/gymnasium_agent/app.py +++ b/responses_api_agents/gymnasium_agent/app.py @@ -235,13 +235,7 @@ async def _run_open_episode( ) 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__": diff --git a/responses_api_agents/non_executing_simple_agent/app.py b/responses_api_agents/non_executing_simple_agent/app.py index aae53536d2..9457cc3628 100644 --- a/responses_api_agents/non_executing_simple_agent/app.py +++ b/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__": diff --git a/responses_api_agents/remote_agent/app.py b/responses_api_agents/remote_agent/app.py index 2d2ec4d085..2382d02553 100644 --- a/responses_api_agents/remote_agent/app.py +++ b/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__": diff --git a/responses_api_agents/remote_agent/tests/test_app.py b/responses_api_agents/remote_agent/tests/test_app.py index d818acfe25..6430f628f0 100644 --- a/responses_api_agents/remote_agent/tests/test_app.py +++ b/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 @@ -990,25 +990,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) diff --git a/responses_api_agents/simple_agent/app.py b/responses_api_agents/simple_agent/app.py index 26a3c5c5ec..4f72a520a1 100644 --- a/responses_api_agents/simple_agent/app.py +++ b/responses_api_agents/simple_agent/app.py @@ -353,6 +353,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.""" + return await self.proxy_aggregate_metrics(self.config.resources_server.name, body) if self.config.skip_verification: return await super().aggregate_metrics(body) diff --git a/responses_api_agents/speed_bench_agent/app.py b/responses_api_agents/speed_bench_agent/app.py index 39357ca258..e212916533 100644 --- a/responses_api_agents/speed_bench_agent/app.py +++ b/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__": diff --git a/responses_api_agents/stirrup_agent/app.py b/responses_api_agents/stirrup_agent/app.py index 9088789fa8..a80b0d3cc6 100644 --- a/responses_api_agents/stirrup_agent/app.py +++ b/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__": diff --git a/tests/unit_tests/test_aggregate_metrics.py b/tests/unit_tests/test_aggregate_metrics.py index 3175dcb062..43528c1e8b 100644 --- a/tests/unit_tests/test_aggregate_metrics.py +++ b/tests/unit_tests/test_aggregate_metrics.py @@ -12,10 +12,14 @@ # 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 unittest.mock import AsyncMock, MagicMock import warnings from unittest.mock import MagicMock +import orjson import pytest +from pydantic import ValidationError from nemo_gym.base_resources_server import ( AggregateMetrics, @@ -330,6 +334,138 @@ async def run(self, body=None): assert "mean/reward" in result.key_metrics +class TestProxyAggregateMetrics: + @staticmethod + def _make_agent(server_client): + class TestAgent(SimpleResponsesAPIAgent): + async def responses(self, body=None): + pass + + async def run(self, body=None): + pass + + config = BaseResponsesAPIAgentConfig(host="127.0.0.1", port=12345, entrypoint="app.py", name="test_agent") + return TestAgent(config=config, server_client=server_client) + + @staticmethod + def _make_response(payload): + response = MagicMock() + response.ok = True + response.read = AsyncMock(return_value=orjson.dumps(payload)) + return response + + @pytest.mark.asyncio + async def test_success(self) -> None: + payload = { + "group_level_metrics": [], + "agent_metrics": {"mean/reward": 1.0}, + "key_metrics": {"mean/reward": 1.0}, + } + server_client = MagicMock(spec=ServerClient) + server_client.post = AsyncMock(return_value=self._make_response(payload)) + agent = self._make_agent(server_client) + body = AggregateMetricsRequest(verify_responses=[]) + + result = await agent.proxy_aggregate_metrics("resources", body) + + assert result == AggregateMetrics.model_validate(payload) + server_client.post.assert_awaited_once_with( + server_name="resources", + url_path="/aggregate_metrics", + json=body, + ) + + @pytest.mark.asyncio + async def test_http_error(self) -> None: + response = MagicMock() + response.ok = False + response.content.read = AsyncMock(return_value=b"failure") + response.raise_for_status.side_effect = RuntimeError("http error") + server_client = MagicMock(spec=ServerClient) + server_client.post = AsyncMock(return_value=response) + agent = self._make_agent(server_client) + + with pytest.raises(RuntimeError, match="http error"): + await agent.proxy_aggregate_metrics("resources", AggregateMetricsRequest(verify_responses=[])) + + @pytest.mark.asyncio + async def test_invalid_response(self) -> None: + server_client = MagicMock(spec=ServerClient) + server_client.post = AsyncMock(return_value=self._make_response({"agent_metrics": []})) + agent = self._make_agent(server_client) + + with pytest.raises(ValidationError): + await agent.proxy_aggregate_metrics("resources", AggregateMetricsRequest(verify_responses=[])) + + @pytest.mark.asyncio + async def test_timeout(self) -> None: + async def hang(*args, **kwargs): + await asyncio.Event().wait() + + server_client = MagicMock(spec=ServerClient) + server_client.post = AsyncMock(side_effect=hang) + agent = self._make_agent(server_client) + + with pytest.raises(asyncio.TimeoutError): + await agent.proxy_aggregate_metrics( + "resources", + AggregateMetricsRequest(verify_responses=[]), + timeout_secs=0.01, + ) + + @pytest.mark.asyncio + async def test_default_timeout_bounds_the_hop(self, monkeypatch: pytest.MonkeyPatch) -> None: + payload = { + "group_level_metrics": [], + "agent_metrics": {}, + "key_metrics": {}, + } + server_client = MagicMock(spec=ServerClient) + server_client.post = AsyncMock(return_value=self._make_response(payload)) + agent = self._make_agent(server_client) + + recorded = [] + original_wait_for = asyncio.wait_for + + async def recording_wait_for(awaitable, timeout=None): + recorded.append(timeout) + return await original_wait_for(awaitable, timeout=timeout) + + monkeypatch.setattr(asyncio, "wait_for", recording_wait_for) + + result = await agent.proxy_aggregate_metrics( + "resources", + AggregateMetricsRequest(verify_responses=[]), + ) + + assert result == AggregateMetrics.model_validate(payload) + assert recorded == [600.0] + + @pytest.mark.asyncio + async def test_none_timeout_bypasses_wait_for(self, monkeypatch: pytest.MonkeyPatch) -> None: + payload = { + "group_level_metrics": [], + "agent_metrics": {}, + "key_metrics": {}, + } + server_client = MagicMock(spec=ServerClient) + server_client.post = AsyncMock(return_value=self._make_response(payload)) + agent = self._make_agent(server_client) + + def unexpected_wait_for(*args, **kwargs): + raise AssertionError("asyncio.wait_for should not be called") + + monkeypatch.setattr(asyncio, "wait_for", unexpected_wait_for) + + result = await agent.proxy_aggregate_metrics( + "resources", + AggregateMetricsRequest(verify_responses=[]), + timeout_secs=None, + ) + + assert result == AggregateMetrics.model_validate(payload) + + class TestTaskIndexInGroupMetrics: def test_task_index_preserved(self) -> None: responses = [