@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4ed51308c8

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 6, 13:15 UTC

[CI] Fix GPU memory leak when RemoteOpenAIServer fails to start in init (#37230)

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

Andreas Karatzas committed Mar 17, 2026 at 11:08 UTC 4ed51308c8826619459be858a6dc4333206f41c1

1 file changed +123 -52

tests/utils.py

+123 -52

@@ -225,13 +225,31 @@ class RemoteVLLMServer:

225 )

226

227 self._start_server(model, vllm_serve_args, env_dict)

228 - max_wait_seconds = max_wait_seconds or 360

229 - self._wait_for_server(url=self.url_for("health"), timeout=max_wait_seconds)

228 + max_wait_seconds = max_wait_seconds or 480

229 + try:

230 + self._wait_for_server(url=self.url_for("health"), timeout=max_wait_seconds)

231 + except Exception:

232 + # If the server never became healthy, we must still clean up

233 + # the subprocess tree. Without this, a timeout in __init__

234 + # leaks the server + EngineCore processes (and their GPU

235 + # memory), because __exit__ is never called when __init__

236 + # raises inside a ``with`` statement.

237 + self._shutdown()

238 + raise

239

240 def __enter__(self):

241 return self

242

243 def __exit__(self, exc_type, exc_value, traceback):

244 + self._shutdown()

245 +

246 + def _shutdown(self) -> None:

247 + """Kill the server process tree and wait for GPU memory release.

248 +

249 + Called from both ``__exit__`` (normal path) and ``__init__``

250 + (when the server fails to start). Must be safe to call even if

251 + the process is already dead.

252 + """

253 pid = self.proc.pid

254

255 # Get the process group ID. Because we used

@@ -265,33 +283,92 @@ class RemoteVLLMServer:

283 self.proc.wait(timeout=10)

284 print(f"[RemoteOpenAIServer] Server {pid} killed")

285 except subprocess.TimeoutExpired:

268 - # Phase 3: last resort - find and kill any orphaned children

269 - self._kill_orphaned_children(pid)

286 + pass

287

271 - # Wait for GPU memory to actually be *freed*, not just

288 + # After killing the root process, ensure all children in the

289 + # process group (e.g. EngineCore workers) are also dead.

290 + # On ROCm especially, surviving children hold GPU contexts and

291 + # prevent VRAM from being reclaimed by the driver.

292 + self._kill_process_group_survivors(pgid)

293 +

294 + # Wait for GPU memory to actually be freed, not just

295 # "stabilized at whatever level it's at".

296 self._wait_for_gpu_memory_release()

297

275 - def _kill_orphaned_children(self, parent_pid: int) -> None:

276 - """Best-effort cleanup of any lingering child processes."""

277 - try:

278 - import psutil

298 + def _kill_process_group_survivors(

299 + self, pgid: int | None, timeout: float = 15.0

300 + ) -> None:

301 + """SIGKILL any processes still in the server's process group

302 + and wait for them to exit.

303

280 - parent = psutil.Process(parent_pid)

281 - children = parent.children(recursive=True)

282 - for child in children:

283 - print(

284 - f"[RemoteOpenAIServer] Killing orphaned child "

285 - f"pid={child.pid} name={child.name()}"

286 - )

287 - child.kill()

288 - psutil.wait_procs(children, timeout=5)

289 - except Exception as e:

290 - # psutil may not be installed, or processes already gone

291 - print(f"[RemoteOpenAIServer] Orphan cleanup failed: {e}")

292 - # Fallback: try to kill by pgid one more time

293 - with contextlib.suppress(ProcessLookupError, OSError):

294 - os.killpg(parent_pid, signal.SIGKILL)

304 + Because the server is launched with ``start_new_session=True``,

305 + all its children (EngineCore, workers, etc.) share the same

306 + pgid. After the root process is killed, stragglers -- especially

307 + on ROCm where GPU contexts linger until the *process* exits --

308 + must be reaped explicitly.

309 +

310 + Uses ``/proc`` to scan for pgid members so this works even after

311 + the parent has been reaped (unlike ``psutil.Process.children``).

312 + """

313 + if pgid is None:

314 + return

315 +

316 + # Send SIGKILL to the entire process group one more time.

317 + # This is cheap and harmless if everyone is already dead.

318 + with contextlib.suppress(ProcessLookupError, OSError):

319 + os.killpg(pgid, signal.SIGKILL)

320 +

321 + # Collect surviving PIDs by scanning /proc for matching pgid.

322 + # This works on Linux even after the parent has been waited on

323 + # and is more reliable than psutil.Process(parent).children().

324 + survivor_pids = self._find_pgid_members(pgid)

325 +

326 + if not survivor_pids:

327 + return

328 +

329 + print(

330 + f"[RemoteOpenAIServer] {len(survivor_pids)} process(es) still "

331 + f"in pgid {pgid} after SIGKILL: {survivor_pids}"

332 + )

333 +

334 + # Wait for each survivor to actually exit so the GPU driver

335 + # releases its VRAM.

336 + deadline = time.time() + timeout

337 + while survivor_pids and time.time() < deadline:

