From e65d207bc2c1798d63118f035b7965542662d764 Mon Sep 17 00:00:00 2001
From: OutstanderWang <wangweiyanster@gmail.com>
Date: Tue, 13 Oct 2026 19:19:39 +0800
Subject: [PATCH 01/12] [feat] kv_update: add missing keys/fields and defer
 readiness with finalize

kv_update is now a superset of kv_put: a missing key is created, and a field
the key does not hold yet is stored as given without calling merge_fn. Merge
results are validated against the stored old value, so the update request no
longer carries a controller field_schema.

finalize=False stores the update but keeps the fields not ready, so samplers,
the dataloader and kv_batch_get do not see intermediate values; the update
with finalize=True (the default) marks them ready. Selective dump/restore
keeps each field's readiness as it was when dumped.

Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
---
 README.md                                     | 18 ++++--
 tests/e2e/test_data_dump_e2e.py               | 33 +++++++++-
 tests/e2e/test_kv_interface_e2e.py            | 63 +++++++++++++-----
 tests/test_controller_data_partitions.py      | 12 ++++
 tests/test_data_dump.py                       | 18 +++++-
 tests/test_kv_update.py                       | 41 ++++++------
 transfer_queue/client.py                      | 19 ++++--
 transfer_queue/controller.py                  | 44 ++++++++-----
 transfer_queue/data_dump.py                   |  4 +-
 transfer_queue/interface.py                   | 53 +++++++++------
 transfer_queue/storage/managers/base.py       | 11 +++-
 .../managers/simple_storage_manager.py        | 13 ++--
 transfer_queue/storage/simple_storage.py      | 64 +++++++++----------
 tutorial/02_kv_interface.py                   |  2 +-
 14 files changed, 266 insertions(+), 129 deletions(-)

diff --git a/README.md b/README.md
index 4df55109..23367a11 100644
--- a/README.md
+++ b/README.md
@@ -116,15 +116,21 @@ To simplify the usage of TransferQueue, we provide a Redis-style high-level API
 - **(async_)kv_put**: Insert/Update a multi-column sample by key, with an optional metadata tag.
 - **(async_)kv_batch_put**: Put multiple key-value pairs efficiently in batches.
 - **(async_)kv_batch_get**: Retrieve samples (by keys), supporting column selection (by fields).
-- **(async_)kv_update / (async_)kv_batch_update**: Merge new values into already-produced fields with `merge_fn(old, new)`. SimpleStorage only.
+- **(async_)kv_update / (async_)kv_batch_update**: Merge new values into fields with `merge_fn(old, new)`; a superset of put that adds missing keys and fields. SimpleStorage only.
 - **(async_)kv_list**: List keys and tags (metadata) in a partition.
 - **(async_)kv_clear**: Remove key-value pairs from storage, or only selected `fields` while keeping the keys.
 
-`kv_update` preserves each field's tensor/non-tensor type and tensor dtype. A
-timeout has an unknown outcome because a storage unit may commit after the
-caller stops waiting, so non-idempotent merge operations must not be retried
-blindly. Updates are atomic within one storage-unit request; a batch spanning
-multiple units is not a distributed transaction.
+`kv_update` stores a field the key does not hold yet as given, without calling
+`merge_fn`; a merge must preserve the stored value's tensor/non-tensor type and
+tensor dtype. A timeout has an unknown outcome because a storage unit may commit
+after the caller stops waiting, so non-idempotent merge operations must not be
+retried blindly. Updates are atomic within one storage-unit request; a batch
+spanning multiple units is not a distributed transaction.
+
+To build a field over several updates without exposing intermediate values, pass
+`finalize=False` to every `kv_update` except the last one, which uses the default
+`finalize=True`. Until then the field is stored but not ready, so samplers, the
+dataloader and `kv_batch_get` do not see it.
 
 **Key Features**
 
diff --git a/tests/e2e/test_data_dump_e2e.py b/tests/e2e/test_data_dump_e2e.py
index a4fd5b9c..07fd56f8 100644
--- a/tests/e2e/test_data_dump_e2e.py
+++ b/tests/e2e/test_data_dump_e2e.py
@@ -243,6 +243,35 @@ def test_heterogeneous_field_sets_are_grouped(self, tq_system, dump_dir, control
         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_restore_keeps_unfinalized_fields_hidden(self, tq_system, dump_dir, controller):
+        partition_id = "d_unfinalized"
+        _put_rows(partition_id, ["u0", "u1"])
+
+        def concat(old, new):
+            return torch.cat([old, new])
+
+        tq.kv_update(
+            key="u1",
+            partition_id=partition_id,
+            fields={"response": torch.tensor([1, 2])},
+            merge_fn=concat,
+            finalize=False,
+        )
+
+        tq.dump_data_by_key(dump_dir, ["u0", "u1"], partition_id)
+        assert _read_row_index(dump_dir)["rows"]["u1"]["unfinalized"] == ["response"]
+        ray.get(controller.clear_partition.remote(partition_id))
+        tq.load_data_by_key(dump_dir)
+
+        snapshot = ray.get(controller.get_partition_snapshot.remote(partition_id))
+        row = snapshot.keys_mapping["u1"]
+        assert snapshot.production_status[row, snapshot.field_name_mapping["input_ids"]] == 1
+        assert snapshot.production_status[row, snapshot.field_name_mapping["response"]] == 0
+
+        tq.kv_update(key="u1", partition_id=partition_id, fields={"response": torch.tensor([3])}, merge_fn=concat)
+        retrieved = tq.kv_batch_get(keys="u1", partition_id=partition_id, select_fields="response")
+        _assert_rows_equal(retrieved["response"], [torch.tensor([1, 2, 3])])
+
     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
@@ -681,10 +710,10 @@ 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)]
+            updates = [{"global_indexes": [i], "field_schema": {}, "ready": True} for i in range(num_updates)]
             return {"bytes_read": 0, "updates": updates}
 
-        async def notify(partition_id, global_indexes, field_schema):
+        async def notify(partition_id, global_indexes, field_schema, **_kwargs):
             started.append(global_indexes)
             if len(started) == num_updates:
                 ready.set()
