Preserve checkpoint storage failures across Ray transport

README.md

@@ -17,6 +17,10 @@
 
 <h2 id="overview">🎉 Overview</h2>
 
+Checkpoint storage failures can cross Ray process boundaries without transferring
+their local writer Future. See [checkpoint error ownership](docs/checkpoint.md#tqsave_checkpoint)
+for the settlement and retirement contract.
+
 TransferQueue is a high-performance data storage and transfer module with panoramic data visibility and streaming scheduling capabilities, optimized for efficient dataflow in post-training workflows.
 
 <p align="center">

docs/checkpoint.md

@@ -50,6 +50,12 @@ def save_checkpoint(
 
 **Raises**: `RuntimeError` if `tq.init()` has not been called, or if any write step fails.
 
+`CheckpointStorageError` and `CheckpointStorageTimeoutError` retain their type,
+message, storage failures and timeout flag when transported through Ray or
+pickle. Their `checkpoint_future` is process-local: only the original error
+retains it for settlement and retirement; a transported copy has `None`.
+Receiving an error does not settle writers or authorize checkpoint publication.
+
 ---
 
 ### `tq.load_checkpoint`
@@ -167,4 +173,4 @@ client.load_controller_checkpoint(...)    # (2) controller restored second
 
 All admitted shard loads settle before returning, including on failure or caller cancellation. If step (1) partially succeeds or step (2) fails, the system is left in a mixed state: some storage units hold checkpoint data while the controller still reflects its pre-restore state. There is no rollback path.
 
-**Workaround**: If `load_checkpoint` raises, stop consumers/producers, close the owned system with `tq.close()`, and create fresh services with `tq.init()` before retrying. Do not continue using partially restored services.
+**Workaround**: If `load_checkpoint` raises, stop consumers/producers, close the owned system with `tq.close()`, and create fresh services with `tq.init()` before retrying. Do not continue using partially restored services.

tests/test_checkpoint_error_transport.py

@@ -0,0 +1,131 @@
+# Copyright 2026 Zyphra Technologies Inc.
+# Licensed under the Apache License, Version 2.0 (the "License").
+"""Checkpoint failures cross Ray without transferring local writer ownership."""
+
+import pickle
+from concurrent.futures import Future
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import ray
+import ray.cloudpickle
+
+import transfer_queue.interface as interface
+from transfer_queue.storage.managers.simple_storage_manager import AsyncSimpleStorageManager
+
+
+@pytest.mark.parametrize("serializer", [pickle, ray.cloudpickle], ids=["pickle", "ray-cloudpickle"])
+@pytest.mark.parametrize("timed_out", [False, True], ids=["write-error", "deadline"])
+def test_error_copy_omits_future_without_changing_local_ownership(serializer, timed_out):
+    kind = interface.CheckpointStorageTimeoutError if timed_out else interface.CheckpointStorageError
+    failures = () if timed_out else ({"storage_unit_id": "unit-2", "storage_io": True, "type": "OSError"},)
+    error = kind("original storage failure", failures=failures, timed_out=timed_out)
+    future = Future()
+    error.checkpoint_future = future
+    future.set_exception(error)
+
+    copied = serializer.loads(serializer.dumps(error))
+
+    assert type(copied) is kind
+    assert copied.args == error.args
+    assert copied.failures == failures
+    assert copied.timed_out is timed_out
+    assert copied.checkpoint_future is None
+    assert error.checkpoint_future is future
+    assert future.exception() is error
+
+
+@pytest.fixture(scope="module")
+def private_ray(tmp_path_factory):
+    assert not ray.is_initialized()
+    ray.init(
+        num_cpus=1,
+        num_gpus=0,
+        include_dashboard=False,
+        object_store_memory=128 * 1024 * 1024,
+        _memory=1024 * 1024 * 1024,
+        _node_ip_address="127.0.0.1",
+        _temp_dir=str(tmp_path_factory.mktemp("error-ray")),
+    )
+    try:
+        yield
+    finally:
+        ray.shutdown()
+
+
+@ray.remote(num_cpus=1)
+class CheckpointOwner:
+    def fail(self, checkpoint_dir, timed_out):
+        self.path = Path(checkpoint_dir)
+        self.settled = False
+        self.failures = (
+            []
+            if timed_out
+            else [{"storage_unit_id": "unit-2", "storage_io": True, "type": "OSError", "message": "write failed"}]
+        )
+
+        def controller(path, **kwargs):
+            Path(path).write_bytes(b"native-controller-cut")
+            return {"indexes": []}
+
+        def storage(path, *, phase, **kwargs):
+            if phase == "prepare":
+                return {"prepared": True, "retained_bytes": 0}
+            assert phase == "status"
+            return {"done": self.settled, "failures": self.failures}
+
+        manager = object.__new__(AsyncSimpleStorageManager)
+        manager.controller_handshake_socket = None
+        manager._owns_zmq_context = False
+        manager.zmq_context = None
+        interface._TQ_CONTROLLER = object()
+        interface._TQ_CLIENT = SimpleNamespace(
+            storage_manager=manager, save_controller_checkpoint=controller, save_storage_checkpoint=storage
+        )
+        self.future = interface.save_checkpoint(
+            self.path, async_storage=True, max_snapshot_bytes=1024, timeout_s=0.05 if timed_out else 5
+        )
+        self.future.result(timeout=10)
+
+    def settle_and_retire(self):
+        import shutil
+
+        error = self.future.exception()
+        assert error.checkpoint_future is self.future
+        assert interface.checkpoint_status(self.path, self.future) == {"settled": False}
+        with pytest.raises(RuntimeError, match="unsettled"):
+            interface.retire_checkpoint(self.path, self.future)
+        with pytest.raises(ValueError, match="owned pending"):
+            interface.checkpoint_status(self.path, Future())
+        self.settled = True
+        assert interface.checkpoint_status(self.path, self.future) == {"settled": True}
+        staging = self.path.parent / f"{self.path.name}.tmp"
+        assert staging.is_dir() and not self.path.exists()
+        shutil.rmtree(staging)
+        interface.retire_checkpoint(self.path, self.future)
+        assert error.checkpoint_future is self.future
+        assert self.future.exception() is error
+        assert not self.path.exists() and not staging.exists()
+        return True
+
+
+@pytest.mark.parametrize("timed_out", [False, True], ids=["write-error", "deadline"])
+def test_public_save_error_crosses_real_ray_and_keeps_local_retirement(private_ray, tmp_path, timed_out):
+    owner = CheckpointOwner.remote()
+    kind = interface.CheckpointStorageTimeoutError if timed_out else interface.CheckpointStorageError
+    try:
+        with pytest.raises(ray.exceptions.RayTaskError) as raised:
+            ray.get(owner.fail.remote(str(tmp_path / "failed-cut"), timed_out), timeout=30)
+        error = raised.value.cause
+        assert type(error) is kind
+        assert error.timed_out is timed_out
+        assert error.failures == (
+            ()
+            if timed_out
+            else ({"storage_unit_id": "unit-2", "storage_io": True, "type": "OSError", "message": "write failed"},)
+        )
+        assert error.checkpoint_future is None
+        assert ray.get(owner.settle_and_retire.remote(), timeout=30)
+    finally:
+        ray.kill(owner, no_restart=True)

transfer_queue/interface.py

@@ -1004,6 +1004,10 @@ def __init__(self, message, *, failures=(), timed_out=False):
         self.timed_out = timed_out
         self.checkpoint_future = None
 
+    def __reduce__(self):
+        # Writer ownership stays local; the failure can cross process boundaries.
+        return type(self), self.args, {**self.__dict__, "checkpoint_future": None}
+
 
 class CheckpointStorageTimeoutError(CheckpointStorageError, TimeoutError):
     """A storage deadline does not attest that its writers have settled."""