From 86b68c52448ee89a0f86306dc4c25f4438065daa Mon Sep 17 00:00:00 2001 From: Behrooz Azarkhalili Date: Sat, 29 Aug 2026 14:38:10 -0700 Subject: [PATCH 1/3] fix: bound the check_server health probe so connection_timeout can fire check_server polls /health with a bare requests.get, and its deadline test lives only inside the except RequestException branch. A server that accepts the connection and then stalls never raises, so the loop never reaches that test and blocks forever. The connection_timeout parameter documents that a ConnectionError is raised once the budget is spent, and against a stalled server that promise could not be kept. Passing timeout=retry_interval bounds each probe without adding a knob: the loop already sleeps that long between attempts, so a probe outliving one interval is late by the method's own measure. A refused connection behaves as before, since it raised straight away already. Adds two hermetic regression tests. One points check_server at a socket that accepts and never answers, and requires it to give up inside the budget. The other is the control: a healthy server must still be accepted. Without the fix the first test hangs rather than failing, which is the bug. Refs #6973 --- tests/test_vllm_client_server.py | 53 ++++++++++++++++++++++++++++++++ trl/generation/vllm_client.py | 4 ++- 2 files changed, 56 insertions(+), 1 deletion(-) diff --git a/tests/test_vllm_client_server.py b/tests/test_vllm_client_server.py index d697692ca91..1fd325bbc56 100644 --- a/tests/test_vllm_client_server.py +++ b/tests/test_vllm_client_server.py @@ -12,11 +12,17 @@ # See the License for the specific language governing permissions and # limitations under the License. +import http.server import os +import socket +import socketserver import subprocess +import threading +import time from types import SimpleNamespace import pytest +import requests from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer from transformers.testing_utils import torch_device @@ -127,6 +133,53 @@ def test_extract_logprobs_returns_none_token_ids_when_logprobs_missing(self): assert all_token_ids is None +class TestCheckServerHealthProbe(TrlTestCase): + """`check_server` must give up after `total_timeout` even when the server stalls rather than refusing.""" + + @staticmethod + def _client_for(port: int) -> VLLMClient: + # `__init__` needs a live vLLM server; `check_server` is exercised on its own here. + client = VLLMClient.__new__(VLLMClient) + client.base_url = f"http://127.0.0.1:{port}" + return client + + def test_stalled_server_raises_rather_than_blocking(self): + # A socket that accepts connections and never answers. The request succeeds at the TCP level and then hangs, + # so an unbounded probe never raises and `total_timeout` is never reached. + server = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + server.bind(("127.0.0.1", 0)) + server.listen(8) + try: + client = self._client_for(server.getsockname()[1]) + start = time.time() + with pytest.raises(requests.exceptions.ConnectionError): + client.check_server(total_timeout=0.1, retry_interval=0.5) + assert time.time() - start < 5.0, "check_server did not honor total_timeout against a stalled server" + finally: + server.close() + + def test_healthy_server_still_accepted(self): + class Handler(http.server.BaseHTTPRequestHandler): + def do_GET(self): + self.send_response(200) + self.end_headers() + + def log_message(self, *args): + pass + + httpd = socketserver.TCPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=httpd.serve_forever, daemon=True) + thread.start() + try: + client = self._client_for(httpd.server_address[1]) + client.check_server(total_timeout=5.0, retry_interval=0.5) # must return without raising + finally: + httpd.shutdown() + httpd.server_close() + thread.join(timeout=5.0) + + @pytest.mark.slow @require_torch_multi_accelerator @require_vllm diff --git a/trl/generation/vllm_client.py b/trl/generation/vllm_client.py index e045675eab7..186ca3d729f 100644 --- a/trl/generation/vllm_client.py +++ b/trl/generation/vllm_client.py @@ -270,7 +270,9 @@ def check_server(self, total_timeout: float = 0.0, retry_interval: float = 2.0): while True: try: - response = requests.get(url) + # The probe must be bounded: a server that accepts the connection and then stalls would otherwise + # block here forever, and `total_timeout` is only checked in the except branch below. + response = requests.get(url, timeout=retry_interval) except requests.exceptions.RequestException as exc: # Check if the total timeout duration has passed elapsed_time = time.time() - start_time From fe8d46eb836eb8d63b1babdccd630d87f28f59b0 Mon Sep 17 00:00:00 2001 From: Behrooz Azarkhalili Date: Sat, 29 Aug 2026 20:46:35 -0700 Subject: [PATCH 2/3] fix(vllm): honor check_server's total_timeout for a reachable but unready server The elapsed-time check sat inside `except RequestException`, so it ran only when the probe raised. A vLLM server that accepts the connection and answers HTTP 503 while it loads weights returns normally, so the deadline was never consulted and the loop retried forever. Against a local 503 server, `total_timeout=0.1` was still running after 5 seconds; 404 reproduces it identically, so the defect is any non-200 status rather than 503 specifically. Checking the deadline where the raising and non-raising branches converge bounds both. Two further bounds follow from the same contract: one probe is capped by the remaining budget, because `retry_interval=5.0` against a stalled socket overshot a 0.1s deadline by 4.9s, and the retry sleep is capped the same way. Deadline arithmetic moves to `time.monotonic()` so a system clock adjustment cannot shift it. `total_timeout` defaults to 0.0, meaning one attempt and then give up. That attempt stays bounded by `retry_interval` rather than by the deadline, and the docstring now says so. --- tests/test_vllm_client_server.py | 63 +++++++++++++++++++++++++------- trl/generation/vllm_client.py | 50 ++++++++++++++++--------- 2 files changed, 81 insertions(+), 32 deletions(-) diff --git a/tests/test_vllm_client_server.py b/tests/test_vllm_client_server.py index 1fd325bbc56..099d935d606 100644 --- a/tests/test_vllm_client_server.py +++ b/tests/test_vllm_client_server.py @@ -134,7 +134,13 @@ def test_extract_logprobs_returns_none_token_ids_when_logprobs_missing(self): class TestCheckServerHealthProbe(TrlTestCase): - """`check_server` must give up after `total_timeout` even when the server stalls rather than refusing.""" + """`check_server` must give up after `total_timeout` whenever the server is not up, however it fails.""" + + # A retry interval far larger than the deadline: every guard that must clamp to the remaining budget shows up + # as the difference between finishing near TOTAL_TIMEOUT and finishing near RETRY_INTERVAL. Asserting against + # RETRY_INTERVAL / 2 keeps the bound derived from these two values rather than a hand-picked constant. + TOTAL_TIMEOUT = 0.1 + RETRY_INTERVAL = 5.0 @staticmethod def _client_for(port: int) -> VLLMClient: @@ -143,6 +149,22 @@ def _client_for(port: int) -> VLLMClient: client.base_url = f"http://127.0.0.1:{port}" return client + @staticmethod + def _serve(status: int) -> tuple[socketserver.TCPServer, threading.Thread]: + # Parameterised over the status code so the healthy and the not-ready cases share one server. + class Handler(http.server.BaseHTTPRequestHandler): + def do_GET(self): + self.send_response(status) + self.end_headers() + + def log_message(self, *args): + pass + + httpd = socketserver.TCPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=httpd.serve_forever, daemon=True) + thread.start() + return httpd, thread + def test_stalled_server_raises_rather_than_blocking(self): # A socket that accepts connections and never answers. The request succeeds at the TCP level and then hangs, # so an unbounded probe never raises and `total_timeout` is never reached. @@ -154,23 +176,36 @@ def test_stalled_server_raises_rather_than_blocking(self): client = self._client_for(server.getsockname()[1]) start = time.time() with pytest.raises(requests.exceptions.ConnectionError): - client.check_server(total_timeout=0.1, retry_interval=0.5) - assert time.time() - start < 5.0, "check_server did not honor total_timeout against a stalled server" + client.check_server(total_timeout=self.TOTAL_TIMEOUT, retry_interval=self.RETRY_INTERVAL) + elapsed = time.time() - start + assert elapsed < self.RETRY_INTERVAL / 2, ( + f"check_server took {elapsed:.3f}s against a stalled server: the probe is bounded by the retry " + f"interval instead of the {self.TOTAL_TIMEOUT}s deadline" + ) finally: server.close() - def test_healthy_server_still_accepted(self): - class Handler(http.server.BaseHTTPRequestHandler): - def do_GET(self): - self.send_response(200) - self.end_headers() - - def log_message(self, *args): - pass + def test_unavailable_server_raises_rather_than_looping(self): + # A server that is reachable but reports it is not ready. `requests` returns normally, so a deadline checked + # only on the exception path never fires and the probe retries forever. vLLM answers 503 while it loads. + httpd, thread = self._serve(503) + try: + client = self._client_for(httpd.server_address[1]) + start = time.time() + with pytest.raises(requests.exceptions.ConnectionError): + client.check_server(total_timeout=self.TOTAL_TIMEOUT, retry_interval=self.RETRY_INTERVAL) + elapsed = time.time() - start + assert elapsed < self.RETRY_INTERVAL / 2, ( + f"check_server took {elapsed:.3f}s against a server answering 503: it slept the full retry " + f"interval instead of stopping at the {self.TOTAL_TIMEOUT}s deadline" + ) + finally: + httpd.shutdown() + httpd.server_close() + thread.join(timeout=5.0) - httpd = socketserver.TCPServer(("127.0.0.1", 0), Handler) - thread = threading.Thread(target=httpd.serve_forever, daemon=True) - thread.start() + def test_healthy_server_still_accepted(self): + httpd, thread = self._serve(200) try: client = self._client_for(httpd.server_address[1]) client.check_server(total_timeout=5.0, retry_interval=0.5) # must return without raising diff --git a/trl/generation/vllm_client.py b/trl/generation/vllm_client.py index 186ca3d729f..04dbc7364d4 100644 --- a/trl/generation/vllm_client.py +++ b/trl/generation/vllm_client.py @@ -257,30 +257,33 @@ def _post(self, url: str, **kwargs) -> dict: def check_server(self, total_timeout: float = 0.0, retry_interval: float = 2.0): """ Check server availability with retries on failure, within a total timeout duration. If the server is not up - after the total timeout duration, raise a `ConnectionError`. + after the total timeout duration, raise a `ConnectionError`. A server that is reachable but not ready, which + vLLM reports as HTTP 503 while it loads, counts as not up. Args: - retry_interval (`float`, *optional*, defaults to `2.0`): - Interval in seconds between retries. total_timeout (`float`, *optional*, defaults to `0.0`): - Total timeout duration in seconds. + Total timeout duration in seconds, measured from the first attempt. The default of `0.0` makes a single + attempt and then gives up, so that attempt is still bounded by `retry_interval` rather than by the + deadline. + retry_interval (`float`, *optional*, defaults to `2.0`): + Interval in seconds between retries, and the upper bound on how long any single attempt may take. """ url = f"{self.base_url}/health" - start_time = time.time() # Record the start time + start_time = time.monotonic() # Deadlines use the monotonic clock so a system clock change cannot move them while True: + # Bound each probe by the retry interval so a server that accepts the connection and then stalls cannot + # block forever, and by whatever is left of the deadline so a single slow probe cannot overshoot it. A + # non-positive remainder means the deadline has already passed, and the check below ends the loop right + # after this attempt; `total_timeout` defaults to 0.0, which is exactly "try once, then give up". + remaining = total_timeout - (time.monotonic() - start_time) + cause = None try: - # The probe must be bounded: a server that accepts the connection and then stalls would otherwise - # block here forever, and `total_timeout` is only checked in the except branch below. - response = requests.get(url, timeout=retry_interval) + response = requests.get( + url, timeout=min(retry_interval, remaining) if remaining > 0 else retry_interval + ) except requests.exceptions.RequestException as exc: - # Check if the total timeout duration has passed - elapsed_time = time.time() - start_time - if elapsed_time >= total_timeout: - raise ConnectionError( - f"The vLLM server can't be reached at {self.base_url} after {total_timeout} seconds. Make " - "sure the server is running by running `vllm serve`." - ) from exc + cause = exc else: if response.status_code == 200: if "X-Forwarded-For" in response.headers: @@ -288,9 +291,20 @@ def check_server(self, total_timeout: float = 0.0, retry_interval: float = 2.0): logger.info("Server is up!") return - # Retry logic: wait before trying again - logger.info(f"Server is not up yet. Retrying in {retry_interval} seconds...") - time.sleep(retry_interval) + # The attempt failed, either by raising or by answering with a non-200 status. A server that is + # reachable but still loading answers 503, so checking the deadline here rather than in the except + # branch is what bounds that case. + elapsed_time = time.monotonic() - start_time + if elapsed_time >= total_timeout: + raise ConnectionError( + f"The vLLM server at {self.base_url} was not ready within {total_timeout} seconds. Make sure " + "the server is running by running `vllm serve`, and that it has finished loading." + ) from cause + + # Retry logic: wait before trying again, never past the deadline + delay = min(retry_interval, total_timeout - elapsed_time) + logger.info(f"Server is not up yet. Retrying in {delay} seconds...") + time.sleep(delay) def get_world_size(self) -> int: """ From dfcf95b9cf9d8b1580c9b62249320792efee120f Mon Sep 17 00:00:00 2001 From: Behrooz Azarkhalili Date: Thu, 3 Sep 2026 12:36:04 -0700 Subject: [PATCH 3/3] fix(vllm): stop the readiness wait before a post-deadline probe and align the async clients `check_server` slept up to `retry_interval`, then probed again without checking the deadline. When the sleep landed on the deadline the next probe ran for a full `retry_interval` and accepted a server that came up after the budget (measured: a 0.05 s budget accepted a 200 at 0.251 s). The wait now checks the deadline once more after sleeping. `retry_interval` must be positive: `requests` rejects a zero timeout with a ValueError from inside the loop, and a negative value would have skipped the sleep. The docstring states the remaining limit, that `requests` applies the timeout to the connect and to each read separately. The async GRPO and async distillation clients keep their own `wait_for_server_ready`, which probed with a fixed `poll_interval_s` timeout, slept the full interval, and checked the deadline only afterwards, so it could overshoot `server_timeout` by up to one poll interval plus one probe. Both now use the monotonic clock, bound every probe and every sleep by what is left of the deadline, and stop before a probe that would start past it. The two copies stay identical apart from the config name. Tests: the health-probe cases gain a retry case (503 then 200, two requests seen), a late-200 case that must be refused, and a rejected non-positive interval; the one-shot mutant that passed the previous cases now fails three. A new experimental test file exercises both async clients against a stalled socket and a 503 server and bounds the wait by the deadline. --- tests/experimental/test_async_vllm_client.py | 93 +++++++++++++++++++ tests/test_vllm_client_server.py | 58 ++++++++++-- .../async_distillation/vllm_client.py | 31 +++++-- trl/experimental/async_grpo/vllm_client.py | 31 +++++-- trl/generation/vllm_client.py | 21 +++-- 5 files changed, 202 insertions(+), 32 deletions(-) create mode 100644 tests/experimental/test_async_vllm_client.py diff --git a/tests/experimental/test_async_vllm_client.py b/tests/experimental/test_async_vllm_client.py new file mode 100644 index 00000000000..9f11db07aa9 --- /dev/null +++ b/tests/experimental/test_async_vllm_client.py @@ -0,0 +1,93 @@ +# Copyright 2020-2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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 http.server +import socket +import socketserver +import threading +import time + +import pytest +from accelerate import PartialState + +from trl.experimental.async_distillation.vllm_client import VLLMClient as AsyncDistillationClient +from trl.experimental.async_grpo.vllm_client import VLLMClient as AsyncGRPOClient + +from ..testing_utils import TrlTestCase + + +@pytest.mark.parametrize("client_cls", [AsyncGRPOClient, AsyncDistillationClient]) +class TestAsyncClientReadiness(TrlTestCase): + """`wait_for_server_ready` must give up after `server_timeout`, however the server fails. The two clients are + copies of each other, so both are exercised.""" + + @pytest.fixture(autouse=True) + def accelerate_state(self): + # The clients log through `accelerate.logging`, which refuses to log before the process state exists; the + # trainers always create it before constructing a client. + PartialState() + + # A poll interval far larger than the deadline, as in `TestCheckServerHealthProbe`: a wait that is bounded by the + # poll interval, or by a fixed per-probe timeout, instead of the deadline overshoots the allowance. + SERVER_TIMEOUT = 0.1 + POLL_INTERVAL = 5.0 + ALLOWANCE = 0.5 + + def test_stalled_server_raises_within_the_deadline(self, client_cls): + # A socket that accepts connections and never answers, so the probe itself has to be bounded. + server = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + server.bind(("127.0.0.1", 0)) + server.listen(8) + try: + client = client_cls(f"http://127.0.0.1:{server.getsockname()[1]}", server_timeout=self.SERVER_TIMEOUT) + start = time.time() + with pytest.raises(TimeoutError): + client.wait_for_server_ready(poll_interval_s=self.POLL_INTERVAL) + elapsed = time.time() - start + assert elapsed < self.SERVER_TIMEOUT + self.ALLOWANCE, ( + f"wait_for_server_ready took {elapsed:.3f}s against a stalled server: the probe is not bounded by " + f"the {self.SERVER_TIMEOUT}s deadline" + ) + finally: + server.close() + + def test_not_ready_server_raises_within_the_deadline(self, client_cls): + # A server answering 503 while it loads: `requests` returns normally, so the sleep between polls has to be + # clamped to the deadline rather than the poll interval. + class Handler(http.server.BaseHTTPRequestHandler): + def do_GET(self): + self.send_response(503) + self.end_headers() + + def log_message(self, *args): + pass + + httpd = socketserver.TCPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=httpd.serve_forever, daemon=True) + thread.start() + try: + client = client_cls(f"http://127.0.0.1:{httpd.server_address[1]}", server_timeout=self.SERVER_TIMEOUT) + start = time.time() + with pytest.raises(TimeoutError): + client.wait_for_server_ready(poll_interval_s=self.POLL_INTERVAL) + elapsed = time.time() - start + assert elapsed < self.SERVER_TIMEOUT + self.ALLOWANCE, ( + f"wait_for_server_ready took {elapsed:.3f}s against a server answering 503: it slept the full poll " + f"interval instead of stopping at the {self.SERVER_TIMEOUT}s deadline" + ) + finally: + httpd.shutdown() + httpd.server_close() + thread.join(timeout=5.0) diff --git a/tests/test_vllm_client_server.py b/tests/test_vllm_client_server.py index 099d935d606..84a7f828d19 100644 --- a/tests/test_vllm_client_server.py +++ b/tests/test_vllm_client_server.py @@ -137,10 +137,12 @@ class TestCheckServerHealthProbe(TrlTestCase): """`check_server` must give up after `total_timeout` whenever the server is not up, however it fails.""" # A retry interval far larger than the deadline: every guard that must clamp to the remaining budget shows up - # as the difference between finishing near TOTAL_TIMEOUT and finishing near RETRY_INTERVAL. Asserting against - # RETRY_INTERVAL / 2 keeps the bound derived from these two values rather than a hand-picked constant. + # as the difference between finishing near TOTAL_TIMEOUT and finishing near RETRY_INTERVAL. The bound is the + # deadline plus an allowance for socket setup on a loaded runner, so a probe that ignored the deadline and used + # any fixed timeout of a second or more fails it as well. TOTAL_TIMEOUT = 0.1 RETRY_INTERVAL = 5.0 + ALLOWANCE = 0.5 @staticmethod def _client_for(port: int) -> VLLMClient: @@ -150,10 +152,19 @@ def _client_for(port: int) -> VLLMClient: return client @staticmethod - def _serve(status: int) -> tuple[socketserver.TCPServer, threading.Thread]: - # Parameterised over the status code so the healthy and the not-ready cases share one server. + def _serve(*statuses: int, delay_before_last: float = 0.0) -> tuple[socketserver.TCPServer, threading.Thread]: + # One status per request in order, the last one repeated; the last one can be held back by a delay so a test + # can place a healthy answer after the deadline. + remaining = list(statuses) + class Handler(http.server.BaseHTTPRequestHandler): def do_GET(self): + if len(remaining) > 1: + status = remaining.pop(0) + else: + status = remaining[0] + time.sleep(delay_before_last) + self.server.requests_seen += 1 self.send_response(status) self.end_headers() @@ -161,6 +172,7 @@ def log_message(self, *args): pass httpd = socketserver.TCPServer(("127.0.0.1", 0), Handler) + httpd.requests_seen = 0 thread = threading.Thread(target=httpd.serve_forever, daemon=True) thread.start() return httpd, thread @@ -178,7 +190,7 @@ def test_stalled_server_raises_rather_than_blocking(self): with pytest.raises(requests.exceptions.ConnectionError): client.check_server(total_timeout=self.TOTAL_TIMEOUT, retry_interval=self.RETRY_INTERVAL) elapsed = time.time() - start - assert elapsed < self.RETRY_INTERVAL / 2, ( + assert elapsed < self.TOTAL_TIMEOUT + self.ALLOWANCE, ( f"check_server took {elapsed:.3f}s against a stalled server: the probe is bounded by the retry " f"interval instead of the {self.TOTAL_TIMEOUT}s deadline" ) @@ -195,7 +207,7 @@ def test_unavailable_server_raises_rather_than_looping(self): with pytest.raises(requests.exceptions.ConnectionError): client.check_server(total_timeout=self.TOTAL_TIMEOUT, retry_interval=self.RETRY_INTERVAL) elapsed = time.time() - start - assert elapsed < self.RETRY_INTERVAL / 2, ( + assert elapsed < self.TOTAL_TIMEOUT + self.ALLOWANCE, ( f"check_server took {elapsed:.3f}s against a server answering 503: it slept the full retry " f"interval instead of stopping at the {self.TOTAL_TIMEOUT}s deadline" ) @@ -214,6 +226,40 @@ def test_healthy_server_still_accepted(self): httpd.server_close() thread.join(timeout=5.0) + def test_server_that_becomes_ready_is_accepted_on_retry(self): + # 503 once, then 200 within the budget: the loop must retry rather than give up on the first non-200, which + # is the case a single-attempt implementation gets wrong. + httpd, thread = self._serve(503, 200) + try: + client = self._client_for(httpd.server_address[1]) + client.check_server(total_timeout=5.0, retry_interval=0.05) # must return without raising + assert httpd.requests_seen == 2 + finally: + httpd.shutdown() + httpd.server_close() + thread.join(timeout=5.0) + + def test_no_probe_starts_after_the_deadline(self): + # 503 once, then a 200 held back until after the deadline but inside one retry interval. The retry sleep is + # clamped to the deadline, so the second probe would start right on it; a probe started there runs for a + # full `retry_interval` and accepts the late 200. The deadline must be checked again before probing. + httpd, thread = self._serve(503, 200, delay_before_last=0.15) + try: + client = self._client_for(httpd.server_address[1]) + with pytest.raises(requests.exceptions.ConnectionError): + client.check_server(total_timeout=self.TOTAL_TIMEOUT, retry_interval=0.2) + finally: + httpd.shutdown() + httpd.server_close() + thread.join(timeout=5.0) + + def test_non_positive_retry_interval_is_rejected(self): + # `retry_interval` bounds every attempt, so zero cannot bound anything; `requests` would raise its own + # `ValueError` about a zero timeout from inside the loop otherwise. + client = self._client_for(9) + with pytest.raises(ValueError, match="retry_interval"): + client.check_server(total_timeout=self.TOTAL_TIMEOUT, retry_interval=0.0) + @pytest.mark.slow @require_torch_multi_accelerator diff --git a/trl/experimental/async_distillation/vllm_client.py b/trl/experimental/async_distillation/vllm_client.py index b12a46ce4fb..a7ccc154b95 100644 --- a/trl/experimental/async_distillation/vllm_client.py +++ b/trl/experimental/async_distillation/vllm_client.py @@ -42,27 +42,38 @@ def __init__(self, server_url: str, server_timeout: float = 240.0): self.server_timeout = server_timeout def wait_for_server_ready(self, poll_interval_s: float = 2.0) -> None: - """Block until the server answers `/health`, or raise `TimeoutError` after `server_timeout` seconds.""" + """Block until the server answers `/health`, or raise `TimeoutError` after `server_timeout` seconds. Each probe + is bounded by `poll_interval_s` and by what is left of the deadline, so a server that accepts the connection + and then stalls, or one that keeps answering 503 while it loads, cannot hold the wait past it.""" logger.info(f"Waiting for vLLM server at {self.server_url} ...") - start = time.time() + start = time.monotonic() # Deadlines use the monotonic clock so a system clock change cannot move them + timed_out = ( + f"Timed out after {self.server_timeout:.0f}s waiting for vLLM server at {self.server_url}. " + "Make sure the vLLM server is running and reachable. If the server needs more time to load the " + "model, increase `vllm_server_timeout` in your AsyncDistillationConfig." + ) while True: - elapsed = time.time() - start + elapsed = time.monotonic() - start + remaining = self.server_timeout - elapsed try: - response = requests.get(f"{self.server_url}/health", timeout=5) + response = requests.get( + f"{self.server_url}/health", + timeout=min(poll_interval_s, remaining) if remaining > 0 else poll_interval_s, + ) if response.status_code == 200: logger.info(f"vLLM server ready after {elapsed:.1f}s") return except (requests.ConnectionError, requests.Timeout, OSError): pass + elapsed = time.monotonic() - start if elapsed >= self.server_timeout: - raise TimeoutError( - f"Timed out after {self.server_timeout:.0f}s waiting for vLLM server at {self.server_url}. " - "Make sure the vLLM server is running and reachable. If the server needs more time to load the " - "model, increase `vllm_server_timeout` in your AsyncDistillationConfig." - ) + raise TimeoutError(timed_out) if int(elapsed) % 10 < poll_interval_s: logger.info(f"Still waiting for vLLM server... ({elapsed:.0f}s)") - time.sleep(poll_interval_s) + # Never sleep past the deadline, and never start a probe once it has passed + time.sleep(min(poll_interval_s, self.server_timeout - elapsed)) + if time.monotonic() - start >= self.server_timeout: + raise TimeoutError(timed_out) def get_max_model_len(self) -> int: """Return the served model's `max_model_len` (the cap on prompt + completion tokens).""" diff --git a/trl/experimental/async_grpo/vllm_client.py b/trl/experimental/async_grpo/vllm_client.py index 278bc5ca3d5..556289bb756 100644 --- a/trl/experimental/async_grpo/vllm_client.py +++ b/trl/experimental/async_grpo/vllm_client.py @@ -42,27 +42,38 @@ def __init__(self, server_url: str, server_timeout: float = 240.0): self.server_timeout = server_timeout def wait_for_server_ready(self, poll_interval_s: float = 2.0) -> None: - """Block until the server answers `/health`, or raise `TimeoutError` after `server_timeout` seconds.""" + """Block until the server answers `/health`, or raise `TimeoutError` after `server_timeout` seconds. Each probe + is bounded by `poll_interval_s` and by what is left of the deadline, so a server that accepts the connection + and then stalls, or one that keeps answering 503 while it loads, cannot hold the wait past it.""" logger.info(f"Waiting for vLLM server at {self.server_url} ...") - start = time.time() + start = time.monotonic() # Deadlines use the monotonic clock so a system clock change cannot move them + timed_out = ( + f"Timed out after {self.server_timeout:.0f}s waiting for vLLM server at {self.server_url}. " + "Make sure the vLLM server is running and reachable. If the server needs more time to load " + "the model, increase `vllm_server_timeout` in your AsyncGRPOConfig." + ) while True: - elapsed = time.time() - start + elapsed = time.monotonic() - start + remaining = self.server_timeout - elapsed try: - response = requests.get(f"{self.server_url}/health", timeout=5) + response = requests.get( + f"{self.server_url}/health", + timeout=min(poll_interval_s, remaining) if remaining > 0 else poll_interval_s, + ) if response.status_code == 200: logger.info(f"vLLM server ready after {elapsed:.1f}s") return except (requests.ConnectionError, requests.Timeout, OSError): pass + elapsed = time.monotonic() - start if elapsed >= self.server_timeout: - raise TimeoutError( - f"Timed out after {self.server_timeout:.0f}s waiting for vLLM server at {self.server_url}. " - "Make sure the vLLM server is running and reachable. If the server needs more time to load " - "the model, increase `vllm_server_timeout` in your AsyncGRPOConfig." - ) + raise TimeoutError(timed_out) if int(elapsed) % 10 < poll_interval_s: logger.info(f"Still waiting for vLLM server... ({elapsed:.0f}s)") - time.sleep(poll_interval_s) + # Never sleep past the deadline, and never start a probe once it has passed + time.sleep(min(poll_interval_s, self.server_timeout - elapsed)) + if time.monotonic() - start >= self.server_timeout: + raise TimeoutError(timed_out) def get_max_model_len(self) -> int: """Return the served model's `max_model_len` (the cap on prompt + completion tokens).""" diff --git a/trl/generation/vllm_client.py b/trl/generation/vllm_client.py index 04dbc7364d4..3b8ed95a8ec 100644 --- a/trl/generation/vllm_client.py +++ b/trl/generation/vllm_client.py @@ -266,10 +266,18 @@ def check_server(self, total_timeout: float = 0.0, retry_interval: float = 2.0): attempt and then gives up, so that attempt is still bounded by `retry_interval` rather than by the deadline. retry_interval (`float`, *optional*, defaults to `2.0`): - Interval in seconds between retries, and the upper bound on how long any single attempt may take. + Interval in seconds between retries, and the upper bound on how long any single attempt may take. Must + be positive. `requests` applies it to the connect and to each read separately, so a server that keeps + sending bytes can hold one attempt past it. """ + if retry_interval <= 0.0: + raise ValueError(f"`retry_interval` must be positive, got {retry_interval}: it bounds every attempt.") url = f"{self.base_url}/health" start_time = time.monotonic() # Deadlines use the monotonic clock so a system clock change cannot move them + not_ready = ( + f"The vLLM server at {self.base_url} was not ready within {total_timeout} seconds. Make sure the server " + "is running by running `vllm serve`, and that it has finished loading." + ) while True: # Bound each probe by the retry interval so a server that accepts the connection and then stalls cannot @@ -296,15 +304,16 @@ def check_server(self, total_timeout: float = 0.0, retry_interval: float = 2.0): # branch is what bounds that case. elapsed_time = time.monotonic() - start_time if elapsed_time >= total_timeout: - raise ConnectionError( - f"The vLLM server at {self.base_url} was not ready within {total_timeout} seconds. Make sure " - "the server is running by running `vllm serve`, and that it has finished loading." - ) from cause + raise ConnectionError(not_ready) from cause - # Retry logic: wait before trying again, never past the deadline + # Retry logic: wait before trying again, never past the deadline. The sleep can still land on the + # deadline, and a probe started then would run for a full `retry_interval` and could accept a server that + # came up after the budget, so check once more before probing. delay = min(retry_interval, total_timeout - elapsed_time) logger.info(f"Server is not up yet. Retrying in {delay} seconds...") time.sleep(delay) + if time.monotonic() - start_time >= total_timeout: + raise ConnectionError(not_ready) from cause def get_world_size(self) -> int: """