src/cism/compile_cache.py

"""On-disk cache helper for compiled-execution manifests (Phase 0).

No execution-path code lives here: canonical config rendering plus
atomic manifest read/write helpers only.
"""

from __future__ import annotations

import json
import os
import re
import struct
from pathlib import Path
from typing import Union

MANIFEST_FILENAME = "manifest.json"

_SANITIZE_RE = re.compile(r"[^A-Za-z0-9_.-]")

PathLike = Union[str, os.PathLike]


def _sanitize_model_id(model_id: str) -> str:
    return _SANITIZE_RE.sub("_", model_id)


def canonical_config(config: dict) -> str:
    """Render a model config canonically for compile-key hashing.

    Legacy sorted-JSON rendering for arbitrary dicts (used by early
    Phase-0 tests). For validated model configs prefer
    :func:`canonical_model_config`, which matches the native 12-field
    pipe string in ``Model::refresh_compile_key`` exactly.
    """
    return json.dumps(config, sort_keys=True, separators=(",", ":"), default=str)


def _fp32_hex(value: object) -> str:
    """Hexfloat of the FP32-rounded value, matching ``std::hexfloat``.

    Native ``Config.eps/rope_theta`` are ``float``; Python floats are F64,
    so round-trip through FP32 first (``struct`` avoids a NumPy import).
    """
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        raise TypeError("rms_norm_eps and rope_theta must be real numbers")
    rounded = struct.unpack("f", struct.pack("f", float(value)))[0]
    return rounded.hex()


def canonical_model_config(config: dict) -> str:
    """Render a validated model config as the native 12-field pipe string.

    Field order mirrors ``native/runtime.cpp`` ``refresh_compile_key``:
    ``model_type|hidden|intermediate|layers|heads|kv_heads|head_dim|
    vocab|context|hexfloat(eps)|hexfloat(rope_theta)|tied``.
    Keys are the post-``validate_config`` HF names (``hidden_size``,
    ``num_hidden_layers``, ...). ``tie_word_embeddings`` renders as 1/0.
    """
    try:
        model_type = config["model_type"]
        hidden = config["hidden_size"]
        intermediate = config["intermediate_size"]
        layers = config["num_hidden_layers"]
        heads = config["num_attention_heads"]
        kv_heads = config["num_key_value_heads"]
        head_dim = config["head_dim"]
        vocab = config["vocab_size"]
        context = config["max_position_embeddings"]
        eps = config["rms_norm_eps"]
        theta = config["rope_theta"]
        tied = config["tie_word_embeddings"]
    except KeyError as exc:
        raise KeyError(f"canonical_model_config missing validated field: {exc}") from exc
    if not isinstance(model_type, str) or not model_type:
        raise TypeError("model_type must be a nonempty string")
    for name, value in (("hidden_size", hidden), ("intermediate_size", intermediate),
                        ("num_hidden_layers", layers), ("num_attention_heads", heads),
                        ("num_key_value_heads", kv_heads), ("head_dim", head_dim),
                        ("vocab_size", vocab), ("max_position_embeddings", context)):
        if type(value) is not int or value <= 0:
            raise TypeError(f"{name} must be a positive integer")
    if type(tied) is not bool:
        raise TypeError("tie_word_embeddings must be a boolean")
    return "|".join((model_type, str(hidden), str(intermediate), str(layers),
                     str(heads), str(kv_heads), str(head_dim), str(vocab),
                     str(context), _fp32_hex(eps), _fp32_hex(theta),
                     "1" if tied else "0"))


def cache_dir_for(base_dir: PathLike, model_id: str, revision: object) -> Path:
    """Return ``<base_dir>/<sanitized model_id>/<revision>`` as a Path."""
    rev = "none" if revision is None else str(revision)
    return Path(base_dir) / _sanitize_model_id(str(model_id)) / rev


def artifact_dir(*args, **kwargs) -> Path:
    """Return the artifact directory for a compile plan.

    Supported forms (positional or keyword):

    * ``artifact_dir(cache_dir, key)`` — join an existing cache dir with key.
    * ``artifact_dir(base_dir, model_id, revision, key)`` — build from
      :func:`cache_dir_for` then join with key.
    """
    # Keyword aliases.
    cache_dir = kwargs.pop("cache_dir", None)
    key = kwargs.pop("key", None)
    base_dir = kwargs.pop("base_dir", None)
    model_id = kwargs.pop("model_id", None)
    revision = kwargs.pop("revision", None)
    if kwargs:
        raise TypeError(f"unexpected keyword arguments: {sorted(kwargs)}")

    positional = list(args)

    if cache_dir is not None and key is not None and base_dir is None and model_id is None:
        if positional:
            raise TypeError("artifact_dir() called with both positional and keyword cache args")
        return Path(cache_dir) / str(key)

    if base_dir is not None or model_id is not None or revision is not None:
        # Keyword 4-arg form.
        if positional:
            raise TypeError("artifact_dir() called with both positional and keyword args")
        if base_dir is None or model_id is None or key is None:
            raise TypeError("artifact_dir() 4-arg form needs base_dir, model_id, revision, key")
        rev = revision  # may be None; cache_dir_for handles it
        return cache_dir_for(base_dir, model_id, rev) / str(key)

    # Positional-only handling.
    if len(positional) == 2:
        cache, k = positional
        return Path(cache) / str(k)
    if len(positional) == 4:
        base, mid, rev, k = positional
        return cache_dir_for(base, mid, rev) / str(k)
    raise TypeError(
        "artifact_dir() expects (cache_dir, key) or (base_dir, model_id, revision, key)"
    )


def write_artifact(directory: PathLike, manifest: str) -> Path:
    """Atomically write ``manifest`` into ``directory/manifest.json``.

    Creates parents, writes to a temp file in the same directory, then
    ``os.replace`` for atomicity. Returns the artifact file path.
    """
    target_dir = Path(directory)
    target_dir.mkdir(parents=True, exist_ok=True)
    target = target_dir / MANIFEST_FILENAME
    tmp_name = f".{MANIFEST_FILENAME}.tmp-{os.getpid()}"
    tmp_path = target_dir / tmp_name
    # Write + fsync + atomic replace; cleanup temp on failure.
    try:
        with open(tmp_path, "w", encoding="utf-8", newline="\n") as handle:
            handle.write(manifest)
            handle.flush()
            try:
                os.fsync(handle.fileno())
            except OSError:
                pass
        os.replace(tmp_path, target)
    finally:
        try:
            if tmp_path.exists():
                tmp_path.unlink()
        except OSError:
            pass
    return target


def read_artifact(directory: PathLike) -> Union[str, None]:
    """Read ``directory/manifest.json``; return str or None on any error.

    Missing files, undecodable bytes, and invalid-JSON payloads are cache
    misses (None) and never raise.
    """
    try:
        target = Path(directory) / MANIFEST_FILENAME
        raw = target.read_bytes()
        text = raw.decode("utf-8")
        # Corrupt payloads (e.g. garbage bytes that happen to decode) are
        # misses too: manifests must be valid JSON documents.
        json.loads(text)
        return text
    except Exception:
        return None