[CI] Add PID namespace and ps auxf diagnostics to killall.py (#21401)

python/sglang/cli/killall.py

@@ -16,15 +16,52 @@ Exit codes:
 """
 
 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::"
+)
 
+# 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"
@@ -39,32 +76,31 @@ def _run_smi(query, query_type="gpu"):
         return []
 
 
-def _get_pid_cmdline(pid):
+def _get_smi_version():
-    """Get command line for a PID. Linux-only via /proc."""
+    """Return nvidia-smi driver version and CUDA version, or None on failure."""
     try:
-        cmdline = Path(f"/proc/{pid}/cmdline").read_bytes()
+        out = subprocess.check_output(
-        cmdline = cmdline.decode("utf-8", errors="replace").replace("\x00", " ").strip()
+            [
-        # Truncate long command lines
+                "nvidia-smi",
-        return cmdline[:120] + ("..." if len(cmdline) > 120 else "")
+                "--query-gpu=driver_version",
-    except (FileNotFoundError, PermissionError):
+                "--format=csv,noheader,nounits",
-        return "<unknown>"
+            ],
-
+            text=True,
-
+            timeout=10,
-def _kill_pids(pids, label=""):
+        )
-    """Send SIGKILL to PIDs, skipping self and init. Logs to _LOG_LINES."""
+        driver = out.strip().splitlines()[0].strip() if out.strip() else "unknown"
-    my_pid = os.getpid()
+    except (subprocess.SubprocessError, FileNotFoundError, IndexError):
-    pids = {p for p in pids if p != my_pid and p > 1}
+        return None
-    if not pids:
+    try:
-        return
+        out = subprocess.check_output(
-    if label:
+            ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
-        _log(f"  Killing {label}:")
+            text=True,
-    for pid in sorted(pids):
+            timeout=10,
-        cmdline = _get_pid_cmdline(pid)
+        )
-        _log(f"    PID {pid}: {cmdline}")
+        gpu_name = out.strip().splitlines()[0].strip() if out.strip() else "unknown"
-        try:
+    except (subprocess.SubprocessError, FileNotFoundError, IndexError):
-            os.kill(pid, signal.SIGKILL)
+        gpu_name = "unknown"
-        except (ProcessLookupError, PermissionError):
+    return f"driver {driver}, {gpu_name}"
-            _log(f"    PID {pid}: failed (already dead or no permission)")
 
 
 def _get_target_gpus():
@@ -97,24 +133,122 @@ def _get_gpu_pids(gpu_indices):
     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.
     """
-    patterns = ["run_suite.py", "run_tests.py"]
+    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)
-            try:
+            cmdline = _read_proc_cmdline(current)
-                cmdline = Path(f"/proc/{current}/cmdline").read_bytes()
+            if cmdline is None:
-                cmdline = cmdline.decode("utf-8", errors="replace").replace("\x00", " ")
-                if any(p in cmdline for p in patterns):
-                    ancestors.add(current)
-            except (FileNotFoundError, PermissionError):
                 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):
@@ -122,80 +256,134 @@ def _get_orchestrator_ancestors(pids):
     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 _check_gpu_memory(gpu_indices):
+def _verify_gpu_clean(gpu_indices):
-    """Check memory usage for target GPUs. Returns list of dirty GPU descriptions."""
+    """Retry loop: wait for GPUs to become clean.
-    dirty = []
-    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
-        if pct >= MEMORY_THRESHOLD_PCT:
-            dirty.append(f"GPU {idx} ({pct:.0f}%)")
-    return dirty
 
+    Returns (dirty_list, unkillable_pids, elapsed_seconds).
+    """
+    max_wait_secs = 100
+    retry_interval = 10
+    elapsed = 0
+    dirty = None
+    unkillable_pids = {}
 
-def _log_gpu_memory(gpu_indices):
+    while True:
-    """Log memory usage for target GPUs. Returns list of dirty GPU descriptions."""
+        dirty = _get_dirty_gpus(gpu_indices)
-    dirty = []
+        remaining_pids = _get_gpu_pids(gpu_indices)
-    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
-        _log(f"  GPU {idx}: {used} MiB / {total} MiB ({pct:.0f}%)")
-        if pct >= MEMORY_THRESHOLD_PCT:
-            dirty.append(f"GPU {idx} ({pct:.0f}%)")
-    return dirty
 
+        if not dirty:
+            _log(f"Check at {elapsed}s: GPUs clean")
+            break
 
-_LOG_LINES = []
+        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
 
-def _log(msg=""):
+        # Kill remaining processes before waiting (silently for retries)
-    """Buffer a line for boxed output."""
+        if remaining_pids:
-    _LOG_LINES.append(msg)
+            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
 
-def _flush_box(title, status=""):
+    if unkillable_pids:
-    """Print all buffered lines inside a box, then clear buffer."""
+        parts = [f"{p} ({unkillable_pids[p]})" for p in sorted(unkillable_pids)]
-    lines = _LOG_LINES.copy()
+        _log(f"  Unkillable PIDs: {', '.join(parts)}")
-    _LOG_LINES.clear()
 
-    # Build content width from title, status, and all lines
+    return dirty, unkillable_pids, elapsed
-    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}┘")
 
 
 def _ci_mode():
@@ -209,6 +397,9 @@ def _ci_mode():
     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. "
@@ -219,7 +410,7 @@ def _ci_mode():
         _log(f"CUDA_VISIBLE_DEVICES={cvd}")
     _log()
 
-    # Before cleanup
+    # Log pre-cleanup state
     _log("Before cleanup:")
     _log_gpu_memory(gpu_indices)
     gpu_pids = _get_gpu_pids(gpu_indices)
@@ -231,54 +422,11 @@ def _ci_mode():
             _log(f"    PID {pid}: {_get_pid_cmdline(pid)}")
     _log()
 
-    # Kill orchestrator ancestors first, then GPU processes (retry once)
+    # Kill phase
-    if gpu_pids:
+    _kill_all_targets(gpu_indices, gpu_pids)
-        _kill_pids(_get_orchestrator_ancestors(gpu_pids), "orchestrator ancestors")
-        time.sleep(1)
-        for attempt in range(2):
-            gpu_pids = _get_gpu_pids(gpu_indices)
-            if not gpu_pids:
-                break
-            label = "GPU processes" if attempt == 0 else "stubborn GPU processes"
-            _kill_pids(gpu_pids, label)
-            time.sleep(3)
-    _log()
 
-    # Verify with retry: wait 10s per attempt, up to 100s total
+    # Verify phase
-    max_wait_secs = 100
+    dirty, unkillable_pids, elapsed = _verify_gpu_clean(gpu_indices)
-    retry_interval = 10
-    elapsed = 0
-    dirty = None
-
-    while True:
-        dirty = _check_gpu_memory(gpu_indices)
-        remaining_pids = _get_gpu_pids(gpu_indices)
-
-        if not dirty:
-            _log(f"Check at {elapsed}s: GPUs clean")
-            break
-
-        # Log summary for this attempt
-        remaining_info = (
-            f", {len(remaining_pids)} processes remaining" if remaining_pids else ""
-        )
-        dirty_summary = ", ".join(dirty)
-        _log(f"Check at {elapsed}s: still dirty [{dirty_summary}]{remaining_info}")
-
-        if elapsed >= max_wait_secs:
-            break
-
-        # Kill remaining processes before waiting
-        if remaining_pids:
-            _kill_pids(remaining_pids, "retry kill")
-
-        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 dirty:
         _log()
@@ -287,17 +435,14 @@ def _ci_mode():
         _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()