From ec63611dd2546dbd70ee09ac6a717ca8a9e7febf Mon Sep 17 00:00:00 2001 From: yaodong-shen Date: Mon, 27 Jul 2026 10:28:21 +0800 Subject: [PATCH] fix: retry transient vLLM adapter reload failures Retry connection failures and transient gateway responses while reloading dedicated LoRA adapters. Stop immediately when the supervised child exits and report runtime diagnostics after exhaustion. Fixes #678 --- src/art/unsloth/service.py | 19 ++-- src/art/vllm_runtime.py | 91 ++++++++++++++++++ tests/unit/test_vllm_runtime.py | 165 ++++++++++++++++++++++++++++++++ 3 files changed, 264 insertions(+), 11 deletions(-) create mode 100644 tests/unit/test_vllm_runtime.py diff --git a/src/art/unsloth/service.py b/src/art/unsloth/service.py index e478e441f..b769cbf06 100644 --- a/src/art/unsloth/service.py +++ b/src/art/unsloth/service.py @@ -566,9 +566,6 @@ async def _sync_merged_weights( async def _reload_adapter(self, checkpoint_path: str, step: int) -> None: """Reload LoRA adapter in vLLM subprocess via HTTP.""" - import httpx - - self._raise_if_child_failed() lora_name = f"{self.model_name}@{step}" logger.info( f"[DEDICATED] _reload_adapter START: lora_name={lora_name} " @@ -580,14 +577,14 @@ async def _reload_adapter(self, checkpoint_path: str, step: int) -> None: } if self.serving_capabilities.inplace_lora_load: payload["load_inplace"] = True - async with httpx.AsyncClient() as client: - response = await client.post( - f"{self._vllm_base_url}/v1/load_lora_adapter", - json=payload, - **self._runtime_request_kwargs(), - timeout=60.0, - ) - response.raise_for_status() + response = await self._vllm_runtime.post_with_retry( + "/v1/load_lora_adapter", + json=payload, + timeout=60.0, + before_attempt=self._raise_if_child_failed, + operation="LoRA adapter reload", + failure_context=f"step={step}; checkpoint={checkpoint_path}", + ) logger.info( f"[DEDICATED] _reload_adapter DONE: lora_name={lora_name} " f"status={response.status_code}" diff --git a/src/art/vllm_runtime.py b/src/art/vllm_runtime.py index 2133db127..5037c7002 100644 --- a/src/art/vllm_runtime.py +++ b/src/art/vllm_runtime.py @@ -3,6 +3,7 @@ import fcntl import hashlib import json +import logging import math import os from pathlib import Path @@ -39,6 +40,9 @@ _FLASHINFER_WORKSPACE_ENV = "FLASHINFER_WORKSPACE_BASE" _ART_FLASHINFER_WORKSPACE_ENV = "ART_VLLM_RUNTIME_FLASHINFER_WORKSPACE_BASE" VLLM_RUNTIME_CLOSE_TIMEOUT = process_shutdown_timeout(1) +TRANSIENT_VLLM_STATUS_CODES = frozenset({502, 503, 504}) + +logger = logging.getLogger(__name__) class VllmRuntimeLaunchConfig(BaseModel): @@ -235,6 +239,93 @@ def request_kwargs(self) -> VllmRuntimeRequestKwargs: return {} return {"headers": {"Authorization": f"Bearer {self.api_key}"}} + def _process_status(self) -> str: + if self.process is None: + return "not managed by this process" + returncode = self.process.poll() + if returncode is None: + return "running" + return f"exited with code {returncode}" + + async def post_with_retry( + self, + path: str, + *, + json: Mapping[str, Any] | None = None, + timeout: float, + max_attempts: int = 4, + retry_base_delay: float = 1.0, + before_attempt: Callable[[], None] | None = None, + operation: str = "vLLM request", + failure_context: str | None = None, + http_client: httpx.AsyncClient | None = None, + ) -> httpx.Response: + """POST to vLLM, retrying failures that are safe to replay. + + Connection establishment failures and gateway/service-unavailable + responses happen before the local runtime can reliably process a LoRA + reload. Read/write failures are not retried because the server may have + already applied a non-idempotent request. + """ + if max_attempts < 1: + raise ValueError("max_attempts must be at least 1") + if retry_base_delay < 0: + raise ValueError("retry_base_delay must be non-negative") + + url = f"{self.base_url}{path}" + owns_client = http_client is None + client = http_client or httpx.AsyncClient() + last_error: Exception | None = None + try: + for attempt in range(1, max_attempts + 1): + if before_attempt is not None: + before_attempt() + try: + response = await client.post( + url, + json=json, + **self.request_kwargs(), + timeout=timeout, + ) + except (httpx.ConnectError, httpx.ConnectTimeout) as exc: + last_error = exc + else: + if response.status_code not in TRANSIENT_VLLM_STATUS_CODES: + response.raise_for_status() + return response + try: + response.raise_for_status() + except httpx.HTTPStatusError as exc: + last_error = exc + await response.aclose() + + if attempt == max_attempts: + break + logger.warning( + "%s failed on attempt %d/%d; retrying in %.1fs: %s", + operation, + attempt, + max_attempts, + retry_base_delay * (2 ** (attempt - 1)), + last_error, + ) + await asyncio.sleep(retry_base_delay * (2 ** (attempt - 1))) + finally: + if owns_client: + await client.aclose() + + assert last_error is not None + details = [ + f"{operation} failed after {max_attempts} attempts", + f"url={url}", + f"process={self._process_status()}", + f"log={self.log_path or 'unavailable'}", + f"last_error={last_error}", + ] + if failure_context is not None: + details.append(failure_context) + raise RuntimeError("; ".join(details)) from last_error + async def start( self, *, diff --git a/tests/unit/test_vllm_runtime.py b/tests/unit/test_vllm_runtime.py new file mode 100644 index 000000000..346c3b670 --- /dev/null +++ b/tests/unit/test_vllm_runtime.py @@ -0,0 +1,165 @@ +from __future__ import annotations + +import httpx +import pytest + +from art.vllm_runtime import ManagedVllmRuntime + + +@pytest.mark.asyncio +async def test_post_with_retry_recovers_from_connect_errors() -> None: + attempts = 0 + process_checks = 0 + + async def handler(request: httpx.Request) -> httpx.Response: + nonlocal attempts + attempts += 1 + if attempts < 3: + raise httpx.ConnectError("runtime restarting", request=request) + return httpx.Response(200, json={"loaded": True}) + + def check_process() -> None: + nonlocal process_checks + process_checks += 1 + + runtime = ManagedVllmRuntime() + runtime.port = 8000 + client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + try: + response = await runtime.post_with_retry( + "/v1/load_lora_adapter", + json={"lora_name": "model@1"}, + timeout=1.0, + max_attempts=3, + retry_base_delay=0, + before_attempt=check_process, + http_client=client, + ) + finally: + await client.aclose() + + assert response.status_code == 200 + assert attempts == 3 + assert process_checks == 3 + + +@pytest.mark.asyncio +async def test_post_with_retry_stops_when_runtime_process_fails() -> None: + attempts = 0 + process_checks = 0 + + async def handler(request: httpx.Request) -> httpx.Response: + nonlocal attempts + attempts += 1 + raise httpx.ConnectError("runtime restarting", request=request) + + def check_process() -> None: + nonlocal process_checks + process_checks += 1 + if process_checks == 2: + raise RuntimeError("vLLM runtime exited with code 9") + + runtime = ManagedVllmRuntime() + runtime.port = 8000 + client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + try: + with pytest.raises(RuntimeError, match="vLLM runtime exited with code 9"): + await runtime.post_with_retry( + "/v1/load_lora_adapter", + timeout=1.0, + retry_base_delay=0, + before_attempt=check_process, + http_client=client, + ) + finally: + await client.aclose() + + assert attempts == 1 + assert process_checks == 2 + + +@pytest.mark.asyncio +async def test_post_with_retry_recovers_from_transient_status() -> None: + attempts = 0 + + async def handler(request: httpx.Request) -> httpx.Response: + nonlocal attempts + attempts += 1 + if attempts == 1: + return httpx.Response(503, text="runtime unavailable") + return httpx.Response(200, json={"loaded": True}) + + runtime = ManagedVllmRuntime() + runtime.port = 8000 + client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + try: + response = await runtime.post_with_retry( + "/v1/load_lora_adapter", + timeout=1.0, + max_attempts=2, + retry_base_delay=0, + http_client=client, + ) + finally: + await client.aclose() + + assert response.status_code == 200 + assert attempts == 2 + + +@pytest.mark.asyncio +async def test_post_with_retry_does_not_retry_client_errors() -> None: + attempts = 0 + + async def handler(request: httpx.Request) -> httpx.Response: + nonlocal attempts + attempts += 1 + return httpx.Response(400, text="adapter already loaded") + + runtime = ManagedVllmRuntime() + runtime.port = 8000 + client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + try: + with pytest.raises(httpx.HTTPStatusError): + await runtime.post_with_retry( + "/v1/load_lora_adapter", + timeout=1.0, + retry_base_delay=0, + http_client=client, + ) + finally: + await client.aclose() + + assert attempts == 1 + + +@pytest.mark.asyncio +async def test_post_with_retry_reports_runtime_diagnostics() -> None: + async def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("runtime unavailable", request=request) + + runtime = ManagedVllmRuntime() + runtime.port = 8123 + runtime.log_path = "/tmp/art/vllm-runtime.log" + client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + try: + with pytest.raises(RuntimeError) as exc_info: + await runtime.post_with_retry( + "/v1/load_lora_adapter", + timeout=1.0, + max_attempts=2, + retry_base_delay=0, + operation="LoRA adapter reload", + failure_context="step=7; checkpoint=/tmp/checkpoints/0007", + http_client=client, + ) + finally: + await client.aclose() + + message = str(exc_info.value) + assert "LoRA adapter reload failed after 2 attempts" in message + assert "url=http://127.0.0.1:8123/v1/load_lora_adapter" in message + assert "process=not managed by this process" in message + assert "log=/tmp/art/vllm-runtime.log" in message + assert "last_error=runtime unavailable" in message + assert "step=7; checkpoint=/tmp/checkpoints/0007" in message