singularic 1.0.0__py3-none-any.whl

singularic/__init__.py

@@ -0,0 +1,13 @@
+"""Singularic's intentionally minimal package marker.
+
+Stable APIs are imported from their owning modules.  Child directories are
+PEP 420 namespaces and deliberately have no package-barrel initializers.
+"""
+
+from importlib.metadata import version as _distribution_version
+from typing import Final
+
+from .core.apply import apply as apply
+from .core.state import State as State
+
+__version__: Final = _distribution_version("singularic")

singularic/checkpoint/bundle.py

@@ -0,0 +1,59 @@
+"""Generic portable checkpoint values independent of training orchestration."""
+
+import json
+from collections.abc import Mapping
+from dataclasses import dataclass
+from types import MappingProxyType
+from typing import Any
+
+from singularic.core.state import State
+
+
+def _json_value(value: Mapping[str, Any] | None, description: str) -> Mapping[str, Any]:
+    resolved = {} if value is None else dict(value)
+    try:
+        json.dumps(resolved, allow_nan=False, sort_keys=True)
+    except (TypeError, ValueError) as error:
+        raise TypeError(
+            f"CheckpointBundle {description} must be JSON-compatible"
+        ) from error
+    return MappingProxyType(resolved)
+
+
+@dataclass(frozen=True, slots=True)
+class CheckpointBundle:
+    """Named portable :class:`State` groups plus JSON metadata and cursor data.
+
+    The bundle has no knowledge of models, optimizers, or trainers.  Group names
+    are stable storage ownership boundaries, allowing a caller to restore just
+    the local state it owns.
+    """
+
+    groups: Mapping[str, State]
+    metadata: Mapping[str, Any] | None = None
+    dataset_cursor: Mapping[str, Any] | None = None
+
+    def __post_init__(self) -> None:
+        normalized = dict(self.groups)
+        if not normalized or any(
+            not isinstance(name, str) or not name for name in normalized
+        ):
+            raise ValueError("CheckpointBundle groups must have non-empty string names")
+        if any(not isinstance(state, State) for state in normalized.values()):
+            raise TypeError("CheckpointBundle groups must be State values")
+        object.__setattr__(self, "groups", MappingProxyType(normalized))
+        object.__setattr__(self, "metadata", _json_value(self.metadata, "metadata"))
+        object.__setattr__(
+            self, "dataset_cursor", _json_value(self.dataset_cursor, "dataset cursor")
+        )
+
+    def subset(self, names: tuple[str, ...]) -> CheckpointBundle:
+        if not names or len(names) != len(set(names)):
+            raise ValueError(
+                "CheckpointBundle subset names must be non-empty and unique"
+            )
+        try:
+            groups = {name: self.groups[name] for name in names}
+        except KeyError as error:
+            raise KeyError(f"Unknown checkpoint group {error.args[0]!r}") from error
+        return CheckpointBundle(groups, self.metadata, self.dataset_cursor)

singularic/checkpoint/manager.py

