gahow/aituner
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Trace P0 batch-level DP coordination

Browse Source
This commit is contained in:
Gahow Wang
2026-07-13 17:45:08 +08:00
parent f1cd859eea
commit 7f4ae1708b
3 changed files with 82 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
scripts/collectivespec/p0_overlay/collectivespec_p0/worker.py
View File

@@ -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

1
scripts/collectivespec/run_p0_phase_trace.py
View File

@@ -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"],

65
scripts/collectivespec/summarize_p0_phase_trace.py
View File

@@ -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))