coffee/DeepInfer
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepInfer.git synced 2026-10-31 00:36:20 +08:00
Code Issues Packages Projects Releases Wiki Activity
DeepInfer/deepinfer/hicache/host_cache.py
whn09 3657e00592 hicache L3: circuit breaker for Store transfer failures, TCP queue limits for Mooncake <= 0.3.13 
A Store transfer that cannot complete blocks for Mooncake's fixed 60 s timeout. Open a breaker on the first
transfer failure or a get slower than DEEPINFER_HICACHE_L3_SLOW_S (5 s); while open, skip reads and drop queued
puts; close after DEEPINFER_HICACHE_L3_COOLDOWN_S (30 s, doubling on consecutive trips, capped at 10 min).

Mooncake <= 0.3.13 queues at most 1024 TCP transfers per peer, fewer than one long prompt's batch put; default
MC_TCP_MAX_QUEUED_TRANSFERS_PER_PEER / MC_TCP_MAX_PENDING_ADMISSIONS_PER_PEER to 65535 (the default since
kvcache-ai/Mooncake#3710) unless set.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-23 21:56:26 +08:00

643 lines
29 KiB
Python
Raw Blame History

"""Host-memory prefix cache for DeepInfer, in the style of SGLang's HiCache.
Layout
Two pinned host pools, both organized in pages:
- KV pages: one page = one paged-KV block (block_size tokens, all layers), same bytes as a GPU block.
- State slots: the SWA "stored region" of a sequence (the part DeepSeek persists), taken at a
128-token-aligned position with GPUCacheInterface.select_state_slice_to_store.
Pages are indexed by a hash chain over the prompt tokens (page i's key covers tokens [0, (i+1)*block_size)),
so a lookup is a walk from the first page until the first miss. Eviction is LRU, separately per pool.
Transfer
"kernel" method only: a small CUDA kernel copies pages between a GPU pool and the pinned host pool through
UVA pointers, on the engine's current stream, so it is ordered with the model's own kernels and never
blocks the CPU. Completion is tracked with CUDA events polled from the engine loop.
Policy
Write-through at the end of every prefill: all full blocks of the prompt plus the state at the last
128-aligned position are stored (pages already present are only touched). On a new request, the longest
page-chain match is capped so that at least `keep_tokens` (the decoder bounded-replay rectangle) remain to
be recomputed, then snapped down to the nearest position that has a stored state; the sequence is
pre-allocated, its blocks and state are filled from the host, and it enters the scheduler with
num_cached_tokens set. Requires DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE=compact (the decoder SWA state is
rebuilt from the recomputed suffix) and DEEPINFER_PREFILL_STORE_EXTRA_STATE=1 (ring slack for the store).
Not covered: tokens generated during decode are not stored (they are re-prefilled as part of the next turn's
prompt and stored then); the cache is per engine process (route a conversation's turns to the same rank).
L3 (optional): with DEEPINFER_HICACHE_MOONCAKE_STORE set, every stored page group and state is also written behind
to a Mooncake Store (deepinfer/hicache/mooncake_store.py), and a host miss is extended from the Store, so prefixes
are shared across DP ranks, server instances and restarts.
"""
from __future__ import annotations
import hashlib
import os
import time
from collections import OrderedDict, deque
from dataclasses import dataclass, field
from typing import TYPE_CHECKING
import numpy as np
import torch
from deepinfer import envs
from deepinfer.logger import init_logger
from deepinfer.worker.cache_engine import KVSTORE_GRAIN, GPUCacheInterface
if TYPE_CHECKING:
from deepinfer._C.engine import Sequence
from deepinfer.engine.llm_engine import LLMEngine
logger = init_logger(__name__)
class _Entry:
__slots__ = ("key", "slot", "refs", "ready")
def __init__(self, key: bytes, slot: int) -> None:
self.key = key
self.slot = slot
self.refs = 0 # in-flight transfers (store or load) touching this slot
self.ready = False # host copy complete and hit-able
@dataclass
class _Chain:
tokens: np.ndarray
hashes: list[bytes] = field(default_factory=list)
counted: bool = False # lookup stats recorded (a deferred request is looked up repeatedly)
l3_tried: bool = False # L3 consulted once for this request
@dataclass
class Stats:
lookups: int = 0
hits: int = 0
hit_tokens: int = 0
lookup_tokens: int = 0
deferred: int = 0 # hit found but GPU allocation not possible yet
stores: int = 0
store_skipped: int = 0
pages_stored: int = 0
pages_loaded: int = 0
states_stored: int = 0
states_loaded: int = 0
evicted_pages: int = 0
evicted_states: int = 0
errors: int = 0
l3_lookups: int = 0
l3_hits: int = 0
l3_pages_loaded: int = 0
l3_states_loaded: int = 0
l3_get_failed: int = 0
l3_get_seconds: float = 0.0
l3_errors: int = 0
def as_dict(self) -> dict[str, int]:
return dict(self.__dict__)
class HostPrefixCache:
def __init__(self, engine: LLMEngine, *, size_bytes: int, num_state_slots: int, staging_slots: int = 16) -> None:
assert envs.DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE == "compact", "host prefix cache needs decoder bounded replay (compact)"
assert envs.DEEPINFER_PREFILL_STORE_EXTRA_STATE, "host prefix cache needs DEEPINFER_PREFILL_STORE_EXTRA_STATE=1"
assert GPUCacheInterface.cache_info_ready, "cache engine not initialized"
self.engine = engine
self.scheduler = engine.scheduler
cache_engine = engine.driver_worker.cache_engine
self.gpu_blocks: torch.Tensor = cache_engine.gpu_blocks
self.gpu_states: torch.Tensor = cache_engine.gpu_states
assert self.gpu_blocks.dim() == 2 and self.gpu_blocks.is_contiguous() and self.gpu_blocks.dtype == torch.uint8
self.device = self.gpu_blocks.device
self.block_size: int = engine.cache_config.block_size
assert self.block_size == KVSTORE_GRAIN, (self.block_size, KVSTORE_GRAIN)
self.page_bytes: int = self.gpu_blocks.size(1)
self.state_bytes: int = GPUCacheInterface.state_size_to_disk
assert self.page_bytes % 16 == 0 and self.state_bytes % 16 == 0, (self.page_bytes, self.state_bytes)
self.keep_tokens: int = envs.DEEPINFER_YOCO_PREFILL_KEEP_TOKENS
# The state kernels gather [b - window_ring, b); keep b at least two blocks in so that range never goes negative.
self.min_hit_tokens = 2 * self.block_size
# A follow-up prompt can only resume at a position <= its length - keep_tokens, so a state stored only
# at the end of a prompt is unusable by any prompt that extends it by fewer than keep_tokens (+ up to a
# block). Store the state at the last few 128-boundaries instead; how far back is bounded by the SWA
# ring: the window ending at b must still be in the ring, b - ring_inference >= L - ring_kvstore.
self.state_lookback_blocks = int(os.environ.get("DEEPINFER_HICACHE_STATE_LOOKBACK", "4"))
self.ring_slack = GPUCacheInterface.swa_block_size_kvstore - GPUCacheInterface.swa_block_size_inference
self.num_pages = max(0, size_bytes // self.page_bytes)
self.num_state_slots = num_state_slots
self.staging_slots = staging_slots
assert self.num_pages > 0 and num_state_slots > 0
self.gpu_blocks_ptr = self.gpu_blocks.data_ptr()
# Host pools, staging ring and the copy kernel are allocated lazily at the first store (nothing can be
# cached before a prefill completes, and an unused cache then costs no pinned memory).
self._allocated = False
self.host_pages = self.host_states = None
self.host_pages_ptr = self.host_states_ptr = 0
self._staging: list[torch.Tensor] = []
self._staging_idx = 0
self.kernel = None
self.pages: OrderedDict[bytes, _Entry] = OrderedDict()
self.states: OrderedDict[bytes, _Entry] = OrderedDict()
self.free_pages: list[int] = list(range(self.num_pages - 1, -1, -1))
self.free_states: list[int] = list(range(num_state_slots - 1, -1, -1))
self._inflight: deque[tuple[torch.cuda.Event, list[_Entry]]] = deque()
self._chains: dict[str, _Chain] = {}
self.stats = Stats()
self.enabled = True
self.l3 = self._make_l3()
self.l3_group = self.l3.group_pages if self.l3 is not None else 0
self._l3_uploaded: OrderedDict[str, None] = OrderedDict() # recently put/seen L3 keys (skip re-puts)
# DEEPINFER_HICACHE_DEBUG_SYNC=1: synchronize after every launch and log it (pinpoints a faulting kernel).
self._debug_sync = os.environ.get("DEEPINFER_HICACHE_DEBUG_SYNC", "0") == "1"
logger.info(
"HostPrefixCache: %d pages x %d B (%.1f GiB) + %d state slots x %d B (%.2f GiB), allocated at first store; keep_tokens=%d",
self.num_pages,
self.page_bytes,
self.num_pages * self.page_bytes / 2**30,
num_state_slots,
self.state_bytes,
num_state_slots * self.state_bytes / 2**30,
self.keep_tokens,
)
def _make_l3(self): # noqa: ANN202
from deepinfer.hicache import mooncake_store
config = mooncake_store.config_from_env()
if config is None:
return None
model_cfg = os.path.join(str(self.engine.model_config.model), "config.json")
with open(model_cfg, "rb") as f:
cfg_hash = hashlib.sha256(f.read()).hexdigest()
fp = mooncake_store.layout_fingerprint(
dict(
config=cfg_hash,
page_bytes=self.page_bytes,
state_bytes=self.state_bytes,
block_size=self.block_size,
extra_state=int(envs.DEEPINFER_PREFILL_STORE_EXTRA_STATE),
extra_size=int(envs.DEEPINFER_PREFILL_STORE_EXTRA_SIZE),
backtracking=int(envs.DEEPINFER_MAX_BACKTRACKING_TOKENS),
swa_kvstore=int(GPUCacheInterface.swa_block_size_kvstore),
)
)
group = int(os.environ.get("DEEPINFER_HICACHE_L3_GROUP_PAGES", "4"))
return mooncake_store.MooncakeL3(config, fp, group)
def _allocate(self) -> None:
if self._allocated:
return
from deepinfer.hicache.kernels import copy_pages_module
t0 = time.perf_counter()
self.kernel = copy_pages_module()
self.host_pages = torch.empty((self.num_pages, self.page_bytes), dtype=torch.uint8, pin_memory=True)
self.host_states = torch.empty((self.num_state_slots, self.state_bytes), dtype=torch.uint8, pin_memory=True)
self.host_pages_ptr = self.host_pages.data_ptr()
self.host_states_ptr = self.host_states.data_ptr()
# Staging ring on the GPU for state slices; reuse is safe because every transfer runs on the same stream.
self._staging = [torch.empty(self.state_bytes, dtype=torch.uint8, device=self.device) for _ in range(self.staging_slots)]
if self.l3 is not None:
self.l3.register(self.host_pages_ptr, self.host_pages.numel())
self.l3.register(self.host_states_ptr, self.host_states.numel())
self._allocated = True
logger.info("HostPrefixCache: host pools pinned in %.1fs", time.perf_counter() - t0)
# ------------------------------------------------------------------ hashing
def _chain(self, request_id: str, tokens: np.ndarray | None) -> _Chain | None:
chain = self._chains.get(request_id)
if chain is None and tokens is not None:
chain = self._chains[request_id] = _Chain(np.ascontiguousarray(tokens, dtype=np.int32))
return chain
def _extend(self, chain: _Chain, num_pages: int) -> None:
hashes, toks, bs = chain.hashes, chain.tokens, self.block_size
prev = hashes[-1] if hashes else b""
for i in range(len(hashes), num_pages):
prev = hashlib.blake2b(prev + toks[i * bs : (i + 1) * bs].tobytes(), digest_size=16).digest()
hashes.append(prev)
def forget(self, request_id: str) -> None:
self._chains.pop(request_id, None)
# ------------------------------------------------------------------ transfers
def _ids(self, ids: list[int]) -> torch.Tensor:
return torch.from_numpy(np.asarray(ids, dtype=np.int32)).pin_memory().to(self.device, non_blocking=True)
def _next_staging(self) -> torch.Tensor:
t = self._staging[self._staging_idx]
self._staging_idx = (self._staging_idx + 1) % len(self._staging)
return t
def _track(self, entries: list[_Entry], uploads: list | None = None) -> None:
ev = torch.cuda.Event()
ev.record()
self._inflight.append((ev, entries, uploads or []))
def _checkpoint(self, what: str) -> None:
if self._debug_sync:
from deepinfer.utils import allow_cuda_sync
with allow_cuda_sync():
torch.cuda.synchronize()
logger.info("hicache debug: ok after %s", what)
def poll(self) -> None:
"""Retire completed transfers (events on one stream complete in order, so only the head is checked)."""
while self._inflight and self._inflight[0][0].query():
_, entries, uploads = self._inflight.popleft()
for e in entries:
e.refs -= 1
e.ready = True
for up in uploads:
self._l3_submit(*up)
if self.l3 is not None:
for held in self.l3.drain_done():
for e in held:
e.refs -= 1
def _evict(self, pool: OrderedDict[bytes, _Entry], free: list[int], need: int) -> int:
"""Evict `need` LRU entries that are idle; evicts nothing (returns 0) if that many cannot be found."""
victims, busy = [], []
for key, e in pool.items():
if len(victims) >= need:
break
if e.refs or not e.ready:
busy.append(key)
continue
victims.append(key)
for key in busy:
pool.move_to_end(key)
if len(victims) < need:
return 0
for key in victims:
free.append(pool.pop(key).slot)
return len(victims)
def _reserve(self, pool: OrderedDict[bytes, _Entry], free: list[int], need: int) -> bool:
if need > len(free) + len(pool):
return False
if len(free) < need:
evicted = self._evict(pool, free, need - len(free))
if pool is self.pages:
self.stats.evicted_pages += evicted
else:
self.stats.evicted_states += evicted
return len(free) >= need
# ------------------------------------------------------------------ load (hit)
def try_resume(self, seq: Sequence, request_id: str, prompt_ids: np.ndarray) -> int:
"""Fill `seq` (not yet in the scheduler) from the host cache.
Returns the number of cached tokens (0: miss, sequence untouched; -1: hit but GPU blocks/state not
allocatable right now, sequence untouched, retry later).
"""
if not self.enabled:
return 0
try:
return self._try_resume(seq, request_id, prompt_ids)
except Exception: # noqa: BLE001 - never take the engine down for a cache problem
logger.exception("host prefix cache: resume failed, request %s continues uncached", request_id)
self._fail()
return 0
def _try_resume(self, seq: Sequence, request_id: str, prompt_ids: np.ndarray) -> int:
bs = self.block_size
n = len(prompt_ids)
chain = self._chain(request_id, prompt_ids)
if not chain.counted:
chain.counted = True
self.stats.lookups += 1
self.stats.lookup_tokens += n
limit = ((n - self.keep_tokens) // bs) * bs # the replay rectangle [n - keep, n) must be recomputed
if limit < self.min_hit_tokens:
return 0
max_pages = limit // bs
self._extend(chain, max_pages)
if self.l3 is not None and not chain.l3_tried and self.l3.readable():
chain.l3_tried = True
try:
self._l3_fill(chain, max_pages, limit)
except Exception: # noqa: BLE001 - fall back to the host-only lookup
logger.exception("host prefix cache: L3 fetch failed for %s", request_id)
self.stats.l3_errors += 1
self.l3.breaker.trip("L3 fetch raised")
matched = 0
for h in chain.hashes[:max_pages]:
e = self.pages.get(h)
if e is None or not e.ready:
break
matched += 1
b = min(matched * bs, limit)
while b >= self.min_hit_tokens:
st = self.states.get(chain.hashes[b // bs - 1])
if st is not None and st.ready:
break
b -= bs
else:
return 0
try:
allocated = self.scheduler.try_allocate_seq(seq)
except RuntimeError: # can never fit (prompt longer than the whole KV pool): let the scheduler reject it
return 0
if not allocated:
self.stats.deferred += 1
return -1
num_pages = b // bs
page_entries = [self.pages[chain.hashes[i]] for i in range(num_pages)]
for e in page_entries:
e.refs += 1
self.pages.move_to_end(e.key)
st.refs += 1
self.states.move_to_end(st.key)
block_table = seq.block_table
self.kernel.copy_pages(
self.host_pages_ptr,
self.gpu_blocks_ptr,
self._ids([e.slot for e in page_entries]),
self._ids(list(block_table[:num_pages])),
self.page_bytes,
)
self._checkpoint(f"load: page copy H2D ({num_pages} pages, b={b})")
staging = self._next_staging()
self.kernel.copy_pages(self.host_states_ptr, staging.data_ptr(), self._ids([st.slot]), self._ids([0]), self.state_bytes)
self._checkpoint("load: state copy H2D")
GPUCacheInterface.load_state_from_store_to_prefill_cache(seq.state_id, b, staging, self.gpu_states)
self._checkpoint("load: load_state_to_cache")
self._track(page_entries + [st])
seq.num_cached_tokens = b
seq.cached_hit_tokens = b
self.stats.hits += 1
self.stats.hit_tokens += b
self.stats.pages_loaded += num_pages
self.stats.states_loaded += 1
return b
# ------------------------------------------------------------------ store (write-through at prefill end)
def on_prefill_step(self, scheduled: list[Sequence], finished: list[Sequence]) -> None:
"""Engine hook (pe_do_save_kvcache): called after a prefill step's GPU work has completed."""
if not self.enabled:
return
try:
for seq in scheduled:
if seq.chunked_rest_len == 0: # last chunk of this prompt
self._store(seq)
except Exception: # noqa: BLE001
logger.exception("host prefix cache: store failed")
self._fail()
def _store(self, seq: Sequence) -> None:
chain = self._chains.get(seq.request_id)
if chain is None: # request did not go through try_resume: hash the prompt from the sequence itself
chain = self._chain(seq.request_id, np.asarray(seq.data.get_token_ids(), dtype=np.int32))
self._allocate()
bs = self.block_size
b = (seq.get_len() // bs) * bs
if b < self.min_hit_tokens:
return
num_pages = b // bs
self._extend(chain, num_pages)
missing = [i for i in range(num_pages) if chain.hashes[i] not in self.pages]
for i in range(num_pages):
if i not in missing:
self.pages.move_to_end(chain.hashes[i])
# State boundaries: b, b - bs, ... while the window ending there is still in the ring.
seq_len = seq.get_len()
new_states = []
for k in range(self.state_lookback_blocks):
pos = b - k * bs
if pos < self.min_hit_tokens or pos < seq_len - self.ring_slack:
break
key = chain.hashes[pos // bs - 1]
if key in self.states:
self.states.move_to_end(key)
else:
new_states.append((pos, key))
uploads = self._l3_plan(chain, num_pages, [key for _, key in new_states])
if not missing and not new_states:
for up in uploads: # everything is already on the host: upload right away
self._l3_submit(*up)
return
if not self._reserve(self.pages, self.free_pages, len(missing)):
self.stats.store_skipped += 1
return
# States are best effort: keep the newest boundaries that fit.
while new_states and not self._reserve(self.states, self.free_states, len(new_states)):
new_states.pop()
if not missing and not new_states:
self.stats.store_skipped += 1
return
block_table = seq.block_table
entries: list[_Entry] = []
if missing:
new = []
for i in missing:
e = _Entry(chain.hashes[i], self.free_pages.pop())
e.refs = 1
self.pages[e.key] = e
new.append(e)
self._checkpoint(f"store: before page copy ({len(new)} pages, blocks {block_table[:3]}..., b={b}, state_id={seq.state_id})")
self.kernel.copy_pages(
self.gpu_blocks_ptr,
self.host_pages_ptr,
self._ids([block_table[i] for i in missing]),
self._ids([e.slot for e in new]),
self.page_bytes,
)
self._checkpoint("store: page copy D2H")
entries += new
self.stats.pages_stored += len(new)
for pos, key in new_states:
st = _Entry(key, self.free_states.pop())
st.refs = 1
self.states[key] = st
staging = self._next_staging()
GPUCacheInterface.select_state_slice_to_store(seq.state_id, pos, self.gpu_states, state_tensor=staging)
self._checkpoint(f"store: select_state_slice at {pos}")
self.kernel.copy_pages(staging.data_ptr(), self.host_states_ptr, self._ids([0]), self._ids([st.slot]), self.state_bytes)
self._checkpoint("store: state copy D2H")
entries.append(st)
self.stats.states_stored += 1
self._track(entries, uploads)
self.stats.stores += 1
# ------------------------------------------------------------------ L3 (Mooncake Store)
def _l3_seen(self, key: str) -> bool:
if key in self._l3_uploaded:
self._l3_uploaded.move_to_end(key)
return True
return False
def _l3_mark(self, key: str) -> None:
self._l3_uploaded[key] = None
while len(self._l3_uploaded) > 1_000_000:
self._l3_uploaded.popitem(last=False)
def _l3_plan(self, chain: _Chain, num_pages: int, new_state_keys: list[bytes]) -> list:
"""Page groups of this prompt and its new states the Store has not seen yet: [(page hashes, groups, state keys)]."""
if self.l3 is None:
return []
g = self.l3_group
groups = [i for i in range(num_pages // g) if not self._l3_seen(self.l3.group_key(chain.hashes[(i + 1) * g - 1]))]
state_keys = [k for k in new_state_keys if not self._l3_seen(self.l3.state_key(k))]
if not groups and not state_keys:
return []
return [(chain.hashes[:num_pages], groups, state_keys)]
def _l3_submit(self, hashes: list[bytes], groups: list[int], state_keys: list[bytes]) -> None:
"""Pin the host slots of the planned groups/states and queue a write-behind put. Entries that are gone (evicted)
or not ready are skipped."""
g = self.l3_group
keys, ptrs, sizes, held = [], [], [], []
for i in groups:
es = [self.pages.get(h) for h in hashes[i * g : (i + 1) * g]]
if any(e is None or not e.ready for e in es):
continue
keys.append(self.l3.group_key(hashes[(i + 1) * g - 1]))
ptrs.append([self.host_pages_ptr + e.slot * self.page_bytes for e in es])
sizes.append([self.page_bytes] * g)
held += es
for k in state_keys:
st = self.states.get(k)
if st is None or not st.ready:
continue
keys.append(self.l3.state_key(k))
ptrs.append([self.host_states_ptr + st.slot * self.state_bytes])
sizes.append([self.state_bytes])
held.append(st)
if not keys:
return
for e in held:
e.refs += 1
for k in keys:
self._l3_mark(k)
self.l3.submit_put(keys, ptrs, sizes, held)
def _l3_fill(self, chain: _Chain, max_pages: int, limit: int) -> None:
"""Extend the host cache with page groups and a state from the Store, so the host lookup that follows hits
deeper. No-op when the Store adds nothing beyond what the host already has."""
bs, g, l3 = self.block_size, self.l3_group, self.l3
self._allocate() # pools are lazy; a fresh instance can hit L3 before it ever stored anything
host_pages = 0
for h in chain.hashes[:max_pages]:
e = self.pages.get(h)
if e is None or not e.ready:
break
host_pages += 1
first_group = host_pages // g
n_groups = max_pages // g - first_group
if n_groups <= 0:
return
self.stats.l3_lookups += 1
group_keys = [l3.group_key(chain.hashes[(first_group + i + 1) * g - 1]) for i in range(n_groups)]
present = l3.exists(group_keys)
k = 0
while k < n_groups and present[k]:
k += 1
if k == 0:
return
l3_pages = (first_group + k) * g # pages available in host + L3
# deepest boundary <= min(l3_pages * bs, limit) beyond the host-only reach that has a state (host or L3)
top = min(l3_pages * bs, limit) // bs
candidates = [p for p in range(top, host_pages, -1) if p * bs >= self.min_hit_tokens]
if not candidates:
return
host_state = {p: (self.states.get(chain.hashes[p - 1]) is not None) for p in candidates}
need_l3 = [p for p in candidates if not host_state[p]]
in_l3 = dict(zip(need_l3, l3.exists([l3.state_key(chain.hashes[p - 1]) for p in need_l3]), strict=True))
pick = next((p for p in candidates if host_state[p] or in_l3.get(p)), None)
if pick is None:
return
state_from_l3 = not host_state[pick]
fetch_groups = list(range(first_group, (pick + g - 1) // g)) # groups covering pages [first_group*g, pick)
# Pin the host prefix so reserving slots below cannot evict it, then reserve.
prefix = [self.pages[h] for h in chain.hashes[:host_pages]]
for e in prefix:
e.refs += 1
try:
need_pages = len(fetch_groups) * g
if not self._reserve(self.pages, self.free_pages, need_pages):
return
if state_from_l3 and not self._reserve(self.states, self.free_states, 1):
return
slots = [self.free_pages.pop() for _ in range(need_pages)]
keys = [l3.group_key(chain.hashes[(i + 1) * g - 1]) for i in fetch_groups]
ptrs = [[self.host_pages_ptr + s * self.page_bytes for s in slots[j * g : (j + 1) * g]] for j in range(len(fetch_groups))]
sizes = [[self.page_bytes] * g for _ in fetch_groups]
st_slot = None
if state_from_l3:
st_slot = self.free_states.pop()
keys.append(l3.state_key(chain.hashes[pick - 1]))
ptrs.append([self.host_states_ptr + st_slot * self.state_bytes])
sizes.append([self.state_bytes])
t0 = time.perf_counter()
ok = l3.get(keys, ptrs, sizes)
self.stats.l3_get_seconds += time.perf_counter() - t0
if not all(ok): # evicted from the Store in between, or a transfer error: drop the whole fetch
self.stats.l3_get_failed += 1
self.free_pages.extend(slots)
if st_slot is not None:
self.free_states.append(st_slot)
return
loaded = 0
for j, gi in enumerate(fetch_groups):
for t in range(g):
page = gi * g + t
slot = slots[j * g + t]
h = chain.hashes[page]
if page >= pick or h in self.pages: # beyond the hit, or the host already holds this page
self.free_pages.append(slot)
continue
e = _Entry(h, slot)
e.ready = True
self.pages[h] = e
loaded += 1
if st_slot is not None:
st = _Entry(chain.hashes[pick - 1], st_slot)
st.ready = True
self.states[st.key] = st
self.stats.l3_states_loaded += 1
self.stats.l3_hits += 1
self.stats.l3_pages_loaded += loaded
finally:
for e in prefix:
e.refs -= 1
# ------------------------------------------------------------------ misc
def _fail(self) -> None:
self.stats.errors += 1
if self.stats.errors >= 3:
logger.error("host prefix cache: disabled after %d errors", self.stats.errors)
self.enabled = False
def stats_dict(self) -> dict[str, int | float]:
d = self.stats.as_dict()
d.update(
enabled=int(self.enabled),
pages_used=len(self.pages),
pages_total=self.num_pages,
states_used=len(self.states),
states_total=len(self.free_states) + len(self.states),
inflight=len(self._inflight),
hit_rate_tokens=(self.stats.hit_tokens / self.stats.lookup_tokens) if self.stats.lookup_tokens else 0.0,
)
if self.l3 is not None:
d.update(self.l3.stats_dict())
return d
Reference in New Issue View Git Blame Copy Permalink