blut-core 0.1.0__py3-none-any.whl

blut_core/__init__.py

@@ -0,0 +1,54 @@
+"""Shared Python runtime and ingredient registry for BLUT cookbooks.
+
+The package contains domain-agnostic building blocks reusable by any cookbook:
+
+- ``async_io``     — bounded, backpressured delivery for persistence work
+- ``runctx``       — run identity + filesystem anchors from the BLUT env (P10)
+- ``MetricLog``    — atomic per-epoch CSV/Parquet metric writer
+- ``read_metric``  — verbatim metric/log reader (no LLM in the path; ADR 0038)
+- ``status``       — emit/read the StatusUpdate wire protocol (P4)
+- ``RunManifest``  — run provenance, written even on crash (P8)
+- ``checkpoint``   — corruption-safe save/load + resume payload contract (P7)
+- ``sysgauge``     — best-effort GPU/host snapshot for the metric stream (P10)
+- ingredient specs and builders used by generic training recipes
+
+``torch`` is lazy-imported only inside ``checkpoint``/``sysgauge`` functions, so
+``import blut_core`` stays cheap and dependency-light. See ``blut/docs/metrics.md``
+and decisions/0037, 0044.
+"""
+
+from blut_core.registry import (
+    build_ingredient,
+    get_spec,
+    list_ingredients,
+    register_ingredient,
+)
+from blut_core.spec import KINDS, IngredientSpec
+
+# Register generic ingredient specs on package import.
+from blut_core.ingredients import _specs as _ingredient_specs  # noqa: F401
+
+from . import async_io, checkpoint, runctx, status, sysgauge
+from .metric_log import MetricLog
+from .run_manifest import RunManifest
+
+# read_metric is a CLI module (`python -m blut_core.read_metric`); it is NOT
+# eager-imported here (that triggers a runpy double-import warning under -m).
+# `from blut_core import read_metric` still works (submodule import).
+
+__all__ = [
+    "IngredientSpec",
+    "KINDS",
+    "MetricLog",
+    "RunManifest",
+    "async_io",
+    "build_ingredient",
+    "checkpoint",
+    "get_spec",
+    "list_ingredients",
+    "read_metric",
+    "register_ingredient",
+    "runctx",
+    "status",
+    "sysgauge",
+]

blut_core/async_io.py

