[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, possibly onto a different number of storage units.

The controller handles only metadata. All payload bytes are written and read by the storage units that own them, in parallel, and never pass through the client. Load uses the ordinary put contract: kv_retrieve_meta(create=True) resolves target indexes, the owner units write the rows, then notify_data_update and set_custom_meta publish metadata.

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) calls kv_batch_get on the client and serializes the result to disk, then restores with kv_batch_put. Every payload byte travels from the storage units to the client and is serialized by a single process, and restore reverses this.

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 load, and all units work in parallel.

Usage

import transfer_queue as tq

# At checkpoint time: save only the in-flight rollout samples. Writers for these keys must be paused.
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()
tags = tq.load_data_by_key(f"{ckpt_dir}/tq_rollout")  # {key: tag}; the partition comes from the dump

How load_data_by_key merges into the running system:

• Target indexes come from the current controller by key; a dumped index only locates a record inside its shard file, so it can never conflict with live indexes.
• Keys that already exist keep their current index: dumped fields overwrite fields of the same name, other fields stay, and tags are merged.
• New keys get fresh indexes. Unrelated keys and partitions are left untouched.
• Sampler and consumption state are not restored, by design.
• Incompatible field types are rejected before any key is registered.

dump_dir must be on a filesystem that every storage unit can reach (any local directory works on a single node). Loading unpickles the dump, so only load directories you trust.

What's here

Public API: dump_data_by_key and load_data_by_key in transfer_queue/interface.py, exported from transfer_queue. The dump format lives in transfer_queue/data_dump.py.

Controller (control plane only)

• KV_DESCRIBE (TransferQueueClient.kv_describe(keys, partition_id)): each key's global index, its own produced fields and tag, plus the declared types of those fields. Unlike kv_retrieve_meta, rows keep their own field sets and a missing key raises instead of being dropped.
• VALIDATE_DUMP_SCHEMA rejects destination type conflicts before any payload is written.

Storage manager and SimpleStorage units (data plane)

• Dump: the manager groups the selected rows by owner unit and calls the units concurrently. Each unit writes one shard of independent per-row records and returns only byte offsets and each row's dtype and shape. Tensors are cloned while pickling (utils/compact_pickle.py), 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 target indexes the controller assigned. Each target unit reads only its assigned byte ranges, in batches of at most 128 rows, and merges them into local storage. Metadata updates are then published concurrently, after every unit has succeeded.
• DUMP_ROWS and LOAD_ROWS use their own socket pool with TQ_SIMPLE_STORAGE_DUMP_TIMEOUT (3600 s by default), because each unit receives one request covering all of its rows. Both are added to the storage operation metrics.
• Selective dump and load require SimpleStorage; other backends raise NotImplementedError before any key is registered.

Format (format_version: 3, the only version)