@@ -0,0 +1,350 @@
+"""Immutable generic checkpoint bundles over a minimal byte-store protocol."""
+
+import hashlib
+import json
+import re
+import threading
+import uuid
+from collections.abc import Callable, Mapping
+from concurrent.futures import Future, ThreadPoolExecutor
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Final, Self
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from safetensors.numpy import load, save
+
+from singularic.core.state import State, StateSchema
+
+from .bundle import CheckpointBundle
+from .store import ByteStore, PosixCheckpointStore, normalize_storage_key
+
+_FORMAT: Final = "singularic-checkpoint-bundle"
+
+
+def _json_bytes(value: Mapping[str, Any]) -> bytes:
+    return json.dumps(value, sort_keys=True, separators=(",", ":")).encode()
+
+
+def _json_object(payload: bytes, description: str) -> dict[str, Any]:
+    try:
+        value = json.loads(payload)
+    except (UnicodeDecodeError, ValueError) as error:
+        raise ValueError(f"Invalid {description} JSON") from error
+    if not isinstance(value, dict):
+        raise ValueError(f"{description} must be a JSON object")
+    return value
+
+
+@dataclass(frozen=True, slots=True, order=True)
+class CheckpointRef:
+    """One atomically published immutable bundle generation."""
+
+    step: int
+    generation: str
+    commit_key: str
+
+    def __post_init__(self) -> None:
+        if self.step < 0 or not re.fullmatch(r"[0-9a-f]{32}", self.generation):
+            raise ValueError("Invalid checkpoint reference")
+        normalize_storage_key(self.commit_key)
+
+
+@dataclass(frozen=True, slots=True)
+class CheckpointRestore:
+    bundle: CheckpointBundle
+    checkpoint: CheckpointRef
+
+
+class SaveHandle:
+    """Completion boundary for one asynchronous immutable publication."""
+
+    def __init__(
+        self, future: Future[CheckpointRef], consume: Callable[[], None]
+    ) -> None:
+        self._future = future
+        self._consume = consume
+
+    @property
+    def done(self) -> bool:
+        return self._future.done()
+
+    def result(self, timeout: float | None = None) -> CheckpointRef:
+        try:
+            return self._future.result(timeout=timeout)
+        finally:
+            if self._future.done():
+                self._consume()
+
+    wait = result
+
+
+class CheckpointManager:
+    """Persist and restore generic named state groups without Trainer knowledge."""
+
+    def __init__(
+        self,
+        directory: str | Path | None = None,
+        *,
+        store: ByteStore | None = None,
+        max_to_keep: int | None = None,
+    ) -> None:
+        if (directory is None) == (store is None):
+            raise ValueError("Provide exactly one of directory or store")
+        if max_to_keep is not None and max_to_keep <= 0:
+            raise ValueError("max_to_keep must be positive")
+        if store is None:
+            assert directory is not None
+            self.store: ByteStore = PosixCheckpointStore(Path(directory))
+        else:
+            self.store = store
+        if not isinstance(self.store, ByteStore):
+            raise TypeError("store must implement ByteStore")
+        self.max_to_keep = max_to_keep
+        self._lock = threading.Lock()
+        self._executor: ThreadPoolExecutor | None = None
+        self._pending: set[Future[CheckpointRef]] = set()
+        self._closed = False
+
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(self, *_: object) -> None:
+        self.close()
+
+    @staticmethod
+    def _commit_key(step: int) -> str:
+        return f"commits/step-{step:012d}.json"
+
+    def _put(self, key: str, payload: bytes) -> None:
+        if not self.store.put_if_absent(key, payload):
+            raise FileExistsError(f"Immutable checkpoint key already exists: {key}")
+
+    def save(self, bundle: CheckpointBundle, *, step: int) -> CheckpointRef:
+        """Write all named groups, then publish one conditional commit record."""
+
+        if not isinstance(bundle, CheckpointBundle):
+            raise TypeError("bundle must be a CheckpointBundle")
+        if step < 0:
+            raise ValueError("Checkpoint step must be nonnegative")
+        commit_key = self._commit_key(step)
+        if self.store.exists(commit_key):
+            raise FileExistsError(f"Checkpoint step {step} is already committed")
+        generation = uuid.uuid4().hex
+        prefix = f"generations/{generation}"
+        arrays: dict[str, np.ndarray[Any, Any]] = {}
+        groups: dict[str, Any] = {}
+        for group_name, state in bundle.groups.items():
+            tensors: list[dict[str, Any]] = []
+            for index, value in enumerate(state.values):
+                name = f"{group_name}:{index}"
+                typed_key = bool(jnp.issubdtype(value.dtype, jax.dtypes.prng_key))
+                physical = jax.random.key_data(value) if typed_key else value
+                arrays[name] = np.asarray(jax.device_get(physical))
+                tensors.append(
+                    {
+                        "name": name,
+                        "typed_key": typed_key,
+                        "key_impl": str(jax.random.key_impl(value))
+                        if typed_key
+                        else None,
+                    }
+                )
+            groups[group_name] = {"schema": state.schema.to_wire(), "tensors": tensors}
+        payload = save(arrays)
+        metadata = _json_bytes(
+            {
+                "format": _FORMAT,
+                "step": step,
+                "generation": generation,
+                "groups": groups,
+                "metadata": dict(bundle.metadata or {}),
+                "dataset_cursor": dict(bundle.dataset_cursor or {}),
+                "arrays_sha256": hashlib.sha256(payload).hexdigest(),
+            }
+        )
+        arrays_key = f"{prefix}/arrays.safetensors"
+        metadata_key = f"{prefix}/metadata.json"
+        commit = _json_bytes(
+            {
+                "format": _FORMAT,
+                "step": step,
+                "generation": generation,
+                "metadata_key": metadata_key,
+                "metadata_sha256": hashlib.sha256(metadata).hexdigest(),
+            }
+        )
+        self._put(arrays_key, payload)
+        self._put(metadata_key, metadata)
+        self._put(commit_key, commit)
+        reference = CheckpointRef(step, generation, commit_key)
+        self._prune()
+        return reference
+
+    def save_async(self, bundle: CheckpointBundle, *, step: int) -> SaveHandle:
+        with self._lock:
+            if self._closed:
+                raise RuntimeError("Checkpoint manager is closed")
+            if self._executor is None:
+                self._executor = ThreadPoolExecutor(max_workers=1)
+            future = self._executor.submit(self.save, bundle, step=step)
+            self._pending.add(future)
+
+        def consume() -> None:
+            with self._lock:
+                self._pending.discard(future)
+
+        return SaveHandle(future, consume)
+
+    def wait_for_pending(self) -> None:
+        with self._lock:
+            pending = tuple(self._pending)
+        for future in pending:
+            try:
+                future.result()
+            finally:
+                with self._lock:
+                    self._pending.discard(future)
+
+    def close(self) -> None:
+        with self._lock:
+            if self._closed:
+                return
+            self._closed = True
+            executor = self._executor
+        self.wait_for_pending()
+        if executor is not None:
+            executor.shutdown(wait=True)
+
+    def checkpoints(self) -> tuple[CheckpointRef, ...]:
+        records: list[CheckpointRef] = []
+        for key in self.store.list("commits/"):
+            record = _json_object(self.store.get(key), "checkpoint commit")
+            if record.get("format") != _FORMAT:
+                raise ValueError("Unsupported checkpoint commit")
+            records.append(
+                CheckpointRef(int(record["step"]), str(record["generation"]), key)
+            )
+        return tuple(sorted(records))
+
+    def latest(self) -> CheckpointRef:
+        records = self.checkpoints()
+        if not records:
+            raise FileNotFoundError("Checkpoint store contains no committed bundles")
+        return records[-1]
+
+    def restore(
+        self,
+        checkpoint: CheckpointRef | str | None = None,
+        *,
+        targets: Mapping[str, State] | None = None,
+    ) -> CheckpointRestore:
+        reference = self.latest() if checkpoint is None else self._reference(checkpoint)
+        commit = _json_object(self.store.get(reference.commit_key), "checkpoint commit")
+        if commit.get("format") != _FORMAT:
+            raise ValueError("Unsupported checkpoint commit")
+        if (
+            commit.get("step") != reference.step
+            or commit.get("generation") != reference.generation
+        ):
+            raise ValueError("Checkpoint reference does not match its commit")
+        metadata_key = str(commit["metadata_key"])
+        if metadata_key != f"generations/{reference.generation}/metadata.json":
+            raise ValueError("Checkpoint metadata key does not match its generation")
+        metadata_payload = self.store.get(normalize_storage_key(metadata_key))
+        if hashlib.sha256(metadata_payload).hexdigest() != commit.get(
+            "metadata_sha256"
+        ):
+            raise ValueError("Checkpoint metadata checksum mismatch")
+        metadata = _json_object(metadata_payload, "checkpoint metadata")
+        if metadata.get("format") != _FORMAT:
+            raise ValueError("Unsupported checkpoint metadata")
+        if (
+            metadata.get("step") != reference.step
+            or metadata.get("generation") != reference.generation
+        ):
+            raise ValueError("Checkpoint metadata does not match its commit")
+        arrays_key = f"generations/{reference.generation}/arrays.safetensors"
+        arrays_payload = self.store.get(arrays_key)
+        if hashlib.sha256(arrays_payload).hexdigest() != metadata.get("arrays_sha256"):
+            raise ValueError("Checkpoint array checksum mismatch")
+        decoded = load(arrays_payload)
+        restored: dict[str, State] = {}
+        raw_groups = metadata.get("groups")
+        if not isinstance(raw_groups, Mapping):
+            raise ValueError("Checkpoint groups must be an object")
+        expected_names: set[str] = set()
+        for name, raw_group in raw_groups.items():
+            if not isinstance(name, str) or not isinstance(raw_group, Mapping):
+                raise ValueError("Invalid checkpoint group")
+            schema = StateSchema.from_wire(raw_group["schema"])
+            tensors = raw_group["tensors"]
+            if not isinstance(tensors, list) or len(tensors) != len(schema):
+                raise ValueError("Checkpoint tensors do not match their schema")
+            values: list[jax.Array] = []
+            target = None if targets is None else targets.get(name)
+            if target is not None and target.schema != schema:
+                raise ValueError(f"Checkpoint group {name!r} does not match its target")
+            for index, item in enumerate(tensors):
+                if not isinstance(item, Mapping) or not isinstance(
+                    item.get("name"), str
+                ):
+                    raise ValueError("Invalid checkpoint tensor record")
+                tensor_name = item["name"]
+                if tensor_name != f"{name}:{index}" or tensor_name in expected_names:
+                    raise ValueError("Checkpoint tensor names are inconsistent")
+                expected_names.add(tensor_name)
+                try:
+                    value = jnp.asarray(decoded[tensor_name])
+                except KeyError as error:
+                    raise ValueError(
+                        "Checkpoint array payload is incomplete"
+                    ) from error
+                if bool(item.get("typed_key")):
+                    impl = item.get("key_impl")
+                    if not isinstance(impl, str):
+                        raise ValueError(
+                            "Typed key checkpoint is missing its implementation"
+                        )
+                    value = jax.random.wrap_key_data(value, impl=impl)
+                if target is not None:
+                    value = jax.device_put(value, target.values[index].sharding)
+                values.append(value)
+            restored[name] = State(tuple(values), schema)
+        if set(decoded) != expected_names:
+            raise ValueError("Checkpoint array payload does not match its metadata")
+        if targets is not None:
+            unknown = set(targets) - set(restored)
+            if unknown:
+                raise KeyError(
+                    f"Checkpoint is missing target groups: {sorted(unknown)}"
+                )
+        bundle_metadata = metadata.get("metadata", {})
+        dataset_cursor = metadata.get("dataset_cursor", {})
+        if not isinstance(bundle_metadata, Mapping) or not isinstance(
+            dataset_cursor, Mapping
+        ):
+            raise ValueError("Checkpoint bundle metadata must be objects")
+        return CheckpointRestore(
+            CheckpointBundle(restored, bundle_metadata, dataset_cursor), reference
+        )
+
+    def _reference(self, checkpoint: CheckpointRef | str) -> CheckpointRef:
+        if isinstance(checkpoint, CheckpointRef):
+            return checkpoint
+        key = normalize_storage_key(checkpoint)
+        record = _json_object(self.store.get(key), "checkpoint commit")
+        if record.get("format") != _FORMAT:
+            raise ValueError("Unsupported checkpoint commit")
+        return CheckpointRef(int(record["step"]), str(record["generation"]), key)
+
+    def _prune(self) -> None:
+        if self.max_to_keep is None:
+            return
+        expired = self.checkpoints()[: -self.max_to_keep]
+        for reference in expired:
+            for key in self.store.list(f"generations/{reference.generation}/"):
+                self.store.delete(key)
+            self.store.delete(reference.commit_key)

