Correct CollectiveSpec P0 phase accounting

scripts/collectivespec/p0_overlay/collectivespec_p0/worker.py

@@ -25,6 +25,9 @@ class P0Worker(Worker):
         original = model_runner._determine_batch_execution_and_padding
 
         def traced_determine(*args: Any, **kwargs: Any):  # type: ignore[no-untyped-def]
+            local_unpadded_rows = kwargs.get("num_tokens")
+            if local_unpadded_rows is None and len(args) > 1:
+                local_unpadded_rows = args[1]
             output = original(*args, **kwargs)
             self._p0_batch_epoch += 1
             scheduler_output = kwargs.get("scheduler_output")
@@ -38,6 +41,7 @@ class P0Worker(Worker):
                 **self._p0_rank_fields(),
                 **plan_summary(scheduler_output),
                 cudagraph_mode=getattr(cudagraph_mode, "value", str(cudagraph_mode)),
+                local_unpadded_rows=local_unpadded_rows,
                 physical_batch_rows=getattr(batch_desc, "num_tokens", None),
                 should_ubatch=should_ubatch,
                 rows_across_dp=rows_across_dp,
@@ -54,7 +58,9 @@ class P0Worker(Worker):
             "global_rank": getattr(self, "rank", None),
             "local_rank": getattr(self, "local_rank", None),
             "data_parallel_rank": getattr(config, "data_parallel_rank", None),
-            "tensor_parallel_rank": getattr(config, "tensor_parallel_rank", None),
+            # Worker.rank is the tensor-parallel rank within a DP replica in
+            # this vLLM executor. ParallelConfig does not expose it here.
+            "tensor_parallel_rank": getattr(self, "rank", None),
             "expert_parallel_size": getattr(config, "expert_parallel_size", None),
             "tensor_parallel_size": getattr(config, "tensor_parallel_size", None),
             "data_parallel_size": getattr(config, "data_parallel_size", None),

scripts/collectivespec/summarize_p0_phase_trace.py

