[Feat] Add HMA connector metrics

Author: dante159753

test/test_ucm_connector_metrics.py

@@ -275,10 +275,26 @@ def get_config(self):
 
 
 def _install_stubs():
+    _install_module(
+        "numpy",
+        ndarray=object,
+        bool_=bool,
+        uint64="uint64",
+        int64="int64",
+        zeros=lambda shape, *args, **kwargs: [
+            [0 for _ in range(shape[1])] for _ in range(shape[0])
+        ],
+        vstack=lambda rows: list(rows),
+        concatenate=lambda arrays, axis=0: [
+            item for array in arrays for item in list(array)
+        ],
+        asarray=lambda values, dtype=None: values,
+        arange=lambda *args, **kwargs: list(range(*args)),
+        isscalar=lambda value: isinstance(value, (int, float, bool, str, bytes)),
+    )
     _install_package("ucm", REPO_ROOT / "ucm")
     _install_package("ucm.integration", REPO_ROOT / "ucm" / "integration")
     _install_package("ucm.integration.vllm", REPO_ROOT / "ucm" / "integration" / "vllm")
-    _install_module("numpy", ndarray=type("ndarray", (), {}))
     _install_module("torch", Tensor=type("Tensor", (), {}))
     _install_module(
         "prometheus_client",
@@ -354,7 +370,13 @@ def _install_stubs():
 
 _install_stubs()
 
+import ucm.integration.vllm.hma_connector as hma_connector_module
 import ucm.integration.vllm.ucm_connector as ucm_connector_module
+from ucm.integration.vllm.hma_connector import (
+    FAWARequestDispatchMeta,
+    UCMFAWAConnector,
+    UCMFAWAConnectorMetadata,
+)
 from ucm.integration.vllm.metrics import UCMConnectorStats, UCMPromMetrics
 from ucm.integration.vllm.ucm_connector import (
     PendingDumpTask,
@@ -460,6 +482,38 @@ def _reset_fakes():
     dispatcher_module._DISPATCHER = None
 
 
+class FakeHMAStore:
+    def __init__(self, label, fail_wait=False):
+        self.label = label
+        self.fail_wait = fail_wait
+        self.loaded = []
+        self.dumped = []
+        self.waited = []
+
+    def load_data(self, keys, shard_indices, ptrs):
+        task = f"{self.label}-load-{len(keys)}-{len(self.loaded)}"
+        self.loaded.append((list(keys), list(shard_indices), ptrs, task))
+        return task
+
+    def dump_data(self, keys, shard_indices, ptrs, event_handle):
+        task = f"{self.label}-dump-{len(keys)}-{len(self.dumped)}"
+        self.dumped.append((list(keys), list(shard_indices), ptrs, event_handle, task))
+        return task
+
+    def wait(self, task):
+        self.waited.append(task)
+        if self.fail_wait:
+            raise RuntimeError(f"{self.label} wait failed")
+
+
+class FakeHMADevice:
+    def __init__(self):
+        self.destroyed = []
+
+    def destroy_event_handle(self, event_handle):
+        self.destroyed.append(event_handle)
+
+
 def test_config_definitions_register_enable_list_and_metric_names():
     _reset_fakes()
     config = _metrics_config()
@@ -969,6 +1023,212 @@ def destroy_event_handle(self, event_handle):
     ]
 
 
+def _fawa_request(load_keys=None, dump_keys=None):
+    return FAWARequestDispatchMeta(
+        load_keys=list(load_keys or []),
+        load_hash_start=0,
+        load_hash_end=len(load_keys or []),
+        load_vllm_block_ids=([1, 2],),
+        dump_keys=list(dump_keys or []),
+        dump_hash_start=0,
+        dump_hash_end=len(dump_keys or []),
+        dump_vllm_block_ids=([1, 2],),
+    )
+
+
+def _new_hma_connector(metadata, fa_store=None, wa_store=None):
+    connector = object.__new__(UCMFAWAConnector)
+    connector.fa_store = fa_store or FakeHMAStore("FA")
+    connector.wa_store = wa_store or FakeHMAStore("WA")
+    connector._fawa_store_block_sizes = {"FA": 100, "WA": 20}
+    connector._invalid_block_ids = set()
+    connector.tp_size = 1
+    connector.tp_rank = 0
+    connector.wa_dump_block_wise = True
+    connector.fa_group_ids = set()
+    connector.window_group_ids = set()
+    connector.group_metas = []
+    connector.tp_dump_tasks = {}
+    connector._fawa_dump_group_metrics = {}
+    connector._fawa_next_dump_group_id = 0
+    connector._direct_task_has_load = False
+    connector._direct_task_load_duration_ms = 0.0
+    connector._direct_task_compute_start_ms = None
+    connector.device = FakeHMADevice()
+    connector._get_connector_metadata = lambda: metadata
+    connector._extract_fa_ptr = lambda *args, **kwargs: [[0], [0]]
+    connector._extract_wa_ptr = lambda *args, **kwargs: [[0], [0]]
+    return connector
+
+
+def test_hma_connector_records_load_metrics(monkeypatch):
+    _reset_fakes()
+    times = iter([1.0, 1.1, 1.11])
+    monkeypatch.setattr(
+        hma_connector_module.time,
+        "perf_counter",
+        lambda: next(times),
+    )
+    metadata = UCMFAWAConnectorMetadata(
+        request_meta={"req-1": _fawa_request(load_keys=[b"a", b"b"])}
+    )
+    connector = _new_hma_connector(metadata)
+
+    connector.start_load_kv(None)
+
+    stats = fake_ucmmetrics.updated[-1]
+    assert stats["load_requests_num"] == 1
+    assert stats["load_blocks_num"] == 3
+    assert stats["load_duration"] == pytest.approx(100.0)
+    assert stats["load_bytes_total"] == 220
+    assert stats["load_speed"] == pytest.approx(220 / 100 / 1024 / 1024)
+    assert connector._direct_task_has_load is True
+    assert connector._direct_task_load_duration_ms == pytest.approx(100.0)
+    assert connector._direct_task_compute_start_ms == pytest.approx(1110.0)
+
+
+def test_hma_connector_load_metrics_ignore_wait_failure_for_counting(monkeypatch):
+    _reset_fakes()
+    times = iter([1.0, 1.1, 1.11])
+    monkeypatch.setattr(
+        hma_connector_module.time,
+        "perf_counter",
+        lambda: next(times),
+    )
+    metadata = UCMFAWAConnectorMetadata(
+        request_meta={"req-1": _fawa_request(load_keys=[b"a", b"b"])}
+    )
+    connector = _new_hma_connector(metadata, wa_store=FakeHMAStore("WA", True))
+
+    connector.start_load_kv(None)
+
+    stats = fake_ucmmetrics.updated[-1]
+    assert stats["load_requests_num"] == 1
+    assert stats["load_blocks_num"] == 3
+    assert stats["load_bytes_total"] == 220
+    assert stats["load_speed"] == pytest.approx(220 / 100 / 1024 / 1024)
+    assert connector._invalid_block_ids == {1, 2}
+
+
+@pytest.mark.parametrize("drain_method", ["get_finished", "preemption"])
+def test_hma_connector_records_grouped_save_completion_metrics(
+    monkeypatch, drain_method
+):
+    _reset_fakes()
+    times = iter([1.0, 1.01, 1.02, 1.10, 1.12, 1.13, 1.15])
+    monkeypatch.setattr(
+        hma_connector_module.time,
+        "perf_counter",
+        lambda: next(times),
+    )
+    metadata = UCMFAWAConnectorMetadata(
+        request_meta={"req-1": _fawa_request(dump_keys=[b"a", b"b"])}
+    )
+    connector = _new_hma_connector(metadata)
+    handles = iter([11, 12])
+    connector._get_dump_event_handle = lambda: next(handles)
+    connector._direct_task_has_load = True
+    connector._direct_task_load_duration_ms = 10.0
+    connector._direct_task_compute_start_ms = 900.0
+
+    connector.wait_for_save()
+
+    assert fake_ucmmetrics.updated == [
+        {
+            "save_requests_num": 1,
+            "save_blocks_num": 4,
+            "save_bytes_total": 240,
+        }
+    ]
+    if drain_method == "get_finished":
+        connector.get_finished({"req-1"})
+    else:
+        connector.handle_preemptions(
+            UCMFAWAConnectorMetadata(preempted_req_ids={"req-1"})
+        )
+
+    completion_stats = fake_ucmmetrics.updated[-2]
+    assert completion_stats["save_duration"] == pytest.approx(140.0)
+    assert completion_stats["save_speed"] == pytest.approx(240 / 140 / 1024 / 1024)
+    assert completion_stats["save_completion_wait_duration"] == pytest.approx(50.0)
+    assert completion_stats["save_completion_wait_speed"] == pytest.approx(
+        240 / 50 / 1024 / 1024
+    )
+    task_stats = fake_ucmmetrics.updated[-1]
+    assert task_stats["connector_task_load_stage_load_dump_ms"] == 10.0
+    assert task_stats["connector_task_compute_stage_load_dump_ms"] == pytest.approx(
+        100.0
+    )
+    assert task_stats["connector_task_dump_stage_load_dump_ms"] == pytest.approx(140.0)
+    assert task_stats["connector_task_total_load_dump_ms"] == pytest.approx(250.0)
+    assert connector.device.destroyed == [11, 12]
+    assert connector.tp_dump_tasks == {}
+    assert connector._fawa_dump_group_metrics == {}
+
+
+def test_hma_connector_save_metrics_count_only_submitted_requests_without_changing_tp_keys(
+    monkeypatch,
+):
+    _reset_fakes()
+    times = iter([1.0, 1.01, 1.02])
+    monkeypatch.setattr(
+        hma_connector_module.time,
+        "perf_counter",
+        lambda: next(times),
+    )
+    metadata = UCMFAWAConnectorMetadata(
+        request_meta={
+            "req-no-rank-work": _fawa_request(dump_keys=[b"a"]),
+            "req-rank-work": _fawa_request(dump_keys=[b"b", b"c"]),
+        }
+    )
+    connector = _new_hma_connector(metadata)
+    connector.tp_size = 2
+    connector.tp_rank = 0
+    handles = iter([11, 12])
+    connector._get_dump_event_handle = lambda: next(handles)
+
+    connector.wait_for_save()
+
+    assert list(connector.tp_dump_tasks) == [("req-no-rank-work", "req-rank-work")]
+    assert fake_ucmmetrics.updated == [
+        {
+            "save_requests_num": 1,
+            "save_blocks_num": 2,
+            "save_bytes_total": 120,
+        }
+    ]
+
+
+def test_hma_connector_no_dump_records_load_only_breakdown(monkeypatch):
+    _reset_fakes()
+    times = iter([1.0])
+    monkeypatch.setattr(
+        hma_connector_module.time,
+        "perf_counter",
+        lambda: next(times),
+    )
+    metadata = UCMFAWAConnectorMetadata(
+        request_meta={"req-1": _fawa_request(load_keys=[b"a"])}
+    )
+    connector = _new_hma_connector(metadata)
+    connector._direct_task_has_load = True
+    connector._direct_task_load_duration_ms = 7.0
+    connector._direct_task_compute_start_ms = 900.0
+
+    connector.wait_for_save()
+
+    task_stats = fake_ucmmetrics.updated[-1]
+    assert task_stats["connector_task_load_stage_load_only_ms"] == 7.0
+    assert task_stats["connector_task_compute_stage_load_only_ms"] == pytest.approx(
+        100.0
+    )
+    assert task_stats["connector_task_total_load_only_ms"] == pytest.approx(107.0)
+    assert connector._direct_task_has_load is False
+    assert connector._direct_task_load_duration_ms == 0.0
+    assert connector._direct_task_compute_start_ms is None
+
+
 def test_multiproc_logger_uses_prefix_and_dispatcher_snapshot(tmp_path):
     _reset_fakes()
     import ucm.observability as observability