[UnifiedTree] feat: support runtime attach/detach (#35269)
Co-authored-by: hzh0425 <hzh0425@apache.org>
This commit is contained in:
@@ -321,6 +321,11 @@ class HiCacheController:
|
||||
# Dedicated stop event for storage background threads (prefetch/backup).
|
||||
self.storage_stop_event = threading.Event()
|
||||
|
||||
# Storage control queues, (re)created whenever the storage threads start.
|
||||
self.prefetch_hit_queue: Optional[Queue[StorageOperation]] = None
|
||||
self.ack_backup_queue: Optional[Queue[StorageOperation]] = None
|
||||
self.host_mem_release_queue: Optional[Queue[torch.Tensor]] = None
|
||||
|
||||
self.device = self.mem_pool_device.device
|
||||
self.layer_num = self.mem_pool_device.layer_num
|
||||
self.layer_done_counter = LayerDoneCounter(self.layer_num)
|
||||
@@ -417,9 +422,9 @@ class HiCacheController:
|
||||
self.prefetch_queue = Queue()
|
||||
self.backup_queue = Queue()
|
||||
|
||||
self.prefetch_hit_queue: Queue[StorageOperation] = Queue()
|
||||
self.ack_backup_queue: Queue[StorageOperation] = Queue()
|
||||
self.host_mem_release_queue: Queue[torch.Tensor] = Queue()
|
||||
self.prefetch_hit_queue = Queue()
|
||||
self.ack_backup_queue = Queue()
|
||||
self.host_mem_release_queue = Queue()
|
||||
|
||||
self.prefetch_thread.start()
|
||||
self.backup_thread.start()
|
||||
|
||||
@@ -0,0 +1,400 @@
|
||||
"""Runtime attach / detach of the HiCache (L3) storage backend.
|
||||
|
||||
``UnifiedRadixCache`` owns the tree; this component owns the lifecycle of the
|
||||
storage backend behind it: validating the requested policies, applying the storage
|
||||
runtime config, starting and stopping the controller's storage threads, and
|
||||
cleaning up the bookkeeping a half-finished prefetch / backup leaves behind.
|
||||
|
||||
Keeping it beside the tree rather than inside it gives the three entry points --
|
||||
startup, the admin API, and the atexit hook -- one implementation.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import (
|
||||
HybridCacheController,
|
||||
)
|
||||
from sglang.srt.observability.metrics_collector import (
|
||||
STAT_LOGGER_ROLE_STORAGE,
|
||||
StorageMetricsCollector,
|
||||
resolve_collector_class,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Kept in sync with the server-args choices; validated here so an admin request
|
||||
# carrying a bad policy is rejected before it can touch the controller.
|
||||
_PREFETCH_POLICIES = ("best_effort", "wait_complete", "timeout")
|
||||
_WRITE_POLICIES = ("write_back", "write_through", "write_through_selective")
|
||||
|
||||
|
||||
class StorageAttachment:
|
||||
"""Attach / detach the storage backend of one ``UnifiedRadixCache``."""
|
||||
|
||||
def __init__(self, cache: UnifiedRadixCache):
|
||||
self._cache = cache
|
||||
|
||||
# ---- Lifecycle entry points ----
|
||||
|
||||
def attach(
|
||||
self,
|
||||
storage_backend: str,
|
||||
storage_backend_extra_config_json: Optional[str] = None,
|
||||
served_model_name: Optional[str] = None,
|
||||
hicache_storage_prefetch_policy: Optional[str] = None,
|
||||
hicache_write_policy: Optional[str] = None,
|
||||
) -> tuple[bool, str]:
|
||||
"""Enable the storage backend at runtime.
|
||||
|
||||
Starts the storage threads inside the cache controller and turns on the
|
||||
prefetch / backup paths. The caller must ensure there are no running or
|
||||
queued requests, so this cannot race the scheduler thread.
|
||||
"""
|
||||
cache = self._cache
|
||||
|
||||
# Validate first: a rejected request must have no side effects.
|
||||
invalid = self._validate_policies(
|
||||
hicache_storage_prefetch_policy, hicache_write_policy
|
||||
)
|
||||
if invalid is not None:
|
||||
return False, invalid
|
||||
|
||||
controller = cache.cache_controller
|
||||
if controller is None:
|
||||
return (
|
||||
False,
|
||||
"HiCache is not initialized (no cache controller); "
|
||||
"launch with --enable-hierarchical-cache to attach a backend.",
|
||||
)
|
||||
|
||||
if cache.enable_storage:
|
||||
current_backend = controller.storage_backend_type
|
||||
if current_backend != storage_backend:
|
||||
return (
|
||||
False,
|
||||
f"HiCache storage backend is already enabled with backend "
|
||||
f"'{current_backend}'. Cannot attach different backend "
|
||||
f"'{storage_backend}'. Detach first.",
|
||||
)
|
||||
# Same backend: the request degenerates to a policy update.
|
||||
self._apply_policies(hicache_storage_prefetch_policy, hicache_write_policy)
|
||||
return (
|
||||
True,
|
||||
"HiCache storage backend already enabled with same backend; "
|
||||
"policies updated.",
|
||||
)
|
||||
|
||||
# Apply policies before the controller attach, so the storage threads
|
||||
# observe the new values as soon as they start.
|
||||
self._apply_policies(hicache_storage_prefetch_policy, hicache_write_policy)
|
||||
|
||||
logger.info(f"Attaching HiCache storage backend: {storage_backend}")
|
||||
try:
|
||||
(
|
||||
extra_config,
|
||||
prefetch_threshold,
|
||||
prefetch_timeout_base,
|
||||
prefetch_timeout_per_ki_token,
|
||||
hicache_storage_pass_prefix_keys,
|
||||
) = HybridCacheController.parse_storage_backend_extra_config(
|
||||
storage_backend_extra_config_json
|
||||
)
|
||||
except Exception as e:
|
||||
logger.exception(f"Failed to parse storage_backend_extra_config_json: {e}")
|
||||
return (
|
||||
False,
|
||||
f"Failed to parse storage_backend_extra_config_json "
|
||||
f"'{storage_backend_extra_config_json}': {e}",
|
||||
)
|
||||
|
||||
try:
|
||||
controller.attach_storage_backend(
|
||||
storage_backend=storage_backend,
|
||||
prefetch_threshold=prefetch_threshold,
|
||||
model_name=served_model_name,
|
||||
storage_backend_extra_config=extra_config,
|
||||
host_pools=controller.mem_pool_host.entries,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.exception(
|
||||
f"Failed to attach storage backend '{storage_backend}': {e}"
|
||||
)
|
||||
return False, f"Failed to attach storage backend '{storage_backend}': {e}"
|
||||
|
||||
self.apply_runtime_config(
|
||||
storage_backend=storage_backend,
|
||||
prefetch_threshold=prefetch_threshold,
|
||||
prefetch_timeout_base=prefetch_timeout_base,
|
||||
prefetch_timeout_per_ki_token=prefetch_timeout_per_ki_token,
|
||||
hicache_storage_pass_prefix_keys=hicache_storage_pass_prefix_keys,
|
||||
enable_storage=True,
|
||||
enable_storage_metrics=cache._enable_metrics_flag,
|
||||
extra_metric_labels=cache.extra_metric_labels,
|
||||
)
|
||||
return True, "Attached HiCache storage backend successfully."
|
||||
|
||||
def detach(self) -> tuple[bool, str]:
|
||||
"""Disable the storage backend at runtime.
|
||||
|
||||
The caller must ensure there are no running or queued requests. Ordering
|
||||
matters and is the reason this is not just ``controller.detach()``:
|
||||
|
||||
1. drain the control queues while the bookkeeping is still intact, so
|
||||
in-flight acks / releases can still be matched to their nodes --
|
||||
otherwise host pages and host locks leak;
|
||||
2. stop the storage threads;
|
||||
3. only then release whatever prefetch / backup is still tracked, since
|
||||
nothing can race the bookkeeping once the threads are gone;
|
||||
4. drain once more, because step 3 only *queues* the host pages -- this
|
||||
is what hands them back to the pool and its sidecars.
|
||||
"""
|
||||
cache = self._cache
|
||||
controller = cache.cache_controller
|
||||
if controller is None:
|
||||
return False, "HiCache storage backend is not initialized."
|
||||
|
||||
try:
|
||||
cache.drain_storage_control_queues_local()
|
||||
# Idempotent: ask the controller to clean up even when `enable_storage`
|
||||
# is already False, since that may be leftover state from an earlier
|
||||
# partial detach.
|
||||
controller.detach_storage_backend()
|
||||
except Exception as e:
|
||||
logger.exception("Failed to detach storage backend.")
|
||||
# Never crash the server for an admin operation. The controller raises
|
||||
# while its threads are still alive, so leave `ongoing_*` untouched --
|
||||
# a retry must still be able to match their acks.
|
||||
return False, f"Failed to detach HiCache storage backend: {e}"
|
||||
|
||||
try:
|
||||
self._release_pending_storage_ops()
|
||||
cache.drain_storage_control_queues_local()
|
||||
except Exception:
|
||||
logger.exception("Failed post-detach cleanup of storage bookkeeping.")
|
||||
|
||||
cache.enable_storage = False
|
||||
cache.enable_storage_metrics = False
|
||||
return True, "Detached HiCache storage backend successfully."
|
||||
|
||||
def shutdown(self) -> None:
|
||||
"""Best-effort auto-detach on process shutdown.
|
||||
|
||||
Keeps startup and runtime behavior consistent: a backend attached either
|
||||
via CLI args or via the admin API is detached on exit.
|
||||
"""
|
||||
try:
|
||||
if self._cache.enable_storage:
|
||||
self.detach()
|
||||
except Exception:
|
||||
logger.exception("Failed to detach storage backend on process shutdown.")
|
||||
|
||||
def clear(self) -> bool:
|
||||
"""Drop everything the backend has stored, keeping it attached."""
|
||||
try:
|
||||
ok = self._cache.cache_controller.clear_storage_backend()
|
||||
except Exception as e:
|
||||
logger.error("Failed to clear hierarchical cache storage backend: %s", e)
|
||||
return False
|
||||
if ok:
|
||||
logger.info("Hierarchical cache storage backend cleared successfully!")
|
||||
return ok
|
||||
|
||||
# ---- Runtime config ----
|
||||
|
||||
def apply_runtime_config(
|
||||
self,
|
||||
*,
|
||||
storage_backend: Optional[str],
|
||||
prefetch_threshold: int,
|
||||
prefetch_timeout_base: float,
|
||||
prefetch_timeout_per_ki_token: float,
|
||||
hicache_storage_pass_prefix_keys: bool,
|
||||
enable_storage: bool,
|
||||
enable_storage_metrics: bool,
|
||||
extra_metric_labels: Optional[dict[str, str]],
|
||||
) -> None:
|
||||
"""Publish the storage knobs onto the tree; the single storage-enable point.
|
||||
|
||||
Both startup and runtime attach funnel through here, so anything that must
|
||||
happen exactly when storage turns on belongs in this method.
|
||||
"""
|
||||
cache = self._cache
|
||||
|
||||
# Nodes already in the tree were built with hashing off. Fill them in as
|
||||
# storage turns on: a node hashed against an unhashed parent restarts the
|
||||
# page hash chain mid-sequence, so its L3 keys would cover only a suffix of
|
||||
# the prefix they claim to represent.
|
||||
if enable_storage and not cache.enable_storage:
|
||||
filled = cache.tree_core.backfill_missing_hash_values()
|
||||
if filled:
|
||||
logger.info(
|
||||
"Hashed %d radix nodes that predate the storage backend.", filled
|
||||
)
|
||||
|
||||
cache.enable_storage = enable_storage
|
||||
cache.prefetch_threshold = prefetch_threshold
|
||||
cache.prefetch_timeout_base = prefetch_timeout_base
|
||||
cache.prefetch_timeout_per_page = (
|
||||
cache.page_size / 1024 * prefetch_timeout_per_ki_token
|
||||
)
|
||||
cache.hicache_storage_pass_prefix_keys = hicache_storage_pass_prefix_keys
|
||||
cache.enable_storage_metrics = enable_storage_metrics
|
||||
|
||||
if enable_storage_metrics:
|
||||
cache.storage_metrics_collector = self._resolve_metrics_collector(
|
||||
storage_backend, extra_metric_labels
|
||||
)
|
||||
else:
|
||||
cache.storage_metrics_collector = None
|
||||
|
||||
def _resolve_metrics_collector(
|
||||
self,
|
||||
storage_backend: Optional[str],
|
||||
extra_metric_labels: Optional[dict[str, str]],
|
||||
) -> Optional[StorageMetricsCollector]:
|
||||
"""Build the storage metrics collector, or relabel the existing one.
|
||||
|
||||
A collector is created once and kept across detach / re-attach: building a
|
||||
second one with the same labels would register duplicate metrics.
|
||||
"""
|
||||
cache = self._cache
|
||||
controller = cache.cache_controller
|
||||
attn_cp_rank, attn_cp_size = controller.get_attn_cp_rank_and_size()
|
||||
labels = {
|
||||
"storage_backend": storage_backend,
|
||||
"tp_rank": controller.tp_rank,
|
||||
"dp_rank": controller.dp_rank,
|
||||
"pp_rank": controller.pp_rank,
|
||||
"pp_size": controller.pp_size,
|
||||
"attn_cp_rank": attn_cp_rank,
|
||||
"attn_cp_size": attn_cp_size,
|
||||
}
|
||||
if extra_metric_labels:
|
||||
labels.update(extra_metric_labels)
|
||||
|
||||
existing_collector = cache.storage_metrics_collector
|
||||
if existing_collector is None:
|
||||
from sglang.srt.runtime_context import get_server_args
|
||||
|
||||
storage_cls = resolve_collector_class(
|
||||
get_server_args(),
|
||||
STAT_LOGGER_ROLE_STORAGE,
|
||||
StorageMetricsCollector,
|
||||
)
|
||||
return storage_cls(labels=labels)
|
||||
|
||||
if set(existing_collector.labels.keys()) == set(labels.keys()):
|
||||
existing_collector.labels = labels
|
||||
else:
|
||||
logger.warning(
|
||||
"Storage metrics labels changed (%s -> %s). Keep existing labels to avoid duplicate metric registration.",
|
||||
sorted(existing_collector.labels.keys()),
|
||||
sorted(labels.keys()),
|
||||
)
|
||||
return existing_collector
|
||||
|
||||
# ---- Internals ----
|
||||
|
||||
@staticmethod
|
||||
def _validate_policies(
|
||||
hicache_storage_prefetch_policy: Optional[str],
|
||||
hicache_write_policy: Optional[str],
|
||||
) -> Optional[str]:
|
||||
"""The rejection reason for an invalid policy, or None when both are ok."""
|
||||
if (
|
||||
hicache_storage_prefetch_policy is not None
|
||||
and hicache_storage_prefetch_policy not in _PREFETCH_POLICIES
|
||||
):
|
||||
return (
|
||||
f"Invalid hicache_storage_prefetch_policy: "
|
||||
f"{hicache_storage_prefetch_policy!r}. "
|
||||
f"Expected one of {list(_PREFETCH_POLICIES)}."
|
||||
)
|
||||
if (
|
||||
hicache_write_policy is not None
|
||||
and hicache_write_policy not in _WRITE_POLICIES
|
||||
):
|
||||
return (
|
||||
f"Invalid hicache_write_policy: {hicache_write_policy!r}. "
|
||||
f"Expected one of {list(_WRITE_POLICIES)}."
|
||||
)
|
||||
return None
|
||||
|
||||
def _apply_policies(
|
||||
self,
|
||||
hicache_storage_prefetch_policy: Optional[str],
|
||||
hicache_write_policy: Optional[str],
|
||||
) -> None:
|
||||
cache = self._cache
|
||||
if hicache_storage_prefetch_policy is not None:
|
||||
cache.prefetch_stop_policy = hicache_storage_prefetch_policy
|
||||
logger.info(
|
||||
f"Set hicache_storage_prefetch_policy to "
|
||||
f"{hicache_storage_prefetch_policy}"
|
||||
)
|
||||
if hicache_write_policy is not None:
|
||||
cache.cache_controller.write_policy = hicache_write_policy
|
||||
cache.write_through_threshold = (
|
||||
1 if hicache_write_policy == "write_through" else 2
|
||||
)
|
||||
cache.is_write_back = hicache_write_policy == "write_back"
|
||||
logger.info(f"Set hicache_write_policy to {hicache_write_policy}")
|
||||
|
||||
def _release_pending_storage_ops(self) -> None:
|
||||
"""Release the host pages and locks that still-tracked ops hold.
|
||||
|
||||
A safety net only: the scheduler refuses a detach unless `ongoing_prefetch`
|
||||
and `ongoing_backup` are already empty, so this normally finds nothing. It
|
||||
earns its keep on the atexit path, which has no such guard.
|
||||
|
||||
Two deliberate departures from `release_aborted_request`, which does the
|
||||
same bookkeeping while the storage threads are still running:
|
||||
|
||||
* no cross-rank barrier -- at process exit the other ranks may already be
|
||||
gone, and a barrier there would hang;
|
||||
* only the completed prefix of a prefetch is freed. The prefetch IO thread
|
||||
owns ``host_indices[completed_tokens:]`` and frees it as it drains, so
|
||||
freeing the whole range here would double-free whatever it already
|
||||
returned. Leaking the tail of an interrupted prefetch is the safer
|
||||
failure of the two.
|
||||
"""
|
||||
cache = self._cache
|
||||
controller = cache.cache_controller
|
||||
|
||||
for req_id in list(cache.ongoing_prefetch):
|
||||
info = cache.ongoing_prefetch[req_id]
|
||||
try:
|
||||
if info.host_indices is None:
|
||||
# Host pages were never allocated for this operation.
|
||||
cache.revoke_pending_prefetch(req_id)
|
||||
continue
|
||||
completed_tokens, _ = controller.terminate_prefetch(info.operation)
|
||||
del cache.ongoing_prefetch[req_id]
|
||||
cache.dec_host_lock_ref(info.anchor_node_id, info.anchor_lock_params)
|
||||
controller.append_host_mem_release(
|
||||
host_indices=info.host_indices[:completed_tokens],
|
||||
extra_pools=[
|
||||
x for xfers in info.comp_xfers.values() for x in xfers
|
||||
],
|
||||
)
|
||||
controller.prefetch_tokens_occupied = max(
|
||||
0, controller.prefetch_tokens_occupied - len(info.prefetch_key)
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Failed to release pending prefetch %s", req_id)
|
||||
cache.ongoing_prefetch.pop(req_id, None)
|
||||
|
||||
for ack_id in list(cache.ongoing_backup):
|
||||
node_id, lock_params = cache.ongoing_backup.pop(ack_id)
|
||||
try:
|
||||
cache.dec_host_lock_ref(node_id, lock_params)
|
||||
except Exception:
|
||||
logger.exception("Failed to release host lock for backup op %s", ack_id)
|
||||
|
||||
cache.prefetch_loaded_tokens_by_reqid.clear()
|
||||
@@ -515,6 +515,27 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
|
||||
"""The hash values owned by this node, excluding its ancestors."""
|
||||
return self.node_by_id(node_id).hash_value or []
|
||||
|
||||
def backfill_missing_hash_values(self) -> int:
|
||||
"""Hash every node that was built while storage was disabled.
|
||||
|
||||
Page hashes chain from the parent's last hash, so a node whose parent has
|
||||
none restarts the chain mid-sequence: its keys then encode only a suffix
|
||||
of the prefix they claim to represent, which would alias unrelated
|
||||
requests that happen to start with those tokens. Walks parent-before-child
|
||||
so each node hashes against an already-filled parent. Returns the number
|
||||
of nodes filled.
|
||||
"""
|
||||
filled = 0
|
||||
stack = [self.root_node]
|
||||
while stack:
|
||||
node = stack.pop()
|
||||
# The root anchors every chain and is seeded with an empty hash list.
|
||||
if node is not self.root_node and node.hash_value is None:
|
||||
node.hash_value = compute_node_hash_values(node, self.page_size)
|
||||
filled += 1
|
||||
stack.extend(node.children.values())
|
||||
return filled
|
||||
|
||||
def root_node_handle(self, extra_key: Optional[str] = None) -> NodeId:
|
||||
"""The NodeId anchoring matches; the single root serves every namespace."""
|
||||
return self.root_node.id
|
||||
|
||||
@@ -173,6 +173,16 @@ class UnifiedTreeCoreInterface(ABC):
|
||||
"""The hash values owned by this node, excluding its ancestors."""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def backfill_missing_hash_values(self) -> int:
|
||||
"""Hash every node built while storage was disabled; return how many.
|
||||
|
||||
Called when a storage backend is attached at runtime: nodes already in
|
||||
the tree carry no hash, and hashing their descendants against them would
|
||||
restart the page hash chain mid-sequence.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def root_node_handle(self, extra_key: Optional[str] = None) -> NodeId:
|
||||
"""The NodeId anchoring matches for the namespace."""
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import atexit
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
@@ -69,6 +70,7 @@ from sglang.srt.mem_cache.unified_cache.components import (
|
||||
from sglang.srt.mem_cache.unified_cache.session_ref_tracker import (
|
||||
UnifiedSessionRefTracker,
|
||||
)
|
||||
from sglang.srt.mem_cache.unified_cache.storage_attachment import StorageAttachment
|
||||
from sglang.srt.mem_cache.unified_cache.tree_core_registry import create_tree_core
|
||||
from sglang.srt.mem_cache.unified_cache.unified_tree_core import ( # noqa: F401
|
||||
NodeId,
|
||||
@@ -77,10 +79,8 @@ from sglang.srt.mem_cache.unified_cache.unified_tree_core import ( # noqa: F401
|
||||
UnifiedTreeNode,
|
||||
)
|
||||
from sglang.srt.observability.metrics_collector import (
|
||||
STAT_LOGGER_ROLE_STORAGE,
|
||||
StorageMetrics,
|
||||
StorageMetricsCollector,
|
||||
resolve_collector_class,
|
||||
)
|
||||
from sglang.srt.session.streaming_session import StreamingSession
|
||||
from sglang.srt.utils.common import ceil_align
|
||||
@@ -230,6 +230,8 @@ class UnifiedRadixCache(BasePrefixCache):
|
||||
# HiCache D↔H defaults (overridden by init_hicache)
|
||||
self.cache_controller: Optional[HybridCacheController] = None
|
||||
self.host_pool_group = None # set by attach_hybrid_pool_to_unified_cache
|
||||
# Owns the storage backend lifecycle; built by init_hicache.
|
||||
self._storage_attachment: Optional[StorageAttachment] = None
|
||||
self.prefetch_stop_policy = "best_effort"
|
||||
self.prefetch_threshold = 256
|
||||
self.prefetch_timeout_base = 1.0
|
||||
@@ -457,8 +459,12 @@ class UnifiedRadixCache(BasePrefixCache):
|
||||
self.load_back_threshold = 10
|
||||
self.prefetch_stop_policy = server_args.hicache_storage_prefetch_policy
|
||||
|
||||
# Runtime attach/detach of the L3 backend (startup, admin API, atexit).
|
||||
self._storage_attachment = StorageAttachment(self)
|
||||
atexit.register(self.shutdown)
|
||||
|
||||
if storage_backend is not None:
|
||||
self._apply_storage_runtime_config(
|
||||
self._storage_attachment.apply_runtime_config(
|
||||
storage_backend=storage_backend,
|
||||
prefetch_threshold=storage_prefetch_threshold,
|
||||
prefetch_timeout_base=prefetch_timeout_base,
|
||||
@@ -1734,7 +1740,7 @@ class UnifiedRadixCache(BasePrefixCache):
|
||||
return False
|
||||
if operation.host_indices is None:
|
||||
self.cache_controller.terminate_prefetch(operation)
|
||||
self._revoke_pending_prefetch(req_id)
|
||||
self.revoke_pending_prefetch(req_id)
|
||||
return True
|
||||
|
||||
completed_tokens, hash_value = self.cache_controller.terminate_prefetch(
|
||||
@@ -1955,7 +1961,7 @@ class UnifiedRadixCache(BasePrefixCache):
|
||||
) = self.ongoing_prefetch[rid]
|
||||
if operation.host_indices is None:
|
||||
self.cache_controller.terminate_prefetch(operation)
|
||||
self._revoke_pending_prefetch(rid)
|
||||
self.revoke_pending_prefetch(rid)
|
||||
return
|
||||
|
||||
completed_tokens, _ = self.cache_controller.terminate_prefetch(operation)
|
||||
@@ -2028,7 +2034,7 @@ class UnifiedRadixCache(BasePrefixCache):
|
||||
return len(host_indices) if host_indices is not None else 0
|
||||
return len(prefetch_key)
|
||||
|
||||
def _revoke_pending_prefetch(self, req_id: str) -> None:
|
||||
def revoke_pending_prefetch(self, req_id: str) -> None:
|
||||
info = self.ongoing_prefetch.pop(req_id, None)
|
||||
if info is None:
|
||||
return
|
||||
@@ -2090,7 +2096,7 @@ class UnifiedRadixCache(BasePrefixCache):
|
||||
if info is None:
|
||||
return True # aborted/cleaned; nothing to retry
|
||||
if operation.is_terminated():
|
||||
self._revoke_pending_prefetch(req_id)
|
||||
self.revoke_pending_prefetch(req_id)
|
||||
return True
|
||||
|
||||
if buffer_mode and cc.prefetch_rate_limited():
|
||||
@@ -2116,7 +2122,7 @@ class UnifiedRadixCache(BasePrefixCache):
|
||||
if host_indices is None:
|
||||
if buffer_mode:
|
||||
return False
|
||||
self._revoke_pending_prefetch(req_id)
|
||||
self.revoke_pending_prefetch(req_id)
|
||||
return True
|
||||
|
||||
operation.storage_hit_count = alloc_len
|
||||
@@ -2146,13 +2152,13 @@ class UnifiedRadixCache(BasePrefixCache):
|
||||
continue
|
||||
if operation.is_terminated():
|
||||
# Aborted while the storage query was in flight.
|
||||
self._revoke_pending_prefetch(req_id)
|
||||
self.revoke_pending_prefetch(req_id)
|
||||
continue
|
||||
if operation.storage_hit_count < self.prefetch_threshold:
|
||||
# Below-threshold hit: classify + feed the L3 miss
|
||||
# accounting, then revoke (not enough benefit).
|
||||
self._account_prefetch_outcome(operation, revoked=True)
|
||||
self._revoke_pending_prefetch(req_id)
|
||||
self.revoke_pending_prefetch(req_id)
|
||||
continue
|
||||
self._invalidate_absent_from_hit_query(operation)
|
||||
self._account_prefetch_outcome(operation, revoked=False)
|
||||
@@ -2250,62 +2256,30 @@ class UnifiedRadixCache(BasePrefixCache):
|
||||
log_metrics=True,
|
||||
)
|
||||
|
||||
def _apply_storage_runtime_config(
|
||||
self,
|
||||
*,
|
||||
storage_backend: Optional[str],
|
||||
prefetch_threshold: int,
|
||||
prefetch_timeout_base: float,
|
||||
prefetch_timeout_per_ki_token: float,
|
||||
hicache_storage_pass_prefix_keys: bool,
|
||||
enable_storage: bool,
|
||||
enable_storage_metrics: bool,
|
||||
extra_metric_labels: Optional[dict[str, str]],
|
||||
) -> None:
|
||||
self.enable_storage = enable_storage
|
||||
self.prefetch_threshold = prefetch_threshold
|
||||
self.prefetch_timeout_base = prefetch_timeout_base
|
||||
self.prefetch_timeout_per_page = (
|
||||
self.page_size / 1024 * prefetch_timeout_per_ki_token
|
||||
def drain_storage_control_queues_local(self) -> None:
|
||||
"""Drain the storage control queues without cross-rank synchronization.
|
||||
|
||||
For the detach / shutdown path, where best-effort cleanup matters more than
|
||||
keeping the drained counts identical across ranks. The prefetch-hit queue is
|
||||
deliberately skipped: servicing it would allocate host pages for a prefetch
|
||||
that can no longer complete.
|
||||
"""
|
||||
cc = self.cache_controller
|
||||
# The storage queues are created by the controller when the storage threads
|
||||
# start, so they are still None when a backend was never attached.
|
||||
if cc is None or cc.prefetch_hit_queue is None:
|
||||
return
|
||||
self._drain_storage_control_queues_impl(
|
||||
n_storage_hit=0,
|
||||
n_backup=None,
|
||||
n_release=None,
|
||||
extra_release_counts={
|
||||
name: None for name in cc.extra_host_mem_release_queues
|
||||
},
|
||||
log_metrics=False,
|
||||
)
|
||||
self.hicache_storage_pass_prefix_keys = hicache_storage_pass_prefix_keys
|
||||
self.enable_storage_metrics = enable_storage_metrics
|
||||
|
||||
if self.enable_storage_metrics:
|
||||
attn_cp_rank, attn_cp_size = (
|
||||
self.cache_controller.get_attn_cp_rank_and_size()
|
||||
)
|
||||
labels = {
|
||||
"storage_backend": storage_backend,
|
||||
"tp_rank": self.cache_controller.tp_rank,
|
||||
"dp_rank": self.cache_controller.dp_rank,
|
||||
"pp_rank": self.cache_controller.pp_rank,
|
||||
"pp_size": self.cache_controller.pp_size,
|
||||
"attn_cp_rank": attn_cp_rank,
|
||||
"attn_cp_size": attn_cp_size,
|
||||
}
|
||||
if extra_metric_labels:
|
||||
labels.update(extra_metric_labels)
|
||||
existing_collector = self.storage_metrics_collector
|
||||
if existing_collector is None:
|
||||
from sglang.srt.runtime_context import get_server_args
|
||||
|
||||
storage_cls = resolve_collector_class(
|
||||
get_server_args(),
|
||||
STAT_LOGGER_ROLE_STORAGE,
|
||||
StorageMetricsCollector,
|
||||
)
|
||||
self.storage_metrics_collector = storage_cls(labels=labels)
|
||||
elif set(existing_collector.labels.keys()) == set(labels.keys()):
|
||||
existing_collector.labels = labels
|
||||
else:
|
||||
logger.warning(
|
||||
"Storage metrics labels changed (%s -> %s). Keep existing labels to avoid duplicate metric registration.",
|
||||
sorted(existing_collector.labels.keys()),
|
||||
sorted(labels.keys()),
|
||||
)
|
||||
else:
|
||||
self.storage_metrics_collector = None
|
||||
# ---- HiCache: Storage backend lifecycle (delegated) ----
|
||||
|
||||
def attach_storage_backend(
|
||||
self,
|
||||
@@ -2315,30 +2289,40 @@ class UnifiedRadixCache(BasePrefixCache):
|
||||
hicache_storage_prefetch_policy: Optional[str] = None,
|
||||
hicache_write_policy: Optional[str] = None,
|
||||
) -> tuple[bool, str]:
|
||||
return (
|
||||
False,
|
||||
"UnifiedRadixCache does not support runtime HiCache storage attach yet. "
|
||||
"Configure hicache_storage_backend at startup instead.",
|
||||
"""Attach (enable) the HiCache storage backend at runtime."""
|
||||
if self._storage_attachment is None:
|
||||
return (
|
||||
False,
|
||||
"HiCache is not initialized; launch with "
|
||||
"--enable-hierarchical-cache to attach a storage backend.",
|
||||
)
|
||||
return self._storage_attachment.attach(
|
||||
storage_backend=storage_backend,
|
||||
storage_backend_extra_config_json=storage_backend_extra_config_json,
|
||||
served_model_name=served_model_name,
|
||||
hicache_storage_prefetch_policy=hicache_storage_prefetch_policy,
|
||||
hicache_write_policy=hicache_write_policy,
|
||||
)
|
||||
|
||||
def detach_storage_backend(self) -> tuple[bool, str]:
|
||||
return (
|
||||
False,
|
||||
"UnifiedRadixCache does not support runtime HiCache storage detach yet. "
|
||||
"Restart without hicache_storage_backend to disable it.",
|
||||
)
|
||||
"""Detach (disable) the HiCache storage backend at runtime."""
|
||||
if self._storage_attachment is None:
|
||||
return False, "HiCache storage backend is not initialized."
|
||||
return self._storage_attachment.detach()
|
||||
|
||||
def shutdown(self) -> None:
|
||||
"""Best-effort auto-detach of the storage backend on process shutdown."""
|
||||
if self._storage_attachment is not None:
|
||||
self._storage_attachment.shutdown()
|
||||
|
||||
def clear_storage_backend(self) -> bool:
|
||||
try:
|
||||
ok = self.cache_controller.clear_storage_backend()
|
||||
except Exception as e:
|
||||
logger.error("Failed to clear hierarchical cache storage backend: %s", e)
|
||||
if self._storage_attachment is None:
|
||||
return False
|
||||
ok = self._storage_attachment.clear()
|
||||
if ok:
|
||||
# L3 is empty now: every storage-presence belief is stale, and a
|
||||
# retained positive would skip that page's backup forever.
|
||||
self.storage_existence_cache.clear()
|
||||
logger.info("Hierarchical cache storage backend cleared successfully!")
|
||||
return ok
|
||||
|
||||
# ---- HiCache: Async Event Management ----
|
||||
|
||||
Reference in New Issue
Block a user