@@ -105,48 +105,264 @@ def rank_key(event: dict[str, Any]) -> str | 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") == "batch_execution_plan"]
+def _event_order(event: dict[str, Any], epoch_field: str) -> tuple[int, float]:
+    epoch = integer(event.get(epoch_field))
+    monotonic_s = event.get("monotonic_s")
+    return (
+        epoch if epoch is not None else -1,
+        float(monotonic_s) if isinstance(monotonic_s, (int, float)) else 0.0,
+    )
+
+
+def _canonical(value: Any) -> str:
+    return json.dumps(value, ensure_ascii=True, separators=(",", ":"), sort_keys=True)
+
+
+def _target_event(event: dict[str, Any]) -> bool:
+    """Exclude profile/dummy calls with an empty SchedulerOutput.
+
+    The worker hook also wraps vLLM's profile and DP dummy runs. Those calls
+    have a physical batch descriptor but no target request, so their ranks and
+    counts are not a valid collective-plan invariant.
+    """
+
+    return (
+        (integer(event.get("request_count")) or 0) > 0
+        and (integer(event.get("total_scheduled_rows")) or 0) > 0
+    )
+
+
+def _target_signature(event: dict[str, Any]) -> str:
+    return _canonical(
+        {
+            "cudagraph_mode": event.get("cudagraph_mode"),
+            "ordered_plan_digest": event.get("ordered_plan_digest"),
+            "physical_batch_rows": event.get("physical_batch_rows"),
+            "rows_across_dp": event.get("rows_across_dp"),
+            "scheduled_dp_metadata": event.get("scheduled_dp_metadata"),
+            "semantic_plan_digest": event.get("semantic_plan_digest"),
+            "should_ubatch": event.get("should_ubatch"),
+            "total_scheduled_rows": event.get("total_scheduled_rows"),
+        }
+    )
+
+
+def _target_metadata_check(event: dict[str, Any]) -> tuple[bool, int | None, bool | None]:
+    """Check the local, post-DP-coordination shape invariants."""
+
+    metadata = event.get("scheduled_dp_metadata")
+    dp_rank = integer(event.get("data_parallel_rank"))
+    dp_size = integer(event.get("data_parallel_size"))
+    total_rows = integer(event.get("total_scheduled_rows"))
+    physical_rows = integer(event.get("physical_batch_rows"))
+    rows_across_dp = event.get("rows_across_dp")
+    if (
+        not isinstance(metadata, dict)
+        or dp_rank is None
+        or dp_size is None
+        or dp_rank < 0
+        or dp_rank >= dp_size
+        or total_rows is None
+        or physical_rows is None
+        or not isinstance(rows_across_dp, list)
+    ):
+        return False, None, None
+    raw_rows = metadata.get("num_tokens_per_rank")
+    padded_rows = metadata.get("num_tokens_across_dp")
+    if (
+        not isinstance(raw_rows, list)
+        or not isinstance(padded_rows, list)
+        or len(raw_rows) != dp_size
+        or len(padded_rows) != dp_size
+        or len(rows_across_dp) != dp_size
+        or not all(integer(value) is not None and integer(value) >= 0 for value in raw_rows)
+        or not all(integer(value) is not None and integer(value) >= 0 for value in padded_rows)
+        or not all(integer(value) is not None and integer(value) >= 0 for value in rows_across_dp)
+    ):
+        return False, None, None
+    raw = [int(value) for value in raw_rows]
+    padded = [int(value) for value in padded_rows]
+    coordinated = [int(value) for value in rows_across_dp]
+    valid = (
+        raw[dp_rank] == total_rows
+        and physical_rows == coordinated[dp_rank]
+        and padded == coordinated
+        and all(padded_value >= raw_value for padded_value, raw_value in zip(padded, raw))
+    )
+    return valid, padded[dp_rank] - raw[dp_rank], len(set(raw)) > 1
+
+
+def _dp_pair_key(event: dict[str, Any]) -> tuple[str, str] | None:
+    """Return the TP-shard key shared by the two actual DP peers.
+
+    New P0 logs carry ``tensor_parallel_rank``. The v2 artifact predates that
+    field, where ``Worker.rank`` was logged as ``global_rank`` and is the same
+    TP-shard identifier within each DP replica. The fallback is explicit in
+    the output so it cannot be confused with a true global rank.
+    """
+
+    tp_rank = integer(event.get("tensor_parallel_rank"))
+    if tp_rank is not None:
+        return f"tp{tp_rank}", "tensor_parallel_rank"
+    legacy_worker_rank = integer(event.get("global_rank"))
+    if legacy_worker_rank is not None:
+        return f"legacy-worker-rank{legacy_worker_rank}", "legacy_worker_rank"
+    local_rank = integer(event.get("local_rank"))
+    tp_size = integer(event.get("tensor_parallel_size"))
+    if local_rank is not None and tp_size is not None and tp_size > 0:
+        return f"local-rank-mod-tp{local_rank % tp_size}", "local_rank_mod_tp_size"
+    return None
+
+
+def _dp_coordination_signature(event: dict[str, Any]) -> str:
+    metadata = event.get("scheduled_dp_metadata")
+    if not isinstance(metadata, dict):
+        metadata = {}
+    return _canonical(
+        {
+            "all_decode": metadata.get("all_decode"),
+            "all_prefill": metadata.get("all_prefill"),
+            "cudagraph_mode": metadata.get("cudagraph_mode"),
+            "input_fits_in_drafter": metadata.get("input_fits_in_drafter"),
+            "is_prompt_batch": metadata.get("is_prompt_batch"),
+            "num_reqs_per_rank": metadata.get("num_reqs_per_rank"),
+            "num_tokens_across_dp": metadata.get("num_tokens_across_dp"),
+            "num_tokens_per_rank": metadata.get("num_tokens_per_rank"),
+            "rows_across_dp": event.get("rows_across_dp"),
+            "should_ubatch": metadata.get("should_ubatch"),
+        }
+    )
+
+
+def dp_pair_coordination(target_events: list[dict[str, Any]]) -> dict[str, Any]:
+    """Compare scalar DP coordination seen by the two peers of each TP shard."""
+
+    grouped: dict[str, dict[int, dict[int, dict[str, Any]]]] = defaultdict(
+        lambda: defaultdict(dict)
+    )
+    sources: dict[str, set[str]] = defaultdict(set)
+    expected_sizes: dict[str, set[int]] = defaultdict(set)
+    for event in target_events:
+        key = _dp_pair_key(event)
+        dp_rank = integer(event.get("data_parallel_rank"))
+        epoch = integer(event.get("batch_phase_epoch"))
+        dp_size = integer(event.get("data_parallel_size"))
+        if key is None or dp_rank is None or epoch is None:
+            continue
+        group, source = key
+        grouped[group][dp_rank][epoch] = event
+        sources[group].add(source)
+        if dp_size is not None:
+            expected_sizes[group].add(dp_size)
+
+    groups: dict[str, Any] = {}
+    for group, by_dp_rank in grouped.items():
+        dp_ranks = sorted(by_dp_rank)
+        epoch_sets = [set(by_epoch) for by_epoch in by_dp_rank.values()]
+        shared_epochs = sorted(set.intersection(*epoch_sets)) if epoch_sets else []
+        signatures = {
+            dp_rank: {
+                epoch: _dp_coordination_signature(event)
+                for epoch, event in by_epoch.items()
+            }
+            for dp_rank, by_epoch in by_dp_rank.items()
+        }
+        shared_equal = all(
+            len({signatures[dp_rank][epoch] for dp_rank in dp_ranks}) == 1
+            for epoch in shared_epochs
+        )
+        expected_values = expected_sizes.get(group, set())
+        expected_dp_size = next(iter(expected_values)) if len(expected_values) == 1 else None
+        groups[group] = {
+            "grouping_source": sorted(sources[group]),
+            "data_parallel_ranks": dp_ranks,
+            "expected_data_parallel_size": expected_dp_size,
+            "target_epoch_count_by_dp_rank": {
+                str(dp_rank): len(by_dp_rank[dp_rank]) for dp_rank in dp_ranks
+            },
+            "shared_target_epoch_count": len(shared_epochs),
+            "all_shared_epoch_signatures_identical": bool(shared_epochs) and shared_equal,
+            "all_expected_dp_peers_observed": expected_dp_size is not None
+            and len(dp_ranks) == expected_dp_size,
+        }
+    return {
+        "group_count": len(groups),
+        "groups": groups,
+        "all_expected_dp_peers_observed": bool(groups)
+        and all(value["all_expected_dp_peers_observed"] for value in groups.values()),
+        "all_shared_epoch_signatures_identical": bool(groups)
+        and all(value["all_shared_epoch_signatures_identical"] for value in groups.values()),
+    }
+
+
+def target_execution_agreement(worker_events: list[dict[str, Any]]) -> dict[str, Any]:
+    batch_events = [event for event in worker_events if event.get("event") == "batch_execution_plan"]
+    target_events = [event for event in batch_events if _target_event(event)]
     per_rank: dict[str, list[dict[str, Any]]] = defaultdict(list)
