From 3657e00592b3a74d73011a458cd14e9e62ad50c3 Mon Sep 17 00:00:00 2001 From: whn09 Date: Fri, 23 Oct 2026 21:56:26 +0800 Subject: [PATCH] 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 --- deepinfer/hicache/host_cache.py | 3 +- deepinfer/hicache/mooncake_store.py | 76 +++++++++++++++++++++++++++-- tests/hicache/test_host_cache_l3.py | 53 +++++++++++++++++++- 3 files changed, 124 insertions(+), 8 deletions(-) diff --git a/deepinfer/hicache/host_cache.py b/deepinfer/hicache/host_cache.py index b395b83..dd80146 100644 --- a/deepinfer/hicache/host_cache.py +++ b/deepinfer/hicache/host_cache.py @@ -322,13 +322,14 @@ class HostPrefixCache: return 0 max_pages = limit // bs self._extend(chain, max_pages) - if self.l3 is not None and not chain.l3_tried: + 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) diff --git a/deepinfer/hicache/mooncake_store.py b/deepinfer/hicache/mooncake_store.py index 00e0575..5fb899d 100644 --- a/deepinfer/hicache/mooncake_store.py +++ b/deepinfer/hicache/mooncake_store.py @@ -18,6 +18,17 @@ Objects Data moves between the Store and the (registered) pinned host pools with batch_put_from_multi_buffers / batch_get_into_multi_buffers, so no extra copies. Puts run on a background thread (write-behind); the engine thread learns about completions through drain_done(). Gets are synchronous (called from try_resume). + +Failure handling. A Store transfer that cannot complete blocks for Mooncake's fixed 60 s transfer timeout, so one bad +get stalls the engine thread that long. A circuit breaker therefore opens on the first transfer failure (any put or +get error other than "not found" / "already exists") or on a get slower than DEEPINFER_HICACHE_L3_SLOW_S (default 5 s). +While it is open, reads are skipped (host-only lookup) and queued puts are dropped. It closes after +DEEPINFER_HICACHE_L3_COOLDOWN_S (default 30 s, doubled on each consecutive trip up to 10 min). + +TCP transport limits. Mooncake <= 0.3.13 queues at most 1024 transfers per peer, and one long prompt stores more +buffers than that in a single batch put; the rejected transfers hang until the 60 s timeout. Unless they are already +set, MC_TCP_MAX_QUEUED_TRANSFERS_PER_PEER and MC_TCP_MAX_PENDING_ADMISSIONS_PER_PEER default to 65535, which is the +default since kvcache-ai/Mooncake#3710. """ from __future__ import annotations @@ -37,6 +48,8 @@ logger = init_logger(__name__) ENV = "DEEPINFER_HICACHE_MOONCAKE_STORE" FORMAT_VERSION = 1 +OBJECT_NOT_FOUND, OBJECT_ALREADY_EXISTS = -704, -705 +TCP_ENV_DEFAULTS = {"MC_TCP_MAX_QUEUED_TRANSFERS_PER_PEER": "65535", "MC_TCP_MAX_PENDING_ADMISSIONS_PER_PEER": "65535"} def config_from_env() -> dict | None: @@ -55,8 +68,36 @@ def layout_fingerprint(parts: dict[str, Any]) -> str: return f"v{FORMAT_VERSION}-{hashlib.sha256(blob).hexdigest()[:16]}" +class _Breaker: + """Opens on a transfer failure; closes after a cooldown that doubles on consecutive trips.""" + + def __init__(self, cooldown_s: float, max_cooldown_s: float = 600.0) -> None: + self.base, self.max = cooldown_s, max_cooldown_s + self.cooldown = cooldown_s + self.open_until = 0.0 + self.trips = 0 + self._lock = threading.Lock() + + def allow(self) -> bool: + return time.monotonic() >= self.open_until + + def trip(self, why: str) -> None: + with self._lock: + now = time.monotonic() + if now < self.open_until: # already open (e.g. the put thread and a get failed together) + return + if now - self.open_until > self.cooldown: # healthy for a while since the last trip: reset the backoff + self.cooldown = self.base + self.open_until = now + self.cooldown + self.trips += 1 + logger.warning("HostPrefixCache L3: %s; Store disabled for %.0f s", why, self.cooldown) + self.cooldown = min(self.cooldown * 2, self.max) + + class MooncakeL3: def __init__(self, config: dict, fingerprint: str, group_pages: int) -> None: + for k, v in TCP_ENV_DEFAULTS.items(): # read by Mooncake's TCP transport when setup() creates it + os.environ.setdefault(k, v) from mooncake.store import MooncakeDistributedStore # optional dependency (pip install mooncake-transfer-engine) cfg = dict(config) @@ -71,8 +112,10 @@ class MooncakeL3: self._state_prefix = f"state/{fingerprint}/" self._puts: queue.Queue = queue.Queue() self._done: collections.deque = collections.deque() # tokens of finished puts (appended by the worker thread) - self.put_objects = self.put_failed = self.put_bytes = 0 + self.put_objects = self.put_failed = self.put_bytes = self.put_dropped = 0 self.put_seconds = 0.0 + self.slow_get_s = float(os.environ.get("DEEPINFER_HICACHE_L3_SLOW_S", "5")) + self.breaker = _Breaker(float(os.environ.get("DEEPINFER_HICACHE_L3_COOLDOWN_S", "30"))) threading.Thread(target=self._put_worker, name="hicache-l3-put", daemon=True).start() logger.info("HostPrefixCache L3: Mooncake Store ready (fingerprint %s, %d pages per object)", fingerprint, group_pages) @@ -102,22 +145,33 @@ class MooncakeL3: def _put_worker(self) -> None: while True: keys, ptrs, sizes, token = self._puts.get() + if not self.breaker.allow(): # Store unhealthy: drop the put, release the host entries right away + self.put_dropped += len(keys) + self._done.append(token) + continue t0 = time.perf_counter() try: rcs = self.store.batch_put_from_multi_buffers(keys, ptrs, sizes) # an object that already exists (another rank or instance stored it first) is not a failure - failed = sum(1 for rc in rcs if rc not in (0, -705)) - self.put_objects += len(keys) - failed - self.put_failed += failed + bad = [rc for rc in rcs if rc not in (0, OBJECT_ALREADY_EXISTS)] + self.put_objects += len(keys) - len(bad) + self.put_failed += len(bad) self.put_bytes += sum(sum(s) for s, rc in zip(sizes, rcs, strict=True) if rc == 0) + if bad: + self.breaker.trip(f"batch put failed for {len(bad)}/{len(keys)} objects (rc {bad[0]})") except Exception: # noqa: BLE001 - a Store problem must not take the engine down logger.exception("HostPrefixCache L3: batch put failed") self.put_failed += len(keys) + self.breaker.trip("batch put raised") self.put_seconds += time.perf_counter() - t0 self._done.append(token) # ---------------------------------------------------------------- read (engine thread) + def readable(self) -> bool: + """False while the breaker is open: skip the Store and use the host cache only.""" + return self.breaker.allow() + def exists(self, keys: list[str]) -> list[bool]: if not keys: return [] @@ -127,8 +181,17 @@ class MooncakeL3: """Read each object into its list of buffers; True where the whole object arrived.""" if not keys: return [] + t0 = time.perf_counter() rcs = self.store.batch_get_into_multi_buffers(keys, ptrs, sizes) - return [rc == sum(s) for rc, s in zip(rcs, sizes, strict=True)] + dt = time.perf_counter() - t0 + ok = [rc == sum(s) for rc, s in zip(rcs, sizes, strict=True)] + # "not found" just means evicted since exists(); anything else is a transfer problem + bad = [rc for rc, good in zip(rcs, ok, strict=True) if not good and rc != OBJECT_NOT_FOUND] + if bad: + self.breaker.trip(f"batch get failed for {len(bad)}/{len(keys)} objects (rc {bad[0]}, {dt:.1f} s)") + elif dt > self.slow_get_s: + self.breaker.trip(f"batch get of {len(keys)} objects took {dt:.1f} s") + return ok def stats_dict(self) -> dict[str, int | float]: return dict( @@ -137,4 +200,7 @@ class MooncakeL3: l3_put_gib=round(self.put_bytes / 2**30, 3), l3_put_seconds=round(self.put_seconds, 3), l3_put_queue=self._puts.qsize(), + l3_put_dropped=self.put_dropped, + l3_breaker_trips=self.breaker.trips, + l3_breaker_open=int(not self.breaker.allow()), ) diff --git a/tests/hicache/test_host_cache_l3.py b/tests/hicache/test_host_cache_l3.py index 93d620b..6e1ce07 100644 --- a/tests/hicache/test_host_cache_l3.py +++ b/tests/hicache/test_host_cache_l3.py @@ -1,15 +1,18 @@ """HostPrefixCache L3 (Mooncake Store) logic against a dict-backed fake store: write-behind of page groups and -states, a second cache (another rank / instance) extending its empty host pool from L3, no re-puts, and a fetch that -loses an object between exists() and get(). Same fake engine as test_host_cache.py. Needs one GPU.""" +states, a second cache (another rank / instance) extending its empty host pool from L3, no re-puts, a fetch that +loses an object between exists() and get(), and the circuit breaker after a transfer failure. Same fake engine as +test_host_cache.py. Needs one GPU.""" from __future__ import annotations import ctypes +import time import numpy as np import pytest import torch +from deepinfer.hicache.mooncake_store import _Breaker from tests.hicache.test_host_cache import BLOCK, PAGE, STATE, FakeSeq, _prefill, _tokens, engine # noqa: F401 pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") @@ -25,8 +28,14 @@ class FakeL3: self.objects: dict[str, bytes] = {} self.puts = 0 self.drop_on_get: set[str] = set() # keys that vanish between exists() and get() + self.fail_gets = False # every get is a transfer failure (trips the breaker, like MooncakeL3.get) + self.gets = 0 + self.breaker = _Breaker(30.0) self._done: list[object] = [] + def readable(self) -> bool: + return self.breaker.allow() + def register(self, ptr: int, nbytes: int) -> None: pass @@ -51,6 +60,10 @@ class FakeL3: return [k in self.objects for k in keys] def get(self, keys, ptrs, sizes) -> list[bool]: # noqa: ANN001 + self.gets += 1 + if self.fail_gets: + self.breaker.trip("test transfer failure") + return [False] * len(keys) ok = [] for k, ps, ss in zip(keys, ptrs, sizes, strict=True): data = None if k in self.drop_on_get else self.objects.get(k) @@ -142,3 +155,39 @@ def test_object_lost_between_exists_and_get_falls_back(engine, l3): # noqa: F81 assert b_cache.try_resume(b, "b", b.tokens) == 0 # miss, sequence untouched assert b.block_table == [] and b_cache.stats.l3_get_failed == 1 assert len(b_cache.pages) == 0 and len(b_cache.free_pages) == b_cache.num_pages + + +def test_transfer_failure_opens_the_breaker(engine, l3): # noqa: F811 + from deepinfer.hicache import HostPrefixCache + + a_cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=8) + a = FakeSeq("a", _tokens(1000, 1)) + _prefill(engine, a_cache, a) + _settle(a_cache) + + l3.fail_gets = True + b_cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=8) + b = FakeSeq("b", np.concatenate([a.tokens, _tokens(500, 2)])) + assert b_cache.try_resume(b, "b", b.tokens) == 0 and l3.gets == 1 + assert not l3.readable() and l3.breaker.trips == 1 + # while open, later requests do not touch the Store at all (no 60 s stalls), even once it is healthy again + l3.fail_gets = False + c = FakeSeq("c", np.concatenate([a.tokens, _tokens(500, 3)])) + assert b_cache.try_resume(c, "c", c.tokens) == 0 and l3.gets == 1 + # after the cooldown the Store is used again + l3.breaker.open_until = 0.0 + d = FakeSeq("d", np.concatenate([a.tokens, _tokens(500, 4)])) + assert b_cache.try_resume(d, "d", d.tokens) == 4 * BLOCK and l3.gets == 2 + + +def test_breaker_backoff() -> None: + br = _Breaker(10.0, max_cooldown_s=25.0) + assert br.allow() + br.trip("x") + assert not br.allow() and br.trips == 1 + br.trip("y") # already open: no second trip + assert br.trips == 1 + for want in (20.0, 25.0): # consecutive trips right after reopening double the cooldown, capped + br.open_until = time.monotonic() - 1e-3 + br.trip("z") + assert abs(br.open_until - time.monotonic() - want) < 1.0