[Fix] Release hierarchical cache host pool on graceful shutdown (#31746)
This commit is contained in:
@@ -1460,6 +1460,11 @@ class Scheduler(
|
|||||||
# HostKVCache.destroy. Called from run_scheduler_process's finally.
|
# HostKVCache.destroy. Called from run_scheduler_process's finally.
|
||||||
if self.hisparse_coordinator is not None:
|
if self.hisparse_coordinator is not None:
|
||||||
self.hisparse_coordinator.destroy()
|
self.hisparse_coordinator.destroy()
|
||||||
|
# A plain HiRadixCache (no hisparse) also holds a large pinned host KV
|
||||||
|
# pool; unregister it here too, else the kernel unpins it during reclaim.
|
||||||
|
host_pool = getattr(self.tree_cache, "token_to_kv_pool_host", None)
|
||||||
|
if host_pool is not None:
|
||||||
|
host_pool.destroy()
|
||||||
|
|
||||||
def run_event_loop(self) -> None:
|
def run_event_loop(self) -> None:
|
||||||
"""Run the scheduler's event loop.
|
"""Run the scheduler's event loop.
|
||||||
|
|||||||
@@ -3,6 +3,7 @@ from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
|||||||
register_cuda_ci(est_time=99, stage="base-b", runner_config="1-gpu-small")
|
register_cuda_ci(est_time=99, stage="base-b", runner_config="1-gpu-small")
|
||||||
register_amd_ci(est_time=300, suite="stage-b-test-1-gpu-small-amd")
|
register_amd_ci(est_time=300, suite="stage-b-test-1-gpu-small-amd")
|
||||||
|
|
||||||
|
import subprocess
|
||||||
import time
|
import time
|
||||||
import unittest
|
import unittest
|
||||||
|
|
||||||
@@ -47,6 +48,14 @@ class TestHiCache(CustomTestCase, MMLUMixin):
|
|||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def tearDownClass(cls):
|
def tearDownClass(cls):
|
||||||
|
# Graceful stop first so the server unregisters its large pinned host KV
|
||||||
|
# pool in userspace; a bare SIGKILL leaves the kernel to unpin it during
|
||||||
|
# reclaim, holding GPU memory long enough to fail the next test.
|
||||||
|
cls.process.terminate()
|
||||||
|
try:
|
||||||
|
cls.process.wait(timeout=60)
|
||||||
|
except subprocess.TimeoutExpired:
|
||||||
|
pass
|
||||||
kill_process_tree(cls.process.pid)
|
kill_process_tree(cls.process.pid)
|
||||||
time.sleep(5)
|
time.sleep(5)
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user