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

[OutstanderWang]

Summary

This PR adds a way to save a chosen set of keys from TransferQueue and later merge them back into a running deployment:

• tq.dump_data_by_key(dump_dir, keys, partition_id) saves the selected keys, all fields produced for them, and their tags.
• tq.load_data_by_key(dump_dir) merges those rows back into the running system.

The controller handles only metadata: which keys, indexes and storage units are involved, the field schemas, and restore reservations. All payload bytes are read and written by the storage units that own them, in parallel. Payloads never pass through the client.

In a real RL training job with partial (streaming) rollout, saving the rollout state at a training checkpoint went from 3.00 s to 0.90 s, about 3.3× faster. Loading was also faster even though the restored state was about 1.5× larger.

Motivation

The existing checkpoint saves and restores the whole system. save_checkpoint / load_checkpoint capture the entire controller (every partition, the global index manager and the sampler) plus each storage unit's full contents. On load, the storage-unit count must match and the system must be in a clean state. That is right for restarting a whole job. It is the wrong tool for "put these N samples back into the TransferQueue I am already running".

Partial rollout needs exactly that. With partial or streaming rollout, rollout samples that are unfinished or not yet consumed live only in TransferQueue when the trainer takes a checkpoint. If they are lost, resuming throws away generation work that was already paid for. So at every checkpoint the trainer must save exactly those keys next to its model checkpoint, and on resume it must put them back:

• only those keys, without dumping or restoring the whole store;
• into a TransferQueue that is already running, without disturbing unrelated rows, partitions or sampler state;
• possibly onto a different number of storage units than when they were saved.

The save also runs on the training critical path at every checkpoint, so it has to be cheap.

The current workaround is slow because all payloads go through one process. The workaround (called V1 below):

• Save: call kv_batch_get for the keys on the client and serialize the result to disk.
• Restore: read the file back and call kv_batch_put.

Every payload byte travels from the storage units to the client over ZMQ and is deserialized there. A single client process then serializes all of it to disk. Restore reverses this, so the payload crosses the client twice per direction. The cost grows with the buffer size and blocks the trainer while it runs.

This PR (V2) keeps the client on the metadata path only. Each owner storage unit writes its own rows during a dump and reads only its own rows during a restore, and all units work in parallel.

Usage

import transfer_queue as tq

# At checkpoint time: save only the in-flight rollout samples.
pending_keys = [...]  # keys of unfinished / unconsumed samples
report = tq.dump_data_by_key(f"{ckpt_dir}/tq_rollout", pending_keys, partition_id="train")
# report == {"keys": ..., "rows_with_data": ..., "shards": ..., "bytes": ...}

# On resume. The number of storage units may differ from when the dump was written.
tq.init()
index = tq.read_row_index(f"{ckpt_dir}/tq_rollout")  # keys, fields and tags; no payload is read
tq.load_data_by_key(f"{ckpt_dir}/tq_rollout")         # partition comes from the dump

# If a load was interrupted (tq.RestorePendingError), settle it before retrying or clearing:
tq.recover_data_load(f"{ckpt_dir}/tq_rollout")        # or cancel=True to abandon it

How load_data_by_key merges into the running system:

• Keys that already exist keep their current global index.
• New keys get fresh indexes from the current controller.
• Unrelated keys, partitions and tag entries are left untouched.
• Sampler and consumption state are not restored, by design.

dump_dir must be on a filesystem that every participating storage unit can reach (any local directory works on a single node).

What's here

Public API (transfer_queue/data_dump.py, exported from transfer_queue)

• dump_data_by_key, load_data_by_key, read_row_index, recover_data_load and RestorePendingError.
• Publishing a dump is crash-safe. The new dump is written to <dir>.tmp and synced. The previous dump is moved to <dir>.old and is deleted only after the new directory has been published and its parent directory synced. If a crash leaves the main directory missing, the next access recovers .old.
• All calls on the same dump path are serialized by an exclusive lock on a sibling .lock file, so this also works across processes. A sibling .restore marker stops a dump from being replaced while a load may still be reading it.

Controller (control plane only)

• New requests:
  • DESCRIBE_ROWS_BY_KEY returns each key's global index, produced fields, tag and original field schema, with no payload.
  • VALIDATE_DUMP_SCHEMA rejects destination type conflicts before any payload is written.
• Restore reservation lifecycle (BEGIN_RESTORE, RESTORE_UNIT claim/complete, FINISH_RESTORE, LIST_RESTORES). Each load has a unique ID. A storage unit must claim permission before writing.
• While a load is unresolved, the controller blocks these operations on the destination partition: clears, reusing its indexes, conflicting KV puts that would create keys, and full checkpoint save/load. A receive timeout never releases a writer that has already claimed permission.
• RestorePendingError explains why a load is still unresolved: per-unit states (pending / running), unit report failures, or an unknown ID after a restart. The controller remembers final outcomes, so retrying recovery after a lost commit reply is safe.

