diff --git a/python/sglang/cli/killall.py b/python/sglang/cli/killall.py index b0ee7f8c2..39a10f053 100755 --- a/python/sglang/cli/killall.py +++ b/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): - """Get command line for a PID. Linux-only via /proc.""" +def _get_smi_version(): + """Return nvidia-smi driver version and CUDA version, or None on failure.""" try: - cmdline = Path(f"/proc/{pid}/cmdline").read_bytes() - cmdline = cmdline.decode("utf-8", errors="replace").replace("\x00", " ").strip() - # Truncate long command lines - return cmdline[:120] + ("..." if len(cmdline) > 120 else "") - except (FileNotFoundError, PermissionError): - return "" - - -def _kill_pids(pids, label=""): - """Send SIGKILL to PIDs, skipping self and init. Logs to _LOG_LINES.""" - my_pid = os.getpid() - pids = {p for p in pids if p != my_pid and p > 1} - if not pids: - return - if label: - _log(f" Killing {label}:") - for pid in sorted(pids): - cmdline = _get_pid_cmdline(pid) - _log(f" PID {pid}: {cmdline}") - try: - os.kill(pid, signal.SIGKILL) - except (ProcessLookupError, PermissionError): - _log(f" PID {pid}: failed (already dead or no permission)") + out = subprocess.check_output( + [ + "nvidia-smi", + "--query-gpu=driver_version", + "--format=csv,noheader,nounits", + ], + text=True, + timeout=10, + ) + driver = out.strip().splitlines()[0].strip() if out.strip() else "unknown" + except (subprocess.SubprocessError, FileNotFoundError, IndexError): + 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(): @@ -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 "" + 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 = Path(f"/proc/{current}/cmdline").read_bytes() - cmdline = cmdline.decode("utf-8", errors="replace").replace("\x00", " ") - if any(p in cmdline for p in patterns): - ancestors.add(current) - except (FileNotFoundError, PermissionError): + 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): @@ -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): - """Check memory usage for target GPUs. Returns list of dirty GPU descriptions.""" - 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 +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 = {} -def _log_gpu_memory(gpu_indices): - """Log memory usage for target GPUs. Returns list of dirty GPU descriptions.""" - 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 - _log(f" GPU {idx}: {used} MiB / {total} MiB ({pct:.0f}%)") - if pct >= MEMORY_THRESHOLD_PCT: - dirty.append(f"GPU {idx} ({pct:.0f}%)") - return dirty + 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 -_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=""): - """Buffer a line for boxed output.""" - _LOG_LINES.append(msg) + # 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 -def _flush_box(title, status=""): - """Print all buffered lines inside a box, then clear buffer.""" - lines = _LOG_LINES.copy() - _LOG_LINES.clear() + if unkillable_pids: + parts = [f"{p} ({unkillable_pids[p]})" for p in sorted(unkillable_pids)] + _log(f" Unkillable PIDs: {', '.join(parts)}") - # Build content width from title, status, and all lines - 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}┘") + return dirty, unkillable_pids, elapsed 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) - if 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() + # Kill phase + _kill_all_targets(gpu_indices, gpu_pids) - # Verify with retry: wait 10s per attempt, up to 100s total - max_wait_secs = 100 - 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 + # Verify phase + dirty, unkillable_pids, elapsed = _verify_gpu_clean(gpu_indices) 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()