fix: retry transient vLLM adapter reload failures

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}"

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,
         *,

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