diff --git a/docs/checkpoint.md b/docs/checkpoint.md index 8c0a2829..38cf0edb 100644 --- a/docs/checkpoint.md +++ b/docs/checkpoint.md @@ -101,6 +101,9 @@ checkpoint_dir/ ] ``` +Loading unpickles the controller state and every storage unit file, which can run +arbitrary code. Load only checkpoints from directories that you trust. + ## Atomic Checkpoint Replacement `save_checkpoint` writes to `.tmp`, renames the existing `checkpoint_dir` to `.old`, renames `.tmp` into `checkpoint_dir`, then deletes `.old`. This keeps the old checkpoint recoverable until the new one is fully in place; a failure partway through restores `.old` automatically, and the directory stays a plain folder (no symlink). @@ -167,4 +170,11 @@ 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 +**Workaround**: If `load_checkpoint` raises, call `tq.init()` again to reset the system to a clean state before retrying. + +## Exporting selected keys + +Use [selective data dumps](data_dump.md) when restoring selected keys into an +existing system or a different number of storage units. SimpleStorage restores read +assigned row ranges directly on the current storage owners and merge values instead +of replacing entire unit and controller state. diff --git a/docs/data_dump.md b/docs/data_dump.md new file mode 100644 index 00000000..f3829b5b --- /dev/null +++ b/docs/data_dump.md @@ -0,0 +1,123 @@ +# Selective data dump + +`dump_data_by_key` persists selected keys, their produced fields and tags. +`load_data_by_key` merges them into a running TransferQueue. It preserves existing +key indexes and unrelated rows, fields and tag entries; new keys receive indexes +from the current controller. It does not restore sampler or consumption state. + +```python +import transfer_queue as tq + +tq.init() +tq.dump_data_by_key("/shared/dumps/selected", ["sample-1", "sample-2"], "train") +tags = tq.load_data_by_key("/shared/dumps/selected") # {key: tag} of the restored keys +``` + +Pause writes and clears for these keys during both operations. A dump is not an +atomic snapshot of concurrent writers. Each dump goes to a new directory: an existing +target is refused, so a published dump never changes and can be loaded from a +read-only location. + +## Distributed I/O + +On export, the storage manager groups source indexes by their current storage +owner and concurrently asks those units to write their records. Only units holding +selected rows participate. Tensor storage is compacted during serialization, so a +row view cannot include the rest of its original batch, including inside tags. + +On SimpleStorage restore: + +1. The caller reads the row index and shard manifest and validates all file ranges. +2. The controller resolves existing keys and allocates indexes for new keys. +3. The storage manager routes records by the **current** indexes, then sends one + load request to each participating unit concurrently. +4. Each target unit reads only its assigned byte ranges and merges those values + into local storage. Records are processed in batches of at most 128 rows per + shard; the caller never reads or forwards their payloads. +5. After every unit has succeeded, the storage manager publishes the saved schemas + to the controller, as an ordinary put does, and the client then writes the tags. + +The number of source units can differ from the number of destination units. +Even a dump with one source shard can restore across several target units because +records are independently addressable. Empty rows are recreated from metadata. + +The dump directory must be on a filesystem accessible to every participating +storage unit. Local temporary storage suffices for single-node deployments. + +| Operation | State | Payload I/O | Unit count on restore | +| --- | --- | --- | --- | +| Checkpoint | Entire controller and storage state | Each unit reads/writes its whole file | Must match | +| Selective dump | Selected fields and tags, merged by key | Each owner unit reads/writes its records | May differ | + +`DUMP_ROWS` and `LOAD_ROWS` are included in storage operation metrics. Unit logs +record loaded rows and bytes; the manager logs total bytes and participating units. +These count application reads, not filesystem read-ahead or physical disk traffic. + +## Format and compatibility + +Dumps use `format_version: 3`, the only version this build reads: + +```text +dump_info.json +row_index.pt +shards/ + shard_info.json + shard_0_.pkl + ... +``` + +Each shard is a sequence of independent pickle records containing a source global +index and a field/value mapping. `shard_info.json` records each shard's file name +and each source index's `[offset, length]`. Source indexes only locate records; they are never reused as +current indexes without controller resolution. + +Loading unpickles `row_index.pt` and every shard record, which can run arbitrary +code. Load only dumps from directories that you trust. + +A dump also saves a schema for each selected field. The controller supplies the +declared type; while writing their records, the owner units report each row's dtype +and shape, which the caller merges without seeing payloads. Controller metadata can +trail the stored values, for example when a later put wraps rows of a tensor field in +`NonTensorStack`, so a field is saved as non-tensor (with a warning) unless every +selected row is a tensor of one dtype, and as nested if row shapes differ. A field +declared non-tensor stays non-tensor. Restore uses that schema regardless of target +topology or batch boundaries; destination type conflicts are rejected before payload +writes. + +Dump and load currently require SimpleStorage. Loading into another backend raises +`NotImplementedError` before any key is registered. + +## Failure behavior + +A dump is staged in a uniquely named sibling `.tmp-` directory, with +`dump_info.json` written last, and renamed into place once everything is synced. +A failed dump removes its staging directory; a crash can leave one behind, which is +never read and can be deleted. If two dumps race to one path, only the first rename +succeeds. Readers need no lock and no write access. + +Restore has the failure semantics of `kv_batch_put`: it is not transactional, and +payload writes before a failure remain. Metadata is published only after every unit +has succeeded, so writes to new keys stay invisible after a failure, but fields that +existing keys have already produced are overwritten in place and are readable at once. +New keys stay registered: `kv_list` shows them with empty tags and no readable +fields. Every key keeps its index, so retrying the same load is idempotent and +restores the tags; clearing the keys abandons it. +Like an ordinary put, a load does not fence late writes against indexes that are +cleared and reused while it runs, so keep writers and clears for these keys paused. + +Each storage unit serves requests on one worker thread: other partitions using that +unit can wait behind a load. The 128-row batches bound memory, not request latency. +Dump and load requests use their own connection pool, whose timeout is +`TQ_SIMPLE_STORAGE_DUMP_TIMEOUT` (3600 seconds by default) rather than the put/get +timeout, because each unit receives a single request covering all of its rows. + +## Tests + +Run the selective E2E suite with its default pytest-managed temporary directory: + +```bash +python -m pytest -q tests/e2e/test_data_dump_e2e.py tests/e2e/test_data_dump_cross_topology_e2e.py +``` + +For a multi-node Ray cluster, set `TQ_DUMP_TEST_ROOT` to an existing shared directory. +Tests create and remove only their own child directories beneath that root. diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py new file mode 100644 index 00000000..b68de91a --- /dev/null +++ b/tests/e2e/conftest.py @@ -0,0 +1,32 @@ +# 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. + +import os +import shutil +import tempfile +from pathlib import Path + +import pytest + + +@pytest.fixture(scope="module") +def dump_test_root(tmp_path_factory): + shared_root = os.environ.get("TQ_DUMP_TEST_ROOT") + if shared_root: + root = Path(tempfile.mkdtemp(prefix="tq-dump-", dir=shared_root)) + else: + root = tmp_path_factory.mktemp("tq-dump") + yield root + shutil.rmtree(root) diff --git a/tests/e2e/test_data_dump_cross_topology_e2e.py b/tests/e2e/test_data_dump_cross_topology_e2e.py new file mode 100644 index 00000000..445bdbd3 --- /dev/null +++ b/tests/e2e/test_data_dump_cross_topology_e2e.py @@ -0,0 +1,164 @@ +# 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. + +"""A dump taken with N storage units must restore into a system with M storage units. + +This is the property that separates a data dump from a checkpoint. ``load_checkpoint`` +sends each storage unit file back to the unit at the same position, so it requires the +same unit count; ``load_data_by_key`` writes rows back by key and lets TransferQueue +route them for the current topology. + +Each test restarts TransferQueue with a different unit count, so this lives apart from +``test_data_dump_e2e.py``, whose fixtures hold one system for the whole module. + +Run with: + pytest tests/e2e/test_data_dump_cross_topology_e2e.py -v +""" + +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" + + +def _tq_config(num_storage_units: int) -> OmegaConf: + return OmegaConf.create( + { + "controller": {"polling_mode": True}, + "backend": { + "storage_backend": "SimpleStorage", + "SimpleStorage": { + "total_storage_size": 400, + "num_data_storage_units": num_storage_units, + }, + }, + } + ) + + +@pytest.fixture(scope="module") +def ray_init(): + if not ray.is_initialized(): + ray.init(namespace="TestDataDumpCrossTopology") + yield + if ray.is_initialized(): + ray.shutdown() + + +@pytest.fixture +def dump_dir(dump_test_root, request): + return dump_test_root / request.node.name / "dump" + + +def _row_input_ids(row: int) -> torch.Tensor: + return torch.tensor([row * 10, row * 10 + 1, row * 10 + 2]) + + +def _put_rows(partition_id: str, keys: list[str]) -> None: + tq.kv_batch_put( + keys=keys, + partition_id=partition_id, + fields=TensorDict( + {"input_ids": torch.stack([_row_input_ids(row) for row in range(len(keys))])}, + batch_size=len(keys), + ), + tags=[{"idx": row} for row in range(len(keys))], + ) + + +def _assert_rows_equal(actual: torch.Tensor, expected_rows: list[torch.Tensor]) -> None: + actual_rows = list(actual.unbind()) if actual.is_nested else list(actual) + assert len(actual_rows) == len(expected_rows) + for actual_row, expected_row in zip(actual_rows, expected_rows, strict=True): + assert torch.equal(actual_row, expected_row) + + +# The controller and each storage unit reserve one CPU, and CI runners have four, +# so more than three units never get scheduled there. +@pytest.mark.parametrize( + ("dump_units", "load_units"), + [(3, 2), (2, 3), (1, 3)], +) +def test_dump_restores_across_storage_unit_counts(ray_init, dump_dir, dump_units, load_units): + # Define test data + partition_id = "cross" + keys = [f"c{i}" for i in range(8)] + + # Dump with one topology + tq.init(_tq_config(dump_units)) + try: + _put_rows(partition_id, keys) + report = tq.dump_data_by_key(dump_dir, keys, partition_id) + assert report["rows_with_data"] == len(keys) + assert report["shards"] <= dump_units + finally: + tq.close() + + # Restore into a different topology + tq.init(_tq_config(load_units)) + try: + _put_rows("bystander", ["unrelated"]) + _put_rows(partition_id, [keys[3]]) + old_index = tq.get_client().kv_retrieve_meta([keys[3]], partition_id).global_indexes[0] + tq.load_data_by_key(dump_dir) + assert tq.get_client().kv_retrieve_meta([keys[3]], partition_id).global_indexes == [old_index] + assert tq.kv_batch_get(["unrelated"], "bystander", ["input_ids"]).batch_size[0] == 1 + + # Check restored state: every row readable, payload and tag intact + retrieved = tq.kv_batch_get(keys=keys, partition_id=partition_id, select_fields=["input_ids"]) + _assert_rows_equal(retrieved["input_ids"], [_row_input_ids(row) for row in range(len(keys))]) + + controller = ray.get_actor("TransferQueueController", namespace="transfer_queue") + snapshot = ray.get(controller.get_partition_snapshot.remote(partition_id)) + for row, key in enumerate(keys): + assert snapshot.custom_meta[snapshot.keys_mapping[key]]["idx"] == row + finally: + tq.close() + + +def test_preserves_nontensor_schema_across_topology(ray_init, dump_dir): + # Spans two 128-row load batches; test_data_dump_e2e.py covers each batch boundary on one system. + row_count = 130 + keys = [f"k{i}" for i in range(row_count)] + values = [torch.tensor([i], dtype=torch.int64) for i in range(row_count - 1)] + [torch.tensor([1.5])] + values[-2] = None + tq.init(_tq_config(1)) + try: + tq.kv_batch_put(keys, "mixed", TensorDict({"x": NonTensorStack(*values)}, batch_size=row_count)) + tq.dump_data_by_key(dump_dir, keys, "mixed") + finally: + tq.close() + tq.init(_tq_config(2)) + try: + tq.load_data_by_key(dump_dir) + schema = tq.get_client().kv_retrieve_meta(keys, "mixed").field_schema["x"] + assert schema["is_non_tensor"] + assert schema["dtype"] is None + assert not schema["is_nested"] + for key, expected in zip(keys, values, strict=True): + value = tq.kv_batch_get([key], "mixed", ["x"])["x"][0] + if expected is None: + assert value is None + else: + torch.testing.assert_close(value, expected) + finally: + tq.close() diff --git a/tests/e2e/test_data_dump_e2e.py b/tests/e2e/test_data_dump_e2e.py new file mode 100644 index 00000000..a4fd5b9c --- /dev/null +++ b/tests/e2e/test_data_dump_e2e.py @@ -0,0 +1,734 @@ +# 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 dump_data_by_key / load_data_by_key. + +Each storage unit writes its own shard from the node it runs on, so +``TQ_DUMP_TEST_ROOT`` must point at a filesystem shared by the whole cluster. +Single-node runs default to pytest-managed temporary storage. + +Run with: + pytest tests/e2e/test_data_dump_e2e.py -v +""" + +import asyncio +import builtins +import json +import os +import shutil +from pathlib import Path +from unittest.mock import AsyncMock + +import pytest +import ray +import torch +from omegaconf import OmegaConf +from tensordict import NonTensorStack, TensorDict + +import transfer_queue as tq +from transfer_queue.data_dump import _read_row_index + +os.environ["RAY_DEDUP_LOGS"] = "0" + +# CI runners have four CPUs, and the controller and each storage unit reserve one. +_NUM_STORAGE_UNITS = 3 + + +def _tq_config(num_storage_units: int) -> OmegaConf: + return OmegaConf.create( + { + "controller": {"polling_mode": True}, + "backend": { + "storage_backend": "SimpleStorage", + "SimpleStorage": { + "total_storage_size": 400, + "num_data_storage_units": num_storage_units, + }, + }, + } + ) + + +# --------------------------------------------------------------------------- +# fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def ray_init(): + if not ray.is_initialized(): + ray.init(namespace="TestDataDumpE2E") + yield + if ray.is_initialized(): + ray.shutdown() + + +@pytest.fixture(scope="module") +def tq_system(ray_init): + tq.init(_tq_config(_NUM_STORAGE_UNITS)) + 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 dump_dir(dump_test_root, request): + case = dump_test_root / request.node.name.replace("/", "_") + yield case / "dump" + shutil.rmtree(case, ignore_errors=True) + + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + + +def _row_input_ids(row: int) -> torch.Tensor: + """Deterministic per-row payload so a restored row can be traced to its key.""" + return torch.tensor([row * 10, row * 10 + 1, row * 10 + 2]) + + +def _put_rows(partition_id: str, keys: list[str]) -> None: + tq.kv_batch_put( + keys=keys, + partition_id=partition_id, + fields=TensorDict( + { + "input_ids": torch.stack([_row_input_ids(row) for row in range(len(keys))]), + "attention_mask": torch.ones(len(keys), 3), + }, + batch_size=len(keys), + ), + tags=[{"idx": row} for row in range(len(keys))], + ) + + +def _assert_rows_equal(actual: torch.Tensor, expected_rows: list[torch.Tensor]) -> None: + """Compare a retrieved field row by row. + + TransferQueue returns a batched field as a nested tensor, which ``torch.equal`` + cannot consume directly. + """ + actual_rows = list(actual.unbind()) if actual.is_nested else list(actual) + assert len(actual_rows) == len(expected_rows) + for actual_row, expected_row in zip(actual_rows, expected_rows, strict=True): + assert torch.equal(actual_row, expected_row) + + +def _keys_mapping(controller, partition_id: str) -> dict[str, int]: + snapshot = ray.get(controller.get_partition_snapshot.remote(partition_id)) + return dict(snapshot.keys_mapping) + + +# --------------------------------------------------------------------------- +# dump / load roundtrip +# --------------------------------------------------------------------------- + + +class TestDumpLoadRoundtrip: + def test_only_selected_keys_are_restored(self, tq_system, dump_dir, controller): + # Define test data + partition_id = "d_basic" + keys = [f"k{i}" for i in range(6)] + selected = ["k1", "k4"] + _put_rows(partition_id, keys) + + # Dump + report = tq.dump_data_by_key(dump_dir, selected, partition_id) + assert report == {"keys": 2, "rows_with_data": 2, "shards": report["shards"], "bytes": report["bytes"]} + + # Wipe, then restore + ray.get(controller.clear_partition.remote(partition_id)) + tq.load_data_by_key(dump_dir) + + # Check restored state + assert sorted(_keys_mapping(controller, partition_id)) == selected + retrieved = tq.kv_batch_get(keys=selected, partition_id=partition_id, select_fields=["input_ids"]) + _assert_rows_equal(retrieved["input_ids"], [_row_input_ids(keys.index(key)) for key in selected]) + + def test_tags_survive_the_roundtrip(self, tq_system, dump_dir, controller): + # Define test data + partition_id = "d_tags" + keys = [f"t{i}" for i in range(5)] + selected = ["t0", "t3"] + _put_rows(partition_id, keys) + + # Dump + wipe + load + tq.dump_data_by_key(dump_dir, selected, partition_id) + ray.get(controller.clear_partition.remote(partition_id)) + assert tq.load_data_by_key(dump_dir) == {key: {"idx": keys.index(key)} for key in selected} + + # Check restored state + snapshot = ray.get(controller.get_partition_snapshot.remote(partition_id)) + for key in selected: + assert snapshot.custom_meta[snapshot.keys_mapping[key]]["idx"] == keys.index(key) + + def test_jagged_and_non_tensor_fields_survive(self, tq_system, dump_dir, controller): + """The shard stores per-row values; packing them back must reproduce the container.""" + # Define test data: variable-length rows plus a string field + partition_id = "d_jagged" + keys = ["j0", "j1", "j2"] + for row, key in enumerate(keys): + tq.kv_put( + key=key, + partition_id=partition_id, + fields=TensorDict( + { + "seq": torch.arange(row + 1, dtype=torch.float).unsqueeze(0), + "text": NonTensorStack(f"row-{row}"), + }, + batch_size=1, + ), + tag=None, + ) + + # Dump + wipe + load + tq.dump_data_by_key(dump_dir, keys, partition_id) + ray.get(controller.clear_partition.remote(partition_id)) + tq.load_data_by_key(dump_dir) + + # Check restored state + retrieved = tq.kv_batch_get(keys=keys, partition_id=partition_id, select_fields=["seq", "text"]) + for row, component in enumerate(retrieved["seq"].unbind()): + assert torch.equal(component, torch.arange(row + 1, dtype=torch.float)) + assert list(retrieved["text"]) == [f"row-{row}" for row in range(len(keys))] + + def test_heterogeneous_field_sets_are_grouped(self, tq_system, dump_dir, controller): + """A selective dump routinely mixes rows that finished different fields.""" + # Define test data: h1 has an extra field the others lack + partition_id = "d_hetero" + _put_rows(partition_id, ["h0", "h1", "h2"]) + tq.kv_put( + key="h1", + partition_id=partition_id, + fields=TensorDict({"routed_experts": torch.tensor([[7, 8]])}, batch_size=1), + tag=None, + ) + + # Dump + wipe + load + tq.dump_data_by_key(dump_dir, ["h0", "h1", "h2"], partition_id) + ray.get(controller.clear_partition.remote(partition_id)) + tq.load_data_by_key(dump_dir) + + # Check restored state: the extra field came back only on its own row + rows = _read_row_index(dump_dir)["rows"] + assert "routed_experts" in rows["h1"]["fields"] + assert "routed_experts" not in rows["h0"]["fields"] + retrieved = tq.kv_batch_get(keys=["h1"], partition_id=partition_id, select_fields=["routed_experts"]) + _assert_rows_equal(retrieved["routed_experts"], [torch.tensor([7, 8])]) + + def test_other_partitions_are_untouched(self, tq_system, dump_dir, controller): + """Restoring merges by key, unlike a checkpoint load which replaces everything.""" + # Define test data + _put_rows("d_target", ["a0", "a1"]) + _put_rows("d_bystander", ["b0", "b1"]) + + # Dump one partition, then restore it without wiping the other + tq.dump_data_by_key(dump_dir, ["a0"], "d_target") + ray.get(controller.clear_partition.remote("d_target")) + tq.load_data_by_key(dump_dir) + + # Check state: the bystander partition and its rows survived + assert sorted(ray.get(controller.list_partitions.remote())) == ["d_bystander", "d_target"] + retrieved = tq.kv_batch_get(keys=["b0", "b1"], partition_id="d_bystander", select_fields=["input_ids"]) + _assert_rows_equal(retrieved["input_ids"], [_row_input_ids(0), _row_input_ids(1)]) + + def test_row_without_produced_fields_keeps_its_key(self, tq_system, dump_dir, controller): + # Define test data: retrieve_meta with create=True registers a key with no fields + partition_id = "d_keyonly" + _put_rows(partition_id, ["p0"]) + client = tq.get_client() + client.kv_retrieve_meta(keys=["empty0"], partition_id=partition_id, create=True) + + # Dump both rows + report = tq.dump_data_by_key(dump_dir, ["p0", "empty0"], partition_id) + assert report["keys"] == 2 + assert report["rows_with_data"] == 1 + + # Wipe + load + ray.get(controller.clear_partition.remote(partition_id)) + tq.load_data_by_key(dump_dir) + + # Check restored state: the field-less row still exists + assert sorted(_keys_mapping(controller, partition_id)) == ["empty0", "p0"] + + def test_duplicate_keys_are_deduplicated(self, tq_system, dump_dir): + # Define test data + partition_id = "d_dupes" + _put_rows(partition_id, ["d0", "d1", "d2"]) + + # Dump with a repeated key + report = tq.dump_data_by_key(dump_dir, ["d1", "d1", "d2"], partition_id) + + # Check report and dump info + assert report["keys"] == 2 + with open(dump_dir / "dump_info.json", encoding="utf-8") as f: + assert json.load(f)["num_keys"] == 2 + + def test_empty_key_set_writes_a_readable_dump(self, tq_system, dump_dir, controller): + # Dump nothing + report = tq.dump_data_by_key(dump_dir, [], "d_empty") + + # Check saved state: still a complete, loadable dump + assert report["keys"] == 0 + assert report["rows_with_data"] == 0 + assert (dump_dir / "dump_info.json").exists() + + # Check that loading it is a no-op rather than an error + assert tq.load_data_by_key(dump_dir) == {} + + def test_live_partition_survives_a_dump(self, tq_system, dump_dir, controller): + # Define test data + partition_id = "d_nonmutating" + keys = [f"l{i}" for i in range(4)] + _put_rows(partition_id, keys) + before = _keys_mapping(controller, partition_id) + + # Dump a subset + tq.dump_data_by_key(dump_dir, ["l1"], partition_id) + + # Check live state: untouched rows keep their indexes and payloads + assert _keys_mapping(controller, partition_id) == before + retrieved = tq.kv_batch_get(keys=keys, partition_id=partition_id, select_fields=["input_ids"]) + _assert_rows_equal(retrieved["input_ids"], [_row_input_ids(row) for row in range(len(keys))]) + + def test_dump_refuses_an_existing_directory(self, tq_system, dump_dir): + partition_id = "d_existing" + _put_rows(partition_id, ["s0"]) + tq.dump_data_by_key(dump_dir, ["s0"], partition_id) + with pytest.raises(FileExistsError): + tq.dump_data_by_key(dump_dir, ["s0"], partition_id) + assert sorted(path.name for path in dump_dir.parent.iterdir()) == [dump_dir.name] + assert _read_row_index(dump_dir)["partition_id"] == partition_id + + def test_published_dump_loads_from_a_read_only_directory(self, tq_system, dump_dir): + partition_id = "d_read_only" + _put_rows(partition_id, ["s0"]) + tq.dump_data_by_key(dump_dir, ["s0"], partition_id) + tq.kv_clear(["s0"], partition_id) + dump_dir.parent.chmod(0o555) + try: + assert sorted(_read_row_index(dump_dir)["rows"]) == ["s0"] + tq.load_data_by_key(dump_dir) + assert sorted(path.name for path in dump_dir.parent.iterdir()) == [dump_dir.name] + finally: + dump_dir.parent.chmod(0o755) + _assert_rows_equal(tq.kv_batch_get(["s0"], partition_id, ["input_ids"])["input_ids"], [_row_input_ids(0)]) + + +# --------------------------------------------------------------------------- +# row index +# --------------------------------------------------------------------------- + + +class TestRowIndex: + def test_row_index_describes_keys_without_reading_payload(self, tq_system, dump_dir): + # Define test data + partition_id = "d_index" + keys = ["i0", "i1"] + _put_rows(partition_id, keys) + + # Dump + tq.dump_data_by_key(dump_dir, keys, partition_id) + + # Check the index + row_index = _read_row_index(dump_dir) + assert row_index["partition_id"] == partition_id + assert sorted(row_index["rows"]) == keys + for row, key in enumerate(keys): + assert row_index["rows"][key]["fields"] == ["attention_mask", "input_ids"] + assert row_index["rows"][key]["tag"] == {"idx": row} + + +# --------------------------------------------------------------------------- +# error handling +# --------------------------------------------------------------------------- + + +class TestDumpErrors: + def test_unknown_key_raises_and_leaves_no_directory(self, tq_system, dump_dir): + # Define test data + partition_id = "d_err_key" + _put_rows(partition_id, ["e0"]) + + # Dump a key that was never put + with pytest.raises(RuntimeError, match="keys not found"): + tq.dump_data_by_key(dump_dir, ["e0", "nope"], partition_id) + + # Check saved state: no partial directory left behind + assert not list(dump_dir.parent.iterdir()) + + def test_unknown_partition_raises(self, tq_system, dump_dir): + _put_rows("d_err_part", ["e0"]) + with pytest.raises(RuntimeError, match="does not exist"): + tq.dump_data_by_key(dump_dir, ["e0"], "d_never_created") + + def test_load_rejects_a_dump_without_info(self, tq_system, dump_dir): + dump_dir.mkdir(parents=True) + with pytest.raises(FileNotFoundError, match="dump_info.json"): + tq.load_data_by_key(dump_dir) + + def test_load_rejects_a_missing_shard(self, tq_system, dump_dir): + # Define test data + dump + partition_id = "d_err_shard" + _put_rows(partition_id, ["m0", "m1"]) + tq.dump_data_by_key(dump_dir, ["m0", "m1"], partition_id) + + # Tamper: delete one shard file + shard = next((dump_dir / "shards").glob("shard_*.pkl")) + shard.unlink() + + with pytest.raises(FileNotFoundError, match="Missing dump shard"): + tq.load_data_by_key(dump_dir) + + def test_load_rejects_an_unknown_format_version(self, tq_system, dump_dir): + # Define test data + dump + partition_id = "d_err_version" + _put_rows(partition_id, ["v0"]) + tq.dump_data_by_key(dump_dir, ["v0"], partition_id) + + # Tamper: bump the format version beyond what this build reads + info_path = dump_dir / "dump_info.json" + with open(info_path, encoding="utf-8") as f: + info = json.load(f) + info["format_version"] = tq.data_dump.DUMP_FORMAT_VERSION + 1 + with open(info_path, "w", encoding="utf-8") as f: + json.dump(info, f) + + with pytest.raises(ValueError, match="Unsupported dump format version"): + tq.load_data_by_key(dump_dir) + + +def test_direct_load_bypasses_caller_payload_io(tq_system, dump_dir, controller, monkeypatch): + partition = "direct_load" + keys = [f"key-{i}" for i in range(16)] + _put_rows(partition, keys) + tq.dump_data_by_key(dump_dir, keys, partition) + before = _keys_mapping(controller, partition) + tq.kv_put(keys[0], partition, fields=TensorDict({"extra": torch.tensor([[42]])}, batch_size=1), tag={"keep": True}) + _put_rows(partition, ["bystander"]) + client = tq.get_client() + manager = client.storage_manager + original_load = manager._load_selected_rows + responses = [] + + async def load(*args, **kwargs): + response = await original_load(*args, **kwargs) + responses.append((kwargs["target_storage_unit"], response)) + return response + + real_open = builtins.open + + def no_payload_open(path, *args, **kwargs): + if isinstance(path, str | Path) and Path(path).name.startswith("shard_") and str(path).endswith(".pkl"): + raise AssertionError("Caller opened a payload shard") + return real_open(path, *args, **kwargs) + + monkeypatch.setattr(builtins, "open", no_payload_open) + monkeypatch.setattr(manager, "_load_selected_rows", load) + monkeypatch.setattr(manager, "put_data", AsyncMock(side_effect=AssertionError("Caller sent payload through put"))) + tq.load_data_by_key(dump_dir) + assert len({unit for unit, _ in responses}) == _NUM_STORAGE_UNITS + shard_bytes = sum(path.stat().st_size for path in (dump_dir / "shards").glob("shard_*.pkl")) + assert sum(response["bytes_read"] for _, response in responses) == shard_bytes + after = _keys_mapping(controller, partition) + assert all(after[key] == before[key] for key in keys) + assert "bystander" in after + actual = tq.kv_batch_get(keys, partition, select_fields=["input_ids"]) + _assert_rows_equal(actual["input_ids"], [_row_input_ids(i) for i in range(16)]) + extra = tq.kv_batch_get([keys[0]], partition, select_fields=["extra"]) + _assert_rows_equal(extra["extra"], [torch.tensor([42])]) + snapshot = ray.get(controller.get_partition_snapshot.remote(partition)) + assert snapshot.custom_meta[after[keys[0]]]["keep"] + assert snapshot.custom_meta[after[keys[0]]]["idx"] == 0 + + +@pytest.mark.parametrize("row_count", [3, 127, 128, 129, 130]) +@pytest.mark.parametrize("last_kind", ["tensor", "none", "object"]) +def test_chunks_with_a_wrapped_last_row_roundtrip(tq_system, dump_dir, row_count, last_kind, monkeypatch): + partition = "wrapped_last_row" + keys = [f"k{i}" for i in range(row_count)] + tensors = [torch.arange(i % 3 + 1, dtype=torch.int64) for i in range(row_count)] + last = tensors[-1] if last_kind == "tensor" else None if last_kind == "none" else {"pixels": tensors[-1]} + expected = [*tensors[:-1], last] + chunk = dump_dir.parent / "legacy-chunk.pt" + chunk.parent.mkdir(parents=True) + first = TensorDict( + { + "input_ids": torch.nested.as_nested_tensor(tensors[:-1], layout=torch.jagged), + "multi_modal_inputs#images": torch.nested.as_nested_tensor(tensors[:-1], layout=torch.jagged), + "wrapped": NonTensorStack(*tensors[:-1]), + }, + batch_size=row_count - 1, + ) + tail = TensorDict( + { + name: NonTensorStack(value) + for name, value in { + "input_ids": last, + "multi_modal_inputs#images": last, + "wrapped": tensors[-1], + }.items() + }, + batch_size=1, + ) + # Legacy wrappers reload each saved TensorDict without normalizing its container type. + for chunk_keys, fields in [(keys[:-1], first), (keys[-1:], tail)]: + torch.save({"keys": chunk_keys, "fields": fields}, chunk) + saved = torch.load(chunk, weights_only=False) + tq.kv_batch_put(saved["keys"], partition, saved["fields"], tags=[{"key": key} for key in chunk_keys]) + + open_file = builtins.open + + def no_shard_read(path, *args, **kwargs): + if isinstance(path, str | Path) and Path(path).name.startswith("shard_") and str(path).endswith(".pkl"): + pytest.fail("The dump schema merge read a payload shard in the caller") + return open_file(path, *args, **kwargs) + + for target in [dump_dir, dump_dir.parent / "second-dump"]: + with monkeypatch.context() as patcher: + patcher.setattr(builtins, "open", no_shard_read) + tq.dump_data_by_key(target, keys, partition) + index = _read_row_index(target) + for name in ["input_ids", "multi_modal_inputs#images"]: + schema = index["field_schema"][name] + assert schema["is_non_tensor"] == (last_kind != "tensor") + if last_kind == "tensor": + assert schema["is_nested"] + for key, value in zip(keys, expected, strict=True): + assert tuple(schema["per_sample_shapes"][index["rows"][key]["global_index"]]) == tuple(value.shape) + assert index["field_schema"]["wrapped"]["is_non_tensor"] + # The controller still declares a tensor field, which the dump does not change. + if last_kind != "tensor" and target == dump_dir: + with pytest.raises(RuntimeError, match="tensor/non-tensor type mismatch"): + tq.load_data_by_key(target) + tq.kv_clear(keys, partition) + tq.load_data_by_key(target) + for start in range(0, row_count, 128): + end = min(start + 128, row_count) + restored = tq.kv_batch_get( + keys[start:end], partition, ["input_ids", "multi_modal_inputs#images", "wrapped"] + ) + for name in ["input_ids", "multi_modal_inputs#images"]: + for value, reference in zip(restored[name], expected[start:end], strict=True): + torch.testing.assert_close(value, reference) + for value, reference in zip(restored["wrapped"], tensors[start:end], strict=True): + torch.testing.assert_close(value, reference) + tags = tq.kv_list(partition)[partition] + assert tags == {key: {"key": key} for key in keys} + + +def test_corrupt_shard_keeps_new_rows_unproduced(tq_system, dump_dir, controller): + partition = "corrupt" + _put_rows(partition, ["key"]) + tq.dump_data_by_key(dump_dir, ["key"], partition) + path = next((dump_dir / "shards").glob("shard_*.pkl")) + path.write_bytes(b"!" * path.stat().st_size) + ray.get(controller.clear_partition.remote(partition)) + with pytest.raises(RuntimeError, match="failed to load rows"): + tq.load_data_by_key(dump_dir) + snapshot = ray.get(controller.get_partition_snapshot.remote(partition)) + assert "key" in snapshot.keys_mapping + assert not snapshot.field_metadata + + +def test_incompatible_schema_rejected_before_writes(tq_system, dump_dir, controller): + tq.kv_batch_put(["k"], "schema", TensorDict({"x": torch.tensor([[3]], dtype=torch.int64)}, batch_size=1)) + tq.dump_data_by_key(dump_dir, ["k"], "schema") + tq.get_client().clear_partition("schema") + tq.kv_batch_put(["k"], "schema", TensorDict({"x": torch.tensor([[1.5]])}, batch_size=1)) + with pytest.raises(RuntimeError, match="dtype mismatch"): + tq.load_data_by_key(dump_dir) + _assert_rows_equal(tq.kv_batch_get(["k"], "schema", ["x"])["x"], [torch.tensor([1.5])]) + + +@pytest.mark.parametrize("last", ["text", torch.tensor([3.0, 4.0, 5.0])], ids=["string", "ragged"]) +def test_dense_field_with_a_wrapped_later_row_roundtrips(tq_system, dump_dir, last): + partition = "dense_then_wrapped" + first = torch.tensor([[1.0, 2.0]]) + tq.kv_batch_put(["first"], partition, TensorDict({"x": first}, batch_size=1)) + tq.kv_batch_put(["last"], partition, TensorDict({"x": NonTensorStack(last)}, batch_size=1)) + tq.dump_data_by_key(dump_dir, ["first", "last"], partition) + tq.kv_clear(["first", "last"], partition) + tq.load_data_by_key(dump_dir) + restored = list(tq.kv_batch_get(["first", "last"], partition, ["x"])["x"]) + torch.testing.assert_close(restored[0], first[0]) + if isinstance(last, str): + assert restored[1] == last + else: + torch.testing.assert_close(restored[1], last) + + +def test_failed_load_publishes_no_metadata_and_retry_is_idempotent(tq_system, dump_dir, controller, monkeypatch): + _put_rows("retry", ["key"]) + tq.dump_data_by_key(dump_dir, ["key"], "retry") + client = tq.get_client() + client.clear_partition("retry") + manager = client.storage_manager + original_load = manager._load_selected_rows + + async def lose_reply(*args, **kwargs): + await original_load(*args, **kwargs) + raise RuntimeError("lost reply") + + with monkeypatch.context() as patcher: + patcher.setattr(manager, "_load_selected_rows", lose_reply) + with pytest.raises(RuntimeError, match="lost reply"): + tq.load_data_by_key(dump_dir) + snapshot = ray.get(controller.get_partition_snapshot.remote("retry")) + index = snapshot.keys_mapping["key"] + assert not snapshot.field_metadata + assert tq.kv_list("retry")["retry"] == {"key": {}} + + tq.load_data_by_key(dump_dir) + snapshot = ray.get(controller.get_partition_snapshot.remote("retry")) + assert snapshot.keys_mapping["key"] == index + _assert_rows_equal(tq.kv_batch_get(["key"], "retry", ["input_ids"])["input_ids"], [_row_input_ids(0)]) + assert tq.kv_list("retry")["retry"] == {"key": {"idx": 0}} + + +def test_load_reaches_every_owner_unit_concurrently(tq_system, monkeypatch): + manager = tq.get_client().storage_manager + units = list(manager.storage_unit_infos) + seen = [] + + async def load_all(): + started, ready = set(), asyncio.Event() + + async def load(shards, target_storage_unit): + started.add(target_storage_unit) + if len(started) == len(units): + ready.set() + # Every unit waits for all the others, so this finishes only if they run at once. + await asyncio.wait_for(ready.wait(), timeout=2) + for shard in shards: + for row in shard["records"]: + assert target_storage_unit == units[row["target_index"] % len(units)] + seen.append(row["source_index"]) + return {"updates": [], "bytes_read": 0} + + monkeypatch.setattr(manager, "_load_selected_rows", load) + records = [{"source_index": i, "target_index": 31 - i} for i in range(16)] + return await manager.load_rows_by_index("p", [{"path": "shard.pkl", "records": records}]) + + assert asyncio.run(load_all()) == 0 + assert sorted(seen) == list(range(16)) + + +def test_load_waits_for_other_units_before_raising(tq_system, monkeypatch): + manager = tq.get_client().storage_manager + units = list(manager.storage_unit_infos) + finished, notified = [], [] + + async def load_with_one_failure(): + failed = asyncio.Event() + + async def load(shards, target_storage_unit): + if target_storage_unit == units[0]: + failed.set() + raise RuntimeError("unit failed") + await failed.wait() + await asyncio.sleep(0) + finished.append(target_storage_unit) + return {"bytes_read": 0, "updates": [{"global_indexes": [1], "field_schema": {}}]} + + async def notify(*args): + notified.append(args) + + monkeypatch.setattr(manager, "_load_selected_rows", load) + monkeypatch.setattr(manager, "notify_data_update", notify) + await manager.load_rows_by_index( + "p", [{"path": "shard", "records": [{"target_index": 0}, {"target_index": 1}]}] + ) + + with pytest.raises(RuntimeError, match="unit failed"): + asyncio.run(load_with_one_failure()) + assert finished == [units[1]] + assert not notified + + +def test_load_publishes_metadata_updates_concurrently(tq_system, monkeypatch): + manager = tq.get_client().storage_manager + num_updates = 3 + + async def publish_all(): + started, ready = [], asyncio.Event() + + async def load(shards, target_storage_unit): + updates = [{"global_indexes": [i], "field_schema": {}} for i in range(num_updates)] + return {"bytes_read": 0, "updates": updates} + + async def notify(partition_id, global_indexes, field_schema): + started.append(global_indexes) + if len(started) == num_updates: + ready.set() + # Every update waits for all the others, so this finishes only if they run at once. + await asyncio.wait_for(ready.wait(), timeout=2) + + monkeypatch.setattr(manager, "_load_selected_rows", load) + monkeypatch.setattr(manager, "notify_data_update", notify) + await manager.load_rows_by_index("p", [{"path": "shard", "records": [{"target_index": 0}]}]) + + asyncio.run(publish_all()) + + +def test_dump_waits_for_writers_before_cleanup_can_start(tq_system, dump_dir, monkeypatch): + manager = tq.get_client().storage_manager + units = list(manager.storage_unit_infos) + completed = [] + + async def dump_with_one_failure(): + failed = asyncio.Event() + + async def dump(path, target_storage_unit, fields_by_index): + if target_storage_unit == units[0]: + failed.set() + raise OSError("write failed") + await failed.wait() + await asyncio.sleep(0) + completed.append(target_storage_unit) + return {"row_offsets": {1: [0, 1]}, "row_schema": {}} + + monkeypatch.setattr(manager, "_dump_single_shard", dump) + await manager.dump_rows_by_index(str(dump_dir), {0: ["x"], 1: ["x"]}) + + with pytest.raises(OSError, match="write failed"): + asyncio.run(dump_with_one_failure()) + assert completed == [units[1]] + + +def test_dump_and_load_do_not_use_the_put_get_timeout_pool(tq_system, dump_dir, monkeypatch): + _put_rows("own_pool", ["key"]) + manager = tq.get_client().storage_manager + # Any lease from the ordinary put/get pool now fails, so only the dump pool can serve these. + monkeypatch.setattr(manager, "storage_rpc_pool", None) + tq.dump_data_by_key(dump_dir, ["key"], "own_pool") + tq.load_data_by_key(dump_dir) + monkeypatch.undo() + _assert_rows_equal(tq.kv_batch_get(["key"], "own_pool", ["input_ids"])["input_ids"], [_row_input_ids(0)]) diff --git a/tests/test_data_dump.py b/tests/test_data_dump.py new file mode 100644 index 00000000..b6eecb55 --- /dev/null +++ b/tests/test_data_dump.py @@ -0,0 +1,284 @@ +# 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. + +"""Selective dump integrity and publication tests.""" + +import builtins +import io +import pickle +from types import SimpleNamespace + +import pytest +import torch + +from transfer_queue import data_dump, interface +from transfer_queue.storage.simple_storage import SimpleStorageUnit, StorageUnitData, validate_dump_values +from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType + + +@pytest.fixture +def unit(): + unit = SimpleStorageUnit.__ray_metadata__.modified_class({}) + yield unit + unit.shutdown() + + +@pytest.mark.parametrize("row_count", [1, 32]) +def test_dump_excludes_unselected_tensor_storage(unit, tmp_path, row_count): + batch = torch.arange(64 * 4096, dtype=torch.float32).reshape(64, 4096) + unit.storage_data.put_data({"x": batch}, list(range(64))) + path = tmp_path / "shard.pkl" + indexes = list(range(row_count)) + reply = unit._handle_dump_rows( + ZMQMessage.create( + request_type=ZMQRequestType.DUMP_ROWS, + sender_id="test", + body={"path": str(path), "fields_by_index": {index: ["x"] for index in indexes}}, + ) + ) + assert reply.body["success"] + assert set(reply.body["row_offsets"]) == set(indexes) + with path.open("rb") as f: + for index, (offset, length) in reply.body["row_offsets"].items(): + f.seek(offset) + row = pickle.loads(f.read(length)) + assert row["global_index"] == index + value = row["fields"]["x"] + torch.testing.assert_close(value, batch[index]) + assert value.untyped_storage().nbytes() == value.numel() * value.element_size() + assert path.stat().st_size < row_count * batch[0].numel() * batch.element_size() * 2 + assert unit.storage_data.field_data["x"][0].untyped_storage().nbytes() == batch.numel() * batch.element_size() + + +def test_row_index_compacts_tensors_inside_tags(monkeypatch, tmp_path): + batch = torch.arange(64 * 4096).reshape(64, 4096) + tag = {"nested": [SimpleNamespace(value=batch[0])]} + client = SimpleNamespace( + kv_describe=lambda *_: { + "rows": {"key": {"global_index": 0, "fields": [], "tag": tag}}, + "field_schema": {}, + } + ) + monkeypatch.setattr(interface, "_TQ_CONTROLLER", object()) + monkeypatch.setattr(interface, "_maybe_create_tq_client", lambda: client) + interface.dump_data_by_key(tmp_path / "dump", ["key"], "p") + restored = data_dump._read_row_index(tmp_path / "dump")["rows"]["key"]["tag"]["nested"][0].value + torch.testing.assert_close(restored, batch[0]) + assert restored.untyped_storage().nbytes() == restored.numel() * restored.element_size() + + +@pytest.fixture +def empty_dump_client(monkeypatch): + monkeypatch.setattr(interface, "_TQ_CONTROLLER", object()) + monkeypatch.setattr( + interface, "_maybe_create_tq_client", lambda: SimpleNamespace(validate_dump_schema=lambda *_: None) + ) + + +def test_failed_dump_leaves_no_staging_directory(empty_dump_client, monkeypatch, tmp_path): + def fail(path): + raise OSError("disk full") + + monkeypatch.setattr(data_dump, "_fsync_directory", fail) + with pytest.raises(OSError, match="disk full"): + interface.dump_data_by_key(tmp_path / "dump", [], "p") + assert not list(tmp_path.iterdir()) + + +def test_racing_dumps_to_one_path_publish_only_the_first(empty_dump_client, monkeypatch, tmp_path): + dump = tmp_path / "dump" + fsync = data_dump._fsync_directory + raced = [] + + def publish_another_dump_first(path): + if path.name.startswith("dump.tmp-") and not raced: + raced.append(path) + interface.dump_data_by_key(dump, [], "first") + fsync(path) + + monkeypatch.setattr(data_dump, "_fsync_directory", publish_another_dump_first) + with pytest.raises(OSError): + interface.dump_data_by_key(dump, [], "second") + assert data_dump._read_row_index(dump)["partition_id"] == "first" + assert [path.name for path in tmp_path.iterdir()] == ["dump"] + + +def test_unit_reads_only_assigned_ranges_and_merges(unit, tmp_path, monkeypatch): + batch = torch.arange(8 * 4096, dtype=torch.float32).reshape(8, 4096) + unit.storage_data.put_data({"x": batch, "stale": batch}, list(range(8))) + path = tmp_path / "shard.pkl" + reply = unit._handle_dump_rows( + ZMQMessage.create( + request_type=ZMQRequestType.DUMP_ROWS, + sender_id="test", + body={"path": str(path), "fields_by_index": {index: ["x"] for index in range(8)}}, + ) + ) + assert reply.body["success"] + content = path.read_bytes() + reads = [] + + class TrackedFile(io.BytesIO): + name = str(path) + + def read(self, size=-1): + reads.append((self.tell(), size)) + return super().read(size) + + open_file = builtins.open + monkeypatch.setattr( + builtins, + "open", + lambda name, *a, **kw: TrackedFile(content) if str(name) == str(path) else open_file(name, *a, **kw), + ) + unit.storage_data = StorageUnitData() + unit.storage_data.put_data({"keep": ["value"]}, [101]) + records = [] + for source, target in [(1, 101), (6, 106)]: + offset, length = reply.body["row_offsets"][source] + records.append( + {"source_index": source, "target_index": target, "fields": ["x"], "offset": offset, "length": length} + ) + loaded = unit._handle_load_rows( + ZMQMessage.create( + request_type=ZMQRequestType.LOAD_ROWS, + sender_id="test", + body={ + "shards": [ + { + "path": str(path), + "records": records, + "field_schema": { + "x": {"dtype": torch.float32, "shape": (4096,), "is_nested": False, "is_non_tensor": False} + }, + } + ] + }, + ) + ) + assert loaded.body["success"], loaded.body + assert reads == [(row["offset"], row["length"]) for row in records] + assert loaded.body["bytes_read"] == sum(row["length"] for row in records) + assert set(unit.storage_data.field_data) == {"x", "keep"} + assert unit.storage_data.field_data["keep"][101] == "value" + for row in records: + torch.testing.assert_close(unit.storage_data.field_data["x"][row["target_index"]], batch[row["source_index"]]) + assert sorted(index for update in loaded.body["updates"] for index in update["global_indexes"]) == [101, 106] + + +@pytest.mark.parametrize("problem", ["wrong_index", "truncated", "missing_field"]) +def test_unit_rejects_invalid_records(unit, tmp_path, problem): + path = tmp_path / "row.pkl" + path.write_bytes(pickle.dumps({"global_index": 1, "fields": {"x": "value"}})) + record = {"source_index": 1, "target_index": 2, "offset": 0, "length": path.stat().st_size, "fields": ["x"]} + if problem == "wrong_index": + record["source_index"] = 3 + elif problem == "truncated": + record["length"] += 1 + else: + record["fields"] = ["missing"] + reply = unit._handle_load_rows( + ZMQMessage.create( + request_type=ZMQRequestType.LOAD_ROWS, + sender_id="test", + body={"shards": [{"path": str(path), "records": [record], "field_schema": {"x": {"is_non_tensor": True}}}]}, + ) + ) + assert not reply.body["success"] + assert not unit.storage_data._active_keys + + +def test_load_refuses_other_backends_before_registering_keys(monkeypatch, tmp_path): + client = SimpleNamespace( + storage_manager=object(), + kv_retrieve_meta=lambda *_, **__: pytest.fail("registered keys on an unsupported backend"), + ) + monkeypatch.setattr(interface, "_TQ_CONTROLLER", object()) + monkeypatch.setattr(interface, "_maybe_create_tq_client", lambda: client) + with pytest.raises(NotImplementedError, match="does not support selective data load"): + interface.load_data_by_key(tmp_path) + + +def test_dump_reports_every_rows_stored_types(unit, tmp_path): + unit.storage_data.put_data({"x": [torch.arange(2), torch.arange(3)], "y": [None, {"key": "value"}]}, [9, 10]) + response = unit._handle_dump_rows( + ZMQMessage.create( + request_type=ZMQRequestType.DUMP_ROWS, + sender_id="test", + body={"path": str(tmp_path / "shard.pkl"), "fields_by_index": {9: ["x", "y"], 10: ["x", "y"]}}, + ) + ) + assert response.body["success"] + assert response.body["row_schema"] == { + 9: {"x": (torch.int64, (2,)), "y": None}, + 10: {"x": (torch.int64, (3,)), "y": None}, + } + + +_DENSE = {"is_nested": False, "is_non_tensor": False} +_NON_TENSOR = {"dtype": None, "shape": None, "is_nested": False, "is_non_tensor": True} + + +@pytest.mark.parametrize( + ("declared", "rows", "expected"), + [ + # A dense field whose later row a put wrapped as a string. + (_DENSE, [(torch.float32, (2,)), None], _NON_TENSOR), + # ... or as a tensor of another dtype. + (_DENSE, [(torch.float32, (2,)), (torch.int64, (2,))], _NON_TENSOR), + # ... or as a tensor of another shape. + ( + _DENSE, + [(torch.float32, (2,)), (torch.float32, (3,))], + { + "dtype": torch.float32, + "shape": None, + "is_nested": True, + "is_non_tensor": False, + "per_sample_shapes": {0: (2,), 1: (3,)}, + }, + ), + (_DENSE, [(torch.int64, ()), (torch.int64, ())], {**_DENSE, "dtype": torch.int64, "shape": (1,)}), + ( + {"is_nested": True, "is_non_tensor": False}, + [(torch.int64, (2,)), (torch.int64, (2,))], + { + "dtype": torch.int64, + "shape": None, + "is_nested": True, + "is_non_tensor": False, + "per_sample_shapes": {0: (2,), 1: (2,)}, + }, + ), + ({"is_nested": False, "is_non_tensor": True}, [(torch.int64, (2,)), (torch.int64, (2,))], _NON_TENSOR), + ], +) +def test_dump_schema_follows_the_stored_rows(declared, rows, expected): + row_schema = {index: {"x": meta} for index, meta in enumerate(rows)} + assert data_dump.dump_field_schema({"x": declared}, row_schema) == {"x": expected} + + +def test_saved_missing_tensor_shape_is_reported_as_invalid_dump(): + schema = { + "x": { + "dtype": torch.int64, + "shape": None, + "is_nested": True, + "is_non_tensor": False, + "per_sample_shapes": {1: None}, + } + } + with pytest.raises(ValueError, match="has no saved shape at row 1"): + validate_dump_values({"x": torch.arange(2)}, schema, 1) diff --git a/transfer_queue/__init__.py b/transfer_queue/__init__.py index 754bb4d8..da7a540c 100644 --- a/transfer_queue/__init__.py +++ b/transfer_queue/__init__.py @@ -25,6 +25,7 @@ async_kv_list, async_kv_put, close, + dump_data_by_key, get_client, get_metrics_endpoint, init, @@ -35,6 +36,7 @@ kv_list, kv_put, load_checkpoint, + load_data_by_key, save_checkpoint, ) from .metadata import BatchMeta, KVBatchMeta @@ -70,6 +72,11 @@ "save_checkpoint", "load_checkpoint", ] + + [ + # Selective Data Dump Interface + "dump_data_by_key", + "load_data_by_key", + ] + [ # High-Level StreamingDataLoader Interface "StreamingDataset", diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 0f7f13f1..05a361ab 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -1067,6 +1067,30 @@ async def async_kv_list( except Exception as e: raise RuntimeError(f"[{self.client_id}]: Error in kv_list: {str(e)}") from e + @with_controller_socket + async def async_kv_describe( + self, + keys: list[str], + partition_id: str, + socket: zmq.asyncio.Socket | None = None, + ) -> dict[str, Any]: + """Asynchronously describe each key's row and its fields' declared types, without payloads. + + Unlike ``kv_retrieve_meta``, every row keeps its own produced fields, and a missing + key raises instead of being dropped. + + Returns: + ``{"rows": {key: {"global_index", "fields", "tag"}}, "field_schema": {field: + {"is_nested", "is_non_tensor"}}}``. + """ + response = await self._request_controller( + socket=socket, + request_type=ZMQRequestType.KV_DESCRIBE, + response_type=ZMQRequestType.KV_DESCRIBE_RESPONSE, + body={"keys": keys, "partition_id": partition_id}, + ) + return {"rows": response.body["rows"], "field_schema": response.body["field_schema"]} + def close(self) -> None: """Close the client and cleanup resources including storage manager. @@ -1124,6 +1148,77 @@ def _can_destroy_zmq_context(self) -> bool: return False return True + # ==================== Selective Data Dump API ==================== + @with_controller_socket + async def async_validate_dump_schema( + self, + partition_id: str, + field_schema: dict, + socket: zmq.asyncio.Socket | None = None, + ) -> None: + """Reject incompatible destination fields before restoring payloads.""" + await self._request_controller( + socket=socket, + request_type=ZMQRequestType.VALIDATE_DUMP_SCHEMA, + response_type=ZMQRequestType.VALIDATE_DUMP_SCHEMA_RESPONSE, + body={"partition_id": partition_id, "field_schema": field_schema}, + ) + + async def async_dump_rows_by_index(self, shard_dir: str, fields_by_index: dict[int, list[str]]) -> dict[str, Any]: + """Asynchronously dump the given rows into per-storage-unit shards. + + Args: + shard_dir: Directory to write shard files into. + fields_by_index: Global index of each row to dump -> its produced fields. + + Returns: + ``{"shards", "row_schema"}``: one entry per written shard, and each row's + stored field types. + + Raises: + RuntimeError: If the storage manager is not initialized, or a unit holds + no data for a row it was asked to dump. + NotImplementedError: If the storage backend does not support dumping. + """ + 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 dump operations." + ) + return await self.storage_manager.dump_rows_by_index(shard_dir, fields_by_index) + + async def async_load_rows_by_key( + self, + partition_id: str, + rows: dict[str, dict[str, Any]], + shards: list[dict[str, Any]], + ) -> int: + """Restore selected payloads at the indexes their keys resolve to now; return bytes read. + + Failure semantics match ``kv_batch_put``: new keys stay registered and payload + writes may be partial, and retrying is idempotent because keys keep their indexes. + + Raises: + RuntimeError: If the storage manager is not initialized, or a unit fails. + NotImplementedError: If the storage backend does not support selective loads. + """ + 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 load operations." + ) + if not rows: + return 0 + metadata = await self.async_kv_retrieve_meta(list(rows), partition_id, create=True) + target_indexes = dict(zip(rows, metadata.global_indexes, strict=True)) + for shard in shards: + for record in shard["records"]: + record["target_index"] = target_indexes[record["key"]] + bytes_read = await self.storage_manager.load_rows_by_index(partition_id, shards) + metadata.update_custom_meta([row["tag"] for row in rows.values()]) + await self.async_set_custom_meta(metadata) + return bytes_read + # ==================== Checkpoint API ==================== @with_controller_socket async def async_save_controller_checkpoint( @@ -1308,6 +1403,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._kv_describe = _make_sync(self.async_kv_describe) + self._validate_dump_schema = _make_sync(self.async_validate_dump_schema) + self._dump_rows_by_index = _make_sync(self.async_dump_rows_by_index) + self._load_rows_by_key = _make_sync(self.async_load_rows_by_key) 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) @@ -1744,6 +1843,41 @@ def kv_list( return self._kv_list(partition_id=partition_id) + def kv_describe(self, keys: list[str], partition_id: str) -> dict[str, Any]: + """Synchronously describe each key's row and its fields' declared types, without payloads. + + See ``async_kv_describe``. + """ + return self._kv_describe(keys, partition_id) + + # ==================== Selective Data Dump API ==================== + + def validate_dump_schema(self, partition_id: str, field_schema: dict) -> None: + """Reject incompatible destination fields before restoring payloads.""" + return self._validate_dump_schema(partition_id, field_schema) + + def dump_rows_by_index(self, shard_dir: str, fields_by_index: dict[int, list[str]]) -> dict[str, Any]: + """Synchronously dump the given rows into per-storage-unit shards. + + Args: + shard_dir: Directory to write shard files into. + fields_by_index: Global index of each row to dump -> its produced fields. + + Returns: + ``{"shards", "row_schema"}``: one entry per written shard, and each row's + stored field types. + + Raises: + RuntimeError: If the storage manager is not initialized, or a unit holds + no data for a row it was asked to dump. + NotImplementedError: If the storage backend does not support dumping. + """ + return self._dump_rows_by_index(shard_dir, fields_by_index) + + def load_rows_by_key(self, partition_id: str, rows: dict, shards: list[dict]) -> int: + """Restore selected payloads at the indexes their keys resolve to now; return bytes read.""" + return self._load_rows_by_key(partition_id, rows, shards) + # ==================== Checkpoint API ==================== def save_controller_checkpoint(self, path: str) -> None: """Synchronously save controller state to a file via ZMQ RPC. diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index b45a85a6..c86c0824 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -557,6 +557,18 @@ def update_production_status( logger.error(f"Error updating production status for partition {self.partition_id}: {e}") return False + def validate_field_schema(self, field_schema: dict[str, dict[str, Any]]) -> None: + """Reject incompatible field types without changing metadata or readiness.""" + for name, incoming in field_schema.items(): + existing = self.field_metadata.get(name) + if existing is None: + continue + if bool(existing.is_non_tensor) != bool(incoming.get("is_non_tensor", False)): + raise ValueError(f"Field {name!r} tensor/non-tensor type mismatch") + dtype = incoming.get("dtype") + if dtype is not None and existing.dtype is not None and dtype != existing.dtype: + raise ValueError(f"Field {name!r} dtype mismatch: {existing.dtype} != {dtype}") + def _update_field_metadata( self, global_indexes: list[int], @@ -1725,6 +1737,54 @@ def kv_retrieve_keys( return keys + def kv_describe(self, keys: list[str], partition_id: str) -> dict[str, Any]: + """Describe each key's row and the declared types of its fields, without payloads. + + Unlike ``kv_retrieve_meta``, which keeps only the fields every row has produced and + drops missing keys, this reports each row's own field set and fails on a missing + key. Avoids ``to_snapshot``: a handful of lookups per key, not a partition copy. + + Args: + keys: Keys to describe, already deduplicated by the caller. + partition_id: Partition that owns ``keys``. + + Returns: + ``{"rows": {key: {"global_index", "fields", "tag"}}, "field_schema": {field: + {"is_nested", "is_non_tensor"}}}``. ``fields`` is empty for a row with no + produced field yet; ``field_schema`` covers the fields of the described rows. + + Raises: + KeyError: The partition or any key does not exist. + """ + partition = self._get_partition(partition_id) + if partition is None: + raise KeyError(f"partition {partition_id!r} does not exist; existing partitions: {sorted(self.partitions)}") + + global_indexes = partition.kv_retrieve_indexes(keys) + missing_keys = [key for key, index in zip(keys, global_indexes, strict=True) if index is None] + if missing_keys: + raise KeyError(f"keys not found in partition {partition_id!r}: {missing_keys}") + + rows: dict[str, dict[str, Any]] = { + key: { + "global_index": global_index, + "fields": sorted( + field_name + for field_name, field_meta in partition.field_metadata.items() + if global_index in field_meta.global_indexes + ), + "tag": partition.custom_meta.get(global_index, {}), + } + for key, global_index in zip(keys, cast(list[int], global_indexes), strict=True) + } + # Only the declared types: row dtypes and shapes live with the storage units. + field_schema = { + name: {"is_nested": bool(meta.is_nested), "is_non_tensor": bool(meta.is_non_tensor)} + for name, meta in partition.field_metadata.items() + if any(name in row["fields"] for row in rows.values()) + } + return {"rows": rows, "field_schema": field_schema} + def _init_zmq_socket(self): """Initialize ZMQ sockets for communication.""" self.zmq_context = zmq.Context() @@ -1948,6 +2008,8 @@ def _handle_request(self, request_msg: ZMQMessage, monitor: Any) -> ZMQMessage | ZMQRequestType.KV_RETRIEVE_META: self._handle_kv_retrieve_meta_request, ZMQRequestType.KV_RETRIEVE_KEYS: self._handle_kv_retrieve_keys_request, ZMQRequestType.KV_LIST: self._handle_kv_list_request, + ZMQRequestType.KV_DESCRIBE: self._handle_kv_describe_request, + ZMQRequestType.VALIDATE_DUMP_SCHEMA: self._handle_validate_dump_schema_request, ZMQRequestType.SAVE_CONTROLLER_CHECKPOINT: self._handle_save_controller_checkpoint_request, ZMQRequestType.LOAD_CONTROLLER_CHECKPOINT: self._handle_load_controller_checkpoint_request, } @@ -2145,6 +2207,11 @@ def _handle_kv_retrieve_keys_request(self, request_msg: ZMQMessage) -> ZMQMessag ) return self._make_response(request_msg, ZMQRequestType.KV_RETRIEVE_KEYS_RESPONSE, {"keys": keys}) + def _handle_kv_describe_request(self, request_msg: ZMQMessage) -> ZMQMessage: + params = request_msg.body + description = self.kv_describe(params["keys"], params["partition_id"]) + return self._make_response(request_msg, ZMQRequestType.KV_DESCRIBE_RESPONSE, description) + def _handle_kv_list_request(self, request_msg: ZMQMessage) -> ZMQMessage: requested_partition_id = request_msg.body["partition_id"] partition_ids = list(self.partitions.keys()) if requested_partition_id is None else [requested_partition_id] @@ -2181,6 +2248,13 @@ def _handle_load_controller_checkpoint_request(self, request_msg: ZMQMessage) -> {"success": True}, ) + def _handle_validate_dump_schema_request(self, request_msg: ZMQMessage) -> ZMQMessage: + params = request_msg.body + partition = self._get_partition(params["partition_id"]) + if partition is not None: + partition.validate_field_schema(params["field_schema"]) + return self._make_response(request_msg, ZMQRequestType.VALIDATE_DUMP_SCHEMA_RESPONSE, {"success": True}) + def get_zmq_server_info(self) -> ZMQServerInfo: """Get ZMQ server connection information.""" return self.zmq_server_info diff --git a/transfer_queue/data_dump.py b/transfer_queue/data_dump.py new file mode 100644 index 00000000..80fa53ca --- /dev/null +++ b/transfer_queue/data_dump.py @@ -0,0 +1,226 @@ +# 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. + +"""File format of selective dumps written by ``dump_data_by_key``. + +Field schemas are saved alongside independent row records in each shard. The manifest +maps source indexes to byte offsets, so current owner units read only their rows when +restoring into a different topology. + +Layout:: + + / + dump_info.json # version, partition, counts + row_index.pt # key -> {global_index, fields, tag} + shards/ + shard_info.json # file, unit, row count, source index -> [offset, length] + shard__.pkl # independent {global_index, fields} records +""" + +import json +import os +from collections import defaultdict +from pathlib import Path +from typing import Any + +import torch + +from transfer_queue.utils import compact_pickle +from transfer_queue.utils.logging_utils import get_logger + +logger = get_logger(__name__) + +DUMP_FORMAT_VERSION = 3 +SHARD_SUBDIR = "shards" + +_DUMP_INFO_FILE = "dump_info.json" +_ROW_INDEX_FILE = "row_index.pt" +_SHARD_INFO_FILE = "shard_info.json" + + +def _fsync_file(file_object: Any) -> None: + """Push a just-written file out of page cache before the caller moves on.""" + file_object.flush() + os.fsync(file_object.fileno()) + + +def _fsync_directory(path: Path) -> None: + """Make a directory's own entries durable. + + fsync on a file says nothing about the directory entry naming it, so a crash can + lose a file that was itself fully synced. The rename that publishes the dump needs + the same treatment. + """ + fd = os.open(path, os.O_RDONLY) + try: + os.fsync(fd) + finally: + os.close(fd) + + +def dump_field_schema(declared: dict, row_schema: dict[int, dict]) -> dict: + """Combine each field's declared type with the values its owner units hold. + + A put that wraps rows in ``NonTensorStack`` leaves a tensor field's metadata + unchanged, so dtypes and shapes come from the stored rows: a field is saved as + non-tensor unless every row is a tensor of one dtype, and nested if shapes differ. + """ + rows_by_field: defaultdict[str, dict[int, Any]] = defaultdict(dict) + for index, fields in row_schema.items(): + for name, meta in fields.items(): + rows_by_field[name][index] = meta + schema: dict[str, dict[str, Any]] = {} + for name, rows in rows_by_field.items(): + metas = list(rows.values()) + if declared[name]["is_non_tensor"] or None in metas or len({dtype for dtype, _ in metas}) > 1: + if not declared[name]["is_non_tensor"]: + logger.warning("Dump field %r holds rows that are not tensors of one dtype; saving as non-tensor", name) + schema[name] = {"dtype": None, "shape": None, "is_nested": False, "is_non_tensor": True} + continue + shapes = {index: tuple(shape) for index, (_, shape) in rows.items()} + nested = declared[name]["is_nested"] or len(set(shapes.values())) > 1 + schema[name] = { + "dtype": metas[0][0], + # A dense field of scalars is declared with shape (1,), as a put records it. + "shape": None if nested else shapes[next(iter(shapes))] or (1,), + "is_nested": nested, + "is_non_tensor": False, + } + if nested: + schema[name]["per_sample_shapes"] = shapes + return schema + + +def publish_dump(tmp_dir: Path, dump_dir: Path, row_index: dict[str, Any], shard_records: list[dict]) -> None: + """Write the manifests next to the unit-written shards, then rename into place. + + ``dump_info.json`` is written last and the staging directory is renamed only once + everything is durable, so a published dump is always complete. + """ + shard_dir = tmp_dir / SHARD_SUBDIR + shard_dir.mkdir(parents=True, exist_ok=True) + with open(shard_dir / _SHARD_INFO_FILE, "w", encoding="utf-8") as f: + json.dump(shard_records, f) + _fsync_file(f) + # The shards themselves were synced by the units that wrote them, but their + # directory entries were created here, on this node. + _fsync_directory(shard_dir) + + # torch.save rather than json: a tag is an arbitrary picklable dict, and this + # path must not fail on a tag that happens to hold a tensor. + with open(tmp_dir / _ROW_INDEX_FILE, "wb") as f: + torch.save(row_index, f, pickle_module=compact_pickle) + _fsync_file(f) + + rows = row_index["rows"] + with open(tmp_dir / _DUMP_INFO_FILE, "w", encoding="utf-8") as f: + json.dump( + { + "format_version": DUMP_FORMAT_VERSION, + "partition_id": row_index["partition_id"], + "num_keys": len(rows), + "num_rows_with_data": sum(1 for row in rows.values() if row["fields"]), + "num_shards": len(shard_records), + }, + f, + indent=2, + ) + _fsync_file(f) + # Everything the dump claims is now durable, so the staging directory can be + # published. Syncing the parent makes the rename itself survive a crash. + _fsync_directory(tmp_dir) + + # rename() refuses a non-empty target, so of two dumps racing to one path only + # the first is published. + tmp_dir.rename(dump_dir) + _fsync_directory(dump_dir.parent) + + +def _read_row_index(dump_dir: Path) -> dict[str, Any]: + row_index_path = dump_dir / _ROW_INDEX_FILE + if not row_index_path.exists(): + raise FileNotFoundError(f"{_ROW_INDEX_FILE} not found in {dump_dir}") + return torch.load(row_index_path, weights_only=False) + + +def read_dump(dump_dir: Path) -> tuple[dict[str, Any], list[dict[str, Any]]]: + """Validate a dump's manifests and return its row index and per-shard records. + + Only manifests are read; each record locates one row's byte range for the unit + that will load it. + + Raises: + FileNotFoundError: The dump is incomplete. + ValueError: The format version is unsupported, or a manifest disagrees with + the row index. + """ + info_path = dump_dir / _DUMP_INFO_FILE + if not info_path.exists(): + raise FileNotFoundError(f"{_DUMP_INFO_FILE} not found in {dump_dir}") + with open(info_path, encoding="utf-8") as f: + dump_info = json.load(f) + if dump_info["format_version"] != DUMP_FORMAT_VERSION: + raise ValueError( + f"Unsupported dump format version {dump_info['format_version']} in {dump_dir}; " + f"this build reads version {DUMP_FORMAT_VERSION}" + ) + + row_index = _read_row_index(dump_dir) + rows = row_index["rows"] + + shard_dir = dump_dir / SHARD_SUBDIR + with open(shard_dir / _SHARD_INFO_FILE, encoding="utf-8") as f: + shard_records = json.load(f) + if ( + len(rows) != dump_info["num_keys"] + or len(shard_records) != dump_info["num_shards"] + or row_index["partition_id"] != dump_info["partition_id"] + ): + raise ValueError("Dump manifest disagrees with the row index") + keys_by_index = {row["global_index"]: key for key, row in rows.items() if row["fields"]} + if len(keys_by_index) != dump_info["num_rows_with_data"]: + raise ValueError("Dump row count disagrees with the row index") + shards = [] + seen = set() + for record in shard_records: + path = shard_dir / record["file"] + if not path.is_file(): + raise FileNotFoundError(f"Missing dump shard: {path}") + records = [] + size = path.stat().st_size + offsets = record["row_offsets"] + if len(offsets) != record["rows"]: + raise ValueError(f"Dump shard row count mismatch: {path}") + for source_index, (offset, length) in offsets.items(): + source_index = int(source_index) + if source_index not in keys_by_index or source_index in seen: + raise ValueError(f"Unexpected or duplicated row {source_index} in {path}") + if offset < 0 or length <= 0 or offset + length > size: + raise ValueError(f"Invalid row range for {source_index} in {path}") + key = keys_by_index[source_index] + records.append( + { + "key": key, + "source_index": source_index, + "fields": rows[key]["fields"], + "offset": offset, + "length": length, + } + ) + seen.add(source_index) + shards.append({"path": str(path), "records": records, "field_schema": row_index["field_schema"]}) + if seen != set(keys_by_index): + raise ValueError("Dump shards do not contain every produced row") + return row_index, shards diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index eca8762e..95a8d683 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -36,6 +36,7 @@ from importlib import resources from pathlib import Path from typing import Any, Callable +from uuid import uuid4 import ray import torch @@ -43,6 +44,7 @@ from tensordict import TensorDict from tensordict.tensorclass import NonTensorStack +from transfer_queue import data_dump from transfer_queue.client import TransferQueueClient from transfer_queue.controller import TransferQueueController from transfer_queue.metadata import KVBatchMeta @@ -1141,3 +1143,117 @@ def load_checkpoint( client.load_controller_checkpoint(str(controller_path)) logger.info(f"Checkpoint loaded from {checkpoint_dir}") + + +# ==================== Selective Data Dump API ==================== + + +def dump_data_by_key(dump_dir: str | Path, keys: list[str], partition_id: str) -> dict[str, int]: + """Dump the rows addressed by ``keys`` into ``dump_dir``. + + Each storage unit pickles the rows it owns in its own process, so the payload never + passes through the caller. The caller writes only a small row index. + + The dump is staged in a uniquely named sibling directory and renamed into place + once durable. An existing ``dump_dir`` is refused, so a published dump is never + replaced: loads need no lock and readers no write access. + + .. note:: + **Multi-node limitation**: dump_dir must reside on a shared network filesystem + (e.g. NFS, GPFS, Lustre) reachable from every storage unit, because each unit + writes its own shard from its own node. + + Callers must freeze writers for ``keys`` for the duration: TransferQueue has no + atomic snapshot, so a concurrent put can land between the row index and the shards. + + Args: + dump_dir: Directory to write the dump into. + keys: Keys to dump. Duplicates are dropped, first occurrence wins. + partition_id: Partition that owns ``keys``. + + Returns: + ``{"keys", "rows_with_data", "shards", "bytes"}``. + + Raises: + FileExistsError: ``dump_dir`` already exists. + RuntimeError: TransferQueue is not initialized, the partition or a key does not + exist, or a storage unit holds no data for a row it was asked to dump. + NotImplementedError: The storage backend does not support selective dumps. + """ + if _TQ_CONTROLLER is None: + raise RuntimeError("TransferQueue is not initialized. Call tq.init() first.") + + dump_dir = Path(dump_dir).resolve() + if dump_dir.exists(): + raise FileExistsError(f"{dump_dir} already exists; write each dump to a new directory") + unique_keys = list(dict.fromkeys(keys)) + client = _maybe_create_tq_client() + # A unique staging name keeps concurrent dumps to the same path from sharing files. + tmp_dir = dump_dir.with_name(f"{dump_dir.name}.tmp-{uuid4().hex}") + tmp_dir.mkdir(parents=True) + + try: + description = client.kv_describe(unique_keys, partition_id) if unique_keys else {"rows": {}, "field_schema": {}} + row_index = {"partition_id": partition_id, **description} + rows = row_index["rows"] + # A row whose fields are all still unproduced has nothing for a storage unit to + # dump, but it keeps its key and tag so the restore can recreate the row. + fields_by_index = {row["global_index"]: row["fields"] for row in rows.values() if row["fields"]} + shard_records = [] + if fields_by_index: + dumped = client.dump_rows_by_index(str(tmp_dir / data_dump.SHARD_SUBDIR), fields_by_index) + shard_records = dumped["shards"] + row_index["field_schema"] = data_dump.dump_field_schema(row_index["field_schema"], dumped["row_schema"]) + data_dump.publish_dump(tmp_dir, dump_dir, row_index, shard_records) + except BaseException: + shutil.rmtree(tmp_dir, ignore_errors=True) + raise + + total_bytes = sum(path.stat().st_size for path in dump_dir.rglob("*") if path.is_file()) + logger.info(f"Dumped {len(unique_keys)} keys of partition {partition_id} to {dump_dir}") + return { + "keys": len(unique_keys), + "rows_with_data": len(fields_by_index), + "shards": len(shard_records), + "bytes": total_bytes, + } + + +def load_data_by_key(dump_dir: str | Path) -> dict[str, dict]: + """Merge selected rows into the running system, preserving existing key indexes. + + SimpleStorage units read their assigned indexed records directly and in parallel. + New keys receive current indexes; unrelated rows and fields remain untouched. + Writers and clears for these keys must be paused during restore. As with + ``kv_batch_put``, a failure may leave new keys registered and existing keys' fields + partially overwritten; retrying is idempotent. + + Args: + dump_dir: Directory previously written by ``dump_data_by_key``. It must be + accessible from every storage unit. + + Returns: + ``{key: tag}`` for every restored key. + + Raises: + RuntimeError: TransferQueue is not initialized or a storage unit fails. + NotImplementedError: The storage backend is not SimpleStorage. + FileNotFoundError: The dump is incomplete. + ValueError: The manifest or a row disagrees with the row index. + """ + if _TQ_CONTROLLER is None: + raise RuntimeError("TransferQueue is not initialized. Call tq.init() first.") + + client = _maybe_create_tq_client() + # Checked up front: the client registers the keys before any unit reads a shard. + if not isinstance(client.storage_manager, AsyncSimpleStorageManager): + raise NotImplementedError(f"{type(client.storage_manager).__name__} does not support selective data load") + + dump_dir = Path(dump_dir).resolve() + row_index, shards = data_dump.read_dump(dump_dir) + partition_id = row_index["partition_id"] + rows = row_index["rows"] + client.validate_dump_schema(partition_id, row_index["field_schema"]) + bytes_read = client.load_rows_by_key(partition_id, rows, shards) + logger.info(f"Restored {len(rows)} keys into partition {partition_id} from {dump_dir}, reading {bytes_read} bytes") + return {key: row["tag"] for key, row in rows.items()} diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index 369f800e..15d75627 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -373,6 +373,22 @@ async def load_checkpoint(self, checkpoint_dir: str) -> None: """ raise NotImplementedError(f"{self.__class__.__name__} does not support checkpoint") + async def dump_rows_by_index(self, shard_dir: str, fields_by_index: dict[int, list[str]]) -> dict[str, Any]: + """Have the owner units write the given rows into shards under shard_dir. + + Raises: + NotImplementedError: If this storage backend does not support selective dumps. + """ + raise NotImplementedError(f"{self.__class__.__name__} does not support selective data dump") + + async def load_rows_by_index(self, partition_id: str, shards: list[dict[str, Any]]) -> int: + """Have the current owner units read dumped rows at their target indexes. + + Raises: + NotImplementedError: If this storage backend does not support selective loads. + """ + raise NotImplementedError(f"{self.__class__.__name__} does not support selective data load") + 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 c89d1cfb..2d9a7a4c 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -58,6 +58,10 @@ # Timeout for the post-failure probe, which only has to answer whether the unit still serves. TQ_SIMPLE_STORAGE_PROBE_TIMEOUT = int(os.environ.get("TQ_SIMPLE_STORAGE_PROBE_TIMEOUT", 10)) +# A selective dump or load sends one request per unit covering all of its rows, so its duration +# grows with the selection and shared-filesystem speed rather than with a batch. +TQ_SIMPLE_STORAGE_DUMP_TIMEOUT = int(os.environ.get("TQ_SIMPLE_STORAGE_DUMP_TIMEOUT", 3600)) + class StorageUnitTimeout(RuntimeError): """A storage unit did not answer within the send/recv timeout. @@ -104,6 +108,12 @@ def _describe_unit_state(body: dict[str, Any]) -> str: resolve_target=lambda args, kwargs: kwargs.get("target_storage_unit"), ) +with_storage_unit_dump_socket = with_zmq_socket( + get_peer=lambda self, target: self.storage_unit_infos[target], + get_pool=lambda self: self.storage_dump_pool, + resolve_target=lambda args, kwargs: kwargs.get("target_storage_unit"), +) + class RoutingGroup(NamedTuple): """Routing result for a single storage unit.""" @@ -141,6 +151,12 @@ def __init__( timeout=TQ_SIMPLE_STORAGE_PROBE_TIMEOUT, maxsize=1, ) + self.storage_dump_pool = ZMQSocketPool( + self.zmq_context, + f"{self.storage_manager_id}_dump", + "put_get_socket", + timeout=TQ_SIMPLE_STORAGE_DUMP_TIMEOUT, + ) self.config = config server_infos: ZMQServerInfo | dict[str, ZMQServerInfo] | None = config.get("zmq_info", None) @@ -813,6 +829,175 @@ async def _load_single_storage_unit( f"[{self.storage_manager_id}]: Error restoring for storage unit {target_storage_unit}: {str(e)}" ) from e + @with_storage_unit_dump_socket + async def _dump_single_shard( + self, + path: str, + target_storage_unit: str, + fields_by_index: dict[int, list[str]], + socket: zmq.Socket = None, + ) -> dict[str, Any]: + """Ask one storage unit to write the rows it owns into a shard file.""" + try: + request_msg = ZMQMessage.create( + request_type=ZMQRequestType.DUMP_ROWS, # type: ignore[arg-type] + sender_id=self.storage_manager_id, + receiver_id=target_storage_unit, + body={"path": path, "fields_by_index": fields_by_index}, + ) + 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.DUMP_ROWS_RESPONSE or not response_msg.body.get("success"): + raise RuntimeError( + f"Storage unit {target_storage_unit} failed to dump rows to {path}: " + f"{response_msg.body.get('message', 'unknown error')}" + ) + missing_rows = response_msg.body["missing_rows"] + if missing_rows: + # The controller reported these rows as produced, so a unit that has no + # data for them means the two disagree. Never write a half table. + raise RuntimeError( + f"Storage unit {target_storage_unit} holds no data for requested rows: {missing_rows[:20]}" + ) + return response_msg.body + except Exception as e: + raise RuntimeError( + f"[{self.storage_manager_id}]: Error dumping shard from storage unit {target_storage_unit}: {str(e)}" + ) from e + + async def dump_rows_by_index( + self, + shard_dir: str, + fields_by_index: dict[int, list[str]], + ) -> dict[str, Any]: + """Dump the given rows into one shard per storage unit, in parallel. + + Each unit pickles its own rows in its own process, so the payload never passes + through the caller. A unit that owns none of the rows is skipped rather than + writing an empty shard. + + Args: + shard_dir: Directory to write shard files into. + fields_by_index: Global index of each row to dump -> its produced fields. + + Returns: + ``{"shards", "row_schema"}``: one ``{"file", "storage_unit_id", "rows", + "row_offsets"}`` entry per written shard, and each dumped row's + ``{field: (dtype, shape) or None}`` as the owner unit holds it. + + Raises: + RuntimeError: A unit holds no data for a row it was asked to dump. + """ + shard_dir_path = Path(shard_dir) + shard_dir_path.mkdir(parents=True, exist_ok=True) + + routing = self._group_by_hash(sorted(fields_by_index)) + targets = [(su_id, group.global_indexes) for su_id, group in routing.items()] + paths = [str(shard_dir_path / f"shard_{pos}_{su_id}.pkl") for pos, (su_id, _) in enumerate(targets)] + + results = await asyncio.gather( + *( + self._dump_single_shard( + path, + target_storage_unit=su_id, + fields_by_index={index: fields_by_index[index] for index in indexes}, + ) + for path, (su_id, indexes) in zip(paths, targets, strict=True) + ), + return_exceptions=True, + ) + shards = [] + row_schema = {} + total_rows = 0 + for path, (su_id, _), result in zip(paths, targets, results, strict=True): + if isinstance(result, BaseException): + raise result + offsets = result["row_offsets"] + total_rows += len(offsets) + row_schema.update(result["row_schema"]) + shards.append( + { + "file": Path(path).name, + "storage_unit_id": su_id, + "rows": len(offsets), + "row_offsets": offsets, + } + ) + + logger.info( + f"[{self.storage_manager_id}]: dumped {total_rows} rows across {len(targets)} shards to {shard_dir_path}" + ) + return {"shards": shards, "row_schema": row_schema} + + async def load_rows_by_index(self, partition_id: str, shards: list[dict[str, Any]]) -> int: + """Have current owner units read assigned byte ranges concurrently, then publish metadata. + + Metadata is published only after every unit succeeded, so writes to new keys stay + invisible after a failure, while fields existing keys already produced are + overwritten in place. Retrying rewrites the same target indexes. + + Returns: + Payload bytes read by the units. + """ + assignments = defaultdict(list) + for shard in shards: + rows = shard["records"] + for unit_id, group in self._group_by_hash([row["target_index"] for row in rows]).items(): + assignments[unit_id].append( + { + **shard, + "records": [rows[pos] for pos in group.batch_positions], + } + ) + results = await asyncio.gather( + *(self._load_selected_rows(shards, target_storage_unit=unit_id) for unit_id, shards in assignments.items()), + return_exceptions=True, + ) + loaded = [] + for result in results: + if isinstance(result, BaseException): + raise result + loaded.append(result) + # One update per (shard, unit, batch, field set): sequential round trips would + # dominate a wide cross-topology load. The notify pool bounds the concurrency. + await asyncio.gather( + *( + self.notify_data_update(partition_id, update["global_indexes"], update["field_schema"]) + for result in loaded + for update in result["updates"] + ) + ) + bytes_read = sum(result["bytes_read"] for result in loaded) + logger.info( + "[%s]: loaded %s bytes across %s units", + self.storage_manager_id, + bytes_read, + len(assignments), + ) + return bytes_read + + @with_storage_unit_dump_socket + async def _load_selected_rows( + self, + shards: list[dict[str, Any]], + target_storage_unit: str, + socket: zmq.Socket = None, + ) -> dict[str, Any]: + request = ZMQMessage.create( + request_type=ZMQRequestType.LOAD_ROWS, + sender_id=self.storage_manager_id, + receiver_id=target_storage_unit, + body={"shards": shards}, + ) + await socket.send_multipart(request.serialize(), copy=False) + response = ZMQMessage.deserialize(await socket.recv_multipart(copy=False)) + if response.request_type != ZMQRequestType.LOAD_ROWS_RESPONSE or not response.body.get("success"): + raise RuntimeError( + f"Storage unit {target_storage_unit} failed to load rows: {response.body.get('message')}" + ) + return response.body + async def save_checkpoint(self, checkpoint_dir: str) -> None: """Dump all storage units to the storage_units/ subdirectory of checkpoint_dir. @@ -878,4 +1063,7 @@ def close(self) -> None: probe_pool = getattr(self, "storage_probe_pool", None) if probe_pool is not None: probe_pool.close() + dump_pool = getattr(self, "storage_dump_pool", None) + if dump_pool is not None: + dump_pool.close() super().close() diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py index df28521e..258d9b7e 100644 --- a/transfer_queue/storage/simple_storage.py +++ b/transfer_queue/storage/simple_storage.py @@ -34,6 +34,7 @@ import shutil import time import weakref +from collections import defaultdict from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from pathlib import Path @@ -47,6 +48,7 @@ import torch import zmq +from transfer_queue.utils import compact_pickle from transfer_queue.utils.common import limit_pytorch_auto_parallel_threads, log_heavy_operation from transfer_queue.utils.enum_utils import Role from transfer_queue.utils.logging_utils import get_logger @@ -719,6 +721,53 @@ def close(self) -> None: self._ssd_store.close() +# Readers for the selective dump records that SimpleStorageUnit._handle_dump_rows writes. +def read_dump_row(file, offset: int, length: int, global_index: int, fields: list[str]) -> dict: + """Read and validate exactly one record against its expected index and fields.""" + file.seek(offset) + payload = file.read(length) + if len(payload) != length: + raise ValueError(f"Truncated dump row {global_index} in {file.name}") + row = pickle.loads(payload) + if row["global_index"] != global_index or set(row["fields"]) != set(fields): + raise ValueError(f"Dump row {global_index} disagrees with the row index in {file.name}") + return row["fields"] + + +def validate_dump_values(values: dict, schema: dict, source_index: int) -> None: + """Validate persisted tensor values without changing their type or dtype.""" + for name, value in values.items(): + field = schema[name] + if field["is_non_tensor"]: + continue + shape = field.get("per_sample_shapes", {}).get(source_index) if field["is_nested"] else field["shape"] + if shape is None: + raise ValueError(f"Dump field {name!r} has no saved shape at row {source_index}") + actual_shape = tuple(value.shape) if isinstance(value, torch.Tensor) else None + # Existing dense scalar fields use a one-element metadata shape. + scalar = actual_shape == () and tuple(shape) == (1,) and not field["is_nested"] + if ( + not isinstance(value, torch.Tensor) + or value.dtype != field["dtype"] + or (actual_shape != tuple(shape) and not scalar) + ): + raise ValueError(f"Dump field {name!r} disagrees with its saved schema at row {source_index}") + + +def select_dump_schema(schema: dict, source_indexes: list[int], target_indexes: list[int], names: tuple) -> dict: + """Remap only the selected nested shapes to the current destination indexes.""" + selected = {} + for name in names: + field = dict(schema[name]) + if field["is_nested"]: + field["per_sample_shapes"] = { + target: field["per_sample_shapes"][source] + for source, target in zip(source_indexes, target_indexes, strict=True) + } + selected[name] = field + return selected + + @ray.remote(num_cpus=1) class SimpleStorageUnit: """A storage unit that provides distributed data storage functionality. @@ -1003,6 +1052,12 @@ def _process_one_worker_request(self, worker_socket: zmq.Socket, monitor: Any) - 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.DUMP_ROWS: # type: ignore[arg-type] + with monitor.measure(op_type="DUMP_ROWS"): + response_msg = self._handle_dump_rows(request_msg) + elif operation == ZMQRequestType.LOAD_ROWS: # type: ignore[arg-type] + with monitor.measure(op_type="LOAD_ROWS"): + response_msg = self._handle_load_rows(request_msg) elif operation == ZMQRequestType.LOAD_STORAGE_CHECKPOINT: # type: ignore[arg-type] response_msg = self._handle_load_checkpoint(request_msg) else: @@ -1238,7 +1293,7 @@ def _handle_get_metrics(self) -> ZMQMessage: # Include per-operation stats if Prometheus metrics are enabled if self._metrics is not None: op_stats = {} - for op_type in ("PUT_DATA", "GET_DATA", "CLEAR_DATA"): + for op_type in ("PUT_DATA", "GET_DATA", "CLEAR_DATA", "DUMP_ROWS", "LOAD_ROWS"): try: hist = self._metrics.request_duration.labels(op_type=op_type) counter = self._metrics.request_total.labels(op_type=op_type) @@ -1293,6 +1348,100 @@ def _handle_save_checkpoint(self, data_parts) -> ZMQMessage: body={"success": False, "message": str(e)}, ) + def _handle_dump_rows(self, request: ZMQMessage) -> ZMQMessage: + """Write independent row records and return their offsets, never their payloads.""" + path = request.body["path"] + fields_by_index = request.body["fields_by_index"] + try: + missing = fields_by_index.keys() - self.storage_data._active_keys + if missing: + raise ValueError(f"Storage holds no data for requested rows: {sorted(missing)[:20]}") + row_offsets = {} + row_schema = {} + with open(path, "wb") as f: + for index in sorted(fields_by_index): + # Reused global indexes can retain fields no longer present in metadata. + fields = {name: self.storage_data.field_data[name][index] for name in fields_by_index[index]} + # Controller metadata can trail the stored values, so report what each + # row actually holds; only type metadata leaves the unit. + row_schema[index] = { + name: (value.dtype, tuple(value.shape)) if isinstance(value, torch.Tensor) else None + for name, value in fields.items() + } + offset = f.tell() + compact_pickle.dump({"global_index": index, "fields": fields}, f) + row_offsets[index] = [offset, f.tell() - offset] + f.flush() + os.fsync(f.fileno()) + logger.info("[%s]: dumped %s rows to %s", self.storage_unit_id, len(fields_by_index), path) + return ZMQMessage.create( + request_type=ZMQRequestType.DUMP_ROWS_RESPONSE, + sender_id=self.storage_unit_id, + body={ + "success": True, + "dumped_rows": len(fields_by_index), + "missing_rows": [], + "row_offsets": row_offsets, + "row_schema": row_schema, + }, + ) + except Exception as e: + logger.error("[%s]: dump rows failed: %s", self.storage_unit_id, e) + return ZMQMessage.create( + request_type=ZMQRequestType.DUMP_ROWS_RESPONSE, + sender_id=self.storage_unit_id, + body={"success": False, "message": str(e)}, + ) + + def _handle_load_rows(self, request: ZMQMessage) -> ZMQMessage: + """Seek to assigned records and merge them into local storage at current indexes.""" + updates = [] + bytes_read = 0 + try: + with limit_pytorch_auto_parallel_threads(TQ_NUM_THREADS): + for shard in request.body["shards"]: + records = sorted(shard["records"], key=lambda row: row["offset"]) + with open(shard["path"], "rb", buffering=0) as f: + # Bound temporary payload memory while retaining batched schema updates. + for start in range(0, len(records), 128): + groups = defaultdict(list) + for row in records[start : start + 128]: + fields = read_dump_row( + f, row["offset"], row["length"], row["source_index"], row["fields"] + ) + validate_dump_values(fields, shard["field_schema"], row["source_index"]) + groups[tuple(row["fields"])].append((row, fields)) + bytes_read += row["length"] + for signature, rows in groups.items(): + indexes = [row["target_index"] for row, _ in rows] + values = {name: [fields[name] for _, fields in rows] for name in signature} + schema = select_dump_schema( + shard["field_schema"], + [row["source_index"] for row, _ in rows], + indexes, + signature, + ) + self.storage_data.put_data(values, indexes) + updates.append({"global_indexes": indexes, "field_schema": schema}) + logger.info( + "[%s]: loaded %s rows (%s bytes)", + self.storage_unit_id, + sum(len(update["global_indexes"]) for update in updates), + bytes_read, + ) + return ZMQMessage.create( + request_type=ZMQRequestType.LOAD_ROWS_RESPONSE, + sender_id=self.storage_unit_id, + body={"success": True, "updates": updates, "bytes_read": bytes_read}, + ) + except Exception as e: + logger.error("[%s]: load rows failed: %s", self.storage_unit_id, e) + return ZMQMessage.create( + request_type=ZMQRequestType.LOAD_ROWS_RESPONSE, + 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 its checkpoint path. diff --git a/transfer_queue/utils/compact_pickle.py b/transfer_queue/utils/compact_pickle.py new file mode 100644 index 00000000..f83a5d60 --- /dev/null +++ b/transfer_queue/utils/compact_pickle.py @@ -0,0 +1,37 @@ +# 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. + +"""Pickle selected values without retaining the storage of unselected tensor rows.""" + +import pickle + +import torch + + +class Pickler(pickle.Pickler): + """Clone tensor leaves during pickling, including those inside arbitrary containers.""" + + def reducer_override(self, value): + """Serialize tensor values independently of their source storage.""" + # Tensor views can retain an entire batch, including unselected keys. This + # also handles tensors nested in picklable payload objects or metadata. + if isinstance(value, torch.Tensor): + return value.clone().__reduce_ex__(pickle.HIGHEST_PROTOCOL) + return NotImplemented + + +def dump(value, file): + """Write a standard pickle containing only the selected tensor values.""" + Pickler(file, protocol=pickle.HIGHEST_PROTOCOL).dump(value) diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index a94cd3ab..06c8d792 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -120,6 +120,8 @@ class ZMQRequestType(ExplicitEnum): KV_RETRIEVE_KEYS_RESPONSE = "KV_RETRIEVE_KEYS_RESPONSE" KV_LIST = "KV_LIST" KV_LIST_RESPONSE = "KV_LIST_RESPONSE" + KV_DESCRIBE = "KV_DESCRIBE" + KV_DESCRIBE_RESPONSE = "KV_DESCRIBE_RESPONSE" # METRICS GET_METRICS = "GET_METRICS" @@ -135,6 +137,14 @@ class ZMQRequestType(ExplicitEnum): LOAD_STORAGE_CHECKPOINT = "LOAD_STORAGE_CHECKPOINT" LOAD_STORAGE_CHECKPOINT_RESPONSE = "LOAD_STORAGE_CHECKPOINT_RESPONSE" + # SELECTIVE DATA DUMP + DUMP_ROWS = "DUMP_ROWS" + DUMP_ROWS_RESPONSE = "DUMP_ROWS_RESPONSE" + LOAD_ROWS = "LOAD_ROWS" + LOAD_ROWS_RESPONSE = "LOAD_ROWS_RESPONSE" + VALIDATE_DUMP_SCHEMA = "VALIDATE_DUMP_SCHEMA" + VALIDATE_DUMP_SCHEMA_RESPONSE = "VALIDATE_DUMP_SCHEMA_RESPONSE" + class ZMQServerInfo: """