#!/usr/bin/env python3 """Token-exact fixed-duration client for the OpProf Phase-3 protocol.""" from __future__ import annotations import argparse import asyncio import gzip import hashlib import json import math import os import time from dataclasses import dataclass from pathlib import Path from typing import Any import aiohttp SCHEMA = 1 TOKEN_BASE = 1000 TOKEN_SPAN = 100000 class ManifestExhausted(RuntimeError): pass def sha256_file(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as f: for chunk in iter(lambda: f.read(1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def atomic_json(path: Path, value: Any, mode: int = 0o640) -> None: path.parent.mkdir(parents=True, exist_ok=True) tmp = path.with_name(path.name + f".tmp.{os.getpid()}") fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_EXCL, mode) with os.fdopen(fd, "w", encoding="utf-8") as f: json.dump(value, f, sort_keys=True, indent=2) f.write("\n") f.flush() os.fsync(f.fileno()) os.replace(tmp, path) def atomic_jsonl(path: Path, rows: list[dict[str, Any]], mode: int = 0o640) -> None: path.parent.mkdir(parents=True, exist_ok=True) tmp = path.with_name(path.name + f".tmp.{os.getpid()}") fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_EXCL, mode) with os.fdopen(fd, "w", encoding="utf-8") as f: for row in rows: f.write(json.dumps(row, sort_keys=True, separators=(",", ":")) + "\n") f.flush() os.fsync(f.fileno()) os.replace(tmp, path) def parse_range(value: str) -> tuple[int, int]: lo_text, hi_text = value.split(":", 1) lo, hi = int(lo_text), int(hi_text) if lo <= 0 or hi < lo: raise argparse.ArgumentTypeError(f"invalid positive range: {value}") return lo, hi def _integer_counts(weights: list[float], total: int) -> list[int]: raw = [w * total for w in weights] counts = [math.floor(x) for x in raw] order = sorted( range(len(raw)), key=lambda i: raw[i] - counts[i], reverse=True ) for idx in order[: total - sum(counts)]: counts[idx] += 1 return counts def numeric_sanity(values: list[float | int]) -> dict[str, Any]: finite = [float(x) for x in values if math.isfinite(float(x))] return { "n": len(values), "finite_n": len(finite), "missing_n": len(values) - len(finite), "min": min(finite) if finite else None, "max": max(finite) if finite else None, "distinct_n": len(set(finite)), } def manifest_summary(rows: list[dict[str, Any]]) -> dict[str, Any]: return { "schema": SCHEMA, "rows": len(rows), "input_tokens": numeric_sanity([int(r["input_tokens"]) for r in rows]), "output_tokens": numeric_sanity([int(r["output_tokens"]) for r in rows]), "arrival_values": sorted({str(r["arrival"]) for r in rows}), "pattern_values": sorted({str(r["pattern_id"]) for r in rows}), } def materialize(args: argparse.Namespace) -> dict[str, Any]: import numpy as np rng = np.random.default_rng(args.workload_seed) n = args.num_requests if args.kind == "prefix-pool": if args.num_prefixes <= 0 or args.prefix_len <= 0 or args.suffix_fixed <= 0: raise ValueError("prefix-pool requires positive pool/prefix/suffix") lengths = np.full(n, args.prefix_len + args.suffix_fixed, dtype=np.int64) prefix_ids = np.arange(n, dtype=np.int64) % args.num_prefixes rng.shuffle(prefix_ids) else: prefix_ids = np.full(n, -1, dtype=np.int64) if args.input_uniform: lo, hi = parse_range(args.input_uniform) lengths = rng.integers(lo, hi + 1, n, dtype=np.int64) elif args.input_fixed: lengths = np.full(n, args.input_fixed, dtype=np.int64) elif args.input_mixture: spec = json.loads(args.input_mixture) if not isinstance(spec, dict) or not spec: raise ValueError("input mixture must be a non-empty JSON object") keys = list(spec) weights = [float(spec[key]) for key in keys] if any(w < 0 for w in weights) or not math.isclose(sum(weights), 1.0): raise ValueError("mixture weights must be non-negative and sum to 1") pieces = [] for key, count in zip( keys, _integer_counts(weights, n), strict=True ): kind, lo_text, hi_text = key.split(":") if kind != "uniform": raise ValueError(f"unsupported mixture component: {key}") pieces.append( rng.integers( int(lo_text), int(hi_text) + 1, count, dtype=np.int64 ) ) lengths = np.concatenate(pieces) rng.shuffle(lengths) else: raise ValueError("exactly one input distribution is required") if args.output_fixed <= 0 or args.arrival not in {"steady", "burst:8"}: raise ValueError("invalid output length or arrival class") rows = [] for i in range(n): row = { "schema": SCHEMA, "request_id": f"{args.id}-{i:05d}", "pattern_id": args.id, "kind": args.kind, "input_tokens": int(lengths[i]), "output_tokens": args.output_fixed, "arrival": args.arrival, "token_seed": int(args.workload_seed * 1000003 + i), } if args.kind == "prefix-pool": row.update( { "prefix_id": int(prefix_ids[i]), "num_prefixes": args.num_prefixes, "prefix_tokens": args.prefix_len, } ) rows.append(row) out = Path(args.out) atomic_jsonl(out, rows, mode=0o600) summary = manifest_summary(rows) summary.update({"sha256": sha256_file(out), "path": str(out)}) atomic_json(out.with_suffix(out.suffix + ".summary.json"), summary, mode=0o600) print(json.dumps(summary, sort_keys=True)) return summary def materialize_private(args: argparse.Namespace) -> dict[str, Any]: from transformers import AutoTokenizer source = Path(args.source) selected: list[dict[str, Any]] = [] with source.open(encoding="utf-8") as f: for source_index, line in enumerate(f): row = json.loads(line) if ( float(row["sampling_u"]) <= args.sampling_u_max and int(row["input_length"]) <= args.max_input_tokens ): selected.append( { "schema": SCHEMA, "request_id": f"{args.id}-{len(selected):05d}", "pattern_id": args.id, "kind": "private-trace", "input_tokens": int(row["input_length"]), "output_tokens": min( int(row["output_length"]), args.output_cap ), "arrival": args.arrival, "source_index": source_index, "prompt": row["prompt"], } ) tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True) diffs = [ len(tokenizer.encode(row["prompt"], add_special_tokens=False)) - row["input_tokens"] for row in selected ] exact = sum(diff == 0 for diff in diffs) exact_fraction = exact / len(diffs) if diffs else 0.0 max_abs = max((abs(diff) for diff in diffs), default=-1) if exact_fraction < 0.99 or max_abs > 1: raise RuntimeError( "tokenizer parity gate failed: " f"exact_fraction={exact_fraction:.6f} max_abs_error={max_abs}" ) out = Path(args.out) atomic_jsonl(out, selected, mode=0o600) summary = manifest_summary(selected) summary.update( { "sha256": sha256_file(out), "source_sha256": sha256_file(source), "tokenizer_exact_n": exact, "tokenizer_exact_fraction": exact_fraction, "tokenizer_max_abs_error": max_abs, "path": str(out), } ) atomic_json(out.with_suffix(out.suffix + ".summary.json"), summary, mode=0o600) print(json.dumps(summary, sort_keys=True)) return summary def load_manifest(path: Path) -> list[dict[str, Any]]: rows = [json.loads(line) for line in path.read_text().splitlines() if line] required = { "request_id", "pattern_id", "input_tokens", "output_tokens", "arrival", } if not rows: raise ValueError("empty manifest") for row in rows: if not required.issubset(row): raise ValueError(f"manifest row lacks {sorted(required - set(row))}") if len({row["request_id"] for row in rows}) != len(rows): raise ValueError("duplicate request_id") return rows def _token_stream(seed: int, count: int) -> list[int]: state = seed & 0xFFFFFFFF out = [] for _ in range(count): state = (1664525 * state + 1013904223) & 0xFFFFFFFF out.append(TOKEN_BASE + state % TOKEN_SPAN) return out def synthetic_prompt(row: dict[str, Any]) -> list[int]: length = int(row["input_tokens"]) seed = int(row.get("token_seed", 0)) if row.get("kind") == "prefix-pool": prefix_n = int(row["prefix_tokens"]) tokens = _token_stream(0xA5A50000 + int(row["prefix_id"]), prefix_n) tokens += _token_stream(seed, length - prefix_n) offset = prefix_n else: tokens = _token_stream(seed, length) offset = 0 if length - offset >= 3: index = int(row["request_id"].rsplit("-", 1)[1]) tokens[offset : offset + 3] = [ TOKEN_BASE + index % 100, TOKEN_BASE + (index // 100) % 100, TOKEN_BASE + (index // 10000) % 100, ] return tokens @dataclass class RunContext: args: argparse.Namespace rows: list[dict[str, Any]] t0: float clean_end: float stop_event: asyncio.Event lock: asyncio.Lock next_index: int = 0 in_flight: int = 0 max_in_flight: int = 0 exhausted: bool = False admission_stop_s: float | None = None async def next_row(self) -> dict[str, Any]: async with self.lock: if self.next_index >= len(self.rows): self.exhausted = True raise ManifestExhausted( f"manifest exhausted after {self.next_index} admissions" ) row = self.rows[self.next_index] self.next_index += 1 return row async def request_one( ctx: RunContext, session: aiohttp.ClientSession, row: dict[str, Any], scheduled: float, ) -> dict[str, Any]: loop = asyncio.get_running_loop() admitted = loop.time() ctx.in_flight += 1 ctx.max_in_flight = max(ctx.max_in_flight, ctx.in_flight) status = 0 actual_output: int | None = None first_token: float | None = None error_kind: str | None = None try: prompt: str | list[int] = ( row["prompt"] if row.get("kind") == "private-trace" else synthetic_prompt(row) ) if not isinstance(prompt, str) and len(prompt) != int(row["input_tokens"]): raise AssertionError("synthetic prompt length drift") payload = { "model": ctx.args.model, "prompt": prompt, "max_tokens": int(row["output_tokens"]), "temperature": ctx.args.temperature, "ignore_eos": ctx.args.ignore_eos, "stream": True, "stream_options": {"include_usage": True}, "add_special_tokens": False, "seed": ctx.args.server_seed, } headers = { "Content-Type": "application/json", "x-request-id": str(row["request_id"]), } async with session.post( ctx.args.base_url.rstrip("/") + "/v1/completions", json=payload, headers=headers, ) as response: status = response.status if status != 200: error_kind = f"http_{status}" else: buf = b"" async for chunk in response.content.iter_any(): buf += chunk while b"\n" in buf: line, buf = buf.split(b"\n", 1) line = line.strip() if not line.startswith(b"data:"): continue data = line[5:].strip() if data == b"[DONE]": continue event = json.loads(data) if event.get("choices") and first_token is None: first_token = loop.time() if event.get("usage") is not None: actual_output = int(event["usage"]["completion_tokens"]) if actual_output is None: error_kind = "missing_usage" except (aiohttp.ClientError, asyncio.TimeoutError) as exc: error_kind = type(exc).__name__ except Exception as exc: error_kind = type(exc).__name__ finally: completed = loop.time() ctx.in_flight -= 1 success = ( status == 200 and error_kind is None and actual_output == int(row["output_tokens"]) ) if status == 200 and actual_output is not None and not success: error_kind = "output_token_mismatch" return { "schema": SCHEMA, "request_id": row["request_id"], "scheduled_s": scheduled - ctx.t0, "admitted_s": admitted - ctx.t0, "first_token_s": None if first_token is None else first_token - ctx.t0, "completed_s": completed - ctx.t0, "input_tokens": int(row["input_tokens"]), "requested_output_tokens": int(row["output_tokens"]), "actual_output_tokens": actual_output, "http_status": status, "success": success, "error_kind": error_kind, } async def saturation_load( ctx: RunContext, session: aiohttp.ClientSession ) -> list[dict[str, Any]]: results: list[dict[str, Any]] = [] async def worker() -> None: while not ctx.stop_event.is_set(): try: row = await ctx.next_row() except ManifestExhausted: ctx.stop_event.set() return results.append( await request_one(ctx, session, row, asyncio.get_running_loop().time()) ) tasks = [ asyncio.create_task(worker()) for _ in range(ctx.args.max_concurrency) ] await asyncio.gather(*tasks) return results async def finite_load( ctx: RunContext, session: aiohttp.ClientSession, rate: float ) -> list[dict[str, Any]]: sem = asyncio.Semaphore(ctx.args.max_concurrency) tasks: list[asyncio.Task[dict[str, Any]]] = [] batch = 8 if str(ctx.rows[0]["arrival"]) == "burst:8" else 1 period = batch / rate event_index = 0 async def limited(row: dict[str, Any], scheduled: float) -> dict[str, Any]: async with sem: return await request_one(ctx, session, row, scheduled) while not ctx.stop_event.is_set(): scheduled = ctx.t0 + event_index * period delay = scheduled - asyncio.get_running_loop().time() if delay > 0: try: await asyncio.wait_for(ctx.stop_event.wait(), timeout=delay) break except asyncio.TimeoutError: pass if ctx.stop_event.is_set(): break try: for _ in range(batch): tasks.append( asyncio.create_task(limited(await ctx.next_row(), scheduled)) ) except ManifestExhausted: ctx.stop_event.set() break event_index += 1 return await asyncio.gather(*tasks) if tasks else [] async def post_profile( session: aiohttp.ClientSession, base_url: str, endpoint: str ) -> tuple[float, float, int]: loop = asyncio.get_running_loop() before = loop.time() async with session.post(base_url.rstrip("/") + endpoint) as response: status = response.status await response.read() return before, loop.time(), status def _trace_loadable(path: Path) -> bool: try: opener = gzip.open if path.suffix == ".gz" else open with opener(path, "rt", encoding="utf-8") as f: parsed = json.load(f) return isinstance(parsed, dict) and isinstance(parsed.get("traceEvents"), list) except (OSError, EOFError, json.JSONDecodeError): return False async def wait_new_trace( trace_dir: Path, before: set[Path], timeout: float ) -> Path: deadline = asyncio.get_running_loop().time() + timeout while asyncio.get_running_loop().time() < deadline: for path in sorted(set(trace_dir.glob("*.pt.trace.json*")) - before): if _trace_loadable(path): return path await asyncio.sleep(0.25) raise TimeoutError(f"no new loadable trace within {timeout}s") async def timeline( ctx: RunContext, session: aiohttp.ClientSession ) -> list[dict[str, Any]]: args = ctx.args profiles: list[dict[str, Any]] = [] await asyncio.sleep(max(0, ctx.clean_end - asyncio.get_running_loop().time())) if args.profile_after_clean: trace_dir = Path(args.profile_trace_dir) for window in range(args.num_profile_windows): prior = set(trace_dir.glob("*.pt.trace.json*")) start_before, start_after, start_status = await post_profile( session, args.base_url, "/start_profile" ) trace = await wait_new_trace( trace_dir, prior, args.profile_timeout_seconds ) trace_ready = asyncio.get_running_loop().time() stop_before, stop_after, stop_status = await post_profile( session, args.base_url, "/stop_profile" ) profiles.append( { "window": window + 1, "start_call_s": start_before - ctx.t0, "start_return_s": start_after - ctx.t0, "trace_ready_s": trace_ready - ctx.t0, "stop_call_s": stop_before - ctx.t0, "stop_return_s": stop_after - ctx.t0, "start_status": start_status, "stop_status": stop_status, "trace_file": trace.name, "trace_sha256": sha256_file(trace), } ) if start_status != 200 or stop_status != 200: raise RuntimeError("profile endpoint returned non-200") await asyncio.sleep(args.recovery_seconds) else: await asyncio.sleep(args.post_clean_seconds) ctx.admission_stop_s = asyncio.get_running_loop().time() - ctx.t0 ctx.stop_event.set() return profiles def segment_summary( records: list[dict[str, Any]], start: float, end: float ) -> dict[str, Any]: admitted = [r for r in records if start <= r["admitted_s"] < end] completed = [r for r in records if start <= r["completed_s"] < end] successes = [r for r in completed if r["success"]] duration = end - start return { "start_s": start, "end_s": end, "duration_s": duration, "admitted": len(admitted), "completed": len(successes), "failed": len(completed) - len(successes), "offered_rps": len(admitted) / duration, "completed_throughput_rps": len(successes) / duration, "input_tokens": sum(r["input_tokens"] for r in successes), "output_tokens": sum(r["actual_output_tokens"] or 0 for r in successes), } async def run_load(args: argparse.Namespace) -> dict[str, Any]: manifest = Path(args.manifest) rows = load_manifest(manifest) arrivals = {row["arrival"] for row in rows} if len(arrivals) != 1: raise ValueError("a manifest must have one arrival class") if args.load_point == "saturation": if args.request_rate != "inf": raise ValueError("saturation requires --request-rate inf") rate = math.inf else: if not args.saturation_result: raise ValueError("moderate requires --saturation-result") sat = json.loads(Path(args.saturation_result).read_text()) rate = args.rate_fraction * float(sat["clean"]["completed_throughput_rps"]) if not math.isfinite(rate) or rate <= 0: raise ValueError("derived moderate rate must be positive and finite") loop = asyncio.get_running_loop() t0 = loop.time() t0_mono_ns = int(t0 * 1e9) t0_wall_ns = time.time_ns() clean_seconds = args.clean_segment_seconds * args.num_clean_segments ctx = RunContext( args=args, rows=rows, t0=t0, clean_end=t0 + args.warmup_seconds + clean_seconds, stop_event=asyncio.Event(), lock=asyncio.Lock(), ) timeout = aiohttp.ClientTimeout(total=None, connect=30, sock_read=600) connector = aiohttp.TCPConnector(limit=args.max_concurrency) control_connector = aiohttp.TCPConnector(limit=2) async with ( aiohttp.ClientSession(timeout=timeout, connector=connector) as session, aiohttp.ClientSession( timeout=timeout, connector=control_connector ) as control_session, ): profile_task = asyncio.create_task(timeline(ctx, control_session)) load_task = asyncio.create_task( saturation_load(ctx, session) if math.isinf(rate) else finite_load(ctx, session, rate) ) try: profiles = await profile_task except Exception: ctx.stop_event.set() await load_task raise records = await load_task clean_start = args.warmup_seconds clean_end = clean_start + clean_seconds clean = segment_summary(records, clean_start, clean_end) segments = [] for i in range(args.num_clean_segments): start = clean_start + i * args.clean_segment_seconds segments.append( { "name": chr(ord("A") + i), **segment_summary( records, start, start + args.clean_segment_seconds ), } ) successful = [r for r in records if r["success"]] elapsed_seconds = loop.time() - t0 if ctx.admission_stop_s is None: raise RuntimeError("admission stop timestamp was not recorded") drain_seconds = elapsed_seconds - ctx.admission_stop_s result = { "schema": SCHEMA, "manifest_sha256": sha256_file(manifest), "manifest_rows": len(rows), "manifest_admitted": ctx.next_index, "manifest_wrapped": False, "manifest_exhausted": ctx.exhausted, "load_point": args.load_point, "t0_mono_ns": t0_mono_ns, "t0_wall_ns": t0_wall_ns, "request_rate": "inf" if math.isinf(rate) else rate, "rate_fraction": None if math.isinf(rate) else args.rate_fraction, "arrival": next(iter(arrivals)), "warmup_seconds": args.warmup_seconds, "clean_segment_seconds": args.clean_segment_seconds, "num_clean_segments": args.num_clean_segments, "elapsed_seconds": elapsed_seconds, "admission_stop_s": ctx.admission_stop_s, "drain_seconds": drain_seconds, "max_in_flight": ctx.max_in_flight, "records": len(records), "successful_records": len(successful), "failed_records": len(records) - len(successful), "clean": clean, "segments": segments, "profiles": profiles, } sanity = { "schema": SCHEMA, "numeric": { "input_tokens": numeric_sanity([r["input_tokens"] for r in records]), "requested_output_tokens": numeric_sanity( [r["requested_output_tokens"] for r in records] ), "actual_output_tokens": numeric_sanity( [ r["actual_output_tokens"] for r in records if r["actual_output_tokens"] is not None ] ), "scheduled_s": numeric_sanity([r["scheduled_s"] for r in records]), "admitted_s": numeric_sanity([r["admitted_s"] for r in records]), "completed_s": numeric_sanity([r["completed_s"] for r in records]), }, "invariants": { "clean_duration_exact": math.isclose(clean["duration_s"], clean_seconds), "segment_count_exact": len(segments) == args.num_clean_segments, "manifest_no_wrap": ctx.next_index <= len(rows), "manifest_not_exhausted": not ctx.exhausted, "concurrency_bounded": ctx.max_in_flight <= args.max_concurrency, "drain_within_timeout": drain_seconds <= args.drain_timeout_seconds, "output_tokens_exact": all( r["actual_output_tokens"] == r["requested_output_tokens"] for r in successful ), "clean_failures_zero": clean["failed"] == 0, "profile_count_exact": len(profiles) == (args.num_profile_windows if args.profile_after_clean else 0), "profile_status_ok": all( p["start_status"] == 200 and p["stop_status"] == 200 for p in profiles ), }, } if not math.isinf(rate): sanity["invariants"]["moderate_offered_within_5pct"] = ( abs(clean["offered_rps"] / rate - 1) <= 0.05 ) out = Path(args.result_dir) out.mkdir(parents=True, exist_ok=True) atomic_jsonl(out / "requests.jsonl", sorted(records, key=lambda r: r["admitted_s"])) atomic_jsonl(out / "segments.jsonl", segments) atomic_json(out / "result.json", result) atomic_json(out / "sanity.json", sanity) if ctx.exhausted: raise ManifestExhausted("manifest exhausted; result retained for diagnosis") failed = [name for name, ok in sanity["invariants"].items() if not ok] if failed: raise RuntimeError(f"client sanity failure: {failed}") return result def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser() sub = parser.add_subparsers(dest="command", required=True) mat = sub.add_parser("materialize") mat.add_argument("--id", required=True) mat.add_argument("--kind", choices=("synthetic", "prefix-pool"), required=True) group = mat.add_mutually_exclusive_group() group.add_argument("--input-uniform") group.add_argument("--input-fixed", type=int) group.add_argument("--input-mixture") mat.add_argument("--output-fixed", type=int, required=True) mat.add_argument("--prefix", default="none") mat.add_argument("--arrival", required=True) mat.add_argument("--num-requests", type=int, required=True) mat.add_argument("--workload-seed", type=int, required=True) mat.add_argument("--num-prefixes", type=int, default=0) mat.add_argument("--prefix-len", type=int, default=0) mat.add_argument("--suffix-fixed", type=int, default=0) mat.add_argument("--out", required=True) private = sub.add_parser("materialize-private") private.add_argument("--id", required=True) private.add_argument("--source", required=True) private.add_argument("--sampling-u-max", type=float, required=True) private.add_argument("--max-input-tokens", type=int, required=True) private.add_argument("--output-cap", type=int, required=True) private.add_argument("--preserve-prompts", action="store_true", required=True) private.add_argument("--disable-shuffle", action="store_true", required=True) private.add_argument("--arrival", required=True) private.add_argument("--model", required=True) private.add_argument("--out", required=True) run = sub.add_parser("run") run.add_argument("--manifest", required=True) run.add_argument("--base-url", required=True) run.add_argument("--model", required=True) run.add_argument("--load-point", choices=("saturation", "moderate"), required=True) run.add_argument("--request-rate") run.add_argument("--saturation-result") run.add_argument("--rate-fraction", type=float, default=0.60) run.add_argument("--max-concurrency", type=int, default=256) run.add_argument("--ignore-eos", action="store_true") run.add_argument("--temperature", type=float, default=0.0) run.add_argument("--warmup-seconds", type=float, default=60) run.add_argument("--clean-segment-seconds", type=float, default=80) run.add_argument("--num-clean-segments", type=int, default=3) run.add_argument("--profile-after-clean", action="store_true") run.add_argument("--num-profile-windows", type=int, default=0) run.add_argument("--profile-warmup-iterations", type=int, default=2) run.add_argument("--profile-active-iterations", type=int, default=8) run.add_argument("--profile-trace-dir") run.add_argument("--profile-timeout-seconds", type=float, default=120) run.add_argument("--recovery-seconds", type=float, default=30) run.add_argument("--post-clean-seconds", type=float, default=0) run.add_argument("--drain-timeout-seconds", type=float, default=120) run.add_argument("--workload-seed", type=int, default=20260712) run.add_argument("--server-seed", type=int, default=20260712) run.add_argument("--result-dir", required=True) return parser def main() -> None: args = build_parser().parse_args() if args.command == "materialize": materialize(args) elif args.command == "materialize-private": materialize_private(args) else: if args.profile_after_clean and not args.profile_trace_dir: raise ValueError("--profile-after-clean requires --profile-trace-dir") print(json.dumps(asyncio.run(run_load(args)), sort_keys=True)) if __name__ == "__main__": main()