From f1cd859eea0bcbedc553e91c97045f37753da782 Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Mon, 13 Jul 2026 17:33:01 +0800 Subject: [PATCH] Add CollectiveSpec P0 phase trace harness --- .../p0_overlay/collectivespec_p0/__init__.py | 6 + .../p0_overlay/collectivespec_p0/common.py | 158 ++++++++++ .../p0_overlay/collectivespec_p0/policy.json | 6 + .../p0_overlay/collectivespec_p0/scheduler.py | 81 +++++ .../p0_overlay/collectivespec_p0/worker.py | 44 +++ scripts/collectivespec/run_p0_phase_trace.py | 298 ++++++++++++++++++ .../summarize_p0_phase_trace.py | 293 +++++++++++++++++ 7 files changed, 886 insertions(+) create mode 100644 scripts/collectivespec/p0_overlay/collectivespec_p0/__init__.py create mode 100644 scripts/collectivespec/p0_overlay/collectivespec_p0/common.py create mode 100644 scripts/collectivespec/p0_overlay/collectivespec_p0/policy.json create mode 100644 scripts/collectivespec/p0_overlay/collectivespec_p0/scheduler.py create mode 100644 scripts/collectivespec/p0_overlay/collectivespec_p0/worker.py create mode 100644 scripts/collectivespec/run_p0_phase_trace.py create mode 100644 scripts/collectivespec/summarize_p0_phase_trace.py diff --git a/scripts/collectivespec/p0_overlay/collectivespec_p0/__init__.py b/scripts/collectivespec/p0_overlay/collectivespec_p0/__init__.py new file mode 100644 index 0000000..a255406 --- /dev/null +++ b/scripts/collectivespec/p0_overlay/collectivespec_p0/__init__.py @@ -0,0 +1,6 @@ +"""Non-invasive instrumentation for the CollectiveSpec P0 premise test. + +This package is loaded only when its parent ``p0_overlay`` directory is added +to ``PYTHONPATH`` and vLLM is explicitly given the custom scheduler/worker +classes. It is intentionally not imported by normal aituner runs. +""" diff --git a/scripts/collectivespec/p0_overlay/collectivespec_p0/common.py b/scripts/collectivespec/p0_overlay/collectivespec_p0/common.py new file mode 100644 index 0000000..a0eda0f --- /dev/null +++ b/scripts/collectivespec/p0_overlay/collectivespec_p0/common.py @@ -0,0 +1,158 @@ +"""Small, dependency-free helpers shared by the P0 vLLM extensions.""" + +from __future__ import annotations + +import hashlib +import json +import os +import threading +import time +from pathlib import Path +from typing import Any + + +_LOCK = threading.Lock() +_POLICY: dict[str, Any] | None = None + + +def _canonical(value: Any) -> str: + return json.dumps(value, ensure_ascii=True, separators=(",", ":"), sort_keys=True) + + +def digest(value: Any) -> str: + return hashlib.blake2b(_canonical(value).encode("utf-8"), digest_size=16).hexdigest() + + +def policy() -> dict[str, Any]: + global _POLICY + if _POLICY is not None: + return _POLICY + path_value = os.environ.get("COLLECTIVESPEC_P0_POLICY_PATH") + if not path_value: + raise RuntimeError("COLLECTIVESPEC_P0_POLICY_PATH is required for P0") + path = Path(path_value) + value = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(value, dict): + raise RuntimeError("P0 policy must be a JSON object") + max_k = value.get("max_k") + seed = value.get("seed") + kind = value.get("policy") + if isinstance(max_k, bool) or not isinstance(max_k, int) or max_k < 0: + raise RuntimeError("P0 policy max_k must be a non-negative integer") + if isinstance(seed, bool) or not isinstance(seed, int): + raise RuntimeError("P0 policy seed must be an integer") + if kind not in {"blake2b-request-static-k", "constant-k"}: + raise RuntimeError(f"Unsupported P0 policy: {kind!r}") + if kind == "constant-k": + constant_k = value.get("constant_k") + if isinstance(constant_k, bool) or not isinstance(constant_k, int): + raise RuntimeError("constant-k policy requires integer constant_k") + if not 0 <= constant_k <= max_k: + raise RuntimeError("constant_k must be within [0, max_k]") + _POLICY = value + return value + + +def requested_k(request_id: str, candidate_count: int) -> int: + value = policy() + max_k = int(value["max_k"]) + if candidate_count < 0: + raise RuntimeError("candidate_count must be non-negative") + if value["policy"] == "constant-k": + chosen = int(value["constant_k"]) + else: + material = f"{value['seed']}|{request_id}".encode("utf-8") + chosen = int.from_bytes(hashlib.blake2b(material, digest_size=8).digest(), "big") % (max_k + 1) + return min(chosen, candidate_count) + + +def histogram(values: list[int]) -> dict[str, int]: + counts: dict[str, int] = {} + for value in values: + key = str(int(value)) + counts[key] = counts.get(key, 0) + 1 + return dict(sorted(counts.items(), key=lambda item: int(item[0]))) + + +def _jsonable(value: Any) -> Any: + if value is None or isinstance(value, (str, int, float, bool)): + return value + if isinstance(value, (list, tuple)): + return [_jsonable(item) for item in value] + if isinstance(value, dict): + return {str(key): _jsonable(item) for key, item in value.items()} + tolist = getattr(value, "tolist", None) + if callable(tolist): + return _jsonable(tolist()) + item = getattr(value, "item", None) + if callable(item): + try: + return _jsonable(item()) + except (TypeError, ValueError): + pass + return repr(value) + + +def dp_metadata_summary(metadata: Any) -> dict[str, Any] | None: + if metadata is None: + return None + keys = ( + "num_pad_tokens", + "num_tokens_across_dp", + "num_tokens_per_rank", + "num_reqs_per_rank", + "is_prompt_batch", + "all_prefill", + "all_decode", + "should_ubatch", + "input_fits_in_drafter", + "cudagraph_mode", + ) + return {key: _jsonable(getattr(metadata, key, None)) for key in keys} + + +def plan_summary(scheduler_output: Any) -> dict[str, Any]: + scheduled = getattr(scheduler_output, "num_scheduled_tokens", {}) or {} + spec = getattr(scheduler_output, "scheduled_spec_decode_tokens", {}) or {} + ordered = [ + { + "request_id": str(request_id), + "scheduled_tokens": int(tokens), + "spec_k": len(spec.get(request_id, [])), + } + for request_id, tokens in scheduled.items() + ] + semantic = sorted(ordered, key=lambda item: item["request_id"]) + spec_ks = [item["spec_k"] for item in ordered if item["spec_k"] > 0] + return { + "ordered_plan_digest": digest(ordered), + "semantic_plan_digest": digest(semantic), + "request_count": len(ordered), + "total_scheduled_rows": int(getattr(scheduler_output, "total_num_scheduled_tokens", 0)), + "spec_request_count": len(spec), + "spec_k_histogram": histogram(spec_ks), + "verify_logical_rows": sum(1 + item["spec_k"] for item in ordered if item["spec_k"] > 0), + "scheduled_dp_metadata": dp_metadata_summary( + getattr(scheduler_output, "scheduled_dp_metadata", None) + ), + } + + +def log_event(role: str, event: str, **payload: Any) -> None: + root_value = os.environ.get("COLLECTIVESPEC_P0_LOG_DIR") + if not root_value: + raise RuntimeError("COLLECTIVESPEC_P0_LOG_DIR is required for P0") + root = Path(root_value) + root.mkdir(parents=True, exist_ok=True) + record = { + "schema_version": "collectivespec-p0-log-v1", + "event": event, + "monotonic_s": time.monotonic(), + "pid": os.getpid(), + "policy_version": policy().get("version"), + "role": role, + **{key: _jsonable(value) for key, value in payload.items()}, + } + path = root / f"{role}-pid{os.getpid()}.jsonl" + with _LOCK, path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(record, ensure_ascii=True, sort_keys=True) + "\n") diff --git a/scripts/collectivespec/p0_overlay/collectivespec_p0/policy.json b/scripts/collectivespec/p0_overlay/collectivespec_p0/policy.json new file mode 100644 index 0000000..d2e3d40 --- /dev/null +++ b/scripts/collectivespec/p0_overlay/collectivespec_p0/policy.json @@ -0,0 +1,6 @@ +{ + "max_k": 3, + "policy": "blake2b-request-static-k", + "seed": 20260713, + "version": "collectivespec-p0-v1" +} diff --git a/scripts/collectivespec/p0_overlay/collectivespec_p0/scheduler.py b/scripts/collectivespec/p0_overlay/collectivespec_p0/scheduler.py new file mode 100644 index 0000000..cfca6d7 --- /dev/null +++ b/scripts/collectivespec/p0_overlay/collectivespec_p0/scheduler.py @@ -0,0 +1,81 @@ +"""P0 scheduler: inject a pre-fixed per-request verification horizon.""" + +from __future__ import annotations + +from typing import Any + +from vllm.v1.core.sched.scheduler import Scheduler +from vllm.v1.outputs import DraftTokenIds + +from .common import digest, histogram, log_event, plan_summary, requested_k + + +class P0Scheduler(Scheduler): + """A Scheduler that changes only post-drafter candidate visibility. + + EAGLE has already produced its Kmax candidates when this method runs. The + P0 experiment therefore tests verifier-side ragged execution and collective + phase behavior, not a drafter speedup. + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self._p0_schedule_epoch = 0 + self._p0_candidate_epoch = 0 + + def _p0_rank_fields(self) -> dict[str, Any]: + config = self.vllm_config.parallel_config + return { + "data_parallel_rank": getattr(config, "data_parallel_rank", None), + "data_parallel_size": getattr(config, "data_parallel_size", None), + "tensor_parallel_size": getattr(config, "tensor_parallel_size", None), + "expert_parallel_size": getattr(config, "expert_parallel_size", None), + } + + def schedule(self, *args: Any, **kwargs: Any): # type: ignore[no-untyped-def] + output = super().schedule(*args, **kwargs) + self._p0_schedule_epoch += 1 + log_event( + "scheduler", + "schedule", + core_local_epoch=self._p0_schedule_epoch, + **self._p0_rank_fields(), + **plan_summary(output), + ) + return output + + def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None: + self._p0_candidate_epoch += 1 + before = [list(ids) for ids in draft_token_ids.draft_token_ids] + chosen: list[int] = [] + after: list[list[int]] = [] + structured_output_request_ids: list[str] = [] + for request_id, ids in zip(draft_token_ids.req_ids, before): + request = self.requests.get(request_id) + if request is not None and getattr(request, "structured_output_request", None) is not None: + # Do not make grammar behavior part of this premise experiment. + chosen.append(len(ids)) + after.append(ids) + structured_output_request_ids.append(str(request_id)) + continue + k = requested_k(str(request_id), len(ids)) + chosen.append(k) + after.append(ids[:k]) + pairs = [ + {"request_id": str(request_id), "before_k": len(ids), "after_k": after_k} + for request_id, ids, after_k in zip(draft_token_ids.req_ids, before, chosen) + ] + log_event( + "scheduler", + "candidate_truncate", + candidate_epoch=self._p0_candidate_epoch, + **self._p0_rank_fields(), + candidate_count=len(pairs), + candidate_digest=digest(pairs), + before_k_histogram=histogram([item["before_k"] for item in pairs]), + after_k_histogram=histogram(chosen), + structured_output_request_ids=structured_output_request_ids, + ) + super().update_draft_token_ids( + DraftTokenIds(req_ids=draft_token_ids.req_ids, draft_token_ids=after) + ) diff --git a/scripts/collectivespec/p0_overlay/collectivespec_p0/worker.py b/scripts/collectivespec/p0_overlay/collectivespec_p0/worker.py new file mode 100644 index 0000000..565933a --- /dev/null +++ b/scripts/collectivespec/p0_overlay/collectivespec_p0/worker.py @@ -0,0 +1,44 @@ +"""P0 worker: record actual rank-side target execution phases.""" + +from __future__ import annotations + +from typing import Any + +from vllm.v1.worker.gpu_worker import Worker + +from .common import log_event, plan_summary + + +class P0Worker(Worker): + """Log rank-local execution without modifying model or collective calls.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self._p0_execute_epoch = 0 + + def _p0_rank_fields(self) -> dict[str, Any]: + config = self.vllm_config.parallel_config + return { + "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), + "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), + "speculative_kmax": getattr( + getattr(self, "speculative_config", None), "num_speculative_tokens", None + ), + } + + def execute_model(self, scheduler_output: Any): # type: ignore[no-untyped-def] + self._p0_execute_epoch += 1 + fields = { + "worker_phase_epoch": self._p0_execute_epoch, + **self._p0_rank_fields(), + **plan_summary(scheduler_output), + } + log_event("worker", "target_execute_begin", **fields) + output = super().execute_model(scheduler_output) + log_event("worker", "target_execute_end", **fields) + return output diff --git a/scripts/collectivespec/run_p0_phase_trace.py b/scripts/collectivespec/run_p0_phase_trace.py new file mode 100644 index 0000000..11e8815 --- /dev/null +++ b/scripts/collectivespec/run_p0_phase_trace.py @@ -0,0 +1,298 @@ +#!/usr/bin/env python3 +"""Run the minimal CollectiveSpec P0 collective-phase premise experiment. + +P0 is deliberately *not* a performance experiment. It holds the EAGLE +drafter at Kmax=3, injects a pre-fixed per-request verification horizon after +candidate generation, and records the scheduler and worker-side phase traces. +The two policy cells are a constant-K=3 control and a deterministic +request-static heterogeneous policy. +""" + +from __future__ import annotations + +import argparse +import datetime as dt +import hashlib +import json +import os +import subprocess +import sys +from pathlib import Path +from typing import Any + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-spec", type=Path, required=True) + parser.add_argument("--source-root", type=Path, required=True) + parser.add_argument("--source-trace", type=Path, required=True) + parser.add_argument("--output-root", type=Path, required=True) + parser.add_argument("--run-id", required=True) + parser.add_argument("--request-count", type=int, default=64) + parser.add_argument("--completion-tokens", type=int, default=64) + parser.add_argument("--port", type=int, required=True) + parser.add_argument("--seed", type=int, default=20260713) + parser.add_argument("--dry-run", action="store_true") + return parser.parse_args() + + +def read_rows(source_trace: Path, count: int) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + with source_trace.open("r", encoding="utf-8") as handle: + for raw in handle: + if not raw.strip(): + continue + row = json.loads(raw) + if not isinstance(row, dict): + continue + rows.append(row) + if len(rows) == count: + break + if len(rows) != count: + raise SystemExit(f"source trace contains {len(rows)} usable rows, need {count}") + return rows + + +def write_workload(args: argparse.Namespace) -> tuple[Path, Path]: + rows = read_rows(args.source_trace, args.request_count) + trace_path = args.output_root / "p0_trace.jsonl" + with trace_path.open("w", encoding="utf-8") as handle: + for index, row in enumerate(rows): + row = dict(row) + # A single burst makes both DP schedulers exercise heterogeneous + # verification plans concurrently. P0 does not measure QPS. + row["request_id"] = f"p0-{index:04d}" + row["sampling_u"] = 0.0 + row["temperature"] = 0.0 + row["timestamp"] = 0.0 + handle.write(json.dumps(row, ensure_ascii=False, separators=(",", ":")) + "\n") + windows_path = args.output_root / "p0_windows.json" + windows = { + "kind": "collectivespec_p0_fixed_burst", + "source_trace": str(args.source_trace), + "source_trace_sha256": sha256(args.source_trace), + "windows": [ + { + "window_id": "collectivespec_p0_burst", + "trace_file": str(trace_path), + "trace_type": "decode_only", + "window_start": 0.0, + "window_end": 1.0, + "num_requests": args.request_count, + "sampling_strategy": "all_requests_fixed_burst", + } + ], + } + windows_path.write_text( + json.dumps(windows, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + return trace_path, windows_path + + +def derived_spec( + base: dict[str, Any], + *, + args: argparse.Namespace, + windows_path: Path, + policy_name: str, +) -> dict[str, Any]: + spec = json.loads(json.dumps(base)) + spec["study_id"] = f"collectivespec-p0-{policy_name}-{args.run_id}" + spec["engine"]["cwd"] = str(args.source_root) + spec["engine"]["port"] = args.port + flags = spec["engine"]["base_flags"] + flags["port"] = args.port + flags["seed"] = args.seed + flags["scheduler-cls"] = "collectivespec_p0.scheduler.P0Scheduler" + flags["worker-cls"] = "collectivespec_p0.worker.P0Worker" + raw_speculative = flags.get("speculative-config") + if not isinstance(raw_speculative, str): + raise SystemExit("base spec must have speculative-config for P0") + speculative = json.loads(raw_speculative) + speculative["num_speculative_tokens"] = 3 + flags["speculative-config"] = json.dumps(speculative, separators=(",", ":")) + + trace = spec["trace"] + trace.update( + { + "completion_tokens_override": args.completion_tokens, + "early_stop_max_elapsed_s": 1800.0, + "early_stop_max_lag_s": None, + "max_concurrency": args.request_count, + "max_requests_per_probe": None, + "replay_time_scale": 1.0, + "request_mode": "decode_only", + "timestamp_field": "timestamp", + "u_field": "sampling_u", + "window_id": "collectivespec_p0_burst", + "windows_path": str(windows_path), + } + ) + spec["slo"] = {"target_pass_rate": 1.0} + spec["search"] = { + "low": 0.0, + "high": 1.0, + "tolerance": 1e-6, + "max_probes": 1, + "sample_seed": args.seed, + } + return spec + + +def policy_payload(policy_name: str, seed: int) -> dict[str, Any]: + if policy_name == "control_k3": + return { + "version": "collectivespec-p0-v1-control-k3", + "seed": seed, + "max_k": 3, + "policy": "constant-k", + "constant_k": 3, + } + if policy_name == "heterogeneous": + return { + "version": "collectivespec-p0-v1-heterogeneous", + "seed": seed, + "max_k": 3, + "policy": "blake2b-request-static-k", + } + raise ValueError(f"unknown P0 policy {policy_name}") + + +def main() -> int: + args = parse_args() + if args.request_count < 8: + raise SystemExit("request-count must be at least 8 to exercise all K values") + if args.completion_tokens < 8: + raise SystemExit("completion-tokens must be at least 8") + if not args.base_spec.is_file(): + raise SystemExit(f"base spec does not exist: {args.base_spec}") + if not args.source_root.is_dir(): + raise SystemExit(f"source root does not exist: {args.source_root}") + if not args.source_trace.is_file(): + raise SystemExit(f"source trace does not exist: {args.source_trace}") + overlay = args.source_root / "scripts" / "collectivespec" / "p0_overlay" + if not overlay.is_dir(): + raise SystemExit(f"P0 overlay does not exist: {overlay}") + + args.output_root.mkdir(parents=True, exist_ok=True) + trace_path, windows_path = write_workload(args) + base = json.loads(args.base_spec.read_text(encoding="utf-8")) + manifest = { + "kind": "collectivespec_p0_phase_trace", + "created_at_utc": dt.datetime.now(dt.timezone.utc).isoformat(), + "base_spec": str(args.base_spec), + "base_spec_sha256": sha256(args.base_spec), + "source_root": str(args.source_root), + "source_trace": str(args.source_trace), + "source_trace_sha256": sha256(args.source_trace), + "materialized_trace": str(trace_path), + "materialized_trace_sha256": sha256(trace_path), + "windows_path": str(windows_path), + "overlay": str(overlay), + "run_id": args.run_id, + "topology_claim": "inherited_from_base_spec", + "request_count": args.request_count, + "completion_tokens": args.completion_tokens, + "policies": ["control_k3", "heterogeneous"], + "seed": args.seed, + "port": args.port, + "python": sys.executable, + } + (args.output_root / "manifest.json").write_text( + json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + + results: list[dict[str, Any]] = [] + for ordinal, policy_name in enumerate(manifest["policies"], start=1): + cell = args.output_root / policy_name + cell.mkdir(exist_ok=True) + policy_path = cell / "policy.json" + policy_path.write_text( + json.dumps(policy_payload(policy_name, args.seed), indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + spec_path = cell / "study_spec.json" + spec_path.write_text( + json.dumps( + derived_spec( + base, + args=args, + windows_path=windows_path, + policy_name=policy_name, + ), + indent=2, + sort_keys=True, + ) + + "\n", + encoding="utf-8", + ) + store_root = cell / "store" + logs = cell / "p0_logs" + command = [ + sys.executable, + "-m", + "aituner.cli", + "study", + "tune", + "--spec", + str(spec_path), + "--store-root", + str(store_root), + "--max-trials", + "1", + ] + record = { + "ordinal": ordinal, + "policy": policy_name, + "policy_path": str(policy_path), + "spec_path": str(spec_path), + "p0_logs": str(logs), + "command": command, + } + print(json.dumps(record, sort_keys=True), flush=True) + if args.dry_run: + results.append({**record, "returncode": None}) + continue + env = os.environ.copy() + inherited = env.get("PYTHONPATH", "") + env["PYTHONPATH"] = os.pathsep.join( + part for part in (str(overlay), str(args.source_root / "src"), inherited) if part + ) + env["PYTHONDONTWRITEBYTECODE"] = "1" + env["COLLECTIVESPEC_P0_LOG_DIR"] = str(logs) + env["COLLECTIVESPEC_P0_POLICY_PATH"] = str(policy_path) + log_path = cell / "driver.log" + with log_path.open("w", encoding="utf-8") as handle: + completed = subprocess.run( + command, + cwd=args.source_root, + env=env, + text=True, + stdout=handle, + stderr=subprocess.STDOUT, + ) + results.append({**record, "driver_log": str(log_path), "returncode": completed.returncode}) + if completed.returncode: + break + result = { + "finished_at_utc": dt.datetime.now(dt.timezone.utc).isoformat(), + "results": results, + "status": "ok" if results and all(item["returncode"] in (0, None) for item in results) else "failed", + } + (args.output_root / "driver_result.json").write_text( + json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + return 0 if result["status"] == "ok" else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/collectivespec/summarize_p0_phase_trace.py b/scripts/collectivespec/summarize_p0_phase_trace.py new file mode 100644 index 0000000..20c46eb --- /dev/null +++ b/scripts/collectivespec/summarize_p0_phase_trace.py @@ -0,0 +1,293 @@ +#!/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 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"] + 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") + dp_rank = event.get("data_parallel_rank") + if rank is None: + continue + key = str(rank) + per_rank[key].append(event) + if dp_rank is not None: + per_dp[str(dp_rank)].add(key) + 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, key=int) + 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_global_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) + ] + worker_metadata = [ + event.get("scheduled_dp_metadata") + for event in worker_events + if event.get("event") == "target_execute_begin" + ] + 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, + "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), + "dp_metadata_present_count": metadata_present, + "dp_metadata_missing_count": len(worker_metadata) - metadata_present, + }, + } + + +def data_sanity(cells: dict[str, dict[str, Any]]) -> 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() + ), + }, + } + + +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)} + 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())