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 d697692ca91..84a7f828d19 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,134 @@ 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` 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. 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: + # `__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 + + @staticmethod + 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() + + 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 + + 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=self.TOTAL_TIMEOUT, retry_interval=self.RETRY_INTERVAL) + elapsed = time.time() - start + 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" + ) + finally: + server.close() + + 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.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" + ) + finally: + httpd.shutdown() + httpd.server_close() + thread.join(timeout=5.0) + + 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 + finally: + httpd.shutdown() + 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 @require_vllm 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 e045675eab7..3b8ed95a8ec 100644 --- a/trl/generation/vllm_client.py +++ b/trl/generation/vllm_client.py @@ -257,28 +257,41 @@ 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. 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.time() # Record the start time + 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 + # 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: - response = requests.get(url) + 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: @@ -286,9 +299,21 @@ 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(not_ready) from cause + + # 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: """