[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).
|
# Dedicated stop event for storage background threads (prefetch/backup).
|
||||||
self.storage_stop_event = threading.Event()
|
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.device = self.mem_pool_device.device
|
||||||
self.layer_num = self.mem_pool_device.layer_num
|
self.layer_num = self.mem_pool_device.layer_num
|
||||||
self.layer_done_counter = LayerDoneCounter(self.layer_num)
|
self.layer_done_counter = LayerDoneCounter(self.layer_num)
|
||||||
@@ -417,9 +422,9 @@ class HiCacheController:
|
|||||||
self.prefetch_queue = Queue()
|
self.prefetch_queue = Queue()
|
||||||
self.backup_queue = Queue()
|
self.backup_queue = Queue()
|
||||||
|
|
||||||
self.prefetch_hit_queue: Queue[StorageOperation] = Queue()
|
self.prefetch_hit_queue = Queue()
|
||||||
self.ack_backup_queue: Queue[StorageOperation] = Queue()
|
self.ack_backup_queue = Queue()
|
||||||
self.host_mem_release_queue: Queue[torch.Tensor] = Queue()
|
self.host_mem_release_queue = Queue()
|
||||||
|
|
||||||
self.prefetch_thread.start()
|
self.prefetch_thread.start()
|
||||||
self.backup_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."""
|
"""The hash values owned by this node, excluding its ancestors."""
|
||||||
return self.node_by_id(node_id).hash_value or []
|
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:
|
def root_node_handle(self, extra_key: Optional[str] = None) -> NodeId:
|
||||||
"""The NodeId anchoring matches; the single root serves every namespace."""
|
"""The NodeId anchoring matches; the single root serves every namespace."""
|
||||||
return self.root_node.id
|
return self.root_node.id
|
||||||
|
|||||||
@@ -173,6 +173,16 @@ class UnifiedTreeCoreInterface(ABC):
|
|||||||
"""The hash values owned by this node, excluding its ancestors."""
|
"""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
|
@abstractmethod
|
||||||
def root_node_handle(self, extra_key: Optional[str] = None) -> NodeId:
|
def root_node_handle(self, extra_key: Optional[str] = None) -> NodeId:
|
||||||
"""The NodeId anchoring matches for the namespace."""
|
"""The NodeId anchoring matches for the namespace."""
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import atexit
|
||||||
import logging
|
import logging
|
||||||
import threading
|
import threading
|
||||||
import time
|
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 (
|
from sglang.srt.mem_cache.unified_cache.session_ref_tracker import (
|
||||||
UnifiedSessionRefTracker,
|
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.tree_core_registry import create_tree_core
|
||||||
from sglang.srt.mem_cache.unified_cache.unified_tree_core import ( # noqa: F401
|
from sglang.srt.mem_cache.unified_cache.unified_tree_core import ( # noqa: F401
|
||||||
NodeId,
|
NodeId,
|
||||||
@@ -77,10 +79,8 @@ from sglang.srt.mem_cache.unified_cache.unified_tree_core import ( # noqa: F401
|
|||||||
UnifiedTreeNode,
|
UnifiedTreeNode,
|
||||||
)
|
)
|
||||||
from sglang.srt.observability.metrics_collector import (
|
from sglang.srt.observability.metrics_collector import (
|
||||||
STAT_LOGGER_ROLE_STORAGE,
|
|
||||||
StorageMetrics,
|
StorageMetrics,
|
||||||
StorageMetricsCollector,
|
StorageMetricsCollector,
|
||||||
resolve_collector_class,
|
|
||||||
)
|
)
|
||||||
from sglang.srt.session.streaming_session import StreamingSession
|
from sglang.srt.session.streaming_session import StreamingSession
|
||||||
from sglang.srt.utils.common import ceil_align
|
from sglang.srt.utils.common import ceil_align
|
||||||
@@ -230,6 +230,8 @@ class UnifiedRadixCache(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.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_stop_policy = "best_effort"
|
||||||
self.prefetch_threshold = 256
|
self.prefetch_threshold = 256
|
||||||
self.prefetch_timeout_base = 1.0
|
self.prefetch_timeout_base = 1.0
|
||||||
@@ -457,8 +459,12 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
self.load_back_threshold = 10
|
self.load_back_threshold = 10
|
||||||
self.prefetch_stop_policy = server_args.hicache_storage_prefetch_policy
|
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:
|
if storage_backend is not None:
|
||||||
self._apply_storage_runtime_config(
|
self._storage_attachment.apply_runtime_config(
|
||||||
storage_backend=storage_backend,
|
storage_backend=storage_backend,
|
||||||
prefetch_threshold=storage_prefetch_threshold,
|
prefetch_threshold=storage_prefetch_threshold,
|
||||||
prefetch_timeout_base=prefetch_timeout_base,
|
prefetch_timeout_base=prefetch_timeout_base,
|
||||||
@@ -1734,7 +1740,7 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
return False
|
return False
|
||||||
if operation.host_indices is None:
|
if operation.host_indices is None:
|
||||||
self.cache_controller.terminate_prefetch(operation)
|
self.cache_controller.terminate_prefetch(operation)
|
||||||
self._revoke_pending_prefetch(req_id)
|
self.revoke_pending_prefetch(req_id)
|
||||||
return True
|
return True
|
||||||
|
|
||||||
completed_tokens, hash_value = self.cache_controller.terminate_prefetch(
|
completed_tokens, hash_value = self.cache_controller.terminate_prefetch(
|
||||||
@@ -1955,7 +1961,7 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
) = self.ongoing_prefetch[rid]
|
) = self.ongoing_prefetch[rid]
|
||||||
if operation.host_indices is None:
|
if operation.host_indices is None:
|
||||||
self.cache_controller.terminate_prefetch(operation)
|
self.cache_controller.terminate_prefetch(operation)
|
||||||
self._revoke_pending_prefetch(rid)
|
self.revoke_pending_prefetch(rid)
|
||||||
return
|
return
|
||||||
|
|
||||||
completed_tokens, _ = self.cache_controller.terminate_prefetch(operation)
|
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(host_indices) if host_indices is not None else 0
|
||||||
return len(prefetch_key)
|
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)
|
info = self.ongoing_prefetch.pop(req_id, None)
|
||||||
if info is None:
|
if info is None:
|
||||||
return
|
return
|
||||||
@@ -2090,7 +2096,7 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
if info is None:
|
if info is None:
|
||||||
return True # aborted/cleaned; nothing to retry
|
return True # aborted/cleaned; nothing to retry
|
||||||
if operation.is_terminated():
|
if operation.is_terminated():
|
||||||
self._revoke_pending_prefetch(req_id)
|
self.revoke_pending_prefetch(req_id)
|
||||||
return True
|
return True
|
||||||
|
|
||||||
if buffer_mode and cc.prefetch_rate_limited():
|
if buffer_mode and cc.prefetch_rate_limited():
|
||||||
@@ -2116,7 +2122,7 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
if host_indices is None:
|
if host_indices is None:
|
||||||
if buffer_mode:
|
if buffer_mode:
|
||||||
return False
|
return False
|
||||||
self._revoke_pending_prefetch(req_id)
|
self.revoke_pending_prefetch(req_id)
|
||||||
return True
|
return True
|
||||||
|
|
||||||
operation.storage_hit_count = alloc_len
|
operation.storage_hit_count = alloc_len
|
||||||
@@ -2146,13 +2152,13 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
continue
|
continue
|
||||||
if operation.is_terminated():
|
if operation.is_terminated():
|
||||||
# Aborted while the storage query was in flight.
|
# Aborted while the storage query was in flight.
|
||||||
self._revoke_pending_prefetch(req_id)
|
self.revoke_pending_prefetch(req_id)
|
||||||
continue
|
continue
|
||||||
if operation.storage_hit_count < self.prefetch_threshold:
|
if operation.storage_hit_count < self.prefetch_threshold:
|
||||||
# Below-threshold hit: classify + feed the L3 miss
|
# Below-threshold hit: classify + feed the L3 miss
|
||||||
# accounting, then revoke (not enough benefit).
|
# accounting, then revoke (not enough benefit).
|
||||||
self._account_prefetch_outcome(operation, revoked=True)
|
self._account_prefetch_outcome(operation, revoked=True)
|
||||||
self._revoke_pending_prefetch(req_id)
|
self.revoke_pending_prefetch(req_id)
|
||||||
continue
|
continue
|
||||||
self._invalidate_absent_from_hit_query(operation)
|
self._invalidate_absent_from_hit_query(operation)
|
||||||
self._account_prefetch_outcome(operation, revoked=False)
|
self._account_prefetch_outcome(operation, revoked=False)
|
||||||
@@ -2250,62 +2256,30 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
log_metrics=True,
|
log_metrics=True,
|
||||||
)
|
)
|
||||||
|
|
||||||
def _apply_storage_runtime_config(
|
def drain_storage_control_queues_local(self) -> None:
|
||||||
self,
|
"""Drain the storage control queues without cross-rank synchronization.
|
||||||
*,
|
|
||||||
storage_backend: Optional[str],
|
For the detach / shutdown path, where best-effort cleanup matters more than
|
||||||
prefetch_threshold: int,
|
keeping the drained counts identical across ranks. The prefetch-hit queue is
|
||||||
prefetch_timeout_base: float,
|
deliberately skipped: servicing it would allocate host pages for a prefetch
|
||||||
prefetch_timeout_per_ki_token: float,
|
that can no longer complete.
|
||||||
hicache_storage_pass_prefix_keys: bool,
|
"""
|
||||||
enable_storage: bool,
|
cc = self.cache_controller
|
||||||
enable_storage_metrics: bool,
|
# The storage queues are created by the controller when the storage threads
|
||||||
extra_metric_labels: Optional[dict[str, str]],
|
# start, so they are still None when a backend was never attached.
|
||||||
) -> None:
|
if cc is None or cc.prefetch_hit_queue is None:
|
||||||
self.enable_storage = enable_storage
|
return
|
||||||
self.prefetch_threshold = prefetch_threshold
|
self._drain_storage_control_queues_impl(
|
||||||
self.prefetch_timeout_base = prefetch_timeout_base
|
n_storage_hit=0,
|
||||||
self.prefetch_timeout_per_page = (
|
n_backup=None,
|
||||||
self.page_size / 1024 * prefetch_timeout_per_ki_token
|
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:
|
# ---- HiCache: Storage backend lifecycle (delegated) ----
|
||||||
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
|
|
||||||
|
|
||||||
def attach_storage_backend(
|
def attach_storage_backend(
|
||||||
self,
|
self,
|
||||||
@@ -2315,30 +2289,40 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
hicache_storage_prefetch_policy: Optional[str] = None,
|
hicache_storage_prefetch_policy: Optional[str] = None,
|
||||||
hicache_write_policy: Optional[str] = None,
|
hicache_write_policy: Optional[str] = None,
|
||||||
) -> tuple[bool, str]:
|
) -> tuple[bool, str]:
|
||||||
return (
|
"""Attach (enable) the HiCache storage backend at runtime."""
|
||||||
False,
|
if self._storage_attachment is None:
|
||||||
"UnifiedRadixCache does not support runtime HiCache storage attach yet. "
|
return (
|
||||||
"Configure hicache_storage_backend at startup instead.",
|
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]:
|
def detach_storage_backend(self) -> tuple[bool, str]:
|
||||||
return (
|
"""Detach (disable) the HiCache storage backend at runtime."""
|
||||||
False,
|
if self._storage_attachment is None:
|
||||||
"UnifiedRadixCache does not support runtime HiCache storage detach yet. "
|
return False, "HiCache storage backend is not initialized."
|
||||||
"Restart without hicache_storage_backend to disable it.",
|
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:
|
def clear_storage_backend(self) -> bool:
|
||||||
try:
|
if self._storage_attachment is None:
|
||||||
ok = self.cache_controller.clear_storage_backend()
|
|
||||||
except Exception as e:
|
|
||||||
logger.error("Failed to clear hierarchical cache storage backend: %s", e)
|
|
||||||
return False
|
return False
|
||||||
|
ok = self._storage_attachment.clear()
|
||||||
if ok:
|
if ok:
|
||||||
# L3 is empty now: every storage-presence belief is stale, and a
|
# L3 is empty now: every storage-presence belief is stale, and a
|
||||||
# retained positive would skip that page's backup forever.
|
# retained positive would skip that page's backup forever.
|
||||||
self.storage_existence_cache.clear()
|
self.storage_existence_cache.clear()
|
||||||
logger.info("Hierarchical cache storage backend cleared successfully!")
|
|
||||||
return ok
|
return ok
|
||||||
|
|
||||||
# ---- HiCache: Async Event Management ----
|
# ---- HiCache: Async Event Management ----
|
||||||
|
|||||||
@@ -28,10 +28,13 @@ from sglang.test.test_utils import (
|
|||||||
)
|
)
|
||||||
from sglang.utils import wait_for_http_ready
|
from sglang.utils import wait_for_http_ready
|
||||||
|
|
||||||
register_cuda_ci(est_time=139, stage="base-b", runner_config="2-gpu-large")
|
register_cuda_ci(est_time=210, stage="base-b", runner_config="2-gpu-large")
|
||||||
|
|
||||||
|
|
||||||
class TestHiCacheStorageRuntimeAttachDetach(CustomTestCase):
|
class TestHiCacheStorageRuntimeAttachDetach(CustomTestCase):
|
||||||
|
# Extra server env; subclasses use it to select a tree_cache implementation.
|
||||||
|
extra_env: dict = {}
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def setUpClass(cls):
|
def setUpClass(cls):
|
||||||
cls.temp_dir = tempfile.mkdtemp()
|
cls.temp_dir = tempfile.mkdtemp()
|
||||||
@@ -60,6 +63,7 @@ class TestHiCacheStorageRuntimeAttachDetach(CustomTestCase):
|
|||||||
"SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.temp_dir,
|
"SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.temp_dir,
|
||||||
# Make runs less flaky for CI/dev.
|
# Make runs less flaky for CI/dev.
|
||||||
"SGLANG_ENABLE_DETERMINISTIC_INFERENCE": "1",
|
"SGLANG_ENABLE_DETERMINISTIC_INFERENCE": "1",
|
||||||
|
**cls.extra_env,
|
||||||
}
|
}
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
@@ -221,6 +225,10 @@ class TestHiCacheStorageRuntimeAttachDetach(CustomTestCase):
|
|||||||
kill_process_tree(process1.pid)
|
kill_process_tree(process1.pid)
|
||||||
time.sleep(2)
|
time.sleep(2)
|
||||||
|
|
||||||
|
self._check_attach_detach_lifecycle()
|
||||||
|
|
||||||
|
def _check_attach_detach_lifecycle(self):
|
||||||
|
"""Attach/detach lifecycle against a server that requires an admin key."""
|
||||||
# Phase B: WITH --admin-api-key, must provide Authorization: Bearer <admin_key>.
|
# Phase B: WITH --admin-api-key, must provide Authorization: Bearer <admin_key>.
|
||||||
admin_key = "sglang-test-admin-key"
|
admin_key = "sglang-test-admin-key"
|
||||||
base_url2 = f"http://127.0.0.1:{find_available_port(int(self.base_url.rsplit(':', 1)[1]) + 1)}"
|
base_url2 = f"http://127.0.0.1:{find_available_port(int(self.base_url.rsplit(':', 1)[1]) + 1)}"
|
||||||
@@ -363,5 +371,18 @@ class TestHiCacheStorageRuntimeAttachDetach(CustomTestCase):
|
|||||||
time.sleep(2)
|
time.sleep(2)
|
||||||
|
|
||||||
|
|
||||||
|
class TestUnifiedRadixCacheStorageRuntimeAttachDetach(
|
||||||
|
TestHiCacheStorageRuntimeAttachDetach
|
||||||
|
):
|
||||||
|
"""Same runtime attach/detach lifecycle, backed by UnifiedRadixCache."""
|
||||||
|
|
||||||
|
extra_env = {"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"}
|
||||||
|
|
||||||
|
def test_runtime_attach_detach(self):
|
||||||
|
# Admin-key gating (phase A of the base test) lives in the HTTP layer and is
|
||||||
|
# independent of the tree cache implementation, so only the lifecycle is run.
|
||||||
|
self._check_attach_detach_lifecycle()
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
unittest.main(verbosity=2)
|
unittest.main(verbosity=2)
|
||||||
|
|||||||
@@ -65,6 +65,7 @@ from sglang.srt.mem_cache.unified_cache.components.tree_component import (
|
|||||||
EvictLayer,
|
EvictLayer,
|
||||||
TreeComponent,
|
TreeComponent,
|
||||||
)
|
)
|
||||||
|
from sglang.srt.mem_cache.unified_cache.storage_attachment import StorageAttachment
|
||||||
from sglang.srt.mem_cache.unified_cache.unified_tree_core import UnifiedTreeCore
|
from sglang.srt.mem_cache.unified_cache.unified_tree_core import UnifiedTreeCore
|
||||||
from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import (
|
from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import (
|
||||||
DecSwaLockOnlyResult,
|
DecSwaLockOnlyResult,
|
||||||
@@ -7136,5 +7137,92 @@ class TestSWAWindowUnderBigramKey(CustomTestCase):
|
|||||||
cache.sanity_check()
|
cache.sanity_check()
|
||||||
|
|
||||||
|
|
||||||
|
class TestUnifiedRadixCacheStorageAttachBackfill(CustomTestCase):
|
||||||
|
"""Enabling a storage backend must hash nodes that predate it.
|
||||||
|
|
||||||
|
Page hashes chain off the parent's last hash, so a node hashed against an
|
||||||
|
unhashed parent restarts the chain mid-sequence: its L3 keys then cover only a
|
||||||
|
suffix of the prefix they claim to represent, which aliases any unrelated
|
||||||
|
request that happens to start with those tokens.
|
||||||
|
"""
|
||||||
|
|
||||||
|
cfg = CacheConfig(
|
||||||
|
page_size=4,
|
||||||
|
components=(ComponentType.FULL,),
|
||||||
|
kv_size=64,
|
||||||
|
max_context_len=64,
|
||||||
|
)
|
||||||
|
|
||||||
|
prefix_tokens = array("q", [1, 2, 3, 4, 5, 6, 7, 8])
|
||||||
|
suffix_tokens = array("q", [9, 10, 11, 12])
|
||||||
|
|
||||||
|
def _build_two_level_tree(self, *, storage_on_from_the_start: bool):
|
||||||
|
"""A parent node plus a child extending it, built with storage on or off."""
|
||||||
|
cache, allocator, _ = build_fixture(self.cfg)
|
||||||
|
cache.enable_storage = storage_on_from_the_start
|
||||||
|
for tokens in (self.prefix_tokens, self.prefix_tokens + self.suffix_tokens):
|
||||||
|
value = allocator.alloc(len(tokens))
|
||||||
|
self.assertIsNotNone(value)
|
||||||
|
cache.insert(InsertParams(key=RadixKey(tokens), value=value))
|
||||||
|
return cache
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _hashes_by_token_ids(cache):
|
||||||
|
"""Every non-root node's token ids mapped to its hash chain."""
|
||||||
|
root = cache.tree_core.root_node
|
||||||
|
hashes = {}
|
||||||
|
stack = [root]
|
||||||
|
while stack:
|
||||||
|
node = stack.pop()
|
||||||
|
if node is not root:
|
||||||
|
hashes[tuple(node.key.token_ids)] = node.hash_value
|
||||||
|
stack.extend(node.children.values())
|
||||||
|
return hashes
|
||||||
|
|
||||||
|
def test_backfill_reproduces_hashing_from_the_start(self):
|
||||||
|
expected = self._hashes_by_token_ids(
|
||||||
|
self._build_two_level_tree(storage_on_from_the_start=True)
|
||||||
|
)
|
||||||
|
self.assertEqual(len(expected), 2, "expected a parent and a child node")
|
||||||
|
|
||||||
|
late = self._build_two_level_tree(storage_on_from_the_start=False)
|
||||||
|
self.assertTrue(
|
||||||
|
all(h is None for h in self._hashes_by_token_ids(late).values()),
|
||||||
|
"nodes built while storage was disabled must start unhashed",
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertEqual(late.tree_core.backfill_missing_hash_values(), len(expected))
|
||||||
|
self.assertEqual(
|
||||||
|
self._hashes_by_token_ids(late),
|
||||||
|
expected,
|
||||||
|
"a backfilled chain must be identical to one hashed from the start",
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_backfill_is_idempotent(self):
|
||||||
|
cache = self._build_two_level_tree(storage_on_from_the_start=True)
|
||||||
|
before = self._hashes_by_token_ids(cache)
|
||||||
|
self.assertEqual(cache.tree_core.backfill_missing_hash_values(), 0)
|
||||||
|
self.assertEqual(self._hashes_by_token_ids(cache), before)
|
||||||
|
|
||||||
|
def test_enabling_storage_backfills_the_tree(self):
|
||||||
|
"""The tree is hashed by the time `enable_storage` flips on."""
|
||||||
|
cache = self._build_two_level_tree(storage_on_from_the_start=False)
|
||||||
|
StorageAttachment(cache).apply_runtime_config(
|
||||||
|
storage_backend="file",
|
||||||
|
prefetch_threshold=64,
|
||||||
|
prefetch_timeout_base=1.0,
|
||||||
|
prefetch_timeout_per_ki_token=0.25,
|
||||||
|
hicache_storage_pass_prefix_keys=False,
|
||||||
|
enable_storage=True,
|
||||||
|
enable_storage_metrics=False,
|
||||||
|
extra_metric_labels=None,
|
||||||
|
)
|
||||||
|
self.assertTrue(cache.enable_storage)
|
||||||
|
self.assertTrue(
|
||||||
|
all(h for h in self._hashes_by_token_ids(cache).values()),
|
||||||
|
"every node must carry a hash chain once storage is enabled",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
unittest.main()
|
unittest.main()
|
||||||
|
|||||||
Reference in New Issue
Block a user