aituner/scripts/collectivespec/p0_overlay/collectivespec_p0/common.py

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