chore: import upstream snapshot with attribution

python/sglang/cli/killall.py

@@ -0,0 +1,457 @@
+#!/usr/bin/env python3
+"""Kill SGLang processes on CUDA_VISIBLE_DEVICES GPUs (CI mode only).
+
+Called at the start of every CI job to clean up orphaned processes from
+previous (possibly cancelled) runs. Requires SGLANG_IS_IN_CI=true.
+
+For local/non-CI usage, use scripts/killall_sglang.sh instead.
+
+Usage:
+    python killall.py
+
+Exit codes:
+    0 - Clean: all target GPUs have <10% memory usage after cleanup
+    1 - Dirty: GPU memory still >10% after cleanup, indicating stuck processes
+        or orphaned CUDA contexts that need a container restart
+"""
+
+import os
+import re
+import signal
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+# Constants
+MEMORY_THRESHOLD_PCT = 10
+
+# Patterns matching SGLang process command lines (equivalent to pgrep -f in killall_sglang.sh)
+_SGLANG_PROCESS_PATTERNS = re.compile(
+    r"sglang::|sglang\.launch_server|sglang\.bench|sglang\.data_parallel|sglang\.srt|sgl_diffusion::|sglang serve"
+)
+
+# Boxed output helpers
+_LOG_LINES = []
+
+
+def _log(msg=""):
+    """Buffer a line for boxed output."""
+    _LOG_LINES.append(msg)
+
+
+def _flush_box(title, status=""):
+    """Print all buffered lines inside a box, then clear buffer."""
+    lines = _LOG_LINES.copy()
+    _LOG_LINES.clear()
+
+    all_text = [title] + ([status] if status else []) + lines
+    width = max((len(line) for line in all_text), default=40) + 4
+    width = max(width, 60)
+
+    h_bar = "─" * (width - 2)
+    print(f"\n┌{h_bar}┐")
+    print(f"│ {title:<{width - 3}}│")
+    print(f"├{h_bar}┤")
+    for line in lines:
+        print(f"│ {line:<{width - 3}}│")
+    if status:
+        print(f"├{h_bar}┤")
+        print(f"│ {status:<{width - 3}}│")
+    print(f"└{h_bar}┘")
+
+
+# nvidia-smi helpers
+def _run_smi(query, query_type="gpu"):
+    """Run nvidia-smi query and return raw CSV lines."""
+    flag = "--query-gpu" if query_type == "gpu" else "--query-compute-apps"
+    try:
+        out = subprocess.check_output(
+            ["nvidia-smi", f"{flag}={query}", "--format=csv,noheader,nounits"],
+            text=True,
+            timeout=10,
+        )
+        return [line.strip() for line in out.strip().splitlines() if line.strip()]
+    except (subprocess.SubprocessError, FileNotFoundError):
+        return []
+
+
+def _get_smi_version():
+    """Return nvidia-smi driver version and GPU name, or None on failure."""
+    # Inline nvidia-smi query — killall.py runs before pip install, so sglang
+    # internals may not be importable.
+    try:
+        result = subprocess.run(
+            [
+                "nvidia-smi",
+                "--query-gpu=driver_version",
+                "--format=csv,noheader,nounits",
+            ],
+            capture_output=True,
+            text=True,
+            check=True,
+            timeout=10,
+        )
+        driver = result.stdout.strip().split("\n")[0].strip() or None
+    except (subprocess.SubprocessError, FileNotFoundError):
+        driver = None
+    if driver is None:
+        return None
+    try:
+        out = subprocess.check_output(
+            ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
+            text=True,
+            timeout=10,
+        )
+        gpu_name = out.strip().splitlines()[0].strip() if out.strip() else "unknown"
+    except (subprocess.SubprocessError, FileNotFoundError, IndexError):
+        gpu_name = "unknown"
+    return f"driver {driver}, {gpu_name}"
+
+
+def _get_target_gpus():
+    """Return GPU indices from CUDA_VISIBLE_DEVICES, or all visible GPUs.
+
+    Note: only numeric indices are supported (e.g. "0,1,2").
+    UUID-style CUDA_VISIBLE_DEVICES values (e.g. "GPU-d4f1...") are not handled.
+    """
+    cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
+    if cvd is not None and cvd.strip():
+        return {int(g.strip()) for g in cvd.split(",") if g.strip().isdigit()}
+    return {int(line) for line in _run_smi("index") if line.isdigit()}
+
+
+def _get_gpu_pids(gpu_indices):
+    """Return PIDs using the specified GPUs (by index)."""
+    target_uuids = set()
+    for line in _run_smi("index,uuid"):
+        parts = line.split(",", 1)
+        if len(parts) == 2 and parts[0].strip().isdigit():
+            if int(parts[0].strip()) in gpu_indices:
+                target_uuids.add(parts[1].strip())
+    pids = set()
+    for line in _run_smi("gpu_uuid,pid", query_type="apps"):
+        parts = line.split(",", 1)
+        if len(parts) == 2 and parts[0].strip() in target_uuids:
+            pid = parts[1].strip()
+            if pid.isdigit():
+                pids.add(int(pid))
+    return pids
+
+
+def _get_gpu_memory(gpu_indices):
+    """Query memory usage for target GPUs.
+
+    Returns list of (idx, used_mib, total_mib, pct) tuples.
+    """
+    result = []
+    for line in _run_smi("index,memory.used,memory.total"):
+        parts = line.split(",")
+        if len(parts) != 3 or not parts[0].strip().isdigit():
+            continue
+        idx = int(parts[0].strip())
+        if idx not in gpu_indices:
+            continue
+        try:
+            used, total = int(float(parts[1].strip())), int(float(parts[2].strip()))
+        except ValueError:
+            continue
+        pct = used / total * 100 if total > 0 else 0
+        result.append((idx, used, total, pct))
+    return result
+
+
+def _get_dirty_gpus(gpu_indices):
+    """Return list of dirty GPU description strings (memory >= threshold)."""
+    return [
+        f"GPU {idx} ({pct:.0f}%)"
+        for idx, _, _, pct in _get_gpu_memory(gpu_indices)
+        if pct >= MEMORY_THRESHOLD_PCT
+    ]
+
+
+def _log_gpu_memory(gpu_indices):
+    """Log memory usage for all target GPUs and return dirty GPU descriptions."""
+    dirty = []
+    for idx, used, total, pct in _get_gpu_memory(gpu_indices):
+        _log(f"  GPU {idx}: {used} MiB / {total} MiB ({pct:.0f}%)")
+        if pct >= MEMORY_THRESHOLD_PCT:
+            dirty.append(f"GPU {idx} ({pct:.0f}%)")
+    return dirty
+
+
+# /proc helpers
+def _read_proc_cmdline(pid):
+    """Read /proc/{pid}/cmdline and return as decoded string, or None on failure."""
+    try:
+        raw = Path(f"/proc/{pid}/cmdline").read_bytes()
+        return raw.decode("utf-8", errors="replace").replace("\x00", " ")
+    except (FileNotFoundError, PermissionError):
+        return None
+
+
+def _get_pid_cmdline(pid):
+    """Get truncated command line for a PID."""
+    cmdline = _read_proc_cmdline(pid)
+    if cmdline is None:
+        return "<unknown>"
+    cmdline = cmdline.strip()
+    return cmdline[:120] + ("..." if len(cmdline) > 120 else "")
+
+
+def _find_sglang_pids_by_name():
+    """Find SGLang process PIDs by command-line pattern matching.
+
+    Scans /proc/*/cmdline for patterns matching known SGLang entry points.
+    Equivalent to: pgrep -f 'sglang::|sglang.launch_server|...'
+
+    Safe in shared-GPU containers: without --pid=host, /proc only exposes
+    processes in our own PID namespace, so this cannot kill other containers.
+    """
+    my_pid = os.getpid()
+    pids = set()
+    for entry in Path("/proc").iterdir():
+        if not entry.name.isdigit():
+            continue
+        pid = int(entry.name)
+        if pid <= 1 or pid == my_pid:
+            continue
+        cmdline = _read_proc_cmdline(pid)
+        if cmdline and _SGLANG_PROCESS_PATTERNS.search(cmdline):
+            pids.add(pid)
+    return pids
+
+
+def _check_pid_namespace(pid):
+    """Check if a PID is in our PID namespace. Linux-only via /proc."""
+    try:
+        my_ns = os.readlink("/proc/self/ns/pid")
+    except OSError:
+        return "unknown (can't read self ns)"
+    try:
+        target_ns = os.readlink(f"/proc/{pid}/ns/pid")
+    except FileNotFoundError:
+        return f"NOT in our namespace (pid not in /proc, self={my_ns})"
+    except PermissionError:
+        return "unknown (no permission to read ns)"
+    if my_ns == target_ns:
+        return f"same namespace ({my_ns})"
+    return f"DIFFERENT namespace (self={my_ns}, target={target_ns})"
+
+
+def _get_orchestrator_ancestors(pids):
+    """Walk process tree upward from PIDs, return ancestors that are test orchestrators.
+
+    Linux-only: reads /proc filesystem. Returns empty set on other platforms.
+    """
+    orchestrator_patterns = ["run_suite.py", "run_tests.py"]
+    ancestors, visited = set(), set()
+    for pid in pids:
+        current = pid
+        while current > 1 and current not in visited:
+            visited.add(current)
+            cmdline = _read_proc_cmdline(current)
+            if cmdline is None:
+                break
+            if any(p in cmdline for p in orchestrator_patterns):
+                ancestors.add(current)
+            try:
+                current = int(Path(f"/proc/{current}/stat").read_text().split()[3])
+            except (FileNotFoundError, PermissionError, IndexError, ValueError):
+                break
+    return ancestors
+
+
+# Kill & diagnostic helpers
+def _kill_pids(pids, label="", quiet=False):
+    """Send SIGKILL to PIDs, skipping self and init.
+
+    Returns dict of {pid: exception_name} for PIDs that could not be killed.
+    When quiet=True, does not log individual kill results.
+    """
+    my_pid = os.getpid()
+    pids = {p for p in pids if p != my_pid and p > 1}
+    if not pids:
+        return {}
+    if label and not quiet:
+        _log(f"  Killing {label}:")
+    failed = {}
+    for pid in sorted(pids):
+        try:
+            os.kill(pid, signal.SIGKILL)
+            if not quiet:
+                _log(f"    PID {pid}: killed ({_get_pid_cmdline(pid)})")
+        except (ProcessLookupError, PermissionError) as e:
+            failed[pid] = type(e).__name__
+            if not quiet:
+                _log(f"    PID {pid}: failed ({type(e).__name__})")
+    return failed
+
+
+def _get_ps_diagnostic():
+    """Return ps auxf output filtered for GPU/sglang-related processes."""
+    try:
+        out = subprocess.run(["ps", "auxf"], capture_output=True, text=True, timeout=5)
+        return [
+            line.strip()[:140]
+            for line in out.stdout.splitlines()
+            if any(k in line.lower() for k in ["sglang", "python", "cuda", "gpu"])
+        ][:20]
+    except (subprocess.SubprocessError, FileNotFoundError):
+        return []
+
+
+def _print_diagnostics(unkillable_pids):
+    """Print detailed diagnostics after the FAIL box (to stdout, outside box)."""
+    if unkillable_pids:
+        print("\n[killall] Diagnostic — unkillable PIDs:")
+        for pid in sorted(unkillable_pids):
+            ns_info = _check_pid_namespace(pid)
+            print(f"  PID {pid}: ns: {ns_info}")
+    ps_lines = _get_ps_diagnostic()
+    if ps_lines:
+        print("\n[killall] Diagnostic — processes in this container (ps auxf):")
+        for line in ps_lines:
+            print(f"  {line}")
+    else:
+        print(
+            "\n[killall] Diagnostic — no sglang/python/gpu processes "
+            "in this container"
+        )
+
+
+# CI mode
+def _kill_all_targets(gpu_indices, gpu_pids):
+    """Kill all target processes: name-matched, orchestrator ancestors, GPU processes."""
+    # Kill name-matched SGLang processes (catches processes not visible to nvidia-smi)
+    name_only = _find_sglang_pids_by_name() - gpu_pids
+    if name_only:
+        _kill_pids(name_only, "name-matched SGLang processes")
+        time.sleep(1)
+        _log()
+
+    # Kill orchestrator ancestors first, then GPU processes (retry once)
+    if gpu_pids:
+        _kill_pids(_get_orchestrator_ancestors(gpu_pids), "orchestrator ancestors")
+        time.sleep(1)
+        for attempt in range(2):
+            current_pids = _get_gpu_pids(gpu_indices)
+            if not current_pids:
+                break
+            label = "GPU processes" if attempt == 0 else "stubborn GPU processes"
+            _kill_pids(current_pids, label)
+            time.sleep(3)
+    _log()
+
+
+def _verify_gpu_clean(gpu_indices):
+    """Retry loop: wait for GPUs to become clean.
+
+    Returns (dirty_list, unkillable_pids, elapsed_seconds).
+    """
+    max_wait_secs = 100
+    retry_interval = 10
+    elapsed = 0
+    dirty = None
+    unkillable_pids = {}
+
+    while True:
+        dirty = _get_dirty_gpus(gpu_indices)
+        remaining_pids = _get_gpu_pids(gpu_indices)
+
+        if not dirty:
+            _log(f"Check at {elapsed}s: GPUs clean")
+            break
+
+        dirty_summary = ", ".join(dirty)
+
+        if elapsed >= max_wait_secs:
+            remaining_info = (
+                f", {len(remaining_pids)} processes remaining" if remaining_pids else ""
+            )
+            _log(f"Check at {elapsed}s: still dirty [{dirty_summary}]{remaining_info}")
+            break
+
+        # Kill remaining processes before waiting (silently for retries)
+        if remaining_pids:
+            failed = _kill_pids(remaining_pids, quiet=True)
+            unkillable_pids.update(failed)
+
+        print(
+            f"[killall] GPUs still dirty at {elapsed}s [{dirty_summary}], "
+            f"retrying in {retry_interval}s "
+            f"({elapsed + retry_interval}/{max_wait_secs}s)..."
+        )
+        time.sleep(retry_interval)
+        elapsed += retry_interval
+
+    if unkillable_pids:
+        parts = [f"{p} ({unkillable_pids[p]})" for p in sorted(unkillable_pids)]
+        _log(f"  Unkillable PIDs: {', '.join(parts)}")
+
+    return dirty, unkillable_pids, elapsed
+
+
+def _ci_mode():
+    """GPU-scoped kill, abort if GPUs remain dirty."""
+    gpu_indices = _get_target_gpus()
+    if not gpu_indices:
+        _log("No GPUs detected, skipping cleanup")
+        _flush_box("killall_sglang", status="SKIP")
+        return 0
+
+    cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
+    gpu_list = ", ".join(str(g) for g in sorted(gpu_indices))
+
+    smi_info = _get_smi_version()
+    if smi_info:
+        _log(f"nvidia-smi: {smi_info}")
+    if cvd is None or not cvd.strip():
+        _log(
+            "WARNING: CUDA_VISIBLE_DEVICES is not set. "
+            "Falling back to all visible GPUs."
+        )
+        _log("This may kill processes from other CI jobs on shared hosts.")
+    else:
+        _log(f"CUDA_VISIBLE_DEVICES={cvd}")
+    _log()
+
+    # Log pre-cleanup state
+    _log("Before cleanup:")
+    _log_gpu_memory(gpu_indices)
+    gpu_pids = _get_gpu_pids(gpu_indices)
+    if not gpu_pids:
+        _log("  No processes on target GPUs")
+    else:
+        _log(f"  Processes ({len(gpu_pids)}):")
+        for pid in sorted(gpu_pids):
+            _log(f"    PID {pid}: {_get_pid_cmdline(pid)}")
+    _log()
+
+    # Kill phase
+    _kill_all_targets(gpu_indices, gpu_pids)
+
+    # Verify phase
+    dirty, unkillable_pids, elapsed = _verify_gpu_clean(gpu_indices)
+
+    if dirty:
+        _log()
+        _log("Final GPU memory:")
+        _log_gpu_memory(gpu_indices)
+        _log(f"ERROR: memory >={MEMORY_THRESHOLD_PCT}%: {', '.join(dirty)}")
+        _log(f"Orphaned CUDA contexts after {elapsed}s — container needs restart.")
+        _flush_box(f"killall_sglang: GPUs [{gpu_list}]", status="FAIL — Aborting CI")
+        _print_diagnostics(unkillable_pids)
+        return 1
+
+    _flush_box(f"killall_sglang: GPUs [{gpu_list}]", status="PASS — GPUs clean")
+    return 0
+
+
+# Entry point
+def main():
+    return _ci_mode()
+
+
+if __name__ == "__main__":
+    sys.exit(main())