Trace P0 batch-level DP coordination

scripts/collectivespec/p0_overlay/collectivespec_p0/worker.py

@@ -15,6 +15,38 @@ class P0Worker(Worker):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
         self._p0_execute_epoch = 0
+        self._p0_batch_epoch = 0
+
+    def init_device(self):  # type: ignore[no-untyped-def]
+        result = super().init_device()
+        model_runner = self.model_runner
+        if getattr(model_runner, "_collectivespec_p0_wrapped", False):
+            return result
+        original = model_runner._determine_batch_execution_and_padding
+
+        def traced_determine(*args: Any, **kwargs: Any):  # type: ignore[no-untyped-def]
+            output = original(*args, **kwargs)
+            self._p0_batch_epoch += 1
+            scheduler_output = kwargs.get("scheduler_output")
+            if scheduler_output is None and args:
+                scheduler_output = args[0]
+            cudagraph_mode, batch_desc, should_ubatch, rows_across_dp, _ = output
+            log_event(
+                "worker",
+                "batch_execution_plan",
+                batch_phase_epoch=self._p0_batch_epoch,
+                **self._p0_rank_fields(),
+                **plan_summary(scheduler_output),
+                cudagraph_mode=getattr(cudagraph_mode, "value", str(cudagraph_mode)),
+                physical_batch_rows=getattr(batch_desc, "num_tokens", None),
+                should_ubatch=should_ubatch,
+                rows_across_dp=rows_across_dp,
+            )
+            return output
+
+        model_runner._determine_batch_execution_and_padding = traced_determine
+        model_runner._collectivespec_p0_wrapped = True
+        return result
 
     def _p0_rank_fields(self) -> dict[str, Any]:
         config = self.vllm_config.parallel_config

scripts/collectivespec/run_p0_phase_trace.py

@@ -200,6 +200,7 @@ def main() -> int:
         "overlay": str(overlay),
         "run_id": args.run_id,
         "topology_claim": "inherited_from_base_spec",
+        "gpu_count": base.get("hardware", {}).get("gpu_count"),
         "request_count": args.request_count,
         "completion_tokens": args.completion_tokens,
         "policies": ["control_k3", "heterogeneous"],

scripts/collectivespec/summarize_p0_phase_trace.py

@@ -97,19 +97,26 @@ def probe_integrity(cell: Path, expected_requests: int, expected_tokens: int) ->
     }
 
 
+def rank_key(event: dict[str, Any]) -> str | None:
+    dp_rank = event.get("data_parallel_rank")
+    local_rank = event.get("local_rank")
+    if dp_rank is None or local_rank is None:
+        return None
+    return f"dp{dp_rank}/local{local_rank}"
+
+
 def sequence_agreement(worker_events: list[dict[str, Any]]) -> dict[str, Any]:
-    begins = [event for event in worker_events if event.get("event") == "target_execute_begin"]
+    begins = [event for event in worker_events if event.get("event") == "batch_execution_plan"]
     per_rank: dict[str, list[dict[str, Any]]] = defaultdict(list)
     per_dp: dict[str, set[str]] = defaultdict(set)
     for event in begins:
-        rank = event.get("global_rank")
+        rank = rank_key(event)
         dp_rank = event.get("data_parallel_rank")
         if rank is None:
             continue
-        key = str(rank)
+        per_rank[rank].append(event)
-        per_rank[key].append(event)
         if dp_rank is not None:
-            per_dp[str(dp_rank)].add(key)
+            per_dp[str(dp_rank)].add(rank)
     sequences: dict[str, list[str]] = {}
     for rank, values in per_rank.items():
         values.sort(key=lambda item: (integer(item.get("worker_phase_epoch")) or -1, float(item.get("monotonic_s", 0.0))))
@@ -117,7 +124,7 @@ def sequence_agreement(worker_events: list[dict[str, Any]]) -> dict[str, Any]:
     counts = {rank: len(values) for rank, values in sequences.items()}
     within_dp: dict[str, Any] = {}
     for dp_rank, ranks in per_dp.items():
