[CI] Add retry loop to killall_sglang GPU cleanup verification (#21393)

This commit is contained in:
Liangsheng Yin
2026-03-25 02:16:20 -07:00
committed by GitHub
parent c494e47843
commit 4480e6c237
+56 -9
View File
@@ -127,6 +127,26 @@ def _get_orchestrator_ancestors(pids):
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
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 _log_gpu_memory(gpu_indices): def _log_gpu_memory(gpu_indices):
"""Log memory usage for target GPUs. Returns list of dirty GPU descriptions.""" """Log memory usage for target GPUs. Returns list of dirty GPU descriptions."""
dirty = [] dirty = []
@@ -224,21 +244,48 @@ def _ci_mode():
time.sleep(3) time.sleep(3)
_log() _log()
# Verify # Verify with retry: wait 10s per attempt, up to 100s total
_log("After cleanup:") max_wait_secs = 100
dirty = _log_gpu_memory(gpu_indices) retry_interval = 10
elapsed = 0
dirty = None
while True:
dirty = _check_gpu_memory(gpu_indices)
remaining_pids = _get_gpu_pids(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: if remaining_pids:
_log(f" Remaining processes ({len(remaining_pids)}):") _kill_pids(remaining_pids, "retry kill")
for pid in sorted(remaining_pids):
_log(f" PID {pid}: {_get_pid_cmdline(pid)}") print(
else: f"[killall] GPUs still dirty at {elapsed}s [{dirty_summary}], "
_log(" No processes on target GPUs") f"retrying in {retry_interval}s "
f"({elapsed + retry_interval}/{max_wait_secs}s)..."
)
time.sleep(retry_interval)
elapsed += retry_interval
if dirty: if dirty:
_log() _log()
_log("Final GPU memory:")
_log_gpu_memory(gpu_indices)
_log(f"ERROR: memory >={MEMORY_THRESHOLD_PCT}%: {', '.join(dirty)}") _log(f"ERROR: memory >={MEMORY_THRESHOLD_PCT}%: {', '.join(dirty)}")
_log("Orphaned CUDA contexts — container needs restart.") _log(f"Orphaned CUDA contexts after {elapsed}s — container needs restart.")
_flush_box(f"killall_sglang: GPUs [{gpu_list}]", status="FAIL — Aborting CI") _flush_box(f"killall_sglang: GPUs [{gpu_list}]", status="FAIL — Aborting CI")
return 1 return 1