"""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 ``//`` 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