minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[PP + HiCache] Add PP Prefetch Tickets for eager cross-stage storage prefetch (#36700)

Browse Source 
Co-authored-by: Chao Shi <stepinto@live.com>
This commit is contained in:
huangtingwei
2026-09-20 11:19:31 +08:00
committed by GitHub
co-authored by Chao Shi
parent 113f6f080e
commit 020703923d
10 changed files with 1153 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/managers/test_priority_scheduling_disaggregation.py
+61
View File
@@ -20,6 +20,7 @@ from sglang.srt.managers.schedule_batch import ( # noqa: E402
ReqKvInfo,
)
from sglang.srt.managers.scheduler import Scheduler # noqa: E402
from sglang.srt.mem_cache.base_prefix_cache import CacheRequestOutcome
from sglang.srt.runtime_context import get_context, publish, reset_context # noqa: E402
from sglang.srt.server_args import ServerArgs
from sglang.test.ci.ci_register import register_cpu_ci
@@ -43,8 +44,13 @@ class TestDisaggregationPriorityQueueing(unittest.TestCase):
scheduler.enable_priority_scheduling = True
scheduler.schedule_low_priority_values_first = False
scheduler.abort_on_priority_when_disabled = False
scheduler.enable_hierarchical_cache = False
scheduler.enable_hicache_storage = False
scheduler.enable_unified_cache_external_linker = False
scheduler.processed_tokens_counter = 0
scheduler.waiting_queue = []
scheduler._prefetch_kvcache = MagicMock()
scheduler.tree_cache = MagicMock()
scheduler._abort_on_queued_limit = MagicMock(return_value=False)
scheduler.model_config = SimpleNamespace(num_key_value_heads=8)
scheduler.disagg_prefill_bootstrap_queue = MagicMock()
@@ -100,6 +106,61 @@ class TestDisaggregationPriorityQueueing(unittest.TestCase):
scheduler.ipc_channels.send_to_tokenizer.send_output.assert_called_once()
req.time_stats.trace_ctx.abort.assert_called_once()
def test_priority_rejection_releases_early_prefetch(self):
for mode in DisaggregationMode:
with self.subTest(mode=mode):
scheduler = self._new_scheduler(mode)
scheduler.enable_priority_scheduling = False
scheduler.abort_on_priority_when_disabled = True
scheduler.enable_hicache_storage = True
scheduler.tree_cache = MagicMock(spec=["finish"])
req = self._new_req(priority=10)
with patch(
"sglang.srt.managers.scheduler.get_serving",
return_value=SimpleNamespace(weight_version="v0"),
):
scheduler._add_request_to_queue(req)
scheduler.tree_cache.finish.assert_called_once_with(
req.cache_request_handle, CacheRequestOutcome.ABORT
)
scheduler._prefetch_kvcache.assert_not_called()
self.assertEqual(scheduler.waiting_queue, [])
scheduler.disagg_prefill_bootstrap_queue.add.assert_not_called()
scheduler.disagg_decode_prealloc_queue.add.assert_not_called()
def test_full_queue_releases_only_the_rejected_request(self):
for priority in (0, 2):
with self.subTest(incoming_priority=priority):
scheduler = self._new_scheduler(DisaggregationMode.NULL)
del scheduler._abort_on_queued_limit # Exercise the real queue check.
scheduler.max_queued_requests = 1
scheduler.enable_hicache_storage = True
scheduler.tree_cache = MagicMock(spec=["finish"])
scheduler.beam_coordinator = MagicMock()
queued = self._new_req(priority=1)
queued.rid = "queued"
scheduler.waiting_queue = [queued]
incoming = self._new_req(priority=priority)
with patch(
"sglang.srt.managers.scheduler.get_serving",
return_value=SimpleNamespace(weight_version="v0"),
):
scheduler._add_request_to_queue(incoming)
rejected = incoming if priority == 0 else queued
scheduler.tree_cache.finish.assert_called_once_with(
rejected.cache_request_handle, CacheRequestOutcome.ABORT
)
if priority == 0:
self.assertEqual(scheduler.waiting_queue, [queued])
scheduler._prefetch_kvcache.assert_not_called()
else:
self.assertEqual(scheduler.waiting_queue, [incoming])
scheduler._prefetch_kvcache.assert_called_once_with(incoming)
class TestDecodePreallocQueuePriority(unittest.TestCase):
def setUp(self):
test/registered/unit/mem_cache/test_hiradix_pp_sync_drain.py
+38
View File
@@ -1,6 +1,7 @@
"""Unit tests for HiCache PP synchronization."""
import unittest
from queue import Queue
from types import SimpleNamespace
from unittest.mock import MagicMock
@@ -59,6 +60,7 @@ class TestUnifiedPPSyncBatching(unittest.TestCase):
)
cache.pp_rank = pp_rank
cache.pp_size = 2
cache.host_memory_mode = "cache"
cache.enable_storage_metrics = False
cache.storage_metrics_collector = None
cache.buffer_pipeline = None
@@ -110,6 +112,42 @@ class TestUnifiedPPSyncBatching(unittest.TestCase):
follower.writing_check.assert_called_once_with(finish_count=1)
follower.loading_check.assert_called_once_with(finish_count=1)
def test_buffer_mode_follower_drains_rank_local_completions(self):
cache = self._make_cache(1, [True, False], [True, True])
cache.host_memory_mode = "buffer_only"
cache.tree_core.enable_storage = True
cache._all_reduce_attn_groups = MagicMock()
cache._drain_storage_control_queues_impl = MagicMock()
cc = cache.cache_controller
for size, name in enumerate(
(
"prefetch_hit_queue",
"ack_prefetch_queue",
"ack_backup_queue",
"host_mem_release_queue",
),
start=1,
):
queue = Queue()
for _ in range(size):
queue.put(object())
setattr(cc, name, queue)
cache.check_hicache_events()
cache._all_reduce.assert_not_called()
cache._all_reduce_attn_groups.assert_called_once()
cache.writing_check.assert_called_once_with(finish_count=1)
cache.loading_check.assert_called_once_with(finish_count=2)
cache._drain_storage_control_queues_impl.assert_called_once_with(
n_storage_hit=1,
n_ack_prefetch=2,
n_backup=3,
n_release=4,
extra_release_counts={},
log_metrics=True,
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_pp_prefetch_ticket.py
+414
View File
@@ -0,0 +1,414 @@
"""CPU regressions for PP ticket ordering, admission, and buffer ownership."""
import pickle
import threading
import unittest
from array import array
from queue import Empty, Queue
from unittest.mock import Mock, call, patch
import torch
from sglang.srt.managers.cache_controller import PrefetchAck
from sglang.srt.mem_cache.base_prefix_cache import (
CacheRequestHandle,
CacheRequestOutcome,
)
from sglang.srt.mem_cache.buffer_mode.pipeline import BufferModePipeline
from sglang.srt.mem_cache.hicache_storage import PoolName, PoolTransfer
from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import (
HybridCacheController,
PPPrefetchDecision,
)
from sglang.srt.mem_cache.radix_cache import RadixKey
from sglang.srt.mem_cache.storage_prefetch import StoragePrefetchRetries
from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache
from sglang.srt.mem_cache.utils import get_storage_hash_str
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=3, suite="base-a-test-cpu")
class TestPPPrefetchTicket(unittest.TestCase):
def setUp(self):
self.c = c = HybridCacheController.__new__(HybridCacheController)
c.page_size = c.prefetch_threshold = 4
c.pp_rank = c.tp_rank = 0
c.pp_size, c.tp_size = 4, 2
c.pp_group, c.pp_prefetch_command_group = "pp", "command"
c.pp_prefetch_command_thread = None
c.prefetch_hits_sync_groups = c.prefetch_completion_sync_groups = ["pp", "tp"]
c.pp_prefetch_states, c.pp_prefetch_decisions = {}, {}
c.pp_prefetch_state_lock = threading.Lock()
c.pp_prefetch_command_queue = Queue()
c.prefetch_queue = Queue()
c.prefetch_sync_queue = Queue()
c.ack_prefetch_queue = Queue()
c.host_mem_release_queue = Queue()
c.prefetch_buffer = Queue()
c.storage_stop_event = threading.Event()
c.prefetch_tokens_occupied = 0
c.mem_pool_host = Mock(page_size=4)
c.mem_pool_host.alloc.side_effect = lambda size, **_: torch.arange(size)
c._storage_hit_query = Mock(return_value=(["h0", "h1"], 8))
c._all_reduce = Mock()
ranks = {"pp": [0, 2, 4, 6], "tp": [0, 1], "command": [0, 2, 4, 6]}
patcher = patch.object(
torch.distributed, "get_process_group_ranks", side_effect=ranks.__getitem__
)
patcher.start()
self.addCleanup(patcher.stop)
patcher = patch.object(
torch.distributed, "get_rank", side_effect=lambda: c.pp_rank * 2 + c.tp_rank
)
patcher.start()
self.addCleanup(patcher.stop)
self.cache = cache = UnifiedRadixCache.__new__(UnifiedRadixCache)
cache.pp_rank = 1
cache.tree_core = Mock(enable_storage=True)
cache.cache_controller = c
cache._all_reduce = Mock()
cache.ongoing_prefetch = {}
cache.prefetch_loaded_tokens_by_reqid = {}
cache.prefetch_loaded_storage_start_by_reqid = {}
cache.storage_prefetch_retries = StoragePrefetchRetries()
cache.linker = None
cache.root_node_handle = Mock(return_value=0)
cache.buffer_pipeline = Mock(spec=BufferModePipeline)
cache._handle_prefetch_result = Mock()
def submit(self, rid="hit", pools=None, attempt_id=0, assume_stored=False):
key = RadixKey(
array("q", range(9)), "adapter", is_bigram=True, cache_salt="tenant"
)
return self.c.submit_prefetch(
CacheRequestHandle(rid, attempt_id),
key,
"ab" * 32,
["prefix"],
[10, 11, 12, 13],
pools,
assume_stored=assume_stored,
)
def run_commands(self, *tickets):
commands = iter((*tickets, None))
self.c.pp_prefetch_command_queue.put(None)
def broadcast(objects, *args, **kwargs):
return [next(commands)] if self.c.pp_rank else objects
with patch(f"{HybridCacheController.__module__}.broadcast_pyobj", broadcast):
self.c.pp_prefetch_command_thread_func()
def sync_acks(self, *acks):
for ack in acks:
self.c.prefetch_sync_queue.put(ack)
with patch.object(
self.c.storage_stop_event,
"is_set",
side_effect=[False] * len(acks) + [True],
):
self.c.prefetch_sync_thread_func()
def test_hit_and_miss_are_reused_across_enqueue_retract_and_retry(self):
c = self.c
for hit in (0, 8):
with self.subTest(hit=hit):
rid = str(hit)
handle = CacheRequestHandle(rid, 0)
c._storage_hit_query.return_value = ([], hit)
self.assertEqual(self.submit(rid).decision, bool(hit))
queries = c._storage_hit_query.call_count
for _ in range(2):
self.assertEqual(
self.cache.prefetch_from_storage(handle, 0, []), bool(hit)
)
self.assertEqual(c._storage_hit_query.call_count, queries)
if hit:
state = c.pp_prefetch_states[rid]
state.ready_event.set()
c.take_ready_pp_prefetch(rid)
self.assertFalse(
self.cache.prefetch_from_storage(
CacheRequestHandle(rid, 1), 0, []
)
)
self.assertTrue(self.cache.check_prefetch_progress(handle))
self.assertFalse(c.release_pp_prefetch(rid))
self.assertIsNone(c.get_prefetch_submission(rid))
self.assertEqual(c.pp_prefetch_command_queue.qsize(), 1)
self.cache._all_reduce.assert_not_called()
def test_submission_does_not_wait_for_worker_and_uses_only_tp_consensus(self):
c = self.c
with patch.object(
c.pp_prefetch_command_queue,
"join",
side_effect=AssertionError("scheduler blocked"),
):
self.assertTrue(self.submit(assume_stored=True).decision)
self.assertTrue(self.submit("next").decision)
self.assertEqual(c.pp_prefetch_command_queue.qsize(), 2)
self.assertFalse(c.is_pp_prefetch_ready("hit"))
c.mem_pool_host.alloc.assert_not_called()
self.assertFalse(c.pp_prefetch_states["hit"].operation.assume_stored)
self.assertEqual(
c._all_reduce.call_args.args[1:], (torch.distributed.ReduceOp.MIN, ["tp"])
)
self.assertEqual(
[call.kwargs["pp_rank"] for call in c._storage_hit_query.call_args_list],
[0, 2, 0, 2],
)
def test_tp_only_preserves_normal_prefetch(self):
c = self.c
c.pp_prefetch_command_group = None
result = self.submit(attempt_id=3, assume_stored=True)
self.assertIsNone(result.decision)
self.assertEqual(result.operation.handle, CacheRequestHandle("hit", 3))
self.assertTrue(result.operation.assume_stored)
self.assertIs(c.prefetch_queue.get_nowait(), result.operation)
c._storage_hit_query.assert_not_called()
def test_downstream_request_and_ticket_order_preserves_pp0_admission(self):
c, cache = self.c, self.cache
self.submit(attempt_id=3)
handle = CacheRequestHandle("hit", 3)
ticket = pickle.loads(pickle.dumps(c.pp_prefetch_states["hit"].ticket))
ticket.last_hash = None
c.pp_rank = 1
c.pp_prefetch_states.clear()
cache.bind_prefetch_ticket("hit")
self.assertFalse(cache.check_prefetch_progress(handle))
c._storage_hit_query.reset_mock()
self.run_commands(ticket)
operation = c.prefetch_buffer.get_nowait()
self.assertEqual(operation.handle, handle)
self.assertEqual(
operation.hash_value, get_storage_hash_str(ticket.prefetch_key, page_size=4)
)
self.assertTrue(operation.token_ids.is_bigram)
self.assertEqual(len(operation.token_ids), 8)
c._storage_hit_query.assert_not_called()
self.sync_acks(PrefetchAck("hit", operation, completed_tokens=8))
self.assertFalse(c.is_pp_prefetch_ready("hit"))
self.sync_acks(PrefetchAck("hit", operation, completed_req=True))
self.assertFalse(cache.check_prefetch_progress(handle)) # PP0 has not admitted.
cache._all_reduce.side_effect = lambda tensor, _: tensor.fill_(1)
self.assertTrue(cache.check_prefetch_progress(handle))
cache._handle_prefetch_result.assert_called_once_with(operation)
cache.buffer_pipeline.try_lock_anchor.assert_called_once_with(handle, 8)
key = cache.ongoing_prefetch[handle].prefetch_key
self.assertEqual((key.extra_key, key.cache_salt), ("adapter", "tenant"))
self.assertTrue(key.is_bigram)
self.assertFalse(c.is_pp_prefetch_ready("hit"))
def test_allocation_error_preserves_ack_sequence_and_next_ticket(self):
c = self.c
self.submit(
pools=[
PoolTransfer(PoolName.SWA, host_indices=torch.arange(4), keys=["h1"])
]
)
ticket = c.pp_prefetch_states.pop("hit").ticket
following = pickle.loads(pickle.dumps(ticket))
following.handle = CacheRequestHandle("next", 0)
c.pp_rank = 1
kv = torch.arange(8)
c.mem_pool_host.alloc.side_effect = [
kv,
KeyError(PoolName.SWA),
torch.arange(8),
torch.arange(4),
]
with self.assertLogs(level="ERROR"):
self.run_commands(ticket, following)
failed = c.pp_prefetch_states["hit"].operation
self.assertTrue(failed.is_terminated())
c.mem_pool_host.free.assert_called_once_with(kv, pool=PoolName.KV)
c.page_get_func = Mock(return_value=1)
c.storage_backend = Mock()
c.storage_backend.batch_get_v2.return_value = {"swa": [True]}
with (
patch("sglang.srt.managers.cache_controller.STORAGE_BATCH_SIZE", 1),
patch.object(
c.storage_stop_event, "is_set", side_effect=[False, False, True]
),
):
c.prefetch_io_aux_func()
acks = [c.prefetch_sync_queue.get_nowait() for _ in range(8)]
self.assertEqual([a.rid for a in acks], ["hit"] * 4 + ["next"] * 4)
self.assertEqual(
[a.completed_tokens for a in acks], [0, 0, None, None, 4, 8, None, None]
)
self.assertEqual((acks[2].pool_hits, acks[6].pool_hits), ({}, {"swa": 1}))
self.assertEqual(
c.page_get_func.call_count, 2
) # No reads for the failed ticket.
self.sync_acks(*acks)
self.assertTrue(c.is_pp_prefetch_ready("hit"))
self.assertEqual(c.pp_prefetch_states["next"].operation.completed_tokens, 8)
self.assertEqual(c.prefetch_tokens_occupied, 8)
def test_cancel_before_ticket_defers_free_until_final_ack(self):
c, cache = self.c, self.cache
self.submit(pools=[PoolTransfer(PoolName.SWA, host_indices=torch.arange(4))])
ticket = c.pp_prefetch_states.pop("hit").ticket
c.pp_rank = 1
cache.bind_prefetch_ticket("hit")
cache.finish(ticket.handle, CacheRequestOutcome.ABORT)
cache.finish(ticket.handle, CacheRequestOutcome.ABORT)
self.assertIs(c.pp_prefetch_decisions["hit"], PPPrefetchDecision.CANCELLED)
self.run_commands(ticket)
operation = c.prefetch_buffer.get_nowait()
self.sync_acks(PrefetchAck("hit", operation, completed_tokens=4))
c.mem_pool_host.free.assert_not_called()
self.sync_acks(PrefetchAck("hit", operation, completed_req=True))
self.assertEqual(
[call.kwargs["pool"] for call in c.mem_pool_host.free.call_args_list],
[PoolName.KV, PoolName.SWA],
)
self.assertEqual(c.prefetch_tokens_occupied, 0)
self.assertEqual(c.pp_prefetch_states, {})
self.assertEqual(c.pp_prefetch_decisions, {})
def test_lazy_sidecars_use_hit_pages_and_pool_page_size(self):
c = self.c
c.mem_pool_host.get_pool.side_effect = lambda name: Mock(
page_size=1 if name == PoolName.MAMBA else 4
)
self.submit(
pools=[
PoolTransfer(PoolName.SWA, keys=["pending"] * 3),
PoolTransfer(PoolName.MAMBA, keys=["pending"]),
PoolTransfer(PoolName.DRAFT_SWA, indices_from_pool=PoolName.SWA),
]
)
ticket = pickle.loads(pickle.dumps(c.pp_prefetch_states["hit"].ticket))
for rank in (0, 1):
with self.subTest(rank=rank):
c.pp_rank = rank
if rank:
c.pp_prefetch_states.clear()
c.mem_pool_host.alloc.reset_mock()
self.run_commands(ticket)
operation = c.prefetch_buffer.get_nowait()
self.assertEqual(
c.mem_pool_host.alloc.call_args_list,
[
call(8, pool=PoolName.KV),
call(8, pool=PoolName.SWA),
call(1, pool=PoolName.MAMBA),
],
)
swa, mamba, draft_swa = operation.pool_transfers
self.assertIs(draft_swa.host_indices, swa.host_indices)
self.assertEqual(mamba.host_indices.numel(), 1)
# Missing pool metadata also rolls back KV, before a failed-ticket ACK.
c.pp_prefetch_states.clear()
c.mem_pool_host.get_pool.side_effect = KeyError(PoolName.SWA)
with self.assertLogs(level="ERROR"):
self.run_commands(ticket)
self.assertTrue(c.prefetch_buffer.get_nowait().is_terminated())
self.assertEqual(c.mem_pool_host.free.call_count, 1)
self.assertEqual(c.mem_pool_host.free.call_args.kwargs["pool"], PoolName.KV)
def test_failed_source_allocation_does_not_free_borrowed_sidecar_early(self):
c = self.c
borrowed = PoolTransfer(PoolName.SWA, host_indices=torch.arange(4))
operation = self.submit(pools=[borrowed]).operation
c.mem_pool_host.alloc.side_effect = lambda *args, **kwargs: None
self.run_commands()
self.assertTrue(operation.is_terminated())
c.mem_pool_host.free.assert_not_called()
self.sync_acks(
PrefetchAck("hit", operation, completed_tokens=0, completed_req=True)
)
c.take_ready_pp_prefetch("hit")
self.assertIsNone(borrowed.host_indices)
self.assertEqual(c.mem_pool_host.free.call_args.kwargs["pool"], PoolName.SWA)
def test_command_failure_finishes_queue_task_and_is_reported_on_poll(self):
c = self.c
with (
patch(
f"{HybridCacheController.__module__}.broadcast_pyobj",
side_effect=RuntimeError("broken"),
),
self.assertLogs(level="ERROR"),
):
worker = threading.Thread(
target=c.pp_prefetch_command_thread_func, daemon=True
)
c.pp_prefetch_command_thread = worker
worker.start()
self.submit()
worker.join(timeout=1)
self.assertFalse(worker.is_alive())
self.assertEqual(c.pp_prefetch_command_queue.unfinished_tasks, 0)
for rid in ("hit", "next"):
if rid == "next":
self.submit(rid)
with self.assertRaisesRegex(RuntimeError, "ticket thread exited"):
c.is_pp_prefetch_ready(rid)
def test_idle_source_broadcasts_empty_then_processes_ticket_and_stop(self):
c = self.c
c.tp_rank = 1 # The source is a global rank, not pp_rank=0.
operation = self.submit().operation
c.pp_prefetch_command_queue.put(None)
get = c.pp_prefetch_command_queue.get
idle_polls = [True, True]
def get_after_idle(*, timeout):
self.assertEqual(timeout, 60)
if idle_polls:
idle_polls.pop()
raise Empty
return get(block=False)
with (
patch.object(c.pp_prefetch_command_queue, "get", get_after_idle),
patch.object(
torch.distributed, "get_process_group_ranks", return_value=[1, 3, 5, 7]
),
patch(
f"{HybridCacheController.__module__}.broadcast_pyobj",
side_effect=lambda objects, *args, **kwargs: objects,
) as broadcast,
):
c.pp_prefetch_command_thread_func()
calls = broadcast.call_args_list
self.assertEqual([len(c.args[0]) for c in calls], [0, 0, 1, 1])
for invocation in calls:
self.assertEqual(invocation.args[1:], (1, "command"))
self.assertEqual(invocation.kwargs, {"src": 1})
self.assertEqual(c.pp_prefetch_command_queue.unfinished_tasks, 0)
self.assertTrue(c.pp_prefetch_command_queue.empty())
self.assertIs(c.prefetch_buffer.get_nowait(), operation)
self.assertTrue(c.prefetch_buffer.empty())
c.mem_pool_host.alloc.assert_called_once()
def test_idle_downstream_skips_empty_broadcasts_before_ticket_and_stop(self):
c = self.c
self.submit()
ticket = c.pp_prefetch_states.pop("hit").ticket
c.pp_rank = 1
with (
patch(
f"{HybridCacheController.__module__}.broadcast_pyobj",
side_effect=[[], [], [ticket], [None]],
),
patch.object(c.pp_prefetch_command_queue, "task_done") as task_done,
):
c.pp_prefetch_command_thread_func()
self.assertEqual(c.prefetch_buffer.get_nowait().request_id, "hit")
self.assertTrue(c.prefetch_buffer.empty())
c.mem_pool_host.alloc.assert_called_once()
task_done.assert_not_called()
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_storage_prefetch_lifecycle.py
+1
View File
@@ -71,6 +71,7 @@ def _staged_fixture(full_match=2):
full_available_size=Mock(return_value=100)
)
cc = HybridCacheController.__new__(HybridCacheController)
cc.pp_prefetch_command_group = None
cc.page_size = 2
cc.get_hash_str = get_hash_str
cc.prefetch_queue = Queue()
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
+22 -11
View File
@@ -861,6 +861,9 @@ class TestUnifiedRadixCacheEagleHiCacheStorageKey(CustomTestCase):
)
def test_l3_prefetch_uses_bigram_radix_key(self):
from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import (
PrefetchSubmission,
)
from sglang.srt.mem_cache.utils import get_hash_str
cache, allocator, _ = build_fixture(self.cfg)
@@ -891,23 +894,30 @@ class TestUnifiedRadixCacheEagleHiCacheStorageKey(CustomTestCase):
def prefetch_rate_limited(self):
return False
def prefetch(
def get_prefetch_submission(self, rid):
return None
def submit_prefetch(
self,
request_id,
new_input_tokens,
last_hash=None,
prefix_keys=None,
extra_pools=None,
handle,
prefetch_key,
last_hash,
prefix_keys,
matched_prefix_tokens,
pool_transfers,
assume_stored=False,
):
self.prefetch_args = (
request_id,
new_input_tokens,
handle,
prefetch_key,
last_hash,
prefix_keys,
extra_pools,
matched_prefix_tokens,
pool_transfers,
)
return PrefetchSubmission(
operation=mock.Mock(assume_stored=assume_stored)
)
return mock.Mock()
controller = FakeCacheController()
cache.cache_controller = controller
@@ -915,7 +925,7 @@ class TestUnifiedRadixCacheEagleHiCacheStorageKey(CustomTestCase):
CacheRequestHandle("req", 0), cache.root_node_handle(), tokens
)
_, storage_key, _, _, _ = controller.prefetch_args
_, storage_key, _, _, _, _ = controller.prefetch_args
self.assertIsInstance(storage_key, RadixKey)
self.assertTrue(storage_key.is_bigram)
self.assertEqual(len(storage_key), len(tokens) - 1)
@@ -9664,6 +9674,7 @@ class TestPrefetchCommitOrdering(CustomTestCase):
cache.page_size = 1
cache.enable_storage_metrics = False
cache.buffer_pipeline = None # cache-mode commit path
cache.cache_controller.pp_prefetch_decisions = {}
walk_action = object()
insert_result = mock.MagicMock()
insert_result.cache_actions = [walk_action]