weftcodec 1.2.0__py3-none-any.whl

weftcodec/dcp.py

@@ -0,0 +1,2108 @@
+"""PyTorch Distributed Checkpoint (DCP) storage adapter over WeftFS.
+
+``weftcodec.dcp`` implements the shape of PyTorch's
+:class:`torch.distributed.checkpoint.StorageWriter` and ``StorageReader``
+interfaces so that ``torch.distributed.checkpoint.save`` / ``load`` can write
+and read a checkpoint through a WeftFS store, with the training process never
+needing to mount anything.
+
+Everything durable goes through the shipped Rust CLI:
+
+* tensor bytes become one canonical ``weft.checkpoint.v1`` checkpoint written by
+  ``weftfs checkpoint write`` (one mode-10 shard object per tensor slice, or an
+  exact ``zstd``/``raw`` shard when the caller asks for it);
+* the DCP protocol metadata (state-dict schema, storage keys, planner data) is
+  serialised as canonical WCB and stored as an exact store file object at
+  ``/checkpoints/<storage-name>/dcp.metadata``;
+* reads use ``weftfs checkpoint restore`` and re-verify every recorded BLAKE3
+  content root before handing bytes to DCP. The default restore is read-only in
+  the store; ``import_restored=True`` opts into the CLI's ``--import``, which
+  publishes a second exact copy under ``/dcp-reads/<uuid>/``.
+
+There is no approximate path and no silent fallback. A missing CLI, a missing
+store, a non-canonical manifest, a content-root mismatch, or an unsupported
+dtype raises a typed error.
+
+``torch`` is optional. Importing :mod:`weftcodec.dcp` without PyTorch works and
+exposes the same classes/methods (contract mode); operations that need tensors
+raise :class:`TorchRequired`. When PyTorch is importable the classes subclass the
+real DCP interfaces.
+
+Not claimed: GPU-direct I/O, a multi-rank cluster run, or stock DCP binaries.
+The adapters do not need a read-only store: ``checkpoint restore`` only reads
+unless ``--import`` is requested.
+"""
+
+from __future__ import annotations
+
+import ctypes
+import importlib
+import io
+import json
+import math
+import os
+import shutil
+import struct
+import subprocess
+import tempfile
+import threading
+import time
+import uuid
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
+from pathlib import Path
+
+from . import wcb
+from .hashing import content_root
+from .service import (
+    CheckpointService,
+    ServiceError,
+    ServiceProtocolError,
+    ServiceUnavailable,
+    canonical_element_bytes,
+)
+
+__all__ = [
+    "AsyncSaveHandle",
+    "DCP_METADATA_NAME",
+    "DCP_METADATA_SCHEMA",
+    "DcpError",
+    "ExactnessError",
+    "InvalidCheckpointId",
+    "InvalidMetadata",
+    "MetadataMissing",
+    "RestoredEntry",
+    "SaveFailed",
+    "StagingLimitExceeded",
+    "StoreCommandError",
+    "StoreUnavailable",
+    "TORCH_AVAILABLE",
+    "TORCH_VERSION",
+    "TorchRequired",
+    "UnsupportedDtype",
+    "UnsupportedStateObject",
+    "WeftStorageReader",
+    "WeftStorageWriter",
+    "init_store",
+    "save_checkpoint_async",
+    "verify_checkpoint",
+]
+
+
+DCP_METADATA_SCHEMA = "weft.dcp.metadata.v1"
+DCP_METADATA_NAME = "dcp.metadata"
+DEFAULT_STAGING_BYTES = 256 * 1024 * 1024
+DEFAULT_TIMEOUT_SECONDS = 1800.0
+DEFAULT_CLI = "/work/.docker-cache/target-rust/release/weftfs"
+
+#: torch dtype spelling -> (canonical WEFT-TENSOR dtype, safetensors spelling,
+#: element bytes). Only these dtypes have an exact Rust safetensors mapping.
+_SAFETENSORS_BY_TORCH = {
+    "torch.float16": ("f16", "F16", 2),
+    "torch.bfloat16": ("bf16", "BF16", 2),
+    "torch.float32": ("f32", "F32", 4),
+    "torch.float64": ("f64", "F64", 8),
+    "torch.int8": ("i8", "I8", 1),
+    "torch.uint8": ("u8", "U8", 1),
+    "torch.int16": ("i16", "I16", 2),
+    "torch.int32": ("i32", "I32", 4),
+    "torch.int64": ("i64", "I64", 8),
+}
+_CANONICAL_TO_TORCH = {
+    canonical: torch_name
+    for torch_name, (canonical, _sf, _size) in _SAFETENSORS_BY_TORCH.items()
+}
+
+
+class DcpError(RuntimeError):
+    """Base class for every adapter error."""
+
+
+class StoreUnavailable(DcpError):
+    """The WeftFS CLI or the store it names is not usable."""
+
+
+class StoreCommandError(DcpError):
+    """A WeftFS CLI invocation failed."""
+
+    def __init__(self, message, *, command, returncode, stderr, error_class=None):
+        super().__init__(message)
+        self.command = tuple(command)
+        self.returncode = returncode
+        self.stderr = stderr
+        self.error_class = error_class
+
+
+class InvalidCheckpointId(DcpError):
+    """The requested checkpoint id is not a legal single store name."""
+
+
+class UnsupportedDtype(DcpError):
+    """A tensor dtype has no exact safetensors/WEFT-TENSOR v1 mapping."""
+
+
+class UnsupportedStateObject(DcpError):
+    """The state dict contains an object this adapter does not persist."""
+
+
+class StagingLimitExceeded(DcpError):
+    """A single resolved tensor payment does not fit the staging bound."""
+
+
+class MetadataMissing(DcpError):
+    """No DCP metadata object and manifest synthesis is disabled."""
+
+
+class InvalidMetadata(DcpError):
+    """The stored DCP metadata is malformed or refers to unknown entries."""
+
+
+class ExactnessError(DcpError):
+    """A stored content root disagrees with the bytes/manifests."""
+
+
+class TorchRequired(DcpError):
+    """This operation needs PyTorch but it is not installed."""
+
+
+class SaveFailed(DcpError):
+    """A save did not produce a usable checkpoint."""
+
+
+# ---------------------------------------------------------------------------
+# Optional PyTorch. Importing ``weftcodec`` never imports this module.
+
+
+def _load_torch():
+    try:
+        return importlib.import_module("torch")
+    except ImportError:
+        return None
+
+
+_TORCH = _load_torch()
+TORCH_AVAILABLE = _TORCH is not None
+TORCH_VERSION = getattr(_TORCH, "__version__", None) if _TORCH is not None else None
+
+if TORCH_AVAILABLE:  # pragma: no branch - both paths are exercised by the court
+    try:
+        from torch.distributed.checkpoint.metadata import (
+            BytesStorageMetadata as _BytesStorageMetadata,
+        )
+        from torch.distributed.checkpoint.metadata import (
+            ChunkStorageMetadata as _ChunkStorageMetadata,
+        )
+        from torch.distributed.checkpoint.metadata import (
+            Metadata as _TorchMetadata,
+        )
+        from torch.distributed.checkpoint.metadata import (
+            MetadataIndex as _MetadataIndex,
+        )
+        from torch.distributed.checkpoint.metadata import (
+            StorageMeta as _TorchStorageMeta,
+        )
+        from torch.distributed.checkpoint.metadata import (
+            TensorProperties as _TensorProperties,
+        )
+        from torch.distributed.checkpoint.metadata import (
+            TensorStorageMetadata as _TensorStorageMetadata,
+        )
+        from torch.distributed.checkpoint.storage import (
+            StorageReader as _StorageReaderBase,
+        )
+        from torch.distributed.checkpoint.storage import (
+            StorageWriter as _StorageWriterBase,
+        )
+        from torch.futures import Future as _TorchFuture
+
+        _TORCH_IMPORT_ERROR = None
+    except ImportError as exc:  # pragma: no cover - defensive for odd builds
+        TORCH_AVAILABLE = False
+        TORCH_VERSION = None
+        _TORCH = None
+        _TORCH_IMPORT_ERROR = str(exc)
+else:  # pragma: no cover - the court always has torch
+    _TORCH_IMPORT_ERROR = "torch is not installed"
+
+if not TORCH_AVAILABLE:
+
+    class _StorageWriterBase:  # type: ignore[no-redef]
+        """Contract-mode stand-in when PyTorch is absent."""
+
+    class _StorageReaderBase:  # type: ignore[no-redef]
+        """Contract-mode stand-in when PyTorch is absent."""
+
+    class _TorchFuture:  # type: ignore[no-redef]
+        """Minimal future with the subset of ``torch.futures.Future`` used here."""
+
+        def __init__(self):
+            self._ready = threading.Event()
+            self._value = None
+            self._error = None
+
+        def set_result(self, value):
+            self._value = value
+            self._ready.set()
+
+        def set_exception(self, error):
+            self._error = error
+            self._ready.set()
+
+        def wait(self):
+            self._ready.wait()
+            return self
+
+        def value(self):
+            self._ready.wait()
+            if self._error is not None:
+                raise self._error
+            return self._value
+
+        def done(self):
+            return self._ready.is_set()
+
+        def add_done_callback(self, callback):
+            self._ready.wait()
+            callback(self)
+
+    @dataclass(frozen=True)
+    class _FallbackStorageMeta:  # type: ignore[no-redef]
+        checkpoint_id: str = None
+        save_id: str = None
+        load_id: str = None
+        modules: tuple = ()
+
+    _TorchStorageMeta = _FallbackStorageMeta
+    _MetadataIndex = None
+    _TorchMetadata = None
+    _TensorStorageMetadata = None
+    _BytesStorageMetadata = None
+    _ChunkStorageMetadata = None
+    _TensorProperties = None
+
+
+def _require_torch(operation):
+    if not TORCH_AVAILABLE:
+        raise TorchRequired(
+            "%s requires PyTorch (install with `pip install weftcodec[torch]`); "
+            "weftcodec.dcp imported in contract mode: %s"
+            % (operation, _TORCH_IMPORT_ERROR)
+        )
+
+
+# ---------------------------------------------------------------------------
+# Small canonical helpers
+
+
+def _u(value):
+    return wcb.U64(int(value))
+
+
+def _i(value):
+    if isinstance(value, wcb.U64):
+        return int(value.value)
+    if isinstance(value, wcb.I64):
+        return int(value.value)
+    return int(value)
+
+
+def _u_list(values):
+    return [_u(value) for value in values]
+
+
+def _i_list(values):
+    return tuple(_i(value) for value in values)
+
+
+def _dtype_name(dtype) -> str:
+    """Canonical torch dtype spelling, e.g. ``torch.float32``."""
+    text = str(dtype)
+    if not text.startswith("torch."):
+        raise UnsupportedDtype("unrecognised torch dtype %r" % (dtype,))
+    return text
+
+
+def _canonical_dtype(torch_dtype_name: str) -> str:
+    try:
+        return _SAFETENSORS_BY_TORCH[torch_dtype_name][0]
+    except KeyError:
+        raise UnsupportedDtype(
+            "dtype %s has no exact WEFT-TENSOR/safetensors v1 mapping; "
+            "supported: %s"
+            % (torch_dtype_name, ", ".join(sorted(_SAFETENSORS_BY_TORCH)))
+        )
+
+
+def _safetensors_dtype(torch_dtype_name: str) -> str:
+    try:
+        return _SAFETENSORS_BY_TORCH[torch_dtype_name][1]
+    except KeyError:
+        raise UnsupportedDtype(
+            "dtype %s has no exact safetensors mapping; supported: %s"
+            % (torch_dtype_name, ", ".join(sorted(_SAFETENSORS_BY_TORCH)))
+        )
+
+
+def _element_bytes(torch_dtype_name: str) -> int:
+    try:
+        return _SAFETENSORS_BY_TORCH[torch_dtype_name][2]
+    except KeyError:
+        raise UnsupportedDtype("dtype %s has no element size" % torch_dtype_name)
+
+
+def _torch_dtype(torch_dtype_name: str):
+    _require_torch("resolving a tensor dtype")
+    try:
+        return getattr(_TORCH, torch_dtype_name.split(".", 1)[1])
+    except AttributeError:
+        raise UnsupportedDtype("this PyTorch has no dtype %s" % torch_dtype_name)
+
+
+def _validate_checkpoint_id(checkpoint_id) -> str:
+    if checkpoint_id is None:
+        raise InvalidCheckpointId("checkpoint id is required")
+    text = os.fspath(checkpoint_id)
+    if not isinstance(text, str) or not text:
+        raise InvalidCheckpointId("checkpoint id must be a non-empty path name")
+    if "/" in text or "\\" in text or text in (".", "..") or "\x00" in text:
+        raise InvalidCheckpointId(
+            "checkpoint id %r must be a single store name; the store directory is "
+            "selected with store_root" % text
+        )
+    return text
+
+
+def _generation(name: str, base: str):
+    if name == base:
+        return 1
+    prefix = base + ".v"
+    if name.startswith(prefix):
+        suffix = name[len(prefix) :]
+        if suffix.isdigit() and str(int(suffix)) == suffix and int(suffix) >= 2:
+            return int(suffix)
+    return None
+
+
+def _parse_topology(topology, axis=None):
+    if isinstance(topology, (tuple, list)):
+        if len(topology) != 2:
+            raise ValueError("topology must be an int or a (shards, axis) pair")
+        shards, explicit_axis = topology
+        if axis is not None and int(axis) != int(explicit_axis):
+            raise ValueError("topology axis and the axis argument disagree")
+        axis = explicit_axis
+    else:
+        shards = topology
+    shards = int(shards)
+    axis = 0 if axis is None else int(axis)
+    if shards < 1:
+        raise ValueError("topology shard count must be >= 1")
+    if axis < 0 or axis > 0xFFFFFFFF:
+        raise ValueError("topology axis must be a u32")
+    return shards, axis
+
+
+def _default_cli(cli=None):
+    if cli:
+        return os.fspath(cli)
+    env = os.environ.get("WEFTFS_CLI")
+    if env:
+        return env
+    if os.path.isfile(DEFAULT_CLI) and os.access(DEFAULT_CLI, os.X_OK):
+        return DEFAULT_CLI
+    found = shutil.which("weftfs")
+    if found:
+        return found
+    raise StoreUnavailable(
+        "no weftfs CLI found: set WEFTFS_CLI or pass weftfs_cli=... (looked at "
+        "%s and PATH)" % DEFAULT_CLI
+    )
+
+
+def _parse_error_class(stderr: str):
+    first = stderr.strip().splitlines()[0] if stderr.strip() else ""
+    if ":" in first:
+        head = first.split(":", 1)[0].strip()
+        if head and head.replace("_", "").isalnum() and head[0].isupper():
+            return head
+    return None
+
+
+class _Storage:
+    """Thin, fail-closed wrapper around the ``weftfs`` CLI."""
+
+    def __init__(self, store_root, *, cli=None, timeout=DEFAULT_TIMEOUT_SECONDS):
+        self.root = Path(store_root)
+        self.cli = _default_cli(cli)
+        self.timeout = float(timeout)
+        self._lock = threading.Lock()
+        self.command_ms = 0.0
+        self.calls = 0
+
+    # -- process plumbing ---------------------------------------------------
+
+    def _run(self, args, *, binary=False):
+        command = [self.cli, "--store", os.fspath(self.root), *[str(a) for a in args]]
+        started = time.monotonic()
+        try:
+            proc = subprocess.run(
+                command,
+                capture_output=True,
+                timeout=self.timeout,
+                check=False,
+            )
+        except FileNotFoundError as exc:
+            raise StoreUnavailable("weftfs CLI not executable: %s" % exc) from exc
+        except subprocess.TimeoutExpired as exc:
+            raise StoreCommandError(
+                "weftfs command timed out after %.0fs: %s"
+                % (self.timeout, " ".join(command)),
+                command=command,
+                returncode=-1,
+                stderr="",
+                error_class="Timeout",
+            ) from exc
+        elapsed = (time.monotonic() - started) * 1000.0
+        with self._lock:
+            self.command_ms += elapsed
+            self.calls += 1
+        if proc.returncode != 0:
+            stderr = proc.stderr.decode("utf-8", "replace").strip()
+            error_class = _parse_error_class(stderr)
+            raise StoreCommandError(
+                "weftfs %s failed (%d): %s"
+                % (" ".join(str(a) for a in args), proc.returncode, stderr),
+                command=command,
+                returncode=proc.returncode,
+                stderr=stderr,
+                error_class=error_class,
+            )
+        return proc
+
+    def _json(self, args):
+        proc = self._run([*args, "--json"])
+        text = proc.stdout.decode("utf-8")
+        try:
+            return json.loads(text)
+        except ValueError as exc:
+            raise StoreCommandError(
+                "weftfs %s produced invalid JSON: %s"
+                % (" ".join(str(a) for a in args), text[:400]),
+                command=[self.cli, "--store", os.fspath(self.root), *args],
+                returncode=0,
+                stderr=text[:400],
+                error_class="InvalidJson",
+            ) from exc
+
+    # -- store verbs --------------------------------------------------------
+
+    def requires_store(self):
+        if not (self.root / "FORMAT").is_file():
+            raise StoreUnavailable(
+                "no WeftFS store at %s (missing FORMAT); call init_store() first"
+                % self.root
+            )
+
+    def init(self):
+        if (self.root / "FORMAT").is_file():
+            return
+        self._run(["init"])
+
+    def list_checkpoints(self):
+        payload = self._json(["checkpoint", "list"])
+        return [row["name"] for row in payload.get("checkpoints", [])]
+
+    def checkpoint_show(self, name):
+        payload = self._json(["checkpoint", "show", "--name", name])
+        manifest = payload.get("manifest")
+        if not isinstance(manifest, dict):
+            raise InvalidMetadata("checkpoint show --json has no manifest object")
+        return payload
+
+    def checkpoint_write(self, staging, name, *, shards, axis, strategy):
+        args = [
+            "checkpoint",
+            "write",
+            "--file",
+            os.fspath(staging),
+            "--name",
+            name,
+            "--shards",
+            str(shards),
+            "--axis",
+            str(axis),
+        ]
+        if strategy:
+            args += ["--shard-strategy", strategy]
+        return self._json(args)
+
+    def restore(self, name, *, entry=None, topology=None, axis=0, host_out, import_restored=False):
+        """Restore to ``host_out``; ``import_restored`` opts into a store copy.
+
+        ``checkpoint restore`` is read-only in the store by default. Passing
+        ``import_restored=True`` adds ``--import`` and publishes the exact
+        restored container under ``/dcp-reads/<uuid>/``, which is a real,
+        durable second copy and therefore opt-in.
+        """
+        args = [
+            "checkpoint",
+            "restore",
+            "--name",
+            name,
+            "--host-out",
+            os.fspath(host_out),
+        ]
+        receipt = os.environ.get("WEFT_LOAD_RECEIPT")
+        if receipt:
+            args += [
+                "--load-receipt",
+                receipt,
+                "--verify-load",
+                os.environ.get("WEFT_VERIFY_LOAD", "lazy"),
+                "--replica",
+                os.environ.get("WEFT_REPLICA", "local"),
+            ]
+        if import_restored:
+            out_dir = "/dcp-reads/%s" % uuid.uuid4().hex
+            args += ["--import", "--out", out_dir]
+        if entry is not None:
+            args += ["--tensor", entry]
+        if topology is not None:
+            args += ["--topology", str(topology), "--axis", str(axis)]
+        return self._json(args)
+
+    def put_file(self, store_path, payload: bytes):
+        with tempfile.NamedTemporaryFile(
+            prefix="weft-dcp-put-", suffix=".bin", delete=False
+        ) as handle:
+            handle.write(payload)
+            host_path = handle.name
+        try:
+            return self._json(["create", store_path, "--file", host_path])
+        finally:
+            try:
+                os.unlink(host_path)
+            except OSError:
+                pass
+
+    def get_file(self, store_path):
+        """Exact bytes of a store file object."""
+        proc = self._run(["cat", store_path], binary=True)
+        return proc.stdout
+
+    def file_exists(self, store_path):
+        try:
+            self._run(["stat", store_path])
+            return True
+        except StoreCommandError as exc:
+            if exc.error_class in ("NotFound",):
+                return False
+            raise
+
+    def verify_full(self):
+        return self._json(["verify", "--full"])
+
+    def checkpoint_verify(self, name):
+        return self._json(["checkpoint", "verify", "--name", name])
+
+
+# ---------------------------------------------------------------------------
+# DCP metadata schema (canonical WCB)
+
+
+@dataclass(frozen=True)
+class StorageInfo:
+    """Where one DCP storage index lives inside the checkpoint container."""
+
+    entry: str
+    dtype: str  # torch dtype spelling
+    shape: tuple
+    storage_offsets: tuple
+    lengths: tuple
+    content_root: bytes
+
+    def to_value(self):
+        return {
+            "entry": self.entry,
+            "dtype": self.dtype,
+            "shape": _u_list(self.shape),
+            "storage_offsets": _u_list(self.storage_offsets),
+            "lengths": _u_list(self.lengths),
+            "content_root": self.content_root,
+        }
+
+    @classmethod
+    def from_value(cls, value):
+        try:
+            return cls(
+                entry=value["entry"],
+                dtype=value["dtype"],
+                shape=_i_list(value["shape"]),
+                storage_offsets=_i_list(value.get("storage_offsets", [])),
+                lengths=_i_list(value.get("lengths", [])),
+                content_root=bytes(value["content_root"]),
+            )
+        except (KeyError, TypeError) as exc:
+            raise InvalidMetadata("bad storage entry: %s" % exc) from exc
+
+
+@dataclass(frozen=True)
+class RestoredEntry:
+    """One entry restored (and re-verified) through the Rust CLI."""
+
+    entry: str
+    data: bytes
+    dtype: str
+    shape: tuple
+    content_root: bytes
+    topology_root: "str | None"
+    shards_read: int
+    store_bytes_read: int
+
+    @property
+    def content_root_hex(self) -> str:
+        return self.content_root.hex()
+
+
+def _state_dict_value(state_dict_metadata):
+    rows = []
+    for fqn in sorted(state_dict_metadata):
+        item = state_dict_metadata[fqn]
+        if _TensorStorageMetadata is not None and isinstance(
+            item, _TensorStorageMetadata
+        ):
+            rows.append(
+                {
+                    "fqn": fqn,
+                    "kind": "tensor",
+                    "dtype": _dtype_name(item.properties.dtype),
+                    "size": _u_list(item.size),
+                    "chunks": [
+                        {
+                            "offsets": _u_list(chunk.offsets),
+                            "sizes": _u_list(chunk.sizes),
+                        }
+                        for chunk in item.chunks
+                    ],
+                }
+            )
+        elif _BytesStorageMetadata is not None and isinstance(
+            item, _BytesStorageMetadata
+        ):
+            rows.append({"fqn": fqn, "kind": "bytes"})
+        else:
+            raise UnsupportedStateObject(
+                "state dict entry %r has storage metadata %s; this adapter "
+                "persists plain tensors and byte objects only"
+                % (fqn, type(item).__name__)
+            )
+    return rows
+
+
+def _planner_data_value(planner_data):
+    if planner_data is None:
+        return {"kind": "none"}
+    if isinstance(planner_data, dict) and all(
+        isinstance(key, str)
+        and isinstance(value, (tuple, list))
+        and all(isinstance(part, str) for part in value)
+        for key, value in planner_data.items()
+    ):
+        return {
+            "kind": "dict_str_tuple",
+            "entries": [
+                {"key": key, "value": list(planner_data[key])}
+                for key in sorted(planner_data)
+            ],
+        }
+    raise UnsupportedStateObject(
+        "planner_data of type %s is not persistable by this adapter; only None "
+        "and dict[str, tuple[str, ...]] are supported" % type(planner_data).__name__
+    )
+
+
+def serialize_metadata(metadata, *, checkpoint_id, storage_name) -> bytes:
+    """Canonical WCB bytes for a DCP ``Metadata`` object."""
+    storage_rows = []
+    storage_data = metadata.storage_data
+    if storage_data is None:
+        raise InvalidMetadata("metadata.storage_data is None; write_data/finish "
+                              "must run before metadata is serialised")
+    for key in sorted(storage_data, key=lambda k: (k.fqn, tuple(k.offset or ()))):
+        info = storage_data[key]
+        if not isinstance(info, StorageInfo):
+            raise InvalidMetadata(
+                "storage_data[%r] is %s, expected StorageInfo"
+                % (key, type(info).__name__)
+            )
+        storage_rows.append(
+            {
+                "fqn": key.fqn,
+                "offset": None if key.offset is None else _u_list(key.offset),
+                "index": None if key.index is None else _u(key.index),
+                "info": info.to_value(),
+            }
+        )
+    version = getattr(metadata, "version", None)
+    value = {
+        "schema": DCP_METADATA_SCHEMA,
+        "checkpoint_id": checkpoint_id,
+        "storage_name": storage_name,
+        "dcp_version": None if version is None else str(version),
+        "state_dict": _state_dict_value(metadata.state_dict_metadata),
+        "storage_data": storage_rows,
+        "planner_data": _planner_data_value(getattr(metadata, "planner_data", None)),
+    }
+    return wcb.encode(value)
+
+
+def _metadata_from_value(value, *, storage_name):
+    _require_torch("reconstructing DCP metadata")
+    if not isinstance(value, dict) or value.get("schema") != DCP_METADATA_SCHEMA:
+        raise InvalidMetadata(
+            "DCP metadata schema is %r, expected %r"
+            % (
+                value.get("schema") if isinstance(value, dict) else None,
+                DCP_METADATA_SCHEMA,
+            )
+        )
+    if value.get("storage_name") != storage_name:
+        raise InvalidMetadata(
+            "DCP metadata names storage %r but the resolved checkpoint is %r"
+            % (value.get("storage_name"), storage_name)
+        )
+    state_dict_metadata = {}
+    for row in value.get("state_dict", []):
+        fqn = row.get("fqn")
+        kind = row.get("kind")
+        if kind == "tensor":
+            dtype_name = row.get("dtype")
+            if dtype_name not in _SAFETENSORS_BY_TORCH:
+                raise UnsupportedDtype(
+                    "stored tensor %r has unsupported dtype %r" % (fqn, dtype_name)
+                )
+            state_dict_metadata[fqn] = _TensorStorageMetadata(
+                properties=_TensorProperties(dtype=_torch_dtype(dtype_name)),
+                size=_TORCH.Size(_i_list(row.get("size", []))),
+                chunks=[
+                    _ChunkStorageMetadata(
+                        offsets=_TORCH.Size(_i_list(chunk.get("offsets", []))),
+                        sizes=_TORCH.Size(_i_list(chunk.get("sizes", []))),
+                    )
+                    for chunk in row.get("chunks", [])
+                ],
+            )
+        elif kind == "bytes":
+            state_dict_metadata[fqn] = _BytesStorageMetadata()
+        else:
+            raise InvalidMetadata("unknown state-dict entry kind %r" % (kind,))
+
+    storage_data = {}
+    for row in value.get("storage_data", []):
+        info = StorageInfo.from_value(row.get("info", {}))
+        offset = row.get("offset")
+        index = row.get("index")
+        key = _MetadataIndex(
+            row.get("fqn"),
+            offset=None if offset is None else _i_list(offset),
+            index=None if index is None else _i(index),
+        )
+        storage_data[key] = info
+
+    planner = value.get("planner_data") or {"kind": "none"}
+    if planner.get("kind") == "none":
+        planner_data = None
+    elif planner.get("kind") == "dict_str_tuple":
+        planner_data = {
+            entry["key"]: tuple(entry["value"]) for entry in planner.get("entries", [])
+        }
+    else:
+        raise InvalidMetadata("unknown planner_data kind %r" % (planner.get("kind"),))
+
+    return _TorchMetadata(
+        state_dict_metadata=state_dict_metadata,
+        planner_data=planner_data,
+        storage_data=storage_data,
+        storage_meta=_TorchStorageMeta(
+            checkpoint_id=value.get("checkpoint_id"),
+            save_id=value.get("storage_name"),
+        ),
+        version=value.get("dcp_version"),
+    )
+
+
+def _metadata_from_manifest(show, *, checkpoint_id, storage_name):
+    """Build DCP metadata for a checkpoint written by the Rust CLI alone.
+
+    Every manifest tensor becomes a state-dict key of the same name with one
+    full-shape chunk. This is exact for plain safetensors checkpoints; it is
+    reported as ``metadata_source == "manifest"`` so callers can tell it apart
+    from adapter-written checkpoints that carry DCP metadata.
+    """
+    _require_torch("synthesising DCP metadata from a manifest")
+    manifest = show["manifest"]
+    state_dict_metadata = {}
+    storage_data = {}
+    for tensor in manifest.get("tensors", []):
+        name = tensor["name"]
+        canonical = tensor["dtype"]
+        torch_name = _CANONICAL_TO_TORCH.get(canonical)
+        if torch_name is None:
+            raise UnsupportedDtype(
+                "manifest tensor %r has dtype %r with no exact torch dtype"
+                % (name, canonical)
+            )
+        shape = tuple(int(x) for x in tensor["shape"])
+        storage_offsets = tuple([0] * len(shape))
+        element_bytes = int(tensor["element_bytes"])
+        expected = 1
+        for extent in shape:
+            expected *= extent
+        expected *= _element_bytes(torch_name)
+        if expected != element_bytes:
+            raise InvalidMetadata(
+                "manifest tensor %r element_bytes %d != shape product %d"
+                % (name, element_bytes, expected)
+            )
+        key = _MetadataIndex(name, offset=storage_offsets)
+        state_dict_metadata[name] = _TensorStorageMetadata(
+            properties=_TensorProperties(dtype=_torch_dtype(torch_name)),
+            size=_TORCH.Size(shape),
+            chunks=[
+                _ChunkStorageMetadata(
+                    offsets=_TORCH.Size([0] * len(shape)),
+                    sizes=_TORCH.Size(shape),
+                )
+            ],
+        )
+        storage_data[key] = StorageInfo(
+            entry=name,
+            dtype=torch_name,
+            shape=shape,
+            storage_offsets=storage_offsets,
+            lengths=shape,
+            content_root=bytes.fromhex(tensor["content_root"]),
+        )
+    return _TorchMetadata(
+        state_dict_metadata=state_dict_metadata,
+        planner_data=None,
+        storage_data=storage_data,
+        storage_meta=_TorchStorageMeta(
+            checkpoint_id=checkpoint_id, save_id=storage_name
+        ),
+        version="1.0.0",
+    )
+
+
+# ---------------------------------------------------------------------------
+# safetensors helpers (writing the staging container, reading restored bytes)
+
+
+def _tensor_bytes(tensor):
+    """Exact little-endian row-major bytes of a CPU tensor, without numpy.
+
+    ``data_ptr()`` addresses the tensor's first logical element, so a single
+    ``ctypes.string_at`` copies exactly ``numel * element_size`` bytes for a
+    contiguous tensor, including dtypes numpy cannot represent (bfloat16,
+    float8). This is a host-memory copy, not a zero-copy path.
+    """
+    _require_torch("serialising a tensor")
+    t = tensor.detach()
+    if t.device.type != "cpu":
+        t = t.cpu()
+    if t.numel() == 0:
+        return b""
+    t = t.contiguous()
+    expected = t.numel() * t.element_size()
+    payload = ctypes.string_at(t.data_ptr(), expected)
+    if len(payload) != expected:
+        raise ExactnessError(
+            "tensor byte extraction produced %d bytes, expected %d"
+            % (len(payload), expected)
+        )
+    return payload
+
+
+def _narrow(tensor, offsets, lengths):
+    if not offsets:
+        return tensor
+    slices = tuple(
+        slice(int(offset), int(offset) + int(length))
+        for offset, length in zip(offsets, lengths)
+    )
+    return tensor[slices]
+
+
+def _parse_safetensors(path):
+    with open(path, "rb") as handle:
+        prefix = handle.read(8)
+        if len(prefix) != 8:
+            raise InvalidMetadata("restored container is shorter than 8 bytes")
+        header_len = struct.unpack("<Q", prefix)[0]
+        if header_len > 64 * 1024 * 1024:
+            raise InvalidMetadata(
+                "restored container header length %d exceeds the 64 MiB bound"
+                % header_len
+            )
+        header_raw = handle.read(header_len)
+        if len(header_raw) != header_len:
+            raise InvalidMetadata("restored container header is truncated")
+        try:
+            header = json.loads(header_raw)
+        except ValueError as exc:
+            raise InvalidMetadata("restored container header is not JSON: %s" % exc)
+        data_start = 8 + header_len
+        entries = {}
+        for name, spec in header.items():
+            if name == "__metadata__":
+                continue
+            try:
+                begin, end = spec["data_offsets"]
+                dtype = spec["dtype"]
+                shape = tuple(int(x) for x in spec["shape"])
+            except (KeyError, TypeError, ValueError) as exc:
+                raise InvalidMetadata(
+                    "restored container tensor %r descriptor is invalid: %s"
+                    % (name, exc)
+                )
+            if begin < 0 or end < begin:
+                raise InvalidMetadata(
+                    "restored container tensor %r has invalid data_offsets" % name
+                )
+            handle.seek(data_start + begin)
+            payload = handle.read(end - begin)
+            if len(payload) != end - begin:
+                raise InvalidMetadata(
+                    "restored container tensor %r is truncated" % name
+                )
+            entries[name] = (dtype, shape, payload)
+    return entries
+
+
+# ---------------------------------------------------------------------------
+# Writer
+
+
+class WeftStorageWriter(_StorageWriterBase):
+    """DCP ``StorageWriter`` that publishes a canonical WeftFS checkpoint.
+
+    Parameters
+    ----------
+    path:
+        DCP checkpoint id. It becomes a single WeftFS checkpoint name under
+        ``/checkpoints``; it must not contain ``/``.
+    store_root:
+        Directory of an initialized WeftFS store (or pass ``init_store=True``).
+    topology:
+        ``n`` or ``(n, axis)``: store every tensor split n ways along ``axis``.
+        ``n == 1`` stores whole tensors. The Rust CLI rejects a non-divisible
+        extent ("no silent padding"); the adapter never pads or reinterprets.
+        This applies to non-tensor (``BYTE_IO``) state objects too: they are
+        stored as ``u8[len]`` tensors, so a split topology needs their
+        serialized length to be divisible by the shard count. Use ``topology=1``
+        for state dicts whose byte objects have arbitrary lengths.
+    shard_strategy:
+        ``zstd`` (default; exact, works for any shard length), ``raw`` (exact),
+        or ``tensor-sharded`` (mode-10 byte-word transport). ``tensor-sharded``
+        additionally requires every shard's byte length to be divisible by 16
+        (the smallest mode-10 block); the Rust CLI rejects anything else rather
+        than pad.
+    staging_bytes:
+        Hard bound on one resolved tensor copy staged for the safetensors
+        container. Tensors above the bound raise :class:`StagingLimitExceeded`.
+    overwrite:
+        When True (default) a second save to the same id is written as a new
+        generation ``<id>.v2``, ``<id>.v3``, ... (no store-side delete exists).
+        When False an existing generation raises ``FileExistsError``.
+    """
+
+    def __init__(
+        self,
+        path,
+        *,
+        store_root,
+        topology=1,
+        axis=None,
+        shard_strategy="zstd",
+        staging_bytes=DEFAULT_STAGING_BYTES,
+        overwrite=True,
+        init_store=False,
+        keep_staging=False,
+        weftfs_cli=None,
+        timeout=DEFAULT_TIMEOUT_SECONDS,
+    ):
+        self._path = _validate_checkpoint_id(path)
+        self.store_root = Path(store_root)
+        self._topology = _parse_topology(topology, axis)
+        if shard_strategy not in ("tensor-sharded", "tensor-tiled", "zstd", "raw"):
+            raise ValueError(
+                "shard_strategy must be tensor-sharded, tensor-tiled, zstd, or raw"
+            )
+        if shard_strategy == "tensor-tiled":
+            raise ValueError(
+                "tensor-tiled checkpoint shards need a tile shape, which this "
+                "adapter does not expose; use tensor-sharded, zstd, or raw"
+            )
+        staging_bytes = int(staging_bytes)
+        if staging_bytes <= 0:
+            raise ValueError("staging_bytes must be a positive byte count")
+        self.staging_bytes = staging_bytes
+        self.shard_strategy = shard_strategy
+        self.overwrite = bool(overwrite)
+        self.keep_staging = bool(keep_staging)
+        self._store = _Storage(self.store_root, cli=weftfs_cli, timeout=timeout)
+        if init_store:
+            self._store.init()
+        self._store.requires_store()
+        self._service_cli = _default_cli(weftfs_cli)
+        self._timeout = float(timeout)
+        self._service = None
+        self._lock = threading.RLock()
+        self._save_id = uuid.uuid4().hex
+        self._storage_name = None
+        self._entries = {}
+        self._results = []
+        self._pending_metadata = None
+        self._finished = False
+        self._manifest = None
+        self._stats = {
+            "entries": 0,
+            "tensor_bytes": 0,
+            "staging_peak_bytes": 0,
+            "checkpoint_write_ms": 0.0,
+        }
+
+    # -- DCP interface ------------------------------------------------------
+
+    @classmethod
+    def validate_checkpoint_id(cls, checkpoint_id) -> bool:
+        try:
+            _validate_checkpoint_id(checkpoint_id)
+            return True
+        except (InvalidCheckpointId, TypeError):
+            return False
+
+    @property
+    def checkpoint_id(self):
+        return self._path
+
+    def _get_service(self):
+        """One helper process for every operation of this writer."""
+        if self._service is None:
+            self._service = CheckpointService(
+                self.store_root, cli=self._service_cli, timeout=self._timeout
+            )
+        return self._service
+
+    def close(self):
+        if self._service is not None:
+            self._service.close()
+            self._service = None
+
+    @property
+    def store_name(self):
+        return self._storage_name
+
+    @property
+    def stats(self):
+        payload = dict(self._stats)
+        payload.update(
+            {
+                "checkpoint_id": self._path,
+                "storage_name": self._storage_name,
+                "store_root": os.fspath(self.store_root),
+                "store_command_ms": self._store.command_ms,
+                "store_command_calls": self._store.calls,
+                "service": (
+                    None if self._service is None else dict(self._service.stats)
+                ),
+                "topology": self._topology,
+                "shard_strategy": self.shard_strategy,
+                "staging_bytes": self.staging_bytes,
+            }
+        )
+        return payload
+
+    def reset(self, checkpoint_id=None):
+        with self._lock:
+            if checkpoint_id is not None:
+                self._path = _validate_checkpoint_id(checkpoint_id)
+            self._save_id = uuid.uuid4().hex
+            self._storage_name = None
+            self._entries = {}
+            self._results = []
+            self._pending_metadata = None
+            self._finished = False
+            self._manifest = None
+            self._stats = {
+                "entries": 0,
+                "tensor_bytes": 0,
+                "staging_peak_bytes": 0,
+                "checkpoint_write_ms": 0.0,
+            }
+        return None
+
+    def set_up_storage_writer(self, is_coordinator, *args, **kwargs):
+        self._is_coordinator = bool(is_coordinator)
+        return None
+
+    def prepare_local_plan(self, plan):
+        return plan
+
+    def prepare_global_plan(self, plans):
+        return plans
+
+    def storage_meta(self):
+        return _TorchStorageMeta(checkpoint_id=self._path, save_id=self._save_id)
+
+    def write_metadata(self, metadata):
+        """Record the DCP metadata object.
+
+        PyTorch 2.15's ``save`` does not call this method (it writes metadata in
+        ``finish``); the method is implemented for interface parity and older
+        callers. The metadata is persisted once, in ``finish``, because the
+        storage keys are only final after ``write_data``.
+        """
+        if metadata is None:
+            raise InvalidMetadata("write_metadata(None)")
+        with self._lock:
+            self._pending_metadata = metadata
+        return None
+
+    def write_data(self, plan, planner):
+        future = _TorchFuture()
+        try:
+            results = self._execute(plan, planner)
+        except BaseException as exc:  # propagated through the DCP future
+            future.set_exception(exc)
+            return future
+        future.set_result(results)
+        return future
+
+    def finish(self, metadata, results):
+        with self._lock:
+            if self._storage_name is None:
+                raise SaveFailed(
+                    "finish() called before a successful write_data(); no "
+                    "checkpoint was published"
+                )
+            if self._finished:
+                raise SaveFailed("finish() called twice for one save")
+            merged = {}
+            for group in results:
+                for result in group:
+                    merged[result.index] = result.storage_data
+            metadata.storage_data = merged
+            metadata.storage_meta = self.storage_meta()
+            if getattr(metadata, "version", None) is None:
+                metadata.version = "1.0.0"
+            blob = serialize_metadata(
+                metadata,
+                checkpoint_id=self._path,
+                storage_name=self._storage_name,
+            )
+            try:
+                self._get_service().write_file(
+                    "/checkpoints/%s/%s" % (self._storage_name, DCP_METADATA_NAME),
+                    blob,
+                )
+            except (StoreCommandError, ServiceError) as exc:
+                raise SaveFailed(
+                    "checkpoint %r was published but its DCP metadata was not: %s"
+                    % (self._storage_name, exc)
+                ) from exc
+            self._finished = True
+            self._metadata = metadata
+        return None
+
+    # -- internals ----------------------------------------------------------
+
+    def _allocate_storage_name(self):
+        existing = set(self._get_service().list_checkpoints())
+        if self._path not in existing:
+            return self._path
+        if not self.overwrite:
+            raise FileExistsError(
+                "checkpoint %r already exists in %s and overwrite=False"
+                % (self._path, self.store_root)
+            )
+        highest = 1
+        for name in existing:
+            generation = _generation(name, self._path)
+            if generation is not None and generation > highest:
+                highest = generation
+        return "%s.v%d" % (self._path, highest + 1)
+
+    def _entry_names(self, items):
+        used = set()
+        names = {}
+        for item in items:
+            fqn = item.index.fqn
+            name = fqn
+            suffix = 1
+            while name in used:
+                name = "%s.__weftwp_%d" % (fqn, suffix)
+                suffix += 1
+            used.add(name)
+            names[(fqn, tuple(item.index.offset or ()), item.index.index)] = name
+        return names
+
+    def _execute(self, plan, planner):
+        _require_torch("writing a DCP checkpoint")
+        items = list(plan.items)
+        if not items:
+            raise SaveFailed(
+                "save plan contains no items; refusing to publish an empty "
+                "checkpoint"
+            )
+        items.sort(
+            key=lambda item: (
+                item.index.fqn,
+                tuple(item.index.offset or ()),
+                -1 if item.index.index is None else int(item.index.index),
+            )
+        )
+        names = self._entry_names(items)
+
+        # Pass 1: derive lengths, dtypes and shapes; cache byte payloads.
+        staged = []
+        byte_payloads = {}
+        byte_bytes = 0
+        for item in items:
+            key = (item.index.fqn, tuple(item.index.offset or ()), item.index.index)
+            entry = names[key]
+            if getattr(item.type, "name", str(item.type)) == "BYTE_IO":
+                data = planner.resolve_data(item)
+                if isinstance(data, io.BytesIO):
+                    payload = data.getvalue()
+                elif isinstance(data, (bytes, bytearray, memoryview)):
+                    payload = bytes(data)
+                else:
+                    raise UnsupportedStateObject(
+                        "BYTE_IO write item %r resolved to %s, expected bytes"
+                        % (item.index.fqn, type(data).__name__)
+                    )
+                byte_bytes += len(payload)
+                if byte_bytes > self.staging_bytes:
+                    raise StagingLimitExceeded(
+                        "byte state objects need %d staged bytes, above the %d "
+                        "staging_bytes bound" % (byte_bytes, self.staging_bytes)
+                    )
+                byte_payloads[key] = payload
+                staged.append(
+                    {
+                        "item": item,
+                        "entry": entry,
+                        "kind": "bytes",
+                        "dtype": "torch.uint8",
+                        "shape": (len(payload),),
+                        "payload": payload,
+                        "offsets": (),
+                        "lengths": (),
+                    }
+                )
+                continue
+
+            tensor_data = item.tensor_data
+            if tensor_data is None:
+                raise UnsupportedStateObject(
+                    "write item %r has no tensor_data" % (item.index.fqn,)
+                )
+            chunk = tensor_data.chunk
+            shape = tuple(int(x) for x in tensor_data.size)
+            offsets = tuple(int(x) for x in chunk.offsets)
+            lengths = tuple(int(x) for x in chunk.sizes)
+            if offsets != tuple([0] * len(shape)) or lengths != shape:
+                raise UnsupportedStateObject(
+                    "write item %r is a chunk at %r of shape %r; sharded/"
+                    "chunked state objects are not persisted by this adapter"
+                    % (item.index.fqn, offsets, shape)
+                )
+            dtype = _dtype_name(tensor_data.properties.dtype)
+            _canonical_dtype(dtype)  # validates the exact mapping
+            staged.append(
+                {
+                    "item": item,
+                    "entry": entry,
+                    "kind": "tensor",
+                    "dtype": dtype,
+                    "shape": shape,
+                    "payload": None,
+                    "offsets": offsets,
+                    "lengths": lengths,
+                }
+            )
+
+        # Header is fixed before data is streamed.
+        header = {}
+        cursor = 0
+        for row in staged:
+            nbytes = (
+                len(row["payload"])
+                if row["kind"] == "bytes"
+                else _element_bytes(row["dtype"]) * math.prod(row["shape"])
+            )
+            header[row["entry"]] = {
+                "dtype": _safetensors_dtype(row["dtype"]),
+                "shape": list(row["shape"]),
+                "data_offsets": [cursor, cursor + nbytes],
+            }
+            row["nbytes"] = nbytes
+            cursor += nbytes
+        header_raw = json.dumps(
+            header, sort_keys=True, separators=(",", ":")
+        ).encode("utf-8")
+        if len(header_raw) > 64 * 1024 * 1024:
+            raise SaveFailed("staging safetensors header exceeds 64 MiB")
+
+        # Stage the canonical container in memory. It must exist as
+        # safetensors bytes because the checkpoint manifest records the source
+        # header and SHA-256, but it never reaches the disk, is never a path
+        # argument to a CLI process, and is handed to the helper over a memfd.
+        container = bytearray(struct.pack("<Q", len(header_raw)))
+        container += header_raw
+        infos = []
+        peak = 0
+        for row in staged:
+            if row["kind"] == "bytes":
+                payload = row["payload"]
+            else:
+                tensor = planner.resolve_data(row["item"])
+                if not hasattr(tensor, "detach"):
+                    raise UnsupportedStateObject(
+                        "tensor write item %r resolved to %s"
+                        % (row["item"].index.fqn, type(tensor).__name__)
+                    )
+                actual_dtype = _dtype_name(tensor.dtype)
+                actual_shape = tuple(int(x) for x in tensor.shape)
+                if actual_dtype != row["dtype"] or actual_shape != row["shape"]:
+                    raise SaveFailed(
+                        "planner resolved %r as %s%r, plan said %s%r"
+                        % (
+                            row["item"].index.fqn,
+                            actual_dtype,
+                            actual_shape,
+                            row["dtype"],
+                            row["shape"],
+                        )
+                    )
+                if row["nbytes"] > self.staging_bytes:
+                    raise StagingLimitExceeded(
+                        "tensor %r needs %d staged bytes, above the %d "
+                        "staging_bytes bound"
+                        % (
+                            row["item"].index.fqn,
+                            row["nbytes"],
+                            self.staging_bytes,
+                        )
+                    )
+                payload = _tensor_bytes(tensor)
+                peak = max(peak, len(payload))
+            if len(payload) != row["nbytes"]:
+                raise SaveFailed(
+                    "entry %r produced %d bytes, header declared %d"
+                    % (row["entry"], len(payload), row["nbytes"])
+                )
+            container += payload
+            infos.append(
+                StorageInfo(
+                    entry=row["entry"],
+                    dtype=row["dtype"],
+                    shape=row["shape"],
+                    storage_offsets=row["offsets"],
+                    lengths=row["lengths"],
+                    content_root=content_root(payload),
+                )
+            )
+
+        shards, axis = self._topology
+        storage_name = self._allocate_storage_name()
+        service = self._get_service()
+        started = time.monotonic()
+        write_report = service.write_checkpoint(
+            storage_name,
+            container,
+            shards=shards,
+            axis=axis,
+            strategy=self.shard_strategy,
+        )
+        self._stats["checkpoint_write_ms"] = (time.monotonic() - started) * 1000.0
+        self._stats["container_bytes"] = len(container)
+        manifest_value = write_report.manifest.to_value()
+        self._verify_report(
+            {
+                "complete": write_report.complete,
+                "tensors": manifest_value["tensors"],
+            },
+            manifest_value["tensors"],
+            infos,
+            storage_name,
+        )
+        self._storage_name = storage_name
+        self._manifest = {
+            "manifest_root": write_report.manifest_root,
+            "manifest": manifest_value,
+        }
+        self._stats["entries"] = len(infos)
+        self._stats["tensor_bytes"] = sum(
+            int(t["element_bytes"]) for t in manifest_value["tensors"]
+        )
+        self._stats["staging_peak_bytes"] = peak
+
+        results = []
+        for row, info in zip(staged, infos):
+            results.append(_write_result(row["item"], info, row["nbytes"]))
+        with self._lock:
+            self._results = results
+            self._entries = {info.entry: info for info in infos}
+        return results
+
+    def _verify_report(self, report, manifest_tensors, infos, storage_name):
+        tensors = manifest_tensors
+        if not isinstance(tensors, list) or not tensors:
+            raise SaveFailed(
+                "checkpoint manifest for %r has no tensor list" % storage_name
+            )
+        by_name = {row["name"]: row for row in tensors}
+        if set(by_name) != {info.entry for info in infos}:
+            raise SaveFailed(
+                "checkpoint %r stored tensors %r, expected %r"
+                % (
+                    storage_name,
+                    sorted(by_name),
+                    sorted(info.entry for info in infos),
+                )
+            )
+        for info in infos:
+            row = by_name[info.entry]
+            if row.get("dtype") != _canonical_dtype(info.dtype):
+                raise ExactnessError(
+                    "entry %r dtype %r != canonical %r"
+                    % (info.entry, row.get("dtype"), _canonical_dtype(info.dtype))
+                )
+            if tuple(int(x) for x in row.get("shape", [])) != info.shape:
+                raise ExactnessError(
+                    "entry %r shape %r != %r"
+                    % (info.entry, row.get("shape"), info.shape)
+                )
+            stored = row.get("content_root")
+            if stored != info.content_root.hex():
+                raise ExactnessError(
+                    "entry %r content root %s != the bytes written (%s)"
+                    % (info.entry, stored, info.content_root.hex())
+                )
+        if report.get("complete") is not True:
+            raise SaveFailed(
+                "checkpoint %r is not complete" % storage_name
+            )
+
+
+def _write_result(item, info, nbytes):
+    _require_torch("building a DCP WriteResult")
+    from torch.distributed.checkpoint.storage import WriteResult as _WriteResult
+
+    return _WriteResult(
+        index=item.index, size_in_bytes=int(nbytes), storage_data=info
+    )
+
+
+# ---------------------------------------------------------------------------
+# Reader
+
+
+class _RestoredHandle:
+    """One memfd-backed restored tensor plus its recorded metadata."""
+
+    __slots__ = ("entry", "buffer", "info")
+
+    def __init__(self, entry, buffer, info):
+        self.entry = entry
+        self.buffer = buffer
+        self.info = info
+
+
+class WeftStorageReader(_StorageReaderBase):
+    """DCP ``StorageReader`` over a WeftFS checkpoint.
+
+    The checkpoint manifest (``weft.checkpoint.v1``) is the authority for tensor
+    identity. If ``/checkpoints/<name>/dcp.metadata`` exists it supplies the DCP
+    protocol metadata and is verified against the manifest; otherwise metadata is
+    synthesised from the manifest itself (``metadata_source == "manifest"``,
+    exact for plain-integer/packed safetensors checkpoints written by the Rust
+    CLI).
+    """
+
+    def __init__(
+        self,
+        path,
+        *,
+        store_root,
+        weftfs_cli=None,
+        timeout=DEFAULT_TIMEOUT_SECONDS,
+        allow_manifest_synthesis=True,
+        import_restored=False,
+    ):
+        self._path = _validate_checkpoint_id(path)
+        self.store_root = Path(store_root)
+        self.allow_manifest_synthesis = bool(allow_manifest_synthesis)
+        # Restore is read-only in the store by default. Importing a second,
+        # exact copy into ``/dcp-reads/...`` is opt-in because it is a durable
+        # store mutation.
+        self.import_restored = bool(import_restored)
+        if self.import_restored:
+            raise DcpError(
+                "import_restored is no longer supported: the checkpoint service "
+                "restores directly into caller buffers and never materializes a "
+                "container to import"
+            )
+        self._store = _Storage(self.store_root, cli=weftfs_cli, timeout=timeout)
+        self._store.requires_store()
+        self._service_cli = _default_cli(weftfs_cli)
+        self._timeout = float(timeout)
+        self._service = None
+        self._storage_name = None
+        self._manifest = None
+        self._metadata = None
+        self._metadata_source = None
+        self._storage_data = None
+        self._stats = {"restores": 0, "entries_restored": 0, "restored_bytes": 0}
+        self._lock = threading.RLock()
+
+    # -- DCP interface ------------------------------------------------------
+
+    @classmethod
+    def validate_checkpoint_id(cls, checkpoint_id) -> bool:
+        try:
+            _validate_checkpoint_id(checkpoint_id)
+            return True
+        except (InvalidCheckpointId, TypeError):
+            return False
+
+    @property
+    def checkpoint_id(self):
+        return self._path
+
+    @property
+    def storage_name(self):
+        return self._storage_name
+
+    def _get_service(self):
+        """One helper process for every read of this reader."""
+        if self._service is None:
+            self._service = CheckpointService(
+                self.store_root, cli=self._service_cli, timeout=self._timeout
+            )
+        return self._service
+
+    def close(self):
+        if self._service is not None:
+            self._service.close()
+            self._service = None
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *exc):
+        self.close()
+        return False
+
+    @property
+    def metadata_source(self):
+        return self._metadata_source
+
+    @property
+    def stats(self):
+        payload = dict(self._stats)
+        payload.update(
+            {
+                "checkpoint_id": self._path,
+                "storage_name": self._storage_name,
+                "metadata_source": self._metadata_source,
+                "store_root": os.fspath(self.store_root),
+                "store_command_ms": self._store.command_ms,
+                "store_command_calls": self._store.calls,
+                "service": (
+                    None if self._service is None else dict(self._service.stats)
+                ),
+            }
+        )
+        return payload
+
+    def reset(self, checkpoint_id=None):
+        with self._lock:
+            if checkpoint_id is not None:
+                self._path = _validate_checkpoint_id(checkpoint_id)
+            self._storage_name = None
+            self._manifest = None
+            self._metadata = None
+            self._metadata_source = None
+            self._storage_data = None
+            self._stats = {"restores": 0, "entries_restored": 0, "restored_bytes": 0}
+        return None
+
+    def _resolve_storage_name(self):
+        names = set(self._get_service().list_checkpoints())
+        highest = 1
+        chosen = None
+        for name in names:
+            generation = _generation(name, self._path)
+            if generation is not None and generation >= highest:
+                highest = generation
+                chosen = name
+        if chosen is None:
+            raise MetadataMissing(
+                "checkpoint %r does not exist in %s" % (self._path, self.store_root)
+            )
+        return chosen
+
+    def read_metadata(self, *args, **kwargs):
+        _require_torch("reading DCP metadata")
+        with self._lock:
+            if self._metadata is not None:
+                return self._metadata
+            storage_name = self._resolve_storage_name()
+            service = self._get_service()
+            # A reader is a long-lived object; refresh so commits made by
+            # another process are visible before the manifest is read.
+            service.refresh()
+            descriptor = service.manifest(storage_name)
+            manifest = descriptor.to_value()
+            show = {
+                "manifest_root": descriptor.manifest_root,
+                "manifest": manifest,
+            }
+            manifest_rows = {
+                row["name"]: row for row in manifest.get("tensors", [])
+            }
+            metadata_path = "/checkpoints/%s/%s" % (storage_name, DCP_METADATA_NAME)
+            blob = None
+            exists, _size = service.stat_file(metadata_path)
+            if exists:
+                with service.read_file(metadata_path) as view:
+                    blob = view.bytes
+            if blob is not None:
+                try:
+                    value = wcb.decode(blob)
+                except Exception as exc:
+                    raise InvalidMetadata(
+                        "stored DCP metadata at %s is not canonical WCB: %s"
+                        % (metadata_path, exc)
+                    ) from exc
+                metadata = _metadata_from_value(value, storage_name=storage_name)
+                source = "dcp"
+                self._validate_against_manifest(metadata, manifest_rows)
+            else:
+                if not self.allow_manifest_synthesis:
+                    raise MetadataMissing(
+                        "checkpoint %r has no %s and manifest synthesis is "
+                        "disabled" % (storage_name, DCP_METADATA_NAME)
+                    )
+                metadata = _metadata_from_manifest(
+                    show,
+                    checkpoint_id=self._path,
+                    storage_name=storage_name,
+                )
+                source = "manifest"
+            metadata.storage_meta = _TorchStorageMeta(
+                checkpoint_id=self._path,
+                save_id=storage_name,
+                load_id=uuid.uuid4().hex,
+            )
+            self._storage_name = storage_name
+            self._manifest = show
+            self._metadata = metadata
+            self._metadata_source = source
+            self._storage_data = metadata.storage_data
+            return metadata
+
+    def set_up_storage_reader(self, metadata, is_coordinator, *args, **kwargs):
+        if metadata.storage_data is None:
+            raise InvalidMetadata("metadata.storage_data is None")
+        with self._lock:
+            self._metadata = metadata
+            self._storage_data = metadata.storage_data
+            if self._storage_name is None:
+                # Metadata supplied by the caller rather than read_metadata().
+                self._storage_name = getattr(
+                    getattr(metadata, "storage_meta", None), "save_id", None
+                )
+        return None
+
+    def prepare_local_plan(self, plan):
+        return plan
+
+    def prepare_global_plan(self, plans):
+        return plans
+
+    def read_data(self, plan, planner):
+        future = _TorchFuture()
+        try:
+            self._read_plan(plan, planner)
+        except BaseException as exc:
+            future.set_exception(exc)
+            return future
+        future.set_result(None)
+        return future
+
+    # -- helpers ------------------------------------------------------------
+
+    def _validate_against_manifest(self, metadata, manifest_rows):
+        for info in metadata.storage_data.values():
+            row = manifest_rows.get(info.entry)
+            if row is None:
+                raise InvalidMetadata(
+                    "DCP metadata names entry %r, absent from the checkpoint "
+                    "manifest" % info.entry
+                )
+            if row.get("content_root") != info.content_root.hex():
+                raise ExactnessError(
+                    "DCP metadata content root for %r (%s) != manifest (%s)"
+                    % (info.entry, info.content_root.hex(), row.get("content_root"))
+                )
+            if tuple(int(x) for x in row.get("shape", [])) != info.shape:
+                raise ExactnessError(
+                    "DCP metadata shape for %r (%r) != manifest (%r)"
+                    % (info.entry, info.shape, row.get("shape"))
+                )
+            expected_canonical = _canonical_dtype(info.dtype)
+            if row.get("dtype") != expected_canonical:
+                raise ExactnessError(
+                    "DCP metadata dtype for %r (%s) != manifest (%s)"
+                    % (info.entry, expected_canonical, row.get("dtype"))
+                )
+
+    def _restore(self, entries, *, topology=None, axis=0):
+        """Restore entries straight into caller-owned memfds.
+
+        Returns ``{entry: _RestoredHandle}``. The handle owns a memfd-backed
+        mapping; callers must ``close()`` it (or pass it to
+        :meth:`_release_handles`). No file is written and the store is never
+        mutated.
+        """
+        if self._storage_name is None:
+            raise DcpError("read_data() before read_metadata()")
+        service = self._get_service()
+        manifest_rows = {
+            row["name"]: row for row in self._manifest["manifest"].get("tensors", [])
+        }
+        unknown = [entry for entry in entries if entry not in manifest_rows]
+        if unknown:
+            raise InvalidMetadata(
+                "requested entries %r are not in checkpoint %r"
+                % (sorted(unknown), self._storage_name)
+            )
+        topology_arg = None if topology is None else (int(topology), int(axis))
+        found = {}
+        try:
+            for entry in sorted(entries):
+                info = self._entry_info(entry)
+                try:
+                    buffer = service.restore_tensor(
+                        self._storage_name,
+                        entry,
+                        topology=topology_arg,
+                        expected_size=int(manifest_rows[entry]["element_bytes"]),
+                    )
+                except ServiceError as exc:
+                    raise self._translate(exc) from exc
+                if tuple(buffer.shape) != tuple(info.shape):
+                    buffer.close()
+                    raise ExactnessError(
+                        "entry %r restored with shape %r, expected %r"
+                        % (entry, tuple(buffer.shape), info.shape)
+                    )
+                actual = content_root(buffer.view)
+                if actual != info.content_root:
+                    buffer.close()
+                    raise ExactnessError(
+                        "entry %r restored content root %s != recorded %s"
+                        % (entry, actual.hex(), info.content_root.hex())
+                    )
+                found[entry] = _RestoredHandle(entry, buffer, info)
+        except BaseException:
+            self._release_handles(found)
+            raise
+        with self._lock:
+            self._stats["restores"] += 1
+            self._stats["entries_restored"] += len(found)
+            self._stats["restored_bytes"] += sum(
+                handle.buffer.size for handle in found.values()
+            )
+        return found
+
+    @staticmethod
+    def _release_handles(handles):
+        for handle in handles.values():
+            handle.buffer.close()
+
+    @staticmethod
+    def _translate(exc):
+        """Map a typed service error onto the adapter's error hierarchy."""
+        message = str(exc)
+        if exc.class_name == "NotFound":
+            return MetadataMissing(message)
+        if exc.class_name in ("Corrupt", "Malformed", "NonCanonical"):
+            return ExactnessError(message)
+        if exc.class_name in ("LimitExceeded", "BudgetExceeded"):
+            return StagingLimitExceeded(message)
+        return DcpError(message)
+
+    def _entry_info(self, entry):
+        for info in self._storage_data.values():
+            if info.entry == entry:
+                return info
+        raise InvalidMetadata("no storage entry %r in DCP metadata" % entry)
+
+    def restore_entry(self, entry, *, topology=None, axis=0):
+        """Restore exactly one manifest entry, optionally through a topology.
+
+        The service decodes the entry straight into a caller-owned memfd and
+        verifies the shard roots, the whole-tensor logical root (and the
+        re-layout topology root when requested). ``data`` is an explicit copy
+        of those bytes; the zero-copy path is ``WeftStorageReader.load``.
+        """
+        if self._metadata is None:
+            self.read_metadata()
+        handles = self._restore([entry], topology=topology, axis=axis)
+        handle = handles[entry]
+        try:
+            buffer = handle.buffer
+            return RestoredEntry(
+                entry=entry,
+                data=buffer.bytes,
+                dtype=handle.info.dtype,
+                shape=tuple(buffer.shape),
+                content_root=handle.info.content_root,
+                topology_root=buffer.topology_root,
+                shards_read=buffer.shards_read,
+                store_bytes_read=buffer.store_bytes_read,
+            )
+        finally:
+            handle.buffer.close()
+
+    def _read_plan(self, plan, planner):
+        _require_torch("reading checkpoint data")
+        if not plan.items:
+            return
+        storage = self._storage_data
+        if storage is None:
+            raise DcpError("read_data() before read_metadata()")
+        items_by_entry = {}
+        for read_item in plan.items:
+            key = (
+                read_item.storage_index.fqn,
+                None
+                if read_item.storage_index.offset is None
+                else tuple(int(x) for x in read_item.storage_index.offset),
+            )
+            info = None
+            for stored_key, stored_info in storage.items():
+                stored_offset = (
+                    None
+                    if stored_key.offset is None
+                    else tuple(int(x) for x in stored_key.offset)
+                )
+                if stored_key.fqn == key[0] and stored_offset == key[1]:
+                    info = stored_info
+                    break
+            if info is None:
+                raise InvalidMetadata(
+                    "load plan requests storage index %r, absent from DCP metadata"
+                    % (read_item.storage_index,)
+                )
+            items_by_entry.setdefault(info.entry, []).append((read_item, info))
+
+        entries = sorted(items_by_entry)
+        handles = self._restore(entries)
+        try:
+            for entry in entries:
+                info = items_by_entry[entry][0][1]
+                handle = handles[entry]
+                buffer = handle.buffer
+                for read_item, _info in items_by_entry[entry]:
+                    item_type = getattr(read_item.type, "name", str(read_item.type))
+                    if item_type == "BYTE_IO":
+                        planner.load_bytes(read_item, io.BytesIO(buffer.bytes))
+                        continue
+                    dtype = _torch_dtype(info.dtype)
+                    if buffer.size:
+                        # Zero-copy: torch views the caller-owned memfd.
+                        base = _TORCH.frombuffer(buffer.view, dtype=_TORCH.uint8)
+                        base = base.view(dtype).reshape(tuple(info.shape))
+                    else:
+                        base = _TORCH.empty(tuple(info.shape), dtype=dtype)
+                    offsets = tuple(int(x) for x in read_item.storage_offsets)
+                    lengths = tuple(int(x) for x in read_item.lengths)
+                    narrowed = _narrow(base, offsets, lengths)
+                    target = planner.resolve_tensor(read_item).detach()
+                    if tuple(target.shape) != tuple(narrowed.shape):
+                        raise ExactnessError(
+                            "target %r shape %r != stored slice %r"
+                            % (
+                                read_item.dest_index.fqn,
+                                tuple(target.shape),
+                                tuple(narrowed.shape),
+                            )
+                        )
+                    # DCP allocates the destination tensor, so this copy is the
+                    # interface's own contract, not the adapter's staging.
+                    target.copy_(narrowed)
+                    planner.commit_tensor(read_item, target)
+                    del narrowed, base
+        finally:
+            self._release_handles(handles)
+
+
+# ---------------------------------------------------------------------------
+# Async save
+
+
+class AsyncSaveHandle:
+    """Handle returned by :func:`save_checkpoint_async`.
+
+    ``result()`` blocks and re-raises any save error on the caller's thread;
+    ``done()``/``cancel()`` mirror ``concurrent.futures.Future``. The returned
+    metadata object is the same one ``torch.distributed.checkpoint.save``
+    returns.
+    """
+
+    def __init__(self, future, writer, executor):
+        self._future = future
+        self._writer = writer
+        self._executor = executor
+        self.metadata = None
+
+    def result(self, timeout=None):
+        metadata = self._future.result(timeout=timeout)
+        self.metadata = metadata
+        return metadata
+
+    def done(self):
+        return self._future.done()
+
+    def cancel(self):
+        return self._future.cancel()
+
+    def add_done_callback(self, callback):
+        return self._future.add_done_callback(callback)
+
+    @property
+    def writer(self):
+        return self._writer
+
+    @property
+    def stats(self):
+        return self._writer.stats
+
+    def wait(self, timeout=None):
+        return self.result(timeout=timeout)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        if exc_type is None:
+            self.result()
+        self._executor.shutdown(wait=True, cancel_futures=True)
+        return False
+
+
+def save_checkpoint_async(
+    state_dict,
+    *,
+    checkpoint_id,
+    store_root,
+    topology=1,
+    axis=None,
+    shard_strategy="zstd",
+    staging_bytes=DEFAULT_STAGING_BYTES,
+    overwrite=True,
+    init_store=False,
+    planner=None,
+    weftfs_cli=None,
+    timeout=DEFAULT_TIMEOUT_SECONDS,
+    max_workers=1,
+):
+    """Run ``torch.distributed.checkpoint.save`` on a worker thread.
+
+    The encode/store work (resolving tensors, copying them to host memory,
+    writing the staging safetensors container, running ``weftfs checkpoint
+    write``, fsync, metadata publication) happens on a
+    :class:`~concurrent.futures.ThreadPoolExecutor` worker, so a training loop
+    can keep stepping while the checkpoint is stored.
+
+    This still **copies every tensor through host memory** (one tensor at a time,
+    bounded by ``staging_bytes``); it is not a GPU-direct or zero-copy path. The
+    worker count is bounded (default 1) and every worker uses the same hard
+    staging bound.
+    """
+    _require_torch("save_checkpoint_async")
+    import torch.distributed.checkpoint as dcp
+
+    workers = int(max_workers)
+    if workers < 1 or workers > 4:
+        raise ValueError("max_workers must be between 1 and 4")
+    writer = WeftStorageWriter(
+        checkpoint_id,
+        store_root=store_root,
+        topology=topology,
+        axis=axis,
+        shard_strategy=shard_strategy,
+        staging_bytes=staging_bytes,
+        overwrite=overwrite,
+        init_store=init_store,
+        weftfs_cli=weftfs_cli,
+        timeout=timeout,
+    )
+    executor = ThreadPoolExecutor(
+        max_workers=workers, thread_name_prefix="weft-dcp-save"
+    )
+
+    def _run():
+        try:
+            return dcp.save(
+                state_dict,
+                storage_writer=writer,
+                planner=planner,
+                no_dist=True,
+            )
+        finally:
+            executor.shutdown(wait=False)
+
+    future = executor.submit(_run)
+    return AsyncSaveHandle(future, writer, executor)
+
+
+# ---------------------------------------------------------------------------
+# Convenience helpers used by tests, docs, and the bench harness
+
+
+def init_store(store_root, *, weftfs_cli=None, timeout=DEFAULT_TIMEOUT_SECONDS):
+    """Create a WeftFS store at ``store_root`` if it does not exist."""
+    store = _Storage(store_root, cli=weftfs_cli, timeout=timeout)
+    store.init()
+    store.requires_store()
+    return Path(store_root)
+
+
+def verify_checkpoint(
+    store_root,
+    name,
+    *,
+    full=False,
+    weftfs_cli=None,
+    timeout=DEFAULT_TIMEOUT_SECONDS,
+):
+    """Run ``checkpoint verify`` (and optionally ``verify --full``).
+
+    Returns ``{"checkpoint": ..., "tensors_verified": N, "full": {...} | None}``.
+    Every failure is a :class:`StoreCommandError`; nothing is swallowed.
+    """
+    store = _Storage(store_root, cli=weftfs_cli, timeout=timeout)
+    store.requires_store()
+    result = {"checkpoint": name, "tensors_verified": None, "full": None}
+    payload = store.checkpoint_verify(name)
+    result["tensors_verified"] = payload.get("tensors_verified")
+    if full:
+        result["full"] = store.verify_full()
+    return result