feat: kv update with the finalize flag

README.md

@@ -116,15 +116,22 @@ 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_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` 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**
 

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

@@ -1244,25 +1244,68 @@ 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"], 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]))
+
+        # 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):
+        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",
@@ -1273,6 +1316,27 @@ 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_kv_backend_rejects(self, tq_api, backend_name):
         if backend_name == "SimpleStorage":
             pytest.skip("rejection is the KV-backend contract")

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

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