@@ -0,0 +1,246 @@
+"""Bounded, backpressured asynchronous I/O for cookbook runtimes.
+
+Bounded workers are daemons so an exceptional interpreter shutdown cannot hang
+on an idle worker. Successful callers must still call :meth:`AsyncSink.close`;
+that is the only boundary that drains accepted work and surfaces worker errors.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from queue import Queue
+from threading import BoundedSemaphore, Condition, Thread, local
+from typing import Callable, Generic, TypeVar, cast
+
+
+__all__ = [
+    "AsyncSink",
+    "AsyncSinkWorkerError",
+    "Bounded",
+    "Inline",
+    "ItemTooLargeError",
+    "SinkClosedError",
+]
+
+
+ItemT = TypeVar("ItemT")
+
+
+def _identity(item: ItemT) -> ItemT:
+    return item
+
+
+class ItemTooLargeError(ValueError):
+    """A submitted item exceeds the sink's declared retention limit."""
+
+    def __init__(self, size_bytes: int, max_item_bytes: int) -> None:
+        self.size_bytes = size_bytes
+        self.max_item_bytes = max_item_bytes
+        super().__init__(f"item is {size_bytes} bytes; limit is {max_item_bytes}")
+
+
+class SinkClosedError(RuntimeError):
+    """Work was submitted after sink shutdown began."""
+
+
+class AsyncSinkWorkerError(RuntimeError):
+    """A sink worker failed while delivering an accepted item."""
+
+    def __init__(self, worker_error_type: str, message: str) -> None:
+        self.worker_error_type = worker_error_type
+        super().__init__(message)
+
+
+@dataclass(frozen=True)
+class _WorkerFailure:
+    worker_error_type: str
+    message: str
+
+    @classmethod
+    def capture(cls, error: BaseException) -> _WorkerFailure:
+        error_type = type(error)
+        worker_error_type = f"{error_type.__module__}.{error_type.__qualname__}"
+        try:
+            message = str(error)
+        except BaseException:
+            message = "worker failure with an unprintable message"
+        return cls(worker_error_type=worker_error_type, message=message)
+
+    def to_exception(self) -> AsyncSinkWorkerError:
+        return AsyncSinkWorkerError(self.worker_error_type, self.message)
+
+
+@dataclass(frozen=True)
+class Inline:
+    """Run the installed handler synchronously on the submitting thread."""
+
+
+@dataclass(frozen=True)
+class Bounded:
+    """Run work on one worker with bounded retained residency."""
+
+    capacity: int
+    max_item_bytes: int
+
+    def __post_init__(self) -> None:
+        if isinstance(self.capacity, bool) or not isinstance(self.capacity, int):
+            raise TypeError("capacity must be an integer")
+        if isinstance(self.max_item_bytes, bool) or not isinstance(
+            self.max_item_bytes, int
+        ):
+            raise TypeError("max_item_bytes must be an integer")
+        if self.capacity <= 0:
+            raise ValueError("capacity must be greater than zero")
+        if self.max_item_bytes <= 0:
+            raise ValueError("max_item_bytes must be greater than zero")
+
+
+_STOP = object()
+
+
+class AsyncSink(Generic[ItemT]):
+    """Deliver items to one handler under an explicit, close-required mode."""
+
+    def __init__(
+        self,
+        handler: Callable[[ItemT], None],
+        mode: Inline | Bounded,
+        *,
+        item_size: Callable[[ItemT], int] | None = None,
+        retain: Callable[[ItemT], ItemT] = _identity,
+        retained_size: Callable[[ItemT], int] | None = None,
+        thread_name: str = "blut-async-io",
+    ) -> None:
+        self._handler = handler
+        self._mode = mode
+        self._item_size = item_size
+        self._retain = retain
+        self._retained_size = retained_size
+        self._closed = False
+        self._failure: _WorkerFailure | None = None
+        self._state = Condition()
+        self._active_submits = 0
+        self._stop_enqueued = False
+        self._callback_state = local()
+        self._queue: Queue[ItemT | object] | None = None
+        self._slots: BoundedSemaphore | None = None
+        self._worker: Thread | None = None
+        if isinstance(mode, Bounded):
+            self._queue = Queue(maxsize=mode.capacity)
+            self._slots = BoundedSemaphore(mode.capacity)
+            self._worker = Thread(target=self._run, name=thread_name, daemon=True)
+            self._worker.start()
+
+    def submit(self, item: ItemT, *, size_bytes: int | None = None) -> None:
+        """Deliver one item according to the configured mode."""
+
+        if getattr(self._callback_state, "active", False):
+            raise RuntimeError("submit cannot be called from a sink callback")
+
+        with self._state:
+            self._raise_if_failed_locked()
+            if self._closed:
+                raise SinkClosedError("sink is closed")
+            self._active_submits += 1
+
+        acquired_slot = False
+        try:
+            if isinstance(self._mode, Inline):
+                self._invoke_handler(self._invoke_callback(self._retain, item))
+                return
+
+            if size_bytes is None and self._item_size is not None:
+                size_bytes = self._invoke_callback(self._item_size, item)
+            if size_bytes is None and self._retained_size is None:
+                raise TypeError(
+                    "bounded submit requires size_bytes, item_size, or retained_size"
+                )
+            if size_bytes is not None:
+                if isinstance(size_bytes, bool) or not isinstance(size_bytes, int):
+                    raise TypeError("item size must be an integer byte count")
+                if size_bytes < 0:
+                    raise ValueError("item size must not be negative")
+                if size_bytes > self._mode.max_item_bytes:
+                    raise ItemTooLargeError(size_bytes, self._mode.max_item_bytes)
+
+            slots = cast(BoundedSemaphore, self._slots)
+            queue = cast(Queue[ItemT | object], self._queue)
+            slots.acquire()
+            acquired_slot = True
+            with self._state:
+                self._raise_if_failed_locked()
+                if self._closed:
+                    raise SinkClosedError("sink is closed")
+            retained = self._invoke_callback(self._retain, item)
+            if self._retained_size is not None:
+                retained_bytes = self._invoke_callback(self._retained_size, retained)
+                if isinstance(retained_bytes, bool) or not isinstance(retained_bytes, int):
+                    raise TypeError("retained item size must be an integer byte count")
+                if retained_bytes < 0:
+                    raise ValueError("retained item size must not be negative")
+                if retained_bytes > self._mode.max_item_bytes:
+                    raise ItemTooLargeError(retained_bytes, self._mode.max_item_bytes)
+            queue.put(retained)
+            acquired_slot = False
+        finally:
+            if acquired_slot:
+                cast(BoundedSemaphore, self._slots).release()
+            with self._state:
+                self._active_submits -= 1
+                self._state.notify_all()
+
+    def close(self) -> None:
+        """Drain accepted work and stop the sink."""
+
+        if getattr(self._callback_state, "active", False):
+            raise RuntimeError("close cannot be called from a sink callback")
+
+        with self._state:
+            self._closed = True
+            while self._active_submits:
+                self._state.wait()
+            if self._worker is not None and not self._stop_enqueued:
+                cast(Queue[ItemT | object], self._queue).put(_STOP)
+                self._stop_enqueued = True
+
+        if self._worker is not None:
+            self._worker.join()
+        self._raise_if_failed()
+
+    def _run(self) -> None:
+        queue = cast(Queue[ItemT | object], self._queue)
+        slots = cast(BoundedSemaphore, self._slots)
+        while True:
+            item = queue.get()
+            if item is _STOP:
+                return
+            try:
+                self._invoke_handler(cast(ItemT, item))
+            except BaseException as error:
+                with self._state:
+                    if self._failure is None:
+                        self._failure = _WorkerFailure.capture(error)
+            finally:
+                del item
+                slots.release()
+
+    def _invoke_handler(self, item: ItemT) -> None:
+        self._invoke_callback(self._handler, item)
+
+    def _invoke_callback(self, callback, item):
+        previous = getattr(self._callback_state, "active", False)
+        self._callback_state.active = True
+        try:
+            return callback(item)
+        finally:
+            self._callback_state.active = previous
+
+    def _raise_if_failed_locked(self) -> None:
+        if self._failure is not None:
+            raise self._failure.to_exception()
+
+    def _raise_if_failed(self) -> None:
+        with self._state:
+            failure = self._failure
+        if failure is not None:
+            raise failure.to_exception()