singularic/checkpoint/store.py

@@ -0,0 +1,107 @@
+"""Byte-oriented checkpoint storage and its durable POSIX implementation."""
+
+import os
+import uuid
+from dataclasses import dataclass
+from pathlib import Path, PurePosixPath
+from typing import Protocol, runtime_checkable
+
+
+@runtime_checkable
+class ByteStore(Protocol):
+    """Dependency-free immutable-key checkpoint storage boundary."""
+
+    def get(self, key: str) -> bytes: ...
+    def put_if_absent(self, key: str, value: bytes) -> bool: ...
+    def exists(self, key: str) -> bool: ...
+    def list(self, prefix: str = "") -> tuple[str, ...]: ...
+    def delete(self, key: str) -> None: ...
+
+
+def normalize_storage_key(key: str, *, allow_empty: bool = False) -> str:
+    """Validate and return a normalized relative POSIX storage key."""
+
+    if not isinstance(key, str) or "\\" in key or "\x00" in key:
+        raise ValueError("Storage keys must be normalized POSIX strings")
+    if allow_empty and key == "":
+        return ""
+    path = PurePosixPath(key)
+    normalized = path.as_posix()
+    if (
+        path.is_absolute()
+        or ".." in path.parts
+        or normalized != key
+        or (not allow_empty and normalized in {"", "."})
+    ):
+        raise ValueError("Storage keys must be normalized and relative")
+    return "" if normalized == "." else normalized
+
+
+def fsync_directory(path: Path) -> None:
+    descriptor = os.open(path, os.O_RDONLY)
+    try:
+        os.fsync(descriptor)
+    finally:
+        os.close(descriptor)
+
+
+@dataclass(frozen=True, slots=True)
+class PosixCheckpointStore:
+    """Durable shared-filesystem implementation of :class:`ByteStore`."""
+
+    root: Path
+
+    def __post_init__(self) -> None:
+        object.__setattr__(self, "root", Path(self.root))
+        self.root.mkdir(parents=True, exist_ok=True)
+
+    def path(self, key: str) -> Path:
+        candidate = self.root / normalize_storage_key(key)
+        try:
+            candidate.resolve(strict=False).relative_to(self.root.resolve())
+        except ValueError as error:
+            raise ValueError("Storage key escapes the configured root") from error
+        return candidate
+
+    def get(self, key: str) -> bytes:
+        return self.path(key).read_bytes()
+
+    def put_if_absent(self, key: str, value: bytes) -> bool:
+        path = self.path(key)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        temporary = path.parent / f".{path.name}.{uuid.uuid4().hex}.writing"
+        try:
+            with temporary.open("xb") as file:
+                file.write(value)
+                file.flush()
+                os.fsync(file.fileno())
+            try:
+                os.link(temporary, path)
+            except FileExistsError:
+                return False
+            fsync_directory(path.parent)
+            return True
+        finally:
+            temporary.unlink(missing_ok=True)
+
+    def list(self, prefix: str = "") -> tuple[str, ...]:
+        trailing_separator = prefix.endswith("/")
+        candidate = prefix[:-1] if trailing_separator else prefix
+        normalized = normalize_storage_key(candidate, allow_empty=True)
+        if trailing_separator and normalized:
+            normalized += "/"
+        return tuple(
+            sorted(
+                path.relative_to(self.root).as_posix()
+                for path in self.root.rglob("*")
+                if path.is_file()
+                and not path.name.endswith(".writing")
+                and path.relative_to(self.root).as_posix().startswith(normalized)
+            )
+        )
+
+    def delete(self, key: str) -> None:
+        self.path(key).unlink(missing_ok=True)
+
+    def exists(self, key: str) -> bool:
+        return self.path(key).is_file()

