autoinference-utils 0.2.9__tar.gz → 0.2.10__tar.gz

{autoinference_utils-0.2.9 → autoinference_utils-0.2.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: autoinference-utils
-Version: 0.2.9
+Version: 0.2.10
 Summary: Shared endpoint abstractions for autoinference deployments
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown

{autoinference_utils-0.2.9 → autoinference_utils-0.2.10}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "autoinference-utils"
-version = "0.2.9"
+version = "0.2.10"
 description = "Shared endpoint abstractions for autoinference deployments"
 readme = "README.md"
 requires-python = ">=3.10"

{autoinference_utils-0.2.9 → autoinference_utils-0.2.10}/src/autoinference_utils/pd.py

@@ -5,6 +5,8 @@ from __future__ import annotations
 import asyncio
 import math
 import os
+import socket
+import struct
 import subprocess
 import threading
 import time
@@ -40,6 +42,9 @@ _MANAGED_ROUTER_FLAGS = frozenset(
     }
 )
 
+# Leave headroom inside Modal's 30-second exit-handler deadline.
+_SHUTDOWN_TIMEOUT = 25
+
 
 def _managed_router_tokens(router_args: Mapping[str, str]) -> list[str]:
     """Rendered override tokens that would reach a managed flag.
@@ -177,6 +182,7 @@ class PDEndpoint(Endpoint):
         health_failure_timeout: float = 30,
         router_args: Mapping[str, str] | None = None,
         allow_custom_router: bool = False,
+        shutdown_port: int | None = None,
     ):
         from modal.experimental import get_cluster_info
 
@@ -217,6 +223,13 @@ class PDEndpoint(Endpoint):
                 f"P/D ratio {ratio} requires a {len(roles)}-container Modal cluster"
             )
         self._host_ip = hosts[cluster.rank]
+        self._hosts = hosts
+        self._rank = cluster.rank
+        self._shutdown_port = (
+            router_port + 1 if shutdown_port is None else shutdown_port
+        )
+        if not 1 <= self._shutdown_port <= 65535:
+            raise ValueError("shutdown_port must be between 1 and 65535")
         self.role = roles[cluster.rank]
         super().__init__(_url(hosts[0], router_port))
         self.engine = engine
@@ -266,6 +279,8 @@ class PDEndpoint(Endpoint):
         self._stopped = threading.Event()
         self._lock = threading.Lock()
         self._started = False
+        self._serving = False
+        self._shutdown_listener: socket.socket | None = None
         self._threads: list[threading.Thread] = []
 
     def start(self, *, warmup: Callable[[], None] | None = None) -> None:
@@ -280,6 +295,13 @@ class PDEndpoint(Endpoint):
             SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION="1",
         )
         try:
+            if self._rank == 0:
+                self._shutdown_listener = socket.socket(socket.AF_INET6)
+                self._shutdown_listener.setsockopt(
+                    socket.SOL_SOCKET, socket.SO_REUSEADDR, 1
+                )
+                self._shutdown_listener.bind(("::", self._shutdown_port))
+                self._shutdown_listener.listen(len(self._hosts) - 1)
             for endpoint in (self.engine, self.router):
                 if endpoint is not None:
                     endpoint.start()
@@ -291,6 +313,7 @@ class PDEndpoint(Endpoint):
                 for _, host in self.router.pd_config:
                     self._watch_health(_url(host, self.engine.worker_port) + "/health")
                 self._watch_health(self.router.base_url + "/health")
+            self._serving = True
         except BaseException:
             self.stop()
             raise
@@ -352,7 +375,7 @@ class PDEndpoint(Endpoint):
             if self._stopped.is_set():
                 return
             self._stopped.set()
-        deadline = time.monotonic() + self.drain_timeout + 30
+        deadline = time.monotonic() + _SHUTDOWN_TIMEOUT
         try:
             for endpoint in (self.router, self.engine):
                 if endpoint is None:
@@ -364,9 +387,58 @@ class PDEndpoint(Endpoint):
                         process.wait(timeout=max(0, deadline - time.monotonic()))
                     except subprocess.TimeoutExpired:
                         process.kill()
+                        process.wait()
                 endpoint.stop()
         finally:
             for thread in self._threads:
-                thread.join(
-                    timeout=self.engine.health_request_timeout + self.health_interval
+                thread.join(timeout=max(0, deadline - time.monotonic()))
+            if self._serving:
+                if self._rank == 0:
+                    self._wait_for_followers(deadline)
+                else:
+                    self._notify_leader(deadline)
+            if self._shutdown_listener is not None:
+                self._shutdown_listener.close()
+                self._shutdown_listener = None
+
+    def _notify_leader(self, deadline: float) -> None:
+        while True:
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                print("[pd] Failed to notify rank zero of shutdown", flush=True)
+                return
+            try:
+                with socket.create_connection(
+                    (self._hosts[0], self._shutdown_port), timeout=min(remaining, 1)
+                ) as connection:
+                    connection.sendall(struct.pack("!I", self._rank))
+                return
+            except OSError:
+                time.sleep(min(0.1, max(0, deadline - time.monotonic())))
+
+    def _wait_for_followers(self, deadline: float) -> None:
+        listener = self._shutdown_listener
+        if listener is None:
+            return
+        pending = set(range(1, len(self._hosts)))
+        while pending:
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                print(
+                    f"[pd] Timed out waiting for ranks {sorted(pending)} to stop",
+                    flush=True,
                 )
+                return
+            listener.settimeout(min(remaining, 1))
+            try:
+                connection, _ = listener.accept()
+            except TimeoutError:
+                continue
+            with connection:
+                connection.settimeout(min(remaining, 1))
+                try:
+                    payload = connection.recv(4, socket.MSG_WAITALL)
+                except TimeoutError:
+                    continue
+            if len(payload) == 4:
+                pending.discard(struct.unpack("!I", payload)[0])