#!/usr/bin/env python3 """Summarize the non-performance CollectiveSpec P0 phase-trace artifacts.""" from __future__ import annotations import argparse import json from collections import defaultdict from pathlib import Path from typing import Any def load_json(path: Path) -> Any: return json.loads(path.read_text(encoding="utf-8")) def integer(value: Any) -> int | None: return value if isinstance(value, int) and not isinstance(value, bool) else None def merge_histograms(records: list[dict[str, Any]], field: str) -> dict[str, int]: result: dict[str, int] = {} for record in records: values = record.get(field) if not isinstance(values, dict): continue for key, value in values.items(): count = integer(value) if count is not None and count >= 0: result[str(key)] = result.get(str(key), 0) + count return dict(sorted(result.items(), key=lambda item: int(item[0]))) def load_events(log_dir: Path) -> tuple[list[dict[str, Any]], list[str]]: events: list[dict[str, Any]] = [] errors: list[str] = [] for path in sorted(log_dir.glob("*.jsonl")): for line_number, raw in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1): if not raw: continue try: record = json.loads(raw) except json.JSONDecodeError as exc: errors.append(f"{path.name}:{line_number}: {exc.msg}") continue if not isinstance(record, dict): errors.append(f"{path.name}:{line_number}: event is not an object") continue events.append(record) return events, errors def probe_integrity(cell: Path, expected_requests: int, expected_tokens: int) -> dict[str, Any]: details_paths = sorted(cell.glob("store/*/trials/trial-*/probe_details.jsonl")) result_paths = sorted(cell.glob("store/*/trials/trial-*/result.json")) if len(details_paths) != 1 or len(result_paths) != 1: return { "valid": False, "failures": [ f"probe_details_count={len(details_paths)}", f"result_count={len(result_paths)}", ], } details_rows = [json.loads(line) for line in details_paths[0].read_text(encoding="utf-8").splitlines() if line] result = load_json(result_paths[0]) failures: list[str] = [] if result.get("status") != "completed": failures.append(f"result_status={result.get('status')}") if len(details_rows) != 1: failures.append(f"detail_rows={len(details_rows)}") return {"valid": False, "failures": failures} outcomes = details_rows[0].get("outcomes") if not isinstance(outcomes, list): failures.append("outcomes_not_list") outcomes = [] if len(outcomes) != expected_requests: failures.append(f"outcome_count={len(outcomes)}_expected={expected_requests}") successful = sum(bool(item.get("success")) for item in outcomes if isinstance(item, dict)) verified_tokens = sum( isinstance(item, dict) and item.get("completion_tokens_source") == "usage" and item.get("completion_tokens") == expected_tokens for item in outcomes ) if successful != expected_requests: failures.append(f"success_count={successful}_expected={expected_requests}") if verified_tokens != expected_requests: failures.append(f"usage_token_count={verified_tokens}_expected={expected_requests}") return { "valid": not failures, "failures": failures, "outcome_count": len(outcomes), "success_count": successful, "usage_token_count": verified_tokens, "details_path": str(details_paths[0]), "result_path": str(result_paths[0]), } 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 _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_replica: dict[str, set[str]] = defaultdict(set) for event in target_events: rank = rank_key(event) dp_rank = integer(event.get("data_parallel_rank")) if rank is None or dp_rank is None: continue per_rank[rank].append(event) per_replica[str(dp_rank)].add(rank) sequences: dict[str, list[str]] = {} for rank, values in per_rank.items(): 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()} 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} 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_target_execution_sequence": len(distinct) == 1, } metadata_valid_count = 0 anchor_local_ranks: dict[int, int] = {} for event in target_events: dp_rank = integer(event.get("data_parallel_rank")) local_rank = integer(event.get("local_rank")) if dp_rank is None or local_rank is None: continue previous = anchor_local_ranks.get(dp_rank) if previous is None or local_rank < previous: anchor_local_ranks[dp_rank] = local_rank anchor_events = [ event for event in target_events if integer(event.get("data_parallel_rank")) in anchor_local_ranks and integer(event.get("local_rank")) == anchor_local_ranks[integer(event.get("data_parallel_rank"))] ] metadata_padding: list[int] = [] raw_unequal_count = 0 for event in target_events: valid, _, _ = _target_metadata_check(event) metadata_valid_count += valid for event in anchor_events: _, padding, raw_unequal = _target_metadata_check(event) if padding is not None: metadata_padding.append(padding) raw_unequal_count += bool(raw_unequal) return { "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), "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), "anchor_record_count": len(anchor_events), "raw_dp_counts_unequal_anchor_record_count": raw_unequal_count, "per_dp_replica_anchor_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), }, }, "dp_pair_coordination": dp_pair_coordination(target_events), } def summarize_cell(cell: Path, expected_requests: int, expected_tokens: int) -> dict[str, Any]: events, parse_errors = load_events(cell / "p0_logs") scheduler_events = [event for event in events if event.get("role") == "scheduler"] worker_events = [event for event in events if event.get("role") == "worker"] candidate_events = [event for event in scheduler_events if event.get("event") == "candidate_truncate"] schedule_events = [event for event in scheduler_events if event.get("event") == "schedule"] verify_schedule_events = [ 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"] 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 target_events if integer(event.get("physical_batch_rows")) is not None ] return { "probe_integrity": probe_integrity(cell, expected_requests, expected_tokens), "event_parse_errors": parse_errors, "event_count": len(events), "scheduler": { "schedule_event_count": len(schedule_events), "verify_schedule_event_count": len(verify_schedule_events), "candidate_event_count": len(candidate_events), "before_k_histogram": merge_histograms(candidate_events, "before_k_histogram"), "after_k_histogram": merge_histograms(candidate_events, "after_k_histogram"), "after_k_distinct_value_count": len(merge_histograms(candidate_events, "after_k_histogram")), "dp_rank_count": len( { event.get("data_parallel_rank") for event in scheduler_events if event.get("data_parallel_rank") is not None } ), }, "worker": { **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, "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]], 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 for cell in cells.values() ] return { "n_cells": len(cells), "event_count": { "n": len(event_counts), "min": min(event_counts) if event_counts else None, "max": max(event_counts) if event_counts else None, "distinct_value_count": len(set(event_counts)), "non_negative": all(value >= 0 for value in event_counts), }, "candidate_k_distinct": { "n": len(candidate_distinct), "min": min(candidate_distinct) if candidate_distinct else None, "max": max(candidate_distinct) if candidate_distinct else None, "distinct_value_count": len(set(candidate_distinct)), "non_negative": all(value >= 0 for value in candidate_distinct), }, "invariants": { "all_probe_integrity_valid": all( 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_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_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_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": ( 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 ), }, } def markdown(payload: dict[str, Any]) -> str: lines = [ "# CollectiveSpec P0 phase-trace summary", "", "P0 verifies the collective/liveness premise only; it is not a performance result.", "", "| 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(): probe = cell["probe_integrity"] scheduler = cell["scheduler"] worker = cell["worker"] values = ",".join(sorted(scheduler["after_k_histogram"])) or "—" lines.append( "| {name} | {success}/{count} ({valid}) | {values} | {ranks} | {equal} | {within} |".format( name=name, success=probe.get("success_count", "—"), count=probe.get("outcome_count", "—"), valid=probe.get("valid", False), values=values, ranks=worker.get("worker_rank_count", 0), 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( [ "", "## Data sanity", "", "```json", json.dumps(payload["data_sanity"], indent=2, sort_keys=True), "```", "", ] ) return "\n".join(lines) def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--root", type=Path, required=True) parser.add_argument("--output-json", type=Path, required=True) parser.add_argument("--output-md", type=Path, required=True) args = parser.parse_args() manifest = load_json(args.root / "manifest.json") expected_requests = integer(manifest.get("request_count")) expected_tokens = integer(manifest.get("completion_tokens")) if expected_requests is None or expected_tokens is None: raise SystemExit("manifest must contain integer request_count and completion_tokens") policies = manifest.get("policies") if not isinstance(policies, list) or not all(isinstance(value, str) for value in policies): raise SystemExit("manifest policies must be a list of strings") cells = { policy: summarize_cell(args.root / policy, expected_requests, expected_tokens) for policy in policies } 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)) return 0 if all(payload["data_sanity"]["invariants"].values()) else 2 if __name__ == "__main__": raise SystemExit(main())