338 + still_alive = []

339 + for spid in survivor_pids:

340 + try:

341 + os.kill(spid, 0) # Check if still alive

342 + still_alive.append(spid)

343 + except (ProcessLookupError, OSError):

344 + pass

345 + survivor_pids = still_alive

346 + if survivor_pids:

347 + time.sleep(0.5)

348 +

349 + if survivor_pids:

350 + print(

351 + f"[RemoteOpenAIServer] WARNING: processes {survivor_pids} "

352 + f"in pgid {pgid} could not be killed within {timeout}s"

353 + )

354 +

355 + @staticmethod

356 + def _find_pgid_members(pgid: int) -> list[int]:

357 + """Return PIDs of all living processes whose pgid matches."""

358 + members: list[int] = []

359 + proc_path = Path("/proc")

360 + if not proc_path.is_dir():

361 + return members

362 + for entry in proc_path.iterdir():

363 + if not entry.name.isdigit():

364 + continue

365 + pid = int(entry.name)

366 + try:

367 + if os.getpgid(pid) == pgid:

368 + members.append(pid)

369 + except OSError:

370 + continue

371 + return members

372

373 def _get_gpu_memory_used(self) -> float | None:

374 """Get total GPU memory used across all visible devices in bytes."""

@@ -318,13 +395,16 @@ class RemoteVLLMServer:

395 return None

396 return None

397

321 - def _wait_for_gpu_memory_release(self, timeout: float = 60.0):

398 + def _wait_for_gpu_memory_release(

399 + self, timeout: float = 120.0, log_interval: float = 10.0

400 + ):

401 """Wait for GPU memory to drop back toward pre-server levels.

402

324 - Two-phase strategy:

325 - 1. Try to wait for memory to return close to pre-server baseline.

326 - 2. If that doesn't happen, fall back to waiting for stabilization

327 - and log a warning (the next server might still OOM).

403 + Waits the full timeout for memory to return close to the

404 + pre-server baseline. Does NOT fall back to a "stabilization"

405 + heuristic -- if memory is still held when the timeout expires,

406 + the test fails so the problem is surfaced immediately rather

407 + than causing cascading OOM failures in every subsequent test.

408 """

409 baseline = self._pre_server_gpu_memory

410 if baseline is None:

@@ -337,8 +417,7 @@ class RemoteVLLMServer:

417 target = baseline + headroom_bytes

418

419 start = time.time()

340 - last_used: float | None = None

341 - stable_count = 0

420 + next_log_time = start + log_interval

421

422 while time.time() - start < timeout:

423 used = self._get_gpu_memory_used()

@@ -350,7 +429,6 @@ class RemoteVLLMServer:

429 target_gb = target / 1e9

430 elapsed = time.time() - start

431

353 - # Phase 1: memory dropped to near baseline - we're done.

432 if used <= target:

433 print(

434 f"[RemoteOpenAIServer] GPU memory released to "

@@ -359,28 +437,19 @@ class RemoteVLLMServer:

437 )

438 return

439

362 - # Phase 2 (after 40s): fall back to stabilization check.

363 - # This handles cases where another process is using GPU memory

364 - # and we'll never reach baseline.

365 - if elapsed > 40.0 and last_used is not None:

366 - delta = abs(used - last_used)

367 - if delta < 200 * 1024 * 1024: # 200 MB

368 - stable_count += 1

369 - if stable_count >= 3:

370 - print(

371 - f"[RemoteOpenAIServer] WARNING: GPU memory "

372 - f"stabilized at {used_gb:.2f} GB "

373 - f"(target was {target_gb:.2f} GB). "

374 - f"Proceeding - next server may OOM."

375 - )

376 - return

377 - else:

378 - stable_count = 0

440 + now = time.time()

441 + if now >= next_log_time:

442 + print(

443 + f"[RemoteOpenAIServer] Waiting for GPU memory release: "

444 + f"{used_gb:.2f} GB (target: {target_gb:.2f} GB) "

445 + f"[{elapsed:.0f}s/{timeout:.0f}s]"

446 + )

447 + next_log_time = now + log_interval

448

380 - last_used = used

449 time.sleep(1.0)

450

383 - # Timeout - log clearly so CI failures are diagnosable

451 + # Timeout -- raise so the current test fails with a clear

452 + # message instead of silently poisoning subsequent tests.

453 final_used = self._get_gpu_memory_used()

454 final_gb = final_used / 1e9 if final_used else 0.0

455 raise RuntimeError(

@@ -534,7 +603,9 @@ class RemoteLaunchRenderServer(RemoteVLLMServer):

603 revision=model_config.tokenizer_revision,

604 )

605

537 - def _wait_for_gpu_memory_release(self, timeout: float = 30.0):

606 + def _wait_for_gpu_memory_release(

607 + self, timeout: float = 30.0, log_interval: float = 10.0

608 + ):

609 pass # No GPU used

610

611

[CI] Fix GPU memory leak when RemoteOpenAIServer fails to start in __init__ (#37230)

[CI] Fix GPU memory leak when RemoteOpenAIServer fails to start in init (#37230)