• dump_info.json, row_index.pt (keys, fields, tags, field schema) and shards/shard_info.json (each shard's file name, owner unit and per-row [offset, length]), plus the shard files.
• Field schemas come from the rows the units actually hold, because controller metadata can trail the stored values: a field is saved as non-tensor unless every row is a tensor of one dtype, and as nested if shapes differ.
• A dump is staged in a uniquely named sibling directory, synced, and renamed into place; an existing dump_dir is refused. Readers need no lock and no write access.

Failure behavior: the same as kv_batch_put, so load is not transactional. Writes to new keys stay invisible after a failure, but fields that existing keys have already produced are overwritten in place. Every key keeps its index, so retrying the same load is idempotent.

Docs: new docs/data_dump.md; docs/checkpoint.md points to it and states the trust boundary for both formats.

Performance

Real job (earlier revision). A VLM RL training job with partial (streaming) rollout on 64 H20 GPUs. V1 is the client-side path described in Motivation.

	V1 (client-side KV get + save)	V2 (this PR)
Average rollout-state save time per checkpoint	3.00 s	0.90 s (≈3.3× faster)

These save numbers come from an earlier revision; the dump data path is the same since (owner units write per-row shard records directly). The load path has changed since that run, so its old numbers were removed; load will be re-measured on the same job when the hardware is available.

Local, this head. Single node (18 cores, 48 GB), 8 storage units, rows of 96 KiB (8192-token int64 ids + fp32 log-probs) with a tag. Median of 3 runs; ranges cover two separate runs of the 1.5 GiB case. Every load was checked for payload and tags.

	0.38 GiB (4,096 rows)	1.50 GiB (16,384 rows)
V1 save (kv_batch_get + torch.save)	0.65 s	2.29 s
V2 save (dump_data_by_key)	0.16 s	0.61–0.67 s
V1 load (torch.load + kv_batch_put)	0.40 s	1.36–1.45 s
V2 load (load_data_by_key)	0.16 s	0.61–0.76 s
V1 load into 4 units	0.28 s	1.19–1.29 s
V2 load into 4 units	0.14 s	0.98–1.14 s

Notes on the local numbers:

• Files stay in the page cache and macOS fsync does not force a full flush, so these compare V1 and V2 on one machine; they are not a substitute for the multi-node run.
• With fewer target units, V2 load is bounded by per-unit deserialization: about 0.9 GiB/s per unit in isolation, almost all of it pickle.loads of the per-row tensors. That is why the gap narrows for 8 → 4 units at 1.5 GiB.
• Metadata publication is not a bottleneck: the 1.5 GiB load sends 128 updates, which take about 0.03 s. With 65,536 small rows (512 updates), publishing them concurrently instead of sequentially took load from 1.00 s to 0.88 s.

Testing

• python -m compileall -q transfer_queue tutorial tests: clean. pre-commit run --all-files (ruff, ruff-format, mypy): passed. Docstring and license sanity checks pass.
• python -m pytest -q tests --ignore=tests/test_yuanrong_storage_client_e2e.py: 729 passed, 10 skipped.
• Selective dump tests (65): tests/test_data_dump.py, tests/e2e/test_data_dump_e2e.py, tests/e2e/test_data_dump_cross_topology_e2e.py. They cover:
  • only the selected keys are restored; tags, jagged, non-tensor and heterogeneous field sets survive, 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 (3→2, 2→3, 1→3);
  • rejection of a missing manifest, a missing or corrupt shard, an unknown format version and incompatible schemas, before any write;
  • a failed load publishes no metadata for new keys, and retrying it is idempotent;
  • owner units are called concurrently, metadata updates are published concurrently, and dump and load do not use the put/get socket pool;
  • unit tests build a real storage unit, and manager tests run against the manager tq.init builds.

[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. 👍

[0oshowero0]

[AI Review]

Thanks for the detailed motivation and benchmarks. The core data-plane design is solid: independent per-row records with byte offsets let each current owner unit read only its own rows in parallel, which makes cross-topology restore straightforward, and compact_pickle keeps a row from dragging its whole batch's storage into a shard.

Checked on 2e66a78: python -m compileall -q transfer_queue tutorial tests is clean, and the 142 selective dump/restore tests listed in the description pass locally.

My main concern is the size relative to the feature. Most of the +4416 lines come from three things the feature does not need:

1. A custom restore reservation protocol (BEGIN/RESTORE_UNIT/FINISH/LIST/REPORT_RESTORE, per-unit states, the .restore marker, recover_data_load). While unresolved it blocks all KV writes and clears on the destination partition; it has no lease, so a lost unit requires restarting TQ; it makes storage units call the controller; and it makes the controller depend on SimpleStorage routing.
2. Readers for dump format versions 1 and 2, which exist only in this PR's own history (data_dump.py is new on main).
3. A FieldMeta.update / extract_field_schema change that affects every put, plus legacy shape recovery built on top of it.

Suggested split:

• A separate PR for the FieldMeta / NonTensorStack schema fix, with focused tests.
• The selective dump/load core with one format version, restoring through the existing metadata path: kv_retrieve_meta(create=True) for target indexes, LOAD_ROWS on owner units, then notify_data_update and set_custom_meta. Failure semantics match kv_batch_put (partial writes are possible; retrying the load is idempotent because existing keys keep their indexes), so recover_data_load and RestorePendingError are not needed.
• If late-write protection is still required, a follow-up that designs it together with the regular put path, which has the same late-write window today.

Details are in the inline comments. Two smaller notes without a single anchor:

• Tests build the controller and units with __new__ and set private fields by hand (for example tests/test_restore_lifecycle.py), which couples them tightly to internals.
• Shards and row_index.pt are read with pickle.loads and torch.load(weights_only=False). That matches the existing checkpoint trust model, but docs/data_dump.md should state that dump directories must be trusted.

[0oshowero0]

[AI Review] High: an unresolved restore blocks every KV write to the destination partition, not only conflicting key creation.

kv_put and kv_batch_put always call kv_retrieve_meta(..., create=True) (interface.py), so this check rejects all KV writes to the partition while any restore is unresolved, including writes to existing keys unrelated to the restore. The PR description and docs/data_dump.md say only "conflicting KV puts that would create keys" are blocked.

Minimal reproduction against the controller:

c.kv_retrieve_meta(["existing"], "train", create=True)
c.begin_restore("r1", "/dump", "train", {"restored": {"fields": ["x"], "tag": {}}}, ["u0"], {})
c.kv_retrieve_meta(["existing"], "train", create=True)
# RuntimeError: Restore r1 is unresolved; recover_data_load('/dump') first

In the resume flow, rollout producers that start writing to train before the load settles fail immediately, and after a RestorePendingError the partition stays closed to KV writers until someone runs recover_data_load.

[OutstanderWang]

Fixed by removing the reservation protocol. Loads now go through the ordinary kv_retrieve_meta path and no longer block other writers.

[0oshowero0]

[AI Review] High: the reservation has no lease, so a lost storage unit makes the partition unusable until the whole TQ system restarts.

Cancellation only moves pending units to cancelled. A unit that claimed permission and then died stays running forever, so recover_data_load(..., cancel=True) keeps raising RestorePendingError; docs/data_dump.md confirms a full restart is required. Until then:

• every KV write and clear on the destination partition is rejected (see the comment on kv_retrieve_meta);
• save_checkpoint and load_checkpoint are rejected for the whole system, because _assert_not_restoring() without a partition matches any restore.

The existing kv_batch_put has no such failure mode: a dead unit fails the call and the caller decides what to do. A convenience restore path should not be able to turn one unit failure into system-wide checkpoint unavailability.

[OutstanderWang]

Removed together with the protocol. A failed unit now fails the load call, as with kv_batch_put, and nothing stays reserved.

[0oshowero0]

[AI Review] High: large loads fail by default on slower shared filesystems.

LOAD_ROWS goes through storage_rpc_pool, whose timeout is TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT (200 s by default). A unit reading ~10 GiB from NFS can exceed that, and load_data_by_key then raises RestorePendingError even though every unit is healthy. The caller has to loop over recover_data_load, and REPORT_RESTORE is queued behind the still-running LOAD_ROWS on the unit's single worker thread, so recovery also reports errors until the load finishes. DUMP_ROWS uses the same 200 s timeout on the training critical path.

If the reservation protocol stays, load_data_by_key should wait for the restore to settle itself (for example by polling FINISH_RESTORE) instead of surfacing the ordinary socket timeout, so that recover_data_load is only needed after the initiating process crashes.

[OutstanderWang]

Dump and load requests now use their own socket pool, with timeout TQ_SIMPLE_STORAGE_DUMP_TIMEOUT (3600 s by default), because each unit receives one request covering all of its rows. A timeout no longer leaves anything to recover; the load can simply be retried.

[0oshowero0]

[AI Review] High: dumps in read-only locations cannot be read or loaded.

_dump_lock creates the parent directory and opens a sibling <dump>.lock in a+b mode, and load_data_by_key also creates <dump>.restore. Resuming from a checkpoint directory without write permission therefore fails before reading anything. Reproduced with a read-only parent directory:

read_row_index(dump) -> PermissionError: [Errno 13] Permission denied: '.../dump.lock'

Read paths also call _recover_dump, which can rename .old back into place, so a read can mutate the filesystem. Reads should not need write access; see also the comment on publication below.

[OutstanderWang]

Fixed. Readers take no lock and create no marker, and a new e2e test loads a dump from a read-only parent directory.

[0oshowero0]

[AI Review] High (design): storage units now make synchronous RPCs to the controller.

Before this PR, units only served requests; only clients and storage managers talked to the controller. Here a unit's single worker thread blocks on a claim RPC and a completion RPC to the controller, with a hard-coded 10 s timeout and a controller address taken from the request body. Together with _restore_results and the claim states on the controller, this is the bulk of the PR's complexity, and it crosses the control-plane/data-plane boundary.

The protection it buys (a late LOAD_ROWS cannot overwrite an index that was cleared and reused after a timeout) applies equally to ordinary put / kv_batch_put, which have the same late-write window and no such protocol. A simpler restore that keeps the existing contract:

1. kv_retrieve_meta(keys, partition_id, create=True) to resolve or allocate target indexes;
2. LOAD_ROWS on the owner units (keep this new op), returning the field schema of what they wrote, as _load_rows already does via updates;
3. the storage manager's existing notify_data_update, then set_custom_meta for tags.

Failure semantics then match kv_batch_put: writes may be partial, and retrying the load is idempotent because existing keys keep their indexes. That removes BEGIN/RESTORE_UNIT/FINISH/LIST/REPORT_RESTORE, _restores, _restore_outcomes, _restore_results, _clearing_indexes, the .restore marker, RestorePendingError and recover_data_load, along with most of tests/test_restore_lifecycle.py and tests/e2e/test_restore_timeout_e2e.py. Stronger guarantees, if needed, could be designed in a follow-up together with the regular put path.

[OutstanderWang]

Done as you suggested. Units no longer talk to the controller, and the storage manager publishes metadata through notify_data_update only after every unit succeeded. Late-write fencing is left for a follow-up designed together with the regular put path.

[0oshowero0]

[AI Review] Low: placement of the public API. These top-level functions import private state from interface.py, while every other public API lives in interface.py. Consider defining the public entry points there and keeping the file-format helpers in this module.

[OutstanderWang]

dump_data_by_key and load_data_by_key now live in interface.py next to the checkpoint API. data_dump.py keeps only the file format.

[0oshowero0]

[AI Review] Low: every method that takes _restore_lock runs on the controller's single request thread (none of them is called as a Ray actor method), so the lock protects nothing. It mostly adds re-indentation to mark_clearing, clear_partition, clear_meta, kv_retrieve_meta and save/load_checkpoint. Separately, _restore_outcomes here and _restore_results on the units grow without bound.

[OutstanderWang]

Gone with the protocol.

[0oshowero0]

[AI Review] Low: an index stays in _clearing_indexes if the client fails or exits between MARK_CLEARING and CLEAR_META (for example a CLEAR_META timeout, or a backend whose clear_data raises). begin_restore then rejects the partition with "unfinished clear" until the whole partition is cleared.

[OutstanderWang]

Gone with the protocol.

[0oshowero0]

[AI Review] Low: if a later update_production_status returns False, the earlier updates are already applied. The client then cancels, and the restore is recorded as not committed while part of its metadata is already published.

[OutstanderWang]

Gone with the protocol. Metadata is now published per unit through notify_data_update, as an ordinary put does, and only after every unit has succeeded.

[0oshowero0]

[AI Review] Low: new keys are registered here, before any payload is written, and a cancelled or failed restore does not unregister them, so they show up in kv_list without fields or tags. The docs only say payload writes are not rolled back; this should be documented too, or the keys released on cancel.

[OutstanderWang]

Documented: those keys stay registered, kv_list shows them with empty tags and no readable fields, and a retry reuses their indexes and restores the tags. The e2e failure test asserts this.

[0oshowero0]

What will happen if the indexes of dumped keys conflict to the current TQ indexes?

[OutstanderWang]

Thanks for the thorough review. I took the simpler restore you suggested and addressed the rest as separate commits on top of the reviewed head:

1. refactor: restore selected rows through the ordinary put contract: removes the reservation protocol. Load is now kv_retrieve_meta(create=True), then LOAD_ROWS on the current owner units, then notify_data_update and set_custom_meta.
2. fix: give selective dump and load their own long-timeout socket pool
3. fix: publish dumps once and read them without locks
4. refactor: ship a single selective dump format
5. fix: take dumped field types from the stored rows: the FieldMeta and extract_field_schema changes are gone from this PR.
6. refactor: narrow the selective dump API
7. docs: describe keys left registered by a failed load
8. docs, test: state the pickle trust boundary; build real objects in dump tests

The PR is now +2325/-2 across 17 files, down from +4416/-232 across 23.

On indexes in a dump conflicting with live indexes: they cannot conflict, because load never reuses a dumped index as a target. Target indexes come from kv_retrieve_meta(create=True) by key. A key that already exists keeps its index, the dumped fields overwrite its fields of the same name, its other fields stay, and its tag is merged. New keys get fresh indexes from the current controller. The dumped index only locates a record inside its shard file. For example, if a/b/c are dumped at indexes 0/1/2, and later another partition holds 0-3 and a is re-put at 4, then the load keeps a at 4, gives b/c new indexes, and leaves the other partition untouched. Incompatible field types are rejected before any key is registered.

Two points from the review summary:

• Tests building objects with __new__: test_restore_lifecycle.py is removed. The dump unit tests construct a real storage unit, and the three manager tests now run against the manager that tq.init builds.
• Pickle trust: docs/data_dump.md and docs/checkpoint.md now say that loading unpickles the files and that only trusted directories should be loaded.

[0oshowero0]

[AI Review] Re-review of 14c1804

Thanks for the quick and thorough rework. I checked every earlier finding against the code rather than the replies, and all of them are addressed: load now follows the ordinary put contract (kv_retrieve_meta(create=True) -> LOAD_ROWS -> notify_data_update / set_custom_meta), units no longer call the controller, readers need no lock or write access, only one format remains, and FieldMeta / extract_field_schema are identical to main.

Checked: python -m compileall -q transfer_queue tutorial tests is clean, and python -m pytest -q tests --ignore=tests/test_yuanrong_storage_client_e2e.py gives 733 passed, 10 skipped. The only commit on main since the PR base touches the mooncake client, so there is no interaction.

On the index-conflict question above: the code matches your answer. Target indexes come from kv_retrieve_meta by key, the dumped index only locates a record in its shard, select_dump_schema remaps nested shapes to the target index, and existing tags are merged per key by BatchMeta.update_custom_meta.

Remaining items:

1. The PR description is out of date. The usage example and "What's here" still describe read_row_index, recover_data_load, RestorePendingError, the BEGIN_RESTORE lifecycle and utils/storage_routing.py. The load numbers in the performance table were measured on the previous implementation, and the load path has changed since (see the inline comment on notify_data_update). Since the description usually becomes the merge commit message, please update it and re-measure load.
2. Three small inline comments: one doc correction with a reproduction, one performance note, and one unused code path.

[0oshowero0]

[AI Review] This does not hold for keys that already exist with produced fields.

Those fields are already marked produced in controller metadata, so a unit overwriting them makes the restored values readable immediately, and they stay after the load fails. Reproduced on 14c1804:

put(1)                                        # kv_batch_put "k": x = [1, 1, 1]
tq.dump_data_by_key(dump_dir, ["k"], "p")
put(9)                                        # "k": x = [9, 9, 9]
manager.notify_data_update = failing_notify   # raises RuntimeError("controller ACK lost")
tq.load_data_by_key(dump_dir)                 # raises
tq.kv_batch_get(["k"], "p", ["x"])            # [[1, 1, 1]], not [[9, 9, 9]]

The behavior matches kv_batch_put, so only the wording needs fixing, for example: "writes for new keys stay invisible; fields of existing keys are overwritten in place". Otherwise a reader may conclude that a failed load left existing keys untouched and skip the retry. The same sentence is in the load_rows_by_index docstring in simple_storage_manager.py.

[OutstanderWang]

Thanks for the reproduction, confirmed on 14c1804: the existing key read back [1, 1, 1] after the failed load. Fixed in 5160e32. docs/data_dump.md, the load_rows_by_index docstring and the load_data_by_key docstring now say that writes to new keys stay invisible after a failure, while fields an existing key has already produced are overwritten in place and readable at once, so a failed load has to be retried or its keys cleared.

[0oshowero0]

[AI Review] Low: metadata is published with one sequential round trip per update.

Units return one update per shard, 128-row batch and field set, so a cross-topology load produces at least one update per (source shard, target unit) pair, for example about 1,000 sequential notify_data_update calls for a 32 -> 32 restore. The previous implementation committed in one step, so this may change the load numbers in the description. Worth checking when re-measuring; if it shows up, issuing the calls concurrently or merging updates that share a field set would bring it down.

[OutstanderWang]

Agreed. 41a30d7 issues the updates with asyncio.gather. The notify socket pool still bounds how many are in flight, so this does not open a socket per update. A new test only passes if the updates run concurrently.

I kept one update per (shard, unit, batch, field set) rather than merging them, because this phase no longer shows up in local measurements (single node, so a lower bound on the multi-node cost):

• 8 → 8 units, 16,384 rows (1.5 GiB): the 128 updates take about 0.03 s of a 0.35 s load.
• 65,536 small rows (512 updates): load went from 1.00 s with sequential updates to 0.88 s with concurrent ones.

If it shows up on a multi-node run, merging updates that share a field set is the next step.

[0oshowero0]

[AI Review] Low: the "raw dump" mode has no caller. dump_data_by_key always passes fields_by_index; only a mocked test in test_data_dump_e2e.py omits it. This branch and the optional fields_by_index parameter in StorageManager.dump_rows_by_index, AsyncSimpleStorageManager.dump_rows_by_index / _dump_single_shard and the client's dump_rows_by_index can be removed, making the field list required.

[OutstanderWang]

Removed in eb51033. fields_by_index is now required through the storage unit, both managers and the client. Once it was required, global_indexes only repeated its keys, so I dropped that parameter too. The unit tests that relied on raw mode now pass explicit fields.

[0oshowero0]

describe_data_dump isn't really dump-specific: it is a per-key metadata lookup (key → global_index, the fields each row has produced, and its tag) plus the declared types of those fields. Only the client-side name ties it to dump; the controller method is already the generic describe_rows_by_key.

Suggest naming it for what it does and putting it in the KV family, e.g. kv_describe(keys, partition_id) with a KV_DESCRIBE request type, and moving the controller method out of the "Selective Data Dump API" section. That also fixes the argument order, which is currently (partition_id, keys) while the other kv_* APIs take (keys, partition_id).

I don't think it should be folded into kv_retrieve_meta, though: BatchMeta is columnar, so it only keeps the fields every requested row has produced, and it silently drops missing keys. Per-row field sets and a hard error on missing keys are exactly what this lookup adds. I'd also keep its scope as-is (no extra dtype/shape/consumption options) until there is a second caller.

[OutstanderWang]

Done in 48b9227: TransferQueueClient.kv_describe(keys, partition_id) / async_kv_describe with a KV_DESCRIBE request type. The controller's kv_describe now sits with the other KV methods and also returns the declared field types, which the handler used to compute. As you suggested, the scope is unchanged and it is not folded into kv_retrieve_meta. I have also not exported it as a top-level tq. function for now.

[0oshowero0]

dump_io.py has a single caller, SimpleStorageUnit._handle_load_rows, and the code that writes these same records (the {"global_index", "fields"} layout and the [offset, length] bookkeeping) is already inline in _handle_dump_rows. Suggest moving these three functions into simple_storage.py next to the two handlers and deleting this module, so a shard's record format is written and read in one place. That also matches the existing SimpleStorage checkpoint, whose format already lives in simple_storage.py.

[OutstanderWang]

Done in 8035fb8. The three readers now live in simple_storage.py next to _handle_dump_rows, and dump_io.py is deleted.

[0oshowero0]

The shard file name is built twice: here and in AsyncSimpleStorageManager.dump_rows_by_index (shard_{pos}_{su_id}.pkl), so the two have to be kept in sync by hand. Suggest having the manager record the file name in each shard_info entry and reading it here instead of rebuilding it.

[OutstanderWang]

Done in be2e50e. The manager records file in each shard_info.json entry and read_dump uses it. position only existed to build the name, so it is gone. The format stays at version 3 because it has not shipped.

[OutstanderWang]

Thanks for the re-review. Each of the six inline comments is addressed in its own commit and answered in its thread:

• 5160e32 docs: a failed load can overwrite fields that existing keys already have
• 41a30d7 perf: selective load publishes its metadata updates concurrently
• eb51033 refactor: the dumped field list is required, and the raw dump mode is gone
• 48b9227 refactor: the per-key row lookup is now the KV request kv_describe
• 8035fb8 refactor: dump record reads sit next to the writer in simple_storage.py, and dump_io.py is deleted
• be2e50e refactor: the manifest records each shard's file name

Follow-up commits:

• 4136c3b adds the type annotations mypy asked for.
• 264778e and fdd8335 fix "Run Tests", which had been cancelled at the 15-minute limit on every run of this PR. It was a hang, not slowness. The controller and each storage unit reserve one CPU, and the GitHub runner has four, so the 4-unit topologies in the dump e2e tests could never be scheduled. Those tests now use at most three units: the shared system has 3 units, and the cross-topology cases are 3→2, 2→3 and 1→3.
• cf9e869 drops cross-topology cases that other tests already cover. The 3→3 case is removed. The non-tensor schema test used to run with 2, 127, 128, 129 and 130 rows; it now runs only with 130, which still spans two 128-row load batches, and test_data_dump_e2e.py covers each batch boundary on a single system.

I also rewrote the PR description for the current design.

Performance

I don't have access to the 64-GPU job right now, so the description drops the old load numbers and labels the remaining save numbers as coming from an earlier revision. Here is a local V1 vs V2 comparison instead: 18 cores, 48 GB RAM, 8 storage units, 96 KiB rows, median of 3 runs.

	0.38 GiB	1.50 GiB
V1 save	0.65 s	2.29 s
V2 save	0.16 s	0.61–0.67 s
V1 load	0.40 s	1.36–1.45 s
V2 load	0.16 s	0.61–0.76 s
V1 load, 8→4 units	0.28 s	1.19–1.29 s
V2 load, 8→4 units	0.14 s	0.98–1.14 s

The 1.50 GiB column shows the range over two separate runs.

Caveats:

• One machine, warm page cache, and macOS fsync does not force a full flush to disk, so the absolute numbers are optimistic. The V1/V2 ratio is the useful part.
• Publishing the metadata updates concurrently matters mostly on wide loads: about 0.03 s for 128 updates, and 1.00 s → 0.88 s for 512.
• The 8→4 load gains less because each unit deserializes twice as many rows, at roughly 0.9 GiB/s per unit (mostly per-row pickle.loads).

Checks

• compileall, pre-commit and the sanity scripts pass.
• Locally the full suite gives 729 passed, 10 skipped. The dump e2e files also pass with Ray limited to 4 CPUs.
• CI on fdd8335 is green; "Run Tests" finishes in 11m23s with 762 passed, 11 skipped.