singularic/core/apply.py

@@ -0,0 +1,62 @@
+"""The explicit graph/state execution boundary for Singularic programs."""
+
+from dataclasses import dataclass
+from typing import Any
+
+from .context import execution_context
+from .module import GraphDef, _apply_graph
+from .precision import NumericsPlan, PrecisionPolicy, precision_scope
+from .rng import RngPlan, rng_scope
+from .state import State
+
+
+@dataclass(frozen=True, slots=True)
+class ApplyContext:
+    """Execution inputs that used to be hidden in ambient framework state.
+
+    ``rng_plan`` and ``rng_state`` must either be supplied together or omitted.
+    Callers can use ordinary JAX transformations over :func:`apply` because the
+    complete mutable program state remains an explicit argument/result.
+    """
+
+    precision: PrecisionPolicy | NumericsPlan | None = None
+    rng_plan: RngPlan | None = None
+    rng_state: State | None = None
+
+    def __post_init__(self) -> None:
+        if (self.rng_plan is None) != (self.rng_state is None):
+            raise ValueError("ApplyContext requires both rng_plan and rng_state")
+
+
+def apply(
+    graph: GraphDef,
+    state: State,
+    *args: Any,
+    context: ApplyContext | None = None,
+    method: str = "__call__",
+    **kwargs: Any,
+) -> tuple[Any, State]:
+    """Execute one graph method with all framework policy inputs explicit."""
+
+    resolved_context = ApplyContext() if context is None else context
+    precision = resolved_context.precision
+    numerics = (
+        None
+        if precision is None
+        else precision
+        if isinstance(precision, NumericsPlan)
+        else precision.compile(state.schema)
+    )
+    with execution_context():
+        if numerics is None:
+            if resolved_context.rng_plan is None:
+                return _apply_graph(graph, state, *args, method=method, **kwargs)
+            assert resolved_context.rng_state is not None
+            with rng_scope(resolved_context.rng_plan, resolved_context.rng_state):
+                return _apply_graph(graph, state, *args, method=method, **kwargs)
+        with precision_scope(numerics):
+            if resolved_context.rng_plan is None:
+                return _apply_graph(graph, state, *args, method=method, **kwargs)
+            assert resolved_context.rng_state is not None
+            with rng_scope(resolved_context.rng_plan, resolved_context.rng_state):
+                return _apply_graph(graph, state, *args, method=method, **kwargs)