[CI] Wait for a killed test server's GPU memory before the next launch (#39545)

python/sglang/test/runners.py

@@ -35,7 +35,12 @@ from sglang.srt.entrypoints.engine import Engine
 from sglang.srt.model_loader.ci_weight_validation import ci_validate_and_clean_hf_cache
 from sglang.srt.utils import get_device, is_npu, load_image
 from sglang.srt.utils.hf_transformers_utils import get_tokenizer
-from sglang.test.test_utils import DEFAULT_PORT_FOR_SRT_TEST_RUNNER, calculate_rouge_l
+from sglang.test.test_utils import (
+    DEFAULT_PORT_FOR_SRT_TEST_RUNNER,
+    calculate_rouge_l,
+    collect_process_tree_pids,
+    wait_for_gpu_release,
+)
 
 if is_npu():
     from sglang.srt.hardware_backend.npu.utils import init_npu_backend
@@ -428,12 +433,14 @@ class HFRunner:
         # Fire-and-forget terminate() leaves the child holding the accelerator
         # during teardown; a follow-on SRTRunner on the same device can then
         # deadlock in driver init (observed on Intel XPU B580).
+        pid = self.model_proc.pid
         self.model_proc.terminate()
         self.model_proc.join(timeout=30)
         if self.model_proc.is_alive():
             self.model_proc.kill()
             self.model_proc.join()
         self.in_queue = self.out_queue = None
+        wait_for_gpu_release([pid])
 
     def terminate(self):
         self._stop_model_proc()
@@ -751,8 +758,13 @@ class SRTRunner:
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
+        # Wait only on the pids this shutdown actually killed;
+        # a nested HFRunner or SRTRunner is deliberately still alive.
+        before = collect_process_tree_pids(os.getpid(), include_parent=False)
         self.engine.shutdown()
         del self.engine
+        alive = set(collect_process_tree_pids(os.getpid(), include_parent=False))
+        wait_for_gpu_release([pid for pid in before if pid not in alive])
 
     @staticmethod
     def forward_generation_raw(

python/sglang/test/test_utils.py

@@ -28,6 +28,7 @@ from typing import Any, Awaitable, Callable, List, Optional, Tuple
 import aiohttp
 import msgspec
 import numpy as np
+import psutil
 import requests
 import torch
 import torch.nn.functional as F
@@ -847,14 +848,17 @@ def terminate_and_kill_process_tree(
     and unpin the host memory during process reclaim, which can hold GPU memory
     for minutes on a busy host -- long enough to trip the per-class GPU-idle
     gate in the next ``setUpClass``. SIGTERM first so the server releases those
-    resources in userspace.
+    resources in userspace, then wait for the memory to come back:
+    a reaped tree does not mean the driver is done with it.
     """
+    pids = collect_process_tree_pids(process.pid)
     process.terminate()
     try:
         process.wait(timeout=terminate_timeout)
     except subprocess.TimeoutExpired:
         pass
     kill_process_tree(process.pid, **kill_kwargs)
+    wait_for_gpu_release(pids)
 
 
 def popen_launch_pd_server(
@@ -2001,6 +2005,8 @@ def maybe_stub_sgl_kernel():
 _GPU_IDLE_TIMEOUT_SECS = 30.0
 _GPU_IDLE_POLL_INTERVAL_SECS = 2.0
 _GPU_IDLE_USED_MEMORY_THRESHOLD = 2 << 30  # 2 GiB
+_GPU_RELEASE_TIMEOUT_SECS = 60.0
+_GPU_RELEASE_POLL_INTERVAL_SECS = 0.5
 
 
 def _format_gib(num_bytes: Optional[int]) -> str:
@@ -2102,6 +2108,86 @@ def _wait_for_gpu_idle_in_ci(
             pass
 
 
+def collect_process_tree_pids(pid: int, include_parent: bool = True) -> List[int]:
+    """Snapshot a process tree's pids, for a later ``wait_for_gpu_release``.
+
+    Call it BEFORE the kill; afterwards the tree cannot be walked.
+    """
+    try:
+        pids = [child.pid for child in psutil.Process(pid).children(recursive=True)]
+    except psutil.Error:
+        pids = []
+    if include_parent:
+        pids.append(pid)
+    return pids
+
+
+def _gpu_memory_holders(pynvml, gpu_indices: List[int], pids: set) -> List[str]:
+    reports = []
+    for index in gpu_indices:
+        handle = pynvml.nvmlDeviceGetHandleByIndex(index)
+        try:
+            procs = pynvml.nvmlDeviceGetComputeRunningProcesses(handle)
+        except pynvml.NVMLError:
+            # No per-pid enumeration in this container; nothing to wait on.
+            continue
+        reports.extend(
+            f"GPU {index} pid={proc.pid} {_format_gib(proc.usedGpuMemory)}"
+            for proc in procs
+            if proc.pid in pids
+        )
+    return reports
+
+
+def wait_for_gpu_release(
+    pids: List[int],
+    timeout: float = _GPU_RELEASE_TIMEOUT_SECS,
+    poll_interval: float = _GPU_RELEASE_POLL_INTERVAL_SECS,
+) -> None:
+    """Block until none of ``pids`` is still charged device memory.
+
+    Killing a server only queues the driver-side teardown,
+    so the next launch can OOM against memory charged to a reaped process.
+    Waiting on these pids, rather than on an idle GPU,
+    keeps this usable while other servers of the same test still run.
+    Best effort: a timeout or a dead NVML warns, never raises.
+    """
+    if not pids:
+        return
+    try:
+        import pynvml
+
+        pynvml.nvmlInit()
+    except Exception:
+        # Non-NVIDIA runner (CPU/AMD) or NVML unavailable; nothing to check.
+        return
+    try:
+        gpu_indices = _visible_gpu_indices(pynvml)
+        pending = set(pids)
+        deadline = time.monotonic() + timeout
+        while True:
+            holders = _gpu_memory_holders(pynvml, gpu_indices, pending)
+            if not holders:
+                return
+            if time.monotonic() >= deadline:
+                print(
+                    f"[CI GPU Release] Still charged after {timeout:.0f}s:"
+                    f" {'; '.join(holders)}",
+                    flush=True,
+                )
+                return
+            time.sleep(poll_interval)
+    except Exception as e:
+        # NVML can go away after a successful init (GPU lost, driver reset).
+        # Raising here would fail a teardown whose test already passed.
+        print(f"[CI GPU Release] Giving up, {type(e).__name__}: {e}", flush=True)
+    finally:
+        try:
+            pynvml.nvmlShutdown()
+        except Exception:
+            pass
+
+
 # Names the runner kits stamp onto a record that are not members of it.
 # `ModelRunner` computes `use_mla_backend` on itself; the kits copy that bool
 # onto the record they hand the runner, and `hasattr` cannot see it.