diff --git a/docs/checkpoint.md b/docs/checkpoint.md new file mode 100644 index 00000000..de6e9342 --- /dev/null +++ b/docs/checkpoint.md @@ -0,0 +1,186 @@ +# Checkpoint: Save and Restore System State + +> Last updated: 07/03/2026 + +## Overview + +TransferQueue provides `tq.save_checkpoint` and `tq.load_checkpoint` to persist and restore the full system state — controller metadata and storage data — enabling fault tolerance and training resumption after a restart. + +A checkpoint captures two components: + +- **Controller state** — all partition metadata, the global index manager, and sampler state. +- **Storage data** — the tensor and non-tensor field data held by each `SimpleStorageUnit`. + +## Quick Start + +```python +import transfer_queue as tq + +tq.init(config) + +# ... put/get data during training ... + +# Save at a step boundary +tq.save_checkpoint("/shared/fs/checkpoints/step_1000", metadata={"step": 1000}) + +# --- After restart --- +tq.init(config) +tq.load_checkpoint("/shared/fs/checkpoints/step_1000") +# System state restored; training can resume +``` + +## API Reference + +### `tq.save_checkpoint` + +```python +def save_checkpoint( + checkpoint_dir: str | Path, + *, + include_storage: bool = True, + metadata: dict[str, Any] | None = None, +) -> None +``` + +| Parameter | Description | +|-----------|-------------| +| `checkpoint_dir` | Directory to write the checkpoint. Created if it does not exist. If a checkpoint already exists at this path it is replaced (best-effort; see [Known Limitations §3](#3-replacing-an-existing-checkpoint-is-not-fully-atomic)). | +| `include_storage` | Whether to save storage unit data. For `SimpleStorage` (in-memory), this is forced to `True` regardless of the value passed — skipping storage would cause complete data loss on restart. For persistent external backends, `False` is valid. | +| `metadata` | Optional user-defined key-value pairs written into `metadata.json`. Useful for recording step number, timestamp, etc. | + +**Raises**: `RuntimeError` if `tq.init()` has not been called, or if any write step fails. + +--- + +### `tq.load_checkpoint` + +```python +def load_checkpoint( + checkpoint_dir: str | Path, +) -> None +``` + +| Parameter | Description | +|-----------|-------------| +| `checkpoint_dir` | Path to a directory previously written by `save_checkpoint`. | + +**Raises**: `FileNotFoundError` if the directory or required files are missing; `ValueError` if the number of storage units in the checkpoint does not match the running system; `RuntimeError` if `tq.init()` has not been called or restore fails. + +**Prerequisite**: `tq.init()` must have been called and the system must be in a clean state (no prior data operations). `load_checkpoint` restores state into a running system; it does not launch components. + +## Checkpoint Directory Layout + +``` +checkpoint_dir/ +├── metadata.json # Flags and user metadata +├── controller_state.pkl # Controller state (pickle) +└── simple_storage/ + ├── storage_unit_info.json # Position-to-ID manifest + ├── su_0_.pkl # StorageUnit at position 0 + ├── su_1_.pkl # StorageUnit at position 1 + └── ... +``` + +**`metadata.json`** + +```json +{ + "storage_saved": true, + "user_metadata": {"step": 1000} +} +``` + +**`storage_unit_info.json`** + +```json +[ + {"position": 0, "storage_unit_id": ""}, + {"position": 1, "storage_unit_id": ""} +] +``` + +## Architecture + +``` +tq.save_checkpoint / tq.load_checkpoint + │ + │ ZMQ RPC + ├──────────────► TransferQueueController (Ray Actor) + │ - partition metadata + │ - index_manager state + │ - sampler state + │ + │ ZMQ RPC (concurrent) + └──────────────► StorageManager + ├── SimpleStorageUnit 0 ──► su_0_.pkl + ├── SimpleStorageUnit 1 ──► su_1_.pkl + └── ... +``` + +Both the controller and each storage unit write their data directly to disk from within their own processes. The ZMQ RPC carries only the target file path and an ACK, not the payload — this avoids routing large tensors through the Ray object store. + +## Save Order and Consistency + +The save sequence is: **controller first, storage units second**. + +This ordering is intentional. In the normal data flow, a `put` writes to the storage unit before the controller's `production_status` is updated. At any point in time, storage units hold a superset (or equal set) of what the controller considers produced. Snapshotting the controller first therefore guarantees that every index the controller records as produced is present in the subsequently snapshotted storage units. + +If storage were snapshotted first, a race could produce a checkpoint where the controller records an index as produced but the corresponding storage data was not yet captured — consumers would read a missing entry on restore. + +**Concurrent `clear` during save**: If `clear_partition` or `clear_samples` runs concurrently with `save_checkpoint`, the controller snapshot may reference indexes that are subsequently deleted from storage before the storage snapshot is taken. The resulting checkpoint reflects a mixed view. To avoid this, callers should ensure no concurrent `clear` operations are issued during `save_checkpoint`. In practice, checkpoints are typically taken at step boundaries where no clearing is in progress. + +## Multi-Node Requirements + +`checkpoint_dir` must reside on a shared filesystem accessible from all nodes (NFS, GPFS, Lustre, etc.), because each StorageUnit writes its file directly to that path. Single-node deployments have no such requirement. + +## Storage Unit Count Matching + +On load, the number of storage units in the checkpoint must exactly match the running system. The matching is by **position** (index in the ordered list), not by storage unit ID — since IDs are regenerated on each `tq.init()`, position-based matching supports restart with freshly created actors. A count mismatch raises `ValueError` and aborts the restore. + +## Known Limitations + +### 1. Controller request thread can hang on checkpoint I/O failure + +The controller's request loop currently has no error handling around the checkpoint branches. If `save_checkpoint` or `load_checkpoint` raises (e.g., the path is not writable, or the pickle file is corrupt), the exception propagates up and kills the request thread. The client's `recv_multipart` call will block indefinitely. + +**Workaround**: Ensure the checkpoint path is writable and the target file is not corrupt before calling. Verify disk space and file system permissions ahead of time. + +--- + +### 2. Save consistency is not guaranteed under concurrent clears + +See [Save Order and Consistency](#save-order-and-consistency) above. Concurrent `clear_partition` or `clear_samples` during `save_checkpoint` can produce a checkpoint whose controller view references storage entries that no longer exist. + +**Workaround**: Do not issue `clear_partition` or `clear_samples` while `save_checkpoint` is running. This is naturally satisfied when checkpointing at training step boundaries. + +--- + +### 3. Replacing an existing checkpoint is not fully atomic + +The current save sequence is: + +```python +if checkpoint_dir.exists(): + shutil.rmtree(checkpoint_dir) # (1) old directory deleted +tmp_dir.rename(checkpoint_dir) # (2) new directory moved into place +``` + +If step (2) fails after step (1) (e.g., cross-device rename, disk full), the old checkpoint has already been deleted and the new one is also cleaned up by the exception handler — both copies are lost. + +**Workaround**: Maintain an additional copy of the previous checkpoint (e.g., save to `step_N` while keeping `step_N-1`) so a failure at step N leaves `step_N-1` intact. + +--- + +### 4. Load is not transactional — partial restore has no rollback + +The load sequence is: + +```python +if meta.get("storage_saved"): + client.load_storage_checkpoint(...) # (1) storage restored first +client.load_controller_checkpoint(...) # (2) controller restored second +``` + +If step (1) partially succeeds and 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, call `tq.init()` again to reset the system to a clean state before retrying. \ No newline at end of file diff --git a/tests/e2e/test_checkpoint_e2e.py b/tests/e2e/test_checkpoint_e2e.py new file mode 100644 index 00000000..f2a8da92 --- /dev/null +++ b/tests/e2e/test_checkpoint_e2e.py @@ -0,0 +1,428 @@ +# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# Copyright 2025 The TransferQueue Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""End-to-end tests for save_checkpoint and load_checkpoint. + +Run with: + pytest tests/e2e/test_checkpoint_e2e.py -v +""" + +import json +import os + +import pytest +import ray +import torch +from omegaconf import OmegaConf +from tensordict import NonTensorStack, TensorDict + +import transfer_queue as tq + +os.environ["RAY_DEDUP_LOGS"] = "0" + +_TQ_CONFIG = OmegaConf.create( + { + "controller": {"polling_mode": True}, + "backend": { + "storage_backend": "SimpleStorage", + "SimpleStorage": { + "total_storage_size": 200, + "num_data_storage_units": 2, + }, + }, + } +) + + +# --------------------------------------------------------------------------- +# fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def ray_init(): + if not ray.is_initialized(): + ray.init(namespace="TestCheckpointE2E") + yield + if ray.is_initialized(): + ray.shutdown() + + +@pytest.fixture(scope="module") +def tq_system(ray_init): + tq.init(_TQ_CONFIG) + yield + tq.close() + + +@pytest.fixture +def controller(tq_system): + return ray.get_actor("TransferQueueController", namespace="transfer_queue") + + +@pytest.fixture(autouse=True) +def cleanup_partitions(controller): + yield + try: + for pid in ray.get(controller.list_partitions.remote()): + ray.get(controller.clear_partition.remote(pid)) + except Exception: + pass + + +@pytest.fixture +def checkpoint_dir(tmp_path): + return tmp_path / "checkpoint" + + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + + +def _assert_tensor_equal(a, b, msg=""): + if (isinstance(a, torch.Tensor) and a.is_nested) or (isinstance(b, torch.Tensor) and b.is_nested): + for t1, t2 in zip(list(a), list(b), strict=True): + assert torch.equal(t1, t2), f"{msg} mismatch" + else: + assert torch.equal(a, b), f"{msg} mismatch" + + +# --------------------------------------------------------------------------- +# basic save / load roundtrip +# --------------------------------------------------------------------------- + + +class TestCheckpointRoundtrip: + """Standard data → save → verify files → wipe → load → verify data.""" + + def test_tensor_fields(self, tq_system, checkpoint_dir, controller): + # Define test data + keys = ["k0", "k1"] + partition_id = "p_tensor" + input_ids = torch.tensor([[1, 2], [3, 4]]) + attention_mask = torch.ones(2, 2) + + # Put + tq.kv_batch_put( + keys=keys, + partition_id=partition_id, + fields=TensorDict({"input_ids": input_ids, "attention_mask": attention_mask}, batch_size=len(keys)), + tags=[{} for _ in keys], + ) + + # Save + tq.save_checkpoint(checkpoint_dir) + + # Check saved state: expected files exist + assert (checkpoint_dir / "metadata.json").exists() + assert (checkpoint_dir / "controller_state.pkl").exists() + su_dir = checkpoint_dir / "simple_storage" + assert su_dir.exists() + assert (su_dir / "storage_unit_info.json").exists() + with open(checkpoint_dir / "metadata.json") as f: + meta = json.load(f) + assert meta["storage_saved"] is True + + # Wipe controller state so load has real work to do + ray.get(controller.clear_partition.remote(partition_id)) + assert ray.get(controller.list_partitions.remote()) == [] + + # Load + tq.load_checkpoint(checkpoint_dir) + + # Check loaded state: partition and data restored + assert partition_id in ray.get(controller.list_partitions.remote()) + retrieved = tq.kv_batch_get(keys=keys, partition_id=partition_id) + _assert_tensor_equal(retrieved["input_ids"], input_ids) + _assert_tensor_equal(retrieved["attention_mask"], attention_mask) + + def test_controller_metadata(self, tq_system, checkpoint_dir, controller): + # Define test data + keys = ["a0", "a1", "a2"] + partition_id = "p_meta" + input_ids = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) + tags = [{"idx": i} for i in range(3)] + + # Put + tq.kv_batch_put( + keys=keys, + partition_id=partition_id, + fields=TensorDict({"input_ids": input_ids, "attention_mask": torch.ones(3, 3)}, batch_size=len(keys)), + tags=tags, + ) + + # Save + tq.save_checkpoint(checkpoint_dir) + + # Wipe + ray.get(controller.clear_partition.remote(partition_id)) + + # Load + tq.load_checkpoint(checkpoint_dir) + + # Check loaded state: key mapping and tags intact + snapshot = ray.get(controller.get_partition_snapshot.remote(partition_id)) + for i, key in enumerate(keys): + assert key in snapshot.keys_mapping + gidx = snapshot.keys_mapping[key] + assert snapshot.custom_meta[gidx]["idx"] == i + + def test_multiple_partitions(self, tq_system, checkpoint_dir, controller): + # Define test data + partitions_data = {f"part_{i}": (torch.full((2, 4), i, dtype=torch.long), torch.ones(2, 4)) for i in range(3)} + + # Put + for pid, (iids, mask) in partitions_data.items(): + tq.kv_batch_put( + keys=[f"{pid}_k0", f"{pid}_k1"], + partition_id=pid, + fields=TensorDict({"input_ids": iids, "attention_mask": mask}, batch_size=2), + tags=[{}, {}], + ) + + # Save + tq.save_checkpoint(checkpoint_dir) + + # Wipe + for pid in partitions_data: + ray.get(controller.clear_partition.remote(pid)) + + # Load + tq.load_checkpoint(checkpoint_dir) + + # Check loaded state + for pid, (iids, _) in partitions_data.items(): + retrieved = tq.kv_batch_get(keys=[f"{pid}_k0", f"{pid}_k1"], partition_id=pid, select_fields=["input_ids"]) + _assert_tensor_equal(retrieved["input_ids"], iids) + + def test_user_metadata_preserved(self, tq_system, checkpoint_dir): + # Define test data + keys = ["m0"] + + # Put + tq.kv_batch_put( + keys=keys, + partition_id="p_usermeta", + fields=TensorDict( + {"input_ids": torch.tensor([[10, 20]]), "attention_mask": torch.ones(1, 2)}, batch_size=1 + ), + tags=[{}], + ) + + # Save with user metadata + tq.save_checkpoint(checkpoint_dir, metadata={"iteration": 42, "loss": 0.5}) + + # Check saved state: user metadata written correctly + with open(checkpoint_dir / "metadata.json") as f: + meta = json.load(f) + assert meta["user_metadata"]["iteration"] == 42 + assert meta["user_metadata"]["loss"] == pytest.approx(0.5) + + def test_non_tensor_fields(self, tq_system, checkpoint_dir, controller): + # Define test data + keys = ["t0", "t1"] + partition_id = "p_str" + input_ids = torch.tensor([[1, 2], [3, 4]]) + fields = TensorDict( + {"input_ids": input_ids, "text": NonTensorStack("hello", "world")}, + batch_size=2, + ) + + # Put + tq.kv_batch_put(keys=keys, partition_id=partition_id, fields=fields, tags=[{}, {}]) + + # Save + tq.save_checkpoint(checkpoint_dir) + + # Wipe + ray.get(controller.clear_partition.remote(partition_id)) + + # Load + tq.load_checkpoint(checkpoint_dir) + + # Check loaded state + retrieved = tq.kv_batch_get(keys=keys, partition_id=partition_id, select_fields=["input_ids"]) + _assert_tensor_equal(retrieved["input_ids"], input_ids) + + def test_nested_tensor_fields(self, tq_system, checkpoint_dir, controller): + # Define test data + keys = ["j0", "j1", "j2"] + partition_id = "p_jagged" + + # Put (variable-length sequences) + for i, key in enumerate(keys): + tq.kv_put( + key=key, + partition_id=partition_id, + fields=TensorDict({"seq": torch.arange(i + 1, dtype=torch.float).unsqueeze(0)}, batch_size=1), + tag=None, + ) + + # Save + tq.save_checkpoint(checkpoint_dir) + + # Wipe + ray.get(controller.clear_partition.remote(partition_id)) + + # Load + tq.load_checkpoint(checkpoint_dir) + + # Check loaded state + retrieved = tq.kv_batch_get(keys=keys, partition_id=partition_id, select_fields=["seq"]) + for i, component in enumerate(retrieved["seq"].unbind()): + _assert_tensor_equal(component, torch.arange(i + 1, dtype=torch.float)) + + +# --------------------------------------------------------------------------- +# include_storage=False (SimpleStorage override) +# --------------------------------------------------------------------------- + + +class TestIncludeStorageFalse: + """For SimpleStorage, include_storage=False is silently forced to True.""" + + def test_storage_saved_is_true(self, tq_system, checkpoint_dir): + # Define test data + Put + tq.kv_batch_put( + keys=["n0"], + partition_id="p_nometa", + fields=TensorDict({"input_ids": torch.tensor([[1, 2]]), "attention_mask": torch.ones(1, 2)}, batch_size=1), + tags=[{}], + ) + + # Save with include_storage=False + tq.save_checkpoint(checkpoint_dir, include_storage=False) + + # Check saved state: storage_saved must be True and directory must exist + with open(checkpoint_dir / "metadata.json") as f: + meta = json.load(f) + assert meta["storage_saved"] is True + assert (checkpoint_dir / "simple_storage").exists() + + def test_both_restored_after_load(self, tq_system, checkpoint_dir, controller): + # Define test data + keys = ["n0", "n1"] + partition_id = "p_nometa2" + input_ids = torch.tensor([[5, 6], [7, 8]]) + + # Put + tq.kv_batch_put( + keys=keys, + partition_id=partition_id, + fields=TensorDict({"input_ids": input_ids, "attention_mask": torch.ones(2, 2)}, batch_size=len(keys)), + tags=[{} for _ in keys], + ) + + # Save + tq.save_checkpoint(checkpoint_dir, include_storage=False) + + # Wipe + ray.get(controller.clear_partition.remote(partition_id)) + + # Load + tq.load_checkpoint(checkpoint_dir) + + # Check loaded state: both controller and storage restored + assert partition_id in ray.get(controller.list_partitions.remote()) + snapshot = ray.get(controller.get_partition_snapshot.remote(partition_id)) + for key in keys: + assert key in snapshot.keys_mapping + retrieved = tq.kv_batch_get(keys=keys, partition_id=partition_id) + _assert_tensor_equal(retrieved["input_ids"], input_ids) + + +# --------------------------------------------------------------------------- +# error handling +# --------------------------------------------------------------------------- + + +class TestCheckpointErrors: + def test_save_raises_if_not_initialized(self, tmp_path): + import transfer_queue.interface as iface + + original = iface._TQ_CONTROLLER + try: + iface._TQ_CONTROLLER = None + with pytest.raises(RuntimeError, match="not initialized"): + tq.save_checkpoint(tmp_path / "ck") + finally: + iface._TQ_CONTROLLER = original + + def test_load_raises_if_not_initialized(self, tmp_path): + import transfer_queue.interface as iface + + original = iface._TQ_CONTROLLER + try: + iface._TQ_CONTROLLER = None + with pytest.raises(RuntimeError, match="not initialized"): + tq.load_checkpoint(tmp_path / "ck") + finally: + iface._TQ_CONTROLLER = original + + def test_load_raises_if_dir_missing(self, tq_system, tmp_path): + with pytest.raises(FileNotFoundError): + tq.load_checkpoint(tmp_path / "nonexistent") + + def test_load_raises_if_metadata_missing(self, tq_system, tmp_path): + ck = tmp_path / "ck" + ck.mkdir() + with pytest.raises(FileNotFoundError, match="metadata.json"): + tq.load_checkpoint(ck) + + def test_load_raises_on_storage_unit_count_mismatch(self, tq_system, tmp_path, checkpoint_dir): + # Define test data + Put + Save + tq.kv_batch_put( + keys=["e0"], + partition_id="p_err", + fields=TensorDict({"input_ids": torch.tensor([[1, 2]]), "attention_mask": torch.ones(1, 2)}, batch_size=1), + tags=[{}], + ) + tq.save_checkpoint(checkpoint_dir) + + # Tamper: add a fake extra SU entry so count differs + su_info_path = checkpoint_dir / "simple_storage" / "storage_unit_info.json" + with open(su_info_path) as f: + su_info = json.load(f) + su_info.append({"position": 99, "storage_unit_id": "fake"}) + with open(su_info_path, "w") as f: + json.dump(su_info, f) + + with pytest.raises(ValueError, match="count mismatch"): + tq.load_checkpoint(checkpoint_dir) + + def test_no_partial_state_on_failed_save(self, tq_system, tmp_path): + # Define test data + Put + tq.kv_batch_put( + keys=["f0"], + partition_id="p_fail", + fields=TensorDict({"input_ids": torch.tensor([[1, 2]]), "attention_mask": torch.ones(1, 2)}, batch_size=1), + tags=[{}], + ) + + ck = tmp_path / "ck" + + import unittest.mock as mock + + with mock.patch( + "transfer_queue.client.TransferQueueClient.save_storage_checkpoint", + side_effect=RuntimeError("simulated dump failure"), + ): + with pytest.raises(RuntimeError, match="simulated dump failure"): + tq.save_checkpoint(ck) + + # Check saved state: no partial directory left + assert not ck.exists() + assert not (tmp_path / "ck.tmp").exists() diff --git a/tests/test_client.py b/tests/test_client.py index 2e7e37a0..d760c98d 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -143,6 +143,12 @@ def _handle_requests(self): elif request_msg.request_type == ZMQRequestType.KV_LIST: response_body = self._mock_kv_list(request_msg.body) response_type = ZMQRequestType.KV_LIST_RESPONSE + elif request_msg.request_type == ZMQRequestType.SAVE_CONTROLLER_CHECKPOINT: + response_body = {"message": "checkpoint saved"} + response_type = ZMQRequestType.SAVE_CONTROLLER_CHECKPOINT_RESPONSE + elif request_msg.request_type == ZMQRequestType.LOAD_CONTROLLER_CHECKPOINT: + response_body = {"message": "checkpoint loaded"} + response_type = ZMQRequestType.LOAD_CONTROLLER_CHECKPOINT_RESPONSE else: response_body = {"error": f"Unknown request type: {request_msg.request_type}"} response_type = ZMQRequestType.CLEAR_META_RESPONSE @@ -1319,3 +1325,123 @@ def test_kv_retrieve_keys_type_validation(self, client_setup): global_indexes=[0, "invalid"], partition_id="test_partition", ) + + +# ===================================================== +# Checkpoint Interface Tests +# ===================================================== + + +class TestClientCheckpointInterface: + """Tests for client save/load checkpoint interface methods.""" + + def test_save_controller_checkpoint_sync(self, client_setup, tmp_path): + """Synchronous save_controller_checkpoint sends correct request.""" + client, _, _ = client_setup + ckpt_path = str(tmp_path / "controller.pkl") + + # Should not raise — mock controller acknowledges the request + client.save_controller_checkpoint(ckpt_path) + + def test_load_controller_checkpoint_sync(self, client_setup, tmp_path): + """Synchronous load_controller_checkpoint sends correct request.""" + client, _, _ = client_setup + ckpt_path = str(tmp_path / "controller.pkl") + + client.load_controller_checkpoint(ckpt_path) + + @pytest.mark.asyncio + async def test_async_save_controller_checkpoint(self, client_setup, tmp_path): + """Async save_controller_checkpoint sends correct request.""" + client, _, _ = client_setup + ckpt_path = str(tmp_path / "controller_async.pkl") + + await client.async_save_controller_checkpoint(ckpt_path) + + @pytest.mark.asyncio + async def test_async_load_controller_checkpoint(self, client_setup, tmp_path): + """Async load_controller_checkpoint sends correct request.""" + client, _, _ = client_setup + ckpt_path = str(tmp_path / "controller_async.pkl") + + await client.async_load_controller_checkpoint(ckpt_path) + + def test_save_storage_checkpoint_sync(self, client_setup, tmp_path): + """Synchronous save_storage_checkpoint delegates to storage manager.""" + client, _, _ = client_setup + ckpt_dir = str(tmp_path / "storage_ckpt") + + called_with = [] + + async def mock_save_checkpoint(directory): + called_with.append(directory) + + client.storage_manager.save_checkpoint = mock_save_checkpoint + client.save_storage_checkpoint(ckpt_dir) + assert called_with == [ckpt_dir] + + def test_load_storage_checkpoint_sync(self, client_setup, tmp_path): + """Synchronous load_storage_checkpoint delegates to storage manager.""" + client, _, _ = client_setup + ckpt_dir = str(tmp_path / "storage_ckpt") + + called_with = [] + + async def mock_load_checkpoint(directory): + called_with.append(directory) + + client.storage_manager.load_checkpoint = mock_load_checkpoint + client.load_storage_checkpoint(ckpt_dir) + assert called_with == [ckpt_dir] + + @pytest.mark.asyncio + async def test_async_save_storage_checkpoint(self, client_setup, tmp_path): + """Async save_storage_checkpoint delegates to storage manager.""" + client, _, _ = client_setup + ckpt_dir = str(tmp_path / "storage_ckpt_async") + + called_with = [] + + async def mock_save_checkpoint(directory): + called_with.append(directory) + + client.storage_manager.save_checkpoint = mock_save_checkpoint + await client.async_save_storage_checkpoint(ckpt_dir) + assert called_with == [ckpt_dir] + + @pytest.mark.asyncio + async def test_async_load_storage_checkpoint(self, client_setup, tmp_path): + """Async load_storage_checkpoint delegates to storage manager.""" + client, _, _ = client_setup + ckpt_dir = str(tmp_path / "storage_ckpt_async") + + called_with = [] + + async def mock_load_checkpoint(directory): + called_with.append(directory) + + client.storage_manager.load_checkpoint = mock_load_checkpoint + await client.async_load_storage_checkpoint(ckpt_dir) + assert called_with == [ckpt_dir] + + def test_save_storage_checkpoint_without_manager(self, mock_controller, tmp_path): + """save_storage_checkpoint raises RuntimeError if storage manager not initialized.""" + client = TransferQueueClient( + client_id="no_storage_client", + controller_info=mock_controller.zmq_server_info, + ) + ckpt_dir = str(tmp_path / "should_fail") + + with pytest.raises(RuntimeError, match="Storage manager not initialized"): + client.save_storage_checkpoint(ckpt_dir) + + def test_load_storage_checkpoint_without_manager(self, mock_controller, tmp_path): + """load_storage_checkpoint raises RuntimeError if storage manager not initialized.""" + client = TransferQueueClient( + client_id="no_storage_client2", + controller_info=mock_controller.zmq_server_info, + ) + ckpt_dir = str(tmp_path / "should_fail") + + with pytest.raises(RuntimeError, match="Storage manager not initialized"): + client.load_storage_checkpoint(ckpt_dir) diff --git a/tests/test_controller.py b/tests/test_controller.py index 36732eb3..b7638cc5 100644 --- a/tests/test_controller.py +++ b/tests/test_controller.py @@ -1157,3 +1157,109 @@ def test_controller_kv_retrieve_keys_multiple_partitions(self, ray_setup): # Clean up ray.get(tq_controller.clear_partition.remote(partition_1)) ray.get(tq_controller.clear_partition.remote(partition_2)) + + +class TestTransferQueueControllerCheckpoint: + """Tests for TransferQueueController save/load checkpoint functionality.""" + + def test_controller_checkpoint_round_trip(self, ray_setup, tmp_path): + """Save and load controller state; verify partition data is preserved.""" + gbs = 4 + partition_id = "ckpt_partition" + ckpt_path = str(tmp_path / "controller.pkl") + + tq_controller = TransferQueueController.remote() + + # Insert some data + data_fields = ["prompt_ids", "attention_mask"] + metadata = ray.get( + tq_controller.get_metadata.remote( + data_fields=data_fields, + batch_size=gbs, + partition_id=partition_id, + mode="insert", + ) + ) + field_schema = { + "prompt_ids": {"dtype": "torch.int64", "shape": (32,)}, + "attention_mask": {"dtype": "torch.bool", "shape": (32,)}, + } + ray.get( + tq_controller.update_production_status.remote( + partition_id=partition_id, + global_indexes=metadata.global_indexes, + field_schema=field_schema, + ) + ) + + # Save checkpoint + ray.get(tq_controller.save_checkpoint.remote(ckpt_path)) + assert (tmp_path / "controller.pkl").exists() + + # Create a fresh controller and restore + new_controller = TransferQueueController.remote() + ray.get(new_controller.load_checkpoint.remote(ckpt_path)) + + # Verify partition state was restored + partition = ray.get(new_controller.get_partition_snapshot.remote(partition_id)) + assert partition is not None + assert set(partition.global_indexes) == set(metadata.global_indexes) + assert partition.production_status is not None + + index_range = ray.get(new_controller.get_partition_index_range.remote(partition_id)) + assert sorted(index_range) == sorted(metadata.global_indexes) + + print("✓ Controller checkpoint round-trip preserves partition state") + + def test_controller_checkpoint_restores_index_manager(self, ray_setup, tmp_path): + """Verify that global_index counter is correctly restored after load.""" + ckpt_path = str(tmp_path / "controller.pkl") + + tq_controller = TransferQueueController.remote() + + # Insert into two partitions to advance the index counter + for pid in ["ckpt_p1", "ckpt_p2"]: + meta = ray.get( + tq_controller.get_metadata.remote( + data_fields=["f"], + batch_size=4, + partition_id=pid, + mode="insert", + ) + ) + ray.get( + tq_controller.update_production_status.remote( + partition_id=pid, + global_indexes=meta.global_indexes, + field_schema={"f": {"dtype": "torch.float32", "shape": (1,)}}, + ) + ) + + ray.get(tq_controller.save_checkpoint.remote(ckpt_path)) + + new_controller = TransferQueueController.remote() + ray.get(new_controller.load_checkpoint.remote(ckpt_path)) + + # New inserts should continue from where the counter left off (no index collision) + meta_after = ray.get( + new_controller.get_metadata.remote( + data_fields=["f"], + batch_size=2, + partition_id="ckpt_p3", + mode="insert", + ) + ) + existing_indexes = set(range(8)) # 2 partitions × 4 samples + assert not set(meta_after.global_indexes) & existing_indexes, "Index collision after checkpoint restore" + + print("✓ Index manager counter correctly restored") + + def test_controller_checkpoint_load_nonexistent_file(self, ray_setup, tmp_path): + """Loading from a missing file should raise RuntimeError.""" + missing = str(tmp_path / "does_not_exist.pkl") + tq_controller = TransferQueueController.remote() + + with pytest.raises(ray.exceptions.RayTaskError): + ray.get(tq_controller.load_checkpoint.remote(missing)) + + print("✓ load_checkpoint raises on missing file") diff --git a/tests/test_simple_storage_unit.py b/tests/test_simple_storage_unit.py index 0e30fa55..e36063ee 100644 --- a/tests/test_simple_storage_unit.py +++ b/tests/test_simple_storage_unit.py @@ -64,6 +64,24 @@ def send_clear(self, client_id, global_indexes): self.socket.send_multipart(msg.serialize()) return ZMQMessage.deserialize(self.socket.recv_multipart(copy=False)) + def send_save_checkpoint(self, client_id, path): + msg = ZMQMessage.create( + request_type=ZMQRequestType.SAVE_STORAGE_CHECKPOINT, + sender_id=f"mock_client_{client_id}", + body={"path": path}, + ) + self.socket.send_multipart(msg.serialize()) + return ZMQMessage.deserialize(self.socket.recv_multipart(copy=False)) + + def send_load_checkpoint(self, client_id, path): + msg = ZMQMessage.create( + request_type=ZMQRequestType.LOAD_STORAGE_CHECKPOINT, + sender_id=f"mock_client_{client_id}", + body={"path": path}, + ) + self.socket.send_multipart(msg.serialize()) + return ZMQMessage.deserialize(self.socket.recv_multipart(copy=False)) + def close(self): self.socket.close() self.context.term() @@ -611,3 +629,105 @@ def wrong_len_parser(field_data): assert "data_parser changed the number of elements" in response.body["message"] client.close() + + +def test_storage_unit_checkpoint_round_trip(storage_setup, tmp_path): + """Save storage state to a file, load it into a fresh unit, verify data.""" + _, put_get_address, storage_ip = storage_setup + ckpt_path = str(tmp_path / "storage_unit.pkl") + client = MockStorageClient(put_get_address, storage_ip) + + # 1. Put some data + global_indexes = [10, 11, 12] + field_data = { + "log_probs": [torch.tensor([1.0, 2.0]), torch.tensor([3.0, 4.0]), torch.tensor([5.0, 6.0])], + "rewards": [torch.tensor([0.1]), torch.tensor([0.2]), torch.tensor([0.3])], + } + response = client.send_put(0, global_indexes, field_data) + assert response.request_type == ZMQRequestType.PUT_DATA_RESPONSE + + # 2. Save checkpoint + response = client.send_save_checkpoint(0, ckpt_path) + assert response.request_type == ZMQRequestType.SAVE_STORAGE_CHECKPOINT_RESPONSE + assert response.body["success"] is True + assert (tmp_path / "storage_unit.pkl").exists() + + # 3. Create a fresh storage unit and load the checkpoint into it + fresh_actor = SimpleStorageUnit.options(max_concurrency=50, num_cpus=1).remote(storage_unit_size=10000) + fresh_zmq_info = ray.get(fresh_actor.get_zmq_server_info.remote()) + import time as _time + + _time.sleep(1) + fresh_address = fresh_zmq_info.to_addr("put_get_socket") + fresh_client = MockStorageClient(fresh_address, fresh_zmq_info.ip) + + response = fresh_client.send_load_checkpoint(0, ckpt_path) + assert response.request_type == ZMQRequestType.LOAD_STORAGE_CHECKPOINT_RESPONSE + assert response.body["success"] is True + + # 4. Verify data is accessible in the fresh unit + response = fresh_client.send_get(0, global_indexes, ["log_probs", "rewards"]) + assert response.request_type == ZMQRequestType.GET_DATA_RESPONSE + retrieved = response.body["data"] + torch.testing.assert_close(retrieved["log_probs"][0], torch.tensor([1.0, 2.0])) + torch.testing.assert_close(retrieved["log_probs"][2], torch.tensor([5.0, 6.0])) + torch.testing.assert_close(retrieved["rewards"][1], torch.tensor([0.2])) + + fresh_client.close() + client.close() + ray.kill(fresh_actor) + + +def test_storage_unit_checkpoint_overwrites_existing_data(storage_setup, tmp_path): + """Loading a checkpoint into a unit that already has data replaces it entirely.""" + _, put_get_address, storage_ip = storage_setup + ckpt_path = str(tmp_path / "storage_unit_overwrite.pkl") + client = MockStorageClient(put_get_address, storage_ip) + + # 1. Put original data and save checkpoint + response = client.send_put(0, [20, 21], {"val": [torch.tensor([1.0]), torch.tensor([2.0])]}) + assert response.request_type == ZMQRequestType.PUT_DATA_RESPONSE + response = client.send_save_checkpoint(0, ckpt_path) + assert response.body["success"] is True + + # 2. Create a second unit, pre-populate it with different data, then load the checkpoint + second_actor = SimpleStorageUnit.options(max_concurrency=50, num_cpus=1).remote(storage_unit_size=10000) + second_zmq_info = ray.get(second_actor.get_zmq_server_info.remote()) + import time as _time + + _time.sleep(1) + second_address = second_zmq_info.to_addr("put_get_socket") + second_client = MockStorageClient(second_address, second_zmq_info.ip) + + # 3. Write different data into the second unit before loading + response = second_client.send_put(0, [99], {"val": [torch.tensor([999.0])]}) + assert response.request_type == ZMQRequestType.PUT_DATA_RESPONSE + + # 4. Load checkpoint — should overwrite + response = second_client.send_load_checkpoint(0, ckpt_path) + assert response.body["success"] is True + + # 5. Old data (index 99) should be gone; checkpoint data (indexes 20, 21) should be present + response = second_client.send_get(0, [20, 21], ["val"]) + assert response.request_type == ZMQRequestType.GET_DATA_RESPONSE + retrieved = response.body["data"] + torch.testing.assert_close(retrieved["val"][0], torch.tensor([1.0])) + torch.testing.assert_close(retrieved["val"][1], torch.tensor([2.0])) + + second_client.close() + client.close() + ray.kill(second_actor) + + +def test_storage_unit_checkpoint_load_missing_file(storage_setup, tmp_path): + """Loading from a non-existent file returns success=False.""" + _, put_get_address, storage_ip = storage_setup + client = MockStorageClient(put_get_address, storage_ip) + missing_path = str(tmp_path / "does_not_exist.pkl") + + response = client.send_load_checkpoint(0, missing_path) + assert response.request_type == ZMQRequestType.LOAD_STORAGE_CHECKPOINT_RESPONSE + assert response.body["success"] is False + assert "message" in response.body + + client.close() diff --git a/transfer_queue/__init__.py b/transfer_queue/__init__.py index 06754278..ad8e8ab9 100644 --- a/transfer_queue/__init__.py +++ b/transfer_queue/__init__.py @@ -34,6 +34,8 @@ kv_clear, kv_list, kv_put, + load_checkpoint, + save_checkpoint, ) from .metadata import BatchMeta, KVBatchMeta from .sampler import BaseSampler @@ -62,6 +64,11 @@ "async_kv_clear", "KVBatchMeta", ] + + [ + # Checkpoint Interface + "save_checkpoint", + "load_checkpoint", + ] + [ # High-Level StreamingDataLoader Interface "StreamingDataset", diff --git a/transfer_queue/client.py b/transfer_queue/client.py index bd1e57b7..2c131a77 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -1064,6 +1064,130 @@ def close(self) -> None: except Exception as e: logger.warning(f"Error closing storage manager: {e}") + # ==================== Checkpoint API ==================== + @with_controller_socket + async def async_save_controller_checkpoint( + self, + path: str, + socket: zmq.asyncio.Socket | None = None, + ) -> None: + """Asynchronously save controller state to a file via ZMQ RPC. + + Sends a SAVE_CONTROLLER_CHECKPOINT request to the controller and waits + for acknowledgement. The controller serializes its state directly to + ``path`` in-process. + + Args: + path: Absolute path for the output .pkl file. The caller must + ensure this path is writable from the node running the controller. + socket: ZMQ socket injected by @with_controller_socket. + + Raises: + RuntimeError: If the RPC fails or an unexpected response is received. + """ + try: + assert socket is not None + request_msg = ZMQMessage.create( + request_type=ZMQRequestType.SAVE_CONTROLLER_CHECKPOINT, # type: ignore[arg-type] + sender_id=self.client_id, + receiver_id=self._controller.id, + body={"path": path}, + ) + await socket.send_multipart(request_msg.serialize()) + response_serialized = await socket.recv_multipart(copy=False) + response_msg = ZMQMessage.deserialize(response_serialized) + if response_msg.request_type != ZMQRequestType.SAVE_CONTROLLER_CHECKPOINT_RESPONSE: + raise RuntimeError( + f"[{self.client_id}]: Unexpected response type {response_msg.request_type} " + f"from controller during checkpoint dump" + ) + except Exception as e: + raise RuntimeError(f"[{self.client_id}]: Error in save_controller_checkpoint: {str(e)}") from e + + @with_controller_socket + async def async_load_controller_checkpoint( + self, + path: str, + socket: zmq.asyncio.Socket | None = None, + ) -> None: + """Asynchronously restore controller state from a file via ZMQ RPC. + + Sends a LOAD_CONTROLLER_CHECKPOINT request to the controller and waits + for acknowledgement. The controller deserializes its state directly from + ``path`` in-process, overwriting the current in-memory state. + + Args: + path: Absolute path to a .pkl file previously written by + ``async_save_controller_checkpoint``. The caller must ensure + this path is readable from the node running the controller. + socket: ZMQ socket injected by @with_controller_socket. + + Raises: + RuntimeError: If the RPC fails or an unexpected response is received. + """ + try: + assert socket is not None + request_msg = ZMQMessage.create( + request_type=ZMQRequestType.LOAD_CONTROLLER_CHECKPOINT, # type: ignore[arg-type] + sender_id=self.client_id, + receiver_id=self._controller.id, + body={"path": path}, + ) + await socket.send_multipart(request_msg.serialize()) + response_serialized = await socket.recv_multipart(copy=False) + response_msg = ZMQMessage.deserialize(response_serialized) + if response_msg.request_type != ZMQRequestType.LOAD_CONTROLLER_CHECKPOINT_RESPONSE: + raise RuntimeError( + f"[{self.client_id}]: Unexpected response type {response_msg.request_type} " + f"from controller during checkpoint restore" + ) + except Exception as e: + raise RuntimeError(f"[{self.client_id}]: Error in load_controller_checkpoint: {str(e)}") from e + + async def async_save_storage_checkpoint(self, checkpoint_dir: str) -> None: + """Asynchronously save storage state to a directory via StorageManager. + + Delegates to the underlying StorageManager, which fans out save requests + to all storage units concurrently. + + Args: + checkpoint_dir: Directory under which storage unit files are written. + The StorageManager creates a ``storage_units/`` subdirectory here. + + Raises: + RuntimeError: If the storage manager is not initialized. + NotImplementedError: If the storage backend does not support checkpoint. + """ + if not hasattr(self, "storage_manager") or self.storage_manager is None: + raise RuntimeError( + f"[{self.client_id}]: Storage manager not initialized. " + "Call initialize_storage_manager() before checkpoint operations." + ) + await self.storage_manager.save_checkpoint(checkpoint_dir) + + async def async_load_storage_checkpoint(self, checkpoint_dir: str) -> None: + """Asynchronously restore storage state from a directory via StorageManager. + + Delegates to the underlying StorageManager, which fans out load requests + to all storage units concurrently. Existing in-memory data is overwritten. + + Args: + checkpoint_dir: Directory previously written by + ``async_save_storage_checkpoint``, containing the + ``storage_units/`` subdirectory and manifest. + + Raises: + RuntimeError: If the storage manager is not initialized. + FileNotFoundError: If the storage unit manifest is missing. + ValueError: If the number of storage units does not match the checkpoint. + """ + if not hasattr(self, "storage_manager") or self.storage_manager is None: + raise RuntimeError( + f"[{self.client_id}]: Storage manager not initialized. " + "Call initialize_storage_manager() before checkpoint operations." + ) + await self.storage_manager.load_checkpoint(checkpoint_dir) + class TransferQueueClient(AsyncTransferQueueClient): """Synchronous client wrapper for TransferQueue. @@ -1132,6 +1256,10 @@ def wrapper(*args, **kwargs): self._kv_retrieve_meta = _make_sync(self.async_kv_retrieve_meta) self._kv_retrieve_keys = _make_sync(self.async_kv_retrieve_keys) self._kv_list = _make_sync(self.async_kv_list) + self._save_controller_checkpoint = _make_sync(self.async_save_controller_checkpoint) + self._load_controller_checkpoint = _make_sync(self.async_load_controller_checkpoint) + self._save_storage_checkpoint = _make_sync(self.async_save_storage_checkpoint) + self._load_storage_checkpoint = _make_sync(self.async_load_storage_checkpoint) # ==================== Basic API ==================== def get_meta( @@ -1564,6 +1692,55 @@ def kv_list( return self._kv_list(partition_id=partition_id) + # ==================== Checkpoint API ==================== + def save_controller_checkpoint(self, path: str) -> None: + """Synchronously save controller state to a file via ZMQ RPC. + + Args: + path: Absolute path for the output .pkl file. + + Raises: + RuntimeError: If the RPC fails or an unexpected response is received. + """ + return self._save_controller_checkpoint(path) + + def load_controller_checkpoint(self, path: str) -> None: + """Synchronously restore controller state from a file via ZMQ RPC. + + Args: + path: Absolute path to a .pkl file previously written by + ``save_controller_checkpoint``. + + Raises: + RuntimeError: If the RPC fails or an unexpected response is received. + """ + return self._load_controller_checkpoint(path) + + def save_storage_checkpoint(self, checkpoint_dir: str) -> None: + """Synchronously save storage state to a directory via StorageManager. + + Args: + checkpoint_dir: Directory under which storage unit files are written. + + Raises: + RuntimeError: If the storage manager is not initialized. + NotImplementedError: If the storage backend does not support checkpoint. + """ + return self._save_storage_checkpoint(checkpoint_dir) + + def load_storage_checkpoint(self, checkpoint_dir: str) -> None: + """Synchronously restore storage state from a directory via StorageManager. + + Args: + checkpoint_dir: Directory previously written by ``save_storage_checkpoint``. + + Raises: + RuntimeError: If the storage manager is not initialized. + FileNotFoundError: If the storage unit manifest is missing. + ValueError: If the number of storage units does not match the checkpoint. + """ + return self._load_storage_checkpoint(checkpoint_dir) + def close(self) -> None: """Close the client and cleanup resources including event loop and thread.""" diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index 7c315223..52dc2ee4 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -15,6 +15,7 @@ import copy import os +import pickle import time from collections import defaultdict from dataclasses import dataclass, field @@ -2056,6 +2057,26 @@ def _process_request(self): body={"partition_info": partition_info, "message": message}, ) + elif request_msg.request_type == ZMQRequestType.SAVE_CONTROLLER_CHECKPOINT: + path = request_msg.body["path"] + self.save_checkpoint(path) + response_msg = ZMQMessage.create( + request_type=ZMQRequestType.SAVE_CONTROLLER_CHECKPOINT_RESPONSE, + sender_id=self.controller_id, + receiver_id=request_msg.sender_id, + body={"success": True}, + ) + + elif request_msg.request_type == ZMQRequestType.LOAD_CONTROLLER_CHECKPOINT: + path = request_msg.body["path"] + self.load_checkpoint(path) + response_msg = ZMQMessage.create( + request_type=ZMQRequestType.LOAD_CONTROLLER_CHECKPOINT_RESPONSE, + sender_id=self.controller_id, + receiver_id=request_msg.sender_id, + body={"success": True}, + ) + self.request_handle_socket.send_multipart([identity, *response_msg.serialize()]) def get_zmq_server_info(self) -> ZMQServerInfo: @@ -2070,6 +2091,64 @@ def get_config(self) -> DictConfig: """Retrieve the global config of TransferQueue.""" return self.tq_config + def save_checkpoint(self, path: str) -> None: + """Serialize controller state directly to a file. + + Writes in-process to avoid transmitting the payload back over the + Ray object store. + + Args: + path: Absolute path for the output .pkl file. + + Raises: + Exception: If serialization or file I/O fails. + """ + try: + state = { + "controller_id": self.controller_id, + "partitions": {pid: p.to_snapshot() for pid, p in self.partitions.items()}, + "index_manager": { + "partition_to_indexes": dict(copy.deepcopy(self.index_manager.partition_to_indexes)), + "reusable_indexes": list(self.index_manager.reusable_indexes), + "global_index_counter": self.index_manager.global_index_counter, + "allocated_indexes": set(self.index_manager.allocated_indexes), + }, + "sampler": self.sampler.save_checkpoint(), + } + with open(path, "wb") as f: + pickle.dump(state, f, protocol=pickle.HIGHEST_PROTOCOL) + logger.info(f"[{self.controller_id}]: dumped to {path}") + except Exception as e: + raise RuntimeError(f"[{self.controller_id}]: save checkpoint failed: {e}") from e + + def load_checkpoint(self, path: str) -> None: + """Restore controller state directly from a file. + + Args: + path: Absolute path to a .pkl file previously written by save_checkpoint(). + + Raises: + Exception: If deserialization or file I/O fails. + """ + try: + with open(path, "rb") as f: + state = pickle.load(f) + + self.controller_id = state["controller_id"] + self.partitions = state["partitions"] + + im = state["index_manager"] + self.index_manager.partition_to_indexes = defaultdict(set, im["partition_to_indexes"]) + self.index_manager.reusable_indexes = im["reusable_indexes"] + self.index_manager.global_index_counter = im["global_index_counter"] + self.index_manager.allocated_indexes = im["allocated_indexes"] + + self.sampler.load_checkpoint(state["sampler"]) + + logger.info(f"[{self.controller_id}]: restored from {path}") + except Exception as e: + raise RuntimeError(f"[{self.controller_id}]: load checkpoint failed: {e}") from e + def register_sampler( self, sampler: BaseSampler | type[BaseSampler] = SequentialSampler, diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index c3dd8018..fd129560 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -13,10 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json import os +import shutil import subprocess import time from importlib import resources +from pathlib import Path from typing import Any, Callable import ray @@ -31,6 +34,7 @@ from transfer_queue.sampler import * # noqa: F401 from transfer_queue.sampler import BaseSampler from transfer_queue.storage.bootstrap import StorageBootstrapProvider +from transfer_queue.storage.managers.simple_storage_manager import AsyncSimpleStorageManager from transfer_queue.utils.logging_utils import get_logger from transfer_queue.utils.yuanrong_utils import cleanup_yuanrong_resources from transfer_queue.utils.zmq_utils import process_zmq_server_info @@ -1058,3 +1062,132 @@ def get_client(): """Get a TransferQueueClient for using low-level API""" assert _TQ_CLIENT is not None, "Please initialize the TransferQueue first by calling `tq.init()`!" return _TQ_CLIENT + + +# ==================== Checkpoint API ==================== + +_METADATA_FILE = "metadata.json" +_CONTROLLER_FILE = "controller_state.pkl" + + +def save_checkpoint( + checkpoint_dir: str | Path, + *, + include_storage: bool = True, + metadata: dict[str, Any] | None = None, +) -> None: + """Save a full checkpoint of the TransferQueue system state. + + .. note:: + **Multi-node limitation**: checkpoint_dir must reside on a shared network + filesystem (e.g. NFS, GPFS, Lustre) accessible from all nodes. + Single-node deployments have no such requirement. + + Args: + checkpoint_dir: Directory to save the checkpoint (created if not exists). + include_storage: Whether to save storage backend state. Set to False to + save controller metadata only (e.g. for KV backends where + data persists externally and does not need to be re-saved). + For SimpleStorage (in-memory), this flag is ignored and + storage is always saved — skipping it would cause data loss + on restart. + metadata: User-defined key-value pairs written into metadata.json. + Example: {"time_stamp": 1234567.891234, "step": 10} + + Raises: + RuntimeError: TransferQueue is not initialized. + OSError: Failed to write checkpoint files. + """ + if _TQ_CONTROLLER is None: + raise RuntimeError("TransferQueue is not initialized. Call tq.init() first.") + + checkpoint_dir = Path(checkpoint_dir) + tmp_dir = checkpoint_dir.parent / (checkpoint_dir.name + ".tmp") + + if tmp_dir.exists(): + shutil.rmtree(tmp_dir) + tmp_dir.mkdir(parents=True) + + try: + client = _maybe_create_tq_client() + + controller_path = tmp_dir / _CONTROLLER_FILE + client.save_controller_checkpoint(str(controller_path)) + logger.info("Controller state saved.") + + if not include_storage and isinstance(client.storage_manager, AsyncSimpleStorageManager): + logger.warning( + "include_storage=False has no effect for SimpleStorage: " + "in-memory data would be lost on restart. Forcing include_storage=True." + ) + include_storage = True + + if include_storage: + try: + client.save_storage_checkpoint(str(tmp_dir)) + except NotImplementedError: + logger.warning("Storage backend does not support checkpoint; storage data will not be saved.") + include_storage = False + + meta_content = { + "storage_saved": include_storage, + "user_metadata": metadata or {}, + } + with open(tmp_dir / _METADATA_FILE, "w") as f: + json.dump(meta_content, f, indent=2) + + if checkpoint_dir.exists(): + shutil.rmtree(checkpoint_dir) + tmp_dir.rename(checkpoint_dir) + + logger.info(f"Checkpoint saved to {checkpoint_dir}") + + except Exception: + if tmp_dir.exists(): + shutil.rmtree(tmp_dir) + raise + + +def load_checkpoint( + checkpoint_dir: str | Path, +) -> None: + """Restore TransferQueue system state from a checkpoint. + + .. note:: + **Multi-node limitation**: checkpoint_dir must reside on a shared network + filesystem (e.g. NFS, GPFS, Lustre) accessible from all nodes. + Single-node deployments have no such requirement. + + Args: + checkpoint_dir: Path to the checkpoint directory. + + Raises: + FileNotFoundError: Checkpoint directory or required files do not exist. + RuntimeError: TransferQueue is not initialized, or restore fails. + """ + if _TQ_CONTROLLER is None: + raise RuntimeError("TransferQueue is not initialized. Call tq.init() first.") + + checkpoint_dir = Path(checkpoint_dir) + if not checkpoint_dir.exists(): + raise FileNotFoundError(f"Checkpoint directory not found: {checkpoint_dir}") + + metadata_path = checkpoint_dir / _METADATA_FILE + if not metadata_path.exists(): + raise FileNotFoundError(f"{_METADATA_FILE} not found in {checkpoint_dir}") + + with open(metadata_path) as f: + meta = json.load(f) + + client = _maybe_create_tq_client() + + controller_path = checkpoint_dir / _CONTROLLER_FILE + if not controller_path.exists(): + raise FileNotFoundError(f"{_CONTROLLER_FILE} not found in {checkpoint_dir}") + + if meta.get("storage_saved"): + client.load_storage_checkpoint(str(checkpoint_dir)) + + client.load_controller_checkpoint(str(controller_path)) + + logger.info(f"Checkpoint loaded from {checkpoint_dir}") diff --git a/transfer_queue/sampler/base.py b/transfer_queue/sampler/base.py index c831ecba..173b19da 100644 --- a/transfer_queue/sampler/base.py +++ b/transfer_queue/sampler/base.py @@ -124,3 +124,29 @@ def clear_cache(self, partition_id: str): """ if partition_id in self._states.keys(): self._states.pop(partition_id) + + def save_checkpoint(self) -> dict: + """Return the sampler's serializable state for checkpointing. + + By default this returns ``_states``, the shared cache used by all + samplers. Subclasses that maintain additional runtime state (e.g. + a separate balanced cache) should override this method and include + their extra fields. + + Returns: + A dict that can be passed to ``restore_state`` to recreate the + current sampling state. + """ + return {"_states": self._states} + + def load_checkpoint(self, state: dict) -> None: + """Restore the sampler's state from a checkpoint. + + By default this restores ``_states``. Subclasses that override + ``save_checkpoint`` must also override this method to restore their extra + fields. + + Args: + state: A dict previously returned by ``save_checkpoint``. + """ + self._states = state["_states"] diff --git a/transfer_queue/sampler/seqlen_balanced_sampler.py b/transfer_queue/sampler/seqlen_balanced_sampler.py index 1dd10740..20c904d7 100644 --- a/transfer_queue/sampler/seqlen_balanced_sampler.py +++ b/transfer_queue/sampler/seqlen_balanced_sampler.py @@ -196,6 +196,17 @@ def clear_cache(self, partition_id: str): for k in keys_to_remove: del self._balanced_cache[k] + def save_checkpoint(self) -> dict: + """Return sampler state including the balanced cache.""" + state = super().save_checkpoint() + state["_balanced_cache"] = self._balanced_cache + return state + + def load_checkpoint(self, state: dict) -> None: + """Restore sampler state including the balanced cache.""" + super().load_checkpoint(state) + self._balanced_cache = state["_balanced_cache"] + # Copied from https://github.com/volcengine/verl/blob/468adf22c43b744348051fccd7a5d830c6c3c36a/verl/utils/seqlen_balancing.py def karmarkar_karp(seqlen_list: list[int], k_partitions: int, equal_size: bool): diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index f4d545da..e6b0faf4 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -352,6 +352,25 @@ async def clear_data(self, metadata: BatchMeta) -> None: """ raise NotImplementedError("Subclasses must implement clear_data") + async def save_checkpoint(self, checkpoint_dir: str) -> None: + """Save storage state into checkpoint_dir. + + The implementation is backend-specific: each backend decides what to + persist and how to lay out files within checkpoint_dir. + + Raises: + NotImplementedError: If this storage backend does not support checkpoint. + """ + raise NotImplementedError(f"{self.__class__.__name__} does not support checkpoint") + + async def load_checkpoint(self, checkpoint_dir: str) -> None: + """Restore storage state from checkpoint_dir. + + Raises: + NotImplementedError: If this storage backend does not support checkpoint. + """ + raise NotImplementedError(f"{self.__class__.__name__} does not support checkpoint") + def close(self) -> None: """Close all ZMQ sockets/contexts and stop the notify loop.""" diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py index 80803522..0b6777fb 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -19,6 +19,7 @@ from collections import defaultdict from collections.abc import Mapping from operator import itemgetter +from pathlib import Path from typing import Any, Callable, NamedTuple import torch @@ -40,6 +41,9 @@ TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT = int(os.environ.get("TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT", 200)) # seconds +_SU_SUBDIR = "simple_storage" +_SU_INFO_FILE = "storage_unit_info.json" + # Pre-bound decorator for storage-unit socket operations. with_storage_unit_socket = with_zmq_socket( "put_get_socket", @@ -524,6 +528,121 @@ def get_zmq_server_info(self) -> dict[str, ZMQServerInfo]: """ return self.storage_unit_infos + @with_storage_unit_socket + async def _save_single_storage_unit( + self, + path: str, + target_storage_unit: str, + socket: zmq.Socket = None, + ): + try: + request_msg = ZMQMessage.create( + request_type=ZMQRequestType.SAVE_STORAGE_CHECKPOINT, # type: ignore[arg-type] + sender_id=self.storage_manager_id, + receiver_id=target_storage_unit, + body={"path": path}, + ) + await socket.send_multipart(request_msg.serialize(), copy=False) + messages = await socket.recv_multipart(copy=False) + response_msg = ZMQMessage.deserialize(messages) + if ( + response_msg.request_type != ZMQRequestType.SAVE_STORAGE_CHECKPOINT_RESPONSE + or not response_msg.body.get("success") + ): + raise RuntimeError( + f"Storage unit {target_storage_unit} failed to save checkpoint to {path}: " + f"{response_msg.body.get('message', 'unknown error')}" + ) + except Exception as e: + raise RuntimeError( + f"[{self.storage_manager_id}]: Error dumping storage unit {target_storage_unit}: {str(e)}" + ) from e + + @with_storage_unit_socket + async def _load_single_storage_unit( + self, + path: str, + target_storage_unit: str, + socket: zmq.Socket = None, + ): + try: + request_msg = ZMQMessage.create( + request_type=ZMQRequestType.LOAD_STORAGE_CHECKPOINT, # type: ignore[arg-type] + sender_id=self.storage_manager_id, + receiver_id=target_storage_unit, + body={"path": path}, + ) + await socket.send_multipart(request_msg.serialize(), copy=False) + messages = await socket.recv_multipart(copy=False) + response_msg = ZMQMessage.deserialize(messages) + if ( + response_msg.request_type != ZMQRequestType.LOAD_STORAGE_CHECKPOINT_RESPONSE + or not response_msg.body.get("success") + ): + raise RuntimeError( + f"Storage unit {target_storage_unit} failed to load checkpoint from {path}: " + f"{response_msg.body.get('message', 'unknown error')}" + ) + except Exception as e: + raise RuntimeError( + f"[{self.storage_manager_id}]: Error restoring for storage unit {target_storage_unit}: {str(e)}" + ) from e + + async def save_checkpoint(self, checkpoint_dir: str) -> None: + """Dump all storage units to the storage_units/ subdirectory of checkpoint_dir. + + Writes a storage_units/su_info.json manifest for load_checkpoint to use. + Validates current SU count matches the manifest on load. + """ + import json + + su_dir = Path(checkpoint_dir) / _SU_SUBDIR + su_dir.mkdir(parents=True, exist_ok=True) + + su_ids = list(self.storage_unit_infos.keys()) + paths = [str(su_dir / f"su_{pos}_{su_id}.pkl") for pos, su_id in enumerate(su_ids)] + tasks = [ + self._save_single_storage_unit(path, target_storage_unit=su_id) + for su_id, path in zip(su_ids, paths, strict=False) + ] + await asyncio.gather(*tasks) + + su_info_list = [{"position": pos, "storage_unit_id": su_id} for pos, su_id in enumerate(su_ids)] + with open(su_dir / _SU_INFO_FILE, "w") as f: + json.dump(su_info_list, f) + + logger.info(f"[{self.storage_manager_id}]: saved {len(su_ids)} storage units to {su_dir}") + + async def load_checkpoint(self, checkpoint_dir: str) -> None: + """Restore all storage units from the storage_units/ subdirectory of checkpoint_dir.""" + import json + + su_dir = Path(checkpoint_dir) / _SU_SUBDIR + su_info_path = su_dir / _SU_INFO_FILE + if not su_info_path.exists(): + raise FileNotFoundError(f"Storage unit manifest not found: {su_info_path}") + + with open(su_info_path) as f: + su_info_list = json.load(f) + + su_ids = list(self.storage_unit_infos.keys()) + if len(su_ids) != len(su_info_list): + raise ValueError( + f"Storage unit count mismatch: checkpoint has {len(su_info_list)}, current system has {len(su_ids)}." + ) + + entries = sorted(su_info_list, key=lambda e: e["position"]) + tasks = [ + self._load_single_storage_unit( + str(su_dir / f"su_{entry['position']}_{entry['storage_unit_id']}.pkl"), + target_storage_unit=su_ids[entry["position"]], + ) + for entry in entries + ] + await asyncio.gather(*tasks) + + logger.info(f"[{self.storage_manager_id}]: restored {len(su_ids)} storage units from {su_dir}") + def close(self) -> None: """Close all ZMQ sockets and context to prevent resource leaks.""" super().close() diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py index 6865a608..aa23855f 100644 --- a/transfer_queue/storage/simple_storage.py +++ b/transfer_queue/storage/simple_storage.py @@ -14,6 +14,7 @@ # limitations under the License. import os +import pickle import time import weakref from threading import Event, Thread @@ -310,6 +311,10 @@ def _worker_routine(self) -> None: response_msg = self._handle_clear(request_msg) elif operation == ZMQRequestType.GET_METRICS: # type: ignore[arg-type] response_msg = self._handle_get_metrics() + elif operation == ZMQRequestType.SAVE_STORAGE_CHECKPOINT: # type: ignore[arg-type] + response_msg = self._handle_save_checkpoint(request_msg) + elif operation == ZMQRequestType.LOAD_STORAGE_CHECKPOINT: # type: ignore[arg-type] + response_msg = self._handle_load_checkpoint(request_msg) else: response_msg = ZMQMessage.create( request_type=ZMQRequestType.PUT_GET_OPERATION_ERROR, # type: ignore[arg-type] @@ -537,11 +542,101 @@ def _handle_get_metrics(self) -> ZMQMessage: metrics["op_stats"] = op_stats return ZMQMessage.create( - request_type=ZMQRequestType.METRICS_RESPONSE, + request_type=ZMQRequestType.METRICS_RESPONSE, # type: ignore[arg-type] sender_id=self.storage_unit_id, body=metrics, ) + def _handle_save_checkpoint(self, data_parts) -> ZMQMessage: + """Serialize storage unit data directly to a file. + + Args: + data_parts: ZMQMessage from client, containing ``path`` in body: + absolute path for the output .pkl file. The caller must ensure + this path is reachable from the node running this actor + (shared filesystem required for multi-node setups). + + Returns: + ZMQMessage with ``success=True`` on success, or ``success=False`` + and ``message`` containing the error string on failure. + """ + path = data_parts.body["path"] + try: + state = { + "storage_unit_id": self.storage_unit_id, + "storage_unit_size": self.storage_unit_size, + "field_data": self.storage_data.field_data, + "active_keys": self.storage_data._active_keys, + } + with open(path, "wb") as f: + pickle.dump(state, f, protocol=pickle.HIGHEST_PROTOCOL) + logger.info(f"[{self.storage_unit_id}]: saved checkpoint to {path}") + return ZMQMessage.create( + request_type=ZMQRequestType.SAVE_STORAGE_CHECKPOINT_RESPONSE, # type: ignore[arg-type] + sender_id=self.storage_unit_id, + body={"success": True}, + ) + except Exception as e: + logger.error(f"[{self.storage_unit_id}]: save checkpoint failed: {e}") + return ZMQMessage.create( + request_type=ZMQRequestType.SAVE_STORAGE_CHECKPOINT_RESPONSE, # type: ignore[arg-type] + sender_id=self.storage_unit_id, + body={"success": False, "message": str(e)}, + ) + + def _handle_load_checkpoint(self, data_parts) -> ZMQMessage: + """Restore storage unit data directly from a file. + + Args: + data_parts: ZMQMessage from client, containing ``path`` in body: + absolute path to a .pkl file previously written by + ``_handle_save_checkpoint``. The caller must ensure this path + is reachable from the node running this actor + (shared filesystem required for multi-node setups). + + Returns: + ZMQMessage with ``success=True`` on success, or ``success=False`` + and ``message`` containing the error string on failure. + """ + path = data_parts.body["path"] + try: + with open(path, "rb") as f: + data = pickle.load(f) + + if data["storage_unit_size"] != self.storage_unit_size: + logger.warning( + f"[{self.storage_unit_id}]: storage_unit_size mismatch — " + f"checkpoint={data['storage_unit_size']}, current={self.storage_unit_size}" + ) + + if self.storage_data._active_keys: + logger.warning( + f"[{self.storage_unit_id}]: overwriting {len(self.storage_data._active_keys)} " + f"existing keys with checkpoint data from {path}" + ) + self.storage_data.field_data.clear() + self.storage_data._active_keys.clear() + self.storage_data.field_data = data["field_data"] + self.storage_data._active_keys = data["active_keys"] + + logger.info( + f"[{self.storage_unit_id}]: loaded checkpoint from {path} — " + f"{len(data['active_keys'])} keys, {len(data['field_data'])} fields" + ) + return ZMQMessage.create( + request_type=ZMQRequestType.LOAD_STORAGE_CHECKPOINT_RESPONSE, # type: ignore[arg-type] + sender_id=self.storage_unit_id, + body={"success": True}, + ) + + except Exception as e: + logger.error(f"[{self.storage_unit_id}]: load checkpoint failed: {e}") + return ZMQMessage.create( + request_type=ZMQRequestType.LOAD_STORAGE_CHECKPOINT_RESPONSE, # type: ignore[arg-type] + sender_id=self.storage_unit_id, + body={"success": False, "message": str(e)}, + ) + @staticmethod def _cumulative_bucket_counts(hist) -> list[float]: """Build cumulative counts from a prometheus_client Histogram's non-cumulative buckets.""" diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index 89608738..a88f68a2 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -101,6 +101,16 @@ class ZMQRequestType(ExplicitEnum): GET_METRICS = "GET_METRICS" METRICS_RESPONSE = "METRICS_RESPONSE" + # CHECKPOINT + SAVE_CONTROLLER_CHECKPOINT = "SAVE_CONTROLLER_CHECKPOINT" + SAVE_CONTROLLER_CHECKPOINT_RESPONSE = "SAVE_CONTROLLER_CHECKPOINT_RESPONSE" + SAVE_STORAGE_CHECKPOINT = "SAVE_STORAGE_CHECKPOINT" + SAVE_STORAGE_CHECKPOINT_RESPONSE = "SAVE_STORAGE_CHECKPOINT_RESPONSE" + LOAD_CONTROLLER_CHECKPOINT = "LOAD_CONTROLLER_CHECKPOINT" + LOAD_CONTROLLER_CHECKPOINT_RESPONSE = "LOAD_CONTROLLER_CHECKPOINT_RESPONSE" + LOAD_STORAGE_CHECKPOINT = "LOAD_STORAGE_CHECKPOINT" + LOAD_STORAGE_CHECKPOINT_RESPONSE = "LOAD_STORAGE_CHECKPOINT_RESPONSE" + class ZMQServerInfo: """