[feat] Support multi-step kv_update with partial values hidden

README.md

@@ -116,15 +116,26 @@ 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) and, on SimpleStorage, an optional `data_parser` run on the storage units.
-- **(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)`, 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.
 
-`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` 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
+`finalize=True`. Until then the field is stored but not ready, so samplers, the
+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.
 
 **Advanced: parsing on get.** Like `kv_batch_put`, `kv_batch_get` takes a
 `data_parser`, here run by each SimpleStorage unit before it replies, so only the

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

tests/e2e/test_data_dump_e2e.py

@@ -243,6 +243,54 @@ 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_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
@@ -681,10 +729,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()

tests/e2e/test_kv_interface_e2e.py

@@ -1320,25 +1320,62 @@ 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)}
+            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": ["prompt", "tokens"], "unfinalized": ["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]))
+        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":
+            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",
@@ -1349,6 +1386,62 @@ def test_rejects_missing_key_unproduced_field_and_incompatible_result(self, tq_a
         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_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")

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

tests/test_controller_data_partitions.py

@@ -1350,6 +1350,28 @@ 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_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()

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":