feat: judge endpoint resiliency

nemo_gym/openai_utils.py

@@ -855,6 +855,14 @@ class NeMoGymAsyncOpenAI(BaseModel):  # pragma: no cover
         description="Set this to true if this particular client is only used to call internal NeMo Gym servers.",
     )
 
+    max_connection_retries: Optional[int] = Field(
+        default=None,
+        description=(
+            "How many connection-error retries per request; None retries forever. "
+            "Allows callers that can resolve a moved endpoint to avoid stalling forever."
+        ),
+    )
+
     default_headers: Dict[str, str] = Field(
         default_factory=dict,
         description="Extra headers to include in every request.",
@@ -867,6 +875,7 @@ async def _request(self, **request_kwargs: Dict) -> ClientResponse:
                 "Authorization": f"Bearer {self.api_key}",
             },
             "_internal": self.internal,
+            "_max_connection_retries": self.max_connection_retries,
         }
         return await self._request_with_retry(**request_kwargs)
 

nemo_gym/server_utils.py

@@ -206,7 +206,11 @@ def global_aiohttp_client_exit():  # pragma: no cover
 
 
 async def request(
-    method: str, url: str, _internal: bool = False, **kwargs: Unpack[_RequestOptions]
+    method: str,
+    url: str,
+    _internal: bool = False,
+    _max_connection_retries: Optional[int] = None,
+    **kwargs: Unpack[_RequestOptions],
 ) -> ClientResponse:  # pragma: no cover
     # Faster JSON dumps than the default aiohttp json
     if kwargs.get("json"):
@@ -232,6 +236,10 @@ async def request(
                     flush=True,
                 )
 
+            # Retrying forever is wrong if the endpoint is expected to sometimes die and move.
+            if _max_connection_retries is not None and retries >= _max_connection_retries:
+                raise
+
             await asyncio.sleep(0.5)
         except ClientOSError:
             global _NUM_CLIENT_OS_ERROR
@@ -244,6 +252,9 @@ async def request(
                     flush=True,
                 )
 
+            if _max_connection_retries is not None and retries >= _max_connection_retries:
+                raise
+
             await asyncio.sleep(0.5)
         except Exception as e:
             if _GLOBAL_AIOHTTP_CLIENT_REQUEST_DEBUG:

responses_api_models/vllm_model/app.py

@@ -19,7 +19,7 @@
 import logging
 import os
 from copy import deepcopy
-from time import time, time_ns
+from time import monotonic, time, time_ns
 from typing import Any, ClassVar, Dict, List, Optional, Union
 
 from aiohttp.client_exceptions import ClientResponseError
@@ -193,6 +193,19 @@ class VLLMModelConfig(BaseResponsesAPIModelConfig):
     extra_body: Optional[Dict[str, Any]] = None
 
     default_headers: Dict[str, str] = Field(default_factory=dict)
+
+    # Optional path to a file that publishes the current backend base_url.
+    # Used for shared serving jobs that move hosts when they restart.
+    endpoint_file: Optional[str] = None
+
+    # How long a missing endpoint file allows for the use of the last-known-good clients.
+    endpoint_stale_grace_s: float = 300.0
+
+    # Connection-error retry bound applied to clients when endpoint_file is set.
+    endpoint_connection_retries: Optional[int] = 8
+
+    # How often endpoint_file may be stat'd; otherwise the `os.stat` results is cached and reused.
+    endpoint_check_interval_s: float = 10.0
     # Optional prefix for resolving relative ``metadata.audio_path`` (or
     # entries in ``metadata.audio_paths``) against. Absolute paths are used
     # as-is. When unset, relative paths raise. Audio is always inlined as a
@@ -269,11 +282,17 @@ def _post_init(self) -> None:
                 base_url=base_url,
                 api_key=self.config.api_key,
                 default_headers=self.config.default_headers,
+                max_connection_retries=(
+                    self.config.endpoint_connection_retries if self.config.endpoint_file else None
+                ),
             )
             for base_url in self.config.base_url
         ]
 
         self._session_id_to_client: Dict[str, NeMoGymAsyncOpenAI] = dict()
+        self._endpoint_file_mtime: Optional[float] = None
+        self._endpoint_missing_since: Optional[float] = None
+        self._endpoint_last_check_at: Optional[float] = None
 
         self._converter = self.get_converter()
         self._transport_call_index = 0
@@ -1270,7 +1289,80 @@ def _create_empty_chat_completion(self) -> NeMoGymChatCompletion:
             ],
         )
 
