feat: TCP keepalive, reverse-channel probe, list_models cache, .pth atomic write, launch exemption

xinference/constants.py

@@ -121,11 +121,18 @@ def get_xinference_home() -> str:
 )
 # get_model -> worker_ref.get_model RPC
 XINFERENCE_GET_MODEL_RPC_TIMEOUT = int(
-    os.environ.get(XINFERENCE_ENV_GET_MODEL_RPC_TIMEOUT, "120")
+    os.environ.get(XINFERENCE_ENV_GET_MODEL_RPC_TIMEOUT, "30")
 )
 XINFERENCE_DISABLE_HEALTH_CHECK = bool(
     int(os.environ.get(XINFERENCE_ENV_DISABLE_HEALTH_CHECK, 0))
 )
+
+# Max concurrent download threads for huggingface_hub.snapshot_download.
+# Default 8 in hf_hub causes GIL contention that starves the actor event loop.
+XINFERENCE_ENV_MODEL_DOWNLOAD_WORKERS = "XINFERENCE_MODEL_DOWNLOAD_WORKERS"
+XINFERENCE_MODEL_DOWNLOAD_WORKERS = int(
+    os.environ.get(XINFERENCE_ENV_MODEL_DOWNLOAD_WORKERS, 2)
+)
 XINFERENCE_DISABLE_METRICS = bool(
     int(os.environ.get(XINFERENCE_ENV_DISABLE_METRICS, 0))
 )

xinference/core/supervisor.py

@@ -127,6 +127,14 @@ def __init__(self):
         self._model_uid_to_replica_info: Dict[str, ReplicaInfo] = {}  # type: ignore
         self._uptime = None
         self._lock = asyncio.Lock()
+        # list_models cache for graceful degradation when a worker is unreachable
+        self._list_models_cache: Dict[str, Dict[str, Dict[str, Any]]] = {}
+        # Reverse-ping failure counter per worker
+        self._reverse_ping_failures: Dict[str, int] = {}
+        # Track workers currently launching models — when a worker has active
+        # launches, reverse-channel dead detection is exempted because
+        # long-running model downloads can starve the actor event loop.
+        self._workers_launching: Dict[str, int] = {}  # address -> active launch count
 
     @classmethod
     def default_uid(cls) -> str:
