diff --git a/nemo_gym/openai_utils.py b/nemo_gym/openai_utils.py index 176760a29d..e50cc7b2be 100644 --- a/nemo_gym/openai_utils.py +++ b/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) diff --git a/nemo_gym/server_utils.py b/nemo_gym/server_utils.py index 6701b8956d..5ebe339e24 100644 --- a/nemo_gym/server_utils.py +++ b/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: diff --git a/responses_api_models/vllm_model/app.py b/responses_api_models/vllm_model/app.py index b2b8205eb5..f7943c9cde 100644 --- a/responses_api_models/vllm_model/app.py +++ b/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 diff --git a/responses_api_models/vllm_model/configs/vllm_model.yaml b/responses_api_models/vllm_model/configs/vllm_model.yaml index bf95088260..8dd93877e6 100644 --- a/responses_api_models/vllm_model/configs/vllm_model.yaml +++ b/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 diff --git a/responses_api_models/vllm_model/tests/test_app.py b/responses_api_models/vllm_model/tests/test_app.py index 3cbbb33c87..1d96befdb0 100644 --- a/responses_api_models/vllm_model/tests/test_app.py +++ b/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() diff --git a/tests/unit_tests/test_server_utils.py b/tests/unit_tests/test_server_utils.py index 35780418fa..cc4cf030e6 100644 --- a/tests/unit_tests/test_server_utils.py +++ b/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