[feat] provide save/load checkpoint interfaces by dodatboii · Pull Request #124 · Ascend/TransferQueue

dodatboii · 2026-06-15T03:25:25Z

Summary

Add tq.save_checkpoint(checkpoint_dir, *, include_storage, metadata) and tq.load_checkpoint(checkpoint_dir) as top-level public APIs
Controller and each storage unit write state directly to file in-process (only a bool ACK goes through Ray object store), avoiding large payload transmission over Ray
Save order: controller first, then storage units in parallel — guarantees consistency without pausing the data flow (storage unit data is always a superset of what the controller has confirmed)
Atomic save via a .tmp directory that is renamed on success and deleted on failure, ensuring no partial checkpoint is left on disk
Load validates storage unit count before touching any state; restoration is position-based (global_idx % num_units) so storage unit IDs regenerated across restarts are handled correctly

Checkpoint layout:

checkpoint_dir/
├── metadata.json          # timestamp, storage unit list, user metadata
├── controller_state.pkl   # TransferQueueController full state
└── storage_units/
    ├── su_0_<id>.pkl
    └── su_1_<id>.pkl

Test plan

pytest tests/e2e/test_checkpoint_e2e.py -v
- save creates expected files and metadata structure
- load restores controller partitions, key mappings, and per-sample tags
- load restores storage data and round-trips tensors correctly across multiple partitions
- include_storage=False saves only controller state
- error cases: uninitialized system, missing directory/metadata, storage unit count mismatch
- failed save leaves no partial directory on disk
- non-tensor fields (NonTensorStack) and variable-length (jagged) tensor fields survive round-trip

ascend-robot · 2026-06-15T03:25:34Z

CLA Signature Pass

dodatboii, thanks for your pull request. All authors of the commits have signed the CLA. 👍

Copilot

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

ascend-robot · 2026-06-15T10:00:58Z

CLA Signature Pass

dodatboii, thanks for your pull request. All authors of the commits have signed the CLA. 👍

ascend-robot · 2026-06-15T10:51:05Z

CLA Signature Pass

dodatboii, thanks for your pull request. All authors of the commits have signed the CLA. 👍

0oshowero0 · 2026-06-16T02:04:17Z

+    CHECKPOINT_DUMP_RESPONSE = "CHECKPOINT_DUMP_RESPONSE"
+    CHECKPOINT_RESTORE = "CHECKPOINT_RESTORE"
+    CHECKPOINT_RESTORE_RESPONSE = "CHECKPOINT_RESTORE_RESPONSE"
+    SAVE_LOAD_CKPT_ERROR = "SAVE_LOAD_CKPT_ERROR"


0oshowero0 · 2026-06-16T02:04:38Z

+    CHECKPOINT_DUMP = "CHECKPOINT_DUMP"
+    CHECKPOINT_DUMP_RESPONSE = "CHECKPOINT_DUMP_RESPONSE"
+    CHECKPOINT_RESTORE = "CHECKPOINT_RESTORE"
+    CHECKPOINT_RESTORE_RESPONSE = "CHECKPOINT_RESTORE_RESPONSE"


Distinguish between controller zmq event & simple storage zmq event

0oshowero0 · 2026-06-16T02:04:50Z

+    @property
+    def storage_checkpoint_required(self) -> bool:
+        """Whether storage contents must be checkpointed for correct restore.
+
+        Returns True for in-memory backends (e.g. SimpleStorage) where data
+        is lost on restart and must be serialized. Returns False for persistent
+        KV backends (e.g. MooncakeStore, Yuanrong) where data survives restarts
+        and only controller metadata needs to be saved.
+
+        Subclasses should override this to reflect their actual persistence model.
+        """
+        return False


This property is strange

0oshowero0 · 2026-06-16T02:05:08Z

+    async def dump_checkpoint(self, output_dir: str) -> list[dict]:
+        """Dump all storage units to files in output_dir.
+
+        Returns:
+            List of dicts, each with keys: position, storage_unit_id.
+
+        Raises:
+            NotImplementedError: If this storage backend does not support checkpoint.
+        """
+        raise NotImplementedError(f"{self.__class__.__name__} does not support checkpoint")
+
+    async def restore_checkpoint(self, checkpoint_dir: str, su_info_list: list[dict]) -> None:
+        """Restore all storage units from files in checkpoint_dir.
+
+        Args:
+            checkpoint_dir: Path to the checkpoint directory containing storage unit files.
+            su_info_list: Ordered list of storage unit info dicts from metadata.json.
+


These interfaces are not universal to other backends

0oshowero0 · 2026-06-16T02:05:25Z

+            logger.error(f"[{self.controller_id}]: dump_to_file failed: {e}")
+            return False
+
+    def restore_from_file(self, path: str) -> bool:


Unify the interface name

Signed-off-by: yxstev <zhangyixiang9@huawei.com>

ascend-robot · 2026-06-16T07:54:01Z

CLA Signature Pass

dodatboii, thanks for your pull request. All authors of the commits have signed the CLA. 👍

Signed-off-by: yxstev <zhangyixiang9@huawei.com>

ascend-robot · 2026-06-17T02:42:19Z

CLA Signature Pass

dodatboii, thanks for your pull request. All authors of the commits have signed the CLA. 👍

0oshowero0 · 2026-06-17T03:51:14Z

+def _put_batch(keys, partition_id, input_ids, attention_mask, tags=None):
+    fields = TensorDict(
+        {"input_ids": input_ids, "attention_mask": attention_mask},
+        batch_size=len(keys),
+    )
+    if tags is None:
+        tags = [{} for _ in keys]
+    tq.kv_batch_put(keys=keys, partition_id=partition_id, fields=fields, tags=tags)
+
+
+def _get_batch(keys, partition_id):
+    return tq.kv_batch_get(keys=keys, partition_id=partition_id)


Just use TQ api. Don't warp it

0oshowero0 · 2026-06-17T03:55:58Z

@@ -0,0 +1,395 @@
+# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved.


Need to refactor the test logics.

Define test data

Put

Save

check saved states

Load

check load states

We can split the unit tests in test_controller.py, test_simple_storage_unit.py

Signed-off-by: yxstev <zhangyixiang9@huawei.com>

ascend-robot · 2026-06-18T02:34:22Z

CLA Signature Pass

dodatboii, thanks for your pull request. All authors of the commits have signed the CLA. 👍

Signed-off-by: yxstev <zhangyixiang9@huawei.com>

ascend-robot · 2026-06-18T03:03:08Z

CLA Signature Pass

dodatboii, thanks for your pull request. All authors of the commits have signed the CLA. 👍

Signed-off-by: yxstev <zhangyixiang9@huawei.com>

ascend-robot · 2026-06-18T03:46:13Z

CLA Signature Pass

dodatboii, thanks for your pull request. All authors of the commits have signed the CLA. 👍

Previously, load_checkpoint restored controller state first, then validated storage unit count. A mismatch would raise ValueError after the controller had already been overwritten, leaving the system in an inconsistent state. Fix: introduce validate_checkpoint on the storage manager that checks file existence and SU count without touching any state. interface.py now runs this validation before any restoration begins. Also adds an assertion to test_load_raises_on_storage_unit_count_mismatch to verify controller state is unchanged after a failed load. Signed-off-by: yxstev <zhangyixiang9@huawei.com>

ascend-robot · 2026-06-26T02:21:26Z

CLA Signature Pass

dodatboii, thanks for your pull request. All authors of the commits have signed the CLA. 👍

ascend-robot added the ascend-cla/yes label Jun 15, 2026

0oshowero0 requested a review from Copilot June 15, 2026 03:47

Copilot AI reviewed Jun 15, 2026

0oshowero0 reviewed Jun 15, 2026

View reviewed changes

Comment thread transfer_queue/interface.py

0oshowero0 reviewed Jun 15, 2026

View reviewed changes

Comment thread transfer_queue/interface.py Outdated

0oshowero0 reviewed Jun 15, 2026

View reviewed changes

Comment thread transfer_queue/interface.py

0oshowero0 reviewed Jun 15, 2026

View reviewed changes

Comment thread transfer_queue/interface.py Outdated

0oshowero0 reviewed Jun 15, 2026

View reviewed changes

Comment thread transfer_queue/interface.py Outdated

0oshowero0 reviewed Jun 15, 2026

View reviewed changes

Comment thread transfer_queue/interface.py Outdated

0oshowero0 reviewed Jun 15, 2026

View reviewed changes

Comment thread transfer_queue/interface.py Outdated

0oshowero0 reviewed Jun 15, 2026

View reviewed changes

Comment thread transfer_queue/interface.py Outdated

0oshowero0 reviewed Jun 15, 2026

View reviewed changes

Comment thread transfer_queue/interface.py Outdated

0oshowero0 reviewed Jun 15, 2026

View reviewed changes

Comment thread transfer_queue/interface.py Outdated

0oshowero0 reviewed Jun 15, 2026

View reviewed changes

Comment thread transfer_queue/storage/simple_storage.py Outdated

0oshowero0 reviewed Jun 15, 2026

View reviewed changes

Comment thread transfer_queue/interface.py Outdated

0oshowero0 reviewed Jun 15, 2026

View reviewed changes

Comment thread transfer_queue/controller.py Outdated

0oshowero0 reviewed Jun 15, 2026

View reviewed changes

Comment thread transfer_queue/controller.py Outdated

0oshowero0 reviewed Jun 15, 2026

View reviewed changes

Comment thread transfer_queue/interface.py

0oshowero0 reviewed Jun 16, 2026

View reviewed changes

provide save/load checkpoint interfaces

255daf8

Signed-off-by: yxstev <zhangyixiang9@huawei.com>

dodatboii force-pushed the dev_dump branch from 41171e5 to 255daf8 Compare June 16, 2026 07:53

fix bug

505e331

Signed-off-by: yxstev <zhangyixiang9@huawei.com>