-    per_dp: dict[str, set[str]] = defaultdict(set)
-    for event in begins:
+    per_replica: dict[str, set[str]] = defaultdict(set)
+    for event in target_events:
         rank = rank_key(event)
-        dp_rank = event.get("data_parallel_rank")
-        if rank is None:
+        dp_rank = integer(event.get("data_parallel_rank"))
+        if rank is None or dp_rank is None:
             continue
         per_rank[rank].append(event)
-        if dp_rank is not None:
-            per_dp[str(dp_rank)].add(rank)
+        per_replica[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))))
-        sequences[rank] = [str(item.get("ordered_plan_digest")) for item in values]
+        values.sort(key=lambda event: _event_order(event, "batch_phase_epoch"))
+        sequences[rank] = [_target_signature(event) for event in values]
     counts = {rank: len(values) for rank, values in sequences.items()}
-    within_dp: dict[str, Any] = {}
-    for dp_rank, ranks in per_dp.items():
+
+    replicas: dict[str, Any] = {}
+    for dp_rank, ranks in per_replica.items():
         rank_list = sorted(ranks)
         distinct = {tuple(sequences[rank]) for rank in rank_list}
-        within_dp[dp_rank] = {
+        phase_counts = [counts[rank] for rank in rank_list]
+        replicas[dp_rank] = {
             "ranks": rank_list,
             "phase_count_by_rank": {rank: counts[rank] for rank in rank_list},
+            "phase_count_equal": bool(phase_counts) and len(set(phase_counts)) == 1,
             "sequence_distinct_count": len(distinct),
-            "identical_execution_sequence": len(distinct) == 1,
+            "identical_target_execution_sequence": len(distinct) == 1,
         }
-    count_values = list(counts.values())
+
+    metadata_valid_count = 0
+    metadata_padding: list[int] = []
+    raw_unequal_count = 0
+    for event in target_events:
+        valid, padding, raw_unequal = _target_metadata_check(event)
+        metadata_valid_count += valid
+        if padding is not None:
+            metadata_padding.append(padding)
+        raw_unequal_count += bool(raw_unequal)
     return {
-        "worker_begin_event_count": len(begins),
+        "batch_execution_plan_count": len(batch_events),
+        "dummy_or_profile_event_count": len(batch_events) - len(target_events),
+        "target_event_count": len(target_events),
         "worker_rank_count": len(sequences),
-        "phase_count_by_rank": counts,
-        "phase_count": {
-            "n": len(count_values),
-            "min": min(count_values) if count_values else None,
-            "max": max(count_values) if count_values else None,
-            "distinct_value_count": len(set(count_values)),
-            "all_equal": bool(count_values) and len(set(count_values)) == 1,
+        "target_phase_count_by_rank": counts,
+        "within_dp_replica_tp_group": replicas,
+        "all_target_phase_counts_equal_within_dp_replica": bool(replicas)
+        and all(value["phase_count_equal"] for value in replicas.values()),
+        "all_target_sequences_identical_within_dp_replica": bool(replicas)
+        and all(value["identical_target_execution_sequence"] for value in replicas.values()),
+        "dp_metadata": {
+            "target_record_count": len(target_events),
+            "valid_record_count": metadata_valid_count,
+            "all_records_valid": len(target_events) > 0
+            and metadata_valid_count == len(target_events),
+            "raw_dp_counts_unequal_record_count": raw_unequal_count,
+            "post_coordinate_minus_local_raw_rows": {
+                "n": len(metadata_padding),
+                "min": min(metadata_padding) if metadata_padding else None,
+                "max": max(metadata_padding) if metadata_padding else None,
+                "sum": sum(metadata_padding),
+                "distinct_value_count": len(set(metadata_padding)),
+                "non_negative": all(value >= 0 for value in metadata_padding),
+            },
         },
-        "within_data_parallel_replica": within_dp,
-        "all_within_dp_sequences_identical": bool(within_dp)
-        and all(value["identical_execution_sequence"] for value in within_dp.values()),
+        "dp_pair_coordination": dp_pair_coordination(target_events),
     }
 
 
@@ -160,10 +376,11 @@ def summarize_cell(cell: Path, expected_requests: int, expected_tokens: int) ->
         event for event in schedule_events if integer(event.get("spec_request_count")) not in (None, 0)
     ]
     batch_events = [event for event in worker_events if event.get("event") == "batch_execution_plan"]
-    rows_across_dp = [event.get("rows_across_dp") for event in batch_events]
+    target_events = [event for event in batch_events if _target_event(event)]
+    rows_across_dp = [event.get("rows_across_dp") for event in target_events]
     physical_rows = [
         integer(event.get("physical_batch_rows"))
-        for event in batch_events
+        for event in target_events
         if integer(event.get("physical_batch_rows")) is not None
     ]
     return {
@@ -186,10 +403,9 @@ def summarize_cell(cell: Path, expected_requests: int, expected_tokens: int) ->
             ),
         },
         "worker": {
-            **sequence_agreement(worker_events),
-            "batch_execution_plan_count": len(batch_events),
-            "rows_across_dp_present_count": sum(value is not None for value in rows_across_dp),
-            "physical_batch_rows": {
+            **target_execution_agreement(worker_events),
+            "target_rows_across_dp_present_count": sum(value is not None for value in rows_across_dp),
+            "target_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,
@@ -227,18 +443,46 @@ def data_sanity(cells: dict[str, dict[str, Any]], expected_gpu_count: int | None
                 bool(cell.get("probe_integrity", {}).get("valid")) for cell in cells.values()
             ),
             "no_event_parse_errors": all(not cell.get("event_parse_errors") for cell in cells.values()),
-            "all_rank_phase_counts_equal": all(
-                bool(cell.get("worker", {}).get("phase_count", {}).get("all_equal"))
+            "all_target_phase_counts_equal_within_dp_replica": all(
+                bool(
+                    cell.get("worker", {}).get(
+                        "all_target_phase_counts_equal_within_dp_replica"
+                    )
+                )
                 for cell in cells.values()
             ),
-            "all_within_dp_sequences_identical": all(
-                bool(cell.get("worker", {}).get("all_within_dp_sequences_identical"))
+            "all_target_sequences_identical_within_dp_replica": all(
+                bool(
+                    cell.get("worker", {}).get(
+                        "all_target_sequences_identical_within_dp_replica"
+                    )
+                )
                 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
+            "all_target_dp_metadata_valid": all(
+                bool(cell.get("worker", {}).get("dp_metadata", {}).get("all_records_valid"))
+                for cell in cells.values()
+            ),
+            "all_target_dp_coordination_records_observed": all(
+                integer(cell.get("worker", {}).get("target_rows_across_dp_present_count"))
+                == integer(cell.get("worker", {}).get("target_event_count"))
+                and integer(cell.get("worker", {}).get("target_event_count", 0)) > 0
+                for cell in cells.values()
+            ),
+            "all_expected_dp_peers_observed": all(
+                bool(
+                    cell.get("worker", {}).get("dp_pair_coordination", {}).get(
+                        "all_expected_dp_peers_observed"
+                    )
+                )
+                for cell in cells.values()
+            ),
+            "all_shared_dp_pair_coordination_signatures_identical": all(
+                bool(
+                    cell.get("worker", {}).get("dp_pair_coordination", {}).get(
+                        "all_shared_epoch_signatures_identical"
+                    )
+                )
                 for cell in cells.values()
             ),
             "all_expected_workers_observed": (
@@ -259,7 +503,7 @@ def markdown(payload: dict[str, Any]) -> str:
         "",
         "P0 verifies the collective/liveness premise only; it is not a performance result.",
         "",
-        "| policy | complete usage-verified requests | candidate K values observed | worker ranks | rank phase counts equal | within-DP execution sequences identical |",
+        "| policy | complete usage-verified requests | candidate K values observed | worker ranks | target phase counts equal within each DP replica | target sequences identical within each DP replica |",
         "|---|---:|---:|---:|---:|---:|",
     ]
     for name, cell in payload["cells"].items():
@@ -275,8 +519,8 @@ def markdown(payload: dict[str, Any]) -> str:
                 valid=probe.get("valid", False),
                 values=values,
                 ranks=worker.get("worker_rank_count", 0),
-                equal=worker.get("phase_count", {}).get("all_equal", False),
-                within=worker.get("all_within_dp_sequences_identical", False),
+                equal=worker.get("all_target_phase_counts_equal_within_dp_replica", False),
+                within=worker.get("all_target_sequences_identical_within_dp_replica", False),
             )
         )
     lines.extend(
@@ -319,7 +563,7 @@ def main() -> int:
     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))
-    return 0 if payload["data_sanity"]["invariants"]["all_probe_integrity_valid"] else 2
+    return 0 if all(payload["data_sanity"]["invariants"].values()) else 2
 
 
 if __name__ == "__main__":