feat: kv update with the finalize flag

[OutstanderWang]

Problem

kv_update (#182) merges new values into stored fields, which suits fields built in several steps, such as a response appended chunk by chunk during rollout. But two problems get in the way:

• Half-built values are visible. Every update marks its fields ready in the controller, so consumers see the field after the first chunk. That includes StreamingDataLoader / sampler fetches, kv_batch_get and kv_retrieve_meta, and they can read an intermediate value.
• The first step has to be a kv_put. kv_update rejects keys or fields that don't exist yet, so callers need a different code path for the first chunk.

Summary

• finalize on kv_update / kv_batch_update (and their async_ variants), default True. With finalize=False, the value is stored and its metadata recorded, but the controller keeps the cells not ready. The field stays hidden from every consumer until an update with finalize=True.
• kv_update is a superset of kv_put. Missing keys are created, and a field the key doesn't hold yet is stored as given, without calling merge_fn. kv_put and restore don't take finalize.
• Dump/restore keeps readiness. A dump records each row's unfinalized fields, and restore keeps them hidden. The dump format becomes version 4; version 3 dumps still load, with every field ready.

Usage

import torch
import transfer_queue as tq

concat = lambda old, new: torch.cat([old, new])

# The first chunk creates the key and field; they stay hidden.
tq.kv_update(key="s0", partition_id="train", fields={"response": torch.tensor([1, 2])},
             merge_fn=concat, finalize=False)
tq.kv_update(key="s0", partition_id="train", fields={"response": torch.tensor([3])},
             merge_fn=concat, finalize=False)

tq.kv_describe(keys="s0", partition_id="train")["s0"]["unfinalized"]  # ['response']
tq.kv_batch_get(keys="s0", partition_id="train", select_fields="response")  # ValueError: not finalized yet

# The last chunk publishes the whole value.
tq.kv_update(key="s0", partition_id="train", fields={"response": torch.tensor([4])}, merge_fn=concat)
# response -> tensor([1, 2, 3, 4]), now visible to kv_batch_get and the dataloader

Implementation

The path is async_kv_batch_update → TransferQueueClient.async_update → AsyncSimpleStorageManager.update_data → StorageUnitData.apply_update on each owner unit → notify_data_update(ready=finalize) → DataPartitionStatus.update_production_status(ready=...).

• Readiness stays in the controller. The new ready flag in the update notification sets production_status to 0 or 1. Field metadata (dtype, shapes) is recorded either way. Storage never decides visibility, so the control-plane/data-plane split is unchanged.
• Creating missing keys and fields. kv_batch_update retrieves metadata with create=True. Each storage unit merges rows that already hold the field and stores the others as given. Every merge and type check finishes before the single put_data, so a failure leaves storage unchanged.
• Column-type check for new rows. kv_retrieve_meta takes include_fields, so the metadata carries the column type of every updated field the partition already knows, even when not every key holds it ready. Units check each new row's tensor/non-tensor kind and dtype against it before writing. A conflict fails before anything is stored or marked ready.
• Dump/restore. kv_describe reports each row's unfinalized fields, and the dump stores them in the row index. On restore, each unit publishes finalized and unfinalized fields in separate updates with the matching ready flag.
• No silent results.
  • kv_batch_get and kv_batch_get_by_meta raise when a selected field isn't ready in every key; before, they returned an empty TensorDict.
  • New kv_describe / async_kv_describe list each key's fields, including hidden ones.
  • kv_batch_update rejects duplicate keys before registering anything. A repeated missing key used to orphan rows, and a repeated existing key used to lose a merge.

Compatibility

• Existing kv_update calls keep their behavior, since finalize defaults to True. The one change: a missing key or field is now created instead of raising.
• kv_batch_get / kv_batch_get_by_meta with a select_fields entry that isn't ready in every key now raise ValueError instead of returning an empty result.
• Dumps are written as version 4. Builds that only read version 3 reject them instead of restoring hidden fields as ready.
• kv_update remains SimpleStorage-only. KV backends still raise NotImplementedError.

Known limitation

Updating a field that is already ready happens in place. This covers finalize=False on a ready field (hiding it again), finalize=True, and kv_put overwrites. Readers that already hold ready metadata can see the new value, or a mix of old and new across storage units. The normal streaming flow isn't affected, because its fields stay hidden until the final update. See the review thread for why that window is narrow and why hiding first wouldn't close it.

Tests

• Unit tests:
  • tests/test_kv_update.py: missing values stored without merge_fn, column-type checks on new rows before writing, duplicate-key rejection.
  • tests/test_controller_data_partitions.py: an unfinalized write records metadata but stays hidden from scans.
  • tests/test_data_dump.py
• e2e (tests/e2e/test_kv_interface_e2e.py, tests/e2e/test_data_dump_e2e.py):
  • unfinalized updates stay hidden until finalized, including the kv_batch_get error and kv_describe;
  • missing keys and fields are created;
  • a new row must match a column type held only by another key's hidden write;
  • restore keeps unfinalized fields hidden;
  • a version 3 dump restores as ready.
• pre-commit run --all-files and python -m pytest -q pass, with the Yuanrong e2e test excluded locally because of an environment-specific TORCH_LIBRARY registration error.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-13T11:34:56.362785Z	936657f	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 936657f949

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Hide fields before writing unfinalized payloads

When finalize=False updates a field that is currently ready, all storage-unit writes have already completed before this notification flips the controller's readiness. A concurrent kv_batch_get in that window can therefore receive ready metadata and read the intermediate payload; if the notification fails, the field can remain visible indefinitely. Hide the cells in the controller before dispatching the write (and account for already in-flight reads), then publish readiness only for a successful finalizing write.

AGENTS.md reference: AGENTS.md:L13-L14

Useful? React with 👍 / 👎.

[OutstanderWang]

Thanks. The window is real, but narrow, so I'm leaving the write order as it is in this PR.

• The flow finalize is meant for has no window. The first finalize=False chunk writes a field whose cells are already not ready (new, or still hidden), so nothing becomes visible early. The controller sets the ready bit only on the finalizing write, after every storage unit has committed.
• A window only exists in two cases. (a) A caller uses finalize=False to hide a field that is currently ready again. (b) The controller notification fails after storage committed. In case (b) the call raises "Storage update committed, but the controller did not accept its metadata" instead of failing silently.
• Hiding first wouldn't close the gap. Readers that already hold ready metadata still read the new value; you point this out yourself ("account for already in-flight reads"). The same exposure exists for any in-place update of a ready field: finalize=True updates, and the existing kv_put overwrite on main. In both, readers can see the new value or an old/new mix across units. A pre-hide RPC would add a round trip and a new state transition but only shrink the window.

Reader isolation for in-place updates would need versioned or leased reads, which is out of scope here. If it becomes a requirement, I'll handle it in a separate PR.

[chatgpt-codex-connector]

Validate missing-row values against the partition schema

When an update adds field f to a row while f already exists elsewhere in the partition with another tensor dtype, present is empty for this row, so its incoming value is stored as-is and only merged rows are type-checked. After the payload is committed, _build_update_field_schema reports the new dtype and the controller rejects it against FieldMeta, but the controller's readiness assignment has already run; the call raises while the conflicting value can be marked ready under the stale schema. Prevalidate missing-row values against the controller's partition schema before writing them.

AGENTS.md reference: AGENTS.md:L13-L14

Useful? React with 👍 / 👎.

[OutstanderWang]

Fixed in 3891aa4. kv_update now passes include_fields to kv_retrieve_meta, so the metadata carries the column schema for every updated field the partition already knows, including unfinalized fields and fields only some keys hold. The manager sends these column_types to each unit, and StorageUnitData.apply_update checks every new row's tensor/non-tensor kind and dtype before the single put_data. A conflict now fails before anything is stored or marked ready. Covered by test_new_row_must_match_column_type_held_by_other_keys (e2e) and test_apply_update_checks_new_rows_against_column_type_before_writing.

[chatgpt-codex-connector]

Reject duplicate keys before creating missing rows

When keys repeats a previously absent key (for example, ['k', 'k']), create=True makes the controller allocate two global indexes and then overwrite keys_mapping['k'] with the second one while retaining both global and reverse mappings. Both rows are written and become eligible for sampling, but only the second remains reachable by key or clear operations, leaking capacity and potentially emitting the orphaned first value; reject duplicate batch keys before registration.

Useful? React with 👍 / 👎.

[OutstanderWang]

Fixed in 63a751d. async_kv_batch_update now raises ValueError on repeated keys before any controller or storage call, so nothing is registered or written. This also covers a repeated existing key: there, both merges would read the same old value and one would be silently lost. Covered by test_batch_update_rejects_duplicate_keys_before_registering. kv_batch_put has the same create-path behavior on main, so a controller-side fix would be a separate change.

[chatgpt-codex-connector]

Bump the dump version for unfinalized readiness

When a dump contains an unfinalized field and is loaded by the immediately previous build, it still advertises DUMP_FORMAT_VERSION = 3, so that reader accepts it but omits unfinalized from load records and publishes every restored field as ready. A downgrade or mixed-version restore therefore silently exposes hidden intermediate data; write a new format version while allowing the current reader to treat legacy v3 dumps as fully finalized.

Useful? React with 👍 / 👎.

[OutstanderWang]

Fixed in 3891aa4. DUMP_FORMAT_VERSION is now 4. This build still reads version 3 dumps, which have no unfinalized, so every field restores as ready. Builds that only read version 3 reject a version 4 dump instead of exposing hidden fields. Covered by test_restores_a_version_3_dump_as_ready; docs/data_dump.md is updated.

[0oshowero0]

[AI Review] kv_update merges into storage values the controller no longer tracks

Problem

StorageUnitData.apply_update chooses between merging and storing from the unit's local state (idx in stored, simple_storage.py#L262-L277). But storage can hold values the controller has already dropped. AsyncSimpleStorageManager.clear_data logs per-unit failures and returns normally (simple_storage_manager.py#L918-L922). The controller then forgets the cell, and for kv_clear(keys) it releases the global index into the pool shared by all partitions.

On main these leftovers are unreachable: kv_update requires keys and fields the controller knows, and kv_put overwrites. With this PR, the next kv_update folds the new value into them:

• Same key: kv_clear(fields="response") fails in storage. kv_describe then lists only prompt, yet kv_update(response=[1]) stores [100, 101, 1].
• Another partition: kv_clear of old in pa fails in storage. The first kv_update of a new key in pb reuses that global index and stores [200, 201, 9], i.e. the other partition's payload.

The same rule makes retrying a first chunk unsafe: after a partial multi-unit failure, the units that already committed merge the chunk a second time.

This is separate from the in-place visibility window discussed here: the stored value itself is wrong, and it stays wrong for every later reader.

Repro (SimpleStorage with one storage unit, after tq.init):

from transfer_queue import interface

def concat(old, new):
    return torch.cat([old, new])

async def fail_clear(*_args, **_kwargs):
    raise RuntimeError("simulated storage unit timeout")

manager = interface._maybe_create_tq_client().storage_manager
real_clear = manager._clear_single_storage_unit

tq.kv_put(key="k", partition_id="p", fields={"response": torch.tensor([100, 101]), "prompt": torch.tensor([7])})
manager._clear_single_storage_unit = fail_clear
tq.kv_clear(keys="k", partition_id="p", fields="response")  # returns normally; the failure is only logged
manager._clear_single_storage_unit = real_clear

tq.kv_update(key="k", partition_id="p", fields={"response": torch.tensor([1])}, merge_fn=concat)
tq.kv_batch_get(keys="k", partition_id="p", select_fields="response")["response"][0]  # tensor([100, 101, 1])

Suggestion: let the controller decide which rows to merge

In async_kv_batch_update, replace kv_retrieve_meta(create=True, include_fields=...) with one controller request specific to kv_update, e.g. kv_prepare_update(keys, partition_id, fields). In a single actor call it registers missing keys and returns the usual BatchMeta. For each updated field it also returns the column type and the batch rows that hold the field. These rows come from FieldMeta.global_indexes, which also covers hidden finalize=False cells.

def kv_prepare_update(self, keys: list[str], partition_id: str, fields: list[str]):
    """Register missing keys and tell kv_update, per field, which rows to merge into."""
    metadata = self.kv_retrieve_meta(keys, partition_id, create=True)
    partition = self._get_partition(partition_id)
    batch_rows = set(metadata.global_indexes)
    columns = {}
    for name in fields:
        meta = partition.field_metadata.get(name)
        columns[name] = {
            "dtype": meta.dtype if meta else None,
            "is_non_tensor": meta.is_non_tensor if meta else None,
            "present": sorted(meta.global_indexes & batch_rows) if meta else [],
        }
    return metadata, columns

The manager passes each unit its slice of present, and apply_update merges exactly those rows:

held = set(columns[field]["present"])
present = [i for i, idx in enumerate(global_indexes) if idx in held]  # was: idx in stored

Example:

• s0 (index 3) holds a hidden response=[1, 2].
• s1 (index 4) holds only prompt, but its unit keeps a leftover response=[200, 201].
• s2 is new and is registered as index 7.

For a batch update of response = [3], [5], [9], the controller returns {"response": {"dtype": torch.int64, "is_non_tensor": False, "present": [3]}}. Row 3 merges to [1, 2, 3], row 7 stores [9], and row 4 stores [5] over the leftover instead of [200, 201, 5].

What this changes:

• Leftovers are overwritten, never merged.
• Retrying a first chunk whose write committed but was never recorded by the controller now overwrites instead of merging twice. Retries of later chunks still must not be blind.
• If the controller holds a row that the unit lacks (e.g. it was cleared concurrently), the update fails with StorageKeyNotFoundError instead of storing the value as new.

Cost:

• No extra RPC: the request replaces the existing kv_retrieve_meta call.
• The controller does one set lookup per row per updated field, which is tens of microseconds for a few thousand rows. The mask is a few KB.
• include_fields and the manager's column_types construction go away. kv_retrieve_meta keeps its "fields ready in every row" meaning, and BatchMeta is unchanged.
• The return value is unaffected, because apply_field_schema already adds the updated fields to field_names.