Storage manager and SimpleStorage units (data plane)

• Dump:
  • The manager groups the selected indexes by their current owner and calls the owner units concurrently.
  • Each unit writes one shard of independent per-row records and returns only byte offsets.
  • Tensors are cloned while pickling, so a row taken from a batch does not drag the rest of that batch's storage into the file.
• Load:
  • The manager routes records by the current indexes the controller assigned. Routing logic is shared with reservations in utils/storage_routing.py.
  • Each target unit reads only its assigned byte ranges, in batches of at most 128 rows, and merges them into local storage.
• The number of source and destination units can differ. Even a single-shard dump can restore across several units, because every record can be addressed on its own.
• DUMP_ROWS and LOAD_ROWS are added to the storage operation metrics.

Schema fidelity

• New dumps use format_version: 3, which saves the original field schemas. A restore reproduces those schemas regardless of target topology or batch boundaries; for example, a non-tensor field stays non-tensor.
• Older chunks can wrap tensor rows in NonTensorStack and leave nested row shapes missing. For those fields, the owner units recover the shape and dtype from the actual values while writing. If the values are mixed, the field is saved as non-tensor, with a warning.

Fallbacks

• Backends without direct selective load restore through kv_batch_put.
• Exporting a non-empty dump currently requires SimpleStorage.

Docs

• New docs/data_dump.md covers the format, distributed I/O, failure handling and recovery.
• docs/checkpoint.md gains a pointer explaining when to use a selective dump instead of a full checkpoint.

Performance

Setup: a real VLM RL training job with partial (streaming) rollout on 64 H20 GPUs, resumed twice in a row:

• from a V1 checkpoint at step 10, trained to step 20;
• then from the V2 checkpoint written by that run at step 20, trained to step 30.

V1 is the client-side path described in Motivation. V2 is dump_data_by_key / load_data_by_key.

	V1 (client-side KV get + save)	V2 (this PR)
Restored successfully, no degraded fallback	yes	yes
Restored keys missing required fields	0	0
First step after resume consumes restored samples	yes	yes
Rollout-state load time	1.77 s (0.91 GB)	1.37 s (1.36 GB)
Average rollout-state save time per checkpoint	3.00 s	0.90 s (≈3.3× faster)

• Save is about 3.3× faster. This is the cost paid on the training critical path at every checkpoint.
• Load is not slower. V2 restored about 1.5× more data in about 23% less time. The two buffers come from different steps, so throughput is the fairer comparison: about 0.51 GB/s for V1 versus about 0.99 GB/s for V2.
• Compatibility: the checkpoint written by the old client-side path still resumed correctly on the upgraded trainer. Training ran normally and the buffered samples were consumed. The V2 checkpoint written by that run then resumed correctly as well.

Correctness was also checked at a larger scale (hundreds of GiB, about 15k keys, 32 shards). After restore, a sample of keys matched in both content and tags, and dumps with missing fields or a missing manifest were rejected.

Testing

• python -m compileall -q transfer_queue tutorial tests: clean.
• Selective dump and restore tests, 142 passed:
  • tests/test_data_dump.py, tests/test_dump_lock.py, tests/test_restore_lifecycle.py
  • tests/e2e/test_data_dump_e2e.py, tests/e2e/test_data_dump_cross_topology_e2e.py, tests/e2e/test_restore_timeout_e2e.py
• Full suite, python -m pytest -q tests --ignore=tests/test_yuanrong_storage_client_e2e.py: 806 passed, 10 skipped.

What the new tests cover:

• Only the selected keys are restored.
• Tags, jagged fields and non-tensor fields survive the round trip, as do rows with no produced fields.
• Other partitions and bystander keys are untouched, and existing keys keep their index.
• Restore across storage-unit counts (4→2, 2→4, 3→3, 1→4).
• Rejection of a missing manifest, a missing or corrupt shard, an unknown format version, and incompatible schemas (before any writes).
• Lock and publication recovery.
• Recovery after a lost reply, and a running restore blocking clears and dumps until it is recovered.
• A claimed load surviving a receive timeout, cancelling an orphaned marker after a whole-system restart, and a cancelled delayed load being unable to overwrite a reused index.
• A test showing that the client never handles payloads during a direct load.

[ascend-robot]

CLA Signature Pass

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

[ascend-robot]

CLA Signature Pass

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

[ascend-robot]

CLA Signature Pass

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