[feat] selective saving and restoring of TQ samples by key instead of the whole

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 `<checkpoint_dir>.tmp`, renames the existing `checkpoint_dir` to `<checkpoint_dir>.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.
+**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.

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_<source-unit-id>.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 `<dump>.tmp-<id>` 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.

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)

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()