327 lines
14 KiB
Python
327 lines
14 KiB
Python
#!/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 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"]
|
|
per_rank: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
|
per_dp: dict[str, set[str]] = defaultdict(set)
|
|
for event in begins:
|
|
rank = rank_key(event)
|
|
dp_rank = event.get("data_parallel_rank")
|
|
if rank is None:
|
|
continue
|
|
per_rank[rank].append(event)
|
|
if dp_rank is not None:
|
|
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))))
|
|
sequences[rank] = [str(item.get("ordered_plan_digest")) for item 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():
|
|
rank_list = sorted(ranks)
|
|
distinct = {tuple(sequences[rank]) for rank in rank_list}
|
|
within_dp[dp_rank] = {
|
|
"ranks": rank_list,
|
|
"phase_count_by_rank": {rank: counts[rank] for rank in rank_list},
|
|
"sequence_distinct_count": len(distinct),
|
|
"identical_execution_sequence": len(distinct) == 1,
|
|
}
|
|
count_values = list(counts.values())
|
|
return {
|
|
"worker_begin_event_count": len(begins),
|
|
"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,
|
|
},
|
|
"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()),
|
|
}
|
|
|
|
|
|
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"]
|
|
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
|
|
]
|
|
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": {
|
|
**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": {
|
|
"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_rank_phase_counts_equal": all(
|
|
bool(cell.get("worker", {}).get("phase_count", {}).get("all_equal"))
|
|
for cell in cells.values()
|
|
),
|
|
"all_within_dp_sequences_identical": all(
|
|
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
|
|
),
|
|
},
|
|
}
|
|
|
|
|
|
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 | rank phase counts equal | within-DP execution sequences identical |",
|
|
"|---|---:|---:|---:|---:|---:|",
|
|
]
|
|
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("phase_count", {}).get("all_equal", False),
|
|
within=worker.get("all_within_dp_sequences_identical", 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 payload["data_sanity"]["invariants"]["all_probe_integrity_valid"] else 2
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|