[HiCache] Rework the buffer-mode storage prefetch pipeline and retry bookkeeping (#39283)
This commit is contained in:
Zhiqiang Xie
@@ -2,6 +2,8 @@ import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import torch
|
||||
|
||||
import sglang.srt.managers.schedule_policy as schedule_policy
|
||||
from sglang.srt.managers.schedule_batch import Req
|
||||
from sglang.srt.managers.schedule_policy import (
|
||||
@@ -57,6 +59,7 @@ class TestPrefillAdder(CustomTestCase):
|
||||
tree_cache.disable = False
|
||||
tree_cache.inc_lock_ref.return_value = IncLockRefResult()
|
||||
tree_cache.dec_lock_ref.return_value = DecLockRefResult()
|
||||
tree_cache.buffer_pipeline = None
|
||||
return tree_cache
|
||||
|
||||
def create_token_allocator(
|
||||
@@ -102,6 +105,7 @@ class TestPrefillAdder(CustomTestCase):
|
||||
req.time_stats = SimpleNamespace(wait_queue_entry_time=wait_time)
|
||||
req.retracted_stain = False
|
||||
req.host_hit_length = 0
|
||||
req.swa_host_hit_length = 0
|
||||
req.storage_hit_length = 0
|
||||
req.storage_hit_start = None
|
||||
req.host_hit_is_storage = False
|
||||
@@ -158,6 +162,24 @@ class TestPrefillAdder(CustomTestCase):
|
||||
req.cache_request_handle, fulfilled_tokens=0, reason="device_capacity"
|
||||
)
|
||||
|
||||
self.mock_tree_cache.finish_storage_prefetch_admission.reset_mock()
|
||||
req.host_hit_length = 4
|
||||
req.host_loaded_length = 4
|
||||
req.storage_hit_length = 8
|
||||
req.storage_hit_start = 4
|
||||
req.materialized_host_hit_len.return_value = 4
|
||||
req.fulfilled_storage_hit_len.return_value = 4
|
||||
req.needs_host_load_back.return_value = True
|
||||
adder._account_prefill_cache_admission(req, prefix_len=8)
|
||||
self.mock_tree_cache.finish_storage_prefetch_admission.assert_called_once_with(
|
||||
req.cache_request_handle,
|
||||
fulfilled_tokens=4,
|
||||
reason="shrunk",
|
||||
)
|
||||
self.assertEqual(adder.log_device_hit_tokens, 8)
|
||||
self.assertEqual(adder.log_host_hit_tokens, 0)
|
||||
self.assertEqual(adder.log_storage_hit_tokens, 12)
|
||||
|
||||
def test_retracted_storage_prefetch_accounting_is_omitted(self):
|
||||
adder = self.create_adder(self.create_running_batch())
|
||||
req = self.create_mock_req(
|
||||
@@ -668,6 +690,169 @@ class TestPrefillAdder(CustomTestCase):
|
||||
adder.add_one_req(req, has_chunked_req=False, truncation_align_size=None)
|
||||
self.assertIn(req, adder.can_run_list)
|
||||
|
||||
def test_load_back_delivery_mismatch_reselects_the_prefill_shape(self):
|
||||
# Two incidents: a load that delivers nothing left the SWA gate sized
|
||||
# for the tail and the allocator OOMed; a cache-mode load that also
|
||||
# surfaces FULL device tokens behind a host-only SWA window tripped a
|
||||
# strict promised==loaded check and crashed the scheduler.
|
||||
WINDOW, PAGE = 128, 8
|
||||
SPAN, HOST_HIT = 1024, 1016
|
||||
self.mock_token_allocator.swa_available_size.return_value = 400
|
||||
self.mock_token_allocator.full_available_size.return_value = 100_000
|
||||
self.mock_token_allocator.available_size.return_value = 100_000
|
||||
self.mock_tree_cache.sliding_window_size = WINDOW
|
||||
self.mock_tree_cache.is_tree_cache.return_value = False
|
||||
|
||||
def run(delivered: int, remaining_after_load: int = 100_000):
|
||||
self.mock_token_allocator.full_available_size.return_value = 100_000
|
||||
self.mock_token_allocator.swa_available_size.return_value = 400
|
||||
adder = self.create_adder(self.create_running_batch(), page_size=PAGE)
|
||||
adder.is_hybrid_swa = True
|
||||
req = self.create_mock_req("dropped-fetch", priority=0, max_new_tokens=8)
|
||||
req.prefix_indices = torch.empty(0, dtype=torch.int64)
|
||||
req.full_untruncated_fill_ids = list(range(SPAN))
|
||||
req.host_hit_length = HOST_HIT
|
||||
req.swa_host_hit_length = WINDOW
|
||||
req.needs_host_load_back.return_value = True
|
||||
req.last_node = MagicMock()
|
||||
req.best_match_node = MagicMock()
|
||||
req.kv = SimpleNamespace(cache_protected_len=0)
|
||||
|
||||
def set_extend_range(start, end):
|
||||
req.extend_range = Range(start, end)
|
||||
|
||||
req.set_extend_range = MagicMock(side_effect=set_extend_range)
|
||||
req.sampling_params = SimpleNamespace(max_new_tokens=8, ignore_eos=False)
|
||||
|
||||
def load_back(params):
|
||||
self.mock_token_allocator.full_available_size.return_value = (
|
||||
remaining_after_load
|
||||
)
|
||||
if remaining_after_load == 0:
|
||||
self.mock_token_allocator.swa_available_size.return_value = 0
|
||||
return torch.arange(delivered, dtype=torch.int64), req.last_node
|
||||
|
||||
self.mock_tree_cache.init_load_back.side_effect = load_back
|
||||
verdict = adder.add_one_req(
|
||||
req, has_chunked_req=False, truncation_align_size=None
|
||||
)
|
||||
return verdict, list(adder.can_run_list), req
|
||||
|
||||
# Promise kept: only the 8-token tail is prefilled, which fits.
|
||||
_, admitted, _ = run(HOST_HIT)
|
||||
self.assertEqual(len(admitted), 1)
|
||||
# Nothing delivered: the whole span is prefilled and no longer fits, so
|
||||
# admission must decline rather than OOM the pool.
|
||||
verdict, admitted, _ = run(0)
|
||||
self.assertIs(verdict, AddReqResult.NO_TOKEN)
|
||||
self.assertEqual(admitted, [])
|
||||
# Over-delivery: admitted with the loaded prefix, not the promise.
|
||||
# The loaded prefix is now pinned and no longer part of the evictable
|
||||
# budget. A successful load must not run admission gates again.
|
||||
_, admitted, req = run(HOST_HIT + 4, remaining_after_load=0)
|
||||
self.assertEqual(len(admitted), 1)
|
||||
self.assertEqual(len(req.prefix_indices), HOST_HIT + 4)
|
||||
self.assertEqual(req.kv.cache_protected_len, HOST_HIT + 4)
|
||||
req.set_extend_range.assert_called_once_with(HOST_HIT + 4, SPAN)
|
||||
# A partial FULL load stays fatal.
|
||||
with self.assertRaisesRegex(RuntimeError, "promised"):
|
||||
run(HOST_HIT // 2)
|
||||
|
||||
def _create_host_hit_req(self, *, prefix_len=0, host_hit=8192, tail=1024):
|
||||
req = self._create_delayer_req(prefix_len + host_hit + tail)
|
||||
req.prefix_indices = torch.arange(prefix_len)
|
||||
req.host_hit_length = host_hit
|
||||
req.needs_host_load_back.return_value = True
|
||||
req.best_match_node = req.last_node
|
||||
req.kv = SimpleNamespace(cache_protected_len=prefix_len)
|
||||
return req
|
||||
|
||||
def test_successful_load_back_commits_the_selected_shape_once(self):
|
||||
cases = (
|
||||
("full", 0, 24, None, None, 8, 8),
|
||||
("full_unaligned", 0, 24, None, None, 7, 8),
|
||||
("retracted_unaligned", 0, 24, None, None, 7, 8),
|
||||
("chunk", 0, 24, 4, None, 4, 0),
|
||||
("aux_only", 24, 0, None, None, 8, 8),
|
||||
("overdelivery_full", 0, 24, None, None, 8, 8),
|
||||
("overdelivery_chunk", 0, 24, 4, None, 4, 0),
|
||||
("overdelivery_chunk_end", 0, 24, 8, None, 8, 8),
|
||||
(
|
||||
"dllm",
|
||||
0,
|
||||
24,
|
||||
None,
|
||||
SimpleNamespace(block_size=4, max_running_requests=2),
|
||||
4,
|
||||
0,
|
||||
),
|
||||
(
|
||||
"overdelivery_dllm",
|
||||
0,
|
||||
24,
|
||||
None,
|
||||
SimpleNamespace(block_size=4, max_running_requests=2),
|
||||
4,
|
||||
0,
|
||||
),
|
||||
)
|
||||
for name, prefix_len, host_hit, chunk, dllm, extend, decode in cases:
|
||||
with self.subTest(mode=name):
|
||||
self.mock_tree_cache.reset_mock()
|
||||
adder = self._create_delayer_adder(
|
||||
available_tokens=100_000,
|
||||
delayer=None,
|
||||
page_size=2,
|
||||
rem_chunk_tokens=chunk,
|
||||
dllm_config=dllm,
|
||||
)
|
||||
req = self._create_host_hit_req(
|
||||
prefix_len=prefix_len,
|
||||
host_hit=host_hit,
|
||||
tail=extend if chunk is None and dllm is None else 8,
|
||||
)
|
||||
req.retracted_stain = name == "retracted_unaligned"
|
||||
if name.startswith("overdelivery"):
|
||||
req.host_hit_length -= 4
|
||||
old_node, restored_node = req.last_node, object()
|
||||
if name == "aux_only":
|
||||
req.swa_host_hit_length = 8
|
||||
|
||||
def load_back(params):
|
||||
self.assertIs(params.req, req)
|
||||
tile_gate.assert_called_once()
|
||||
tile_gate.return_value = AddReqResult.OTHER
|
||||
return torch.arange(host_hit), restored_node
|
||||
|
||||
self.mock_tree_cache.init_load_back.side_effect = load_back
|
||||
with patch.object(
|
||||
adder, "_check_prefill_tile_budget", return_value=None
|
||||
) as tile_gate:
|
||||
adder.add_one_req(req, False, None)
|
||||
tile_gate.assert_called_once()
|
||||
self.mock_tree_cache.init_load_back.assert_called_once()
|
||||
self.assertEqual(adder.can_run_list, [req])
|
||||
req.set_extend_range.assert_called_once_with(24, 24 + extend)
|
||||
self.mock_tree_cache.inc_lock_ref.assert_any_call(restored_node)
|
||||
self.assertIs(
|
||||
self.mock_tree_cache.dec_lock_ref.call_args.args[0], old_node
|
||||
)
|
||||
self.assertEqual(adder.log_hit_tokens, 24)
|
||||
self.assertEqual(adder.log_input_tokens, extend)
|
||||
self.assertEqual(
|
||||
adder.reprocessed_log_input_tokens,
|
||||
extend if req.retracted_stain else 0,
|
||||
)
|
||||
self.assertEqual(
|
||||
adder.rem_total_token_offset,
|
||||
adder.ceil_paged_tokens(extend) + decode + 2,
|
||||
)
|
||||
self.assertEqual(
|
||||
adder.new_chunked_req is req,
|
||||
name in ("chunk", "overdelivery_chunk"),
|
||||
)
|
||||
self.mock_tree_cache.init_load_back.side_effect = None
|
||||
|
||||
def test_swa_new_tokens_clamps_remaining_not_total(self):
|
||||
# Remaining decode headroom must be min(max_new - generated, CLIP)
|
||||
# (subtract-then-clip). The reversed order (clip-then-subtract) zeroes
|
||||
|
||||
@@ -81,6 +81,10 @@ def _scheduler_for_get_next_batch(*, tree_cache, chunked_req) -> Scheduler:
|
||||
s.dllm_manager = None
|
||||
s.enable_hisparse = False
|
||||
s.enable_fpm = False
|
||||
# Exercise the unconditional scheduler-loop HiCache event-drain point.
|
||||
s.enable_hierarchical_cache = True
|
||||
s.enable_hicache_storage = False
|
||||
s.enable_unified_cache_external_linker = False
|
||||
s.last_batch = None
|
||||
s.require_mlp_sync = False
|
||||
s.spec_algorithm = MagicMock()
|
||||
@@ -104,6 +108,7 @@ def _scheduler_for_get_next_batch(*, tree_cache, chunked_req) -> Scheduler:
|
||||
side_effect=lambda batch, **_: batch
|
||||
)
|
||||
s.update_running_batch = MagicMock(side_effect=lambda batch: batch)
|
||||
tree_cache.check_hicache_events = MagicMock()
|
||||
s.tree_cache = tree_cache
|
||||
s.chunked_req = chunked_req
|
||||
s._pending_chunked_abort_req = None
|
||||
|
||||
@@ -0,0 +1,139 @@
|
||||
"""HiCache progress must not depend on whether a prefill batch is admitted."""
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock, call, patch
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import maybe_stub_sgl_kernel
|
||||
|
||||
maybe_stub_sgl_kernel()
|
||||
|
||||
from sglang.srt.managers.schedule_batch import NextBatchPlan
|
||||
from sglang.srt.managers.scheduler import Scheduler
|
||||
|
||||
register_cpu_ci(est_time=2, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
class TestSchedulerHiCacheEvents(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.calls = Mock()
|
||||
self.scheduler = s = Scheduler.__new__(Scheduler)
|
||||
s.scheduler_stage_metrics = None
|
||||
s.enable_hierarchical_cache = True
|
||||
s.enable_unified_cache_external_linker = False
|
||||
s.enable_hicache_storage = True
|
||||
s.tree_cache = SimpleNamespace(check_hicache_events=self.calls.drain)
|
||||
s._process_storage_prefetch_retries = self.calls.retry
|
||||
s.process_pending_chunked_abort = Mock()
|
||||
s.process_prefill_chunk = Mock()
|
||||
s.dp_attn_adapter = Mock()
|
||||
s.dp_attn_adapter.maybe_prepare_mlp_sync_batch.return_value = None
|
||||
self.running_batch = SimpleNamespace(
|
||||
batch_is_full=False, is_prefill_only=False, is_empty=lambda: True
|
||||
)
|
||||
|
||||
def test_feature_gates(self):
|
||||
for hierarchical, flexkv, linker, storage in (
|
||||
(False, False, False, False),
|
||||
(True, False, False, False),
|
||||
(False, True, False, False),
|
||||
(False, False, True, False),
|
||||
(True, False, False, True),
|
||||
):
|
||||
with (
|
||||
self.subTest(
|
||||
hierarchical=hierarchical,
|
||||
flexkv=flexkv,
|
||||
linker=linker,
|
||||
storage=storage,
|
||||
),
|
||||
patch(
|
||||
"sglang.srt.managers.scheduler.get_memory",
|
||||
return_value=SimpleNamespace(enable_flexkv=flexkv),
|
||||
),
|
||||
):
|
||||
self.calls.reset_mock()
|
||||
s = self.scheduler
|
||||
s.enable_hierarchical_cache = hierarchical
|
||||
s.enable_unified_cache_external_linker = linker
|
||||
s.enable_hicache_storage = storage
|
||||
s._process_hicache_events()
|
||||
expected = [call.drain()] if hierarchical or flexkv or linker else []
|
||||
if storage:
|
||||
expected.append(call.retry())
|
||||
self.assertEqual(self.calls.mock_calls, expected)
|
||||
|
||||
def test_pd_prefill_drains_before_admission_even_with_empty_queue(self):
|
||||
s = self.scheduler
|
||||
s.resolve_waiting_queue_bootstrap = Mock()
|
||||
s.get_new_batch_prefill = self.calls.admit
|
||||
s.get_new_batch_prefill.return_value = NextBatchPlan(
|
||||
batch_to_run=None, running_batch=self.running_batch
|
||||
)
|
||||
for waiting_queue in ([], [SimpleNamespace(rid="pending_l3")]):
|
||||
with self.subTest(waiting=bool(waiting_queue)):
|
||||
s.waiting_queue = waiting_queue
|
||||
self.calls.reset_mock()
|
||||
for _ in range(2):
|
||||
plan = s.get_next_disagg_prefill_batch_to_run(
|
||||
running_batch=self.running_batch, last_batch=None
|
||||
)
|
||||
self.assertIsNone(plan.batch_to_run)
|
||||
self.assertEqual(
|
||||
self.calls.mock_calls,
|
||||
[call.drain(), call.retry(), call.admit(self.running_batch)] * 2,
|
||||
)
|
||||
|
||||
def test_unified_drains_when_prefill_is_deferred(self):
|
||||
s = self.scheduler
|
||||
s.enable_fpm = False
|
||||
s._abort_on_waiting_timeout = Mock()
|
||||
s._abort_on_running_timeout = Mock()
|
||||
s.dllm_config = None
|
||||
s.chunked_req = None
|
||||
s.enable_hisparse = False
|
||||
s.require_mlp_sync = False
|
||||
s._should_defer_prefill = self.calls.defer
|
||||
s._should_defer_prefill.return_value = True
|
||||
s.get_new_batch_prefill = Mock()
|
||||
s.dp_attn_adapter.maybe_convert_decode_to_extend.return_value = None
|
||||
s._arm_prefill_decode_interval = Mock()
|
||||
s.ngram_embedding_manager = Mock()
|
||||
s.ngram_embedding_manager.prepare_for_forward.return_value = None
|
||||
|
||||
plan = s.get_next_batch_to_run(self.running_batch, None)
|
||||
|
||||
self.assertIsNone(plan.batch_to_run)
|
||||
s.get_new_batch_prefill.assert_not_called()
|
||||
self.assertEqual(
|
||||
self.calls.mock_calls, [call.drain(), call.retry(), call.defer()]
|
||||
)
|
||||
|
||||
def test_pp_prefill_drains_before_admission(self):
|
||||
s = self.scheduler
|
||||
s.init_pp_loop_state = Mock()
|
||||
s.pp_loop_size = 1
|
||||
s.ps = SimpleNamespace(pp_size=2)
|
||||
s.pp_group = SimpleNamespace(is_last_rank=True)
|
||||
s.running_mbs = [self.running_batch]
|
||||
s.last_mbs = [None]
|
||||
s.ingest_requests = Mock(return_value=[])
|
||||
s._pp_pd_get_bootstrapped_ids = Mock(return_value=[])
|
||||
s._pp_pd_get_prefill_transferred_ids = Mock(return_value=[])
|
||||
s._pp_commit_comm_work = Mock()
|
||||
s.get_new_batch_prefill = self.calls.admit
|
||||
# Stop the infinite loop at admission; no PP transport or GPU is needed.
|
||||
s.get_new_batch_prefill.side_effect = StopIteration
|
||||
|
||||
with self.assertRaises(StopIteration):
|
||||
s.event_loop_pp_disagg_prefill()
|
||||
|
||||
self.assertEqual(
|
||||
self.calls.mock_calls,
|
||||
[call.drain(), call.retry(), call.admit(self.running_batch)],
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -26,6 +26,7 @@ from sglang.srt.mem_cache.unified_cache.components.base import (
|
||||
from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import (
|
||||
BufferBackupSnapshot,
|
||||
)
|
||||
from sglang.srt.mem_cache.unified_radix_cache import _OngoingPrefetch
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=8, suite="base-a-test-cpu")
|
||||
@@ -244,13 +245,13 @@ class TestBufferModeSidecar(unittest.TestCase):
|
||||
cache.page_size = 2
|
||||
cache.cache_controller.prefetch_tokens_occupied = len(host_indices)
|
||||
cache.ongoing_prefetch = {
|
||||
req_id: (
|
||||
0,
|
||||
RadixKey(array("q", [1, 2, 3, 4])),
|
||||
host_indices,
|
||||
operation,
|
||||
None,
|
||||
{ComponentType.SWA: [swa]},
|
||||
req_id: _OngoingPrefetch(
|
||||
anchor_node_id=0,
|
||||
prefetch_key=RadixKey(array("q", [1, 2, 3, 4])),
|
||||
host_indices=host_indices,
|
||||
operation=operation,
|
||||
anchor_lock_params=None,
|
||||
comp_xfers={ComponentType.SWA: [swa]},
|
||||
)
|
||||
}
|
||||
cache.prefetch_loaded_tokens_by_reqid = {}
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
"""Unit tests for HiCache staged write-back host-pool dispatch."""
|
||||
|
||||
import unittest
|
||||
from array import array
|
||||
from contextlib import contextmanager
|
||||
from types import SimpleNamespace
|
||||
from unittest import mock
|
||||
@@ -264,15 +265,22 @@ class TestHiCacheStagedWriteBackDispatch(CustomTestCase):
|
||||
controller._num_tokens_by_pool.assert_called_once_with(merged_op)
|
||||
self.assertEqual(controller.ack_load_queue[0].node_ids, [7, 7])
|
||||
|
||||
def test_short_staged_swa_tail_resolves_device_covered_head(self):
|
||||
def _short_swa_tail_pipeline(self, swa_page_size: int) -> BufferModePipeline:
|
||||
"""Pipeline holding one staged span [2, 8) whose 4-slot trailing SWA
|
||||
window outruns the splice left by a device prefix of 6."""
|
||||
handle = CacheRequestHandle("r", 0)
|
||||
pipeline = BufferModePipeline.__new__(BufferModePipeline)
|
||||
pipeline._cache = mock.Mock()
|
||||
pipeline._cache.cache_controller.mem_pool_host.entry_map = {
|
||||
PoolName.SWA: SimpleNamespace(
|
||||
host_pool=SimpleNamespace(page_size=swa_page_size)
|
||||
)
|
||||
}
|
||||
pipeline.release_staged_hold = mock.Mock(return_value=True)
|
||||
pipeline.staged_prefetches = {
|
||||
handle: SimpleNamespace(
|
||||
request=handle,
|
||||
key_tokens=list(range(8)),
|
||||
key_tokens=array("q", range(8)),
|
||||
extra_key=None,
|
||||
cache_salt=None,
|
||||
matched_len=2,
|
||||
@@ -289,14 +297,42 @@ class TestHiCacheStagedWriteBackDispatch(CustomTestCase):
|
||||
operation_id=1,
|
||||
)
|
||||
}
|
||||
return pipeline
|
||||
|
||||
self.assertEqual(
|
||||
pipeline.plan_staged_splice(handle, device_prefix_len=6), (0, 0)
|
||||
def test_short_staged_swa_tail_keeps_complete_window(self):
|
||||
"""FULL-prefix growth trims only FULL; SWA keeps its complete window."""
|
||||
handle = CacheRequestHandle("r", 0)
|
||||
pipeline = self._short_swa_tail_pipeline(swa_page_size=2)
|
||||
pipeline._cache.tree_core.is_eagle = False
|
||||
pipeline._cache.tree_core.match_full_device_prefix.return_value = (6, 1, 6)
|
||||
pipeline._cache.tree_core.collect_full_device_indices.return_value = _indices(
|
||||
0, 6
|
||||
)
|
||||
pipeline._cache._resolve_storage_prefetch_tokens.assert_called_once_with(
|
||||
handle, 4
|
||||
req = SimpleNamespace(
|
||||
rid="r",
|
||||
cache_request_handle=handle,
|
||||
prefix_indices=_indices(0, 0),
|
||||
kv=SimpleNamespace(cache_protected_len=0),
|
||||
)
|
||||
pipeline.release_staged_hold.assert_called_once_with(handle, reason="shrunk")
|
||||
self.assertTrue(pipeline.prepare_staged_prefetch(req))
|
||||
self.assertEqual((req.host_hit_length, req.swa_host_hit_length), (2, 4))
|
||||
pipeline.release_staged_hold.assert_not_called()
|
||||
|
||||
pipeline = self._short_swa_tail_pipeline(swa_page_size=4)
|
||||
pipeline._cache.tree_core.is_eagle = False
|
||||
pipeline._cache.tree_core.match_full_device_prefix.return_value = (6, 1, 6)
|
||||
pipeline._cache.tree_core.collect_full_device_indices.return_value = _indices(
|
||||
0, 6
|
||||
)
|
||||
req = SimpleNamespace(
|
||||
rid="r",
|
||||
cache_request_handle=handle,
|
||||
prefix_indices=_indices(0, 0),
|
||||
kv=SimpleNamespace(cache_protected_len=0),
|
||||
)
|
||||
self.assertTrue(pipeline.prepare_staged_prefetch(req))
|
||||
self.assertEqual((req.host_hit_length, req.swa_host_hit_length), (2, 4))
|
||||
pipeline.release_staged_hold.assert_not_called()
|
||||
|
||||
def test_l2_transfer_maps_global_layers(self):
|
||||
host_pool = mock.Mock()
|
||||
|
||||
@@ -1833,17 +1833,26 @@ def test_swa_prefetch_commit_end_to_end():
|
||||
core.has_swa_host_pool = True
|
||||
anchor = core.match_prefix(MatchPrefixParams(key=_key([99]))).best_match_node
|
||||
|
||||
# The build wraps the host buffer with placeholder keys, trailing-pages policy.
|
||||
# Without planned staging the SWA pool takes no part in the fetch.
|
||||
assert (
|
||||
core.build_hicache_transfers(
|
||||
ComponentType.SWA, anchor, CacheTransferPhase.PREFETCH
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
# The build carries the planned staging as placeholder keys, trailing-pages
|
||||
# policy; the host buffer is attached once the hit is known.
|
||||
(xfer,) = core.build_hicache_transfers(
|
||||
ComponentType.SWA,
|
||||
anchor,
|
||||
CacheTransferPhase.PREFETCH,
|
||||
host_indices=torch.tensor([30, 31], dtype=torch.int64),
|
||||
staging_tokens=2,
|
||||
)
|
||||
assert xfer.name == PoolName.SWA
|
||||
assert xfer.keys == ["__placeholder__", "__placeholder__"]
|
||||
assert xfer.hit_policy == PoolHitPolicy.TRAILING_PAGES
|
||||
assert xfer.host_indices.tolist() == [30, 31]
|
||||
assert xfer.host_indices is None
|
||||
|
||||
# The prefetched suffix lands as one host node; its SWA host is a tombstone.
|
||||
insert_result = core.insert_host(
|
||||
|
||||
@@ -0,0 +1,498 @@
|
||||
"""Staged L3 prefetch lifecycle through the buffer pipeline; no GPU kernels."""
|
||||
|
||||
import tempfile
|
||||
import unittest
|
||||
from array import array
|
||||
from collections import defaultdict, deque
|
||||
from datetime import timedelta
|
||||
from queue import Queue
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.mem_cache.base_prefix_cache import (
|
||||
CacheRequestHandle,
|
||||
InitLoadBackParams,
|
||||
)
|
||||
from sglang.srt.mem_cache.buffer_mode.pipeline import (
|
||||
BufferModePipeline,
|
||||
_StagedPrefetch,
|
||||
)
|
||||
from sglang.srt.mem_cache.hicache_storage import PoolName, PoolTransfer
|
||||
from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import (
|
||||
HybridCacheController,
|
||||
)
|
||||
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,
|
||||
_OngoingPrefetch,
|
||||
)
|
||||
from sglang.srt.mem_cache.utils import get_hash_str
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=20, suite="base-a-test-cpu")
|
||||
|
||||
_REQ = CacheRequestHandle("r", 0)
|
||||
|
||||
|
||||
def _staged_fixture(full_match=2):
|
||||
cache = UnifiedRadixCache.__new__(UnifiedRadixCache)
|
||||
cache.tree_core = SimpleNamespace(
|
||||
page_size=2,
|
||||
is_eagle=False,
|
||||
enable_storage=True,
|
||||
prefetch_anchor_info=lambda node: (None, None),
|
||||
match_full_device_prefix=Mock(return_value=(full_match, 1, full_match)),
|
||||
collect_full_device_indices=Mock(return_value=torch.arange(8)),
|
||||
inc_full_pin=Mock(),
|
||||
dec_full_pin=Mock(),
|
||||
empty_match_result=SimpleNamespace(
|
||||
last_device_node=0, device_indices=torch.arange(0)
|
||||
),
|
||||
)
|
||||
cache.host_memory_mode = "buffer_only"
|
||||
cache.linker = None
|
||||
cache.storage_prefetch_retries = StoragePrefetchRetries()
|
||||
cache.prefetch_loaded_tokens_by_reqid = {_REQ: 6}
|
||||
cache.prefetch_loaded_storage_start_by_reqid = {_REQ: 2}
|
||||
cache._storage_prefetch_hit_remaining_by_reqid = {}
|
||||
cache.enable_storage_metrics = False
|
||||
cache.storage_metrics_collector = None
|
||||
cache.ongoing_prefetch = {}
|
||||
cache._prefetch_outcome_stats = defaultdict(int)
|
||||
cache.tree_components = []
|
||||
cache.prefetch_threshold = 2
|
||||
cache._build_sidecar_transfers = Mock(return_value=[])
|
||||
cache.supports_swa = lambda: True
|
||||
cache.evict_for_alloc = Mock()
|
||||
cache.token_to_kv_pool_allocator = SimpleNamespace(
|
||||
full_available_size=Mock(return_value=100)
|
||||
)
|
||||
cc = HybridCacheController.__new__(HybridCacheController)
|
||||
cc.page_size = 2
|
||||
cc.get_hash_str = get_hash_str
|
||||
cc.prefetch_queue = Queue()
|
||||
cc.prefetch_tokens_occupied = 6
|
||||
cc.prefetch_rate_limited = lambda: False
|
||||
cc.load = Mock(return_value=None)
|
||||
cc.storage_backend = Mock()
|
||||
cc.storage_backend.batch_exists.return_value = 0
|
||||
cc.mem_pool_host = SimpleNamespace(
|
||||
free=Mock(),
|
||||
entry_map={
|
||||
PoolName.SWA: SimpleNamespace(host_pool=SimpleNamespace(free=Mock()))
|
||||
},
|
||||
)
|
||||
cache.cache_controller = cc
|
||||
pipeline = BufferModePipeline.__new__(BufferModePipeline)
|
||||
pipeline._cache = cache
|
||||
pipeline.reset()
|
||||
cache.buffer_pipeline = pipeline
|
||||
pipeline.staged_prefetches[_REQ] = _StagedPrefetch(
|
||||
request=_REQ,
|
||||
key_tokens=array("q", range(8)),
|
||||
extra_key=None,
|
||||
cache_salt=None,
|
||||
matched_len=2,
|
||||
num_tokens=6,
|
||||
occupied_tokens=6,
|
||||
host_indices=torch.arange(6),
|
||||
aux_xfers=[PoolTransfer(name=PoolName.SWA, host_indices=torch.arange(4))],
|
||||
hash_values=["a", "b", "c"],
|
||||
operation_id=1,
|
||||
)
|
||||
req = SimpleNamespace(
|
||||
rid="r",
|
||||
cache_request_handle=_REQ,
|
||||
prefix_indices=torch.arange(2),
|
||||
last_node=1,
|
||||
kv=SimpleNamespace(cache_protected_len=2),
|
||||
extra_key=None,
|
||||
cache_salt=None,
|
||||
host_hit_length=0,
|
||||
swa_host_hit_length=0,
|
||||
host_hit_is_storage=False,
|
||||
host_loaded_length=0,
|
||||
storage_prefetch_last_match_len=4,
|
||||
storage_prefetch_retry_attempts=0,
|
||||
)
|
||||
return cache, pipeline, req
|
||||
|
||||
|
||||
def _hit_drain_fixture():
|
||||
"""A buffer-mode cache whose hit drain and outcome accounting are real."""
|
||||
cache = UnifiedRadixCache.__new__(UnifiedRadixCache)
|
||||
cache.host_memory_mode = "buffer_only"
|
||||
cache.prefetch_threshold = 2
|
||||
cache.enable_storage_metrics = False
|
||||
cache.storage_metrics_collector = None
|
||||
cache.storage_prefetch_retries = StoragePrefetchRetries()
|
||||
cache._prefetch_outcome_stats = defaultdict(int)
|
||||
cache._storage_prefetch_hit_remaining_by_reqid = {}
|
||||
cache._record_storage_prefetch_hit = Mock()
|
||||
cache.revoke_pending_prefetch = Mock()
|
||||
cache.buffer_pipeline = SimpleNamespace(pending_hit_allocs=deque())
|
||||
cache.cache_controller = SimpleNamespace(
|
||||
prefetch_hit_queue=Queue(),
|
||||
ack_prefetch_queue=Queue(),
|
||||
ack_backup_queue=Queue(),
|
||||
host_mem_release_queue=Queue(),
|
||||
extra_host_mem_release_queues={},
|
||||
)
|
||||
cache.ongoing_prefetch = {}
|
||||
return cache
|
||||
|
||||
|
||||
def _terminated_query(cache, rid, hit_tokens):
|
||||
handle = CacheRequestHandle(rid, 0)
|
||||
operation = SimpleNamespace(
|
||||
request_id=rid,
|
||||
handle=handle,
|
||||
storage_hit_count=hit_tokens,
|
||||
stats_requested_tokens=8,
|
||||
is_terminated=lambda: True,
|
||||
)
|
||||
cache.ongoing_prefetch[handle] = _OngoingPrefetch(
|
||||
0, RadixKey(array("q", range(8))), None, operation, None, {}
|
||||
)
|
||||
cache.cache_controller.prefetch_hit_queue.put(operation)
|
||||
|
||||
|
||||
def _two_rank_retry_trace(rank, rendezvous):
|
||||
torch.distributed.init_process_group(
|
||||
"gloo",
|
||||
init_method=f"file://{rendezvous}",
|
||||
rank=rank,
|
||||
world_size=2,
|
||||
timeout=timedelta(seconds=30),
|
||||
)
|
||||
try:
|
||||
cache, pipeline, req = _staged_fixture()
|
||||
held = pipeline.staged_prefetches.pop(req.cache_request_handle)
|
||||
cache.cache_controller.prefetch_tokens_occupied = 0
|
||||
waiting = [SimpleNamespace(rid="head", storage_prefetch_retry_attempts=0), req]
|
||||
published = False
|
||||
issued = []
|
||||
for step in range(7):
|
||||
# Native completion arrives one pass earlier on rank 0. Only the
|
||||
# agreed completion may publish staging to the admission path.
|
||||
ready = torch.tensor([int(step >= rank + 1)])
|
||||
torch.distributed.all_reduce(ready, op=torch.distributed.ReduceOp.MIN)
|
||||
if ready.item() and not published:
|
||||
pipeline.staged_prefetches[req.cache_request_handle] = held
|
||||
cache.cache_controller.prefetch_tokens_occupied = 6
|
||||
published = True
|
||||
for retry_req, hit_end in cache.storage_prefetch_retries.pop_ready(
|
||||
waiting, 2, 8
|
||||
):
|
||||
issued.append((step, retry_req.rid, hit_end))
|
||||
if pipeline.has_staged(req.cache_request_handle):
|
||||
full_match = 2 if step < 3 else 0
|
||||
cache.tree_core.match_full_device_prefix.return_value = (
|
||||
full_match,
|
||||
1,
|
||||
full_match,
|
||||
)
|
||||
req.prefix_indices = torch.arange(full_match)
|
||||
if cache.buffer_pipeline.prepare_staged_prefetch(req):
|
||||
assert (
|
||||
cache.init_load_back(
|
||||
InitLoadBackParams(None, req.host_hit_length, req=req)
|
||||
)
|
||||
is None
|
||||
)
|
||||
snapshot = (
|
||||
list(issued),
|
||||
pipeline.has_staged(req.cache_request_handle),
|
||||
cache.cache_controller.prefetch_tokens_occupied,
|
||||
)
|
||||
snapshots = [None, None]
|
||||
torch.distributed.all_gather_object(snapshots, snapshot)
|
||||
assert snapshots[0] == snapshots[1], (step, snapshots)
|
||||
assert issued == [(4, "r", 8)], issued
|
||||
assert cache.cache_controller.load.call_count == 1
|
||||
assert cache.cache_controller.prefetch_tokens_occupied == 0
|
||||
finally:
|
||||
torch.distributed.destroy_process_group()
|
||||
|
||||
|
||||
class TestStagedPrefetchLifecycle(unittest.TestCase):
|
||||
def test_trim_and_stage_preserve_raw_token_boundaries(self):
|
||||
for bigram in (False, True):
|
||||
for trims in ((2,), (2, 2), (8,), (2, 6)):
|
||||
with self.subTest(bigram=bigram, trims=trims):
|
||||
cache, pipeline, req = _staged_fixture()
|
||||
pipeline.staged_prefetches.clear()
|
||||
cache.tree_core.is_eagle = bigram
|
||||
tokens = array("q", range(10 + int(bigram)))
|
||||
cache.prefetch_from_storage(
|
||||
req.cache_request_handle,
|
||||
0,
|
||||
tokens[2:],
|
||||
matched_prefix_tokens=tokens[:2],
|
||||
storage_hit_end=10,
|
||||
)
|
||||
info = cache.ongoing_prefetch[req.cache_request_handle]
|
||||
operation = info.operation
|
||||
self.assertEqual(len(info.prefetch_key), 8)
|
||||
self.assertTrue(operation.assume_stored)
|
||||
operation.hash_value = ["h0", "h1", "h2", "h3"]
|
||||
operation.storage_hit_count = 8
|
||||
matched_len, hit_tokens = 2, 8
|
||||
for trim in trims:
|
||||
matched_len += trim
|
||||
info, hit_tokens, aux_tokens = (
|
||||
cache._trim_buffer_prefetch_full_head(
|
||||
req.cache_request_handle,
|
||||
info,
|
||||
operation,
|
||||
matched_len,
|
||||
hit_tokens,
|
||||
)
|
||||
)
|
||||
self.assertEqual(aux_tokens, 8)
|
||||
self.assertEqual(
|
||||
pipeline._prefetch_prefix_ctx[req.cache_request_handle][0],
|
||||
list(tokens[:matched_len]),
|
||||
)
|
||||
self.assertEqual(
|
||||
list(info.prefetch_key.raw_token_ids()),
|
||||
list(tokens[matched_len:]),
|
||||
)
|
||||
# A capacity retry must retain the endpoint even when
|
||||
# there is no FULL suffix and only SWA remains to load.
|
||||
cache.tree_core.match_full_device_prefix.return_value = (
|
||||
matched_len,
|
||||
1,
|
||||
matched_len,
|
||||
)
|
||||
pipeline.anchor_lock_cap_tokens = 100
|
||||
pipeline.try_lock_anchor(req.cache_request_handle, hit_tokens)
|
||||
anchor_key = (
|
||||
cache.tree_core.match_full_device_prefix.call_args.args[0]
|
||||
)
|
||||
self.assertEqual(anchor_key.raw_token_ids(), tokens)
|
||||
pipeline.release_anchor_lock(req.cache_request_handle)
|
||||
cache.tree_core.match_full_device_prefix.reset_mock()
|
||||
swa = PoolTransfer(name=PoolName.SWA, host_indices=torch.arange(4))
|
||||
operation.pool_transfers = [swa]
|
||||
operation.host_indices = torch.arange(hit_tokens)
|
||||
cache.ongoing_prefetch[req.cache_request_handle] = info._replace(
|
||||
host_indices=operation.host_indices, comp_xfers={"swa": [swa]}
|
||||
)
|
||||
cache.cache_controller.prefetch_tokens_occupied = hit_tokens
|
||||
cache.storage_existence_cache = Mock()
|
||||
pipeline.stage_completed_prefetch(
|
||||
req.cache_request_handle, hit_tokens, operation.hash_value
|
||||
)
|
||||
held = pipeline.staged_prefetches[req.cache_request_handle]
|
||||
self.assertEqual(held.key_tokens, tokens)
|
||||
self.assertEqual(held.matched_len, matched_len)
|
||||
self.assertEqual(held.num_tokens, hit_tokens)
|
||||
self.assertEqual(
|
||||
len(RadixKey(held.key_tokens, is_bigram=bigram)), 10
|
||||
)
|
||||
# A joint match counts bigrams, not their extra raw boundary token.
|
||||
req.prefix_indices = torch.arange(10)
|
||||
req.kv.cache_protected_len = 10
|
||||
self.assertTrue(pipeline.prepare_staged_prefetch(req))
|
||||
self.assertFalse(pipeline.has_staged(req.cache_request_handle))
|
||||
cache.tree_core.match_full_device_prefix.assert_not_called()
|
||||
|
||||
def test_two_rank_completion_capacity_and_anchor_loss(self):
|
||||
with tempfile.TemporaryDirectory(prefix="prefetch-rank-test-") as directory:
|
||||
torch.multiprocessing.spawn(
|
||||
_two_rank_retry_trace, args=(f"{directory}/store",), nprocs=2, join=True
|
||||
)
|
||||
|
||||
def test_next_pass_uses_fresh_joint_match(self):
|
||||
cache, pipeline, req = _staged_fixture(full_match=8)
|
||||
self.assertTrue(cache.buffer_pipeline.prepare_staged_prefetch(req))
|
||||
self.assertEqual((req.host_hit_length, req.swa_host_hit_length), (0, 4))
|
||||
self.assertEqual(req.kv.cache_protected_len, 8)
|
||||
plan = req.staged_prefetch_plan
|
||||
self.assertIs(
|
||||
plan.key.token_ids,
|
||||
pipeline.staged_prefetches[req.cache_request_handle].key_tokens,
|
||||
)
|
||||
cache.tree_core.match_full_device_prefix.assert_called_once()
|
||||
# A twin finished first: the next pass's joint match runs past the
|
||||
# staged span and is kept as is (a shrink would strand its recompute).
|
||||
req.prefix_indices = torch.arange(12)
|
||||
req.last_node = 9
|
||||
req.kv.cache_protected_len = 12
|
||||
self.assertTrue(cache.buffer_pipeline.prepare_staged_prefetch(req))
|
||||
self.assertIsNone(req.staged_prefetch_plan)
|
||||
self.assertEqual(len(req.prefix_indices), 12)
|
||||
self.assertEqual((req.last_node, req.kv.cache_protected_len), (9, 12))
|
||||
self.assertEqual((req.host_hit_length, req.swa_host_hit_length), (0, 0))
|
||||
self.assertFalse(pipeline.has_staged(req.cache_request_handle))
|
||||
self.assertEqual(cache.cache_controller.prefetch_tokens_occupied, 0)
|
||||
cache.cache_controller.mem_pool_host.free.assert_called_once()
|
||||
cache.tree_core.match_full_device_prefix.assert_called_once()
|
||||
|
||||
def test_capacity_retry_keeps_buffers_without_a_storage_retry(self):
|
||||
for available in (0, 100):
|
||||
with self.subTest(full_available=available):
|
||||
cache, pipeline, req = _staged_fixture()
|
||||
cache.cache_controller.load.return_value = None
|
||||
cache.token_to_kv_pool_allocator.full_available_size.return_value = (
|
||||
available
|
||||
)
|
||||
held = pipeline.staged_prefetches[req.cache_request_handle]
|
||||
pipeline.anchor_lock_cap_tokens = 8
|
||||
pipeline._prefetch_prefix_ctx[req.cache_request_handle] = (
|
||||
[0, 1],
|
||||
None,
|
||||
None,
|
||||
)
|
||||
self.assertEqual(
|
||||
pipeline.try_lock_anchor(req.cache_request_handle, 0), ("locked", 2)
|
||||
)
|
||||
anchor = pipeline.anchor_locks[req.cache_request_handle]
|
||||
cache.tree_core.match_full_device_prefix.reset_mock()
|
||||
for attempt in range(1, 4):
|
||||
req.prefix_indices = torch.arange(2)
|
||||
self.assertTrue(cache.buffer_pipeline.prepare_staged_prefetch(req))
|
||||
self.assertIsNone(
|
||||
cache.init_load_back(
|
||||
InitLoadBackParams(
|
||||
best_match_node=None,
|
||||
host_hit_length=req.host_hit_length,
|
||||
req=req,
|
||||
)
|
||||
)
|
||||
)
|
||||
self.assertIs(
|
||||
pipeline.staged_prefetches[req.cache_request_handle], held
|
||||
)
|
||||
self.assertIs(
|
||||
pipeline.anchor_locks[req.cache_request_handle], anchor
|
||||
)
|
||||
self.assertEqual(pipeline.anchor_locked_tokens_, 2)
|
||||
self.assertEqual(
|
||||
cache.storage_prefetch_retries.pop_ready([req], 0, 8), []
|
||||
)
|
||||
self.assertEqual(
|
||||
cache.tree_core.match_full_device_prefix.call_count, attempt
|
||||
)
|
||||
self.assertEqual(
|
||||
cache.tree_core.collect_full_device_indices.call_count, attempt
|
||||
)
|
||||
cache.cache_controller.mem_pool_host.free.assert_not_called()
|
||||
self.assertEqual(cache.cache_controller.prefetch_tokens_occupied, 6)
|
||||
cache.tree_core.dec_full_pin.assert_not_called()
|
||||
pipeline.release_staged_hold(req.cache_request_handle)
|
||||
cache.tree_core.dec_full_pin.assert_called_once_with(anchor.node_id)
|
||||
self.assertEqual(pipeline.anchor_locked_tokens_, 0)
|
||||
|
||||
def test_next_pass_replans_growth_and_refetches_anchor_loss_once(self):
|
||||
cache, pipeline, req = _staged_fixture()
|
||||
self.assertTrue(cache.buffer_pipeline.prepare_staged_prefetch(req))
|
||||
self.assertEqual((req.host_hit_length, req.swa_host_hit_length), (6, 4))
|
||||
|
||||
# Tree changes occur while queued, before the next preparation pass.
|
||||
cache.tree_core.match_full_device_prefix.return_value = (6, 1, 6)
|
||||
req.prefix_indices = torch.arange(2)
|
||||
self.assertTrue(cache.buffer_pipeline.prepare_staged_prefetch(req))
|
||||
self.assertEqual((req.host_hit_length, req.swa_host_hit_length), (2, 4))
|
||||
self.assertTrue(pipeline.has_staged(req.cache_request_handle))
|
||||
self.assertEqual(cache.storage_prefetch_retries.pop_ready([req], 0, 8), [])
|
||||
|
||||
cache.tree_core.match_full_device_prefix.return_value = (0, 0, 0)
|
||||
req.prefix_indices = torch.arange(0)
|
||||
self.assertFalse(cache.buffer_pipeline.prepare_staged_prefetch(req))
|
||||
self.assertFalse(pipeline.has_staged(req.cache_request_handle))
|
||||
self.assertEqual(
|
||||
cache.storage_prefetch_retries.pop_ready([req], 0, 8), [(req, 8)]
|
||||
)
|
||||
self.assertEqual(cache.storage_prefetch_retries.pop_ready([req], 0, 8), [])
|
||||
self.assertEqual(cache.cache_controller.prefetch_tokens_occupied, 0)
|
||||
cache.cache_controller.load.assert_not_called()
|
||||
|
||||
def test_retry_budget_bounds_reissues_and_paces_capacity_misses(self):
|
||||
"""Past --hicache-storage-prefetch-retry-max-attempts a request stops
|
||||
re-issuing; a rate-limited cache-mode query is paced, not re-issued."""
|
||||
retries = StoragePrefetchRetries()
|
||||
head = SimpleNamespace(rid="head", storage_prefetch_retry_attempts=0)
|
||||
req = SimpleNamespace(rid="r", storage_prefetch_retry_attempts=8)
|
||||
retries.refetch(req.rid, 8)
|
||||
self.assertEqual(retries.pop_ready([head, req], 0, 8), [])
|
||||
req.storage_prefetch_retry_attempts = 7
|
||||
retries.refetch(req.rid, 8)
|
||||
self.assertEqual(retries.pop_ready([head, req], 0, 8), [(req, 8)])
|
||||
# A paced retry yields to the queue head; an immediate one is re-issued.
|
||||
retries.poll_miss(head.rid)
|
||||
retries.refetch(req.rid, 8)
|
||||
self.assertEqual(retries.pop_ready([head, req], 2, 8), [(req, 8)])
|
||||
|
||||
cache, _, req = _staged_fixture()
|
||||
cache.host_memory_mode = "cache"
|
||||
cache.buffer_pipeline = None
|
||||
cache.cache_controller.prefetch_rate_limited = lambda: True
|
||||
tokens = array("q", range(10))
|
||||
cache.prefetch_from_storage(
|
||||
req.cache_request_handle,
|
||||
0,
|
||||
tokens[2:],
|
||||
matched_prefix_tokens=tokens[:2],
|
||||
storage_hit_end=10,
|
||||
)
|
||||
self.assertEqual(cache.ongoing_prefetch, {})
|
||||
retries = cache.storage_prefetch_retries
|
||||
self.assertEqual(retries.pop_ready([head, req], 2, 8), [])
|
||||
self.assertEqual(retries.pop_ready([head, req], 2, 8), [])
|
||||
self.assertEqual(retries.pop_ready([head, req], 2, 8), [(req, 10)])
|
||||
|
||||
def test_staged_hold_drops_after_bounded_admission_deferrals(self):
|
||||
"""A hold that cannot be materialized after max_staged_admission_defers
|
||||
passes is released, and the request re-plans without a new L3 query."""
|
||||
cache, pipeline, req = _staged_fixture()
|
||||
cache.cache_controller.load.return_value = None
|
||||
pipeline.max_staged_admission_defers = 3
|
||||
params = lambda: InitLoadBackParams(None, req.host_hit_length, req=req)
|
||||
for attempt in range(1, 4):
|
||||
req.prefix_indices = torch.arange(2)
|
||||
self.assertTrue(pipeline.prepare_staged_prefetch(req))
|
||||
self.assertIsNone(cache.init_load_back(params()))
|
||||
self.assertEqual(pipeline.has_staged(req.cache_request_handle), attempt < 3)
|
||||
self.assertEqual(cache.cache_controller.prefetch_tokens_occupied, 0)
|
||||
self.assertEqual(cache.storage_prefetch_retries.pop_ready([req], 0, 8), [])
|
||||
self.assertTrue(pipeline.prepare_staged_prefetch(req))
|
||||
self.assertEqual((req.staged_prefetch_plan, req.storage_hit_length), (None, 0))
|
||||
|
||||
def test_controller_terminated_query_counts_as_an_l3_miss(self):
|
||||
"""A query the controller terminated (store miss or short hit) must feed
|
||||
the L3-miss counters, or a store that lost pages reads as zero misses."""
|
||||
cache = _hit_drain_fixture()
|
||||
_terminated_query(cache, "miss", hit_tokens=0)
|
||||
_terminated_query(cache, "short", hit_tokens=2)
|
||||
cache._drain_storage_control_queues_impl(
|
||||
n_storage_hit=2,
|
||||
n_ack_prefetch=0,
|
||||
n_backup=0,
|
||||
n_release=0,
|
||||
extra_release_counts={},
|
||||
log_metrics=False,
|
||||
)
|
||||
stats = cache._prefetch_outcome_stats
|
||||
self.assertEqual(
|
||||
(
|
||||
stats["revoked_full_miss"],
|
||||
stats["revoked_insufficient"],
|
||||
stats["l3_miss_tokens"],
|
||||
),
|
||||
(1, 1, 14),
|
||||
)
|
||||
self.assertEqual(cache.revoke_pending_prefetch.call_count, 2)
|
||||
head = SimpleNamespace(rid="head", storage_prefetch_retry_attempts=0)
|
||||
req = SimpleNamespace(rid="miss", storage_prefetch_retry_attempts=0)
|
||||
retries = cache.storage_prefetch_retries
|
||||
self.assertEqual(retries.pop_ready([head, req], 1, 8), [])
|
||||
self.assertEqual(retries.pop_ready([head, req], 1, 8), [(req, None)])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -214,6 +214,7 @@ class TestHiCacheMetrics(unittest.TestCase):
|
||||
|
||||
collector.log_storage_prefetch_hit_tokens(21)
|
||||
collector.log_storage_prefetch_unfulfilled_tokens(4, "storage_transfer")
|
||||
collector.log_storage_prefetch_deferred_tokens(7, "device_capacity")
|
||||
|
||||
self.assertEqual(
|
||||
collector.storage_prefetch_hit_tokens_total.increments, [(labels, 21)]
|
||||
@@ -222,6 +223,10 @@ class TestHiCacheMetrics(unittest.TestCase):
|
||||
collector.storage_prefetch_unfulfilled_tokens_total.increments,
|
||||
[({**labels, "reason": "storage_transfer"}, 4)],
|
||||
)
|
||||
self.assertEqual(
|
||||
collector.storage_prefetch_deferred_tokens_total.increments,
|
||||
[({**labels, "reason": "device_capacity"}, 7)],
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user