Add CollectiveSpec P0 phase trace harness
This commit is contained in:
@@ -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.
|
||||
"""
|
||||
158
scripts/collectivespec/p0_overlay/collectivespec_p0/common.py
Normal file
158
scripts/collectivespec/p0_overlay/collectivespec_p0/common.py
Normal file
@@ -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")
|
||||
@@ -0,0 +1,6 @@
|
||||
{
|
||||
"max_k": 3,
|
||||
"policy": "blake2b-request-static-k",
|
||||
"seed": 20260713,
|
||||
"version": "collectivespec-p0-v1"
|
||||
}
|
||||
@@ -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)
|
||||
)
|
||||
@@ -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
|
||||
298
scripts/collectivespec/run_p0_phase_trace.py
Normal file
298
scripts/collectivespec/run_p0_phase_trace.py
Normal file
@@ -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())
|
||||
293
scripts/collectivespec/summarize_p0_phase_trace.py
Normal file
293
scripts/collectivespec/summarize_p0_phase_trace.py
Normal file
@@ -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())
|
||||
Reference in New Issue
Block a user