[Refactor] Simplify HMA metrics completion tracking

Author: dante159753

test/test_ucm_connector_metrics.py

@@ -1049,8 +1049,6 @@ def _new_hma_connector(metadata, fa_store=None, wa_store=None):
     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
@@ -1111,7 +1109,7 @@ def test_hma_connector_load_metrics_ignore_wait_failure_for_counting(monkeypatch
 
 
 @pytest.mark.parametrize("drain_method", ["get_finished", "preemption"])
-def test_hma_connector_records_grouped_save_completion_metrics(
+def test_hma_connector_records_mixed_fa_wa_save_completion_metrics(
     monkeypatch, drain_method
 ):
     _reset_fakes()
@@ -1147,23 +1145,32 @@ def test_hma_connector_records_grouped_save_completion_metrics(
             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 fake_ucmmetrics.updated[1] == {
+        "save_duration": pytest.approx(110.0),
+        "save_speed": pytest.approx(200 / 110 / 1024 / 1024),
+        "save_completion_wait_duration": pytest.approx(20.0),
+        "save_completion_wait_speed": pytest.approx(200 / 20 / 1024 / 1024),
+    }
+    assert fake_ucmmetrics.updated[2] == {
+        "connector_task_load_stage_load_dump_ms": 10.0,
+        "connector_task_compute_stage_load_dump_ms": pytest.approx(100.0),
+        "connector_task_dump_stage_load_dump_ms": pytest.approx(110.0),
+        "connector_task_total_load_dump_ms": pytest.approx(220.0),
+    }
+    assert fake_ucmmetrics.updated[3] == {
+        "save_duration": pytest.approx(130.0),
+        "save_speed": pytest.approx(40 / 130 / 1024 / 1024),
+        "save_completion_wait_duration": pytest.approx(20.0),
+        "save_completion_wait_speed": pytest.approx(40 / 20 / 1024 / 1024),
+    }
+    assert fake_ucmmetrics.updated[4] == {
+        "connector_task_load_stage_load_dump_ms": 10.0,
+        "connector_task_compute_stage_load_dump_ms": pytest.approx(100.0),
+        "connector_task_dump_stage_load_dump_ms": pytest.approx(130.0),
+        "connector_task_total_load_dump_ms": pytest.approx(240.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(

ucm/integration/vllm/hma_connector.py

@@ -268,20 +268,11 @@ class FAWADumpTask:
     task: Task
     key_count: int
     event_handle: int
-    metrics_group_id: int
     byte_count: int = 0
     submit_start_ms: float = 0.0
-
-
-@dataclass
-class FAWADumpGroupMetrics:
-    save_requests: int
-    save_blocks: int
-    save_bytes: int
-    submit_start_ms: float
-    direct_task_is_load: bool
-    direct_task_load_duration_ms: float
-    direct_task_compute_duration_ms: float
+    direct_task_is_load: bool = False
+    direct_task_load_duration_ms: float = 0.0
+    direct_task_compute_duration_ms: float = 0.0
 
 
 class UCMFAWAConnector(UCMDirectConnector, SupportsHMA):
@@ -318,8 +309,6 @@ def __init__(
         self.requests_meta: dict[str, FAWARequestMeta] = {}
         self.tp_dump_tasks: dict[tuple, list[FAWADumpTask]] = {}
         self._fawa_store_block_sizes: dict[str, int] = {}
-        self._fawa_dump_group_metrics: dict[int, FAWADumpGroupMetrics] = {}
-        self._fawa_next_dump_group_id = 0
         self.wa_dump_block_wise = self.launch_config.get("wa_dump_block_wise", True)
 
         if role == KVConnectorRole.SCHEDULER:
@@ -956,7 +945,9 @@ def _submit_dump_task(
         keys: list[bytes],
         ptrs: np.ndarray,
         event_handle: int,
-        metrics_group_id: int,
+        direct_task_is_load: bool,
+        direct_task_load_duration_ms: float,
+        direct_task_compute_duration_ms: float,
     ) -> FAWADumpTask:
         """Submit one store dump for FA or WA rows."""
 
@@ -969,9 +960,11 @@ def _submit_dump_task(
             task=task,
             key_count=len(keys),
             event_handle=event_handle,
-            metrics_group_id=metrics_group_id,
             byte_count=len(keys) * self._fawa_store_block_sizes.get(label, 0),
             submit_start_ms=submit_start_ms,
+            direct_task_is_load=direct_task_is_load,
+            direct_task_load_duration_ms=direct_task_load_duration_ms,
+            direct_task_compute_duration_ms=direct_task_compute_duration_ms,
         )
 
     def _extract_fa_ptr(self, store_keys, hash_start, hash_end, candidate_vllm_ids):
@@ -1234,10 +1227,13 @@ def wait_for_save(self) -> None:
                         )
                     )
 
-        metrics_group_id = -1
-        if fa_dump_keys or wa_dump_keys:
-            metrics_group_id = self._fawa_next_dump_group_id
-            self._fawa_next_dump_group_id += 1
+        direct_task_is_load = getattr(self, "_direct_task_has_load", False)
+        direct_task_load_duration_ms = getattr(
+            self, "_direct_task_load_duration_ms", 0.0
+        )
+        direct_task_compute_duration_ms = self._current_compute_duration_ms(
+            compute_end_time
+        )
 
         if fa_dump_keys:
             event_handle = self._get_dump_event_handle()
@@ -1251,7 +1247,9 @@ def wait_for_save(self) -> None:
                     fa_dump_keys,
                     fa_ptrs,
                     event_handle,
-                    metrics_group_id,
+                    direct_task_is_load,
+                    direct_task_load_duration_ms,
+                    direct_task_compute_duration_ms,
                 )
                 self.tp_dump_tasks[dump_request_ids].append(dump_task)
                 submitted_dump_tasks.append(dump_task)
@@ -1270,7 +1268,9 @@ def wait_for_save(self) -> None:
                     wa_dump_keys,
                     window_ptrs,
                     event_handle,
-                    metrics_group_id,
+                    direct_task_is_load,
+                    direct_task_load_duration_ms,
+                    direct_task_compute_duration_ms,
                 )
                 self.tp_dump_tasks[dump_request_ids].append(dump_task)
                 submitted_dump_tasks.append(dump_task)
@@ -1296,19 +1296,6 @@ def wait_for_save(self) -> None:
                 "save_bytes_total": save_bytes,
             }
         )
-        self._fawa_dump_group_metrics[metrics_group_id] = FAWADumpGroupMetrics(
-            save_requests=save_requests,
-            save_blocks=save_blocks,
-            save_bytes=save_bytes,
-            submit_start_ms=min(task.submit_start_ms for task in submitted_dump_tasks),
-            direct_task_is_load=getattr(self, "_direct_task_has_load", False),
-            direct_task_load_duration_ms=getattr(
-                self, "_direct_task_load_duration_ms", 0.0
-            ),
-            direct_task_compute_duration_ms=self._current_compute_duration_ms(
-                compute_end_time
-            ),
-        )
         self._clear_direct_connector_task_state()
 
     def _drain_best_effort_dump_tasks(self, finished_req_ids: set[str]) -> None:
@@ -1322,77 +1309,57 @@ def _drain_best_effort_dump_tasks(self, finished_req_ids: set[str]) -> None:
             return
 
         finished_chunk_req_ids = []
-        completed_group_ids: set[int] = set()
-        failed_group_ids: set[int] = set()
-        group_wait_start_ms: dict[int, float] = {}
-        group_wait_end_ms: dict[int, float] = {}
         for request_ids, dump_tasks in self.tp_dump_tasks.items():
             if finished_req_ids.intersection(request_ids):
                 finished_chunk_req_ids.append(request_ids)
                 for dump_task in dump_tasks:
-                    group_id = dump_task.metrics_group_id
                     wait_start_ms = time.perf_counter() * 1000
                     try:
                         dump_task.store.wait(dump_task.task)
                     except Exception as e:
-                        failed_group_ids.add(group_id)
                         logger.error(
                             "Best-effort FAWA dump task failed; external cache may miss. "
                             f"label={dump_task.label}, keys={dump_task.key_count}, "
                             f"{type(e).__name__}: {e}"
                         )
-                    finally:
+                    else:
                         wait_end_ms = time.perf_counter() * 1000
-                        completed_group_ids.add(group_id)
-                        group_wait_start_ms[group_id] = min(
-                            group_wait_start_ms.get(group_id, wait_start_ms),
-                            wait_start_ms,
+                        wait_duration_ms = self._non_negative_ms(
+                            wait_end_ms - wait_start_ms
+                        )
+                        async_duration_ms = (
+                            wait_end_ms - dump_task.submit_start_ms
+                            if dump_task.submit_start_ms > 0
+                            else wait_duration_ms
+                        )
+                        async_duration_ms = self._non_negative_ms(async_duration_ms)
+                        ucmmetrics.update_stats(
+                            {
+                                "save_duration": async_duration_ms,
+                                "save_speed": self._bytes_per_ms_to_gb_s(
+                                    dump_task.byte_count,
+                                    async_duration_ms,
+                                ),
+                                "save_completion_wait_duration": wait_duration_ms,
+                                "save_completion_wait_speed": self._bytes_per_ms_to_gb_s(
+                                    dump_task.byte_count,
+                                    wait_duration_ms,
+                                ),
+                            }
                         )
-                        group_wait_end_ms[group_id] = max(
-                            group_wait_end_ms.get(group_id, wait_end_ms),
-                            wait_end_ms,
+                        self._record_direct_connector_task_metrics(
+                            is_load=dump_task.direct_task_is_load,
+                            is_save=True,
+                            load_duration_ms=dump_task.direct_task_load_duration_ms,
+                            compute_duration_ms=dump_task.direct_task_compute_duration_ms,
+                            dump_duration_ms=async_duration_ms,
                         )
+                    finally:
                         self.device.destroy_event_handle(dump_task.event_handle)
 
         for request_ids in finished_chunk_req_ids:
             self.tp_dump_tasks.pop(request_ids, None)
 
-        for group_id in completed_group_ids:
-            group_metrics = self._fawa_dump_group_metrics.pop(group_id, None)
-            if group_metrics is None or group_id in failed_group_ids:
-                continue
-            wait_duration_ms = (
-                group_wait_end_ms[group_id] - group_wait_start_ms[group_id]
-            )
-            async_duration_ms = (
-                group_wait_end_ms[group_id] - group_metrics.submit_start_ms
-                if group_metrics.submit_start_ms > 0
-                else wait_duration_ms
-            )
-            async_duration_ms = self._non_negative_ms(async_duration_ms)
-            wait_duration_ms = self._non_negative_ms(wait_duration_ms)
-            ucmmetrics.update_stats(
-                {
-                    "save_duration": async_duration_ms,
-                    "save_speed": self._bytes_per_ms_to_gb_s(
-                        group_metrics.save_bytes,
-                        async_duration_ms,
-                    ),
-                    "save_completion_wait_duration": wait_duration_ms,
-                    "save_completion_wait_speed": self._bytes_per_ms_to_gb_s(
-                        group_metrics.save_bytes,
-                        wait_duration_ms,
-                    ),
-                }
-            )
-            self._record_direct_connector_task_metrics(
-                is_load=group_metrics.direct_task_is_load,
-                is_save=True,
-                load_duration_ms=group_metrics.direct_task_load_duration_ms,
-                compute_duration_ms=group_metrics.direct_task_compute_duration_ms,
-                dump_duration_ms=async_duration_ms,
-            )
-
     def handle_preemptions(self, kv_connector_metadata: UCMFAWAConnectorMetadata):
         # Worker side method
         self._drain_best_effort_dump_tasks(kv_connector_metadata.preempted_req_ids)