src/pns/common.py

"""Shared constants, paths and IO utilities.

Every path is environment-overridable, and every default is relative to this
repository, so nothing depends on the machine the study was run on.

    PNS_DATA       checkout of the companion dataset repo (default: ./data)
    PNS_CKPT       published checkpoints              (default: ./checkpoints)
    PNS_EVAL       where evaluation output is written (default: ./results/eval)
    PNS_TOKENIZER  tokenizer.json  (default: $PNS_DATA/benchmark/tokenizer.json)

`PNS_DATA` must point at the root of the `pns-world` dataset repo, i.e. the
directory containing `data/`, `benchmark/` and `manifests/`.

GPU selection is left entirely to `CUDA_VISIBLE_DEVICES`; nothing here binds a
device.
"""
from __future__ import annotations

import hashlib
import json
import os
import tempfile
from pathlib import Path

# ---------------------------------------------------------------- paths
PNS_ROOT = Path(os.environ.get("PNS_ROOT", Path(__file__).resolve().parents[2]))


def data_root() -> Path:
    """Root of the pns-world dataset checkout (the dir holding data/ and
    benchmark/). Set PNS_DATA to point at it; otherwise it is located by
    looking in the usual places relative to this repository."""
    env = os.environ.get("PNS_DATA")
    if env:
        return Path(env)
    for cand in (PNS_ROOT / "data", PNS_ROOT, PNS_ROOT.parent / "pns-world"):
        if (cand / "data" / "e1_original").is_dir() or (cand / "manifests").is_dir():
            return cand
    return PNS_ROOT / "data"


def shards_root() -> Path:
    """Directory holding the split directories (see pns.data.view.SPLIT_DIRS)."""
    return data_root() / "data"


def eval_root() -> Path:
    return Path(os.environ.get("PNS_EVAL", PNS_ROOT / "results" / "eval"))


def ckpt_root() -> Path:
    return Path(os.environ.get("PNS_CKPT", PNS_ROOT / "checkpoints"))


def logs_root() -> Path:
    return Path(os.environ.get("PNS_LOGS", PNS_ROOT / "logs"))


def tokenizer_path() -> Path:
    return Path(os.environ.get("PNS_TOKENIZER",
                               data_root() / "benchmark" / "tokenizer.json"))


# ---------------------------------------------------------------- atomic io

def atomic_write_json(path: Path | str, obj) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, tmp = tempfile.mkstemp(dir=path.parent, prefix=path.name + ".tmp")
    try:
        with os.fdopen(fd, "w") as f:
            json.dump(obj, f, indent=1, sort_keys=True, default=str)
            f.flush()
            os.fsync(f.fileno())
        os.replace(tmp, path)
        dirfd = os.open(path.parent, os.O_RDONLY)
        try:
            os.fsync(dirfd)
        finally:
            os.close(dirfd)
    except BaseException:
        if os.path.exists(tmp):
            os.unlink(tmp)
        raise


def sha256_file(path: Path | str) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def sha256_text(text: str) -> str:
    return hashlib.sha256(text.encode()).hexdigest()


# ---------------------------------------------------- metadata firewall classes
# Keys a model forward() may consume vs keys that exist only for supervision or
# for eval bucketing.  tests assert disjointness and that the model consumes
# only MODEL_INPUT_KEYS.
MODEL_INPUT_KEYS = frozenset({
    "tokens", "ev_bounds", "etype", "dt",
    "rec_store", "rec_kind", "rec_ent", "rec_key", "rec_val_toks", "rec_key_toks",
    "rec_birth_ev", "cache_live", "cache_adds",
    "enum_legal",  # the legal answer set is part of the question, runtime-visible
})
SUPERVISION_KEYS = frozenset({
    "mode_gold", "family", "enum_gold", "ptr_gold_slot", "op_gold", "op_arg_slots",
    "gold_val_hash",
})
EVAL_METADATA_KEYS = frozenset({
    "delay", "kappa", "gold_age_rank", "chain_len", "chain_pos", "corrected",
    "world_seed", "lifetime_id", "q_clock",
})

SCHEMA_VERSION = "PNS_V1"