-        rank_list = sorted(ranks, key=int)
+        rank_list = sorted(ranks)
         distinct = {tuple(sequences[rank]) for rank in rank_list}
         within_dp[dp_rank] = {
             "ranks": rank_list,
@@ -129,7 +136,7 @@ def sequence_agreement(worker_events: list[dict[str, Any]]) -> dict[str, Any]:
     return {
         "worker_begin_event_count": len(begins),
         "worker_rank_count": len(sequences),
-        "phase_count_by_global_rank": counts,
+        "phase_count_by_rank": counts,
         "phase_count": {
             "n": len(count_values),
             "min": min(count_values) if count_values else None,
@@ -152,12 +159,13 @@ def summarize_cell(cell: Path, expected_requests: int, expected_tokens: int) ->
     verify_schedule_events = [
         event for event in schedule_events if integer(event.get("spec_request_count")) not in (None, 0)
     ]
-    worker_metadata = [
+    batch_events = [event for event in worker_events if event.get("event") == "batch_execution_plan"]
-        event.get("scheduled_dp_metadata")
+    rows_across_dp = [event.get("rows_across_dp") for event in batch_events]
-        for event in worker_events
+    physical_rows = [
-        if event.get("event") == "target_execute_begin"
+        integer(event.get("physical_batch_rows"))
+        for event in batch_events
+        if integer(event.get("physical_batch_rows")) is not None
     ]
-    metadata_present = sum(value is not None for value in worker_metadata)
     return {
         "probe_integrity": probe_integrity(cell, expected_requests, expected_tokens),
         "event_parse_errors": parse_errors,
@@ -179,13 +187,20 @@ def summarize_cell(cell: Path, expected_requests: int, expected_tokens: int) ->
         },
         "worker": {
             **sequence_agreement(worker_events),
-            "dp_metadata_present_count": metadata_present,
+            "batch_execution_plan_count": len(batch_events),
-            "dp_metadata_missing_count": len(worker_metadata) - metadata_present,
+            "rows_across_dp_present_count": sum(value is not None for value in rows_across_dp),
+            "physical_batch_rows": {
+                "n": len(physical_rows),
+                "min": min(physical_rows) if physical_rows else None,
+                "max": max(physical_rows) if physical_rows else None,
+                "distinct_value_count": len(set(physical_rows)),
+                "non_negative": all(value >= 0 for value in physical_rows),
+            },
         },
     }
 
 
-def data_sanity(cells: dict[str, dict[str, Any]]) -> dict[str, Any]:
+def data_sanity(cells: dict[str, dict[str, Any]], expected_gpu_count: int | None) -> dict[str, Any]:
     event_counts = [integer(cell.get("event_count")) or 0 for cell in cells.values()]
     candidate_distinct = [
         integer(cell.get("scheduler", {}).get("after_k_distinct_value_count")) or 0
@@ -220,6 +235,20 @@ def data_sanity(cells: dict[str, dict[str, Any]]) -> dict[str, Any]:
                 bool(cell.get("worker", {}).get("all_within_dp_sequences_identical"))
                 for cell in cells.values()
             ),
+            "all_dp_coordination_records_observed": all(
+                integer(cell.get("worker", {}).get("rows_across_dp_present_count"))
+                == integer(cell.get("worker", {}).get("batch_execution_plan_count"))
+                and integer(cell.get("worker", {}).get("batch_execution_plan_count", 0)) > 0
+                for cell in cells.values()
+            ),
+            "all_expected_workers_observed": (
+                all(
+                    integer(cell.get("worker", {}).get("worker_rank_count")) == expected_gpu_count
+                    for cell in cells.values()
+                )
+                if expected_gpu_count is not None
+                else None
+            ),
         },
     }
 
@@ -282,7 +311,11 @@ def main() -> int:
         policy: summarize_cell(args.root / policy, expected_requests, expected_tokens)
         for policy in policies
     }
-    payload = {"manifest": manifest, "cells": cells, "data_sanity": data_sanity(cells)}
+    payload = {
+        "manifest": manifest,
+        "cells": cells,
+        "data_sanity": data_sanity(cells, integer(manifest.get("gpu_count"))),
+    }
     args.output_json.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")
     args.output_md.write_text(markdown(payload), encoding="utf-8")
     print(json.dumps(payload, indent=2, sort_keys=True))