[Feat] Add UCM metrics for SGLang HiCache and Direct Linker

test/common/sglang_utils.py

@@ -10,8 +10,10 @@
 import re
 import sys
 import threading
+import time
 import types
-from dataclasses import dataclass, replace
+from contextlib import contextmanager
+from dataclasses import dataclass, field, replace
 from pathlib import Path
 from typing import Any, List
 
@@ -44,6 +46,7 @@ def load_source(relative_path, definitions=None, **dependencies):
     defaults = dict(
         dataclass=dataclass,
         replace=replace,
+        field=field,
         np=np,
         copy=copy,
         hashlib=hashlib,
@@ -64,9 +67,20 @@ def load_source(relative_path, definitions=None, **dependencies):
     return _load_source(relative_path, definitions=definitions, **defaults)
 
 
-adapter_module = load_source("ucm/integration/sglang/ucm_direct_store_adapter.py")
+metrics_module = load_source(
+    "ucm/integration/sglang/metrics.py", contextmanager=contextmanager, time=time
+)
+adapter_module = load_source(
+    "ucm/integration/sglang/ucm_direct_store_adapter.py",
+    AsyncIoObservation=metrics_module.AsyncIoObservation,
+    record_metrics=metrics_module.record_metrics,
+)
 linker_module = load_source(
     "ucm/integration/sglang/ucm_direct_linker.py",
+    LinkerOperation=metrics_module.LinkerOperation,
+    record_metrics=metrics_module.record_metrics,
+    shutdown_metrics_logger=metrics_module.shutdown_metrics_logger,
+    create_metrics_logger=metrics_module.create_metrics_logger,
     **{
         name: getattr(adapter_module, name)
         for name in ("UcmDirectStoreAdapter", "UcmLayerTransfer", "UcmPoolMetadata")

test/test_sglang_dump_partition.py

@@ -47,6 +47,7 @@ def setUp(self):
 
     def linker(self, rank, ranks=(0, 1, 2, 3)):
         obj = object.__new__(linker_module.UcmDirectLinker)
+        obj._metrics_enabled = False
         obj._tp_rank, obj._tp_size = rank, 4
         obj._dump_ranks = ranks
         obj._lookup_owner = True
@@ -243,7 +244,8 @@ def test_offload_resolves_once_and_empty_rank_records_no_events(self):
         expected = obj._partition_offload(transfers)
         self.assertTrue(obj.offload(transfers))
         self.assertEqual(len(expanded), 1)
-        queued, ready = obj._offload_queue.get_nowait()
+        queued, ready, observation = obj._offload_queue.get_nowait()
+        self.assertIsNone(observation)
         self.assertEqual([t.name for t in queued], [t.name for t in expected])
         self.assertEqual(len(events), int(bool(expected)))
         if not expected: