Add CollectiveSpec P0 phase trace harness

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")