+    def _maybe_rebind_endpoint(self) -> None:
+        """Rebind clients when a shared serving job publishes a new endpoint.
+
+        Raises when the endpoints have stayed unpublished for longer than `endpoint_stale_grace_s`.
+        """
+        if not self.config.endpoint_file:
+            return
+        now = monotonic()
+        if (
+            self._endpoint_last_check_at is not None
+            and now - self._endpoint_last_check_at < self.config.endpoint_check_interval_s
+        ):
+            if self._endpoint_missing_since is not None:
+                self._note_endpoint_unpublished()
+            return
+        self._endpoint_last_check_at = now
+        try:
+            mtime = os.stat(self.config.endpoint_file).st_mtime
+        except FileNotFoundError:
+            # Serving jobs remove the endpoint file while rotating;
+            # keep the current clients until the successor publishes.
+            self._note_endpoint_unpublished()
+            return
+        except OSError:
+            # Transient filesystem trouble is not a backend exit; retry the current clients.
+            return
+        if mtime == self._endpoint_file_mtime:
+            if self._endpoint_missing_since is not None:
+                self._note_endpoint_unpublished()
+            return
+        try:
+            with open(self.config.endpoint_file) as endpoint_stream:
+                url = endpoint_stream.read().strip()
+        except OSError:
+            return
+        self._endpoint_file_mtime = mtime
+        if not url:
+            # An empty file is as unpublished as a missing one.
+            self._note_endpoint_unpublished()
+            return
+        self._endpoint_missing_since = None
+        if [url] == self.config.base_url:
+            return
+        print(
+            f"vllm_model '{self.config.name}': backend endpoint changed "
+            f"{self.config.base_url} -> {[url]}; rebinding clients.",
+            flush=True,
+        )
+        self.config.base_url = [url]
+        self._clients = [
+            NeMoGymAsyncOpenAI(
+                base_url=url,
+                api_key=self.config.api_key,
+                default_headers=self.config.default_headers,
+                max_connection_retries=self.config.endpoint_connection_retries,
+            )
+        ]
+        # Every session re-resolves onto the new host.
+        self._session_id_to_client.clear()
+
+    def _note_endpoint_unpublished(self) -> None:
+        now = monotonic()
+        if self._endpoint_missing_since is None:
+            self._endpoint_missing_since = now
+        elif now - self._endpoint_missing_since > self.config.endpoint_stale_grace_s:
+            raise RuntimeError(
+                f"vllm_model endpoint file {self.config.endpoint_file} unpublished (absent "
+                f"or empty) for {now - self._endpoint_missing_since:.0f}s (grace "
+                f"{self.config.endpoint_stale_grace_s:.0f}s); refusing to keep serving "
+                "against a backend that is no longer published."
+            )
+
     def _resolve_client(self, request: Request) -> NeMoGymAsyncOpenAI:
+        self._maybe_rebind_endpoint()
         session_id = request.session[SESSION_ID_KEY]
         if session_id not in self._session_id_to_client:
             # Uvicorn workers do not share this cache. A stable assignment keeps

responses_api_models/vllm_model/configs/vllm_model.yaml

@@ -12,3 +12,7 @@ policy_model:
       sampling_overrides: null
       extra_body: null
       default_headers: {}
+      endpoint_file: null
+      endpoint_stale_grace_s: 300.0
+      endpoint_connection_retries: 8
+      endpoint_check_interval_s: 10.0

responses_api_models/vllm_model/tests/test_app.py

@@ -5104,3 +5104,78 @@ def test_overrides_reach_the_responses_native_path(self) -> None:
         server = self._server({"temperature": 1.0}, is_responses_native=True)
         body = {"model": "dummy_model", "temperature": 0.2}
         assert server._apply_sampling_overrides(body)["temperature"] == 1.0
