[wchwawa]

TLDR

Replace the torch Distributed Checkpoint adapter's persisted pickle metadata and shard values with versioned, non-executable formats, and keep its loads correct on PyTorch 2.2 through 2.10 with the real DCP path covered in CI.

Problem

Checkpoint artifacts could be treated as executable pickle input when loaded. This is a private security follow-up; no public issue was opened because SECURITY.md directs vulnerability details to private reporting.

Review of the first revision found further gaps: every load before PyTorch 2.10 went through the deprecated load_state_dict entry point, nightly or release-candidate builds were refused at import, FP8 and uint16/uint32/uint64 tensors failed at save, and the only real-DCP test skipped in CI because no PyTorch was installed.

Consequences

Loading an untrusted checkpoint could execute code before its contents were validated. On PyTorch 2.4 to 2.9 every load also warned and skipped the Stateful handling that only dcp.load provides, legitimate FP8 checkpoints could not be saved at all, and the four PyTorch planner internals the adapter relies on had no automated coverage.

Example

Previously, the reader deserialized artifact bytes with pickle. Now it accepts only the declared metadata and payload formats and rejects legacy or malformed checkpoint bytes before handing values to PyTorch. A shard whose tensor has a different shape than the read plan now fails with a clear error instead of being broadcast into the target, and a float8_e4m3fn weight saves and loads bit-exactly.

Solution

Use a bounded metadata document and typed shard payloads. Keep the built-in PyTorch planners on a direct checked-value path; preserve the custom planner API using only locally serialized checked values. Add MessagePack to the optional torch dependency set and name that extra when it is missing.

Parse the numeric prefix of torch.__version__ so dev, rc, and local builds work. Load through dcp.load from PyTorch 2.4 on; only 2.2 (its load() gathers keys before honouring no_dist) and 2.3 (its load() discards nested non-tensor values, fixed upstream in 2.4) keep the deprecated entry point. Whitelist the FP8 and unsigned dtypes that the installed PyTorch defines, and require the decoded tensor shape to equal the read plan shape. The wheel CI job installs a pinned CPU-only PyTorch 2.10.0 and requires the real DCP tests, which now fail instead of skipping when PyTorch is present but unusable.

Gain

The standard checkpoint read path no longer interprets executable artifact bytes. Tensor, FP8, unsigned-integer, Stateful, and ordinary optimizer-state round trips are supported, verified on PyTorch 2.2.2, 2.3.0, and 2.10.0, and the planner internals the adapter relies on are exercised by CI on every change.

Tradeoffs

Legacy pickle checkpoints cannot be loaded by this adapter; re-save a trusted original state to migrate. Unsupported Python object types fail closed.

Behavior change relative to main: the reader compared a LoadItemType with WriteItemType.BYTE_IO, so non-tensor leaves such as optimizer step or param_groups could never load; they now load.

Declared limits: the metadata document is at most 64 MiB, and a non-tensor value that MessagePack cannot carry directly (tuples, bytearrays, complex numbers, non-string dictionary keys, tensors nested in Python values) is JSON-encoded and limited to 64 MiB. Tensor shards and plain MessagePack values have no adapter-imposed limit. PyTorch 2.2 and 2.3 still load through the deprecated entry point, so Stateful values need 2.4 or later and single-process loads on those releases need no_dist=True.

The release workflow still runs the SDK tests without PyTorch (macOS x86_64 has no PyTorch wheel after 2.2.2), so the real DCP tests skip there and rely on the pull request job. A 15-trial macOS codec microbenchmark found no regression for a 4 MiB tensor or a 1,000-value optimizer dictionary, but it does not qualify whole-checkpoint throughput or distributed recovery.

Validation

• python -m pytest -q crates/nokv-python/tests/test_torch_unit.py crates/nokv-python/tests/test_torch_dcp_integration.py from the source tree on macOS arm64, Python 3.12: 39 passed, nothing skipped, with PyTorch 2.10.0, 2.3.0, and 2.2.2 (2.2.2 additionally needs numpy for torch.distributed.checkpoint to import).
• Same command without PyTorch: 37 passed, 2 skipped; with NOKV_TEST_REQUIRE_TORCH=1: the 2 real-DCP tests fail with "PyTorch is not installed".
• Installed wheel, as the CI job runs it: maturin build --locked, scripts/release/python_sdk_release.py verify-install, then NOKV_TEST_INSTALLED_WHEEL=1 NOKV_TEST_REQUIRE_TORCH=1 python -m pytest -q crates/nokv-python/tests from a scratch directory with PyTorch 2.10.0: 70 passed.
• pip install --dry-run --no-deps --only-binary=:all: --platform manylinux_2_28_x86_64 --python-version 3.12 --index-url https://download.pytorch.org/whl/cpu torch==2.10.0 resolves to torch-2.10.0+cpu-cp312-cp312-manylinux_2_28_x86_64.whl, the wheel the CI job will install.
• Evidence for the 2.4 gate: on PyTorch 2.3.0, dcp.load with PyTorch's own FileSystemReader leaves {"optimizer": {"step": 0}} unchanged while load_state_dict loads step=4; the same dcp.load probe on 2.10.0 loads it.
• python3 -m py_compile on the three adapter modules, python3 scripts/workbench/workbench_contract_test.py (8 tests OK), and git diff --check: passed.
• cargo test --workspace --locked, cargo clippy --workspace --all-targets --locked -- -D warnings, and cargo fmt --all -- --check: passed on macOS for the first revision; not re-run after the review fixes, which change no Rust files.
• Not run: the pull request CI job itself on Linux, the release workflow, live S3/RustFS, multi-rank fault recovery, and whole-checkpoint latency.