diff --git a/scripts/collectivespec/p0_overlay/collectivespec_p0/worker.py b/scripts/collectivespec/p0_overlay/collectivespec_p0/worker.py index 565933a..c318e36 100644 --- a/scripts/collectivespec/p0_overlay/collectivespec_p0/worker.py +++ b/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 diff --git a/scripts/collectivespec/run_p0_phase_trace.py b/scripts/collectivespec/run_p0_phase_trace.py index 11e8815..9712bbb 100644 --- a/scripts/collectivespec/run_p0_phase_trace.py +++ b/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"], diff --git a/scripts/collectivespec/summarize_p0_phase_trace.py b/scripts/collectivespec/summarize_p0_phase_trace.py index 20c46eb..c21aca6 100644 --- a/scripts/collectivespec/summarize_p0_phase_trace.py +++ b/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[key].append(event) + per_rank[rank].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 = [ - event.get("scheduled_dp_metadata") - for event in worker_events - if event.get("event") == "target_execute_begin" + 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] + physical_rows = [ + 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, - "dp_metadata_missing_count": len(worker_metadata) - metadata_present, + "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": { + "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))