[Fix] Unify pinned host pool release on graceful shutdown (#32029)
This commit is contained in:
@@ -105,6 +105,9 @@ class DecodeKVCacheOffloadManager:
|
|||||||
self.offload_inflight = {}
|
self.offload_inflight = {}
|
||||||
logger.info("Enable offload kv cache for decode side")
|
logger.info("Enable offload kv cache for decode side")
|
||||||
|
|
||||||
|
def release_host_resources(self) -> None:
|
||||||
|
self.decode_host_mem_pool.destroy()
|
||||||
|
|
||||||
def _mark_offload_started(self, rid):
|
def _mark_offload_started(self, rid):
|
||||||
self.offload_inflight[rid] = self.offload_inflight.get(rid, 0) + 1
|
self.offload_inflight[rid] = self.offload_inflight.get(rid, 0) + 1
|
||||||
|
|
||||||
|
|||||||
@@ -1465,11 +1465,9 @@ 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
|
self.tree_cache.release_host_resources()
|
||||||
# pool; unregister it here too, else the kernel unpins it during reclaim.
|
if self.decode_offload_manager is not None:
|
||||||
host_pool = getattr(self.tree_cache, "token_to_kv_pool_host", None)
|
self.decode_offload_manager.release_host_resources()
|
||||||
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.
|
||||||
|
|||||||
@@ -236,6 +236,13 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
|
|||||||
)
|
)
|
||||||
self.metrics_collector.increment_eviction_num_tokens(num_evicted)
|
self.metrics_collector.increment_eviction_num_tokens(num_evicted)
|
||||||
|
|
||||||
|
def release_host_resources(self) -> None:
|
||||||
|
"""Release pinned host buffers in userspace on graceful shutdown.
|
||||||
|
|
||||||
|
Kernel-side unpinning during process reclaim can stall teardown for
|
||||||
|
tens of seconds (see HostKVCache.destroy). Idempotent.
|
||||||
|
"""
|
||||||
|
|
||||||
@abstractmethod
|
@abstractmethod
|
||||||
def reset(self):
|
def reset(self):
|
||||||
pass
|
pass
|
||||||
|
|||||||
@@ -207,6 +207,9 @@ class HiMambaRadixCache(MambaRadixCache):
|
|||||||
)
|
)
|
||||||
super().reset()
|
super().reset()
|
||||||
|
|
||||||
|
def release_host_resources(self) -> None:
|
||||||
|
self.host_pool_group.destroy()
|
||||||
|
|
||||||
def write_backup(self, node: TreeNode, write_back=False) -> int:
|
def write_backup(self, node: TreeNode, write_back=False) -> int:
|
||||||
# Backup invariant (for write-through mode): backed-up nodes must form a
|
# Backup invariant (for write-through mode): backed-up nodes must form a
|
||||||
# contiguous prefix from root — no gaps. Skip if parent isn't backed
|
# contiguous prefix from root — no gaps. Skip if parent isn't backed
|
||||||
|
|||||||
@@ -775,6 +775,10 @@ class HiRadixCache(RadixCache):
|
|||||||
self.evictable_host_leaves.clear()
|
self.evictable_host_leaves.clear()
|
||||||
super().reset()
|
super().reset()
|
||||||
|
|
||||||
|
def release_host_resources(self) -> None:
|
||||||
|
if self.token_to_kv_pool_host is not None:
|
||||||
|
self.token_to_kv_pool_host.destroy()
|
||||||
|
|
||||||
def get_height(self, node: TreeNode):
|
def get_height(self, node: TreeNode):
|
||||||
height = 0
|
height = 0
|
||||||
while node != self.root_node:
|
while node != self.root_node:
|
||||||
|
|||||||
@@ -1573,6 +1573,10 @@ class HostPoolGroup:
|
|||||||
for entry in self.entries:
|
for entry in self.entries:
|
||||||
entry.host_pool.clear()
|
entry.host_pool.clear()
|
||||||
|
|
||||||
|
def destroy(self) -> None:
|
||||||
|
for entry in self.entries:
|
||||||
|
entry.host_pool.destroy()
|
||||||
|
|
||||||
def available_size(self):
|
def available_size(self):
|
||||||
return self.anchor_entry.host_pool.available_size()
|
return self.anchor_entry.host_pool.available_size()
|
||||||
|
|
||||||
|
|||||||
@@ -371,6 +371,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
|
|||||||
|
|
||||||
# HiCache D↔H defaults (overridden by init_hicache)
|
# HiCache D↔H defaults (overridden by init_hicache)
|
||||||
self.cache_controller: Optional[HybridCacheController] = None
|
self.cache_controller: Optional[HybridCacheController] = None
|
||||||
|
self.host_pool_group = None # set by attach_hybrid_pool_to_unified_cache
|
||||||
self.write_through_threshold = 256
|
self.write_through_threshold = 256
|
||||||
self.prefetch_stop_policy = "best_effort"
|
self.prefetch_stop_policy = "best_effort"
|
||||||
self.prefetch_threshold = 256
|
self.prefetch_threshold = 256
|
||||||
@@ -568,6 +569,10 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
|
|||||||
def register_sidecar_pool(self, spec: SidecarPoolSpec) -> None:
|
def register_sidecar_pool(self, spec: SidecarPoolSpec) -> None:
|
||||||
self.sidecar_pool_specs.append(spec)
|
self.sidecar_pool_specs.append(spec)
|
||||||
|
|
||||||
|
def release_host_resources(self) -> None:
|
||||||
|
if self.host_pool_group is not None:
|
||||||
|
self.host_pool_group.destroy()
|
||||||
|
|
||||||
def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
|
def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
|
||||||
result = self.session.try_match_prefix(params)
|
result = self.session.try_match_prefix(params)
|
||||||
if result is not None:
|
if result is not None:
|
||||||
|
|||||||
Reference in New Issue
Block a user