Observability enhancement for HiCache (#32388)

test/registered/unit/mem_cache/test_hicache_load_back_timing.py

@@ -65,6 +65,7 @@ class TestLoadBackDurationMetric(CustomTestCase):
             node_ids=[1, 2],
             num_tokens=1024,
             timing_enabled=True,
+            num_tokens_by_pool={"kv": 1024},
         )
         stub = object.__new__(HiRadixCache)
         stub.cache_controller = SimpleNamespace(ack_load_queue=[ack])
@@ -77,7 +78,7 @@ class TestLoadBackDurationMetric(CustomTestCase):
         stub.loading_check()
 
         stub.metrics_collector.increment_load_back_num_tokens.assert_called_once_with(
-            1024
+            num_tokens=1024, pool="kv"
         )
         stub.metrics_collector.observe_load_back_duration.assert_called_once()
         (observed,), _ = stub.metrics_collector.observe_load_back_duration.call_args
@@ -100,6 +101,7 @@ class TestLoadBackDurationMetric(CustomTestCase):
             node_ids=[7],
             num_tokens=512,
             timing_enabled=False,
+            num_tokens_by_pool={"kv": 512},
         )
         stub = object.__new__(HiRadixCache)
         stub.cache_controller = SimpleNamespace(ack_load_queue=[ack])
@@ -112,7 +114,7 @@ class TestLoadBackDurationMetric(CustomTestCase):
         stub.loading_check()
 
         stub.metrics_collector.increment_load_back_num_tokens.assert_called_once_with(
-            512
+            num_tokens=512, pool="kv"
         )
         stub.metrics_collector.observe_load_back_duration.assert_not_called()
         self.assertEqual(stub.cache_controller.ack_load_queue, [])

test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py

@@ -139,6 +139,9 @@ def _cpu_per_layer_pf_lf_copy(
 
 
 class _FakeEvent:
+    def __init__(self, enable_timing=False):
+        self.enable_timing = enable_timing
+
     def record(self):
         pass
 
@@ -156,6 +159,15 @@ class _FakeDeviceModule:
 
 
 class TestHiCacheStagedWriteBackDispatch(unittest.TestCase):
+    def setUp(self):
+        # start_writing probes timing support via a module-cached check;
+        # clear it on both sides so results from (or against) the fake
+        # device module never leak across tests.
+        manager_cache_controller._timing_events_supported.cache_clear()
+
+    def tearDown(self):
+        manager_cache_controller._timing_events_supported.cache_clear()
+
     def _patched_transfers(self, src_registry=None, module=MEMORY_POOL_HOST_MODULE):
         staged_side_effect = None
         if src_registry is not None:
@@ -678,6 +690,10 @@ class TestHiCacheStagedWriteBackDispatch(unittest.TestCase):
         class FakeHostGroup:
             layout = "page_first"
             can_use_write_back_jit = True
+            anchor_entry = SimpleNamespace(
+                name=PoolName.KV, host_pool=SimpleNamespace(size_per_token=2)
+            )
+            entry_map = {}
 
             def backup_from_device_all_layer(
                 self,
@@ -718,8 +734,13 @@ class TestHiCacheStagedWriteBackDispatch(unittest.TestCase):
             )
         )
 
-        with mock.patch.object(
-            hybrid_cache_controller, "device_module", _FakeDeviceModule
+        with (
+            mock.patch.object(
+                hybrid_cache_controller, "device_module", _FakeDeviceModule
+            ),
+            mock.patch.object(
+                manager_cache_controller, "device_module", _FakeDeviceModule
+            ),
         ):
             controller.start_writing()
 
@@ -733,6 +754,10 @@ class TestHiCacheStagedWriteBackDispatch(unittest.TestCase):
         class FakeHostGroup:
             layout = "page_first"
             can_use_write_back_jit = False
+            anchor_entry = SimpleNamespace(
+                name=PoolName.KV, host_pool=SimpleNamespace(size_per_token=2)
+            )
+            entry_map = {}
 
             def backup_from_device_all_layer(
                 self,
@@ -770,8 +795,13 @@ class TestHiCacheStagedWriteBackDispatch(unittest.TestCase):
             return_value=(op.host_indices, op.device_indices, op.pool_transfers)
         )
 
-        with mock.patch.object(
-            hybrid_cache_controller, "device_module", _FakeDeviceModule
+        with (
+            mock.patch.object(
+                hybrid_cache_controller, "device_module", _FakeDeviceModule
+            ),
+            mock.patch.object(
+                manager_cache_controller, "device_module", _FakeDeviceModule
+            ),
         ):
             controller.start_writing()
 
@@ -785,6 +815,7 @@ class TestHiCacheStagedWriteBackDispatch(unittest.TestCase):
         class FakeHostPool:
             layout = "page_first"
             can_use_write_back_jit = True
+            size_per_token = 2
 
             def backup_from_device_all_layer(
                 self, device_pool, host_indices, device_indices, io_backend
@@ -825,6 +856,7 @@ class TestHiCacheStagedWriteBackDispatch(unittest.TestCase):
         class FakeHostPool:
             layout = "page_first"
             can_use_write_back_jit = False
+            size_per_token = 2
 
             def backup_from_device_all_layer(
                 self, device_pool, host_indices, device_indices, io_backend