sgl-project/sglang

diff --git a/test/registered/unit/mem_cache/test_hicache_load_back_metrics.py b/test/registered/unit/mem_cache/test_hicache_load_back_metrics.pynew file mode 100644index 0000000..a4333b5--- /dev/null+++ b/test/registered/unit/mem_cache/test_hicache_load_back_metrics.py@@ -0,0 +1,152 @@+"""Unit tests for the HiCache load-back duration/token metrics.++These tests exercise the device-event based instrumentation that HiCache uses to+report how long a KV-cache load-back took and how many tokens it moved. They run+on any backend (including CPU-only environments) by using mock device events, so+they do not require a real GPU.+"""++import unittest+from types import SimpleNamespace+from unittest.mock import MagicMock, patch+++class TestMakeTimingEventPair(unittest.TestCase):+    def setUp(self):+        from sglang.srt.managers import cache_controller as cc++        cc._timing_events_supported.cache_clear()+        self.cc = cc++    def tearDown(self):+        self.cc._timing_events_supported.cache_clear()++    def test_timing_supported_uses_enable_timing(self):+        created = []++        def create_event(*, enable_timing=False):+            created.append(enable_timing)+            return MagicMock()++        with patch.object(self.cc.device_module, "Event", side_effect=create_event):+            self.cc._timing_events_supported.cache_clear()+            start, finish, timing_enabled = self.cc.make_timing_event_pair()++        self.assertTrue(timing_enabled)+        self.assertIsNot(start, finish)+        # Every event (probe + the returned pair) was created with enable_timing.+        self.assertTrue(created and all(created))++    def test_timing_fallback_uses_dedicated_events(self):+        events = []++        def create_event(*, enable_timing=False):+            if enable_timing:+                raise TypeError("Event() takes no arguments")+            event = MagicMock()+            events.append(event)+            return event++        with patch.object(self.cc.device_module, "Event", side_effect=create_event):+            self.cc._timing_events_supported.cache_clear()+            start, finish, timing_enabled = self.cc.make_timing_event_pair()++        self.assertFalse(timing_enabled)+        self.assertIs(start, events[0])+        self.assertIs(finish, events[1])+        self.assertIsNot(start, finish)+++class TestHiCacheAckFields(unittest.TestCase):+    def test_num_tokens_and_timing_defaults(self):+        from sglang.srt.managers.cache_controller import HiCacheAck++        ack = HiCacheAck(start_event=None, finish_event=None, node_ids=[1])+        self.assertEqual(ack.num_tokens, 0)+        self.assertFalse(ack.timing_enabled)++    def test_fields_carry_supplied_values(self):+        from sglang.srt.managers.cache_controller import HiCacheAck++        ack = HiCacheAck(+            start_event=None,+            finish_event=None,+            node_ids=[3, 4],+            num_tokens=128,+            timing_enabled=True,+        )+        self.assertEqual(ack.num_tokens, 128)+        self.assertTrue(ack.timing_enabled)+++def _make_load_stub(ack):+    return SimpleNamespace(+        cache_controller=SimpleNamespace(ack_load_queue=[ack]),+        ongoing_load_back={i: object() for i in ack.node_ids},+        dec_lock_ref=MagicMock(),+        metrics_collector=MagicMock(),+        pp_rank=0,+        _all_reduce=MagicMock(),+    )+++class TestLoadingCheckMetrics(unittest.TestCase):+    def test_loading_check_records_duration_and_tokens(self):+        from sglang.srt.managers.cache_controller import HiCacheAck+        from sglang.srt.mem_cache.hiradix_cache import HiRadixCache++        finish_event = MagicMock()+        finish_event.query.return_value = True+        start_event = MagicMock()+        start_event.elapsed_time.return_value = 2000.0  # milliseconds++        ack = HiCacheAck(+            start_event=start_event,+            finish_event=finish_event,+            node_ids=[1, 2],+            num_tokens=1024,+            timing_enabled=True,+        )+        stub = _make_load_stub(ack)++        HiRadixCache.loading_check(stub)++        stub.metrics_collector.increment_load_back_num_tokens.assert_called_once_with(+            1024+        )+        stub.metrics_collector.observe_load_back_duration.assert_called_once()+        (observed,), _ = stub.metrics_collector.observe_load_back_duration.call_args+        # Duration is reported in seconds (2000 ms -> 2.0 s).+        self.assertAlmostEqual(observed, 2.0)+        self.assertEqual(stub.cache_controller.ack_load_queue, [])+        self.assertEqual(stub.ongoing_load_back, {})++    def test_loading_check_skips_duration_when_timing_unsupported(self):+        from sglang.srt.managers.cache_controller import HiCacheAck+        from sglang.srt.mem_cache.hiradix_cache import HiRadixCache++        finish_event = MagicMock()+        finish_event.query.return_value = True+        start_event = MagicMock()++        ack = HiCacheAck(+            start_event=start_event,+            finish_event=finish_event,+            node_ids=[7],+            num_tokens=512,+            timing_enabled=False,+        )+        stub = _make_load_stub(ack)++        HiRadixCache.loading_check(stub)++        stub.metrics_collector.increment_load_back_num_tokens.assert_called_once_with(+            512+        )+        stub.metrics_collector.observe_load_back_duration.assert_not_called()+        start_event.elapsed_time.assert_not_called()+        self.assertEqual(stub.cache_controller.ack_load_queue, [])+++if __name__ == "__main__":+    unittest.main()