+
+
+class TestEndpointFile:
+    def _make_server(self, tmp_path, **overrides) -> VLLMModel:
+        params = dict(
+            host="0.0.0.0",
+            port=8081,
+            base_url="http://placeholder:8712/v1",
+            api_key="dummy_key",  # pragma: allowlist secret
+            model="dummy_model",
+            entrypoint="",
+            name="safety_judge_model",
+            return_token_id_information=False,
+            uses_reasoning_parser=False,
+            endpoint_file=str(tmp_path / "endpoint.txt"),
+        )
+        params.update(overrides)
+        return VLLMModel(config=VLLMModelConfig(**params), server_client=MagicMock(spec=ServerClient))
+
+    def test_publish_rebinds_clients_and_clears_sessions(self, tmp_path) -> None:
+        (tmp_path / "endpoint.txt").write_text("http://new-host:8712/v1\n")
+        server = self._make_server(tmp_path, endpoint_check_interval_s=3600.0)
+        server._session_id_to_client["session-on-old-host"] = server._clients[0]
+
+        server._maybe_rebind_endpoint()
+
+        assert server.config.base_url == ["http://new-host:8712/v1"]
+        assert [client.base_url for client in server._clients] == ["http://new-host:8712/v1"]
+        # endpoint_file-backed clients are retry-bounded so in-flight calls
+        # against a dead host fail fast and re-enter through the rebound
+        # client; sessions re-resolve onto the new host.
+        assert server._clients[0].max_connection_retries == 8
+        assert not server._session_id_to_client
+        # Within endpoint_check_interval_s the filesystem is left alone, so a
+        # fresh publish is only seen once the window is over.
+        (tmp_path / "endpoint.txt").write_text("http://newer-host:8712/v1\n")
+        server._maybe_rebind_endpoint()
+        assert server.config.base_url == ["http://new-host:8712/v1"]
+        server._endpoint_last_check_at = None  # window over: the next call re-checks
+        server._maybe_rebind_endpoint()
+        assert server.config.base_url == ["http://newer-host:8712/v1"]
+        # Static base_url clients keep today's retry-forever behavior.
+        assert self._make_server(tmp_path, endpoint_file=None)._clients[0].max_connection_retries is None
+
+    def test_unpublished_endpoint_grace_lifecycle(self, tmp_path, monkeypatch: MonkeyPatch) -> None:
+        endpoint_file = tmp_path / "endpoint.txt"
+        server = self._make_server(tmp_path)  # grace defaults to 300s
+        now = 1000.0
+        monkeypatch.setattr("responses_api_models.vllm_model.app.monotonic", lambda: now)
+
+        server._maybe_rebind_endpoint()  # absent: the clock starts, last known-good client kept
+        assert server.config.base_url == ["http://placeholder:8712/v1"]
+        now = 1100.0
+        endpoint_file.write_text("")  # empty is as unpublished as missing: no clock reset
+        server._maybe_rebind_endpoint()
+        now = 1301.0  # past the grace COUNTED FROM 1000, proving the empty write reset nothing
+        with raises(RuntimeError, match="no longer published"):
+            server._maybe_rebind_endpoint()
+
+        endpoint_file.write_text("http://placeholder:8712/v1\n")  # republish on the SAME host
+        now = 1301.5  # within the check window: the publish is not seen yet, the raise stays loud
+        with raises(RuntimeError, match="no longer published"):
+            server._maybe_rebind_endpoint()
+        now = 1311.5  # next window: heals with no rebind needed
+        server._maybe_rebind_endpoint()
+        assert server._endpoint_missing_since is None
+
+        now = 1400.0
+        endpoint_file.unlink()
+        server._maybe_rebind_endpoint()  # a fresh absence starts a fresh clock
+        now = 1650.0
+        server._maybe_rebind_endpoint()  # 250s in: within grace, so the republish reset the clock
+        now = 1701.0
+        with raises(RuntimeError, match="no longer published"):
+            server._maybe_rebind_endpoint()

tests/unit_tests/test_server_utils.py

@@ -15,6 +15,7 @@
 import socket
 from unittest.mock import AsyncMock, MagicMock
 
+from aiohttp import ClientOSError
 from pytest import MonkeyPatch, raises
 
 import nemo_gym.global_config
@@ -375,3 +376,23 @@ async def get_session(request: Request) -> dict:
             response = client.get("/session")
             assert response.json()["session_id"]
             assert 1 == len(response.headers.get_list("set-cookie"))
+
+    def _mock_global_client(self, monkeypatch: MonkeyPatch, connection_errors: int) -> MagicMock:
+        """Global-client stand-in whose request() raises ClientOSError `connection_errors` times, then succeeds."""
+        client = MagicMock()
+        client.request = AsyncMock(side_effect=[ClientOSError()] * connection_errors + [client.success_response])
+        monkeypatch.setattr(nemo_gym.server_utils, "get_global_aiohttp_client", lambda: client)
+        monkeypatch.setattr(nemo_gym.server_utils.asyncio, "sleep", AsyncMock())
+        return client
+
+    async def test_request_bounded_connection_retries_surface_dead_endpoint(self, monkeypatch: MonkeyPatch) -> None:
+        client = self._mock_global_client(monkeypatch, connection_errors=10)
+        with raises(ClientOSError):
+            await nemo_gym.server_utils.request("POST", "http://dead-host:1/v1", _max_connection_retries=3)
+        assert client.request.await_count == 3
+
+    async def test_request_connection_retries_unbounded_by_default(self, monkeypatch: MonkeyPatch) -> None:
+        client = self._mock_global_client(monkeypatch, connection_errors=4)
+        response = await nemo_gym.server_utils.request("POST", "http://flaky-host:1/v1")
+        assert response is client.success_response
+        assert client.request.await_count == 5