[feat] add kv_update and field-level kv_clear

README.md

@@ -74,7 +74,7 @@ Specifically, we provide a `StorageManager` abstraction class that defines the c
 
 - `async def put_data(self, data: TensorDict, metadata: BatchMeta) -> None`
 - `async def get_data(self, metadata: BatchMeta) -> TensorDict`
-- `async def clear_data(self, metadata: BatchMeta) -> None`
+- `async def clear_data(self, metadata: BatchMeta, fields: list[str] | None = None) -> None`
 
 This class encapsulates the core interaction logic within the TransferQueue system. You only need to write a simple subclass to integrate your custom storage backend. Refer to the [Customize](#customize) section for details.
 
@@ -116,8 +116,15 @@ 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_list**: List keys and tags (metadata) in a partition.
-- **(async_)kv_clear**: Remove key-value pairs from storage.
+- **(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.
 
 **Key Features**
 

docs/ssd_offload.md

@@ -213,10 +213,22 @@ GET reads each requested SSD file and reconstructs the value in host memory.
 SSD offload reduces long-lived memory use, but it does not remove temporary
 memory use while values are read or encoded.
 
-### `data_parser` outputs should not share backing storage across samples
+### Dense batches can retain shared backing memory
 
-If a `data_parser` returns differently sized tensor or NumPy views backed by
-the same allocation, an in-memory view can keep the entire allocation resident
-after another view is offloaded. When parsing URLs or file paths, return an
-independently owned value for each sample, or copy retained views before
-returning them.
+SimpleStorage slices dense batches without copying. The rows therefore share
+the batch allocation, and replacing or clearing one row does not release its
+share while another stored row still references that allocation. This is the
+memory trade-off for zero-copy batch distribution; clear or replace every
+sibling row before expecting the full allocation to be released.
+
+The same rule applies when a `data_parser` returns tensor or NumPy views backed
+by one allocation. Return independently owned values when partial release is
+more important than zero-copy storage.
+
+### Atomic multi-field writes require temporary SSD headroom
+
+SimpleStorage writes all replacement files before publishing any of them. This
+keeps a failed multi-field put or merge from exposing a partially updated
+sample, but temporarily requires space for the old files and all new files.
+When sizing the SSD tier, allow for that operation-level peak rather than only
+the final active-byte count.

tests/e2e/test_kv_interface_e2e.py

@@ -1070,6 +1070,34 @@ def test_kv_clear_idempotent(self, controller, tq_api):
         # Clean up
         tq_api.kv_clear(keys=keys, partition_id=partition_id)
 
+    def test_kv_clear_fields_keeps_keys_and_other_fields(self, controller, tq_api):
+        """Field-level clear removes only the selected fields; the field can be put again."""
+        partition_id = "test_partition"
+        keys = ["clear_fields_0", "clear_fields_1"]
+        tq_api.kv_batch_put(
+            keys=keys,
+            partition_id=partition_id,
+            fields=TensorDict(
+                {"pixel_values": torch.ones(2, 4), "input_ids": torch.tensor([[1, 2], [3, 4]])},
+                batch_size=2,
+            ),
+            tags=[{"id": 0}, {"id": 1}],
+        )
+        tq_api.kv_put(key="clear_fields_other", partition_id=partition_id, fields={"input_ids": torch.tensor([5, 6])})
+
+        tq_api.kv_clear(keys=keys + ["nonexistent_key"], partition_id=partition_id, fields=["pixel_values", "missing"])
+
+        assert tq_api.kv_list(partition_id=partition_id)[partition_id][keys[0]] == {"id": 0}
+        data = tq_api.kv_batch_get(keys=keys, partition_id=partition_id)
+        assert list(data.keys()) == ["input_ids"]
+        assert_tensor_equal(data["input_ids"], torch.tensor([[1, 2], [3, 4]]))
+        partition = get_controller_partition(controller, partition_id)
+        assert "pixel_values" not in partition.field_metadata
+
+        tq_api.kv_put(key=keys[0], partition_id=partition_id, fields={"pixel_values": torch.zeros(3)})
+        restored = tq_api.kv_batch_get(keys=keys[0], partition_id=partition_id, select_fields="pixel_values")
+        assert_tensor_equal(restored["pixel_values"][0], torch.zeros(3))
+
 
 class TestKVE2ECornerCases:
     """End-to-end tests for corner cases."""
@@ -1110,6 +1138,156 @@ def test_field_expansion_across_samples(self, controller, tq_api):
         tq_api.kv_clear(keys=keys, partition_id=partition_id)
 
 
+class TestKVUpdateE2E:
+    """SimpleStorage merge behavior through sync and async public APIs."""
+
+    def test_update_then_clear_field_keeps_partition_field_metadata(self, controller, tq_api, backend_name):
+        if backend_name != "SimpleStorage":
+            pytest.skip("merge-backed kv_update is implemented only for SimpleStorage")
+
+        partition_id = "test_partition"
+        keys = ["clear_field_0", "clear_field_1", "clear_field_2"]
+        tq_api.kv_batch_put(
+            keys=keys,
+            partition_id=partition_id,
+            fields=TensorDict({"tokens": torch.arange(12).reshape(3, 4)}, batch_size=3),
+        )
+
+        def concat(old, new):
+            return torch.cat([old, new])
+
+        meta = tq_api.kv_update(
+            key=keys[1],
+            partition_id=partition_id,
+            fields={"tokens": torch.tensor([12, 13])},
+            merge_fn=concat,
+        )
+        assert "tokens" in meta.fields
+        tq_api.kv_clear(keys=keys[0], partition_id=partition_id, fields="tokens")
+
+        partition = get_controller_partition(controller, partition_id)
+        col = partition.field_name_mapping["tokens"]
+        global_idx = partition.keys_mapping[keys[0]]
+        assert partition.production_status[global_idx, col] == 0
+        assert global_idx not in partition.field_metadata["tokens"].global_indexes
+        assert partition.field_metadata["tokens"].is_non_tensor is False
+        assert partition.field_metadata["tokens"].dtype == torch.int64
+
+        tq_api.kv_batch_put(
+            keys=["nested_0", "nested_1"],
+            partition_id=partition_id,
+            fields=TensorDict(
+                {
+                    "tokens": torch.nested.as_nested_tensor(
+                        [torch.tensor([20, 21]), torch.tensor([30, 31, 32])],
+                        layout=torch.jagged,
+                    )
+                },
+                batch_size=2,
+            ),
+        )
+        partition = get_controller_partition(controller, partition_id)
+        assert partition.field_metadata["tokens"].is_nested is True
+
+    def test_batch_update_multiple_fields_across_storage_units(self, controller, tq_api, backend_name):
+        if backend_name != "SimpleStorage":
+            pytest.skip("merge-backed kv_update is implemented only for SimpleStorage")
+
+        partition_id = "test_partition"
+        keys = [f"batch_update_{i}" for i in range(8)]
+        lengths = [i + 1 for i in range(len(keys))]
+        tq_api.kv_batch_put(
+            keys=keys,
+            partition_id=partition_id,
+            fields=TensorDict(
+                {
+                    "tokens": torch.nested.as_nested_tensor(
+                        [torch.arange(length) for length in lengths],
+                        layout=torch.jagged,
+                    ),
+                    "score": torch.arange(8).reshape(8, 1),
+                },
+                batch_size=8,
+            ),
+        )
+        tq_api.kv_batch_update(
+            keys=keys,
+            partition_id=partition_id,
+            fields=TensorDict(
+                {
+                    "tokens": torch.arange(100, 108).reshape(8, 1),
+                    "score": torch.arange(200, 208).reshape(8, 1),
+                },
+                batch_size=8,
+            ),
+            merge_fn=lambda old, new: torch.cat([old, new]),
+        )
+
+        retrieved = tq_api.kv_batch_get(keys=keys, partition_id=partition_id)
+        for i, length in enumerate(lengths):
+            assert_tensor_equal(
+                retrieved["tokens"][i],
+                torch.cat([torch.arange(length), torch.tensor([100 + i])]),
+            )
+            assert_tensor_equal(retrieved["score"][i], torch.tensor([i, 200 + i]))
+
+        partition = get_controller_partition(controller, partition_id)
+        tokens = partition.field_metadata["tokens"]
+        assert tokens.is_nested is True
+        for key, length in zip(keys, lengths, strict=True):
+            global_idx = partition.keys_mapping[key]
+            assert tuple(tokens.per_sample_shapes[global_idx]) == (length + 1,)
+
+        tq_api.kv_clear(keys=keys, partition_id=partition_id, fields=["tokens", "score"])
+        partition = get_controller_partition(controller, partition_id)
+        assert "tokens" not in partition.field_metadata
+        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):
+        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,
+            )
+        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",
+                partition_id="test_partition",
+                fields={"tokens": torch.tensor([3])},
+                merge_fn=lambda old, new: torch.cat([old, new]).to(torch.float32),
+            )
+        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_kv_backend_rejects(self, tq_api, backend_name):
+        if backend_name == "SimpleStorage":
+            pytest.skip("rejection is the KV-backend contract")
+
+        tq_api.kv_put(key="k", partition_id="test_partition", fields={"tokens": torch.tensor([1])})
+        with pytest.raises(NotImplementedError, match="kv_update is not supported"):
+            tq_api.kv_update(
+                key="k",
+                partition_id="test_partition",
+                fields={"tokens": torch.tensor([2])},
+                merge_fn=lambda _old, new: new,
+            )
+        tq_api.kv_clear(keys="k", partition_id="test_partition")
+
+
 def run_tests():
     """Run all e2e tests manually for debugging."""
     pytest.main([__file__, "-v", "-s"])

