chore: vendor sglang v0.5.10 snapshot

third_party/sglang/sgl-model-gateway/e2e_test/infra/process_utils.py

@@ -0,0 +1,161 @@
+"""Process management utilities for E2E tests."""
+
+from __future__ import annotations
+
+import logging
+import os
+import signal
+import subprocess
+import time
+
+import requests
+
+logger = logging.getLogger(__name__)
+
+
+def kill_process_tree(pid: int, sig: int = signal.SIGTERM) -> None:
+    """Kill a process and all its children.
+
+    Args:
+        pid: Process ID to kill
+        sig: Signal to send (default: SIGTERM)
+    """
+    try:
+        import psutil
+
+        parent = psutil.Process(pid)
+        children = parent.children(recursive=True)
+        for child in children:
+            try:
+                child.send_signal(sig)
+            except psutil.NoSuchProcess:
+                pass
+        parent.send_signal(sig)
+    except ImportError:
+        # Fallback if psutil not available
+        os.kill(pid, sig)
+    except Exception as e:
+        logger.warning("Failed to kill process tree for PID %d: %s", pid, e)
+
+
+def terminate_process(proc: subprocess.Popen, timeout: float = 30) -> None:
+    """Gracefully terminate a process, kill if needed.
+
+    Args:
+        proc: Process to terminate
+        timeout: Seconds to wait before force-killing
+    """
+    if proc is None or proc.poll() is not None:
+        return
+    proc.terminate()
+    start = time.perf_counter()
+    while proc.poll() is None:
+        if time.perf_counter() - start > timeout:
+            proc.kill()
+            break
+        time.sleep(1)
+
+
+def wait_for_health(
+    url: str,
+    timeout: float = 60,
+    api_key: str | None = None,
+    check_interval: float = 1.0,
+) -> None:
+    """Wait for a server's /health endpoint to return 200.
+
+    Args:
+        url: Base URL of the server
+        timeout: Seconds to wait before timing out
+        api_key: Optional API key for auth header
+        check_interval: Seconds between health checks
+    """
+    start = time.perf_counter()
+    headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
+
+    with requests.Session() as session:
+        while time.perf_counter() - start < timeout:
+            try:
+                resp = session.get(f"{url}/health", headers=headers, timeout=5)
+                if resp.status_code == 200:
+                    logger.info("Service healthy at %s", url)
+                    return
+            except requests.RequestException:
+                pass
+            time.sleep(check_interval)
+
+    raise TimeoutError(f"Server at {url} did not become healthy within {timeout}s")
+
+
+def wait_for_workers_ready(
+    router_url: str,
+    expected_workers: int,
+    timeout: float = 300,
+    api_key: str | None = None,
+) -> None:
+    """Wait for router to have all workers connected.
+
+    Args:
+        router_url: Base URL of the router
+        expected_workers: Number of workers to wait for
+        timeout: Seconds to wait before timing out
+        api_key: Optional API key for auth header
+    """
+    start = time.perf_counter()
+    headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
+
+    while time.perf_counter() - start < timeout:
+        try:
+            resp = requests.get(f"{router_url}/workers", headers=headers, timeout=5)
+            if resp.status_code == 200:
+                data = resp.json()
+                total = data.get("total", len(data.get("workers", [])))
+                if total >= expected_workers:
+                    logger.info(
+                        "All %d workers connected after %.1fs",
+                        expected_workers,
+                        time.perf_counter() - start,
+                    )
+                    return
+        except requests.RequestException:
+            pass
+        time.sleep(2)
+
+    raise TimeoutError(
+        f"Router at {router_url} did not get {expected_workers} workers within {timeout}s"
+    )
+
+
+def detect_ib_device() -> str | None:
+    """Detect first active InfiniBand device (e.g., mlx5_0).
+
+    Returns:
+        Device name if found (e.g., "mlx5_0"), None otherwise.
+    """
+    try:
+        subprocess.run(
+            ["ibv_devinfo", "-l"],
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+            timeout=1,
+        )
+    except (FileNotFoundError, subprocess.TimeoutExpired):
+        return None
+
+    for i in range(12):
+        dev = f"mlx5_{i}"
+        try:
+            res = subprocess.run(
+                ["ibv_devinfo", dev],
+                capture_output=True,
+                text=True,
+                timeout=2,
+            )
+            if res.returncode == 0 and "state:" in res.stdout:
+                for line in res.stdout.splitlines():
+                    if "state:" in line and "PORT_ACTIVE" in line:
+                        logger.info("Detected IB device: %s", dev)
+                        return dev
+        except Exception:
+            pass
+    return None