diff --git a/tests/e2e/test_kv_interface_e2e.py b/tests/e2e/test_kv_interface_e2e.py
index fa5d488b..ed88a995 100644
--- a/tests/e2e/test_kv_interface_e2e.py
+++ b/tests/e2e/test_kv_interface_e2e.py
@@ -1244,25 +1244,58 @@ def test_batch_update_multiple_fields_across_storage_units(self, controller, tq_
         assert "score" not in partition.field_metadata
         assert set(tq_api.kv_list(partition_id=partition_id)[partition_id]) == set(keys)
 
-    def test_rejects_missing_key_unproduced_field_and_incompatible_result(self, tq_api, backend_name):
+    def test_unfinalized_updates_stay_hidden_until_finalize(self, controller, tq_api, backend_name):
         if backend_name != "SimpleStorage":
             pytest.skip("merge-backed kv_update is implemented only for SimpleStorage")
 
-        with pytest.raises(ValueError, match="not found"):
-            tq_api.kv_update(
-                key="no_such",
-                partition_id="test_partition",
-                fields={"tokens": torch.tensor([1])},
-                merge_fn=lambda _old, new: new,
-            )
+        partition_id = "test_partition"
+        key = "streamed"
+
+        def concat(old, new):
+            return torch.cat([old, new])
+
+        def tokens_ready():
+            partition = get_controller_partition(controller, partition_id)
+            row, col = partition.keys_mapping[key], partition.field_name_mapping["tokens"]
+            return bool(partition.production_status[row, col])
+
+        tq_api.kv_put(key=key, partition_id=partition_id, fields={"prompt": torch.tensor([7])})
+        for chunk in ([1], [2]):
+            fields = {"tokens": torch.tensor(chunk)}
+            tq_api.kv_update(key=key, partition_id=partition_id, fields=fields, merge_fn=concat, finalize=False)
+            assert not tokens_ready()
+        assert set(tq_api.kv_batch_get(keys=key, partition_id=partition_id).keys()) == {"prompt"}
+
+        tq_api.kv_update(key=key, partition_id=partition_id, fields={"tokens": torch.tensor([3])}, merge_fn=concat)
+        assert tokens_ready()
+        data = tq_api.kv_batch_get(keys=key, partition_id=partition_id, select_fields="tokens")
+        assert_tensor_equal(data["tokens"][0], torch.tensor([1, 2, 3]))
+
+        # A non-final update hides a field that was already ready.
+        tq_api.kv_update(
+            key=key, partition_id=partition_id, fields={"tokens": torch.tensor([4])}, merge_fn=concat, finalize=False
+        )
+        assert not tokens_ready()
+
+    def test_creates_missing_key_and_field_and_rejects_incompatible_result(self, tq_api, backend_name):
+        if backend_name != "SimpleStorage":
+            pytest.skip("merge-backed kv_update is implemented only for SimpleStorage")
+
+        def fail(_old, _new):
+            raise AssertionError("merge_fn must not run for a missing value")
+
+        meta = tq_api.kv_update(
+            key="new_key", partition_id="test_partition", fields={"tokens": torch.tensor([1, 2])}, merge_fn=fail
+        )
+        assert meta.fields == ["tokens"]
+        tq_api.kv_update(
+            key="new_key", partition_id="test_partition", fields={"fresh": torch.tensor([9])}, merge_fn=fail
+        )
+        created = tq_api.kv_batch_get(keys="new_key", partition_id="test_partition")
+        assert_tensor_equal(created["tokens"][0], torch.tensor([1, 2]))
+        assert_tensor_equal(created["fresh"][0], torch.tensor([9]))
+
         tq_api.kv_put(key="sample", partition_id="test_partition", fields={"tokens": torch.tensor([1, 2])})
-        with pytest.raises(ValueError, match="not produced"):
-            tq_api.kv_update(
-                key="sample",
-                partition_id="test_partition",
-                fields={"fresh": torch.tensor([1])},
-                merge_fn=lambda _old, new: new,
-            )
         with pytest.raises(RuntimeError, match="keep dtype"):
             tq_api.kv_update(
                 key="sample",
diff --git a/tests/test_controller_data_partitions.py b/tests/test_controller_data_partitions.py
index ef8538cc..1a6cced5 100644
--- a/tests/test_controller_data_partitions.py
+++ b/tests/test_controller_data_partitions.py
@@ -1350,6 +1350,18 @@ def test_clear_fields_drops_fieldmeta_when_no_sample_holds_it(self):
         assert "pv" not in partition.field_metadata
         assert partition.global_indexes == {0, 1}
 
+    def test_unfinalized_write_records_metadata_but_stays_hidden(self):
+        partition = self._make_partition()
+        t4 = {"dtype": "torch.float32", "shape": (4,), "is_nested": False, "is_non_tensor": False}
+        partition.update_production_status([0, 1], ["pv"], {"pv": t4})
+
+        assert partition.update_production_status([1], ["pv"], {"pv": t4}, ready=False)
+        assert partition.field_metadata["pv"].global_indexes == {0, 1}
+        assert partition.scan_data_status(["pv"], task_name="trainer") == [0]
+
+        partition.update_production_status([1], ["pv"], {"pv": t4})
+        assert partition.scan_data_status(["pv"], task_name="trainer") == [0, 1]
+
     def test_update_production_status_updates_field_metadata(self):
         """Test that update_production_status correctly updates field_metadata via _update_field_metadata."""
         partition = self._make_partition()
diff --git a/tests/test_data_dump.py b/tests/test_data_dump.py
index b6eecb55..b2de9cfe 100644
--- a/tests/test_data_dump.py
+++ b/tests/test_data_dump.py
@@ -149,7 +149,14 @@ def read(self, size=-1):
     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}
+            {
+                "source_index": source,
+                "target_index": target,
+                "fields": ["x"],
+                "unfinalized": [],
+                "offset": offset,
+                "length": length,
+            }
         )
     loaded = unit._handle_load_rows(
         ZMQMessage.create(
@@ -182,7 +189,14 @@ def read(self, size=-1):
 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"]}
+    record = {
+        "source_index": 1,
+        "target_index": 2,
+        "offset": 0,
+        "length": path.stat().st_size,
+        "fields": ["x"],
+        "unfinalized": [],
+    }
     if problem == "wrong_index":
         record["source_index"] = 3
     elif problem == "truncated":
diff --git a/tests/test_kv_update.py b/tests/test_kv_update.py
index 567722ff..7a97e0cd 100644
--- a/tests/test_kv_update.py
+++ b/tests/test_kv_update.py
@@ -35,10 +35,6 @@ def _concat(old, new):
     return torch.cat([old, new])
 
 
-def _tensor_schema(*fields, dtype=torch.int64):
-    return {field: {"dtype": dtype, "shape": None, "is_nested": True, "is_non_tensor": False} for field in fields}
-
-
 def test_single_update_batch_wraps_tensor_and_python_values():
     batch = _single_update_batch({"tokens": torch.tensor([4, 5]), "meta": {"step": 1}})
 
@@ -57,12 +53,7 @@ def test_apply_update_concatenates_and_keeps_field_name():
     data = StorageUnitData()
     data.put_data({"sequence_ids": [torch.tensor([10, 11, 12])]}, [0])
 
-    described = data.apply_update(
-        [0],
-        {"sequence_ids": [torch.tensor([20, 21])]},
-        _concat,
-        _tensor_schema("sequence_ids"),
-    )
+    described = data.apply_update([0], {"sequence_ids": [torch.tensor([20, 21])]}, _concat)
 
     assert described["sequence_ids"] == {"dtype": torch.int64, "shapes": [(5,)]}
     assert torch.equal(data.field_data["sequence_ids"][0], torch.tensor([10, 11, 12, 20, 21]))
@@ -78,17 +69,28 @@ def fail(_old, _new):
         raise RuntimeError("merge failed")
 
     with pytest.raises(RuntimeError, match="merge failed"):
-        data.apply_update([7], {"tokens": [torch.tensor([9])]}, fail, _tensor_schema("tokens"))
+        data.apply_update([7], {"tokens": [torch.tensor([9])]}, fail)
     assert torch.equal(data.field_data["tokens"][7], original)
 
 
-def test_apply_update_rejects_unproduced_field():
+def test_apply_update_stores_missing_values_without_merging():
     data = StorageUnitData()
     data.put_data({"tokens": [torch.tensor([1])]}, [3])
+    merged = []
+
+    def concat(old, new):
+        merged.append(old)
+        return torch.cat([old, new])
+
+    described = data.apply_update(
+        [3, 4], {"tokens": [torch.tensor([2]), torch.tensor([5])], "fresh": [{"a": 1}, {"a": 2}]}, concat
+    )
 
-    with pytest.raises(ValueError, match="unproduced field"):
-        data.apply_update([3], {"fresh": [torch.tensor([2])]}, lambda _old, new: new, _tensor_schema("fresh"))
-    assert "fresh" not in data.field_data
+    assert len(merged) == 1
+    assert torch.equal(data.field_data["tokens"][3], torch.tensor([1, 2]))
+    assert torch.equal(data.field_data["tokens"][4], torch.tensor([5]))
+    assert data.field_data["fresh"] == {3: {"a": 1}, 4: {"a": 2}}
+    assert described["tokens"]["shapes"] == [(2,), (1,)]
 
 
 @pytest.mark.parametrize(
@@ -104,17 +106,16 @@ def test_apply_update_rejects_incompatible_tensor_results_before_writing(merge_f
     data.put_data({"tokens": [original.clone()]}, [0])
 
     with pytest.raises(TypeError, match=error):
-        data.apply_update([0], {"tokens": [torch.tensor([3])]}, merge_fn, _tensor_schema("tokens"))
+        data.apply_update([0], {"tokens": [torch.tensor([3])]}, merge_fn)
     assert torch.equal(data.field_data["tokens"][0], original)
 
 
 def test_apply_update_rejects_tensor_result_for_non_tensor_field():
     data = StorageUnitData()
     data.put_data({"meta": [{"step": 1}]}, [0])
-    schema = {"meta": {"dtype": None, "shape": None, "is_nested": False, "is_non_tensor": True}}
 
     with pytest.raises(TypeError, match="must remain non-tensor"):
-        data.apply_update([0], {"meta": [{"step": 2}]}, lambda _old, _new: torch.tensor([2]), schema)
+        data.apply_update([0], {"meta": [{"step": 2}]}, lambda _old, _new: torch.tensor([2]))
     assert data.field_data["meta"][0] == {"step": 1}
 
 
@@ -126,7 +127,7 @@ def test_apply_update_decodes_ssd_offloaded_old_value(tmp_path):
     data.put_data({"tokens": [prompt]}, [0])
     assert data.ssd_active_values == 1
 
-    data.apply_update([0], {"tokens": [torch.tensor([99])]}, _concat, _tensor_schema("tokens"))
+    data.apply_update([0], {"tokens": [torch.tensor([99])]}, _concat)
 
     assert torch.equal(data.get_data(["tokens"], [0])["tokens"][0], torch.cat([prompt, torch.tensor([99])]))
     assert data.ssd_active_values == 1
@@ -152,7 +153,7 @@ def fail_second_field(encoded_values):
         monkeypatch.setattr(data._ssd_store, "write_values", fail_second_field)
         new_data = {"tokens": [torch.tensor([99])], "mask": [torch.tensor([1])]}
         with pytest.raises(OSError, match="No space left"):
-            data.apply_update([0], new_data, _concat, _tensor_schema("tokens", "mask"))
+            data.apply_update([0], new_data, _concat)
 
         stored = data.get_data(["tokens", "mask"], [0])
         assert torch.equal(stored["tokens"][0], prompt)
diff --git a/transfer_queue/client.py b/transfer_queue/client.py
index 41ff29d4..b7ee8d27 100644
--- a/transfer_queue/client.py
+++ b/transfer_queue/client.py
@@ -497,13 +497,15 @@ async def async_update(
         metadata: BatchMeta,
         values: TensorDict,
         merge_fn: Callable[[Any, Any], Any],
+        finalize: bool = True,
     ) -> BatchMeta:
-        """Merge new values into produced fields on SimpleStorage units.
+        """Merge new values into fields on SimpleStorage units, storing them where absent.
 
         Args:
-            metadata: Samples to update. The key must already exist.
+            metadata: Samples to update, already registered with the controller.
             values: New values aligned with ``metadata``.
-            merge_fn: Called per sample per field as ``merge_fn(old, new)``.
+            merge_fn: Called per sample per field as ``merge_fn(old, new)`` when a value exists.
+            finalize: If False, keep the updated fields hidden from consumers.
 
         Returns:
             The same metadata with field_schema and production status updated.
@@ -519,8 +521,10 @@ async def async_update(
         if not metadata or metadata.size == 0:
             raise ValueError("metadata cannot be none or empty")
 
-        field_schema = await self.storage_manager.update_data(metadata, values, merge_fn)
-        return metadata.apply_field_schema(field_schema)
+        field_schema = await self.storage_manager.update_data(metadata, values, merge_fn, finalize=finalize)
+        if not finalize:
+            metadata.production_status[:] = 0
+        return metadata.apply_field_schema(field_schema, set_all_ready=finalize)
 
     async def async_get_data(self, metadata: BatchMeta) -> TensorDict:
         """Asynchronously fetch data from storage units and organize into TensorDict.
@@ -1638,12 +1642,13 @@ def update(
         metadata: BatchMeta,
         values: TensorDict,
         merge_fn: Callable[[Any, Any], Any],
+        finalize: bool = True,
     ) -> BatchMeta:
-        """Synchronously merge new values into produced fields on SimpleStorage units.
+        """Synchronously merge or add values on SimpleStorage units.
 
         See ``async_update``.
         """
-        return self._update(metadata, values, merge_fn)
+        return self._update(metadata, values, merge_fn, finalize=finalize)
 
     def get_data(self, metadata: BatchMeta) -> TensorDict:
         """Synchronously fetch data from storage units and organize into TensorDict.
diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py
index c2ba668f..ba8199ff 100644
--- a/transfer_queue/controller.py
+++ b/transfer_queue/controller.py
@@ -500,6 +500,7 @@ def update_production_status(
         field_names: list[str],
         field_schema: dict[str, dict[str, Any]],
         custom_backend_meta: dict[int, dict[str, Any]] | None = None,
+        ready: bool = True,
     ) -> bool:
         """
         Update production status for specific samples and fields.
@@ -515,6 +516,8 @@ def update_production_status(
             field_schema: Columnar field schema {field_name: {dtype, shape, is_nested, ...}}
             custom_backend_meta: Optional per-sample per-field
                 custom metadata provided by storage backend
+            ready: Whether consumers may see these cells. False records the stored
+                value (field metadata) but hides it until a finalizing kv_update.
 
         Returns:
             True if update was successful, False on error
@@ -543,7 +546,7 @@ def update_production_status(
             # Update production status
             if self.production_status is not None and global_indices and field_names:
                 field_indices = [self.field_name_mapping.get(f) for f in field_names]
-                self.production_status[torch.tensor(global_indices)[:, None], torch.tensor(field_indices)] = 1
+                self.production_status[torch.tensor(global_indices)[:, None], torch.tensor(field_indices)] = int(ready)
 
             # Update field metadata
             self._update_field_metadata(global_indices, field_schema, custom_backend_meta)
@@ -1130,6 +1133,7 @@ def update_production_status(
         global_indexes: list[int],
         field_schema: dict[str, dict[str, Any]],
         custom_backend_meta: dict[int, dict[str, Any]] | None = None,
+        ready: bool = True,
     ) -> bool:
         """
         Update production status for specific samples and fields in a partition.
@@ -1140,6 +1144,7 @@ def update_production_status(
             global_indexes: List of sample indices to update
             field_schema: Columnar field schema {field_name: {dtype, shape, is_nested, ...}}
             custom_backend_meta: Optional custom backend metadata
+            ready: False keeps the written cells hidden from consumers (unfinalized kv_update)
 
         Returns:
             True if update was successful, False otherwise
@@ -1150,11 +1155,13 @@ def update_production_status(
             logger.error(f"Partition {partition_id} not found")
             return False
 
-        success = partition.update_production_status(global_indexes, field_names, field_schema, custom_backend_meta)
+        success = partition.update_production_status(
+            global_indexes, field_names, field_schema, custom_backend_meta, ready=ready
+        )
         if success:
             logger.debug(
                 f"[{self.controller_id}]: Updated production status for partition {partition_id}: "
-                f"samples={global_indexes}, fields={field_names}"
+                f"samples={global_indexes}, fields={field_names}, ready={ready}"
             )
         return success
 
@@ -1783,9 +1790,10 @@ def kv_describe(self, keys: list[str], partition_id: str) -> dict[str, Any]:
             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.
+            ``{"rows": {key: {"global_index", "fields", "unfinalized", "tag"}}, "field_schema":
+            {field: {"is_nested", "is_non_tensor"}}}``. ``fields`` is empty for a row with no
+            produced field yet; ``unfinalized`` lists the stored fields not yet ready (written
+            by ``kv_update(finalize=False)``); ``field_schema`` covers the described rows' fields.
 
         Raises:
             KeyError: The partition or any key does not exist.
@@ -1799,18 +1807,23 @@ def kv_describe(self, keys: list[str], partition_id: str) -> dict[str, Any]:
         if missing_keys:
             raise KeyError(f"keys not found in partition {partition_id!r}: {missing_keys}")
 
-        rows: dict[str, dict[str, Any]] = {
-            key: {
+        rows: dict[str, dict[str, Any]] = {}
+        for key, global_index in zip(keys, cast(list[int], global_indexes), strict=True):
+            fields = sorted(
+                field_name
+                for field_name, field_meta in partition.field_metadata.items()
+                if global_index in field_meta.global_indexes
+            )
+            rows[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
-                ),
+                "fields": fields,
+                "unfinalized": [
+                    name
+                    for name in fields
+                    if not partition.production_status[global_index, partition.field_name_mapping[name]]
+                ],
                 "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)}
@@ -2097,6 +2110,7 @@ def _handle_notify_data_update_request(self, request_msg: ZMQMessage) -> ZMQMess
             global_indexes=global_indexes,
             field_schema=params.get("field_schema", {}),
             custom_backend_meta=params.get("custom_backend_meta", {}),
+            ready=params.get("ready", True),
         )
         if success:
             if self._metrics is not None:
diff --git a/transfer_queue/data_dump.py b/transfer_queue/data_dump.py
index 80fa53ca..9d30d5b7 100644
--- a/transfer_queue/data_dump.py
+++ b/transfer_queue/data_dump.py
@@ -23,7 +23,7 @@
 
     <dump_dir>/
         dump_info.json                # version, partition, counts
-        row_index.pt                  # key -> {global_index, fields, tag}
+        row_index.pt                  # key -> {global_index, fields, unfinalized, tag}
         shards/
             shard_info.json           # file, unit, row count, source index -> [offset, length]
             shard_<N>_<su_id>.pkl      # independent {global_index, fields} records
@@ -215,6 +215,8 @@ def read_dump(dump_dir: Path) -> tuple[dict[str, Any], list[dict[str, Any]]]:
                     "key": key,
                     "source_index": source_index,
                     "fields": rows[key]["fields"],
+                    # Older dumps lack this key; every field they hold was ready.
+                    "unfinalized": rows[key].get("unfinalized", []),
                     "offset": offset,
                     "length": length,
                 }
diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py
index b3e0ede4..74f257b9 100644
--- a/transfer_queue/interface.py
+++ b/transfer_queue/interface.py
@@ -629,22 +629,29 @@ def kv_update(
     partition_id: str,
     fields: dict[str, Any],
     merge_fn: Callable[[Any, Any], Any],
+    finalize: bool = True,
 ) -> KVBatchMeta:
-    """Merge new values into produced fields of one existing key.
+    """Merge new values into the fields of one key, adding the key or fields if missing.
+
+    A superset of ``kv_put``: a missing key is created and a field the key does not
+    hold yet is stored as given, without calling ``merge_fn``.
 
     Args:
-        key: Existing user-specified key. The key is not created if it is missing.
+        key: User-specified key. Created if missing.
         partition_id: Partition that holds the key.
-        fields: New values keyed by an already-produced field name.
-        merge_fn: ``merge_fn(old, new) -> stored``. SimpleStorage runs it once
-            per field on the unit that owns the sample.
+        fields: New values keyed by field name.
+        merge_fn: ``merge_fn(old, new) -> stored`` for fields the key already holds.
+            SimpleStorage runs it once per field on the unit that owns the sample.
+        finalize: If False, the updated fields become (or stay) hidden from consumers,
+            so readers never see an intermediate value. Pass True (the default) on the
+            last update to make them ready.
 
     Returns:
         Metadata for the updated sample.
 
     Raises:
-        ValueError: If the key or a field is missing.
-        TypeError: If ``fields`` is not a non-empty dict or ``merge_fn`` is not callable.
+        TypeError: If ``fields`` is not a non-empty dict, ``merge_fn`` is not callable,
+            or a merge result changes the stored value's type or dtype.
         NotImplementedError: If the storage backend is not SimpleStorage.
 
     Note:
@@ -658,6 +665,7 @@ def kv_update(
             partition_id=partition_id,
             fields=fields,
             merge_fn=merge_fn,
+            finalize=finalize,
         )
     )
 
@@ -667,15 +675,16 @@ def kv_batch_update(
     partition_id: str,
     fields: TensorDict,
     merge_fn: Callable[[Any, Any], Any],
+    finalize: bool = True,
 ) -> KVBatchMeta:
-    """Merge batched values into produced fields of existing keys.
+    """Merge batched values into the fields of keys, adding keys or fields if missing.
 
-    A timeout has an unknown outcome because a unit may commit after the caller
-    stops waiting. Do not blindly retry a merge.
+    See ``kv_update``. A timeout has an unknown outcome because a unit may commit
+    after the caller stops waiting. Do not blindly retry a merge.
     """
     tq_client = _maybe_create_tq_client()
     return tq_client._run_coroutine(
-        async_kv_batch_update(keys=keys, partition_id=partition_id, fields=fields, merge_fn=merge_fn)
+        async_kv_batch_update(keys=keys, partition_id=partition_id, fields=fields, merge_fn=merge_fn, finalize=finalize)
     )
 
 
@@ -1091,13 +1100,15 @@ async def async_kv_update(
     partition_id: str,
     fields: dict[str, Any],
     merge_fn: Callable[[Any, Any], Any],
+    finalize: bool = True,
 ) -> KVBatchMeta:
-    """Asynchronously update fields of an existing key. See ``kv_update``."""
+    """Asynchronously merge or add fields of one key. See ``kv_update``."""
     return await async_kv_batch_update(
         keys=[key],
         partition_id=partition_id,
         fields=_single_update_batch(fields),
         merge_fn=merge_fn,
+        finalize=finalize,
     )
 
 
@@ -1106,8 +1117,9 @@ async def async_kv_batch_update(
     partition_id: str,
     fields: TensorDict,
     merge_fn: Callable[[Any, Any], Any],
+    finalize: bool = True,
 ) -> KVBatchMeta:
-    """Asynchronously merge batched values into produced fields of existing keys."""
+    """Asynchronously merge or add batched fields of keys. See ``kv_update``."""
     if not isinstance(keys, list) or not keys or not all(isinstance(key, str) for key in keys):
         raise TypeError("keys must be a non-empty list of strings")
     if not isinstance(fields, TensorDict) or not fields.keys():
@@ -1117,15 +1129,18 @@ async def async_kv_batch_update(
     if not callable(merge_fn):
         raise TypeError("merge_fn must be callable")
     tq_client = _maybe_create_tq_client()
-    batch_meta = await tq_client.async_kv_retrieve_meta(keys=keys, partition_id=partition_id, create=False)
+    # Checked up front: missing keys are registered before storage is asked to merge.
+    if not isinstance(tq_client.storage_manager, AsyncSimpleStorageManager):
+        raise NotImplementedError(
+            f"kv_update is not supported by {type(tq_client.storage_manager).__name__}; "
+            "it requires the SimpleStorage backend"
+        )
+    batch_meta = await tq_client.async_kv_retrieve_meta(keys=keys, partition_id=partition_id, create=True)
 
     if batch_meta.size != len(keys):
-        raise ValueError("Some keys or the partition were not found")
-    missing_fields = sorted(set(fields.keys()) - set(batch_meta.field_names))
-    if missing_fields:
-        raise ValueError(f"Fields are not produced for every key: {missing_fields}")
+        raise RuntimeError(f"Retrieved BatchMeta size {batch_meta.size} does not match {len(keys)} keys")
 
-    batch_meta = await tq_client.async_update(metadata=batch_meta, values=fields, merge_fn=merge_fn)
+    batch_meta = await tq_client.async_update(metadata=batch_meta, values=fields, merge_fn=merge_fn, finalize=finalize)
 
     return KVBatchMeta(
         keys=keys,
diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py
index 1015172e..02fc1575 100644
--- a/transfer_queue/storage/managers/base.py
+++ b/transfer_queue/storage/managers/base.py
@@ -226,15 +226,17 @@ async def notify_data_update(
         global_indexes: list[int],
         field_schema: dict[str, dict[str, Any]],
         custom_backend_meta: dict[int, dict[str, Any]] | None = None,
+        ready: bool = True,
     ) -> bool:
         """
-        Notify controller that new data is ready.
+        Notify controller that new data is stored.
 
         Args:
             partition_id: Current data partition id.
             global_indexes: Data update related global_indexes.
             field_schema: Columnar field schema {field_name: {dtype, shape, is_nested, ...}}.
             custom_backend_meta: Per-field custom_meta for each sample, in {global_index: {field: custom_meta}} format.
+            ready: Whether consumers may see the data now; False for a kv_update with finalize=False.
         """
 
         if not self.controller_info:
@@ -264,6 +266,7 @@ async def notify_data_update(
                 "global_indexes": global_indexes,
                 "field_schema": normalized_field_schema,
                 "custom_backend_meta": custom_backend_meta,
+                "ready": ready,
             },
         ).serialize()
 
@@ -362,13 +365,15 @@ async def update_data(
         metadata: BatchMeta,
         values: TensorDict,
         merge_fn: Callable[[Any, Any], Any],
+        finalize: bool = True,
     ) -> dict[str, dict[str, Any]]:
-        """Merge new values on the unit that holds each sample.
+        """Merge new values on the unit that holds each sample, or store them if absent.
 
         Args:
             metadata: Samples to update.
             values: New values as a TensorDict aligned with ``metadata``.
-            merge_fn: Called per sample per field as ``merge_fn(old, new)``.
+            merge_fn: Called per sample per field as ``merge_fn(old, new)`` when a value exists.
+            finalize: If False, keep the updated fields hidden from consumers.
 
         Returns:
             field_schema of the stored values, keyed by field name.
diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py
index b5bf88b2..1f90fbfe 100644
--- a/transfer_queue/storage/managers/simple_storage_manager.py
+++ b/transfer_queue/storage/managers/simple_storage_manager.py
@@ -544,11 +544,12 @@ async def update_data(
         metadata: BatchMeta,
         values: TensorDict,
         merge_fn: Callable[[Any, Any], Any],
+        finalize: bool = True,
     ) -> dict[str, dict[str, Any]]:
-        """Merge values on each owner unit and notify the controller.
+        """Merge or add values on each owner unit and notify the controller.
 
         Each unit reports the values it stored; this assembles one batch-level
-        field_schema so production status stays ready.
+        field_schema. ``finalize`` decides whether the fields become ready.
         """
         logger.debug(f"[{self.storage_manager_id}]: receive update_data request, updating {metadata.size} samples.")
 
@@ -574,7 +575,6 @@ async def update_data(
                         group.global_indexes,
                         storage_data,
                         merge_fn,
-                        {name: metadata.field_schema[name] for name in field_names},
                         target_storage_unit=su_id,
                     ),
                     # A merge reads stored state, so a replay after a lost answer
@@ -603,6 +603,7 @@ async def update_data(
             metadata.partition_ids[0],
             metadata.global_indexes,
             field_schema,
+            ready=finalize,
         )
         if not published:
             raise RuntimeError(
@@ -674,7 +675,6 @@ async def _update_to_single_storage_unit(
         global_indexes: list[int],
         storage_data: dict[str, Any],
         merge_fn: Callable[[Any, Any], Any],
-        field_schema: dict[str, dict[str, Any]],
         target_storage_unit: str,
         socket: zmq.Socket = None,
     ) -> dict[str, dict[str, Any]]:
@@ -683,7 +683,6 @@ async def _update_to_single_storage_unit(
             "global_indexes": global_indexes,
             "data": storage_data,
             "merge_fn": merge_fn,
-            "field_schema": field_schema,
         }
 
         request_msg = ZMQMessage.create(
@@ -1139,7 +1138,9 @@ async def load_rows_by_index(self, partition_id: str, shards: list[dict[str, Any
         # 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"])
+                self.notify_data_update(
+                    partition_id, update["global_indexes"], update["field_schema"], ready=update["ready"]
+                )
                 for result in loaded
                 for update in result["updates"]
             )
diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py
index ca8ad853..a8d56bb5 100644
--- a/transfer_queue/storage/simple_storage.py
+++ b/transfer_queue/storage/simple_storage.py
@@ -235,11 +235,12 @@ def apply_update(
         global_indexes: list[int],
         new_data: dict[str, Any],
         merge_fn: Callable[[Any, Any], Any],
-        field_schema: dict[str, dict[str, Any]],
     ) -> dict[str, dict[str, Any]]:
-        """Merge produced values, validate their types, then write once.
+        """Merge into stored values, store new ones as is, validate, then write once.
 
-        All merge calls and type checks finish before one ``put_data`` call, so
+        A row that has no value for a field takes ``new`` without calling ``merge_fn``,
+        as a put would. A merge result must keep the old value's tensor/non-tensor kind
+        and tensor dtype. All merges and checks finish before one ``put_data`` call, so
         a merge error or incompatible result leaves storage unchanged.
 
         Returns:
@@ -248,23 +249,21 @@ def apply_update(
         computed: dict[str, list] = {}
         for field, new_values in new_data.items():
             stored = self.field_data.get(field, {})
-            missing = [idx for idx in global_indexes if idx not in stored]
-            if missing:
-                raise ValueError(f"Cannot update unproduced field {field!r} for indexes {missing[:20]}")
-            old_values = self.get_data([field], global_indexes)[field]
-            computed[field] = [merge_fn(old, new_values[i]) for i, old in enumerate(old_values)]
-
-            expected = field_schema[field]
-            if expected.get("is_non_tensor"):
-                if any(isinstance(value, torch.Tensor) for value in computed[field]):
-                    raise TypeError(f"Merge result for non-tensor field {field!r} must remain non-tensor")
-            else:
-                for value in computed[field]:
-                    if not isinstance(value, torch.Tensor) or value.dtype != expected["dtype"]:
-                        actual = value.dtype if isinstance(value, torch.Tensor) else type(value).__name__
+            present = [i for i, idx in enumerate(global_indexes) if idx in stored]
+            old_values = self.get_data([field], [global_indexes[i] for i in present])[field] if present else []
+            values = list(new_values)
+            for i, old in zip(present, old_values, strict=True):
+                merged = merge_fn(old, new_values[i])
+                if isinstance(old, torch.Tensor):
+                    if not isinstance(merged, torch.Tensor) or merged.dtype != old.dtype:
+                        actual = merged.dtype if isinstance(merged, torch.Tensor) else type(merged).__name__
                         raise TypeError(
-                            f"Merge result for tensor field {field!r} must keep dtype {expected['dtype']}, got {actual}"
+                            f"Merge result for tensor field {field!r} must keep dtype {old.dtype}, got {actual}"
                         )
+                elif isinstance(merged, torch.Tensor):
+                    raise TypeError(f"Merge result for non-tensor field {field!r} must remain non-tensor")
+                values[i] = merged
+            computed[field] = values
         self.put_data(computed, global_indexes)
         return {field: _describe_stored_values(values) for field, values in computed.items()}
 
@@ -1243,12 +1242,11 @@ def _handle_put(self, data_parts: ZMQMessage) -> ZMQMessage:
             )
 
     def _handle_update(self, data_parts: ZMQMessage) -> ZMQMessage:
-        """Merge existing fields, write once, and describe what was stored."""
+        """Merge or add fields, write once, and describe what was stored."""
         try:
             global_indexes = data_parts.body["global_indexes"]
             merge_fn = data_parts.body.get("merge_fn")
             new_data = data_parts.body.get("data")
-            field_schema = data_parts.body.get("field_schema")
 
             with limit_pytorch_auto_parallel_threads(
                 target_num_threads=TQ_NUM_THREADS, info=f"[{self.storage_unit_id}] _handle_update"
@@ -1257,9 +1255,7 @@ def _handle_update(self, data_parts: ZMQMessage) -> ZMQMessage:
                     raise TypeError(f"merge_fn must be callable, got {type(merge_fn).__name__}")
                 if not isinstance(new_data, dict) or not new_data:
                     raise TypeError("update data must be a non-empty dict")
-                if not isinstance(field_schema, dict) or set(field_schema) != set(new_data):
-                    raise ValueError("field_schema must describe every update field")
-                stored_shapes = self.storage_data.apply_update(global_indexes, new_data, merge_fn, field_schema)
+                stored_shapes = self.storage_data.apply_update(global_indexes, new_data, merge_fn)
 
             return ZMQMessage.create(
                 request_type=ZMQRequestType.UPDATE_DATA_RESPONSE,  # type: ignore[arg-type]
@@ -1511,19 +1507,23 @@ def _handle_load_rows(self, request: ZMQMessage) -> ZMQMessage:
                                     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))
+                                groups[(tuple(row["fields"]), tuple(row["unfinalized"]))].append((row, fields))
                                 bytes_read += row["length"]
-                            for signature, rows in groups.items():
+                            for (signature, unfinalized), 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})
+                                source_indexes = [row["source_index"] for row, _ in rows]
+                                finalized = tuple(name for name in signature if name not in unfinalized)
+                                # Publish each field with the readiness it had when dumped.
+                                for names, ready in ((finalized, True), (unfinalized, False)):
+                                    if names:
+                                        schema = select_dump_schema(
+                                            shard["field_schema"], source_indexes, indexes, names
+                                        )
+                                        updates.append(
+                                            {"global_indexes": indexes, "field_schema": schema, "ready": ready}
+                                        )
             logger.info(
                 "[%s]: loaded %s rows (%s bytes)",
                 self.storage_unit_id,
diff --git a/tutorial/02_kv_interface.py b/tutorial/02_kv_interface.py
index 999a8193..ef9e62b4 100644
--- a/tutorial/02_kv_interface.py
+++ b/tutorial/02_kv_interface.py
@@ -215,7 +215,7 @@ def main():
         Key Methods:
         1. (async_)kv_put          - Insert/Update a multi-column sample by key, with optional metadata tag
         2. (async_)kv_batch_put    - Put multiple key-value pairs efficiently in batch
-        3. (async_)kv_update       - Merge new values into produced fields with merge_fn(old, new)
+        3. (async_)kv_update       - Merge into fields with merge_fn(old, new), adding missing keys/fields
         4. (async_)kv_batch_get    - Retrieve samples (by keys), supporting column selection (by fields)
         5. (async_)kv_list         - List keys and tags (metadata) in a partition
         6. (async_)kv_clear        - Remove key-value pairs, or only selected fields, from storage

From 3891aa47d18761d1fa9295da32063f5b915ef2cd Mon Sep 17 00:00:00 2001
From: Azency <azency_1128@outlook.com>
Date: Tue, 13 Oct 2026 20:35:19 +0800
Subject: [PATCH 02/12] [fix] kv_update: type-check new rows, surface
 unfinalized fields, bump dump format

- Storage units check rows kv_update adds against the controller's column type,
  so a dtype conflict fails before anything is stored or marked ready.
- kv_batch_get raises when a selected field is not ready in every key instead of
  returning an empty result.
- Add kv_describe / async_kv_describe to list each key's unfinalized fields.
- Dump format version 4 records unfinalized fields; version 3 dumps still load.

Signed-off-by: Azency <azency_1128@outlook.com>
---
 README.md                                     |  1 +
 docs/data_dump.md                             |  5 +-
 tests/e2e/test_data_dump_e2e.py               | 19 +++++++
 tests/e2e/test_kv_interface_e2e.py            | 33 ++++++++++-
 tests/test_kv_update.py                       | 32 ++++++++---
 transfer_queue/__init__.py                    |  4 ++
 transfer_queue/client.py                      |  7 ++-
 transfer_queue/controller.py                  |  7 +++
 transfer_queue/data_dump.py                   | 10 ++--
 transfer_queue/interface.py                   | 57 ++++++++++++++++++-
 .../managers/simple_storage_manager.py        |  9 +++
 transfer_queue/storage/simple_storage.py      | 39 ++++++++-----
 12 files changed, 194 insertions(+), 29 deletions(-)

diff --git a/README.md b/README.md
index 23367a11..242b577c 100644
--- a/README.md
+++ b/README.md
@@ -118,6 +118,7 @@ To simplify the usage of TransferQueue, we provide a Redis-style high-level API
 - **(async_)kv_batch_get**: Retrieve samples (by keys), supporting column selection (by fields).
 - **(async_)kv_update / (async_)kv_batch_update**: Merge new values into fields with `merge_fn(old, new)`; a superset of put that adds missing keys and fields. SimpleStorage only.
 - **(async_)kv_list**: List keys and tags (metadata) in a partition.
+- **(async_)kv_describe**: Show each key's stored fields, including those `kv_update(finalize=False)` keeps hidden.
 - **(async_)kv_clear**: Remove key-value pairs from storage, or only selected `fields` while keeping the keys.
 
 `kv_update` stores a field the key does not hold yet as given, without calling
diff --git a/docs/data_dump.md b/docs/data_dump.md
index f3829b5b..4365ff63 100644
--- a/docs/data_dump.md
+++ b/docs/data_dump.md
@@ -55,7 +55,10 @@ These count application reads, not filesystem read-ahead or physical disk traffi
 
 ## Format and compatibility
 
-Dumps use `format_version: 3`, the only version this build reads:
+Dumps use `format_version: 4`. Version 4 adds each row's `unfinalized` fields, which
+restore keeps hidden. This build also reads version 3 dumps, whose fields restore as
+ready. Builds that only read version 3 reject version 4 dumps instead of exposing
+hidden fields. The layout:
 
 ```text
 dump_info.json
diff --git a/tests/e2e/test_data_dump_e2e.py b/tests/e2e/test_data_dump_e2e.py
index 07fd56f8..c60b50bf 100644
--- a/tests/e2e/test_data_dump_e2e.py
+++ b/tests/e2e/test_data_dump_e2e.py
@@ -272,6 +272,25 @@ def concat(old, new):
         retrieved = tq.kv_batch_get(keys="u1", partition_id=partition_id, select_fields="response")
         _assert_rows_equal(retrieved["response"], [torch.tensor([1, 2, 3])])
 
+    def test_restores_a_version_3_dump_as_ready(self, tq_system, dump_dir, controller):
+        partition_id = "d_version3"
+        _put_rows(partition_id, ["w0"])
+        tq.dump_data_by_key(dump_dir, ["w0"], partition_id)
+
+        # Rewrite it as version 3 wrote it: no ``unfinalized`` in the row index.
+        row_index = _read_row_index(dump_dir)
+        del row_index["rows"]["w0"]["unfinalized"]
+        torch.save(row_index, dump_dir / "row_index.pt")
+        info_path = dump_dir / "dump_info.json"
+        info = json.loads(info_path.read_text(encoding="utf-8"))
+        info["format_version"] = 3
+        info_path.write_text(json.dumps(info), encoding="utf-8")
+
+        ray.get(controller.clear_partition.remote(partition_id))
+        tq.load_data_by_key(dump_dir)
+        retrieved = tq.kv_batch_get(keys="w0", partition_id=partition_id, select_fields="input_ids")
+        _assert_rows_equal(retrieved["input_ids"], [_row_input_ids(0)])
+
     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
diff --git a/tests/e2e/test_kv_interface_e2e.py b/tests/e2e/test_kv_interface_e2e.py
index ed88a995..69c85036 100644
--- a/tests/e2e/test_kv_interface_e2e.py
+++ b/tests/e2e/test_kv_interface_e2e.py
@@ -1262,12 +1262,21 @@ def tokens_ready():
         tq_api.kv_put(key=key, partition_id=partition_id, fields={"prompt": torch.tensor([7])})
         for chunk in ([1], [2]):
             fields = {"tokens": torch.tensor(chunk)}
-            tq_api.kv_update(key=key, partition_id=partition_id, fields=fields, merge_fn=concat, finalize=False)
+            meta = tq_api.kv_update(key=key, partition_id=partition_id, fields=fields, merge_fn=concat, finalize=False)
             assert not tokens_ready()
         assert set(tq_api.kv_batch_get(keys=key, partition_id=partition_id).keys()) == {"prompt"}
+        with pytest.raises(ValueError, match="not finalized"):
+            tq_api.kv_batch_get(keys=key, partition_id=partition_id, select_fields="tokens")
+        with pytest.raises(ValueError, match="not finalized"):
+            tq_api.kv_batch_get_by_meta(meta)
+        row = tq_api.kv_describe(keys=key, partition_id=partition_id)[key]
+        assert (row["fields"], row["unfinalized"]) == (["prompt", "tokens"], ["tokens"])
+        with pytest.raises(RuntimeError, match="not found"):
+            tq_api.kv_describe(keys=[key, "absent"], partition_id=partition_id)
 
         tq_api.kv_update(key=key, partition_id=partition_id, fields={"tokens": torch.tensor([3])}, merge_fn=concat)
         assert tokens_ready()
+        assert tq_api.kv_describe(keys=key, partition_id=partition_id)[key]["unfinalized"] == []
         data = tq_api.kv_batch_get(keys=key, partition_id=partition_id, select_fields="tokens")
         assert_tensor_equal(data["tokens"][0], torch.tensor([1, 2, 3]))
 
@@ -1276,6 +1285,7 @@ def tokens_ready():
             key=key, partition_id=partition_id, fields={"tokens": torch.tensor([4])}, merge_fn=concat, finalize=False
         )
         assert not tokens_ready()
+        tq_api.kv_clear(keys=key, partition_id=partition_id)
 
     def test_creates_missing_key_and_field_and_rejects_incompatible_result(self, tq_api, backend_name):
         if backend_name != "SimpleStorage":
@@ -1306,6 +1316,27 @@ def fail(_old, _new):
         unchanged = tq_api.kv_batch_get(keys="sample", partition_id="test_partition", select_fields="tokens")
         assert_tensor_equal(unchanged["tokens"][0], torch.tensor([1, 2]))
 
+    def test_new_row_must_match_column_type_held_by_other_keys(self, tq_api, backend_name):
+        if backend_name != "SimpleStorage":
+            pytest.skip("merge-backed kv_update is implemented only for SimpleStorage")
+
+        def fail(_old, _new):
+            raise AssertionError("merge_fn must not run for a missing value")
+
+        partition_id = "test_partition"
+        # The column is known only through a hidden write on another key.
+        tq_api.kv_update(
+            key="k1", partition_id=partition_id, fields={"reward": torch.tensor([0.5])}, merge_fn=fail, finalize=False
+        )
+        with pytest.raises(RuntimeError, match="keep dtype"):
+            tq_api.kv_update(key="k2", partition_id=partition_id, fields={"reward": torch.tensor([1])}, merge_fn=fail)
+
+        # Nothing was stored for k2, so a compatible value is still a first write.
+        tq_api.kv_update(key="k2", partition_id=partition_id, fields={"reward": torch.tensor([0.25])}, merge_fn=fail)
+        data = tq_api.kv_batch_get(keys="k2", partition_id=partition_id, select_fields="reward")
+        assert_tensor_equal(data["reward"][0], torch.tensor([0.25]))
+        tq_api.kv_clear(keys=["k1", "k2"], partition_id=partition_id)
+
     def test_kv_backend_rejects(self, tq_api, backend_name):
         if backend_name == "SimpleStorage":
             pytest.skip("rejection is the KV-backend contract")
diff --git a/tests/test_kv_update.py b/tests/test_kv_update.py
index 7a97e0cd..09574289 100644
--- a/tests/test_kv_update.py
+++ b/tests/test_kv_update.py
@@ -53,7 +53,7 @@ def test_apply_update_concatenates_and_keeps_field_name():
     data = StorageUnitData()
     data.put_data({"sequence_ids": [torch.tensor([10, 11, 12])]}, [0])
 
-    described = data.apply_update([0], {"sequence_ids": [torch.tensor([20, 21])]}, _concat)
+    described = data.apply_update([0], {"sequence_ids": [torch.tensor([20, 21])]}, _concat, {})
 
     assert described["sequence_ids"] == {"dtype": torch.int64, "shapes": [(5,)]}
     assert torch.equal(data.field_data["sequence_ids"][0], torch.tensor([10, 11, 12, 20, 21]))
@@ -69,7 +69,7 @@ def fail(_old, _new):
         raise RuntimeError("merge failed")
 
     with pytest.raises(RuntimeError, match="merge failed"):
-        data.apply_update([7], {"tokens": [torch.tensor([9])]}, fail)
+        data.apply_update([7], {"tokens": [torch.tensor([9])]}, fail, {})
     assert torch.equal(data.field_data["tokens"][7], original)
 
 
@@ -83,7 +83,7 @@ def concat(old, new):
         return torch.cat([old, new])
 
     described = data.apply_update(
-        [3, 4], {"tokens": [torch.tensor([2]), torch.tensor([5])], "fresh": [{"a": 1}, {"a": 2}]}, concat
+        [3, 4], {"tokens": [torch.tensor([2]), torch.tensor([5])], "fresh": [{"a": 1}, {"a": 2}]}, concat, {}
     )
 
     assert len(merged) == 1
@@ -106,7 +106,7 @@ def test_apply_update_rejects_incompatible_tensor_results_before_writing(merge_f
     data.put_data({"tokens": [original.clone()]}, [0])
 
     with pytest.raises(TypeError, match=error):
-        data.apply_update([0], {"tokens": [torch.tensor([3])]}, merge_fn)
+        data.apply_update([0], {"tokens": [torch.tensor([3])]}, merge_fn, {})
     assert torch.equal(data.field_data["tokens"][0], original)
 
 
@@ -115,10 +115,28 @@ def test_apply_update_rejects_tensor_result_for_non_tensor_field():
     data.put_data({"meta": [{"step": 1}]}, [0])
 
     with pytest.raises(TypeError, match="must remain non-tensor"):
-        data.apply_update([0], {"meta": [{"step": 2}]}, lambda _old, _new: torch.tensor([2]))
+        data.apply_update([0], {"meta": [{"step": 2}]}, lambda _old, _new: torch.tensor([2]), {})
     assert data.field_data["meta"][0] == {"step": 1}
 
 
+@pytest.mark.parametrize(
+    "column, value, error",
+    [
+        ({"dtype": torch.float32, "is_non_tensor": False}, torch.tensor([1]), "keep dtype"),
+        ({"dtype": torch.float32, "is_non_tensor": False}, {"not": "a tensor"}, "keep dtype"),
+        ({"dtype": None, "is_non_tensor": True}, torch.tensor([1.0]), "must remain non-tensor"),
+    ],
+)
+def test_apply_update_checks_new_rows_against_column_type_before_writing(column, value, error):
+    data = StorageUnitData()
+    data.put_data({"score": [torch.tensor([0.5])]}, [0])
+
+    with pytest.raises(TypeError, match=error):
+        data.apply_update([0, 1], {"score": [torch.tensor([0.25]), value]}, _concat, {"score": column})
+    assert set(data.field_data["score"]) == {0}
+    assert torch.equal(data.field_data["score"][0], torch.tensor([0.5]))
+
+
 def test_apply_update_decodes_ssd_offloaded_old_value(tmp_path):
     data = HybridStorageUnitData(
         storage_size=4, threshold_bytes=64, ssd_path=str(tmp_path), run_id="run", unit_id="unit"
@@ -127,7 +145,7 @@ def test_apply_update_decodes_ssd_offloaded_old_value(tmp_path):
     data.put_data({"tokens": [prompt]}, [0])
     assert data.ssd_active_values == 1
 
-    data.apply_update([0], {"tokens": [torch.tensor([99])]}, _concat)
+    data.apply_update([0], {"tokens": [torch.tensor([99])]}, _concat, {})
 
     assert torch.equal(data.get_data(["tokens"], [0])["tokens"][0], torch.cat([prompt, torch.tensor([99])]))
     assert data.ssd_active_values == 1
@@ -153,7 +171,7 @@ def fail_second_field(encoded_values):
         monkeypatch.setattr(data._ssd_store, "write_values", fail_second_field)
         new_data = {"tokens": [torch.tensor([99])], "mask": [torch.tensor([1])]}
         with pytest.raises(OSError, match="No space left"):
-            data.apply_update([0], new_data, _concat)
+            data.apply_update([0], new_data, _concat, {})
 
         stored = data.get_data(["tokens", "mask"], [0])
         assert torch.equal(stored["tokens"][0], prompt)
diff --git a/transfer_queue/__init__.py b/transfer_queue/__init__.py
index 6c514b2b..05e63cdf 100644
--- a/transfer_queue/__init__.py
+++ b/transfer_queue/__init__.py
@@ -23,6 +23,7 @@
     async_kv_batch_put,
     async_kv_batch_update,
     async_kv_clear,
+    async_kv_describe,
     async_kv_list,
     async_kv_put,
     async_kv_update,
@@ -36,6 +37,7 @@
     kv_batch_put,
     kv_batch_update,
     kv_clear,
+    kv_describe,
     kv_list,
     kv_put,
     kv_update,
@@ -72,6 +74,7 @@
         "kv_batch_get",
         "kv_batch_get_by_meta",
         "kv_list",
+        "kv_describe",
         "kv_clear",
         "kv_update",
         "async_kv_put",
@@ -80,6 +83,7 @@
         "async_kv_batch_get",
         "async_kv_batch_get_by_meta",
         "async_kv_list",
+        "async_kv_describe",
         "async_kv_clear",
         "async_kv_update",
         "KVBatchMeta",
diff --git a/transfer_queue/client.py b/transfer_queue/client.py
index b7ee8d27..85faa74f 100644
--- a/transfer_queue/client.py
+++ b/transfer_queue/client.py
@@ -972,6 +972,7 @@ async def async_kv_retrieve_meta(
         keys: list[str] | str,
         partition_id: str,
         create: bool = False,
+        include_fields: list[str] | None = None,
         socket: zmq.asyncio.Socket | None = None,
     ) -> BatchMeta:
         """Asynchronously retrieve BatchMeta by user-defined keys.
@@ -983,6 +984,7 @@ async def async_kv_retrieve_meta(
             keys: List of keys to retrieve.
             partition_id: The ID of the logical partition to search for keys.
             create:  If True, automatically create entries for missing keys.
+            include_fields: Known fields to include even if not every key holds them ready.
             socket: ZMQ socket injected by @with_controller_socket.
 
         Returns:
@@ -1008,6 +1010,7 @@ async def async_kv_retrieve_meta(
                     "keys": keys,
                     "partition_id": partition_id,
                     "create": create,
+                    "include_fields": include_fields,
                 },
             )
             return response_msg.body.get("metadata", BatchMeta.empty())
@@ -1137,8 +1140,8 @@ async def async_kv_describe(
         key raises instead of being dropped.
 
         Returns:
-            ``{"rows": {key: {"global_index", "fields", "tag"}}, "field_schema": {field:
-            {"is_nested", "is_non_tensor"}}}``.
+            ``{"rows": {key: {"global_index", "fields", "unfinalized", "tag"}}, "field_schema":
+            {field: {"is_nested", "is_non_tensor"}}}``.
         """
         response = await self._request_controller(
             socket=socket,
diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py
index ba8199ff..36c34302 100644
--- a/transfer_queue/controller.py
+++ b/transfer_queue/controller.py
@@ -1652,6 +1652,7 @@ def kv_retrieve_meta(
         keys: list[str],
         partition_id: str,
         create: bool = False,
+        include_fields: list[str] | None = None,
     ) -> BatchMeta:
         """
         Retrieve BatchMeta from the controller using a list of keys.
@@ -1661,6 +1662,8 @@ def kv_retrieve_meta(
             keys: List of keys to retrieve from the controller
             partition_id: Partition id to retrieve from the controller
             create: Whether to register new keys if not found.
+            include_fields: Fields whose column schema to include even if not every key holds
+                them ready, so kv_update can type-check the values it adds to them.
 
         Returns:
             metadata: BatchMeta of the requested keys
@@ -1721,6 +1724,9 @@ def kv_retrieve_meta(
         for field_name, col_idx in partition.field_name_mapping.items():
             if col_idx < len(col_mask) and col_mask[col_idx]:
                 data_fields.append(field_name)
+        for field_name in include_fields or []:
+            if field_name in partition.field_metadata and field_name not in data_fields:
+                data_fields.append(field_name)
 
         return self.generate_batch_meta(partition_id, verified_global_indexes, data_fields, mode="force_fetch")
 
@@ -2241,6 +2247,7 @@ def _handle_kv_retrieve_meta_request(self, request_msg: ZMQMessage) -> ZMQMessag
             keys=params["keys"],
             partition_id=params["partition_id"],
             create=params["create"],
+            include_fields=params.get("include_fields"),
         )
         return self._make_response(
             request_msg,
diff --git a/transfer_queue/data_dump.py b/transfer_queue/data_dump.py
index 9d30d5b7..ad3791fb 100644
--- a/transfer_queue/data_dump.py
+++ b/transfer_queue/data_dump.py
@@ -42,7 +42,10 @@
 
 logger = get_logger(__name__)
 
-DUMP_FORMAT_VERSION = 3
+DUMP_FORMAT_VERSION = 4
+# Version 3 predates ``unfinalized``; every field it holds was ready. A version-3 reader
+# rejects version 4 instead of restoring hidden fields as ready.
+_READABLE_FORMAT_VERSIONS = (3, DUMP_FORMAT_VERSION)
 SHARD_SUBDIR = "shards"
 
 _DUMP_INFO_FILE = "dump_info.json"
@@ -171,10 +174,10 @@ def read_dump(dump_dir: Path) -> tuple[dict[str, Any], list[dict[str, Any]]]:
         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:
+    if dump_info["format_version"] not in _READABLE_FORMAT_VERSIONS:
         raise ValueError(
             f"Unsupported dump format version {dump_info['format_version']} in {dump_dir}; "
-            f"this build reads version {DUMP_FORMAT_VERSION}"
+            f"this build reads versions {list(_READABLE_FORMAT_VERSIONS)}"
         )
 
     row_index = _read_row_index(dump_dir)
@@ -215,7 +218,6 @@ def read_dump(dump_dir: Path) -> tuple[dict[str, Any], list[dict[str, Any]]]:
                     "key": key,
                     "source_index": source_index,
                     "fields": rows[key]["fields"],
-                    # Older dumps lack this key; every field they hold was ready.
                     "unfinalized": rows[key].get("unfinalized", []),
                     "offset": offset,
                     "length": length,
diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py
index 74f257b9..10356651 100644
--- a/transfer_queue/interface.py
+++ b/transfer_queue/interface.py
@@ -539,6 +539,7 @@ def kv_batch_get(keys: list[str] | str, partition_id: str, select_fields: list[s
     Raises:
         ValueError: If keys or partition are not found
         ValueError: If empty fields exist in any key (sample)
+        ValueError: If a field in ``select_fields`` is not ready in every found key
 
     Example:
         >>> import transfer_queue as tq
@@ -596,6 +597,21 @@ def kv_list(partition_id: str | None = None) -> dict[str, dict[str, Any]]:
     return tq_client._run_coroutine(async_kv_list(partition_id=partition_id))
 
 
+def kv_describe(keys: list[str] | str, partition_id: str) -> dict[str, dict[str, Any]]:
+    """Describe which fields each key holds, including unfinalized ones.
+
+    See ``async_kv_describe``.
+
+    Example:
+        >>> import transfer_queue as tq
+        >>> tq.init()
+        >>> rows = tq.kv_describe(keys=["sample_1", "sample_2"], partition_id="train")
+        >>> hidden = {key: row["unfinalized"] for key, row in rows.items() if row["unfinalized"]}
+    """
+    tq_client = _maybe_create_tq_client()
+    return tq_client._run_coroutine(async_kv_describe(keys=keys, partition_id=partition_id))
+
+
 def kv_clear(keys: list[str] | str, partition_id: str, fields: list[str] | str | None = None) -> None:
     """Clear key-value pairs from TransferQueue.
 
@@ -964,6 +980,7 @@ async def async_kv_batch_get(
     Raises:
         ValueError: If keys or partition are not found
         ValueError: If empty fields exist in any key (sample)
+        ValueError: If a field in ``select_fields`` is not ready in every found key
 
     Example:
         >>> import transfer_queue as tq
@@ -989,6 +1006,13 @@ async def async_kv_batch_get(
             fields_to_fetch = [select_fields]
         else:
             fields_to_fetch = select_fields
+        # The meta keeps only fields ready in every key; select_fields would drop the rest silently.
+        not_ready = [f for f in fields_to_fetch if f not in batch_meta.field_names]
+        if not_ready:
+            raise ValueError(
+                f"Fields {not_ready} are not ready in all the requested keys: missing, or written by "
+                "kv_update(finalize=False) and not finalized yet"
+            )
         batch_meta = batch_meta.select_fields(fields_to_fetch)
 
     if not batch_meta.is_ready:
@@ -1040,6 +1064,35 @@ async def async_kv_list(partition_id: str | None = None) -> dict[str, dict[str,
     return partition_info
 
 
+async def async_kv_describe(keys: list[str] | str, partition_id: str) -> dict[str, dict[str, Any]]:
+    """Asynchronously describe which fields each key holds, including unfinalized ones.
+
+    ``kv_batch_get`` returns only fields that are ready in every key, so this is the way
+    to find fields written by ``kv_update(finalize=False)`` and not yet finalized.
+
+    Args:
+        keys: Single key or list of keys to describe.
+        partition_id: Partition containing the keys.
+
+    Returns:
+        ``{key: {"global_index", "fields", "unfinalized", "tag"}}``. ``fields`` lists every
+        stored field of the key and ``unfinalized`` the subset still hidden from consumers.
+
+    Raises:
+        RuntimeError: If the partition or any key does not exist.
+
+    Example:
+        >>> import transfer_queue as tq
+        >>> tq.init()
+        >>> rows = await tq.async_kv_describe(keys="sample_1", partition_id="train")
+        >>> print(rows["sample_1"]["unfinalized"])
+    """
+    keys = [keys] if isinstance(keys, str) else keys
+    tq_client = _maybe_create_tq_client()
+    description = await tq_client.async_kv_describe(list(dict.fromkeys(keys)), partition_id)
+    return description["rows"]
+
+
 async def async_kv_clear(keys: list[str] | str, partition_id: str, fields: list[str] | str | None = None) -> None:
     """Asynchronously clear key-value pairs from TransferQueue.
 
@@ -1135,7 +1188,9 @@ async def async_kv_batch_update(
             f"kv_update is not supported by {type(tq_client.storage_manager).__name__}; "
             "it requires the SimpleStorage backend"
         )
-    batch_meta = await tq_client.async_kv_retrieve_meta(keys=keys, partition_id=partition_id, create=True)
+    batch_meta = await tq_client.async_kv_retrieve_meta(
+        keys=keys, partition_id=partition_id, create=True, include_fields=list(fields.keys())
+    )
 
     if batch_meta.size != len(keys):
         raise RuntimeError(f"Retrieved BatchMeta size {batch_meta.size} does not match {len(keys)} keys")
diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py
index 1f90fbfe..681125a5 100644
--- a/transfer_queue/storage/managers/simple_storage_manager.py
+++ b/transfer_queue/storage/managers/simple_storage_manager.py
@@ -560,6 +560,12 @@ async def update_data(
                 f"Batch size of values ({values.batch_size[0]}) does not match metadata size ({metadata.size})"
             )
         field_names = list(values.keys())
+        # Fields new to the partition have no column type yet; the rest must keep theirs.
+        column_types = {
+            name: {key: metadata.field_schema[name][key] for key in ("dtype", "is_non_tensor")}
+            for name in field_names
+            if name in metadata.field_schema
+        }
 
         routing = self._group_by_hash(metadata.global_indexes)
         tasks = []
@@ -575,6 +581,7 @@ async def update_data(
                         group.global_indexes,
                         storage_data,
                         merge_fn,
+                        column_types,
                         target_storage_unit=su_id,
                     ),
                     # A merge reads stored state, so a replay after a lost answer
@@ -675,6 +682,7 @@ async def _update_to_single_storage_unit(
         global_indexes: list[int],
         storage_data: dict[str, Any],
         merge_fn: Callable[[Any, Any], Any],
+        column_types: dict[str, dict[str, Any]],
         target_storage_unit: str,
         socket: zmq.Socket = None,
     ) -> dict[str, dict[str, Any]]:
@@ -683,6 +691,7 @@ async def _update_to_single_storage_unit(
             "global_indexes": global_indexes,
             "data": storage_data,
             "merge_fn": merge_fn,
+            "column_types": column_types,
         }
 
         request_msg = ZMQMessage.create(
diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py
index a8d56bb5..abc3b3a9 100644
--- a/transfer_queue/storage/simple_storage.py
+++ b/transfer_queue/storage/simple_storage.py
@@ -112,6 +112,16 @@ class _SSDValueRef:
 TQ_ACCEPT_PROBE_INTERVAL = float(os.environ.get("TQ_ACCEPT_PROBE_INTERVAL", 0))
 
 
+def _check_value_type(field: str, value: Any, is_tensor: bool, dtype: Any) -> None:
+    """Reject a value that would change its field's tensor/non-tensor kind or tensor dtype."""
+    if not is_tensor:
+        if isinstance(value, torch.Tensor):
+            raise TypeError(f"Value for non-tensor field {field!r} must remain non-tensor")
+    elif not isinstance(value, torch.Tensor) or (dtype is not None and value.dtype != dtype):
+        actual = value.dtype if isinstance(value, torch.Tensor) else type(value).__name__
+        raise TypeError(f"Value for tensor field {field!r} must keep dtype {dtype}, got {actual}")
+
+
 def _describe_stored_values(values: list) -> dict[str, Any]:
     """Describe per-sample stored values for the manager.
 
@@ -235,13 +245,15 @@ def apply_update(
         global_indexes: list[int],
         new_data: dict[str, Any],
         merge_fn: Callable[[Any, Any], Any],
+        column_types: dict[str, dict[str, Any]],
     ) -> dict[str, dict[str, Any]]:
         """Merge into stored values, store new ones as is, validate, then write once.
 
         A row that has no value for a field takes ``new`` without calling ``merge_fn``,
         as a put would. A merge result must keep the old value's tensor/non-tensor kind
-        and tensor dtype. All merges and checks finish before one ``put_data`` call, so
-        a merge error or incompatible result leaves storage unchanged.
+        and tensor dtype; a new row must match ``column_types``, the controller's column
+        type for fields it already knows. All merges and checks finish before one
+        ``put_data`` call, so a merge error or incompatible value leaves storage unchanged.
 
         Returns:
             Per-sample description of the stored values, keyed by field name.
@@ -253,16 +265,14 @@ def apply_update(
             old_values = self.get_data([field], [global_indexes[i] for i in present])[field] if present else []
             values = list(new_values)
             for i, old in zip(present, old_values, strict=True):
-                merged = merge_fn(old, new_values[i])
-                if isinstance(old, torch.Tensor):
-                    if not isinstance(merged, torch.Tensor) or merged.dtype != old.dtype:
-                        actual = merged.dtype if isinstance(merged, torch.Tensor) else type(merged).__name__
-                        raise TypeError(
-                            f"Merge result for tensor field {field!r} must keep dtype {old.dtype}, got {actual}"
-                        )
-                elif isinstance(merged, torch.Tensor):
-                    raise TypeError(f"Merge result for non-tensor field {field!r} must remain non-tensor")
-                values[i] = merged
+                values[i] = merge_fn(old, new_values[i])
+                _check_value_type(field, values[i], isinstance(old, torch.Tensor), getattr(old, "dtype", None))
+            column = column_types.get(field)
+            if column is not None:
+                merged = set(present)
+                for i, value in enumerate(values):
+                    if i not in merged:
+                        _check_value_type(field, value, not column["is_non_tensor"], column["dtype"])
             computed[field] = values
         self.put_data(computed, global_indexes)
         return {field: _describe_stored_values(values) for field, values in computed.items()}
@@ -1247,6 +1257,7 @@ def _handle_update(self, data_parts: ZMQMessage) -> ZMQMessage:
             global_indexes = data_parts.body["global_indexes"]
             merge_fn = data_parts.body.get("merge_fn")
             new_data = data_parts.body.get("data")
+            column_types = data_parts.body.get("column_types")
 
             with limit_pytorch_auto_parallel_threads(
                 target_num_threads=TQ_NUM_THREADS, info=f"[{self.storage_unit_id}] _handle_update"
@@ -1255,7 +1266,9 @@ def _handle_update(self, data_parts: ZMQMessage) -> ZMQMessage:
                     raise TypeError(f"merge_fn must be callable, got {type(merge_fn).__name__}")
                 if not isinstance(new_data, dict) or not new_data:
                     raise TypeError("update data must be a non-empty dict")
-                stored_shapes = self.storage_data.apply_update(global_indexes, new_data, merge_fn)
+                if not isinstance(column_types, dict):
+                    raise TypeError("update request must carry column_types")
+                stored_shapes = self.storage_data.apply_update(global_indexes, new_data, merge_fn, column_types)
 
             return ZMQMessage.create(
                 request_type=ZMQRequestType.UPDATE_DATA_RESPONSE,  # type: ignore[arg-type]

From 63a751d204c640de17228586e6b2d438b2c1a1be Mon Sep 17 00:00:00 2001
From: OutstanderWang <wangweiyanster@gmail.com>
Date: Tue, 13 Oct 2026 21:16:28 +0800
Subject: [PATCH 03/12] [fix] kv_batch_update: reject duplicate keys before
 registering them

Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
---
 tests/test_kv_update.py     | 11 ++++++++++-
 transfer_queue/interface.py | 10 ++++++++--
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/tests/test_kv_update.py b/tests/test_kv_update.py
index 09574289..94c6edff 100644
--- a/tests/test_kv_update.py
+++ b/tests/test_kv_update.py
@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
 import errno
 from unittest.mock import MagicMock, patch
 
@@ -21,7 +22,7 @@
 import torch
 from tensordict import TensorDict
 
-from transfer_queue.interface import _single_update_batch
+from transfer_queue.interface import _single_update_batch, async_kv_batch_update
 from transfer_queue.metadata import BatchMeta
 from transfer_queue.storage.managers.base import KVStorageManager
 from transfer_queue.storage.managers.mooncake_manager import MooncakeStorageManager
@@ -49,6 +50,14 @@ def test_single_update_batch_rejects_invalid_fields(fields):
         _single_update_batch(fields)
 
 
+@patch("transfer_queue.interface._maybe_create_tq_client", side_effect=AssertionError("must fail before any RPC"))
+def test_batch_update_rejects_duplicate_keys_before_registering(_client):
+    fields = TensorDict({"tokens": torch.tensor([[1], [2], [3]])}, batch_size=3)
+
+    with pytest.raises(ValueError, match=r"duplicated: \['k'\]"):
+        asyncio.run(async_kv_batch_update(["k", "other", "k"], "p", fields, _concat))
+
+
 def test_apply_update_concatenates_and_keeps_field_name():
     data = StorageUnitData()
     data.put_data({"sequence_ids": [torch.tensor([10, 11, 12])]}, [0])
diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py
index 10356651..c3e40431 100644
--- a/transfer_queue/interface.py
+++ b/transfer_queue/interface.py
@@ -695,8 +695,9 @@ def kv_batch_update(
 ) -> KVBatchMeta:
     """Merge batched values into the fields of keys, adding keys or fields if missing.
 
-    See ``kv_update``. A timeout has an unknown outcome because a unit may commit
-    after the caller stops waiting. Do not blindly retry a merge.
+    See ``kv_update``. ``keys`` must be unique; a repeated key raises ``ValueError``
+    before anything is registered or written. A timeout has an unknown outcome because
+    a unit may commit after the caller stops waiting. Do not blindly retry a merge.
     """
     tq_client = _maybe_create_tq_client()
     return tq_client._run_coroutine(
@@ -1175,6 +1176,11 @@ async def async_kv_batch_update(
     """Asynchronously merge or add batched fields of keys. See ``kv_update``."""
     if not isinstance(keys, list) or not keys or not all(isinstance(key, str) for key in keys):
         raise TypeError("keys must be a non-empty list of strings")
+    # A repeated missing key would get one row per occurrence, orphaning all but the last;
+    # a repeated existing key would merge twice from the same old value, losing one.
+    if len(set(keys)) != len(keys):
+        duplicates = sorted({key for key in keys if keys.count(key) > 1})
+        raise ValueError(f"keys must be unique in one kv_batch_update; duplicated: {duplicates}")
     if not isinstance(fields, TensorDict) or not fields.keys():
         raise TypeError("fields must be a non-empty TensorDict")
     if fields.batch_size != torch.Size([len(keys)]):

From 045bdb0a6d17955b26a23bf02dc69e9b920225b7 Mon Sep 17 00:00:00 2001
From: OutstanderWang <wangweiyanster@gmail.com>
Date: Tue, 13 Oct 2026 21:16:28 +0800
Subject: [PATCH 04/12] style: annotate load-rows updates so mypy accepts mixed
 value types

Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
---
 transfer_queue/storage/simple_storage.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py
index abc3b3a9..76fd6821 100644
--- a/transfer_queue/storage/simple_storage.py
+++ b/transfer_queue/storage/simple_storage.py
@@ -1505,7 +1505,7 @@ def _handle_dump_rows(self, request: ZMQMessage) -> ZMQMessage:
 
     def _handle_load_rows(self, request: ZMQMessage) -> ZMQMessage:
         """Seek to assigned records and merge them into local storage at current indexes."""
-        updates = []
+        updates: list[dict[str, Any]] = []
         bytes_read = 0
         try:
             with limit_pytorch_auto_parallel_threads(TQ_NUM_THREADS):

From bfdbd7af3af7459cbfc8b0698652a3fda260b714 Mon Sep 17 00:00:00 2001
From: OutstanderWang <wangweiyanster@gmail.com>
Date: Wed, 14 Oct 2026 14:11:51 +0800
Subject: [PATCH 05/12] [fix] kv_update: let the controller decide which rows
 to merge

apply_update chose between merging and storing from the unit's local state,
but a unit can keep values the controller already dropped (a clear that
failed in storage only logs; kv_clear then frees the global index for any
partition). kv_update folded new chunks into those leftovers, and a retried
first chunk merged twice on units that had committed.

kv_batch_update now sends one KV_PREPARE_UPDATE request instead of
kv_retrieve_meta: it registers missing keys and returns, per field, the
column type and a `held` mask aligned with the batch, built from
FieldMeta.global_indexes (hidden cells included). Units merge exactly the
held rows and overwrite the rest; a held row missing on the unit raises
before anything is written. Every final value is checked against the
column type, replacing the separate old-value check. include_fields and
the manager-side column_types go away, so kv_retrieve_meta and BatchMeta
match main again.

Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
---
 tests/e2e/test_kv_interface_e2e.py            | 35 +++++++++++
 tests/test_controller.py                      | 23 +++++++
 tests/test_kv_update.py                       | 62 +++++++++++++++----
 tests/test_storage_request_retry.py           |  7 ++-
 transfer_queue/client.py                      | 39 ++++++++++--
 transfer_queue/controller.py                  | 41 +++++++++---
 transfer_queue/interface.py                   |  8 ++-
 transfer_queue/storage/managers/base.py       |  5 +-
 .../managers/simple_storage_manager.py        | 16 +++--
 transfer_queue/storage/simple_storage.py      | 36 +++++------
 transfer_queue/utils/zmq_utils.py             |  2 +
 11 files changed, 214 insertions(+), 60 deletions(-)

diff --git a/tests/e2e/test_kv_interface_e2e.py b/tests/e2e/test_kv_interface_e2e.py
index 69c85036..1e8a6f6d 100644
--- a/tests/e2e/test_kv_interface_e2e.py
+++ b/tests/e2e/test_kv_interface_e2e.py
@@ -1337,6 +1337,41 @@ def fail(_old, _new):
         assert_tensor_equal(data["reward"][0], torch.tensor([0.25]))
         tq_api.kv_clear(keys=["k1", "k2"], partition_id=partition_id)
 
+    def test_storage_values_the_controller_does_not_hold_are_overwritten(self, tq_api, backend_name, monkeypatch):
+        if backend_name != "SimpleStorage":
+            pytest.skip("merge-backed kv_update is implemented only for SimpleStorage")
+
+        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")
+
+        async def reject_notify(*_args, **_kwargs):
+            return False
+
+        partition_id = "test_partition"
+        manager = interface._maybe_create_tq_client().storage_manager
+        fields = {"response": torch.tensor([100, 101]), "prompt": torch.tensor([7])}
+        tq_api.kv_put(key="k", partition_id=partition_id, fields=fields)
+        with monkeypatch.context() as patched:
+            patched.setattr(manager, "_clear_single_storage_unit", fail_clear)
+            tq_api.kv_clear(keys="k", partition_id=partition_id, fields="response")
+        tq_api.kv_update(key="k", partition_id=partition_id, fields={"response": torch.tensor([1])}, merge_fn=concat)
+        data = tq_api.kv_batch_get(keys="k", partition_id=partition_id, select_fields="response")
+        assert_tensor_equal(data["response"][0], torch.tensor([1]))
+
+        # A first chunk the controller never recorded is rewritten on retry, not merged twice.
+        with monkeypatch.context() as patched, pytest.raises(RuntimeError, match="did not accept"):
+            patched.setattr(manager, "notify_data_update", reject_notify)
+            tq_api.kv_update(key="r", partition_id=partition_id, fields={"tokens": torch.tensor([5])}, merge_fn=concat)
+        tq_api.kv_update(key="r", partition_id=partition_id, fields={"tokens": torch.tensor([5])}, merge_fn=concat)
+        data = tq_api.kv_batch_get(keys="r", partition_id=partition_id, select_fields="tokens")
+        assert_tensor_equal(data["tokens"][0], torch.tensor([5]))
+        tq_api.kv_clear(keys=["k", "r"], partition_id=partition_id)
+
     def test_kv_backend_rejects(self, tq_api, backend_name):
         if backend_name == "SimpleStorage":
             pytest.skip("rejection is the KV-backend contract")
diff --git a/tests/test_controller.py b/tests/test_controller.py
index 32c1603a..344505ff 100644
--- a/tests/test_controller.py
+++ b/tests/test_controller.py
@@ -1191,6 +1191,29 @@ def test_controller_kv_retrieve_meta_with_production_status(self, ray_setup):
         # Clean up
         ray.get(tq_controller.clear_partition.remote(partition_id))
 
+    def test_controller_kv_prepare_update_reports_held_rows_from_field_metadata(self, ray_setup):
+        """Hidden cells count as held; new keys and fields the partition lacks do not."""
+        tq_controller = TransferQueueController.remote()
+        partition_id = "kv_prepare_update_test"
+        existing = ray.get(
+            tq_controller.kv_retrieve_meta.remote(keys=["s0", "s1"], partition_id=partition_id, create=True)
+        )
+        s0, s1 = existing.global_indexes
+        response = {"response": {"dtype": torch.int64, "shape": (2,), "is_nested": False, "is_non_tensor": False}}
+        prompt = {"prompt": {"dtype": torch.int64, "shape": (1,), "is_nested": False, "is_non_tensor": False}}
+        ray.get(tq_controller.update_production_status.remote(partition_id, [s0], response, ready=False))
+        ray.get(tq_controller.update_production_status.remote(partition_id, [s1], prompt))
+
+        metadata, columns = ray.get(
+            tq_controller.kv_prepare_update.remote(["s0", "s1", "s2"], partition_id, ["response", "fresh"])
+        )
+
+        assert metadata.global_indexes[:2] == [s0, s1]
+        assert columns["response"] == {"dtype": torch.int64, "is_non_tensor": False, "held": [True, False, False]}
+        assert columns["fresh"] == {"dtype": None, "is_non_tensor": None, "held": [False, False, False]}
+
+        ray.get(tq_controller.clear_partition.remote(partition_id))
+
     def test_controller_kv_retrieve_meta_with_custom_meta(self, ray_setup):
         """Test kv_retrieve_meta preserves custom_meta through retrieve."""
         tq_controller = TransferQueueController.remote()
diff --git a/tests/test_kv_update.py b/tests/test_kv_update.py
index 94c6edff..85328f1d 100644
--- a/tests/test_kv_update.py
+++ b/tests/test_kv_update.py
@@ -29,13 +29,17 @@
 from transfer_queue.storage.managers.ray_storage_manager import RayStorageManager
 from transfer_queue.storage.managers.simple_storage_manager import _build_update_field_schema
 from transfer_queue.storage.managers.yuanrong_manager import YuanrongStorageManager
-from transfer_queue.storage.simple_storage import HybridStorageUnitData, StorageUnitData
+from transfer_queue.storage.simple_storage import HybridStorageUnitData, StorageKeyNotFoundError, StorageUnitData
 
 
 def _concat(old, new):
     return torch.cat([old, new])
 
 
+def _column(*held, dtype=torch.int64, is_non_tensor=False):
+    return {"dtype": dtype, "is_non_tensor": is_non_tensor, "held": list(held)}
+
+
 def test_single_update_batch_wraps_tensor_and_python_values():
     batch = _single_update_batch({"tokens": torch.tensor([4, 5]), "meta": {"step": 1}})
 
@@ -62,7 +66,9 @@ def test_apply_update_concatenates_and_keeps_field_name():
     data = StorageUnitData()
     data.put_data({"sequence_ids": [torch.tensor([10, 11, 12])]}, [0])
 
-    described = data.apply_update([0], {"sequence_ids": [torch.tensor([20, 21])]}, _concat, {})
+    described = data.apply_update(
+        [0], {"sequence_ids": [torch.tensor([20, 21])]}, _concat, {"sequence_ids": _column(True)}
+    )
 
     assert described["sequence_ids"] == {"dtype": torch.int64, "shapes": [(5,)]}
     assert torch.equal(data.field_data["sequence_ids"][0], torch.tensor([10, 11, 12, 20, 21]))
@@ -78,7 +84,7 @@ def fail(_old, _new):
         raise RuntimeError("merge failed")
 
     with pytest.raises(RuntimeError, match="merge failed"):
-        data.apply_update([7], {"tokens": [torch.tensor([9])]}, fail, {})
+        data.apply_update([7], {"tokens": [torch.tensor([9])]}, fail, {"tokens": _column(True)})
     assert torch.equal(data.field_data["tokens"][7], original)
 
 
@@ -92,7 +98,10 @@ def concat(old, new):
         return torch.cat([old, new])
 
     described = data.apply_update(
-        [3, 4], {"tokens": [torch.tensor([2]), torch.tensor([5])], "fresh": [{"a": 1}, {"a": 2}]}, concat, {}
+        [3, 4],
+        {"tokens": [torch.tensor([2]), torch.tensor([5])], "fresh": [{"a": 1}, {"a": 2}]},
+        concat,
+        {"tokens": _column(True, False), "fresh": _column(False, False, dtype=None, is_non_tensor=None)},
     )
 
     assert len(merged) == 1
@@ -115,7 +124,7 @@ def test_apply_update_rejects_incompatible_tensor_results_before_writing(merge_f
     data.put_data({"tokens": [original.clone()]}, [0])
 
     with pytest.raises(TypeError, match=error):
-        data.apply_update([0], {"tokens": [torch.tensor([3])]}, merge_fn, {})
+        data.apply_update([0], {"tokens": [torch.tensor([3])]}, merge_fn, {"tokens": _column(True)})
     assert torch.equal(data.field_data["tokens"][0], original)
 
 
@@ -124,16 +133,21 @@ def test_apply_update_rejects_tensor_result_for_non_tensor_field():
     data.put_data({"meta": [{"step": 1}]}, [0])
 
     with pytest.raises(TypeError, match="must remain non-tensor"):
-        data.apply_update([0], {"meta": [{"step": 2}]}, lambda _old, _new: torch.tensor([2]), {})
+        data.apply_update(
+            [0],
+            {"meta": [{"step": 2}]},
+            lambda _old, _new: torch.tensor([2]),
+            {"meta": _column(True, dtype=None, is_non_tensor=True)},
+        )
     assert data.field_data["meta"][0] == {"step": 1}
 
 
 @pytest.mark.parametrize(
     "column, value, error",
     [
-        ({"dtype": torch.float32, "is_non_tensor": False}, torch.tensor([1]), "keep dtype"),
-        ({"dtype": torch.float32, "is_non_tensor": False}, {"not": "a tensor"}, "keep dtype"),
-        ({"dtype": None, "is_non_tensor": True}, torch.tensor([1.0]), "must remain non-tensor"),
+        (_column(True, False, dtype=torch.float32), torch.tensor([1]), "keep dtype"),
+        (_column(True, False, dtype=torch.float32), {"not": "a tensor"}, "keep dtype"),
+        (_column(False, False, dtype=None, is_non_tensor=True), torch.tensor([1.0]), "must remain non-tensor"),
     ],
 )
 def test_apply_update_checks_new_rows_against_column_type_before_writing(column, value, error):
@@ -146,6 +160,30 @@ def test_apply_update_checks_new_rows_against_column_type_before_writing(column,
     assert torch.equal(data.field_data["score"][0], torch.tensor([0.5]))
 
 
+def test_apply_update_overwrites_a_local_value_the_controller_does_not_hold():
+    data = StorageUnitData()
+    data.put_data({"response": [torch.tensor([100, 101])]}, [0])
+
+    def must_not_merge(_old, _new):
+        raise AssertionError("a value the controller does not hold must not be merged")
+
+    data.apply_update([0], {"response": [torch.tensor([1])]}, must_not_merge, {"response": _column(False)})
+
+    assert torch.equal(data.field_data["response"][0], torch.tensor([1]))
+
+
+def test_apply_update_rejects_held_row_missing_locally_before_writing():
+    data = StorageUnitData()
+    data.put_data({"tokens": [torch.tensor([1])]}, [0])
+
+    with pytest.raises(StorageKeyNotFoundError):
+        data.apply_update(
+            [0, 1], {"tokens": [torch.tensor([2]), torch.tensor([3])]}, _concat, {"tokens": _column(True, True)}
+        )
+    assert set(data.field_data["tokens"]) == {0}
+    assert torch.equal(data.field_data["tokens"][0], torch.tensor([1]))
+
+
 def test_apply_update_decodes_ssd_offloaded_old_value(tmp_path):
     data = HybridStorageUnitData(
         storage_size=4, threshold_bytes=64, ssd_path=str(tmp_path), run_id="run", unit_id="unit"
@@ -154,7 +192,7 @@ def test_apply_update_decodes_ssd_offloaded_old_value(tmp_path):
     data.put_data({"tokens": [prompt]}, [0])
     assert data.ssd_active_values == 1
 
-    data.apply_update([0], {"tokens": [torch.tensor([99])]}, _concat, {})
+    data.apply_update([0], {"tokens": [torch.tensor([99])]}, _concat, {"tokens": _column(True)})
 
     assert torch.equal(data.get_data(["tokens"], [0])["tokens"][0], torch.cat([prompt, torch.tensor([99])]))
     assert data.ssd_active_values == 1
@@ -180,7 +218,7 @@ def fail_second_field(encoded_values):
         monkeypatch.setattr(data._ssd_store, "write_values", fail_second_field)
         new_data = {"tokens": [torch.tensor([99])], "mask": [torch.tensor([1])]}
         with pytest.raises(OSError, match="No space left"):
-            data.apply_update([0], new_data, _concat, {})
+            data.apply_update([0], new_data, _concat, {"tokens": _column(True), "mask": _column(True)})
 
         stored = data.get_data(["tokens", "mask"], [0])
         assert torch.equal(stored["tokens"][0], prompt)
@@ -244,4 +282,4 @@ async def test_kv_backends_reject_update(mock_create, manager_cls):
     values = TensorDict({"x": torch.ones(1, 1, dtype=torch.int64)}, batch_size=1)
 
     with pytest.raises(NotImplementedError, match=f"not supported by {manager_cls.__name__}"):
-        await manager.update_data(meta, values, lambda _old, new: new)
+        await manager.update_data(meta, values, lambda _old, new: new, {"x": _column(True)})
diff --git a/tests/test_storage_request_retry.py b/tests/test_storage_request_retry.py
index 70be8985..ff514992 100644
--- a/tests/test_storage_request_retry.py
+++ b/tests/test_storage_request_retry.py
@@ -107,6 +107,9 @@ async def test_put_retries_only_without_a_parser(data_parser, expected_attempts)
     manager.notify_data_update.assert_not_awaited()
 
 
+_HELD_INPUT_IDS = {"input_ids": {"dtype": torch.int64, "is_non_tensor": False, "held": [True]}}
+
+
 @pytest.mark.asyncio
 async def test_update_reports_controller_rejection_after_storage_commit():
     manager = _manager()
@@ -117,7 +120,7 @@ async def test_update_reports_controller_rejection_after_storage_commit():
     values, metadata = _single_sample_batch()
 
     with pytest.raises(RuntimeError, match="Storage update committed.*do not retry"):
-        await manager.update_data(metadata, values, lambda old, new: torch.cat([old, new]))
+        await manager.update_data(metadata, values, lambda old, new: torch.cat([old, new]), _HELD_INPUT_IDS)
 
     manager._update_to_single_storage_unit.assert_awaited_once()
     manager.notify_data_update.assert_awaited_once()
@@ -131,7 +134,7 @@ async def test_update_timeout_is_not_retried_or_published():
     values, metadata = _single_sample_batch()
 
     with patch.object(manager, "_diagnose_storage_unit", return_value="diagnosis"), pytest.raises(StorageUnitTimeout):
-        await manager.update_data(metadata, values, lambda old, new: torch.cat([old, new]))
+        await manager.update_data(metadata, values, lambda old, new: torch.cat([old, new]), _HELD_INPUT_IDS)
 
     manager._update_to_single_storage_unit.assert_awaited_once()
     manager.notify_data_update.assert_not_awaited()
diff --git a/transfer_queue/client.py b/transfer_queue/client.py
index 85faa74f..f6ae1521 100644
--- a/transfer_queue/client.py
+++ b/transfer_queue/client.py
@@ -497,6 +497,7 @@ async def async_update(
         metadata: BatchMeta,
         values: TensorDict,
         merge_fn: Callable[[Any, Any], Any],
+        columns: dict[str, dict[str, Any]],
         finalize: bool = True,
     ) -> BatchMeta:
         """Merge new values into fields on SimpleStorage units, storing them where absent.
@@ -504,7 +505,8 @@ async def async_update(
         Args:
             metadata: Samples to update, already registered with the controller.
             values: New values aligned with ``metadata``.
-            merge_fn: Called per sample per field as ``merge_fn(old, new)`` when a value exists.
+            merge_fn: Called per sample per field as ``merge_fn(old, new)`` on held rows.
+            columns: Per-field column type and held rows, from ``async_kv_prepare_update``.
             finalize: If False, keep the updated fields hidden from consumers.
 
         Returns:
@@ -521,7 +523,7 @@ async def async_update(
         if not metadata or metadata.size == 0:
             raise ValueError("metadata cannot be none or empty")
 
-        field_schema = await self.storage_manager.update_data(metadata, values, merge_fn, finalize=finalize)
+        field_schema = await self.storage_manager.update_data(metadata, values, merge_fn, columns, finalize=finalize)
         if not finalize:
             metadata.production_status[:] = 0
         return metadata.apply_field_schema(field_schema, set_all_ready=finalize)
@@ -972,7 +974,6 @@ async def async_kv_retrieve_meta(
         keys: list[str] | str,
         partition_id: str,
         create: bool = False,
-        include_fields: list[str] | None = None,
         socket: zmq.asyncio.Socket | None = None,
     ) -> BatchMeta:
         """Asynchronously retrieve BatchMeta by user-defined keys.
@@ -984,7 +985,6 @@ async def async_kv_retrieve_meta(
             keys: List of keys to retrieve.
             partition_id: The ID of the logical partition to search for keys.
             create:  If True, automatically create entries for missing keys.
-            include_fields: Known fields to include even if not every key holds them ready.
             socket: ZMQ socket injected by @with_controller_socket.
 
         Returns:
@@ -1010,13 +1010,29 @@ async def async_kv_retrieve_meta(
                     "keys": keys,
                     "partition_id": partition_id,
                     "create": create,
-                    "include_fields": include_fields,
                 },
             )
             return response_msg.body.get("metadata", BatchMeta.empty())
         except Exception as e:
             raise RuntimeError(f"[{self.client_id}] Failed in async_kv_retrieve_meta: {e}") from e
 
+    @with_controller_socket
+    async def async_kv_prepare_update(
+        self, keys: list[str], partition_id: str, fields: list[str], socket=None
+    ) -> tuple[BatchMeta, dict[str, dict[str, Any]]]:
+        """Register missing keys and get, per field, its column type and which rows hold it.
+
+        Returns:
+            The keys' BatchMeta and the ``columns`` that ``async_update`` takes.
+        """
+        response_msg = await self._request_controller(
+            socket=socket,
+            request_type=ZMQRequestType.KV_PREPARE_UPDATE,
+            response_type=ZMQRequestType.KV_PREPARE_UPDATE_RESPONSE,
+            body={"keys": keys, "partition_id": partition_id, "fields": fields},
+        )
+        return response_msg.body["metadata"], response_msg.body["columns"]
+
     async def async_kv_clear_fields(self, keys: list[str], partition_id: str, fields: list[str]) -> None:
         """Clear fields of existing keys while keeping the keys, tags and consumption status.
 
@@ -1462,6 +1478,7 @@ def wrapper(*args, **kwargs):
         self._set_custom_meta = _make_sync(self.async_set_custom_meta)
         self._reset_consumption = _make_sync(self.async_reset_consumption)
         self._kv_retrieve_meta = _make_sync(self.async_kv_retrieve_meta)
+        self._kv_prepare_update = _make_sync(self.async_kv_prepare_update)
         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)
@@ -1645,13 +1662,14 @@ def update(
         metadata: BatchMeta,
         values: TensorDict,
         merge_fn: Callable[[Any, Any], Any],
+        columns: dict[str, dict[str, Any]],
         finalize: bool = True,
     ) -> BatchMeta:
         """Synchronously merge or add values on SimpleStorage units.
 
         See ``async_update``.
         """
-        return self._update(metadata, values, merge_fn, finalize=finalize)
+        return self._update(metadata, values, merge_fn, columns, finalize=finalize)
 
     def get_data(self, metadata: BatchMeta) -> TensorDict:
         """Synchronously fetch data from storage units and organize into TensorDict.
@@ -1867,6 +1885,15 @@ def kv_retrieve_meta(
         """
         return self._kv_retrieve_meta(keys=keys, partition_id=partition_id, create=create)
 
+    def kv_prepare_update(
+        self, keys: list[str], partition_id: str, fields: list[str]
+    ) -> tuple[BatchMeta, dict[str, dict[str, Any]]]:
+        """Synchronously register missing keys and get the ``columns`` for ``update``.
+
+        See ``async_kv_prepare_update``.
+        """
+        return self._kv_prepare_update(keys, partition_id, fields)
+
     def kv_retrieve_keys(
         self,
         global_indexes: list[int] | int,
diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py
index 36c34302..45a9dc11 100644
--- a/transfer_queue/controller.py
+++ b/transfer_queue/controller.py
@@ -1652,7 +1652,6 @@ def kv_retrieve_meta(
         keys: list[str],
         partition_id: str,
         create: bool = False,
-        include_fields: list[str] | None = None,
     ) -> BatchMeta:
         """
         Retrieve BatchMeta from the controller using a list of keys.
@@ -1662,8 +1661,6 @@ def kv_retrieve_meta(
             keys: List of keys to retrieve from the controller
             partition_id: Partition id to retrieve from the controller
             create: Whether to register new keys if not found.
-            include_fields: Fields whose column schema to include even if not every key holds
-                them ready, so kv_update can type-check the values it adds to them.
 
         Returns:
             metadata: BatchMeta of the requested keys
@@ -1724,12 +1721,35 @@ def kv_retrieve_meta(
         for field_name, col_idx in partition.field_name_mapping.items():
             if col_idx < len(col_mask) and col_mask[col_idx]:
                 data_fields.append(field_name)
-        for field_name in include_fields or []:
-            if field_name in partition.field_metadata and field_name not in data_fields:
-                data_fields.append(field_name)
 
         return self.generate_batch_meta(partition_id, verified_global_indexes, data_fields, mode="force_fetch")
 
+    def kv_prepare_update(
+        self, keys: list[str], partition_id: str, fields: list[str]
+    ) -> tuple[BatchMeta, dict[str, dict[str, Any]]]:
+        """Register missing keys and tell kv_update, per field, which rows to merge into.
+
+        Storage units may keep values the controller no longer tracks (e.g. after a failed
+        clear), so ``held`` comes from ``FieldMeta.global_indexes``, which covers hidden cells.
+
+        Returns:
+            The keys' BatchMeta and ``{field: {"dtype", "is_non_tensor", "held"}}``, where
+            ``held`` is aligned with the BatchMeta rows. The types are None for a new field.
+        """
+        metadata = self.kv_retrieve_meta(keys, partition_id, create=True)
+        partition = self._get_partition(partition_id)
+        assert partition is not None
+        columns = {}
+        for name in fields:
+            meta = partition.field_metadata.get(name)
+            held = meta.global_indexes if meta is not None else set()
+            columns[name] = {
+                "dtype": meta.dtype if meta is not None else None,
+                "is_non_tensor": meta.is_non_tensor if meta is not None else None,
+                "held": [idx in held for idx in metadata.global_indexes],
+            }
+        return metadata, columns
+
     def kv_clear_fields(self, keys: list[str], partition_id: str, fields: list[str]) -> BatchMeta:
         """Clear fields of existing keys and return the cleared cells' metadata for storage deletion.
 
@@ -2062,6 +2082,7 @@ def _handle_request(self, request_msg: ZMQMessage, monitor: Any) -> ZMQMessage |
             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.KV_PREPARE_UPDATE: self._handle_kv_prepare_update_request,
             ZMQRequestType.VALIDATE_DUMP_SCHEMA: self._handle_validate_dump_schema_request,
             ZMQRequestType.KV_CLEAR_FIELDS: self._handle_kv_clear_fields_request,
             ZMQRequestType.SAVE_CONTROLLER_CHECKPOINT: self._handle_save_controller_checkpoint_request,
@@ -2247,7 +2268,6 @@ def _handle_kv_retrieve_meta_request(self, request_msg: ZMQMessage) -> ZMQMessag
             keys=params["keys"],
             partition_id=params["partition_id"],
             create=params["create"],
-            include_fields=params.get("include_fields"),
         )
         return self._make_response(
             request_msg,
@@ -2268,6 +2288,13 @@ def _handle_kv_describe_request(self, request_msg: ZMQMessage) -> ZMQMessage:
         description = self.kv_describe(params["keys"], params["partition_id"])
         return self._make_response(request_msg, ZMQRequestType.KV_DESCRIBE_RESPONSE, description)
 
+    def _handle_kv_prepare_update_request(self, request_msg: ZMQMessage) -> ZMQMessage:
+        params = request_msg.body
+        metadata, columns = self.kv_prepare_update(params["keys"], params["partition_id"], params["fields"])
+        return self._make_response(
+            request_msg, ZMQRequestType.KV_PREPARE_UPDATE_RESPONSE, {"metadata": metadata, "columns": columns}
+        )
+
     def _handle_kv_clear_fields_request(self, request_msg: ZMQMessage) -> ZMQMessage:
         params = request_msg.body
         metadata = self.kv_clear_fields(
diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py
index c3e40431..c9447473 100644
--- a/transfer_queue/interface.py
+++ b/transfer_queue/interface.py
@@ -1194,14 +1194,16 @@ async def async_kv_batch_update(
             f"kv_update is not supported by {type(tq_client.storage_manager).__name__}; "
             "it requires the SimpleStorage backend"
         )
-    batch_meta = await tq_client.async_kv_retrieve_meta(
-        keys=keys, partition_id=partition_id, create=True, include_fields=list(fields.keys())
+    batch_meta, columns = await tq_client.async_kv_prepare_update(
+        keys=keys, partition_id=partition_id, fields=list(fields.keys())
     )
 
     if batch_meta.size != len(keys):
         raise RuntimeError(f"Retrieved BatchMeta size {batch_meta.size} does not match {len(keys)} keys")
 
-    batch_meta = await tq_client.async_update(metadata=batch_meta, values=fields, merge_fn=merge_fn, finalize=finalize)
+    batch_meta = await tq_client.async_update(
+        metadata=batch_meta, values=fields, merge_fn=merge_fn, columns=columns, finalize=finalize
+    )
 
     return KVBatchMeta(
         keys=keys,
diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py
index 02fc1575..e2e9c1d2 100644
--- a/transfer_queue/storage/managers/base.py
+++ b/transfer_queue/storage/managers/base.py
@@ -365,6 +365,7 @@ async def update_data(
         metadata: BatchMeta,
         values: TensorDict,
         merge_fn: Callable[[Any, Any], Any],
+        columns: dict[str, dict[str, Any]],
         finalize: bool = True,
     ) -> dict[str, dict[str, Any]]:
         """Merge new values on the unit that holds each sample, or store them if absent.
@@ -372,7 +373,9 @@ async def update_data(
         Args:
             metadata: Samples to update.
             values: New values as a TensorDict aligned with ``metadata``.
-            merge_fn: Called per sample per field as ``merge_fn(old, new)`` when a value exists.
+            merge_fn: Called per sample per field as ``merge_fn(old, new)`` on held rows.
+            columns: ``{field: {"dtype", "is_non_tensor", "held"}}`` from the controller;
+                ``held`` is aligned with ``metadata`` and decides which rows merge.
             finalize: If False, keep the updated fields hidden from consumers.
 
         Returns:
diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py
index 681125a5..2e60b311 100644
--- a/transfer_queue/storage/managers/simple_storage_manager.py
+++ b/transfer_queue/storage/managers/simple_storage_manager.py
@@ -544,6 +544,7 @@ async def update_data(
         metadata: BatchMeta,
         values: TensorDict,
         merge_fn: Callable[[Any, Any], Any],
+        columns: dict[str, dict[str, Any]],
         finalize: bool = True,
     ) -> dict[str, dict[str, Any]]:
         """Merge or add values on each owner unit and notify the controller.
@@ -560,17 +561,14 @@ async def update_data(
                 f"Batch size of values ({values.batch_size[0]}) does not match metadata size ({metadata.size})"
             )
         field_names = list(values.keys())
-        # Fields new to the partition have no column type yet; the rest must keep theirs.
-        column_types = {
-            name: {key: metadata.field_schema[name][key] for key in ("dtype", "is_non_tensor")}
-            for name in field_names
-            if name in metadata.field_schema
-        }
 
         routing = self._group_by_hash(metadata.global_indexes)
         tasks = []
         for su_id, group in routing.items():
             storage_data = {f: self._select_by_positions(values[f], group.batch_positions) for f in field_names}
+            unit_columns = {
+                f: {**columns[f], "held": [columns[f]["held"][p] for p in group.batch_positions]} for f in field_names
+            }
             tasks.append(
                 self._request_with_retry(
                     "update",
@@ -581,7 +579,7 @@ async def update_data(
                         group.global_indexes,
                         storage_data,
                         merge_fn,
-                        column_types,
+                        unit_columns,
                         target_storage_unit=su_id,
                     ),
                     # A merge reads stored state, so a replay after a lost answer
@@ -682,7 +680,7 @@ async def _update_to_single_storage_unit(
         global_indexes: list[int],
         storage_data: dict[str, Any],
         merge_fn: Callable[[Any, Any], Any],
-        column_types: dict[str, dict[str, Any]],
+        columns: dict[str, dict[str, Any]],
         target_storage_unit: str,
         socket: zmq.Socket = None,
     ) -> dict[str, dict[str, Any]]:
@@ -691,7 +689,7 @@ async def _update_to_single_storage_unit(
             "global_indexes": global_indexes,
             "data": storage_data,
             "merge_fn": merge_fn,
-            "column_types": column_types,
+            "columns": columns,
         }
 
         request_msg = ZMQMessage.create(
diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py
index 76fd6821..37d39c49 100644
--- a/transfer_queue/storage/simple_storage.py
+++ b/transfer_queue/storage/simple_storage.py
@@ -245,34 +245,30 @@ def apply_update(
         global_indexes: list[int],
         new_data: dict[str, Any],
         merge_fn: Callable[[Any, Any], Any],
-        column_types: dict[str, dict[str, Any]],
+        columns: dict[str, dict[str, Any]],
     ) -> dict[str, dict[str, Any]]:
-        """Merge into stored values, store new ones as is, validate, then write once.
+        """Merge into held values, store the other rows as is, validate, then write once.
 
-        A row that has no value for a field takes ``new`` without calling ``merge_fn``,
-        as a put would. A merge result must keep the old value's tensor/non-tensor kind
-        and tensor dtype; a new row must match ``column_types``, the controller's column
-        type for fields it already knows. All merges and checks finish before one
-        ``put_data`` call, so a merge error or incompatible value leaves storage unchanged.
+        ``columns[field]["held"]`` comes from the controller, not from this unit: a local
+        value the controller no longer tracks (e.g. after a failed clear) is overwritten,
+        never merged. Every final value must match the field's known column type. All
+        merges and checks finish before one ``put_data`` call, so a merge error, an
+        incompatible value or a held row missing here leaves storage unchanged.
 
         Returns:
             Per-sample description of the stored values, keyed by field name.
         """
         computed: dict[str, list] = {}
         for field, new_values in new_data.items():
-            stored = self.field_data.get(field, {})
-            present = [i for i, idx in enumerate(global_indexes) if idx in stored]
+            column = columns[field]
+            present = [i for i, held in enumerate(column["held"]) if held]
             old_values = self.get_data([field], [global_indexes[i] for i in present])[field] if present else []
             values = list(new_values)
             for i, old in zip(present, old_values, strict=True):
                 values[i] = merge_fn(old, new_values[i])
-                _check_value_type(field, values[i], isinstance(old, torch.Tensor), getattr(old, "dtype", None))
-            column = column_types.get(field)
-            if column is not None:
-                merged = set(present)
-                for i, value in enumerate(values):
-                    if i not in merged:
-                        _check_value_type(field, value, not column["is_non_tensor"], column["dtype"])
+            if column["is_non_tensor"] is not None:
+                for value in values:
+                    _check_value_type(field, value, not column["is_non_tensor"], column["dtype"])
             computed[field] = values
         self.put_data(computed, global_indexes)
         return {field: _describe_stored_values(values) for field, values in computed.items()}
@@ -1257,7 +1253,7 @@ def _handle_update(self, data_parts: ZMQMessage) -> ZMQMessage:
             global_indexes = data_parts.body["global_indexes"]
             merge_fn = data_parts.body.get("merge_fn")
             new_data = data_parts.body.get("data")
-            column_types = data_parts.body.get("column_types")
+            columns = data_parts.body.get("columns")
 
             with limit_pytorch_auto_parallel_threads(
                 target_num_threads=TQ_NUM_THREADS, info=f"[{self.storage_unit_id}] _handle_update"
@@ -1266,9 +1262,9 @@ def _handle_update(self, data_parts: ZMQMessage) -> ZMQMessage:
                     raise TypeError(f"merge_fn must be callable, got {type(merge_fn).__name__}")
                 if not isinstance(new_data, dict) or not new_data:
                     raise TypeError("update data must be a non-empty dict")
-                if not isinstance(column_types, dict):
-                    raise TypeError("update request must carry column_types")
-                stored_shapes = self.storage_data.apply_update(global_indexes, new_data, merge_fn, column_types)
+                if not isinstance(columns, dict) or set(columns) != set(new_data):
+                    raise ValueError("columns must describe every update field")
+                stored_shapes = self.storage_data.apply_update(global_indexes, new_data, merge_fn, columns)
 
             return ZMQMessage.create(
                 request_type=ZMQRequestType.UPDATE_DATA_RESPONSE,  # type: ignore[arg-type]
diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py
index af30be40..9c765cfb 100644
--- a/transfer_queue/utils/zmq_utils.py
+++ b/transfer_queue/utils/zmq_utils.py
@@ -125,6 +125,8 @@ class ZMQRequestType(ExplicitEnum):
     KV_LIST_RESPONSE = "KV_LIST_RESPONSE"
     KV_DESCRIBE = "KV_DESCRIBE"
     KV_DESCRIBE_RESPONSE = "KV_DESCRIBE_RESPONSE"
+    KV_PREPARE_UPDATE = "KV_PREPARE_UPDATE"
+    KV_PREPARE_UPDATE_RESPONSE = "KV_PREPARE_UPDATE_RESPONSE"
     KV_CLEAR_FIELDS = "KV_CLEAR_FIELDS"
     KV_CLEAR_FIELDS_RESPONSE = "KV_CLEAR_FIELDS_RESPONSE"
 

From 33ac43d13a3ca6ff358a0628d1cdf404b763e695 Mon Sep 17 00:00:00 2001
From: OutstanderWang <wangweiyanster@gmail.com>
Date: Wed, 14 Oct 2026 14:29:41 +0800
Subject: [PATCH 06/12] docs: describe controller-decided merge rows and
 kv_update error types

Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
---
 README.md                   | 15 +++++++++------
 transfer_queue/interface.py | 14 ++++++++------
 2 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/README.md b/README.md
index 242b577c..f3719ba7 100644
--- a/README.md
+++ b/README.md
@@ -121,12 +121,15 @@ To simplify the usage of TransferQueue, we provide a Redis-style high-level API
 - **(async_)kv_describe**: Show each key's stored fields, including those `kv_update(finalize=False)` keeps hidden.
 - **(async_)kv_clear**: Remove key-value pairs from storage, or only selected `fields` while keeping the keys.
 
-`kv_update` stores a field the key does not hold yet as given, without calling
-`merge_fn`; a merge must preserve the stored value's tensor/non-tensor type and
-tensor dtype. A timeout has an unknown outcome because a storage unit may commit
-after the caller stops waiting, so non-idempotent merge operations must not be
-retried blindly. Updates are atomic within one storage-unit request; a batch
-spanning multiple units is not a distributed transaction.
+`kv_update` merges only the rows the controller records as holding the field
+(including hidden ones); every other row takes the new value as given, without
+calling `merge_fn`, and overwrites anything a storage unit kept after a failed
+clear. Every stored value must keep the field's tensor/non-tensor type and
+tensor dtype in the partition. A timeout has an unknown outcome because a
+storage unit may commit after the caller stops waiting, so non-idempotent merge
+operations must not be retried blindly. Updates are atomic within one
+storage-unit request; a batch spanning multiple units is not a distributed
+transaction.
 
 To build a field over several updates without exposing intermediate values, pass
 `finalize=False` to every `kv_update` except the last one, which uses the default
diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py
index c9447473..b0ce2e1c 100644
--- a/transfer_queue/interface.py
+++ b/transfer_queue/interface.py
@@ -649,15 +649,16 @@ def kv_update(
 ) -> KVBatchMeta:
     """Merge new values into the fields of one key, adding the key or fields if missing.
 
-    A superset of ``kv_put``: a missing key is created and a field the key does not
-    hold yet is stored as given, without calling ``merge_fn``.
+    A superset of ``kv_put``: a missing key is created and a field the controller does
+    not record for the key yet is stored as given, without calling ``merge_fn``.
 
     Args:
         key: User-specified key. Created if missing.
         partition_id: Partition that holds the key.
         fields: New values keyed by field name.
-        merge_fn: ``merge_fn(old, new) -> stored`` for fields the key already holds.
-            SimpleStorage runs it once per field on the unit that owns the sample.
+        merge_fn: ``merge_fn(old, new) -> stored`` for fields the key already holds,
+            hidden ones included. SimpleStorage runs it once per field on the unit
+            that owns the sample.
         finalize: If False, the updated fields become (or stay) hidden from consumers,
             so readers never see an intermediate value. Pass True (the default) on the
             last update to make them ready.
@@ -666,8 +667,9 @@ def kv_update(
         Metadata for the updated sample.
 
     Raises:
-        TypeError: If ``fields`` is not a non-empty dict, ``merge_fn`` is not callable,
-            or a merge result changes the stored value's type or dtype.
+        TypeError: If ``fields`` is not a non-empty dict or ``merge_fn`` is not callable.
+        RuntimeError: If ``merge_fn`` fails or a stored value would change the field's
+            tensor/non-tensor type or dtype; the failing storage unit writes nothing.
         NotImplementedError: If the storage backend is not SimpleStorage.
 
     Note:

From 912262d2b5f5026abccacdf00a6ed202d89f15ed Mon Sep 17 00:00:00 2001
From: OutstanderWang <wangweiyanster@gmail.com>
Date: Wed, 14 Oct 2026 18:25:26 +0800
Subject: [PATCH 07/12] [fix] controller: check dtype before writing production
 status

A dtype conflict raised in _update_field_metadata after production_status
was already written, so a rejected write still flipped the cells (ready, or
hidden for ready=False) for rows FieldMeta does not list. Run the same dtype
check up front so a conflicting write changes nothing. Tensor/non-tensor
mixing stays allowed, as FieldMeta.update allows it.

Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
---
 tests/test_controller_data_partitions.py | 10 ++++++++++
 transfer_queue/controller.py             |  7 +++++++
 2 files changed, 17 insertions(+)

diff --git a/tests/test_controller_data_partitions.py b/tests/test_controller_data_partitions.py
index 1a6cced5..3d92e175 100644
--- a/tests/test_controller_data_partitions.py
+++ b/tests/test_controller_data_partitions.py
@@ -1362,6 +1362,16 @@ def test_unfinalized_write_records_metadata_but_stays_hidden(self):
         partition.update_production_status([1], ["pv"], {"pv": t4})
         assert partition.scan_data_status(["pv"], task_name="trainer") == [0, 1]
 
+    def test_dtype_conflict_leaves_readiness_unchanged(self):
+        partition = self._make_partition()
+        f32 = {"dtype": "torch.float32", "shape": (4,), "is_nested": False, "is_non_tensor": False}
+        i64 = {**f32, "dtype": "torch.int64"}
+        partition.update_production_status([0, 1], ["pv"], {"pv": f32})
+
+        assert not partition.update_production_status([1], ["pv"], {"pv": i64}, ready=False)
+        assert partition.scan_data_status(["pv"], task_name="trainer") == [0, 1]
+        assert partition.field_metadata["pv"].dtype == "torch.float32"
+
     def test_update_production_status_updates_field_metadata(self):
         """Test that update_production_status correctly updates field_metadata via _update_field_metadata."""
         partition = self._make_partition()
diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py
index 45a9dc11..ba23dc81 100644
--- a/transfer_queue/controller.py
+++ b/transfer_queue/controller.py
@@ -523,6 +523,13 @@ def update_production_status(
             True if update was successful, False on error
         """
         try:
+            # Same dtype check as FieldMeta.update, run first so a conflicting write changes nothing.
+            for name, incoming in field_schema.items():
+                existing = self.field_metadata.get(name)
+                dtype = incoming.get("dtype")
+                if existing is not None and None not in (existing.dtype, dtype) and existing.dtype != dtype:
+                    raise ValueError(f"Field {name!r} dtype mismatch: {existing.dtype} != {dtype}")
+
             # Derive field_names from field_schema to guarantee consistency
             field_names = list(field_schema.keys())
 

From f7b47a39d4dbe7b5d8cb3616a730b7f6b93b7da0 Mon Sep 17 00:00:00 2001
From: OutstanderWang <wangweiyanster@gmail.com>
Date: Wed, 14 Oct 2026 18:25:54 +0800
Subject: [PATCH 08/12] [fix] kv_describe: return only fields and unfinalized
 per key

global_index is internal and reused across partitions once a key is cleared,
and tag duplicates kv_list. The dump keeps reading the full controller rows
through the client.

Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
---
 tests/e2e/test_kv_interface_e2e.py | 2 +-
 transfer_queue/interface.py        | 8 +++++---
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/tests/e2e/test_kv_interface_e2e.py b/tests/e2e/test_kv_interface_e2e.py
index 1e8a6f6d..dd6cfe31 100644
--- a/tests/e2e/test_kv_interface_e2e.py
+++ b/tests/e2e/test_kv_interface_e2e.py
@@ -1270,7 +1270,7 @@ def tokens_ready():
         with pytest.raises(ValueError, match="not finalized"):
             tq_api.kv_batch_get_by_meta(meta)
         row = tq_api.kv_describe(keys=key, partition_id=partition_id)[key]
-        assert (row["fields"], row["unfinalized"]) == (["prompt", "tokens"], ["tokens"])
+        assert row == {"fields": ["prompt", "tokens"], "unfinalized": ["tokens"]}
         with pytest.raises(RuntimeError, match="not found"):
             tq_api.kv_describe(keys=[key, "absent"], partition_id=partition_id)
 
diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py
index b0ce2e1c..40b4053e 100644
--- a/transfer_queue/interface.py
+++ b/transfer_queue/interface.py
@@ -1078,8 +1078,8 @@ async def async_kv_describe(keys: list[str] | str, partition_id: str) -> dict[st
         partition_id: Partition containing the keys.
 
     Returns:
-        ``{key: {"global_index", "fields", "unfinalized", "tag"}}``. ``fields`` lists every
-        stored field of the key and ``unfinalized`` the subset still hidden from consumers.
+        ``{key: {"fields", "unfinalized"}}``. ``fields`` lists every stored field of the
+        key and ``unfinalized`` the subset that is not ready yet.
 
     Raises:
         RuntimeError: If the partition or any key does not exist.
@@ -1093,7 +1093,9 @@ async def async_kv_describe(keys: list[str] | str, partition_id: str) -> dict[st
     keys = [keys] if isinstance(keys, str) else keys
     tq_client = _maybe_create_tq_client()
     description = await tq_client.async_kv_describe(list(dict.fromkeys(keys)), partition_id)
-    return description["rows"]
+    return {
+        key: {"fields": row["fields"], "unfinalized": row["unfinalized"]} for key, row in description["rows"].items()
+    }
 
 
 async def async_kv_clear(keys: list[str] | str, partition_id: str, fields: list[str] | str | None = None) -> None:

From bd4d25f103fe23e6d47fa31c22b5b9798b5b6b87 Mon Sep 17 00:00:00 2001
From: OutstanderWang <wangweiyanster@gmail.com>
Date: Wed, 14 Oct 2026 18:27:35 +0800
Subject: [PATCH 09/12] docs: require one writer per key for kv_update; drop
 the re-hide contract

Readiness follows the last controller notification, so interleaved writers
can re-hide a field another writer finalized. finalize=False on an already
ready field does not isolate in-flight readers, so the e2e test no longer
asserts it and the docstring no longer promises it.

Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
---
 README.md                          | 3 ++-
 tests/e2e/test_kv_interface_e2e.py | 6 ------
 transfer_queue/interface.py        | 7 ++++---
 3 files changed, 6 insertions(+), 10 deletions(-)

diff --git a/README.md b/README.md
index f3719ba7..ef1bdd71 100644
--- a/README.md
+++ b/README.md
@@ -134,7 +134,8 @@ transaction.
 To build a field over several updates without exposing intermediate values, pass
 `finalize=False` to every `kv_update` except the last one, which uses the default
 `finalize=True`. Until then the field is stored but not ready, so samplers, the
-dataloader and `kv_batch_get` do not see it.
+dataloader and `kv_batch_get` do not see it. Updates to one key must come from
+one writer at a time; use `kv_global_lock` if several processes write the same key.
 
 **Key Features**
 
diff --git a/tests/e2e/test_kv_interface_e2e.py b/tests/e2e/test_kv_interface_e2e.py
index dd6cfe31..8edc63c4 100644
--- a/tests/e2e/test_kv_interface_e2e.py
+++ b/tests/e2e/test_kv_interface_e2e.py
@@ -1279,12 +1279,6 @@ def tokens_ready():
         assert tq_api.kv_describe(keys=key, partition_id=partition_id)[key]["unfinalized"] == []
         data = tq_api.kv_batch_get(keys=key, partition_id=partition_id, select_fields="tokens")
         assert_tensor_equal(data["tokens"][0], torch.tensor([1, 2, 3]))
-
-        # A non-final update hides a field that was already ready.
-        tq_api.kv_update(
-            key=key, partition_id=partition_id, fields={"tokens": torch.tensor([4])}, merge_fn=concat, finalize=False
-        )
-        assert not tokens_ready()
         tq_api.kv_clear(keys=key, partition_id=partition_id)
 
     def test_creates_missing_key_and_field_and_rejects_incompatible_result(self, tq_api, backend_name):
diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py
index 40b4053e..2f8a7d42 100644
--- a/transfer_queue/interface.py
+++ b/transfer_queue/interface.py
@@ -659,9 +659,10 @@ def kv_update(
         merge_fn: ``merge_fn(old, new) -> stored`` for fields the key already holds,
             hidden ones included. SimpleStorage runs it once per field on the unit
             that owns the sample.
-        finalize: If False, the updated fields become (or stay) hidden from consumers,
-            so readers never see an intermediate value. Pass True (the default) on the
-            last update to make them ready.
+        finalize: If False, the updated fields stay hidden from consumers until an update
+            with True (the default), so a field built over several updates never shows an
+            intermediate value. Hiding a field that is already ready is not guaranteed:
+            readers that already hold its metadata can still see the new value.
 
     Returns:
         Metadata for the updated sample.

From 8f5b14d5b2469c1009773c9ce302faa0e205bc61 Mon Sep 17 00:00:00 2001
From: OutstanderWang <wangweiyanster@gmail.com>
Date: Wed, 14 Oct 2026 18:27:43 +0800
Subject: [PATCH 10/12] docs: stop calling kv_update a superset of kv_put

It takes no tags or data_parser and runs only on SimpleStorage.

Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
---
 README.md                   | 2 +-
 transfer_queue/interface.py | 5 +++--
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index ef1bdd71..a5745395 100644
--- a/README.md
+++ b/README.md
@@ -116,7 +116,7 @@ To simplify the usage of TransferQueue, we provide a Redis-style high-level API
 - **(async_)kv_put**: Insert/Update a multi-column sample by key, with an optional metadata tag.
 - **(async_)kv_batch_put**: Put multiple key-value pairs efficiently in batches.
 - **(async_)kv_batch_get**: Retrieve samples (by keys), supporting column selection (by fields).
-- **(async_)kv_update / (async_)kv_batch_update**: Merge new values into fields with `merge_fn(old, new)`; a superset of put that adds missing keys and fields. SimpleStorage only.
+- **(async_)kv_update / (async_)kv_batch_update**: Merge new values into fields with `merge_fn(old, new)`, adding missing keys and fields. SimpleStorage only.
 - **(async_)kv_list**: List keys and tags (metadata) in a partition.
 - **(async_)kv_describe**: Show each key's stored fields, including those `kv_update(finalize=False)` keeps hidden.
 - **(async_)kv_clear**: Remove key-value pairs from storage, or only selected `fields` while keeping the keys.
diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py
index 2f8a7d42..465dc815 100644
--- a/transfer_queue/interface.py
+++ b/transfer_queue/interface.py
@@ -649,8 +649,9 @@ def kv_update(
 ) -> KVBatchMeta:
     """Merge new values into the fields of one key, adding the key or fields if missing.
 
-    A superset of ``kv_put``: a missing key is created and a field the controller does
-    not record for the key yet is stored as given, without calling ``merge_fn``.
+    A missing key is created and a field the controller does not record for the key
+    yet is stored as given, without calling ``merge_fn``. Unlike ``kv_put``, it takes
+    no ``tags`` or ``data_parser``.
 
     Args:
         key: User-specified key. Created if missing.

From c042ba9b7a6e1b1fd938540364cc059ab8832452 Mon Sep 17 00:00:00 2001
From: OutstanderWang <wangweiyanster@gmail.com>
Date: Wed, 14 Oct 2026 18:28:37 +0800
Subject: [PATCH 11/12] [refactor] kv_update: rename columns to merge_plan

It carries the controller's plan for the update (which rows merge, and the
column type to check results against), not column data.

Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
---
 tests/test_controller.py                         |  6 +++---
 transfer_queue/client.py                         | 16 ++++++++--------
 transfer_queue/controller.py                     | 10 +++++-----
 transfer_queue/interface.py                      |  4 ++--
 transfer_queue/storage/managers/base.py          |  4 ++--
 .../storage/managers/simple_storage_manager.py   |  9 +++++----
 transfer_queue/storage/simple_storage.py         | 14 +++++++-------
 7 files changed, 32 insertions(+), 31 deletions(-)

diff --git a/tests/test_controller.py b/tests/test_controller.py
index 344505ff..edbfdff2 100644
--- a/tests/test_controller.py
+++ b/tests/test_controller.py
@@ -1204,13 +1204,13 @@ def test_controller_kv_prepare_update_reports_held_rows_from_field_metadata(self
         ray.get(tq_controller.update_production_status.remote(partition_id, [s0], response, ready=False))
         ray.get(tq_controller.update_production_status.remote(partition_id, [s1], prompt))
 
-        metadata, columns = ray.get(
+        metadata, merge_plan = ray.get(
             tq_controller.kv_prepare_update.remote(["s0", "s1", "s2"], partition_id, ["response", "fresh"])
         )
 
         assert metadata.global_indexes[:2] == [s0, s1]
-        assert columns["response"] == {"dtype": torch.int64, "is_non_tensor": False, "held": [True, False, False]}
-        assert columns["fresh"] == {"dtype": None, "is_non_tensor": None, "held": [False, False, False]}
+        assert merge_plan["response"] == {"dtype": torch.int64, "is_non_tensor": False, "held": [True, False, False]}
+        assert merge_plan["fresh"] == {"dtype": None, "is_non_tensor": None, "held": [False, False, False]}
 
         ray.get(tq_controller.clear_partition.remote(partition_id))
 
diff --git a/transfer_queue/client.py b/transfer_queue/client.py
index f6ae1521..7a088b99 100644
--- a/transfer_queue/client.py
+++ b/transfer_queue/client.py
@@ -497,7 +497,7 @@ async def async_update(
         metadata: BatchMeta,
         values: TensorDict,
         merge_fn: Callable[[Any, Any], Any],
-        columns: dict[str, dict[str, Any]],
+        merge_plan: dict[str, dict[str, Any]],
         finalize: bool = True,
     ) -> BatchMeta:
         """Merge new values into fields on SimpleStorage units, storing them where absent.
@@ -506,7 +506,7 @@ async def async_update(
             metadata: Samples to update, already registered with the controller.
             values: New values aligned with ``metadata``.
             merge_fn: Called per sample per field as ``merge_fn(old, new)`` on held rows.
-            columns: Per-field column type and held rows, from ``async_kv_prepare_update``.
+            merge_plan: Per-field column type and held rows, from ``async_kv_prepare_update``.
             finalize: If False, keep the updated fields hidden from consumers.
 
         Returns:
@@ -523,7 +523,7 @@ async def async_update(
         if not metadata or metadata.size == 0:
             raise ValueError("metadata cannot be none or empty")
 
-        field_schema = await self.storage_manager.update_data(metadata, values, merge_fn, columns, finalize=finalize)
+        field_schema = await self.storage_manager.update_data(metadata, values, merge_fn, merge_plan, finalize=finalize)
         if not finalize:
             metadata.production_status[:] = 0
         return metadata.apply_field_schema(field_schema, set_all_ready=finalize)
@@ -1023,7 +1023,7 @@ async def async_kv_prepare_update(
         """Register missing keys and get, per field, its column type and which rows hold it.
 
         Returns:
-            The keys' BatchMeta and the ``columns`` that ``async_update`` takes.
+            The keys' BatchMeta and the ``merge_plan`` that ``async_update`` takes.
         """
         response_msg = await self._request_controller(
             socket=socket,
@@ -1031,7 +1031,7 @@ async def async_kv_prepare_update(
             response_type=ZMQRequestType.KV_PREPARE_UPDATE_RESPONSE,
             body={"keys": keys, "partition_id": partition_id, "fields": fields},
         )
-        return response_msg.body["metadata"], response_msg.body["columns"]
+        return response_msg.body["metadata"], response_msg.body["merge_plan"]
 
     async def async_kv_clear_fields(self, keys: list[str], partition_id: str, fields: list[str]) -> None:
         """Clear fields of existing keys while keeping the keys, tags and consumption status.
@@ -1662,14 +1662,14 @@ def update(
         metadata: BatchMeta,
         values: TensorDict,
         merge_fn: Callable[[Any, Any], Any],
-        columns: dict[str, dict[str, Any]],
+        merge_plan: dict[str, dict[str, Any]],
         finalize: bool = True,
     ) -> BatchMeta:
         """Synchronously merge or add values on SimpleStorage units.
 
         See ``async_update``.
         """
-        return self._update(metadata, values, merge_fn, columns, finalize=finalize)
+        return self._update(metadata, values, merge_fn, merge_plan, finalize=finalize)
 
     def get_data(self, metadata: BatchMeta) -> TensorDict:
         """Synchronously fetch data from storage units and organize into TensorDict.
@@ -1888,7 +1888,7 @@ def kv_retrieve_meta(
     def kv_prepare_update(
         self, keys: list[str], partition_id: str, fields: list[str]
     ) -> tuple[BatchMeta, dict[str, dict[str, Any]]]:
-        """Synchronously register missing keys and get the ``columns`` for ``update``.
+        """Synchronously register missing keys and get the ``merge_plan`` for ``update``.
 
         See ``async_kv_prepare_update``.
         """
diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py
index ba23dc81..d19152cd 100644
--- a/transfer_queue/controller.py
+++ b/transfer_queue/controller.py
@@ -1746,16 +1746,16 @@ def kv_prepare_update(
         metadata = self.kv_retrieve_meta(keys, partition_id, create=True)
         partition = self._get_partition(partition_id)
         assert partition is not None
-        columns = {}
+        merge_plan = {}
         for name in fields:
             meta = partition.field_metadata.get(name)
             held = meta.global_indexes if meta is not None else set()
-            columns[name] = {
+            merge_plan[name] = {
                 "dtype": meta.dtype if meta is not None else None,
                 "is_non_tensor": meta.is_non_tensor if meta is not None else None,
                 "held": [idx in held for idx in metadata.global_indexes],
             }
-        return metadata, columns
+        return metadata, merge_plan
 
     def kv_clear_fields(self, keys: list[str], partition_id: str, fields: list[str]) -> BatchMeta:
         """Clear fields of existing keys and return the cleared cells' metadata for storage deletion.
@@ -2297,9 +2297,9 @@ def _handle_kv_describe_request(self, request_msg: ZMQMessage) -> ZMQMessage:
 
     def _handle_kv_prepare_update_request(self, request_msg: ZMQMessage) -> ZMQMessage:
         params = request_msg.body
-        metadata, columns = self.kv_prepare_update(params["keys"], params["partition_id"], params["fields"])
+        metadata, merge_plan = self.kv_prepare_update(params["keys"], params["partition_id"], params["fields"])
         return self._make_response(
-            request_msg, ZMQRequestType.KV_PREPARE_UPDATE_RESPONSE, {"metadata": metadata, "columns": columns}
+            request_msg, ZMQRequestType.KV_PREPARE_UPDATE_RESPONSE, {"metadata": metadata, "merge_plan": merge_plan}
         )
 
     def _handle_kv_clear_fields_request(self, request_msg: ZMQMessage) -> ZMQMessage:
diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py
index 465dc815..1ba2f36b 100644
--- a/transfer_queue/interface.py
+++ b/transfer_queue/interface.py
@@ -1200,7 +1200,7 @@ async def async_kv_batch_update(
             f"kv_update is not supported by {type(tq_client.storage_manager).__name__}; "
             "it requires the SimpleStorage backend"
         )
-    batch_meta, columns = await tq_client.async_kv_prepare_update(
+    batch_meta, merge_plan = await tq_client.async_kv_prepare_update(
         keys=keys, partition_id=partition_id, fields=list(fields.keys())
     )
 
@@ -1208,7 +1208,7 @@ async def async_kv_batch_update(
         raise RuntimeError(f"Retrieved BatchMeta size {batch_meta.size} does not match {len(keys)} keys")
 
     batch_meta = await tq_client.async_update(
-        metadata=batch_meta, values=fields, merge_fn=merge_fn, columns=columns, finalize=finalize
+        metadata=batch_meta, values=fields, merge_fn=merge_fn, merge_plan=merge_plan, finalize=finalize
     )
 
     return KVBatchMeta(
diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py
index e2e9c1d2..f881a0d5 100644
--- a/transfer_queue/storage/managers/base.py
+++ b/transfer_queue/storage/managers/base.py
@@ -365,7 +365,7 @@ async def update_data(
         metadata: BatchMeta,
         values: TensorDict,
         merge_fn: Callable[[Any, Any], Any],
-        columns: dict[str, dict[str, Any]],
+        merge_plan: dict[str, dict[str, Any]],
         finalize: bool = True,
     ) -> dict[str, dict[str, Any]]:
         """Merge new values on the unit that holds each sample, or store them if absent.
@@ -374,7 +374,7 @@ async def update_data(
             metadata: Samples to update.
             values: New values as a TensorDict aligned with ``metadata``.
             merge_fn: Called per sample per field as ``merge_fn(old, new)`` on held rows.
-            columns: ``{field: {"dtype", "is_non_tensor", "held"}}`` from the controller;
+            merge_plan: ``{field: {"dtype", "is_non_tensor", "held"}}`` from the controller;
                 ``held`` is aligned with ``metadata`` and decides which rows merge.
             finalize: If False, keep the updated fields hidden from consumers.
 
diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py
index 2e60b311..1ce03bf5 100644
--- a/transfer_queue/storage/managers/simple_storage_manager.py
+++ b/transfer_queue/storage/managers/simple_storage_manager.py
@@ -544,7 +544,7 @@ async def update_data(
         metadata: BatchMeta,
         values: TensorDict,
         merge_fn: Callable[[Any, Any], Any],
-        columns: dict[str, dict[str, Any]],
+        merge_plan: dict[str, dict[str, Any]],
         finalize: bool = True,
     ) -> dict[str, dict[str, Any]]:
         """Merge or add values on each owner unit and notify the controller.
@@ -567,7 +567,8 @@ async def update_data(
         for su_id, group in routing.items():
             storage_data = {f: self._select_by_positions(values[f], group.batch_positions) for f in field_names}
             unit_columns = {
-                f: {**columns[f], "held": [columns[f]["held"][p] for p in group.batch_positions]} for f in field_names
+                f: {**merge_plan[f], "held": [merge_plan[f]["held"][p] for p in group.batch_positions]}
+                for f in field_names
             }
             tasks.append(
                 self._request_with_retry(
@@ -680,7 +681,7 @@ async def _update_to_single_storage_unit(
         global_indexes: list[int],
         storage_data: dict[str, Any],
         merge_fn: Callable[[Any, Any], Any],
-        columns: dict[str, dict[str, Any]],
+        merge_plan: dict[str, dict[str, Any]],
         target_storage_unit: str,
         socket: zmq.Socket = None,
     ) -> dict[str, dict[str, Any]]:
@@ -689,7 +690,7 @@ async def _update_to_single_storage_unit(
             "global_indexes": global_indexes,
             "data": storage_data,
             "merge_fn": merge_fn,
-            "columns": columns,
+            "merge_plan": merge_plan,
         }
 
         request_msg = ZMQMessage.create(
diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py
index 37d39c49..09a245f1 100644
--- a/transfer_queue/storage/simple_storage.py
+++ b/transfer_queue/storage/simple_storage.py
@@ -245,11 +245,11 @@ def apply_update(
         global_indexes: list[int],
         new_data: dict[str, Any],
         merge_fn: Callable[[Any, Any], Any],
-        columns: dict[str, dict[str, Any]],
+        merge_plan: dict[str, dict[str, Any]],
     ) -> dict[str, dict[str, Any]]:
         """Merge into held values, store the other rows as is, validate, then write once.
 
-        ``columns[field]["held"]`` comes from the controller, not from this unit: a local
+        ``merge_plan[field]["held"]`` comes from the controller, not from this unit: a local
         value the controller no longer tracks (e.g. after a failed clear) is overwritten,
         never merged. Every final value must match the field's known column type. All
         merges and checks finish before one ``put_data`` call, so a merge error, an
@@ -260,7 +260,7 @@ def apply_update(
         """
         computed: dict[str, list] = {}
         for field, new_values in new_data.items():
-            column = columns[field]
+            column = merge_plan[field]
             present = [i for i, held in enumerate(column["held"]) if held]
             old_values = self.get_data([field], [global_indexes[i] for i in present])[field] if present else []
             values = list(new_values)
@@ -1253,7 +1253,7 @@ def _handle_update(self, data_parts: ZMQMessage) -> ZMQMessage:
             global_indexes = data_parts.body["global_indexes"]
             merge_fn = data_parts.body.get("merge_fn")
             new_data = data_parts.body.get("data")
-            columns = data_parts.body.get("columns")
+            merge_plan = data_parts.body.get("merge_plan")
 
             with limit_pytorch_auto_parallel_threads(
                 target_num_threads=TQ_NUM_THREADS, info=f"[{self.storage_unit_id}] _handle_update"
@@ -1262,9 +1262,9 @@ def _handle_update(self, data_parts: ZMQMessage) -> ZMQMessage:
                     raise TypeError(f"merge_fn must be callable, got {type(merge_fn).__name__}")
                 if not isinstance(new_data, dict) or not new_data:
                     raise TypeError("update data must be a non-empty dict")
-                if not isinstance(columns, dict) or set(columns) != set(new_data):
-                    raise ValueError("columns must describe every update field")
-                stored_shapes = self.storage_data.apply_update(global_indexes, new_data, merge_fn, columns)
+                if not isinstance(merge_plan, dict) or set(merge_plan) != set(new_data):
+                    raise ValueError("merge_plan must describe every update field")
+                stored_shapes = self.storage_data.apply_update(global_indexes, new_data, merge_fn, merge_plan)
 
             return ZMQMessage.create(
                 request_type=ZMQRequestType.UPDATE_DATA_RESPONSE,  # type: ignore[arg-type]

From bd301128ca502f55bd537820904285ea36349c30 Mon Sep 17 00:00:00 2001
From: OutstanderWang <wangweiyanster@gmail.com>
Date: Thu, 15 Oct 2026 10:24:12 +0800
Subject: [PATCH 12/12] fix(tutorial): give the first attention_mask the batch
 dtype

02_kv_interface put a float32 attention_mask for key 0_0, then int64 masks for
the batch. The controller rejected the batch notification on the dtype
mismatch, but used to mark its cells ready anyway, so the tutorial passed by
accident. With the dtype now checked before readiness is written, input_ids
of the batch stayed not ready and Step 7 failed.

Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
---
 tutorial/02_kv_interface.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tutorial/02_kv_interface.py b/tutorial/02_kv_interface.py
index ef9e62b4..01752b5e 100644
--- a/tutorial/02_kv_interface.py
+++ b/tutorial/02_kv_interface.py
@@ -67,7 +67,7 @@ def demonstrate_kv_api():
 
     # Define the data content (The "Value")
     input_ids = torch.tensor([[1, 2, 3]])
-    attention_mask = torch.ones(input_ids.size())
+    attention_mask = torch.ones_like(input_ids)
 
     single_sample = TensorDict(
         {