tests/test_controller_data_partitions.py

@@ -1310,6 +1310,46 @@ def test_clear_data_nested_field_becomes_regular(self):
         assert field_meta.global_indexes == {0, 1}
         assert field_meta.per_sample_shapes == {}
 
+    def test_clear_fields_keeps_samples_and_other_columns(self):
+        """Clearing one sample's field leaves the field's schema for other samples intact."""
+        import torch
+
+        partition = self._make_partition()
+        t4 = {"dtype": "torch.float32", "shape": (4,), "is_nested": False, "is_non_tensor": False}
+        partition.update_production_status([0, 1, 2], ["pv", "ids"], {"pv": t4, "ids": t4})
+        partition.keys_mapping = {"a": 0, "b": 1, "c": 2}
+        partition.custom_meta = {0: {"tag": "a"}}
+        partition.field_custom_backend_meta = {0: {"pv": "m0", "ids": "i0"}}
+        partition.consumption_status["task"] = torch.ones(3, dtype=torch.int8)
+
+        partition.clear_fields([0], ["pv"])
+
+        pv_col = partition.field_name_mapping["pv"]
+        ids_col = partition.field_name_mapping["ids"]
+        assert partition.production_status[:3, pv_col].tolist() == [0, 1, 1]
+        assert partition.production_status[:3, ids_col].tolist() == [1, 1, 1]
+        assert partition.field_metadata["pv"].to_batch_schema([1, 2])["shape"] == (4,)
+        assert partition.field_metadata["pv"].global_indexes == {1, 2}
+        assert partition.field_custom_backend_meta[0] == {"ids": "i0"}
+        assert partition.global_indexes == {0, 1, 2}
+        assert partition.keys_mapping["a"] == 0
+        assert partition.custom_meta[0] == {"tag": "a"}
+        assert partition.consumption_status["task"].tolist() == [1, 1, 1]
+
+        nested = {"dtype": "torch.float32", "is_nested": True, "per_sample_shapes": {3: (3,), 4: (7,)}}
+        assert partition.update_production_status([3, 4], ["pv"], {"pv": nested})
+        assert partition.field_metadata["pv"].per_sample_shapes == {1: (4,), 2: (4,), 3: (3,), 4: (7,)}
+
+    def test_clear_fields_drops_fieldmeta_when_no_sample_holds_it(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})
+
+        partition.clear_fields([0, 1], ["pv", "missing"])
+
+        assert "pv" not in partition.field_metadata
+        assert partition.global_indexes == {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_kv_storage_manager.py

@@ -480,6 +480,24 @@ def test_put_data_without_custom_backend_meta(mock_notify, test_data_for_put_dat
     assert per_field_custom_backend_meta == {}
 
 
+@patch.object(KVStorageManager, "_connect_to_controller", lambda self: None)
+def test_clear_data_with_fields_deletes_only_those_keys(test_data_for_put_data):
+    mock_storage_client = MagicMock()
+    with patch(f"{STORAGE_CLIENT_FACTORY_PATH}.create", return_value=mock_storage_client):
+        manager = KVStorageManager(controller_info=MagicMock(), config={"client_name": "MockClient"})
+    metadata = test_data_for_put_data["metadata"]
+    metadata._custom_backend_meta = [{"text": f"t{i}", "label": f"l{i}"} for i in range(3)]
+
+    asyncio.run(manager.clear_data(metadata, fields=["text"]))
+    mock_storage_client.clear.assert_called_once_with(
+        keys=["0@text", "1@text", "2@text"], custom_backend_meta=["t0", "t1", "t2"]
+    )
+
+    mock_storage_client.clear.reset_mock()
+    asyncio.run(manager.clear_data(metadata, fields=["missing"]))
+    mock_storage_client.clear.assert_not_called()
+
+
 @patch.object(KVStorageManager, "_connect_to_controller", lambda self: None)
 def test_put_data_custom_backend_meta_length_mismatch_raises_error(test_data_for_put_data):
     """Test that put_data raises ValueError when custom_backend_meta length doesn't match keys."""