@@ -1554,6 +1562,11 @@ async def _launch_one_model(
             # LLM as default for compatibility
             model_type = model_type or "LLM"
 
+            # Track this worker as launching so reverse-channel dead
+            # detection is exempted during long model downloads.
+            _addr = worker_ref.address
+            self._workers_launching[_addr] = self._workers_launching.get(_addr, 0) + 1
+
             try:
                 subpool_address = await worker_ref.launch_builtin_model(
                     model_uid=_replica_model_uid,
@@ -1591,6 +1604,13 @@ async def _launch_one_model(
                     {"status": LaunchStatus.ERROR.name, "error_message": str(e)},
                 )
                 raise
+            finally:
+                # Decrement launching counter
+                _cnt = self._workers_launching.get(_addr, 1) - 1
+                if _cnt <= 0:
+                    self._workers_launching.pop(_addr, None)
+                else:
+                    self._workers_launching[_addr] = _cnt
 
         async def _launch_model():
             try:
@@ -2102,6 +2122,77 @@ async def _check_dead_nodes(self):
                 for address in dead_nodes:
                     self._worker_status.pop(address, None)
                     self._worker_address_to_worker.pop(address, None)
+
+                # ---- Reverse-channel probe ----
+                # Heartbeat only covers worker->supervisor; this probes
+                # supervisor->worker so we detect dead reverse channels
+                # even when the worker process is alive.
+                for address, worker_ref in list(self._worker_address_to_worker.items()):
+                    if address in dead_nodes:
+                        self._reverse_ping_failures.pop(address, None)
+                        continue
+                    try:
+                        await xo.wait_for(worker_ref.ping(), timeout=5)
+                        # Reset on success
+                        self._reverse_ping_failures.pop(address, None)
+                    except Exception:
+                        # If the worker is launching a model, exempt it from
+                        # reverse-channel dead detection. Long-running
+                        # downloads can starve the actor event loop for
+                        # minutes, but the worker is still alive.
+                        # Heartbeat path still detects real process crashes.
+                        if address in self._workers_launching:
+                            logger.warning(
+                                "Worker reverse-channel timeout. address: %s "
+                                "(launching model, reverse-channel check skipped)",
+                                address,
+                            )
+                            # Do NOT accumulate failure count; reset it so
+                            # that when launching finishes, the worker starts
+                            # with a clean slate.
+                            self._reverse_ping_failures.pop(address, None)
+                            continue
+
+                        count = self._reverse_ping_failures.get(address, 0) + 1
+                        self._reverse_ping_failures[address] = count
+                        if count >= XINFERENCE_HEALTH_CHECK_FAILURE_THRESHOLD:
+                            # Treat as dead — same cleanup as heartbeat failure
+                            status = self._worker_status.get(address)
+                            if status is not None:
+                                status.failure_remaining_count = 0
+                            dead_models = []
+                            for model_uid in self._replica_model_uid_to_worker:
+                                worker_refs = self._replica_model_uid_to_worker[
+                                    model_uid
+                                ]
+                                if not isinstance(worker_refs, (list, tuple)):
+                                    worker_refs = [worker_refs]
+                                for wr in worker_refs:
+                                    if wr.address == address:
+                                        dead_models.append(model_uid)
+                            logger.error(
+                                "Worker reverse-channel dead. address: %s, "
+                                "influenced models: %s",
+                                address,
+                                dead_models,
+                            )
+                            for replica_model_uid in dead_models:
+                                model_uid, _ = parse_replica_model_uid(
+                                    replica_model_uid
+                                )
+                                self._model_uid_to_replica_info.pop(model_uid, None)
+                                self._replica_model_uid_to_worker.pop(
+                                    replica_model_uid, None
+                                )
+                            dead_nodes.append(address)
+                            self._reverse_ping_failures.pop(address, None)
+                        else:
+                            logger.error(
+                                "Worker reverse-channel timeout. address: %s, "
+                                "check count remaining %s...",
+                                address,
+                                XINFERENCE_HEALTH_CHECK_FAILURE_THRESHOLD - count,
+                            )
             finally:
                 await asyncio.sleep(XINFERENCE_HEALTH_CHECK_INTERVAL)
 
@@ -2332,17 +2423,30 @@ async def _fetch_one(
             worker_address: str, worker_ref: xo.ActorRefType["WorkerActor"]
         ) -> Dict[str, Dict[str, Any]]:
             try:
-                return await xo.wait_for(
+                result = await xo.wait_for(
                     worker_ref.list_models(),
                     XINFERENCE_LIST_MODELS_PER_WORKER_TIMEOUT,
                 )
+                # Update cache on success
+                self._list_models_cache[worker_address] = result
+                return result
             except Exception as ex:
-                logger.warning(
-                    "list_models from worker %s failed or timed out: %s",
-                    worker_address,
-                    ex,
-                )
-                return {}
+                cached = self._list_models_cache.get(worker_address, {})
+                if cached:
+                    logger.warning(
+                        "list_models from worker %s failed or timed out: %s, "
+                        "returning cached result (%d models)",
+                        worker_address,
+                        ex,
+                        len(cached),
+                    )
+                else:
+                    logger.warning(
+                        "list_models from worker %s failed or timed out: %s",
+                        worker_address,
+                        ex,
+                    )
+                return cached
 
         parts = await asyncio.gather(*(_fetch_one(addr, ref) for addr, ref in workers))
         for part in parts:

xinference/core/worker.py

@@ -53,6 +53,7 @@
     XINFERENCE_DISABLE_METRICS,
     XINFERENCE_ENABLE_VIRTUAL_ENV,
     XINFERENCE_HEALTH_CHECK_INTERVAL,
+    XINFERENCE_MODEL_DOWNLOAD_WORKERS,
     XINFERENCE_STATUS_GATHER_TIMEOUT,
     XINFERENCE_STATUS_REPORT_MULTIPLIER,
     XINFERENCE_TCP_REQUEST_TIMEOUT,
@@ -1467,12 +1468,38 @@ def _create_virtual_env_manager(
                     )
                     child_site_packages.mkdir(parents=True, exist_ok=True)
                     pth_file = child_site_packages / "_xinference_parent.pth"
-                    pth_file.write_text(parent_site_packages + "\n")
-                    logger.debug(
-                        "Injected parent site-packages into child venv via %s -> %s",
-                        pth_file,
-                        parent_site_packages,
-                    )
+                    desired_content = parent_site_packages + "\n"
+                    # Avoid truncate race when multiple replicas write the
+                    # same .pth concurrently: skip if content already correct,
+                    # otherwise atomic-write via a temp file + os.replace.
+                    needs_write = True
+                    try:
+                        if (
+                            pth_file.exists()
+                            and pth_file.read_text() == desired_content
+                        ):
+                            needs_write = False
+                    except OSError:
+                        pass
+                    if needs_write:
+                        tmp_file = pth_file.with_suffix(".pth.tmp")
+                        try:
+                            tmp_file.write_text(desired_content)
+                            os.replace(str(tmp_file), str(pth_file))
+                        except OSError:
+                            # Fallback: direct write (still better than no .pth)
+                            pth_file.write_text(desired_content)
+                        logger.debug(
+                            "Injected parent site-packages into child venv "
+                            "via %s -> %s",
+                            pth_file,
+                            parent_site_packages,
+                        )
+                    else:
+                        logger.debug(
+                            "Skipped .pth write (content unchanged): %s",
+                            pth_file,
+                        )
                 else:
                     logger.warning(
                         "Parent site-packages path does not exist: %s — child venv "
@@ -1820,20 +1847,32 @@ async def launch_builtin_model(
                                 self._upload_download_progress, progressor, downloader
                             )
                         )
-                        model = await asyncio.to_thread(
-                            create_model_instance,
-                            model_uid,
-                            model_type,
-                            model_name,
-                            model_engine,
-                            model_format,
-                            model_size_in_billions,
-                            quantization,
-                            peft_model_config,
-                            download_hub,
-                            model_path,
-                            **model_kwargs,
+                        # Limit hf_hub download concurrency to reduce GIL
+                        # contention that starves the event loop.
+                        _orig_hf_workers = os.environ.get("HF_HUB_DOWNLOAD_WORKERS")
+                        os.environ["HF_HUB_DOWNLOAD_WORKERS"] = str(
+                            XINFERENCE_MODEL_DOWNLOAD_WORKERS
                         )
+                        try:
+                            model = await asyncio.to_thread(
+                                create_model_instance,
+                                model_uid,
+                                model_type,
+                                model_name,
+                                model_engine,
+                                model_format,
+                                model_size_in_billions,
+                                quantization,
+                                peft_model_config,
+                                download_hub,
+                                model_path,
+                                **model_kwargs,
+                            )
+                        finally:
+                            if _orig_hf_workers is not None:
+                                os.environ["HF_HUB_DOWNLOAD_WORKERS"] = _orig_hf_workers
+                            else:
+                                os.environ.pop("HF_HUB_DOWNLOAD_WORKERS", None)
                     model.model_family.address = subpool_address
                     model.model_family.accelerators = devices
                     model.model_family.multimodal_projector = model_kwargs.get(
@@ -2173,6 +2212,10 @@ async def report_status(self):
             supervisor_ref = await self.get_supervisor_ref(add_worker=True)
             await supervisor_ref.report_worker_status(self.address, status)
 
+    async def ping(self) -> bool:
+        """Lightweight liveness probe for supervisor reverse-channel check."""
+        return True
+
     async def heartbeat(self):
         """
         Lightweight heartbeat for liveness detection.

xinference/core/xoscar_keepalive.py

@@ -0,0 +1,70 @@
+"""
+Monkey-patch xoscar SocketClient.connect to enable TCP keepalive on every
+RPC connection.
+
+Without this, idle TCP connections between Supervisor and Worker are silently
+dropped by stateful network devices (firewalls, NAT gateways) whose TCP
+session timeout (commonly 15-45 min) is shorter than the application idle
+window.  xoscar's ``Router._cache`` only checks ``writer.is_closing()``
+which stays ``False`` after a silent drop, so the dead connection is reused
+and the next RPC hangs until OS-level TCP retransmission timeout (60-120 s).
+
+Enabling ``SO_KEEPALIVE`` with ``TCP_KEEPIDLE=60 / TCP_KEEPINTVL=10 /
+TCP_KEEPCNT=3`` causes the kernel to probe idle connections every 60 s and
+declare them dead within 90 s -- well below typical firewall timeouts.
+"""
+
+import logging
+import socket
+
+logger = logging.getLogger(__name__)
+
+_patched = False
+
+
+def patch_xoscar_socket_keepalive():
+    """Apply the keepalive patch exactly once (idempotent)."""
+    global _patched
+    if _patched:
+        return
+    _patched = True
+
+    try:
+        from xoscar.backends.communication.socket import SocketClient
+    except ImportError:
+        logger.debug("xoscar not installed, skipping keepalive patch")
+        return
+
+    _original_connect = SocketClient.connect
+
+    @staticmethod
+    async def _patched_connect(dest_address, local_address=None, **kwargs):
+        client = await _original_connect(
+            dest_address, local_address=local_address, **kwargs
+        )
+        try:
+            sock = client.channel.writer.get_extra_info("socket")
+            if sock is not None:
+                sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
+                if hasattr(socket, "TCP_KEEPIDLE"):
+                    # Linux
+                    sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, 60)
+                    sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 10)
+                    sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 3)
+                elif hasattr(socket, "TCP_KEEPALIVE"):
+                    # macOS: TCP_KEEPALIVE is the equivalent of TCP_KEEPIDLE
+                    sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPALIVE, 60)
+        except Exception:
+            logger.debug(
+                "Failed to set TCP keepalive on xoscar connection to %s",
+                dest_address,
+                exc_info=True,
+            )
+        return client
+
+    SocketClient.connect = _patched_connect
+    logger.info("xoscar SocketClient.connect patched with TCP keepalive")
+
+
+# Auto-apply on import
+patch_xoscar_socket_keepalive()

xinference/deploy/local.py

@@ -24,6 +24,8 @@
 import xoscar as xo
 from xoscar.utils import get_next_port
 
+import xinference.core.xoscar_keepalive  # noqa: F401  TCP keepalive patch
+
 from ..constants import (
     XINFERENCE_HEALTH_CHECK_FAILURE_THRESHOLD,
     XINFERENCE_HEALTH_CHECK_INTERVAL,

xinference/deploy/supervisor.py

@@ -22,6 +22,8 @@
 import xoscar as xo
 from xoscar.utils import get_next_port
 
+import xinference.core.xoscar_keepalive  # noqa: F401  TCP keepalive patch
+
 from ..constants import (
     XINFERENCE_HEALTH_CHECK_FAILURE_THRESHOLD,
     XINFERENCE_HEALTH_CHECK_INTERVAL,

xinference/deploy/worker.py

@@ -21,6 +21,8 @@
 import xoscar as xo
 from xoscar import MainActorPoolType
 
+import xinference.core.xoscar_keepalive  # noqa: F401  TCP keepalive patch
+
 from ..core.worker import WorkerActor
 from ..device_utils import get_available_device_env_name, gpu_count