blut_core/checkpoint.py

@@ -0,0 +1,172 @@
+"""checkpoint — corruption-safe torch checkpoint I/O + a resume payload contract
+(ADR 0044 P7, "durability"; the single most important reliability gap).
+
+- ``save(payload, path)`` — free-space preflight → write to a unique tmp in the
+  same dir → ``fsync`` → atomic ``os.replace`` → write a sidecar ``<path>.sha256``.
+  A mid-write kill leaves the PRIOR checkpoint intact, never a truncated file
+  (the failure mode that silently corrupts on a disk-full ``torch.save``).
+  Optional ``keep_last_n`` rotation of ``<stem>_ep*.pt`` siblings.
+- ``load(path, validate=True)`` — verify the sidecar SHA (or, absent a sidecar,
+  a dry ``torch.load`` of the bytes) BEFORE returning → a corrupt/truncated
+  checkpoint raises ``CheckpointError``, never silent garbage.
+- ``validate_payload(payload, required)`` — assert the resume contract keys are
+  present so a resume fails fast at SAVE, not 3 epochs into a relaunch (prior
+  bites: optimizer-state, QAT-alpha). Default contract = the training-resume set.
+
+torch is lazy-imported (only on save/load), so importing ``blut_core`` stays cheap.
+"""
+from __future__ import annotations
+
+import hashlib
+import os
+import shutil
+import tempfile
+from pathlib import Path
+from typing import Any, Dict, Iterable, Optional
+
+__all__ = ["save", "load", "validate_payload", "RESUME_CONTRACT", "CheckpointError"]
+
+# ADR 0044 P7 resume payload contract. A trainer that wants resumability MUST
+# put these in the checkpoint; validate_payload(strict=True) enforces it.
+RESUME_CONTRACT = (
+    "model", "optimizer", "lr_scheduler", "rng_state", "step",
+)
+
+
+class CheckpointError(RuntimeError):
+    """A checkpoint is missing, corrupt, or violates the resume contract."""
+
+
+def _sha256(path: Path) -> str:
+    h = hashlib.sha256()
+    with path.open("rb") as f:
+        for chunk in iter(lambda: f.read(1 << 20), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def validate_payload(payload: Dict[str, Any], required: Iterable[str] = RESUME_CONTRACT,
+                     *, strict: bool = True) -> list:
+    """Return the list of missing contract keys. ``strict`` → raise on any missing.
+    Use before save so a resume can't silently lack optimizer/RNG/step."""
+    if not isinstance(payload, dict):
+        raise CheckpointError(f"payload must be a dict, got {type(payload).__name__}")
+    missing = [k for k in required if k not in payload]
+    if missing and strict:
+        raise CheckpointError(
+            f"checkpoint payload missing resume-contract keys {missing}; "
+            f"present: {sorted(payload)}")
+    return missing
+
+
+def _free_bytes(directory: Path) -> int:
+    return shutil.disk_usage(directory).free
+
+
+def save(payload: Dict[str, Any], path, *, contract: Optional[Iterable[str]] = None,
+         keep_last_n: Optional[int] = None, min_free_mb: int = 256) -> Path:
+    """Atomically save ``payload`` to ``path`` with a sidecar SHA.
+
+    ``contract`` (default None = skip) → validate_payload(strict) first.
+    ``keep_last_n`` → after writing, keep only the newest N ``<stem>_ep*.pt``.
+    Raises ``CheckpointError`` on insufficient free space or a contract miss."""
+    import torch  # lazy
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    if contract is not None:
+        validate_payload(payload, contract, strict=True)
+
+    # Free-space preflight: a disk-full torch.save silently yields garbage.
+    if _free_bytes(path.parent) < min_free_mb * (1 << 20):
+        raise CheckpointError(
+            f"insufficient free space in {path.parent} (<{min_free_mb} MiB) — refusing to save")
+
+    fd, tmp_name = tempfile.mkstemp(dir=str(path.parent), prefix=path.stem + ".",
+                                    suffix=path.suffix + ".tmp")
+    tmp = Path(tmp_name)
+    try:
+        os.close(fd)  # torch.save writes by path
+        torch.save(payload, tmp)
+        with tmp.open("rb") as f:
+            os.fsync(f.fileno())
+        sha = _sha256(tmp)
+        os.replace(tmp, path)
+        # Sidecar: write to a tmp, fsync, then atomic rename — write_text() left
+        # the SHA in the page cache (no fsync), and a non-atomic overwrite could
+        # leave a .sha256 pointing at the PREVIOUS checkpoint after a crash.
+        sha_path = path.with_suffix(path.suffix + ".sha256")
+        sha_tmp = sha_path.with_suffix(sha_path.suffix + ".tmp")
+        payload_bytes = (sha + "\n").encode("utf-8")
+        try:
+            sfd = os.open(str(sha_tmp), os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o644)
+            try:
+                n = os.write(sfd, payload_bytes)
+                if n != len(payload_bytes):  # never short on a local FS at this size
+                    raise OSError(f"short sidecar write: {n}/{len(payload_bytes)}")
+                os.fsync(sfd)
+            finally:
+                os.close(sfd)
+            os.replace(sha_tmp, sha_path)
+        except Exception:
+            try:
+                sha_tmp.unlink()
+            except OSError:
+                pass
+            raise
+    except Exception:
+        try:
+            tmp.unlink()
+        except Exception:
+            pass
+        raise
+
+    if keep_last_n is not None and keep_last_n > 0:
+        _rotate(path, keep_last_n)
+    return path
+
+
+def _rotate(path: Path, keep_last_n: int) -> None:
+    # Rotates the EPOCH-SNAPSHOT siblings `{stem}_ep<N>{suffix}` (the trainer
+    # convention, e.g. snn_4state_best_ep10.pt), keeping the newest N. It does
+    # NOT touch the base `{stem}{suffix}` (the canonical/best checkpoint) — that
+    # is intentional: you keep best + the last N epoch snapshots.
+    sibs = sorted(path.parent.glob(f"{path.stem}_ep*{path.suffix}"),
+                  key=lambda p: p.stat().st_mtime, reverse=True)
+    for old in sibs[keep_last_n:]:
+        for victim in (old, old.with_suffix(old.suffix + ".sha256")):
+            try:
+                victim.unlink()
+            except OSError:
+                pass
+
+
+def load(path, *, validate: bool = True, map_location: Any = "cpu") -> Dict[str, Any]:
+    """Load a checkpoint, verifying integrity first. Raises ``CheckpointError``
+    on a missing/corrupt file rather than returning silent garbage."""
+    import torch  # lazy
+    path = Path(path)
+    if not path.exists():
+        raise CheckpointError(f"no checkpoint: {path}")
+    if validate:
+        sidecar = path.with_suffix(path.suffix + ".sha256")
+        if sidecar.exists():
+            want = sidecar.read_text().strip()
+            got = _sha256(path)
+            if got != want:
+                raise CheckpointError(
+                    f"checkpoint SHA mismatch for {path}: sidecar {want[:12]}… got {got[:12]}…")
+        else:
+            # No sidecar → integrity is UNVERIFIED (e.g. a crash between the
+            # checkpoint rename and the sidecar write; the sidecar is not part
+            # of the atomic unit). Surface it; don't silently accept-as-clean.
+            import sys as _sys
+            _sys.stderr.write(
+                f"[checkpoint] WARNING: no .sha256 sidecar for {path}; integrity unverified\n")
+    try:
+        # NOT weights_only=True: the resume payload carries optimizer/LR/RNG/
+        # scheduler objects (not plain tensors), which weights_only rejects.
+        # Trust model: checkpoints are own-produced under the job dir, not loaded
+        # from untrusted sources — the sidecar SHA guards corruption, not supply chain.
+        return torch.load(path, map_location=map_location)
+    except Exception as e:
+        raise CheckpointError(f"failed to load {path}: {type(e).__name__}: {e}") from e

blut_core/evaluator.py

@@ -0,0 +1,98 @@
+"""Generic ingredient-based evaluator for BLUT core cookbook.
+
+Reads an eval config JSON (written by the EvaluateModel Rust stage),
+builds the eval ingredient, and runs evaluation.
+
+Usage:
+    python -m blut_core.evaluator --config eval_config.json --output eval_report.json
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from blut_core import build_ingredient
+
+
+def load_config(config_path: str) -> dict:
+    with open(config_path) as f:
+        return json.load(f)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="BLUT core generic evaluator")
+    parser.add_argument("--config", required=True, help="Path to eval config JSON")
+    parser.add_argument("--output", required=True, help="Path to output eval report JSON")
+    args = parser.parse_args()
+
+    config = load_config(args.config)
+
+    checkpoint_path = config["checkpoint_path"]
+    dataset_path = config["dataset_path"]
+    eval_cfg = config["eval"]
+    batch_size = config.get("batch_size", 64)
+    device = config.get("device", "cuda")
+
+    # Build eval ingredient
+    print(f"[evaluator] building eval: {eval_cfg['kind']}:{eval_cfg['name']}")
+    evaluate_fn = build_ingredient(eval_cfg["kind"], eval_cfg["name"], eval_cfg.get("config"))
+
+    # Load checkpoint
+    import torch
+    print(f"[evaluator] loading checkpoint from {checkpoint_path}")
+    checkpoint_path = Path(checkpoint_path)
+
+    # Try to load as HuggingFace model
+    model = None
+    hf_dir = checkpoint_path / "hf"
+    if hf_dir.exists():
+        try:
+            from transformers import AutoModelForCausalLM
+            model = AutoModelForCausalLM.from_pretrained(str(hf_dir))
+            print(f"[evaluator] loaded HF model from {hf_dir}")
+        except Exception as e:
+            print(f"[evaluator] warning: could not load HF model: {e}")
+
+    if model is None:
+        # Try loading as torch checkpoint
+        pt_path = checkpoint_path / "model.pt"
+        if pt_path.exists():
+            print(f"[evaluator] found torch checkpoint at {pt_path}")
+            # Need a model architecture to load into — this is a limitation
+            # of the generic evaluator. Real usage should provide the model.
+            print("[evaluator] error: no model architecture available for torch checkpoint")
+            sys.exit(1)
+        else:
+            print(f"[evaluator] error: no checkpoint found at {checkpoint_path}")
+            sys.exit(1)
+
+    # Load dataset
+    print(f"[evaluator] loading dataset from {dataset_path}")
+    try:
+        from datasets import load_dataset
+        ds = load_dataset("json", data_files=str(dataset_path), split="train")
+        print(f"[evaluator] loaded {len(ds)} examples")
+    except Exception as e:
+        print(f"[evaluator] error: could not load dataset: {e}")
+        sys.exit(1)
+
+    # Run evaluation
+    print(f"[evaluator] running evaluation...")
+    if device != "cpu" and torch.cuda.is_available():
+        model = model.to(device)
+
+    metrics = evaluate_fn(model, ds, device=device)
+    print(f"[evaluator] metrics: {metrics}")
+
+    # Write output
+    with open(args.output, "w") as f:
+        json.dump(metrics, f, indent=2)
+    print(f"[evaluator] report written to {args.output}")
+
+
+if __name__ == "__main__":
+    main()

blut_core/ingredients/__init__.py

@@ -0,0 +1,20 @@
+"""Generic ingredient specs — the core training primitives.
+
+Importing this module registers all generic specs into the global registry.
+Domain cookbooks import ``blut_core`` first (registering these), then register
+their own domain-specific specs on top.
+"""
+
+# Import all _specs modules to trigger their @register_ingredient decorators.
+from blut_core.ingredients.data import _specs as _data_specs  # noqa: F401
+from blut_core.ingredients.model import _specs as _model_specs  # noqa: F401
+from blut_core.ingredients.optimizer import _specs as _optimizer_specs  # noqa: F401
+from blut_core.ingredients.scheduler import _specs as _scheduler_specs  # noqa: F401
+from blut_core.ingredients.loss import _specs as _loss_specs  # noqa: F401
+from blut_core.ingredients.step import _specs as _step_specs  # noqa: F401
+from blut_core.ingredients.ema import _specs as _ema_specs  # noqa: F401
+from blut_core.ingredients.checkpoint import _specs as _checkpoint_specs  # noqa: F401
+from blut_core.ingredients.eval import _specs as _eval_specs  # noqa: F401
+from blut_core.ingredients.sampler import _specs as _sampler_specs  # noqa: F401
+from blut_core.ingredients.logging import _specs as _logging_specs  # noqa: F401
+from blut_core.ingredients.forward import _specs as _forward_specs  # noqa: F401

blut_core/ingredients/_specs.py

@@ -0,0 +1,2 @@
+# Compatibility shim — all registrations happen in __init__.py already.
+# This module exists because blut_core.__init__ imports it.