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