From 75bfd9b8fe49c681342c409b6d70c50a1dbd0e59 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Sun, 20 Sep 2026 15:31:59 +0800 Subject: [PATCH 1/9] [test] Add unit tests for kv_update and empty Cover argument normalization, unit-side parser/empty application, cross-unit field_schema assembly, controller non-tensor marking, and KV backend rejection. Signed-off-by: OutstanderWang --- tests/test_kv_update.py | 203 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 203 insertions(+) create mode 100644 tests/test_kv_update.py diff --git a/tests/test_kv_update.py b/tests/test_kv_update.py new file mode 100644 index 00000000..28a6694d --- /dev/null +++ b/tests/test_kv_update.py @@ -0,0 +1,203 @@ +# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# Copyright 2025 The TransferQueue Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest +import torch + +from transfer_queue.controller import DataPartitionStatus +from transfer_queue.interface import _normalize_kv_update_args +from transfer_queue.metadata import BatchMeta +from transfer_queue.storage.managers.base import KVStorageManager +from transfer_queue.storage.managers.simple_storage_manager import _build_update_field_schema +from transfer_queue.storage.simple_storage import StorageUnitData + + +def _concat(old, new): + if old is None: + return new + return torch.cat([old, new]) + + +def test_normalize_empty_rejects_values_and_parser(): + with pytest.raises(ValueError, match="must not specify values"): + _normalize_kv_update_args("tokens", torch.tensor([1]), None, empty=True) + with pytest.raises(ValueError, match="must not specify parser"): + _normalize_kv_update_args("tokens", None, _concat, empty=True) + names, batch, parser, use_empty = _normalize_kv_update_args("tokens", None, None, empty=True) + assert names == ["tokens"] + assert batch is None + assert parser is None + assert use_empty is True + + +def test_normalize_custom_parser_requires_values(): + with pytest.raises(ValueError, match="requires values"): + _normalize_kv_update_args("tokens", None, _concat) + + +def test_normalize_rejects_non_callable_parser(): + with pytest.raises(TypeError, match="parser must be callable unless empty=True"): + _normalize_kv_update_args("tokens", torch.tensor([1]), None) + + +def test_normalize_single_field_wraps_a_batch(): + names, batch, parser, use_empty = _normalize_kv_update_args("tokens", torch.tensor([4, 5]), _concat) + assert names == ["tokens"] + assert use_empty is False + assert parser is _concat + assert batch is not None + assert batch.batch_size == torch.Size([1]) + assert torch.equal(batch["tokens"][0], torch.tensor([4, 5])) + + +def test_normalize_multi_field_requires_matching_dict(): + with pytest.raises(TypeError, match="must be a dict"): + _normalize_kv_update_args(["a", "b"], torch.tensor([1]), _concat) + with pytest.raises(ValueError, match="same columns"): + _normalize_kv_update_args(["a", "b"], {"a": 1}, _concat) + + +def test_apply_update_concat_prompt_and_response_keeps_field_name(): + """Stored prompt_ids plus new response_ids become the sequence; the field name is unchanged.""" + prompt_ids = torch.tensor([10, 11, 12]) + response_ids = torch.tensor([20, 21]) + data = StorageUnitData() + data.put_data({"sequence_ids": [prompt_ids.clone()]}, [0]) + + described = data.apply_update([0], ["sequence_ids"], {"sequence_ids": [response_ids]}, _concat, False) + + assert list(described) == ["sequence_ids"] + assert torch.equal(data.field_data["sequence_ids"][0], torch.tensor([10, 11, 12, 20, 21])) + assert "prompt_ids" not in data.field_data + assert "response_ids" not in data.field_data + + +def test_apply_update_concatenates_and_is_atomic_on_parser_error(): + data = StorageUnitData() + data.put_data({"tokens": [torch.tensor([1, 2, 3])]}, [7]) + + described = data.apply_update([7], ["tokens"], {"tokens": [torch.tensor([4, 5])]}, _concat, False) + assert described["tokens"] == {"dtype": torch.int64, "shapes": [(5,)]} + assert torch.equal(data.field_data["tokens"][7], torch.tensor([1, 2, 3, 4, 5])) + + def boom(old, new): + raise RuntimeError("parser failed") + + with pytest.raises(RuntimeError, match="parser failed"): + data.apply_update([7], ["tokens"], {"tokens": [torch.tensor([9])]}, boom, False) + assert torch.equal(data.field_data["tokens"][7], torch.tensor([1, 2, 3, 4, 5])) + + +def test_apply_update_empty_stores_none(): + data = StorageUnitData() + data.put_data({"tokens": [torch.tensor([1])], "keep": [torch.tensor([2])]}, [3]) + described = data.apply_update([3], ["tokens"], None, None, True) + assert described["tokens"]["shapes"] is None + assert data.field_data["tokens"][3] is None + assert torch.equal(data.field_data["keep"][3], torch.tensor([2])) + + +def test_apply_update_missing_field_passes_none_as_old(): + data = StorageUnitData() + seen = [] + + def record(old, new): + seen.append(old) + return new + + data.apply_update([1], ["fresh"], {"fresh": [torch.tensor([8])]}, record, False) + assert seen == [None] + assert torch.equal(data.field_data["fresh"][1], torch.tensor([8])) + + +def test_build_update_field_schema_orders_shapes_across_units(): + """Units describe only their own rows; the batch schema must follow metadata order.""" + described = _build_update_field_schema( + [0, 1, 2, 3], + [ + ([0, 2], {"tokens": {"dtype": torch.int64, "shapes": [(4,), (4,)]}}), + ([1, 3], {"tokens": {"dtype": torch.int64, "shapes": [(9,), (9,)]}}), + ], + ) + + assert described["tokens"]["is_nested"] is True + assert described["tokens"]["shape"] is None + assert described["tokens"]["per_sample_shapes"] == [(4,), (9,), (4,), (9,)] + + +def test_build_update_field_schema_keeps_uniform_column_flat(): + described = _build_update_field_schema( + [0, 1], + [ + ([0], {"tokens": {"dtype": torch.int64, "shapes": [(4,)]}}), + ([1], {"tokens": {"dtype": torch.int64, "shapes": [(4,)]}}), + ], + ) + + assert described["tokens"] == { + "dtype": torch.int64, + "shape": (4,), + "is_nested": False, + "is_non_tensor": False, + } + + +def test_build_update_field_schema_marks_column_non_tensor_if_any_unit_is(): + described = _build_update_field_schema( + [0, 1], + [ + ([0], {"tokens": {"dtype": torch.int64, "shapes": [(4,)]}}), + ([1], {"tokens": {"dtype": None, "shapes": None}}), + ], + ) + + assert described["tokens"]["is_non_tensor"] is True + assert described["tokens"]["shape"] is None + + +def test_empty_marks_the_controller_field_non_tensor(): + """tq.kv_empty stores None, so the controller must stop describing the column as a tensor.""" + partition = DataPartitionStatus(partition_id="p") + partition._update_field_metadata( + [0], {"tokens": {"dtype": torch.int64, "shape": (5,), "is_nested": False, "is_non_tensor": False}} + ) + + partition._update_field_metadata( + [0], {"tokens": {"dtype": None, "shape": None, "is_nested": False, "is_non_tensor": True}} + ) + + tokens = partition.field_metadata["tokens"] + assert tokens.is_non_tensor is True + assert tokens.shape is None + assert tokens.to_batch_schema([0])["is_non_tensor"] is True + + +@pytest.mark.asyncio +@patch("transfer_queue.storage.managers.base.StorageClientFactory.create") +@patch.object(KVStorageManager, "_connect_to_controller", lambda self: None) +async def test_kv_backend_rejects_update(mock_create): + mock_create.return_value = MagicMock() + manager = KVStorageManager(controller_info=MagicMock(), config={"client_name": "YuanrongStorageClient"}) + meta = BatchMeta( + global_indexes=[0], + partition_ids=["p"], + field_schema={"x": {"dtype": torch.int64, "shape": (1,), "is_nested": False, "is_non_tensor": False}}, + production_status=np.ones(1, dtype=np.int8), + ) + with pytest.raises(NotImplementedError, match="kv_update is not supported for KV-based backends"): + await manager.update_data(meta, ["x"], empty=True) From bb0e3aeb99d3ddfa69ed0bae33360d123d468688 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Sun, 20 Sep 2026 15:34:52 +0800 Subject: [PATCH 2/9] [feat] Add kv_update for read-modify-write on an existing key kv_update rewrites selected fields of an existing key by running a custom parser(old, new) on the SimpleStorage unit that holds the row, so the payload never travels to the client and back. kv_empty stores None instead, keeping the key and its production status. The unit computes every value before writing, so a failing parser leaves the row unchanged, and it reports only the shape it stored per sample; the manager assembles one batch-level field_schema ordered to match the notified indexes. FieldMeta now accepts a column turning non-tensor, so an emptied field stops being described as a tensor. KV-based backends reject the call. Signed-off-by: OutstanderWang --- README.md | 1 + tests/e2e/test_kv_interface_e2e.py | 127 ++++++++++++ tests/test_kv_update.py | 17 +- transfer_queue/__init__.py | 8 + transfer_queue/client.py | 51 +++++ transfer_queue/controller.py | 50 +++-- transfer_queue/interface.py | 147 ++++++++++++++ transfer_queue/metadata.py | 16 +- transfer_queue/storage/managers/base.py | 35 ++++ .../managers/simple_storage_manager.py | 186 ++++++++++++++++++ transfer_queue/storage/simple_storage.py | 87 +++++++- transfer_queue/utils/zmq_utils.py | 3 + tutorial/02_kv_interface.py | 41 ++-- 13 files changed, 729 insertions(+), 40 deletions(-) diff --git a/README.md b/README.md index 3c08a6ad..62610857 100644 --- a/README.md +++ b/README.md @@ -116,6 +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**: Rewrite selected fields of an existing key with `parser(old, new)`, or `empty=True` to store `None`. `kv_empty()` is the same as `kv_update(..., empty=True)`. SimpleStorage only. - **(async_)kv_list**: List keys and tags (metadata) in a partition. - **(async_)kv_clear**: Remove key-value pairs from storage. diff --git a/tests/e2e/test_kv_interface_e2e.py b/tests/e2e/test_kv_interface_e2e.py index 271758ee..c3f9c447 100644 --- a/tests/e2e/test_kv_interface_e2e.py +++ b/tests/e2e/test_kv_interface_e2e.py @@ -1110,6 +1110,133 @@ def test_field_expansion_across_samples(self, controller, tq_api): tq_api.kv_clear(keys=keys, partition_id=partition_id) +class TestKVUpdateE2E: + """kv_update: SimpleStorage runs parsers on the unit; KV backends reject.""" + + def test_concat_then_empty(self, controller, tq_api, backend_name): + if backend_name != "SimpleStorage": + pytest.skip("parser-backed kv_update is implemented only for SimpleStorage") + + partition_id = "test_partition" + key = "sample_update" + + tq_api.kv_put(key=key, partition_id=partition_id, fields={"tokens": torch.tensor([1, 2, 3])}, tag={"step": 1}) + + def concat(old, new): + return torch.cat([old, new]) + + meta = tq_api.kv_update( + key=key, partition_id=partition_id, fields="tokens", values=torch.tensor([4, 5]), parser=concat + ) + assert "tokens" in meta.fields + + retrieved = tq_api.kv_batch_get(keys=key, partition_id=partition_id, select_fields="tokens") + assert_tensor_equal(retrieved["tokens"][0], torch.tensor([1, 2, 3, 4, 5])) + + tq_api.kv_empty(key=key, partition_id=partition_id, fields="tokens") + emptied = tq_api.kv_batch_get(keys=key, partition_id=partition_id, select_fields="tokens") + assert emptied["tokens"][0] is None + + partition = get_controller_partition(controller, partition_id) + col = partition.field_name_mapping["tokens"] + global_idx = partition.keys_mapping[key] + assert partition.production_status[global_idx, col] == 1 + # The column now holds None, so the controller must not still describe it as a tensor. + assert partition.field_metadata["tokens"].is_non_tensor is True + + tq_api.kv_clear(keys=key, partition_id=partition_id) + + def test_concat_prompt_and_response_keeps_field_name(self, tq_api, backend_name): + """prompt_ids already stored; update appends response_ids; field stays sequence_ids.""" + if backend_name != "SimpleStorage": + pytest.skip("parser-backed kv_update is implemented only for SimpleStorage") + + partition_id = "test_partition" + key = "sample_seq" + prompt_ids = torch.tensor([10, 11, 12]) + response_ids = torch.tensor([20, 21]) + + tq_api.kv_put(key=key, partition_id=partition_id, fields={"sequence_ids": prompt_ids}) + tq_api.kv_update( + key=key, + partition_id=partition_id, + fields="sequence_ids", + values=response_ids, + parser=lambda old, new: torch.cat([old, new]), + ) + + retrieved = tq_api.kv_batch_get(keys=key, partition_id=partition_id) + assert "sequence_ids" in retrieved + assert "prompt_ids" not in retrieved + assert "response_ids" not in retrieved + assert_tensor_equal(retrieved["sequence_ids"][0], torch.tensor([10, 11, 12, 20, 21])) + + tq_api.kv_clear(keys=key, partition_id=partition_id) + + def test_update_multiple_fields_at_once(self, tq_api, backend_name): + if backend_name != "SimpleStorage": + pytest.skip("parser-backed kv_update is implemented only for SimpleStorage") + + partition_id = "test_partition" + key = "sample_multi" + + tq_api.kv_put( + key=key, + partition_id=partition_id, + fields={"a": torch.tensor([1, 2]), "b": torch.tensor([10])}, + ) + tq_api.kv_update( + key=key, + partition_id=partition_id, + fields=["a", "b"], + values={"a": torch.tensor([3]), "b": torch.tensor([20, 30])}, + parser=lambda old, new: torch.cat([old, new]), + ) + + retrieved = tq_api.kv_batch_get(keys=key, partition_id=partition_id) + assert_tensor_equal(retrieved["a"][0], torch.tensor([1, 2, 3])) + assert_tensor_equal(retrieved["b"][0], torch.tensor([10, 20, 30])) + + tq_api.kv_empty(key=key, partition_id=partition_id, fields=["a", "b"]) + emptied = tq_api.kv_batch_get(keys=key, partition_id=partition_id) + assert emptied["a"][0] is None + assert emptied["b"][0] is None + + tq_api.kv_clear(keys=key, partition_id=partition_id) + + def test_missing_key_and_empty_with_values(self, tq_api, backend_name): + if backend_name != "SimpleStorage": + pytest.skip("parser-backed kv_update is implemented only for SimpleStorage") + + with pytest.raises(ValueError, match="must not specify values"): + tq_api.kv_update( + key="no_such", partition_id="test_partition", fields="tokens", values=torch.tensor([1]), empty=True + ) + with pytest.raises(ValueError, match="not found"): + tq_api.kv_update( + key="no_such", + partition_id="test_partition", + fields="tokens", + values=torch.tensor([1]), + parser=lambda old, new: new, + ) + + 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", + values=torch.tensor([2]), + parser=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"]) diff --git a/tests/test_kv_update.py b/tests/test_kv_update.py index 28a6694d..aed9d4a8 100644 --- a/tests/test_kv_update.py +++ b/tests/test_kv_update.py @@ -24,7 +24,7 @@ from transfer_queue.metadata import BatchMeta from transfer_queue.storage.managers.base import KVStorageManager from transfer_queue.storage.managers.simple_storage_manager import _build_update_field_schema -from transfer_queue.storage.simple_storage import StorageUnitData +from transfer_queue.storage.simple_storage import HybridStorageUnitData, StorageUnitData def _concat(old, new): @@ -125,6 +125,21 @@ def record(old, new): assert torch.equal(data.field_data["fresh"][1], torch.tensor([8])) +def test_apply_update_decodes_ssd_offloaded_old_value(tmp_path): + """With SSD offload on, the parser must see the stored tensor, not its file reference.""" + data = HybridStorageUnitData( + storage_size=4, threshold_bytes=64, ssd_path=str(tmp_path), run_id="run", unit_id="unit" + ) + prompt = torch.arange(32) + data.put_data({"tokens": [prompt]}, [0]) + assert data.ssd_active_values == 1, "precondition: the prompt must live on SSD" + + data.apply_update([0], ["tokens"], {"tokens": [torch.tensor([99])]}, _concat, False) + + assert torch.equal(data.get_data(["tokens"], [0])["tokens"][0], torch.cat([prompt, torch.tensor([99])])) + assert data.ssd_active_values == 1, "the replaced SSD file must be released, not leaked" + + def test_build_update_field_schema_orders_shapes_across_units(): """Units describe only their own rows; the batch schema must follow metadata order.""" described = _build_update_field_schema( diff --git a/transfer_queue/__init__.py b/transfer_queue/__init__.py index da7a540c..7f034662 100644 --- a/transfer_queue/__init__.py +++ b/transfer_queue/__init__.py @@ -22,8 +22,10 @@ async_kv_batch_get_by_meta, async_kv_batch_put, async_kv_clear, + async_kv_empty, async_kv_list, async_kv_put, + async_kv_update, close, dump_data_by_key, get_client, @@ -33,8 +35,10 @@ kv_batch_get_by_meta, kv_batch_put, kv_clear, + kv_empty, kv_list, kv_put, + kv_update, load_checkpoint, load_data_by_key, save_checkpoint, @@ -59,12 +63,16 @@ "kv_batch_get_by_meta", "kv_list", "kv_clear", + "kv_update", + "kv_empty", "async_kv_put", "async_kv_batch_put", "async_kv_batch_get", "async_kv_batch_get_by_meta", "async_kv_list", "async_kv_clear", + "async_kv_update", + "async_kv_empty", "KVBatchMeta", ] + [ diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 05a361ab..17ca3b97 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -492,6 +492,42 @@ async def async_put( return metadata + async def async_update( + self, + metadata: BatchMeta, + field_names: list[str], + values: TensorDict | None = None, + parser: Callable[[Any, Any], Any] | None = None, + empty: bool = False, + ) -> BatchMeta: + """Apply empty or parser(old, new) on SimpleStorage units for the named fields. + + Args: + metadata: Samples to update. The key must already exist. + field_names: Fields to rewrite. + values: New values aligned with ``metadata``, or None when empty. + parser: ``parser(old, new) -> stored``. Ignored when empty. + empty: Store None for each named field. + + Returns: + The same metadata with field_schema and production status updated. + + Raises: + NotImplementedError: If the storage backend is not SimpleStorage. + """ + if not hasattr(self, "storage_manager") or self.storage_manager is None: + raise RuntimeError( + f"[{self.client_id}]: Storage manager not initialized. " + "Call initialize_storage_manager() before performing storage operations." + ) + if not metadata or metadata.size == 0: + raise ValueError("metadata cannot be none or empty") + + field_schema = await self.storage_manager.update_data( + metadata, field_names, values=values, parser=parser, empty=empty + ) + return metadata.apply_field_schema(field_schema) + async def async_get_data(self, metadata: BatchMeta) -> TensorDict: """Asynchronously fetch data from storage units and organize into TensorDict. @@ -1389,6 +1425,7 @@ def wrapper(*args, **kwargs): # Bind internal sync wrappers. Public methods are defined explicitly below # to ensure proper type hints and documentation. self._put = _make_sync(self.async_put) + self._update = _make_sync(self.async_update) self._get_meta = _make_sync(self.async_get_meta) self._get_data = _make_sync(self.async_get_data) self._clear_partition = _make_sync(self.async_clear_partition) @@ -1579,6 +1616,20 @@ def put( """ return self._put(data=data, metadata=metadata, partition_id=partition_id, data_parser=data_parser) + def update( + self, + metadata: BatchMeta, + field_names: list[str], + values: TensorDict | None = None, + parser: Callable[[Any, Any], Any] | None = None, + empty: bool = False, + ) -> BatchMeta: + """Synchronously apply empty or parser(old, new) on SimpleStorage units. + + See ``async_update``. + """ + return self._update(metadata, field_names, values=values, parser=parser, empty=empty) + 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 c86c0824..ae8691a2 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -253,28 +253,36 @@ def update(self, incoming: dict[str, Any], incoming_global_indexes: list[int]) - # Update newly provided per_sample_shapes self.per_sample_shapes.update(new_per_sample_shapes) + elif new_is_non_tensor: + # Some samples now hold plain objects (e.g. kv_update empty stores None), so no + # single dtype/shape describes the column any more. dtype is kept so a later + # tensor write still has to agree with the one this field was created with. + self.is_non_tensor = True + self.is_nested = False + self.shape = None + self.per_sample_shapes.clear() + else: - if not new_is_non_tensor: - # newly input is regular tensor - new_shape = incoming.get("shape", None) - if new_shape is None: - raise ValueError("Receiving a regular tensor without 'shape'!") - if self.is_nested: - # we need to update incoming shape into per_sample_shapes - for gi in incoming_global_indexes: - self.per_sample_shapes[gi] = new_shape - else: - if self.is_non_tensor is not None and not self.is_non_tensor: - # original data is also regular tensor - assert self.shape is not None - if self.shape != new_shape: - for gi in self.global_indexes: - self.per_sample_shapes[gi] = self.shape - for gi in incoming_global_indexes: - self.per_sample_shapes[gi] = new_shape - - self.shape = None - self.is_nested = True + # newly input is regular tensor + new_shape = incoming.get("shape", None) + if new_shape is None: + raise ValueError("Receiving a regular tensor without 'shape'!") + if self.is_nested: + # we need to update incoming shape into per_sample_shapes + for gi in incoming_global_indexes: + self.per_sample_shapes[gi] = new_shape + else: + if self.is_non_tensor is not None and not self.is_non_tensor: + # original data is also regular tensor + assert self.shape is not None + if self.shape != new_shape: + for gi in self.global_indexes: + self.per_sample_shapes[gi] = self.shape + for gi in incoming_global_indexes: + self.per_sample_shapes[gi] = new_shape + + self.shape = None + self.is_nested = True self.global_indexes.update(incoming_global_indexes) diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 95a8d683..64f1d9a8 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -608,6 +608,61 @@ def kv_clear(keys: list[str] | str, partition_id: str) -> None: tq_client._run_coroutine(async_kv_clear(keys=keys, partition_id=partition_id)) +def kv_update( + key: str, + partition_id: str, + fields: str | list[str], + *, + values: Any | dict[str, Any] | None = None, + parser: Callable[[Any, Any], Any] | None = None, + empty: bool = False, +) -> KVBatchMeta: + """Update fields of an existing key by combining the stored value with a new one. + + A custom ``parser(old, new)`` runs once per sample per field on the SimpleStorage + unit that holds the row, then writes the return value back. ``empty=True`` + stores ``None`` instead and must not be given ``values`` or ``parser``. + ``kv_empty()`` is the same as ``kv_update(..., empty=True)``. + + Args: + key: Existing user-specified key. The key is not created if it is missing. + partition_id: Partition that holds the key. + fields: One field name, or a list of field names, to update. + values: New value for a single field, or ``{field: new}`` for several. + Must be omitted when ``empty=True``. + parser: ``parser(old, new) -> stored``. Required unless ``empty=True``. + Only SimpleStorage executes parsers; KV backends reject the call. + ``old`` is ``None`` when the field has never been written. + empty: If True, store None for each named field. + + Returns: + KVBatchMeta for the updated sample, including every field stored for it. + + Raises: + ValueError: If the key is missing, ``empty=True`` is given values or a + parser, a custom parser is given no values, or ``fields`` / ``values`` + do not line up. + TypeError: If ``parser`` is missing or not callable when ``empty`` is False. + NotImplementedError: If the storage backend is not SimpleStorage. + """ + tq_client = _maybe_create_tq_client() + return tq_client._run_coroutine( + async_kv_update( + key=key, + partition_id=partition_id, + fields=fields, + values=values, + parser=parser, + empty=empty, + ) + ) + + +def kv_empty(key: str, partition_id: str, fields: str | list[str]) -> KVBatchMeta: + """Store None for the named fields of an existing key. Same as kv_update(..., empty=True).""" + return kv_update(key=key, partition_id=partition_id, fields=fields, empty=True) + + # ==================== KV Interface API ==================== async def async_kv_put( key: str, @@ -989,6 +1044,98 @@ async def async_kv_clear(keys: list[str] | str, partition_id: str) -> None: await tq_client.async_clear_samples(batch_meta) +def _normalize_kv_update_args( + fields: str | list[str], + values: Any | dict[str, Any] | None, + parser: Callable[[Any, Any], Any] | None, + empty: bool = False, +) -> tuple[list[str], TensorDict | None, Callable[[Any, Any], Any] | None, bool]: + """Validate kv_update arguments and wrap new values as a one-row TensorDict.""" + if isinstance(fields, str): + field_names = [fields] + elif isinstance(fields, list) and fields and all(isinstance(f, str) for f in fields): + field_names = list(fields) + else: + raise TypeError("fields must be a field name or a non-empty list of field names") + + if empty: + if values is not None: + raise ValueError("kv_update with empty=True must not specify values") + if parser is not None: + raise ValueError("kv_update with empty=True must not specify parser") + return field_names, None, None, True + + if not callable(parser): + raise TypeError("parser must be callable unless empty=True") + if values is None: + raise ValueError("kv_update with a custom parser requires values") + + if isinstance(fields, str): + value_map = {fields: values} + else: + if not isinstance(values, dict): + raise TypeError("values must be a dict mapping field name to new value when updating multiple fields") + missing = [name for name in field_names if name not in values] + extra = [name for name in values if name not in field_names] + if missing or extra: + raise ValueError( + f"fields and values must name the same columns; fields={field_names}, extra={extra}, missing={missing}" + ) + value_map = values + + batch: dict[str, Any] = {} + for field_name, value in value_map.items(): + if isinstance(value, torch.Tensor): + if value.is_nested: + raise ValueError("nested tensors are not supported for single-key kv_update") + batch[field_name] = value.unsqueeze(0) + else: + batch[field_name] = NonTensorStack(value) + return field_names, TensorDict(batch, batch_size=[1]), parser, False + + +async def async_kv_update( + key: str, + partition_id: str, + fields: str | list[str], + *, + values: Any | dict[str, Any] | None = None, + parser: Callable[[Any, Any], Any] | None = None, + empty: bool = False, +) -> KVBatchMeta: + """Asynchronously update fields of an existing key. See ``kv_update``.""" + field_names, value_batch, bound_parser, use_empty = _normalize_kv_update_args(fields, values, parser, empty) + + tq_client = _maybe_create_tq_client() + batch_meta = await tq_client.async_kv_retrieve_meta(keys=[key], partition_id=partition_id, create=False) + + if batch_meta.size == 0: + raise ValueError("keys or partition were not found!") + if batch_meta.size != 1: + raise RuntimeError(f"Retrieved BatchMeta size {batch_meta.size} does not match with input `key` size of 1!") + + batch_meta = await tq_client.async_update( + metadata=batch_meta, + field_names=field_names, + values=value_batch, + parser=bound_parser, + empty=use_empty, + ) + + return KVBatchMeta( + keys=[key], + tags=batch_meta.custom_meta, + partition_id=partition_id, + fields=batch_meta.field_names, + extra_info=batch_meta.extra_info, + ) + + +async def async_kv_empty(key: str, partition_id: str, fields: str | list[str]) -> KVBatchMeta: + """Store None for the named fields of an existing key. Same as async_kv_update(..., empty=True).""" + return await async_kv_update(key=key, partition_id=partition_id, fields=fields, empty=True) + + # ==================== Low-Level Native API ==================== # For low-level API support, please refer to transfer_queue/client.py for details. def get_client(): diff --git a/transfer_queue/metadata.py b/transfer_queue/metadata.py index e89c1531..ab395e07 100644 --- a/transfer_queue/metadata.py +++ b/transfer_queue/metadata.py @@ -461,10 +461,20 @@ def add_fields(self, tensor_dict: TensorDict, set_all_ready: bool = True) -> "Ba if batch_size != self.size: raise ValueError(f"add_fields batch size mismatch: self.size={self.size} vs tensor_dict={batch_size}") - field_schema = extract_field_schema(tensor_dict) + return self.apply_field_schema(extract_field_schema(tensor_dict), set_all_ready=set_all_ready) - for key, value in field_schema.items(): - self.field_schema[key] = value + def apply_field_schema(self, field_schema: dict[str, dict[str, Any]], set_all_ready: bool = True) -> "BatchMeta": + """Merge a field_schema into this batch in place. + + The only place that keeps ``field_names`` and ``is_ready`` in step with + ``field_schema`` and ``production_status``; callers holding a schema from storage + should go through here rather than writing the derived state themselves. + + Args: + field_schema: Per-field metadata to add or overwrite. + set_all_ready (bool): If True, set all production_status to READY_FOR_CONSUME. Default is True. + """ + self.field_schema.update(field_schema) if set_all_ready: self.production_status[:] = 1 diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index 15d75627..4b5e8897 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -354,6 +354,28 @@ async def clear_data(self, metadata: BatchMeta) -> None: """ raise NotImplementedError("Subclasses must implement clear_data") + async def update_data( + self, + metadata: BatchMeta, + field_names: list[str], + values: TensorDict | None = None, + parser: Callable[[Any, Any], Any] | None = None, + empty: bool = False, + ) -> dict[str, dict[str, Any]]: + """Apply empty or parser(old, new) on the unit that holds each sample. + + Args: + metadata: Samples to update. + field_names: Fields to rewrite. + values: New values as a TensorDict aligned with ``metadata``, or None for empty. + parser: Called per sample per field as ``parser(old, new)``. Ignored when empty. + empty: If True, store None for each named field and ignore values/parser. + + Returns: + field_schema of the stored values, keyed by field name. + """ + raise NotImplementedError(f"{self.__class__.__name__} does not support update_data") + async def save_checkpoint(self, checkpoint_dir: str) -> None: """Save storage state into checkpoint_dir. @@ -800,6 +822,19 @@ async def put_data( per_field_custom_backend_meta, ) + async def update_data( + self, + metadata: BatchMeta, + field_names: list[str], + values: TensorDict | None = None, + parser: Callable[[Any, Any], Any] | None = None, + empty: bool = False, + ) -> dict[str, dict[str, Any]]: + """kv_update is only implemented for SimpleStorage.""" + raise NotImplementedError( + "kv_update is not supported for KV-based backends (MooncakeStore, Yuanrong, RayStore)." + ) + async def get_data(self, metadata: BatchMeta) -> TensorDict: """ Retrieve tensor data from the backend storage. diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py index 2d9a7a4c..f9c52844 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -91,6 +91,48 @@ def _describe_unit_state(body: dict[str, Any]) -> str: return f"verdict=unit_serving_again ({' '.join(parts)})" +def _build_update_field_schema( + global_indexes: list[int], + per_unit: list[tuple[list[int], dict[str, dict[str, Any]]]], +) -> dict[str, dict[str, Any]]: + """Assemble one batch-level field_schema from each unit's per-sample description. + + A unit only describes its own rows, so shapes are collected per global index and + then ordered to match ``global_indexes``: the controller reads per_sample_shapes + positionally against the indexes it is notified about. + """ + dtypes: dict[str, Any] = {} + shapes: dict[str, dict[int, tuple]] = defaultdict(dict) + non_tensor: set[str] = set() + for unit_indexes, described in per_unit: + for field_name, description in described.items(): + unit_shapes = description["shapes"] + if unit_shapes is None: + non_tensor.add(field_name) + continue + dtypes[field_name] = description["dtype"] + # Serialization turns the unit's tuples into lists; shapes are compared and hashed here. + shapes[field_name].update(zip(unit_indexes, (tuple(s) for s in unit_shapes), strict=True)) + + field_schema: dict[str, dict[str, Any]] = {} + for field_name in non_tensor | set(shapes): + # A column is only a tensor column when every unit stored tensors for it. + if field_name in non_tensor: + field_schema[field_name] = {"dtype": None, "shape": None, "is_nested": False, "is_non_tensor": True} + continue + ordered = [shapes[field_name][gi] for gi in global_indexes] + is_nested = len(set(ordered)) > 1 + field_schema[field_name] = { + "dtype": dtypes[field_name], + "shape": None if is_nested else ordered[0], + "is_nested": is_nested, + "is_non_tensor": False, + } + if is_nested: + field_schema[field_name]["per_sample_shapes"] = ordered + return field_schema + + _SU_SUBDIR = "simple_storage" _SU_INFO_FILE = "storage_unit_info.json" @@ -497,6 +539,86 @@ async def put_data( field_schema, ) + async def update_data( + self, + metadata: BatchMeta, + field_names: list[str], + values: TensorDict | None = None, + parser: Callable[[Any, Any], Any] | None = None, + empty: bool = False, + ) -> dict[str, dict[str, Any]]: + """Send an update to each hashed storage unit and notify the controller. + + The unit reads the old value, applies empty or parser(old, new), and reports + the shape it stored per sample; this assembles one batch-level field_schema + so production status stays ready. + """ + logger.debug(f"[{self.storage_manager_id}]: receive update_data request, updating {metadata.size} samples.") + + if metadata.size == 0: + return {} + if empty: + if values is not None or parser is not None: + raise ValueError("empty update must not include values or a parser") + else: + if values is None or parser is None: + raise ValueError("update requires values and a parser unless empty is set") + if values.batch_size[0] != metadata.size: + raise ValueError( + f"Batch size of values ({values.batch_size[0]}) does not match metadata size ({metadata.size})" + ) + + routing = self._group_by_hash(metadata.global_indexes) + # Parser-backed updates are not replayed: parser(old, new) reads the stored value, so a + # second attempt after a lost answer would fold the new value in twice. Empty is a + # plain overwrite and stays retryable. + max_attempts = 1 if parser is not None else None + tasks = [] + for su_id, group in routing.items(): + storage_data = None + if values is not None: + storage_data = {f: self._select_by_positions(values[f], group.batch_positions) for f in field_names} + tasks.append( + self._request_with_retry( + "update", + su_id, + f"samples={len(group.global_indexes)} fields={field_names}", + partial( + self._update_to_single_storage_unit, + group.global_indexes, + field_names, + storage_data, + parser, + empty, + target_storage_unit=su_id, + ), + max_attempts=max_attempts, + ) + ) + + try: + described = await asyncio.gather(*tasks) + except Exception as e: + logger.error( + f"[{self.storage_manager_id}]: update_data failed. " + f"partition_id={metadata.partition_ids[0]}, " + f"num_samples={metadata.size}, " + f"num_storage_units={len(routing)}, " + f"error={type(e).__name__}: {e}" + ) + raise + + field_schema = _build_update_field_schema( + metadata.global_indexes, list(zip([g.global_indexes for g in routing.values()], described, strict=True)) + ) + + await self.notify_data_update( + metadata.partition_ids[0], + metadata.global_indexes, + field_schema, + ) + return field_schema + @with_storage_unit_socket async def _put_to_single_storage_unit( self, @@ -554,6 +676,70 @@ async def _put_to_single_storage_unit( ) raise RuntimeError(f"Error in put to storage unit {target_storage_unit}: {type(e).__name__}: {e}") from e + @with_storage_unit_socket + async def _update_to_single_storage_unit( + self, + global_indexes: list[int], + field_names: list[str], + storage_data: dict[str, Any] | None, + parser: Callable[[Any, Any], Any] | None, + empty: bool, + target_storage_unit: str, + socket: zmq.Socket = None, + ) -> dict[str, dict[str, Any]]: + """Send an update to one storage unit and return what it stored, per field.""" + body: dict[str, Any] = { + "global_indexes": global_indexes, + "fields": field_names, + "empty": empty, + } + if not empty: + body["data"] = storage_data + body["parser"] = parser + + request_msg = ZMQMessage.create( + request_type=ZMQRequestType.UPDATE_DATA, # type: ignore[arg-type] + sender_id=self.storage_manager_id, + receiver_id=target_storage_unit, + body=body, + ) + + serialized_bytes = 0 + started = time.perf_counter() + try: + frames = request_msg.serialize() + serialized_bytes = sum(frame_nbytes(frame) or 0 for frame in frames) + await socket.send_multipart(frames, copy=False) + messages = await socket.recv_multipart(copy=False) + response_msg = ZMQMessage.deserialize(messages) + + if response_msg.request_type != ZMQRequestType.UPDATE_DATA_RESPONSE: + raise RuntimeError( + f"Failed to update data on storage unit {target_storage_unit}: " + f"{response_msg.body.get('message', 'Unknown error')}" + ) + log_heavy_operation( + self.storage_manager_id, + "update", + time.perf_counter() - started, + serialized_bytes, + f"to {target_storage_unit} at {self._describe_storage_unit(target_storage_unit)} " + f"samples={len(global_indexes)} fields={field_names}", + ) + return response_msg.body.get("stored_shapes", {}) + except zmq.error.Again as e: + raise StorageUnitTimeout( + f"no answer in {TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT}s during update to storage unit " + f"{target_storage_unit} at {self._describe_storage_unit(target_storage_unit)}; " + f"serialized_mb={serialized_bytes / 2**20:.1f}" + ) from e + except Exception as e: + logger.error( + f"[{self.storage_manager_id}]: Unexpected error during update to storage unit " + f"{target_storage_unit}: {type(e).__name__}: {e}" + ) + raise RuntimeError(f"Error in update to storage unit {target_storage_unit}: {type(e).__name__}: {e}") from e + @staticmethod def _pack_field_values(values: list) -> torch.Tensor | NonTensorStack: """ diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py index 258d9b7e..d7bc184d 100644 --- a/transfer_queue/storage/simple_storage.py +++ b/transfer_queue/storage/simple_storage.py @@ -39,7 +39,7 @@ from dataclasses import dataclass from pathlib import Path from threading import Event, Thread -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Callable from uuid import uuid4 import numpy as np @@ -112,6 +112,18 @@ class _SSDValueRef: TQ_ACCEPT_PROBE_INTERVAL = float(os.environ.get("TQ_ACCEPT_PROBE_INTERVAL", 0)) +def _describe_stored_values(values: list) -> dict[str, Any]: + """Describe per-sample stored values for the manager. + + A unit only holds part of a batch, so it reports one shape per sample and + leaves ``shapes=None`` for a column it cannot describe as tensors; only the + manager sees the whole batch and can turn this into a field_schema. + """ + if not all(isinstance(v, torch.Tensor) for v in values): + return {"dtype": None, "shapes": None} + return {"dtype": values[0].dtype, "shapes": [tuple(v.shape) for v in values]} + + class StorageKeyNotFoundError(KeyError): """Raised when a requested global index is absent from a storage unit. @@ -217,6 +229,40 @@ def put_data(self, field_data: dict[str, Any], global_indexes: list) -> None: field_dict[key] = val self._active_keys.update(global_indexes) + def apply_update( + self, + global_indexes: list[int], + fields: list[str], + new_data: dict[str, Any] | None, + parser: Callable[[Any, Any], Any] | None, + use_empty: bool, + ) -> dict[str, dict[str, Any]]: + """Read each (field, index), compute the stored value, then write once. + + Missing fields yield ``old=None``. All parser calls finish before any write, + so a raise leaves storage unchanged. + + Returns: + Per-sample description of the stored values, keyed by field name. + """ + computed: dict[str, list] = {} + if use_empty: + computed = {field: [None] * len(global_indexes) for field in fields} + else: + if parser is None or new_data is None: + raise TypeError("apply_update requires a parser and new values unless use_empty is set") + for field in fields: + # Read through get_data so SSD-offloaded values reach the parser decoded, + # not as file references; indexes never written keep old=None. + stored = self.field_data.get(field, {}) + present = [idx for idx in global_indexes if idx in stored] + old_values = dict(zip(present, self.get_data([field], present)[field], strict=True)) if present else {} + computed[field] = [ + parser(old_values.get(idx), new_data[field][i]) for i, idx in enumerate(global_indexes) + ] + self.put_data(computed, global_indexes) + return {field: _describe_stored_values(values) for field, values in computed.items()} + def clear(self, keys: list[int]) -> None: """Remove data at given global index keys, immediately freeing memory. @@ -1042,6 +1088,9 @@ def _process_one_worker_request(self, worker_socket: zmq.Socket, monitor: Any) - if operation == ZMQRequestType.PUT_DATA: # type: ignore[arg-type] with monitor.measure(op_type="PUT_DATA"): response_msg = self._handle_put(request_msg) + elif operation == ZMQRequestType.UPDATE_DATA: # type: ignore[arg-type] + with monitor.measure(op_type="UPDATE_DATA"): + response_msg = self._handle_update(request_msg) elif operation == ZMQRequestType.GET_DATA: # type: ignore[arg-type] with monitor.measure(op_type="GET_DATA"): response_msg = self._handle_get(request_msg) @@ -1175,6 +1224,40 @@ def _handle_put(self, data_parts: ZMQMessage) -> ZMQMessage: }, ) + def _handle_update(self, data_parts: ZMQMessage) -> ZMQMessage: + """Read each field, apply empty or parser(old, new), write once, describe what was stored.""" + try: + global_indexes = data_parts.body["global_indexes"] + fields = data_parts.body["fields"] + use_empty = bool(data_parts.body.get("empty", False)) + parser = data_parts.body.get("parser") + new_data = data_parts.body.get("data") + + with limit_pytorch_auto_parallel_threads( + target_num_threads=TQ_NUM_THREADS, info=f"[{self.storage_unit_id}] _handle_update" + ): + if not use_empty: + if not callable(parser): + raise TypeError(f"update parser must be callable, got {type(parser).__name__}") + if not isinstance(new_data, dict): + raise TypeError("update data must be a dict of new field values") + stored_shapes = self.storage_data.apply_update(global_indexes, fields, new_data, parser, use_empty) + + return ZMQMessage.create( + request_type=ZMQRequestType.UPDATE_DATA_RESPONSE, # type: ignore[arg-type] + sender_id=self.storage_unit_id, + body={"stored_shapes": stored_shapes}, + ) + except Exception as e: + return ZMQMessage.create( + request_type=ZMQRequestType.UPDATE_ERROR, # type: ignore[arg-type] + sender_id=self.storage_unit_id, + body={ + "message": f"Failed to update data in storage unit id " + f"#{self.storage_unit_id}, detail error message: {str(e)}" + }, + ) + def _handle_get(self, data_parts: ZMQMessage) -> ZMQMessage: """ Handle get request, return data from storage unit. @@ -1293,7 +1376,7 @@ def _handle_get_metrics(self) -> ZMQMessage: # Include per-operation stats if Prometheus metrics are enabled if self._metrics is not None: op_stats = {} - for op_type in ("PUT_DATA", "GET_DATA", "CLEAR_DATA", "DUMP_ROWS", "LOAD_ROWS"): + for op_type in ("PUT_DATA", "UPDATE_DATA", "GET_DATA", "CLEAR_DATA", "DUMP_ROWS", "LOAD_ROWS"): try: hist = self._metrics.request_duration.labels(op_type=op_type) counter = self._metrics.request_total.labels(op_type=op_type) diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index 06c8d792..06ee0f8e 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -69,8 +69,10 @@ class ZMQRequestType(ExplicitEnum): # DATA_OPERATION GET_DATA = "GET" PUT_DATA = "PUT" + UPDATE_DATA = "UPDATE" GET_DATA_RESPONSE = "GET_DATA_RESPONSE" PUT_DATA_RESPONSE = "PUT_DATA_RESPONSE" + UPDATE_DATA_RESPONSE = "UPDATE_DATA_RESPONSE" CLEAR_DATA = "CLEAR_DATA" CLEAR_DATA_RESPONSE = "CLEAR_DATA_RESPONSE" @@ -78,6 +80,7 @@ class ZMQRequestType(ExplicitEnum): PUT_GET_ERROR = "PUT_GET_ERROR" PUT_ERROR = "PUT_ERROR" GET_ERROR = "GET_ERROR" + UPDATE_ERROR = "UPDATE_ERROR" CLEAR_DATA_ERROR = "CLEAR_DATA_ERROR" # META_OPERATION diff --git a/tutorial/02_kv_interface.py b/tutorial/02_kv_interface.py index 778739c2..601f80d8 100644 --- a/tutorial/02_kv_interface.py +++ b/tutorial/02_kv_interface.py @@ -56,10 +56,10 @@ def demonstrate_kv_api(): """ Demonstrate the Key-Value (KV) semantic API: - kv_put & kv_batch_put -> kv_list -> kv_batch_get -> kv_clear + kv_put & kv_batch_put -> kv_update -> kv_list -> kv_batch_get -> kv_clear """ print("=" * 80) - print("Key-Value Semantic API Demo: kv_put/kv_batch_put → kv_list → kv_batch_get → kv_clear") + print("Key-Value Semantic API Demo: kv_put/kv_batch_put → kv_update → kv_list → kv_batch_get → kv_clear") print("=" * 80) # Step 1: Put a single key-value pair with kv_put @@ -153,8 +153,22 @@ def demonstrate_kv_api(): tq.kv_put(key=key_for_update_tags, partition_id=partition_id, fields=None, tag=tag_update) print(f" ✓ Update success: Samples '0_0' now has tag as {tag_update}.") - # Step 5: List all keys and tags in a partition - print("\n[Step 5] Listing all keys and tags in partition...") + # Step 5: Concatenate a field with kv_update, then empty it + print("\n[Step 5] Updating a field with kv_update (concat, then empty)...") + tq.kv_put(key=key, partition_id=partition_id, fields={"scratch": torch.tensor([1, 2])}) + tq.kv_update( + key=key, + partition_id=partition_id, + fields="scratch", + values=torch.tensor([3, 4]), + parser=lambda old, new: torch.cat([old, new]), + ) + print(" ✓ kv_update concat: scratch of '0_0' is now [1, 2, 3, 4].") + tq.kv_empty(key=key, partition_id=partition_id, fields="scratch") + print(" ✓ tq.kv_empty: scratch of '0_0' is stored as None (key remains).") + + # Step 6: List all keys and tags in a partition + print("\n[Step 6] Listing all keys and tags in partition...") partition_info = tq.kv_list() print(f" Found {len(partition_info.keys())} partitions: '{list(partition_info.keys())}'") @@ -162,16 +176,16 @@ def demonstrate_kv_api(): for k, t in keys_and_tags.items(): print(f"Partition: {pid}, - key='{k}' | tag={t}") - # Step 6: Retrieve specific fields using kv_batch_get - print("\n[Step 6] Retrieving specific fields (Column) with kv_batch_get...") + # Step 7: Retrieve specific fields using kv_batch_get + print("\n[Step 7] Retrieving specific fields (Column) with kv_batch_get...") print(" Fetching only 'input_ids' to save bandwidth (ignoring 'attention_mask' and 'response').") all_keys = list(partition_info[partition_id].keys()) retrieved_input_ids = tq.kv_batch_get(keys=all_keys, partition_id=partition_id, select_fields="input_ids") print(f" ✓ Successfully retrieved only {list(retrieved_input_ids.keys())} field for all samples.") - # # Step 7: Retrieve all fields using kv_batch_get - print("\n[Step 7] Retrieving all fields with kv_batch_get...") + # Step 8: Retrieve all fields using kv_batch_get + print("\n[Step 8] Retrieving all fields with kv_batch_get...") retrieved_all = tq.kv_batch_get(keys=all_keys, partition_id=partition_id) print(f" Retrieved all fields for {all_keys}:") print(f" Fields: {list(retrieved_all.keys())}") @@ -179,8 +193,8 @@ def demonstrate_kv_api(): f" Note: We cannot retrieve fields {list(response_batch.keys())}, since they only available in {append_keys}" ) - # Step 8: Clear specific keys - print("\n[Step 8] Clearing keys from partition...") + # Step 9: Clear specific keys + print("\n[Step 9] Clearing keys from partition...") keys_to_clear = all_keys[:2] # Delete the first 2 keys tq.kv_clear(keys=keys_to_clear, partition_id=partition_id) print(f" ✓ Cleared keys: {keys_to_clear}") @@ -202,9 +216,10 @@ 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_batch_get - Retrieve samples (by keys), supporting column selection (by fields) - 4. (async_)kv_list - List keys and tags (metadata) in a partition - 5. (async_)kv_clear - Remove key-value pairs from storage + 3. (async_)kv_update - Rewrite fields with parser(old, new), or empty=True / tq.kv_empty to store None + 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 from storage Key Features: ✓ Redis-style Semantics - Familiar KV interface (Put/Get/List) for zero learning curve From 8eedd22c29c9e3bf6ff4d1ac3e43be6e387792e5 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Wed, 23 Sep 2026 14:55:08 +0800 Subject: [PATCH 3/9] [fix] Name the unsupported backend when kv_update is rejected kv_update needs to run parser(old, new) where the value lives, which only SimpleStorage can do. The refusal used to list backend names in a string that would go stale as soon as another KV backend was added, and it never told the caller which backend they were actually on. Report the concrete manager class instead, and cover Mooncake, Yuanrong and Ray individually rather than only the shared base class. Signed-off-by: OutstanderWang --- tests/test_kv_update.py | 18 +++++++++++++++--- transfer_queue/storage/managers/base.py | 8 ++++++-- 2 files changed, 21 insertions(+), 5 deletions(-) diff --git a/tests/test_kv_update.py b/tests/test_kv_update.py index aed9d4a8..6fd4e741 100644 --- a/tests/test_kv_update.py +++ b/tests/test_kv_update.py @@ -23,7 +23,10 @@ from transfer_queue.interface import _normalize_kv_update_args from transfer_queue.metadata import BatchMeta from transfer_queue.storage.managers.base import KVStorageManager +from transfer_queue.storage.managers.mooncake_manager import MooncakeStorageManager +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 @@ -203,16 +206,25 @@ def test_empty_marks_the_controller_field_non_tensor(): @pytest.mark.asyncio +@pytest.mark.parametrize( + "manager_cls", [MooncakeStorageManager, YuanrongStorageManager, RayStorageManager, KVStorageManager] +) @patch("transfer_queue.storage.managers.base.StorageClientFactory.create") @patch.object(KVStorageManager, "_connect_to_controller", lambda self: None) -async def test_kv_backend_rejects_update(mock_create): +async def test_kv_backends_reject_update(mock_create, manager_cls): + """Every KV backend must refuse kv_update by name; only SimpleStorage implements it.""" mock_create.return_value = MagicMock() - manager = KVStorageManager(controller_info=MagicMock(), config={"client_name": "YuanrongStorageClient"}) + # Each manager validates its own config before reaching update_data. + config = { + KVStorageManager: {"client_name": "YuanrongStorageClient"}, + YuanrongStorageManager: {"worker_port": 31501}, + }.get(manager_cls, {}) + manager = manager_cls(controller_info=MagicMock(), config=config) meta = BatchMeta( global_indexes=[0], partition_ids=["p"], field_schema={"x": {"dtype": torch.int64, "shape": (1,), "is_nested": False, "is_non_tensor": False}}, production_status=np.ones(1, dtype=np.int8), ) - with pytest.raises(NotImplementedError, match="kv_update is not supported for KV-based backends"): + with pytest.raises(NotImplementedError, match=f"not supported by {manager_cls.__name__}"): await manager.update_data(meta, ["x"], empty=True) diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index 4b5e8897..72585f0f 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -830,9 +830,13 @@ async def update_data( parser: Callable[[Any, Any], Any] | None = None, empty: bool = False, ) -> dict[str, dict[str, Any]]: - """kv_update is only implemented for SimpleStorage.""" + """kv_update is only implemented for SimpleStorage. + + A KV backend stores each sample-field under its own key and offers no hook to run + parser(old, new) where the value lives, so read-modify-write cannot be made atomic. + """ raise NotImplementedError( - "kv_update is not supported for KV-based backends (MooncakeStore, Yuanrong, RayStore)." + f"kv_update is not supported by {type(self).__name__}; it requires the SimpleStorage backend." ) async def get_data(self, metadata: BatchMeta) -> TensorDict: From 2c01da6bda9502c17c8cc03ed23d9324ef927298 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Wed, 23 Sep 2026 14:55:20 +0800 Subject: [PATCH 4/9] [test] Add a cluster end-to-end example for kv_update The unit tests stub out the transport, so they cannot show that a parser really travels to the storage unit or that a batch spanning several units reports its per-sample shapes back to the controller in the right order. This script exercises both against a live Ray cluster, alongside kv_empty, multi-field updates, non-tensor payloads, parser failure atomicity, argument validation, the async entry points, and updates issued from remote workers. It refuses to run when a controller already exists, so it can never attach to somebody else's deployment and clear their keys on the way out. Signed-off-by: OutstanderWang --- recipe/kv_update_e2e.py | 389 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 389 insertions(+) create mode 100644 recipe/kv_update_e2e.py diff --git a/recipe/kv_update_e2e.py b/recipe/kv_update_e2e.py new file mode 100644 index 00000000..c25c6bc9 --- /dev/null +++ b/recipe/kv_update_e2e.py @@ -0,0 +1,389 @@ +# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# Copyright 2025 The TransferQueue Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""End-to-end exercise of kv_update / kv_empty against a real Ray cluster. + +kv_update rewrites fields of an existing key by running ``parser(old, new)`` on the +SimpleStorage unit that holds the row, so nothing here works unless the parser +actually travels to the unit and the resulting schema travels back to the +controller. Every check below reads its result back through the public API, and +the multi-unit ones also inspect controller metadata, which is the part a +data-only assertion cannot see. + +Run it as a Ray job so the driver lives on the cluster: + + ray job submit --address http://: --working-dir . \ + -- python recipe/kv_update_e2e.py + +Locally, ``python recipe/kv_update_e2e.py`` starts its own cluster instead. +""" + +import asyncio +import os +import traceback + +# Disable Ray's cross-worker log deduplication before importing Ray itself, +# otherwise worker-side prints get folded into "[repeated Nx across cluster]". +os.environ.setdefault("RAY_DEDUP_LOGS", "0") + +import ray +import torch +from omegaconf import OmegaConf +from tensordict import TensorDict + +import transfer_queue as tq + +PARTITION = "kv_update_e2e" + +# Several units so the hash routing actually splits a batch; the cross-unit checks +# below are meaningless with a single unit. +CONFIG = { + "controller": {"polling_mode": True}, + "backend": { + "storage_backend": "SimpleStorage", + "SimpleStorage": {"total_storage_size": 512, "num_data_storage_units": 4}, + }, +} + + +def concat(old, new): + """Append new tokens to what is already stored, keeping the field name.""" + return new if old is None else torch.cat([old, new]) + + +# ==================== checks ==================== + + +def check_concat_keeps_field_name(): + """The rollout case: prompt_ids in place, response_ids appended, one field.""" + key = "concat" + tq.kv_put(key=key, partition_id=PARTITION, fields={"sequence_ids": torch.tensor([10, 11, 12])}) + + tq.kv_update( + key=key, + partition_id=PARTITION, + fields="sequence_ids", + values=torch.tensor([20, 21]), + parser=concat, + ) + + got = tq.kv_batch_get(keys=key, partition_id=PARTITION) + assert torch.equal(got["sequence_ids"][0], torch.tensor([10, 11, 12, 20, 21])), got["sequence_ids"][0] + assert list(got.keys()) == ["sequence_ids"], f"update must not add columns: {list(got.keys())}" + + +def check_parser_sees_none_for_a_new_field(): + """A field the key never held reaches the parser as old=None. + + The parser runs in the storage unit's process, so what it observed can only + come back inside the value it returns; a closure variable would stay untouched + here in the driver. + """ + key = "fresh_field" + tq.kv_put(key=key, partition_id=PARTITION, fields={"anchor": torch.tensor([1])}) + + def report_old(old, new): + return {"old_was_none": old is None, "new": new.tolist()} + + tq.kv_update(key=key, partition_id=PARTITION, fields="logprobs", values=torch.tensor([7, 8]), parser=report_old) + + got = tq.kv_batch_get(keys=key, partition_id=PARTITION, select_fields="logprobs") + assert got["logprobs"][0] == {"old_was_none": True, "new": [7, 8]}, got["logprobs"][0] + + +def check_multiple_fields_in_one_call(): + key = "multi_field" + tq.kv_put( + key=key, + partition_id=PARTITION, + fields={"a": torch.tensor([1, 2]), "b": torch.tensor([10])}, + ) + + tq.kv_update( + key=key, + partition_id=PARTITION, + fields=["a", "b"], + values={"a": torch.tensor([3]), "b": torch.tensor([20, 30])}, + parser=concat, + ) + + got = tq.kv_batch_get(keys=key, partition_id=PARTITION) + assert torch.equal(got["a"][0], torch.tensor([1, 2, 3])), got["a"][0] + assert torch.equal(got["b"][0], torch.tensor([10, 20, 30])), got["b"][0] + + +def check_non_tensor_payload(): + """Parsers are not tensor-only: a plain object round-trips the same way.""" + key = "non_tensor" + tq.kv_put(key=key, partition_id=PARTITION, fields={"stats": {"step": 1}}) + + tq.kv_update( + key=key, + partition_id=PARTITION, + fields="stats", + values={"reward": 0.5}, + parser=lambda old, new: {**old, **new}, + ) + + got = tq.kv_batch_get(keys=key, partition_id=PARTITION, select_fields="stats") + assert got["stats"][0] == {"step": 1, "reward": 0.5}, got["stats"][0] + + +def check_empty_keeps_the_key_and_updates_the_controller(): + """kv_empty stores None; the key stays produced and stops being a tensor column.""" + key = "to_empty" + tq.kv_put(key=key, partition_id=PARTITION, fields={"tokens": torch.tensor([1, 2, 3])}) + + tq.kv_empty(key=key, partition_id=PARTITION, fields="tokens") + + got = tq.kv_batch_get(keys=key, partition_id=PARTITION, select_fields="tokens") + assert got["tokens"][0] is None, got["tokens"][0] + assert key in tq.kv_list(partition_id=PARTITION)[PARTITION], "kv_empty must not delete the key" + + partition = partition_snapshot() + col = partition.field_name_mapping["tokens"] + row = partition.keys_mapping[key] + assert partition.production_status[row, col] == 1, "an emptied field stays produced" + assert partition.field_metadata["tokens"].is_non_tensor is True, "controller still calls the column a tensor" + + +def check_failed_parser_leaves_the_row_unchanged(): + """The unit computes every value before writing, so a raise is a no-op.""" + key = "atomic" + tq.kv_put(key=key, partition_id=PARTITION, fields={"tokens": torch.tensor([1, 2, 3])}) + + def boom(old, new): + raise RuntimeError("parser failed on purpose") + + try: + tq.kv_update(key=key, partition_id=PARTITION, fields="tokens", values=torch.tensor([4]), parser=boom) + except Exception as e: + assert "parser failed on purpose" in str(e), f"unexpected error: {e}" + else: + raise AssertionError("a raising parser must not succeed") + + got = tq.kv_batch_get(keys=key, partition_id=PARTITION, select_fields="tokens") + assert torch.equal(got["tokens"][0], torch.tensor([1, 2, 3])), got["tokens"][0] + + +def check_rejected_arguments(): + key = "rejects" + tq.kv_put(key=key, partition_id=PARTITION, fields={"tokens": torch.tensor([1])}) + + def expect(exc_type, match, fn): + try: + fn() + except exc_type as e: + assert match in str(e), f"expected {match!r} in {e!r}" + else: + raise AssertionError(f"expected {exc_type.__name__} containing {match!r}") + + expect( + ValueError, + "must not specify values", + lambda: tq.kv_update(key=key, partition_id=PARTITION, fields="tokens", values=torch.tensor([1]), empty=True), + ) + expect( + ValueError, + "must not specify parser", + lambda: tq.kv_update(key=key, partition_id=PARTITION, fields="tokens", parser=concat, empty=True), + ) + expect( + TypeError, + "parser must be callable", + lambda: tq.kv_update(key=key, partition_id=PARTITION, fields="tokens", values=torch.tensor([1])), + ) + expect( + ValueError, + "requires values", + lambda: tq.kv_update(key=key, partition_id=PARTITION, fields="tokens", parser=concat), + ) + expect( + ValueError, + "same columns", + lambda: tq.kv_update( + key=key, partition_id=PARTITION, fields=["tokens", "other"], values={"tokens": 1}, parser=concat + ), + ) + expect( + ValueError, + "not found", + lambda: tq.kv_update( + key="never_written", partition_id=PARTITION, fields="tokens", values=torch.tensor([1]), parser=concat + ), + ) + + +def check_async_variants(): + key = "async_key" + tq.kv_put(key=key, partition_id=PARTITION, fields={"tokens": torch.tensor([1, 2])}) + + asyncio.run( + tq.async_kv_update(key=key, partition_id=PARTITION, fields="tokens", values=torch.tensor([3]), parser=concat) + ) + got = tq.kv_batch_get(keys=key, partition_id=PARTITION, select_fields="tokens") + assert torch.equal(got["tokens"][0], torch.tensor([1, 2, 3])), got["tokens"][0] + + asyncio.run(tq.async_kv_empty(key=key, partition_id=PARTITION, fields="tokens")) + got = tq.kv_batch_get(keys=key, partition_id=PARTITION, select_fields="tokens") + assert got["tokens"][0] is None, got["tokens"][0] + + +def check_multi_sample_update_across_units(): + """One update spanning several units must report per-sample shapes in batch order. + + Each key starts with a different length, so concatenating turns the column + nested. The controller's per_sample_shapes is the only place a mis-ordered or + overwritten merge shows up; the payload alone would still look correct. + """ + keys = [f"span_{i}" for i in range(8)] + lengths = [i + 1 for i in range(len(keys))] + for key, length in zip(keys, lengths, strict=True): + tq.kv_put(key=key, partition_id=PARTITION, fields={"tokens": torch.ones(length, dtype=torch.int64)}) + + client = tq.get_client() + metadata = client.kv_retrieve_meta(keys=keys, partition_id=PARTITION, create=False) + # Same addition for every row, so this holds whichever order the metadata came back in. + values = TensorDict({"tokens": torch.full((len(keys), 1), 9, dtype=torch.int64)}, batch_size=[len(keys)]) + client.update(metadata, ["tokens"], values=values, parser=concat) + + got = tq.kv_batch_get(keys=keys, partition_id=PARTITION, select_fields="tokens") + for i, (key, length) in enumerate(zip(keys, lengths, strict=True)): + row = got["tokens"][i] + assert row.shape == (length + 1,), f"{key}: expected len {length + 1}, got {tuple(row.shape)}" + assert row[-1].item() == 9, f"{key}: appended value missing" + + partition = partition_snapshot() + tokens = partition.field_metadata["tokens"] + assert tokens.is_nested is True, "mixed lengths must promote the column to nested" + for key, length in zip(keys, lengths, strict=True): + recorded = tuple(tokens.per_sample_shapes[partition.keys_mapping[key]]) + assert recorded == (length + 1,), f"{key}: controller recorded {recorded}, expected {(length + 1,)}" + + +@ray.remote(num_cpus=0) +class Updater: + """A worker that attaches to the running TransferQueue and updates its own key.""" + + def __init__(self): + tq.init() + + def run(self, key: str, addition: int) -> str: + tq.kv_update( + key=key, + partition_id=PARTITION, + fields="tokens", + values=torch.tensor([addition]), + parser=concat, + ) + return key + + +def check_updates_from_remote_workers(): + """Updates issued by other processes on the cluster, not just the driver.""" + keys = [f"worker_{i}" for i in range(4)] + for i, key in enumerate(keys): + tq.kv_put(key=key, partition_id=PARTITION, fields={"tokens": torch.tensor([i])}) + + updaters = [Updater.remote() for _ in keys] + ray.get([u.run.remote(key, 100 + i) for i, (u, key) in enumerate(zip(updaters, keys, strict=True))]) + for u in updaters: + ray.kill(u) + + got = tq.kv_batch_get(keys=keys, partition_id=PARTITION, select_fields="tokens") + for i, key in enumerate(keys): + expected = torch.tensor([i, 100 + i]) + assert torch.equal(got["tokens"][i], expected), f"{key}: {got['tokens'][i]} != {expected}" + + +# ==================== harness ==================== + + +def partition_snapshot(): + controller = ray.get_actor("TransferQueueController", namespace="transfer_queue") + return ray.get(controller.get_partition_snapshot.remote(PARTITION)) + + +CHECKS = [ + ("concat keeps the field name", check_concat_keeps_field_name), + ("parser sees old=None for a new field", check_parser_sees_none_for_a_new_field), + ("several fields in one call", check_multiple_fields_in_one_call), + ("non-tensor payload", check_non_tensor_payload), + ("kv_empty keeps the key, updates the controller", check_empty_keeps_the_key_and_updates_the_controller), + ("a failed parser leaves the row unchanged", check_failed_parser_leaves_the_row_unchanged), + ("rejected arguments", check_rejected_arguments), + ("async_kv_update / async_kv_empty", check_async_variants), + ("multi-sample update across units", check_multi_sample_update_across_units), + ("updates from remote workers", check_updates_from_remote_workers), +] + + +def refuse_if_already_deployed() -> None: + """Never attach to a TransferQueue somebody else is using. + + tq.init() silently joins an existing controller by name, so on a shared cluster + this script would otherwise run its checks against a live deployment and clear + its keys on the way out. + """ + try: + ray.get_actor("TransferQueueController", namespace="transfer_queue") + except ValueError: + return + raise SystemExit( + "A TransferQueueController is already running on this cluster. This script needs a " + "deployment of its own; stop the owning job or run it on an idle cluster." + ) + + +def main() -> int: + if not ray.is_initialized(): + ray.init(namespace="transfer_queue") + + print("=" * 78) + print("kv_update end-to-end checks") + print(f"cluster: {ray.get_runtime_context().gcs_address} nodes: {len(ray.nodes())}") + print("=" * 78) + + refuse_if_already_deployed() + tq.init(OmegaConf.create(CONFIG)) + failures = [] + try: + for name, check in CHECKS: + try: + check() + except Exception as e: + failures.append(name) + print(f" FAIL {name}: {type(e).__name__}: {e}") + traceback.print_exc() + else: + print(f" ok {name}") + finally: + # Each check owns its keys; clear between them so one failure cannot cascade. + leftover = list(tq.kv_list(partition_id=PARTITION).get(PARTITION, {})) + if leftover: + tq.kv_clear(keys=leftover, partition_id=PARTITION) + finally: + tq.close() + + print("=" * 78) + print(f"{len(CHECKS) - len(failures)}/{len(CHECKS)} checks passed") + if failures: + print("failed: " + ", ".join(failures)) + print("=" * 78) + return 1 if failures else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 2077102bdf360a4309a2b0e2b306a383874d41df Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Sat, 3 Oct 2026 13:32:20 +0800 Subject: [PATCH 5/9] [fix] Make SSD-offloaded put all-or-nothing across fields With SSD offload on, put_data wrote, committed and cleaned up one field at a time. If writing a later field's files failed (for example a full disk), the earlier fields were already replaced and their old files deleted, yet the caller got an error. A plain put converges on retry, but kv_update parsers such as concat are not idempotent and the manager does not retry updates, so retrying applied the earlier fields twice. Write every field's files first and only then swap all references, update the SSD accounting and delete superseded files; any failure removes just the files this call wrote. The base put_data now checks every field's length before writing any, so the single commit cannot fail halfway either. apply_update's "a raise leaves storage unchanged" now holds for SSD write errors as well as parser errors. Signed-off-by: OutstanderWang --- tests/test_kv_update.py | 34 ++++++++++++++ tests/test_simple_storage_unit.py | 34 ++++++++++++++ transfer_queue/storage/simple_storage.py | 59 ++++++++++++++---------- 3 files changed, 103 insertions(+), 24 deletions(-) diff --git a/tests/test_kv_update.py b/tests/test_kv_update.py index 6fd4e741..0ab2f83f 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 errno from unittest.mock import MagicMock, patch import numpy as np @@ -143,6 +144,39 @@ def test_apply_update_decodes_ssd_offloaded_old_value(tmp_path): assert data.ssd_active_values == 1, "the replaced SSD file must be released, not leaked" +def test_apply_update_is_atomic_when_a_later_field_fails_to_reach_ssd(tmp_path, monkeypatch): + """Concat is not idempotent, so a partly applied update would double-apply on retry.""" + data = HybridStorageUnitData( + storage_size=4, threshold_bytes=64, ssd_path=str(tmp_path), run_id="run", unit_id="unit" + ) + try: + prompt, mask = torch.arange(32), torch.ones(32, dtype=torch.int64) + data.put_data({"tokens": [prompt], "mask": [mask]}, [0]) + old_files = set(tmp_path.rglob("*.bin")) + + write_values = data._ssd_store.write_values + calls = [] + + def fail_second_field(encoded_values): + calls.append(encoded_values) + if len(calls) == 2: + raise OSError(errno.ENOSPC, "No space left on device") + return write_values(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], ["tokens", "mask"], new_data, _concat, False) + + stored = data.get_data(["tokens", "mask"], [0]) + assert torch.equal(stored["tokens"][0], prompt) + assert torch.equal(stored["mask"][0], mask) + assert set(tmp_path.rglob("*.bin")) == old_files + assert (data.ssd_active_values, data.ssd_active_bytes) == (2, prompt.nbytes + mask.nbytes) + finally: + data.close() + + def test_build_update_field_schema_orders_shapes_across_units(): """Units describe only their own rows; the batch schema must follow metadata order.""" described = _build_update_field_schema( diff --git a/tests/test_simple_storage_unit.py b/tests/test_simple_storage_unit.py index 99ceb4ec..fe52f175 100644 --- a/tests/test_simple_storage_unit.py +++ b/tests/test_simple_storage_unit.py @@ -28,6 +28,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import errno import pickle import time from pathlib import Path @@ -603,6 +604,39 @@ def fail_write(_fd, _data): storage.close() +def test_failed_ssd_write_on_second_field_leaves_every_field_unchanged(tmp_path, monkeypatch): + storage = HybridStorageUnitData( + storage_size=10, + threshold_bytes=64, + ssd_path=str(tmp_path), + run_id="test-run", + unit_id="test-unit", + ) + try: + storage.put_data({"a": [b"a" * 100], "b": [b"b" * 100]}, [1]) + old_files = set(tmp_path.rglob("*.bin")) + + write_values = storage._ssd_store.write_values + calls = [] + + def fail_second_field(encoded_values): + calls.append(encoded_values) + if len(calls) == 2: + raise OSError(errno.ENOSPC, "No space left on device") + return write_values(encoded_values) + + monkeypatch.setattr(storage._ssd_store, "write_values", fail_second_field) + with pytest.raises(OSError, match="No space left"): + storage.put_data({"a": [b"A" * 100], "b": [b"B" * 100]}, [1]) + + assert storage.get_data(["a", "b"], [1]) == {"a": [b"a" * 100], "b": [b"b" * 100]} + assert set(tmp_path.rglob("*.bin")) == old_files + assert storage.ssd_active_values == 2 + assert storage.ssd_active_bytes == 200 + finally: + storage.close() + + @pytest.mark.parametrize("threshold_bytes", [0, -1]) def test_hybrid_storage_requires_positive_threshold(tmp_path, threshold_bytes): with pytest.raises(ValueError, match="threshold must be greater than zero"): diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py index d7bc184d..63002519 100644 --- a/transfer_queue/storage/simple_storage.py +++ b/transfer_queue/storage/simple_storage.py @@ -222,6 +222,7 @@ def put_data(self, field_data: dict[str, Any], global_indexes: list) -> None: f"StorageUnitData put_data: field '{f}' values length {len(values)} " f"!= global_indexes length {len(global_indexes)}, length mismatch" ) + for f, values in field_data.items(): if f not in self.field_data: self.field_data[f] = {} field_dict = self.field_data[f] @@ -239,8 +240,9 @@ def apply_update( ) -> dict[str, dict[str, Any]]: """Read each (field, index), compute the stored value, then write once. - Missing fields yield ``old=None``. All parser calls finish before any write, - so a raise leaves storage unchanged. + Missing fields yield ``old=None``. All parser calls finish before one + ``put_data`` that writes every field or none, so a raise from a parser or + an SSD write leaves storage unchanged. Returns: Per-sample description of the stored values, keyed by field name. @@ -597,38 +599,47 @@ def put_data( unique_global_indexes = set(global_indexes) has_duplicate_indexes = len(unique_global_indexes) != len(global_indexes) - for field, values in field_data.items(): - prepared_values, entries, fallback_values = self._prepare_field_values(values) - + old_ssd_values = [] + for field in field_data: stored_field = self.field_data.get(field, {}) - old_ssd_values = [] for global_index in unique_global_indexes: old_value = stored_field.get(global_index) if isinstance(old_value, _SSDValueRef): old_ssd_values.append(old_value) - try: - super().put_data({field: prepared_values}, global_indexes) - except Exception: - for entry in entries: - self._ssd_store.unlink(entry) - raise - - obsolete_ssd_values = old_ssd_values - if has_duplicate_indexes: - retained_ssd_paths = set() + + # Write every field's files before replacing any value: kv_update parsers are not + # idempotent, so a failure on a later field must not leave earlier fields committed. + prepared_data = {} + entries: list[_SSDValueRef] = [] + fallback_values = 0 + try: + for field, values in field_data.items(): + prepared_data[field], field_entries, field_fallback_values = self._prepare_field_values(values) + entries.extend(field_entries) + fallback_values += field_fallback_values + super().put_data(prepared_data, global_indexes) + except Exception: + for entry in entries: + self._ssd_store.unlink(entry) + raise + + obsolete_ssd_values = old_ssd_values + if has_duplicate_indexes: + retained_ssd_paths = set() + for field in field_data: for global_index in unique_global_indexes: retained_value = self.field_data[field][global_index] if isinstance(retained_value, _SSDValueRef): retained_ssd_paths.add(retained_value.path) - obsolete_ssd_values.extend(entry for entry in entries if entry.path not in retained_ssd_paths) + obsolete_ssd_values.extend(entry for entry in entries if entry.path not in retained_ssd_paths) - self._ssd_active_values += len(entries) - len(obsolete_ssd_values) - self._ssd_active_bytes += sum(entry.size_bytes for entry in entries) - sum( - value.size_bytes for value in obsolete_ssd_values - ) - self._ssd_fallback_values_total += fallback_values - for obsolete_value in obsolete_ssd_values: - self._ssd_store.unlink(obsolete_value) + self._ssd_active_values += len(entries) - len(obsolete_ssd_values) + self._ssd_active_bytes += sum(entry.size_bytes for entry in entries) - sum( + value.size_bytes for value in obsolete_ssd_values + ) + self._ssd_fallback_values_total += fallback_values + for obsolete_value in obsolete_ssd_values: + self._ssd_store.unlink(obsolete_value) def get_data(self, fields: list[str], global_indexes: list) -> dict[str, list]: """Read mixed memory- and SSD-backed samples in request order.""" From 7d925168f41a5771c9eb98fa92f8ffe317ecf7b8 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Thu, 8 Oct 2026 22:01:17 +0800 Subject: [PATCH 6/9] [fix] Address kv_update consistency issues Signed-off-by: OutstanderWang --- transfer_queue/__init__.py | 4 + transfer_queue/client.py | 28 +-- transfer_queue/controller.py | 50 ++--- transfer_queue/interface.py | 204 ++++++++++-------- transfer_queue/storage/managers/base.py | 51 ++--- .../managers/simple_storage_manager.py | 69 +++--- transfer_queue/storage/simple_storage.py | 69 +++--- 7 files changed, 228 insertions(+), 247 deletions(-) diff --git a/transfer_queue/__init__.py b/transfer_queue/__init__.py index 7f034662..4f3c19ad 100644 --- a/transfer_queue/__init__.py +++ b/transfer_queue/__init__.py @@ -21,6 +21,7 @@ async_kv_batch_get, async_kv_batch_get_by_meta, async_kv_batch_put, + async_kv_batch_update, async_kv_clear, async_kv_empty, async_kv_list, @@ -34,6 +35,7 @@ kv_batch_get, kv_batch_get_by_meta, kv_batch_put, + kv_batch_update, kv_clear, kv_empty, kv_list, @@ -59,6 +61,7 @@ "get_metrics_endpoint", "kv_put", "kv_batch_put", + "kv_batch_update", "kv_batch_get", "kv_batch_get_by_meta", "kv_list", @@ -67,6 +70,7 @@ "kv_empty", "async_kv_put", "async_kv_batch_put", + "async_kv_batch_update", "async_kv_batch_get", "async_kv_batch_get_by_meta", "async_kv_list", diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 17ca3b97..41667747 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -495,19 +495,15 @@ async def async_put( async def async_update( self, metadata: BatchMeta, - field_names: list[str], - values: TensorDict | None = None, - parser: Callable[[Any, Any], Any] | None = None, - empty: bool = False, + values: TensorDict, + merge_fn: Callable[[Any, Any], Any], ) -> BatchMeta: - """Apply empty or parser(old, new) on SimpleStorage units for the named fields. + """Merge new values into produced fields on SimpleStorage units. Args: metadata: Samples to update. The key must already exist. - field_names: Fields to rewrite. - values: New values aligned with ``metadata``, or None when empty. - parser: ``parser(old, new) -> stored``. Ignored when empty. - empty: Store None for each named field. + values: New values aligned with ``metadata``. + merge_fn: Called per sample per field as ``merge_fn(old, new)``. Returns: The same metadata with field_schema and production status updated. @@ -523,9 +519,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, field_names, values=values, parser=parser, empty=empty - ) + field_schema = await self.storage_manager.update_data(metadata, values, merge_fn) return metadata.apply_field_schema(field_schema) async def async_get_data(self, metadata: BatchMeta) -> TensorDict: @@ -1619,16 +1613,14 @@ def put( def update( self, metadata: BatchMeta, - field_names: list[str], - values: TensorDict | None = None, - parser: Callable[[Any, Any], Any] | None = None, - empty: bool = False, + values: TensorDict, + merge_fn: Callable[[Any, Any], Any], ) -> BatchMeta: - """Synchronously apply empty or parser(old, new) on SimpleStorage units. + """Synchronously merge new values into produced fields on SimpleStorage units. See ``async_update``. """ - return self._update(metadata, field_names, values=values, parser=parser, empty=empty) + return self._update(metadata, values, merge_fn) 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 ae8691a2..c86c0824 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -253,36 +253,28 @@ def update(self, incoming: dict[str, Any], incoming_global_indexes: list[int]) - # Update newly provided per_sample_shapes self.per_sample_shapes.update(new_per_sample_shapes) - elif new_is_non_tensor: - # Some samples now hold plain objects (e.g. kv_update empty stores None), so no - # single dtype/shape describes the column any more. dtype is kept so a later - # tensor write still has to agree with the one this field was created with. - self.is_non_tensor = True - self.is_nested = False - self.shape = None - self.per_sample_shapes.clear() - else: - # newly input is regular tensor - new_shape = incoming.get("shape", None) - if new_shape is None: - raise ValueError("Receiving a regular tensor without 'shape'!") - if self.is_nested: - # we need to update incoming shape into per_sample_shapes - for gi in incoming_global_indexes: - self.per_sample_shapes[gi] = new_shape - else: - if self.is_non_tensor is not None and not self.is_non_tensor: - # original data is also regular tensor - assert self.shape is not None - if self.shape != new_shape: - for gi in self.global_indexes: - self.per_sample_shapes[gi] = self.shape - for gi in incoming_global_indexes: - self.per_sample_shapes[gi] = new_shape - - self.shape = None - self.is_nested = True + if not new_is_non_tensor: + # newly input is regular tensor + new_shape = incoming.get("shape", None) + if new_shape is None: + raise ValueError("Receiving a regular tensor without 'shape'!") + if self.is_nested: + # we need to update incoming shape into per_sample_shapes + for gi in incoming_global_indexes: + self.per_sample_shapes[gi] = new_shape + else: + if self.is_non_tensor is not None and not self.is_non_tensor: + # original data is also regular tensor + assert self.shape is not None + if self.shape != new_shape: + for gi in self.global_indexes: + self.per_sample_shapes[gi] = self.shape + for gi in incoming_global_indexes: + self.per_sample_shapes[gi] = new_shape + + self.shape = None + self.is_nested = True self.global_indexes.update(incoming_global_indexes) diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 64f1d9a8..e496f0dc 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -611,39 +611,29 @@ def kv_clear(keys: list[str] | str, partition_id: str) -> None: def kv_update( key: str, partition_id: str, - fields: str | list[str], - *, - values: Any | dict[str, Any] | None = None, - parser: Callable[[Any, Any], Any] | None = None, - empty: bool = False, + fields: dict[str, Any], + merge_fn: Callable[[Any, Any], Any], ) -> KVBatchMeta: - """Update fields of an existing key by combining the stored value with a new one. - - A custom ``parser(old, new)`` runs once per sample per field on the SimpleStorage - unit that holds the row, then writes the return value back. ``empty=True`` - stores ``None`` instead and must not be given ``values`` or ``parser``. - ``kv_empty()`` is the same as ``kv_update(..., empty=True)``. + """Merge new values into produced fields of one existing key. Args: key: Existing user-specified key. The key is not created if it is missing. partition_id: Partition that holds the key. - fields: One field name, or a list of field names, to update. - values: New value for a single field, or ``{field: new}`` for several. - Must be omitted when ``empty=True``. - parser: ``parser(old, new) -> stored``. Required unless ``empty=True``. - Only SimpleStorage executes parsers; KV backends reject the call. - ``old`` is ``None`` when the field has never been written. - empty: If True, store None for each named field. + 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. Returns: - KVBatchMeta for the updated sample, including every field stored for it. + Metadata for the updated sample. Raises: - ValueError: If the key is missing, ``empty=True`` is given values or a - parser, a custom parser is given no values, or ``fields`` / ``values`` - do not line up. - TypeError: If ``parser`` is missing or not callable when ``empty`` is False. + ValueError: If the key or a field is missing. + TypeError: If ``fields`` is not a non-empty dict or ``merge_fn`` is not callable. NotImplementedError: If the storage backend is not SimpleStorage. + + Note: + A timeout has an unknown outcome because the storage unit may finish the + merge after the caller stops waiting. Do not blindly retry a merge. """ tq_client = _maybe_create_tq_client() return tq_client._run_coroutine( @@ -651,16 +641,32 @@ def kv_update( key=key, partition_id=partition_id, fields=fields, - values=values, - parser=parser, - empty=empty, + merge_fn=merge_fn, ) ) -def kv_empty(key: str, partition_id: str, fields: str | list[str]) -> KVBatchMeta: - """Store None for the named fields of an existing key. Same as kv_update(..., empty=True).""" - return kv_update(key=key, partition_id=partition_id, fields=fields, empty=True) +def kv_batch_update( + keys: list[str], + partition_id: str, + fields: TensorDict, + merge_fn: Callable[[Any, Any], Any], +) -> KVBatchMeta: + """Merge batched values into produced fields of existing keys. + + 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) + ) + + +def kv_empty(keys: str | list[str], partition_id: str, fields: str | list[str]) -> KVBatchMeta: + """Release produced fields by storing ``None`` while keeping the keys ready.""" + tq_client = _maybe_create_tq_client() + return tq_client._run_coroutine(async_kv_empty(keys=keys, partition_id=partition_id, fields=fields)) # ==================== KV Interface API ==================== @@ -1044,86 +1050,66 @@ async def async_kv_clear(keys: list[str] | str, partition_id: str) -> None: await tq_client.async_clear_samples(batch_meta) -def _normalize_kv_update_args( - fields: str | list[str], - values: Any | dict[str, Any] | None, - parser: Callable[[Any, Any], Any] | None, - empty: bool = False, -) -> tuple[list[str], TensorDict | None, Callable[[Any, Any], Any] | None, bool]: - """Validate kv_update arguments and wrap new values as a one-row TensorDict.""" - if isinstance(fields, str): - field_names = [fields] - elif isinstance(fields, list) and fields and all(isinstance(f, str) for f in fields): - field_names = list(fields) - else: - raise TypeError("fields must be a field name or a non-empty list of field names") - - if empty: - if values is not None: - raise ValueError("kv_update with empty=True must not specify values") - if parser is not None: - raise ValueError("kv_update with empty=True must not specify parser") - return field_names, None, None, True - - if not callable(parser): - raise TypeError("parser must be callable unless empty=True") - if values is None: - raise ValueError("kv_update with a custom parser requires values") - - if isinstance(fields, str): - value_map = {fields: values} - else: - if not isinstance(values, dict): - raise TypeError("values must be a dict mapping field name to new value when updating multiple fields") - missing = [name for name in field_names if name not in values] - extra = [name for name in values if name not in field_names] - if missing or extra: - raise ValueError( - f"fields and values must name the same columns; fields={field_names}, extra={extra}, missing={missing}" - ) - value_map = values - +def _single_update_batch(fields: dict[str, Any]) -> TensorDict: + """Wrap one sample's field mapping in a one-row TensorDict.""" + if not isinstance(fields, dict) or not fields: + raise TypeError("fields must be a non-empty dict") batch: dict[str, Any] = {} - for field_name, value in value_map.items(): + for field_name, value in fields.items(): + if not isinstance(field_name, str): + raise TypeError("field names must be strings") if isinstance(value, torch.Tensor): if value.is_nested: - raise ValueError("nested tensors are not supported for single-key kv_update") + raise ValueError("Use async_kv_batch_update for nested tensors") batch[field_name] = value.unsqueeze(0) else: batch[field_name] = NonTensorStack(value) - return field_names, TensorDict(batch, batch_size=[1]), parser, False + return TensorDict(batch, batch_size=[1]) async def async_kv_update( key: str, partition_id: str, - fields: str | list[str], - *, - values: Any | dict[str, Any] | None = None, - parser: Callable[[Any, Any], Any] | None = None, - empty: bool = False, + fields: dict[str, Any], + merge_fn: Callable[[Any, Any], Any], ) -> KVBatchMeta: """Asynchronously update fields of an existing key. See ``kv_update``.""" - field_names, value_batch, bound_parser, use_empty = _normalize_kv_update_args(fields, values, parser, empty) + return await async_kv_batch_update( + keys=[key], + partition_id=partition_id, + fields=_single_update_batch(fields), + merge_fn=merge_fn, + ) + +async def async_kv_batch_update( + keys: list[str], + partition_id: str, + fields: TensorDict, + merge_fn: Callable[[Any, Any], Any], +) -> KVBatchMeta: + """Asynchronously merge batched values into produced fields of existing keys.""" + 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(): + raise TypeError("fields must be a non-empty TensorDict") + if fields.batch_size != torch.Size([len(keys)]): + raise ValueError(f"fields batch size {fields.batch_size} does not match {len(keys)} keys") + 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=[key], partition_id=partition_id, create=False) + batch_meta = await tq_client.async_kv_retrieve_meta(keys=keys, partition_id=partition_id, create=False) - if batch_meta.size == 0: - raise ValueError("keys or partition were not found!") - if batch_meta.size != 1: - raise RuntimeError(f"Retrieved BatchMeta size {batch_meta.size} does not match with input `key` size of 1!") + 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}") - batch_meta = await tq_client.async_update( - metadata=batch_meta, - field_names=field_names, - values=value_batch, - parser=bound_parser, - empty=use_empty, - ) + batch_meta = await tq_client.async_update(metadata=batch_meta, values=fields, merge_fn=merge_fn) return KVBatchMeta( - keys=[key], + keys=keys, tags=batch_meta.custom_meta, partition_id=partition_id, fields=batch_meta.field_names, @@ -1131,9 +1117,41 @@ async def async_kv_update( ) -async def async_kv_empty(key: str, partition_id: str, fields: str | list[str]) -> KVBatchMeta: - """Store None for the named fields of an existing key. Same as async_kv_update(..., empty=True).""" - return await async_kv_update(key=key, partition_id=partition_id, fields=fields, empty=True) +async def async_kv_empty(keys: str | list[str], partition_id: str, fields: str | list[str]) -> KVBatchMeta: + """Asynchronously release fields by storing ``None`` through the ordinary put path.""" + if isinstance(keys, str): + keys = [keys] + if isinstance(fields, str): + field_names = [fields] + elif isinstance(fields, list) and fields and all(isinstance(name, str) for name in fields): + field_names = list(dict.fromkeys(fields)) + else: + raise TypeError("fields must be a field name or a non-empty list of field names") + if not keys: + raise ValueError("keys must not be empty") + + tq_client = _maybe_create_tq_client() + if not isinstance(tq_client.storage_manager, AsyncSimpleStorageManager): + raise NotImplementedError("kv_empty requires the SimpleStorage backend") + batch_meta = await tq_client.async_kv_retrieve_meta(keys=keys, partition_id=partition_id, create=False) + if batch_meta.size != len(keys): + raise ValueError("Some keys or the partition were not found") + missing_fields = sorted(set(field_names) - set(batch_meta.field_names)) + if missing_fields: + raise ValueError(f"Fields are not produced for every key: {missing_fields}") + + empty_values = TensorDict( + {field_name: NonTensorStack(*([None] * len(keys))) for field_name in field_names}, + batch_size=[len(keys)], + ) + batch_meta = await tq_client.async_put(empty_values, batch_meta) + return KVBatchMeta( + keys=keys, + tags=batch_meta.custom_meta, + partition_id=partition_id, + fields=batch_meta.field_names, + extra_info=batch_meta.extra_info, + ) # ==================== Low-Level Native API ==================== diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index 72585f0f..9ea1fc30 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -226,7 +226,7 @@ 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, - ) -> None: + ) -> bool: """ Notify controller that new data is ready. @@ -239,7 +239,7 @@ async def notify_data_update( if not self.controller_info: logger.warning(f"No controller connected for storage manager {self.storage_manager_id}") - return + return False normalized_field_schema = {} for field_name, field in field_schema.items(): @@ -271,9 +271,9 @@ async def notify_data_update( self._notify_and_wait(request_msg), self._notify_loop, ) - await asyncio.wrap_future(thread_future) + return await asyncio.wrap_future(thread_future) - async def _notify_and_wait(self, request_msg: list) -> None: + async def _notify_and_wait(self, request_msg: list) -> bool: """Send a data status notification to the controller and block until ACK is received.""" # Acquiring the lease sits outside the handler below: a missing socket name or a dead # context is a configuration/lifecycle fault the caller must see, not a slow ACK. @@ -298,16 +298,18 @@ async def _notify_and_wait(self, request_msg: list) -> None: response_msg = ZMQMessage.deserialize(messages) if response_msg.request_type == ZMQRequestType.NOTIFY_DATA_UPDATE_ACK: # type: ignore[arg-type] + success = bool(response_msg.body.get("success")) logger.debug( - f"[{self.storage_manager_id}]: Get data status update ACK response " - f"from controller id #{response_msg.sender_id} successfully." + f"[{self.storage_manager_id}]: Received data status update ACK " + f"from controller id #{response_msg.sender_id}: success={success}." ) - return + return success except Exception as e: # Logged rather than raised, so a slow controller does not fail the put. Close # the socket: a late ACK would otherwise be read as the next lessee's reply. logger.error(f"[{self.storage_manager_id}]: Data status update failed: {type(e).__name__}: {e}") sock.close(linger=0) + return False @abstractmethod async def put_data( @@ -357,24 +359,22 @@ async def clear_data(self, metadata: BatchMeta) -> None: async def update_data( self, metadata: BatchMeta, - field_names: list[str], - values: TensorDict | None = None, - parser: Callable[[Any, Any], Any] | None = None, - empty: bool = False, + values: TensorDict, + merge_fn: Callable[[Any, Any], Any], ) -> dict[str, dict[str, Any]]: - """Apply empty or parser(old, new) on the unit that holds each sample. + """Merge new values on the unit that holds each sample. Args: metadata: Samples to update. - field_names: Fields to rewrite. - values: New values as a TensorDict aligned with ``metadata``, or None for empty. - parser: Called per sample per field as ``parser(old, new)``. Ignored when empty. - empty: If True, store None for each named field and ignore values/parser. + values: New values as a TensorDict aligned with ``metadata``. + merge_fn: Called per sample per field as ``merge_fn(old, new)``. Returns: field_schema of the stored values, keyed by field name. """ - raise NotImplementedError(f"{self.__class__.__name__} does not support update_data") + raise NotImplementedError( + f"kv_update is not supported by {self.__class__.__name__}; it requires the SimpleStorage backend" + ) async def save_checkpoint(self, checkpoint_dir: str) -> None: """Save storage state into checkpoint_dir. @@ -822,23 +822,6 @@ async def put_data( per_field_custom_backend_meta, ) - async def update_data( - self, - metadata: BatchMeta, - field_names: list[str], - values: TensorDict | None = None, - parser: Callable[[Any, Any], Any] | None = None, - empty: bool = False, - ) -> dict[str, dict[str, Any]]: - """kv_update is only implemented for SimpleStorage. - - A KV backend stores each sample-field under its own key and offers no hook to run - parser(old, new) where the value lives, so read-modify-write cannot be made atomic. - """ - raise NotImplementedError( - f"kv_update is not supported by {type(self).__name__}; it requires the SimpleStorage backend." - ) - async def get_data(self, metadata: BatchMeta) -> TensorDict: """ Retrieve tensor data from the backend storage. diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py index f9c52844..dd482c4e 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -542,42 +542,28 @@ async def put_data( async def update_data( self, metadata: BatchMeta, - field_names: list[str], - values: TensorDict | None = None, - parser: Callable[[Any, Any], Any] | None = None, - empty: bool = False, + values: TensorDict, + merge_fn: Callable[[Any, Any], Any], ) -> dict[str, dict[str, Any]]: - """Send an update to each hashed storage unit and notify the controller. + """Merge values on each owner unit and notify the controller. - The unit reads the old value, applies empty or parser(old, new), and reports - the shape it stored per sample; this assembles one batch-level field_schema - so production status stays ready. + Each unit reports the values it stored; this assembles one batch-level + field_schema so production status stays ready. """ logger.debug(f"[{self.storage_manager_id}]: receive update_data request, updating {metadata.size} samples.") if metadata.size == 0: return {} - if empty: - if values is not None or parser is not None: - raise ValueError("empty update must not include values or a parser") - else: - if values is None or parser is None: - raise ValueError("update requires values and a parser unless empty is set") - if values.batch_size[0] != metadata.size: - raise ValueError( - f"Batch size of values ({values.batch_size[0]}) does not match metadata size ({metadata.size})" - ) + if values.batch_size[0] != metadata.size: + raise ValueError( + f"Batch size of values ({values.batch_size[0]}) does not match metadata size ({metadata.size})" + ) + field_names = list(values.keys()) routing = self._group_by_hash(metadata.global_indexes) - # Parser-backed updates are not replayed: parser(old, new) reads the stored value, so a - # second attempt after a lost answer would fold the new value in twice. Empty is a - # plain overwrite and stays retryable. - max_attempts = 1 if parser is not None else None tasks = [] for su_id, group in routing.items(): - storage_data = None - if values is not None: - storage_data = {f: self._select_by_positions(values[f], group.batch_positions) for f in field_names} + storage_data = {f: self._select_by_positions(values[f], group.batch_positions) for f in field_names} tasks.append( self._request_with_retry( "update", @@ -586,13 +572,14 @@ async def update_data( partial( self._update_to_single_storage_unit, group.global_indexes, - field_names, storage_data, - parser, - empty, + merge_fn, + {name: metadata.field_schema[name] for name in field_names}, target_storage_unit=su_id, ), - max_attempts=max_attempts, + # A merge reads stored state, so a replay after a lost answer + # could fold the same new value in twice. + max_attempts=1, ) ) @@ -612,11 +599,16 @@ async def update_data( metadata.global_indexes, list(zip([g.global_indexes for g in routing.values()], described, strict=True)) ) - await self.notify_data_update( + published = await self.notify_data_update( metadata.partition_ids[0], metadata.global_indexes, field_schema, ) + if not published: + raise RuntimeError( + "Storage update committed, but the controller did not accept its metadata; " + "do not retry this merge blindly" + ) return field_schema @with_storage_unit_socket @@ -680,22 +672,19 @@ async def _put_to_single_storage_unit( async def _update_to_single_storage_unit( self, global_indexes: list[int], - field_names: list[str], - storage_data: dict[str, Any] | None, - parser: Callable[[Any, Any], Any] | None, - empty: bool, + 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]]: """Send an update to one storage unit and return what it stored, per field.""" body: dict[str, Any] = { "global_indexes": global_indexes, - "fields": field_names, - "empty": empty, + "data": storage_data, + "merge_fn": merge_fn, + "field_schema": field_schema, } - if not empty: - body["data"] = storage_data - body["parser"] = parser request_msg = ZMQMessage.create( request_type=ZMQRequestType.UPDATE_DATA, # type: ignore[arg-type] @@ -724,7 +713,7 @@ async def _update_to_single_storage_unit( time.perf_counter() - started, serialized_bytes, f"to {target_storage_unit} at {self._describe_storage_unit(target_storage_unit)} " - f"samples={len(global_indexes)} fields={field_names}", + f"samples={len(global_indexes)} fields={list(storage_data)}", ) return response_msg.body.get("stored_shapes", {}) except zmq.error.Again as e: diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py index 63002519..2610792e 100644 --- a/transfer_queue/storage/simple_storage.py +++ b/transfer_queue/storage/simple_storage.py @@ -233,35 +233,38 @@ def put_data(self, field_data: dict[str, Any], global_indexes: list) -> None: def apply_update( self, global_indexes: list[int], - fields: list[str], - new_data: dict[str, Any] | None, - parser: Callable[[Any, Any], Any] | None, - use_empty: bool, + new_data: dict[str, Any], + merge_fn: Callable[[Any, Any], Any], + field_schema: dict[str, dict[str, Any]], ) -> dict[str, dict[str, Any]]: - """Read each (field, index), compute the stored value, then write once. + """Merge produced values, validate their types, then write once. - Missing fields yield ``old=None``. All parser calls finish before one - ``put_data`` that writes every field or none, so a raise from a parser or - an SSD write leaves storage unchanged. + All merge calls and type checks finish before one ``put_data`` call, so + a merge error or incompatible result leaves storage unchanged. Returns: Per-sample description of the stored values, keyed by field name. """ computed: dict[str, list] = {} - if use_empty: - computed = {field: [None] * len(global_indexes) for field in fields} - else: - if parser is None or new_data is None: - raise TypeError("apply_update requires a parser and new values unless use_empty is set") - for field in fields: - # Read through get_data so SSD-offloaded values reach the parser decoded, - # not as file references; indexes never written keep old=None. - stored = self.field_data.get(field, {}) - present = [idx for idx in global_indexes if idx in stored] - old_values = dict(zip(present, self.get_data([field], present)[field], strict=True)) if present else {} - computed[field] = [ - parser(old_values.get(idx), new_data[field][i]) for i, idx in enumerate(global_indexes) - ] + 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__ + raise TypeError( + f"Merge result for tensor field {field!r} must keep dtype {expected['dtype']}, got {actual}" + ) self.put_data(computed, global_indexes) return {field: _describe_stored_values(values) for field, values in computed.items()} @@ -607,7 +610,7 @@ def put_data( if isinstance(old_value, _SSDValueRef): old_ssd_values.append(old_value) - # Write every field's files before replacing any value: kv_update parsers are not + # Write every field's files before replacing any value: kv_update merge functions are not # idempotent, so a failure on a later field must not leave earlier fields committed. prepared_data = {} entries: list[_SSDValueRef] = [] @@ -1236,23 +1239,23 @@ def _handle_put(self, data_parts: ZMQMessage) -> ZMQMessage: ) def _handle_update(self, data_parts: ZMQMessage) -> ZMQMessage: - """Read each field, apply empty or parser(old, new), write once, describe what was stored.""" + """Merge existing fields, write once, and describe what was stored.""" try: global_indexes = data_parts.body["global_indexes"] - fields = data_parts.body["fields"] - use_empty = bool(data_parts.body.get("empty", False)) - parser = data_parts.body.get("parser") + 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" ): - if not use_empty: - if not callable(parser): - raise TypeError(f"update parser must be callable, got {type(parser).__name__}") - if not isinstance(new_data, dict): - raise TypeError("update data must be a dict of new field values") - stored_shapes = self.storage_data.apply_update(global_indexes, fields, new_data, parser, use_empty) + if not callable(merge_fn): + 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) return ZMQMessage.create( request_type=ZMQRequestType.UPDATE_DATA_RESPONSE, # type: ignore[arg-type] From 12c60302c516881aabc4b90496c0781d1a301a72 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Thu, 8 Oct 2026 22:01:23 +0800 Subject: [PATCH 7/9] [test] Cover batch kv_update behavior Signed-off-by: OutstanderWang --- recipe/kv_update_e2e.py | 389 ---------------------------- tests/e2e/test_kv_interface_e2e.py | 171 ++++++------ tests/test_kv_update.py | 173 +++++-------- tests/test_storage_request_retry.py | 63 ++++- 4 files changed, 223 insertions(+), 573 deletions(-) delete mode 100644 recipe/kv_update_e2e.py diff --git a/recipe/kv_update_e2e.py b/recipe/kv_update_e2e.py deleted file mode 100644 index c25c6bc9..00000000 --- a/recipe/kv_update_e2e.py +++ /dev/null @@ -1,389 +0,0 @@ -# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved. -# Copyright 2025 The TransferQueue Team -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""End-to-end exercise of kv_update / kv_empty against a real Ray cluster. - -kv_update rewrites fields of an existing key by running ``parser(old, new)`` on the -SimpleStorage unit that holds the row, so nothing here works unless the parser -actually travels to the unit and the resulting schema travels back to the -controller. Every check below reads its result back through the public API, and -the multi-unit ones also inspect controller metadata, which is the part a -data-only assertion cannot see. - -Run it as a Ray job so the driver lives on the cluster: - - ray job submit --address http://: --working-dir . \ - -- python recipe/kv_update_e2e.py - -Locally, ``python recipe/kv_update_e2e.py`` starts its own cluster instead. -""" - -import asyncio -import os -import traceback - -# Disable Ray's cross-worker log deduplication before importing Ray itself, -# otherwise worker-side prints get folded into "[repeated Nx across cluster]". -os.environ.setdefault("RAY_DEDUP_LOGS", "0") - -import ray -import torch -from omegaconf import OmegaConf -from tensordict import TensorDict - -import transfer_queue as tq - -PARTITION = "kv_update_e2e" - -# Several units so the hash routing actually splits a batch; the cross-unit checks -# below are meaningless with a single unit. -CONFIG = { - "controller": {"polling_mode": True}, - "backend": { - "storage_backend": "SimpleStorage", - "SimpleStorage": {"total_storage_size": 512, "num_data_storage_units": 4}, - }, -} - - -def concat(old, new): - """Append new tokens to what is already stored, keeping the field name.""" - return new if old is None else torch.cat([old, new]) - - -# ==================== checks ==================== - - -def check_concat_keeps_field_name(): - """The rollout case: prompt_ids in place, response_ids appended, one field.""" - key = "concat" - tq.kv_put(key=key, partition_id=PARTITION, fields={"sequence_ids": torch.tensor([10, 11, 12])}) - - tq.kv_update( - key=key, - partition_id=PARTITION, - fields="sequence_ids", - values=torch.tensor([20, 21]), - parser=concat, - ) - - got = tq.kv_batch_get(keys=key, partition_id=PARTITION) - assert torch.equal(got["sequence_ids"][0], torch.tensor([10, 11, 12, 20, 21])), got["sequence_ids"][0] - assert list(got.keys()) == ["sequence_ids"], f"update must not add columns: {list(got.keys())}" - - -def check_parser_sees_none_for_a_new_field(): - """A field the key never held reaches the parser as old=None. - - The parser runs in the storage unit's process, so what it observed can only - come back inside the value it returns; a closure variable would stay untouched - here in the driver. - """ - key = "fresh_field" - tq.kv_put(key=key, partition_id=PARTITION, fields={"anchor": torch.tensor([1])}) - - def report_old(old, new): - return {"old_was_none": old is None, "new": new.tolist()} - - tq.kv_update(key=key, partition_id=PARTITION, fields="logprobs", values=torch.tensor([7, 8]), parser=report_old) - - got = tq.kv_batch_get(keys=key, partition_id=PARTITION, select_fields="logprobs") - assert got["logprobs"][0] == {"old_was_none": True, "new": [7, 8]}, got["logprobs"][0] - - -def check_multiple_fields_in_one_call(): - key = "multi_field" - tq.kv_put( - key=key, - partition_id=PARTITION, - fields={"a": torch.tensor([1, 2]), "b": torch.tensor([10])}, - ) - - tq.kv_update( - key=key, - partition_id=PARTITION, - fields=["a", "b"], - values={"a": torch.tensor([3]), "b": torch.tensor([20, 30])}, - parser=concat, - ) - - got = tq.kv_batch_get(keys=key, partition_id=PARTITION) - assert torch.equal(got["a"][0], torch.tensor([1, 2, 3])), got["a"][0] - assert torch.equal(got["b"][0], torch.tensor([10, 20, 30])), got["b"][0] - - -def check_non_tensor_payload(): - """Parsers are not tensor-only: a plain object round-trips the same way.""" - key = "non_tensor" - tq.kv_put(key=key, partition_id=PARTITION, fields={"stats": {"step": 1}}) - - tq.kv_update( - key=key, - partition_id=PARTITION, - fields="stats", - values={"reward": 0.5}, - parser=lambda old, new: {**old, **new}, - ) - - got = tq.kv_batch_get(keys=key, partition_id=PARTITION, select_fields="stats") - assert got["stats"][0] == {"step": 1, "reward": 0.5}, got["stats"][0] - - -def check_empty_keeps_the_key_and_updates_the_controller(): - """kv_empty stores None; the key stays produced and stops being a tensor column.""" - key = "to_empty" - tq.kv_put(key=key, partition_id=PARTITION, fields={"tokens": torch.tensor([1, 2, 3])}) - - tq.kv_empty(key=key, partition_id=PARTITION, fields="tokens") - - got = tq.kv_batch_get(keys=key, partition_id=PARTITION, select_fields="tokens") - assert got["tokens"][0] is None, got["tokens"][0] - assert key in tq.kv_list(partition_id=PARTITION)[PARTITION], "kv_empty must not delete the key" - - partition = partition_snapshot() - col = partition.field_name_mapping["tokens"] - row = partition.keys_mapping[key] - assert partition.production_status[row, col] == 1, "an emptied field stays produced" - assert partition.field_metadata["tokens"].is_non_tensor is True, "controller still calls the column a tensor" - - -def check_failed_parser_leaves_the_row_unchanged(): - """The unit computes every value before writing, so a raise is a no-op.""" - key = "atomic" - tq.kv_put(key=key, partition_id=PARTITION, fields={"tokens": torch.tensor([1, 2, 3])}) - - def boom(old, new): - raise RuntimeError("parser failed on purpose") - - try: - tq.kv_update(key=key, partition_id=PARTITION, fields="tokens", values=torch.tensor([4]), parser=boom) - except Exception as e: - assert "parser failed on purpose" in str(e), f"unexpected error: {e}" - else: - raise AssertionError("a raising parser must not succeed") - - got = tq.kv_batch_get(keys=key, partition_id=PARTITION, select_fields="tokens") - assert torch.equal(got["tokens"][0], torch.tensor([1, 2, 3])), got["tokens"][0] - - -def check_rejected_arguments(): - key = "rejects" - tq.kv_put(key=key, partition_id=PARTITION, fields={"tokens": torch.tensor([1])}) - - def expect(exc_type, match, fn): - try: - fn() - except exc_type as e: - assert match in str(e), f"expected {match!r} in {e!r}" - else: - raise AssertionError(f"expected {exc_type.__name__} containing {match!r}") - - expect( - ValueError, - "must not specify values", - lambda: tq.kv_update(key=key, partition_id=PARTITION, fields="tokens", values=torch.tensor([1]), empty=True), - ) - expect( - ValueError, - "must not specify parser", - lambda: tq.kv_update(key=key, partition_id=PARTITION, fields="tokens", parser=concat, empty=True), - ) - expect( - TypeError, - "parser must be callable", - lambda: tq.kv_update(key=key, partition_id=PARTITION, fields="tokens", values=torch.tensor([1])), - ) - expect( - ValueError, - "requires values", - lambda: tq.kv_update(key=key, partition_id=PARTITION, fields="tokens", parser=concat), - ) - expect( - ValueError, - "same columns", - lambda: tq.kv_update( - key=key, partition_id=PARTITION, fields=["tokens", "other"], values={"tokens": 1}, parser=concat - ), - ) - expect( - ValueError, - "not found", - lambda: tq.kv_update( - key="never_written", partition_id=PARTITION, fields="tokens", values=torch.tensor([1]), parser=concat - ), - ) - - -def check_async_variants(): - key = "async_key" - tq.kv_put(key=key, partition_id=PARTITION, fields={"tokens": torch.tensor([1, 2])}) - - asyncio.run( - tq.async_kv_update(key=key, partition_id=PARTITION, fields="tokens", values=torch.tensor([3]), parser=concat) - ) - got = tq.kv_batch_get(keys=key, partition_id=PARTITION, select_fields="tokens") - assert torch.equal(got["tokens"][0], torch.tensor([1, 2, 3])), got["tokens"][0] - - asyncio.run(tq.async_kv_empty(key=key, partition_id=PARTITION, fields="tokens")) - got = tq.kv_batch_get(keys=key, partition_id=PARTITION, select_fields="tokens") - assert got["tokens"][0] is None, got["tokens"][0] - - -def check_multi_sample_update_across_units(): - """One update spanning several units must report per-sample shapes in batch order. - - Each key starts with a different length, so concatenating turns the column - nested. The controller's per_sample_shapes is the only place a mis-ordered or - overwritten merge shows up; the payload alone would still look correct. - """ - keys = [f"span_{i}" for i in range(8)] - lengths = [i + 1 for i in range(len(keys))] - for key, length in zip(keys, lengths, strict=True): - tq.kv_put(key=key, partition_id=PARTITION, fields={"tokens": torch.ones(length, dtype=torch.int64)}) - - client = tq.get_client() - metadata = client.kv_retrieve_meta(keys=keys, partition_id=PARTITION, create=False) - # Same addition for every row, so this holds whichever order the metadata came back in. - values = TensorDict({"tokens": torch.full((len(keys), 1), 9, dtype=torch.int64)}, batch_size=[len(keys)]) - client.update(metadata, ["tokens"], values=values, parser=concat) - - got = tq.kv_batch_get(keys=keys, partition_id=PARTITION, select_fields="tokens") - for i, (key, length) in enumerate(zip(keys, lengths, strict=True)): - row = got["tokens"][i] - assert row.shape == (length + 1,), f"{key}: expected len {length + 1}, got {tuple(row.shape)}" - assert row[-1].item() == 9, f"{key}: appended value missing" - - partition = partition_snapshot() - tokens = partition.field_metadata["tokens"] - assert tokens.is_nested is True, "mixed lengths must promote the column to nested" - for key, length in zip(keys, lengths, strict=True): - recorded = tuple(tokens.per_sample_shapes[partition.keys_mapping[key]]) - assert recorded == (length + 1,), f"{key}: controller recorded {recorded}, expected {(length + 1,)}" - - -@ray.remote(num_cpus=0) -class Updater: - """A worker that attaches to the running TransferQueue and updates its own key.""" - - def __init__(self): - tq.init() - - def run(self, key: str, addition: int) -> str: - tq.kv_update( - key=key, - partition_id=PARTITION, - fields="tokens", - values=torch.tensor([addition]), - parser=concat, - ) - return key - - -def check_updates_from_remote_workers(): - """Updates issued by other processes on the cluster, not just the driver.""" - keys = [f"worker_{i}" for i in range(4)] - for i, key in enumerate(keys): - tq.kv_put(key=key, partition_id=PARTITION, fields={"tokens": torch.tensor([i])}) - - updaters = [Updater.remote() for _ in keys] - ray.get([u.run.remote(key, 100 + i) for i, (u, key) in enumerate(zip(updaters, keys, strict=True))]) - for u in updaters: - ray.kill(u) - - got = tq.kv_batch_get(keys=keys, partition_id=PARTITION, select_fields="tokens") - for i, key in enumerate(keys): - expected = torch.tensor([i, 100 + i]) - assert torch.equal(got["tokens"][i], expected), f"{key}: {got['tokens'][i]} != {expected}" - - -# ==================== harness ==================== - - -def partition_snapshot(): - controller = ray.get_actor("TransferQueueController", namespace="transfer_queue") - return ray.get(controller.get_partition_snapshot.remote(PARTITION)) - - -CHECKS = [ - ("concat keeps the field name", check_concat_keeps_field_name), - ("parser sees old=None for a new field", check_parser_sees_none_for_a_new_field), - ("several fields in one call", check_multiple_fields_in_one_call), - ("non-tensor payload", check_non_tensor_payload), - ("kv_empty keeps the key, updates the controller", check_empty_keeps_the_key_and_updates_the_controller), - ("a failed parser leaves the row unchanged", check_failed_parser_leaves_the_row_unchanged), - ("rejected arguments", check_rejected_arguments), - ("async_kv_update / async_kv_empty", check_async_variants), - ("multi-sample update across units", check_multi_sample_update_across_units), - ("updates from remote workers", check_updates_from_remote_workers), -] - - -def refuse_if_already_deployed() -> None: - """Never attach to a TransferQueue somebody else is using. - - tq.init() silently joins an existing controller by name, so on a shared cluster - this script would otherwise run its checks against a live deployment and clear - its keys on the way out. - """ - try: - ray.get_actor("TransferQueueController", namespace="transfer_queue") - except ValueError: - return - raise SystemExit( - "A TransferQueueController is already running on this cluster. This script needs a " - "deployment of its own; stop the owning job or run it on an idle cluster." - ) - - -def main() -> int: - if not ray.is_initialized(): - ray.init(namespace="transfer_queue") - - print("=" * 78) - print("kv_update end-to-end checks") - print(f"cluster: {ray.get_runtime_context().gcs_address} nodes: {len(ray.nodes())}") - print("=" * 78) - - refuse_if_already_deployed() - tq.init(OmegaConf.create(CONFIG)) - failures = [] - try: - for name, check in CHECKS: - try: - check() - except Exception as e: - failures.append(name) - print(f" FAIL {name}: {type(e).__name__}: {e}") - traceback.print_exc() - else: - print(f" ok {name}") - finally: - # Each check owns its keys; clear between them so one failure cannot cascade. - leftover = list(tq.kv_list(partition_id=PARTITION).get(PARTITION, {})) - if leftover: - tq.kv_clear(keys=leftover, partition_id=PARTITION) - finally: - tq.close() - - print("=" * 78) - print(f"{len(CHECKS) - len(failures)}/{len(CHECKS)} checks passed") - if failures: - print("failed: " + ", ".join(failures)) - print("=" * 78) - return 1 if failures else 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/tests/e2e/test_kv_interface_e2e.py b/tests/e2e/test_kv_interface_e2e.py index c3f9c447..875aedaf 100644 --- a/tests/e2e/test_kv_interface_e2e.py +++ b/tests/e2e/test_kv_interface_e2e.py @@ -1111,115 +1111,139 @@ def test_field_expansion_across_samples(self, controller, tq_api): class TestKVUpdateE2E: - """kv_update: SimpleStorage runs parsers on the unit; KV backends reject.""" + """SimpleStorage merge and empty behavior through sync and async public APIs.""" - def test_concat_then_empty(self, controller, tq_api, backend_name): + def test_update_then_empty_preserves_partition_field_metadata(self, controller, tq_api, backend_name): if backend_name != "SimpleStorage": - pytest.skip("parser-backed kv_update is implemented only for SimpleStorage") + pytest.skip("merge-backed kv_update is implemented only for SimpleStorage") partition_id = "test_partition" - key = "sample_update" - - tq_api.kv_put(key=key, partition_id=partition_id, fields={"tokens": torch.tensor([1, 2, 3])}, tag={"step": 1}) + keys = ["empty_0", "empty_1", "empty_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=key, partition_id=partition_id, fields="tokens", values=torch.tensor([4, 5]), parser=concat + key=keys[1], + partition_id=partition_id, + fields={"tokens": torch.tensor([12, 13])}, + merge_fn=concat, ) assert "tokens" in meta.fields - - retrieved = tq_api.kv_batch_get(keys=key, partition_id=partition_id, select_fields="tokens") - assert_tensor_equal(retrieved["tokens"][0], torch.tensor([1, 2, 3, 4, 5])) - - tq_api.kv_empty(key=key, partition_id=partition_id, fields="tokens") - emptied = tq_api.kv_batch_get(keys=key, partition_id=partition_id, select_fields="tokens") + tq_api.kv_empty(keys=keys[0], partition_id=partition_id, fields="tokens") + emptied = tq_api.kv_batch_get(keys=keys[0], partition_id=partition_id, select_fields="tokens") assert emptied["tokens"][0] is None partition = get_controller_partition(controller, partition_id) col = partition.field_name_mapping["tokens"] - global_idx = partition.keys_mapping[key] + global_idx = partition.keys_mapping[keys[0]] assert partition.production_status[global_idx, col] == 1 - # The column now holds None, so the controller must not still describe it as a tensor. - assert partition.field_metadata["tokens"].is_non_tensor is True - - tq_api.kv_clear(keys=key, partition_id=partition_id) + assert partition.field_metadata["tokens"].is_non_tensor is False + assert partition.field_metadata["tokens"].dtype == torch.int64 - def test_concat_prompt_and_response_keeps_field_name(self, tq_api, backend_name): - """prompt_ids already stored; update appends response_ids; field stays sequence_ids.""" - if backend_name != "SimpleStorage": - pytest.skip("parser-backed kv_update is implemented only for SimpleStorage") - - partition_id = "test_partition" - key = "sample_seq" - prompt_ids = torch.tensor([10, 11, 12]) - response_ids = torch.tensor([20, 21]) - - tq_api.kv_put(key=key, partition_id=partition_id, fields={"sequence_ids": prompt_ids}) - tq_api.kv_update( - key=key, + tq_api.kv_batch_put( + keys=["nested_0", "nested_1"], partition_id=partition_id, - fields="sequence_ids", - values=response_ids, - parser=lambda old, new: torch.cat([old, new]), + 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 - retrieved = tq_api.kv_batch_get(keys=key, partition_id=partition_id) - assert "sequence_ids" in retrieved - assert "prompt_ids" not in retrieved - assert "response_ids" not in retrieved - assert_tensor_equal(retrieved["sequence_ids"][0], torch.tensor([10, 11, 12, 20, 21])) - - tq_api.kv_clear(keys=key, partition_id=partition_id) - - def test_update_multiple_fields_at_once(self, tq_api, backend_name): + def test_batch_update_multiple_fields_across_storage_units(self, controller, tq_api, backend_name): if backend_name != "SimpleStorage": - pytest.skip("parser-backed kv_update is implemented only for SimpleStorage") + pytest.skip("merge-backed kv_update is implemented only for SimpleStorage") partition_id = "test_partition" - key = "sample_multi" - - tq_api.kv_put( - key=key, + 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={"a": torch.tensor([1, 2]), "b": torch.tensor([10])}, + 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_update( - key=key, + tq_api.kv_batch_update( + keys=keys, partition_id=partition_id, - fields=["a", "b"], - values={"a": torch.tensor([3]), "b": torch.tensor([20, 30])}, - parser=lambda old, new: torch.cat([old, new]), + 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=key, partition_id=partition_id) - assert_tensor_equal(retrieved["a"][0], torch.tensor([1, 2, 3])) - assert_tensor_equal(retrieved["b"][0], torch.tensor([10, 20, 30])) + 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])) - tq_api.kv_empty(key=key, partition_id=partition_id, fields=["a", "b"]) - emptied = tq_api.kv_batch_get(keys=key, partition_id=partition_id) - assert emptied["a"][0] is None - assert emptied["b"][0] is None + 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=key, partition_id=partition_id) + tq_api.kv_empty(keys=keys, partition_id=partition_id, fields=["tokens", "score"]) + emptied = tq_api.kv_batch_get(keys=keys, partition_id=partition_id) + assert all(value is None for value in emptied["tokens"]) + assert all(value is None for value in emptied["score"]) - def test_missing_key_and_empty_with_values(self, tq_api, backend_name): + def test_rejects_missing_key_unproduced_field_and_incompatible_result(self, tq_api, backend_name): if backend_name != "SimpleStorage": - pytest.skip("parser-backed kv_update is implemented only for SimpleStorage") + pytest.skip("merge-backed kv_update is implemented only for SimpleStorage") - with pytest.raises(ValueError, match="must not specify values"): - tq_api.kv_update( - key="no_such", partition_id="test_partition", fields="tokens", values=torch.tensor([1]), empty=True - ) with pytest.raises(ValueError, match="not found"): tq_api.kv_update( key="no_such", partition_id="test_partition", - fields="tokens", - values=torch.tensor([1]), - parser=lambda old, new: new, + 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": @@ -1230,9 +1254,8 @@ def test_kv_backend_rejects(self, tq_api, backend_name): tq_api.kv_update( key="k", partition_id="test_partition", - fields="tokens", - values=torch.tensor([2]), - parser=lambda old, new: new, + fields={"tokens": torch.tensor([2])}, + merge_fn=lambda _old, new: new, ) tq_api.kv_clear(keys="k", partition_id="test_partition") diff --git a/tests/test_kv_update.py b/tests/test_kv_update.py index 0ab2f83f..567722ff 100644 --- a/tests/test_kv_update.py +++ b/tests/test_kv_update.py @@ -19,9 +19,9 @@ import numpy as np import pytest import torch +from tensordict import TensorDict -from transfer_queue.controller import DataPartitionStatus -from transfer_queue.interface import _normalize_kv_update_args +from transfer_queue.interface import _single_update_batch from transfer_queue.metadata import BatchMeta from transfer_queue.storage.managers.base import KVStorageManager from transfer_queue.storage.managers.mooncake_manager import MooncakeStorageManager @@ -32,120 +32,107 @@ def _concat(old, new): - if old is None: - return new return torch.cat([old, new]) -def test_normalize_empty_rejects_values_and_parser(): - with pytest.raises(ValueError, match="must not specify values"): - _normalize_kv_update_args("tokens", torch.tensor([1]), None, empty=True) - with pytest.raises(ValueError, match="must not specify parser"): - _normalize_kv_update_args("tokens", None, _concat, empty=True) - names, batch, parser, use_empty = _normalize_kv_update_args("tokens", None, None, empty=True) - assert names == ["tokens"] - assert batch is None - assert parser is None - assert use_empty is True +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_normalize_custom_parser_requires_values(): - with pytest.raises(ValueError, match="requires values"): - _normalize_kv_update_args("tokens", None, _concat) +def test_single_update_batch_wraps_tensor_and_python_values(): + batch = _single_update_batch({"tokens": torch.tensor([4, 5]), "meta": {"step": 1}}) - -def test_normalize_rejects_non_callable_parser(): - with pytest.raises(TypeError, match="parser must be callable unless empty=True"): - _normalize_kv_update_args("tokens", torch.tensor([1]), None) - - -def test_normalize_single_field_wraps_a_batch(): - names, batch, parser, use_empty = _normalize_kv_update_args("tokens", torch.tensor([4, 5]), _concat) - assert names == ["tokens"] - assert use_empty is False - assert parser is _concat - assert batch is not None assert batch.batch_size == torch.Size([1]) assert torch.equal(batch["tokens"][0], torch.tensor([4, 5])) + assert batch["meta"][0] == {"step": 1} -def test_normalize_multi_field_requires_matching_dict(): - with pytest.raises(TypeError, match="must be a dict"): - _normalize_kv_update_args(["a", "b"], torch.tensor([1]), _concat) - with pytest.raises(ValueError, match="same columns"): - _normalize_kv_update_args(["a", "b"], {"a": 1}, _concat) +@pytest.mark.parametrize("fields", [{}, [], {"tokens": torch.tensor([1]), 2: "bad"}]) +def test_single_update_batch_rejects_invalid_fields(fields): + with pytest.raises(TypeError): + _single_update_batch(fields) -def test_apply_update_concat_prompt_and_response_keeps_field_name(): - """Stored prompt_ids plus new response_ids become the sequence; the field name is unchanged.""" - prompt_ids = torch.tensor([10, 11, 12]) - response_ids = torch.tensor([20, 21]) +def test_apply_update_concatenates_and_keeps_field_name(): data = StorageUnitData() - data.put_data({"sequence_ids": [prompt_ids.clone()]}, [0]) + data.put_data({"sequence_ids": [torch.tensor([10, 11, 12])]}, [0]) - described = data.apply_update([0], ["sequence_ids"], {"sequence_ids": [response_ids]}, _concat, False) + described = data.apply_update( + [0], + {"sequence_ids": [torch.tensor([20, 21])]}, + _concat, + _tensor_schema("sequence_ids"), + ) - assert list(described) == ["sequence_ids"] + 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])) - assert "prompt_ids" not in data.field_data - assert "response_ids" not in data.field_data + assert set(data.field_data) == {"sequence_ids"} -def test_apply_update_concatenates_and_is_atomic_on_parser_error(): +def test_apply_update_is_atomic_on_merge_error(): data = StorageUnitData() - data.put_data({"tokens": [torch.tensor([1, 2, 3])]}, [7]) + original = torch.tensor([1, 2, 3]) + data.put_data({"tokens": [original.clone()]}, [7]) - described = data.apply_update([7], ["tokens"], {"tokens": [torch.tensor([4, 5])]}, _concat, False) - assert described["tokens"] == {"dtype": torch.int64, "shapes": [(5,)]} - assert torch.equal(data.field_data["tokens"][7], torch.tensor([1, 2, 3, 4, 5])) + def fail(_old, _new): + raise RuntimeError("merge failed") - def boom(old, new): - raise RuntimeError("parser failed") + with pytest.raises(RuntimeError, match="merge failed"): + data.apply_update([7], {"tokens": [torch.tensor([9])]}, fail, _tensor_schema("tokens")) + assert torch.equal(data.field_data["tokens"][7], original) - with pytest.raises(RuntimeError, match="parser failed"): - data.apply_update([7], ["tokens"], {"tokens": [torch.tensor([9])]}, boom, False) - assert torch.equal(data.field_data["tokens"][7], torch.tensor([1, 2, 3, 4, 5])) - -def test_apply_update_empty_stores_none(): +def test_apply_update_rejects_unproduced_field(): data = StorageUnitData() - data.put_data({"tokens": [torch.tensor([1])], "keep": [torch.tensor([2])]}, [3]) - described = data.apply_update([3], ["tokens"], None, None, True) - assert described["tokens"]["shapes"] is None - assert data.field_data["tokens"][3] is None - assert torch.equal(data.field_data["keep"][3], torch.tensor([2])) + data.put_data({"tokens": [torch.tensor([1])]}, [3]) + + 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 -def test_apply_update_missing_field_passes_none_as_old(): +@pytest.mark.parametrize( + "merge_fn, error", + [ + (lambda old, _new: old.to(torch.float32), "keep dtype"), + (lambda _old, _new: {"not": "a tensor"}, "keep dtype"), + ], +) +def test_apply_update_rejects_incompatible_tensor_results_before_writing(merge_fn, error): data = StorageUnitData() - seen = [] + original = torch.tensor([1, 2]) + data.put_data({"tokens": [original.clone()]}, [0]) - def record(old, new): - seen.append(old) - return new + with pytest.raises(TypeError, match=error): + data.apply_update([0], {"tokens": [torch.tensor([3])]}, merge_fn, _tensor_schema("tokens")) + assert torch.equal(data.field_data["tokens"][0], original) - data.apply_update([1], ["fresh"], {"fresh": [torch.tensor([8])]}, record, False) - assert seen == [None] - assert torch.equal(data.field_data["fresh"][1], torch.tensor([8])) + +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) + assert data.field_data["meta"][0] == {"step": 1} def test_apply_update_decodes_ssd_offloaded_old_value(tmp_path): - """With SSD offload on, the parser must see the stored tensor, not its file reference.""" data = HybridStorageUnitData( storage_size=4, threshold_bytes=64, ssd_path=str(tmp_path), run_id="run", unit_id="unit" ) prompt = torch.arange(32) data.put_data({"tokens": [prompt]}, [0]) - assert data.ssd_active_values == 1, "precondition: the prompt must live on SSD" + assert data.ssd_active_values == 1 - data.apply_update([0], ["tokens"], {"tokens": [torch.tensor([99])]}, _concat, False) + data.apply_update([0], {"tokens": [torch.tensor([99])]}, _concat, _tensor_schema("tokens")) assert torch.equal(data.get_data(["tokens"], [0])["tokens"][0], torch.cat([prompt, torch.tensor([99])])) - assert data.ssd_active_values == 1, "the replaced SSD file must be released, not leaked" + assert data.ssd_active_values == 1 def test_apply_update_is_atomic_when_a_later_field_fails_to_reach_ssd(tmp_path, monkeypatch): - """Concat is not idempotent, so a partly applied update would double-apply on retry.""" data = HybridStorageUnitData( storage_size=4, threshold_bytes=64, ssd_path=str(tmp_path), run_id="run", unit_id="unit" ) @@ -153,7 +140,6 @@ def test_apply_update_is_atomic_when_a_later_field_fails_to_reach_ssd(tmp_path, prompt, mask = torch.arange(32), torch.ones(32, dtype=torch.int64) data.put_data({"tokens": [prompt], "mask": [mask]}, [0]) old_files = set(tmp_path.rglob("*.bin")) - write_values = data._ssd_store.write_values calls = [] @@ -166,7 +152,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], ["tokens", "mask"], new_data, _concat, False) + data.apply_update([0], new_data, _concat, _tensor_schema("tokens", "mask")) stored = data.get_data(["tokens", "mask"], [0]) assert torch.equal(stored["tokens"][0], prompt) @@ -178,7 +164,6 @@ def fail_second_field(encoded_values): def test_build_update_field_schema_orders_shapes_across_units(): - """Units describe only their own rows; the batch schema must follow metadata order.""" described = _build_update_field_schema( [0, 1, 2, 3], [ @@ -209,36 +194,6 @@ def test_build_update_field_schema_keeps_uniform_column_flat(): } -def test_build_update_field_schema_marks_column_non_tensor_if_any_unit_is(): - described = _build_update_field_schema( - [0, 1], - [ - ([0], {"tokens": {"dtype": torch.int64, "shapes": [(4,)]}}), - ([1], {"tokens": {"dtype": None, "shapes": None}}), - ], - ) - - assert described["tokens"]["is_non_tensor"] is True - assert described["tokens"]["shape"] is None - - -def test_empty_marks_the_controller_field_non_tensor(): - """tq.kv_empty stores None, so the controller must stop describing the column as a tensor.""" - partition = DataPartitionStatus(partition_id="p") - partition._update_field_metadata( - [0], {"tokens": {"dtype": torch.int64, "shape": (5,), "is_nested": False, "is_non_tensor": False}} - ) - - partition._update_field_metadata( - [0], {"tokens": {"dtype": None, "shape": None, "is_nested": False, "is_non_tensor": True}} - ) - - tokens = partition.field_metadata["tokens"] - assert tokens.is_non_tensor is True - assert tokens.shape is None - assert tokens.to_batch_schema([0])["is_non_tensor"] is True - - @pytest.mark.asyncio @pytest.mark.parametrize( "manager_cls", [MooncakeStorageManager, YuanrongStorageManager, RayStorageManager, KVStorageManager] @@ -246,9 +201,7 @@ def test_empty_marks_the_controller_field_non_tensor(): @patch("transfer_queue.storage.managers.base.StorageClientFactory.create") @patch.object(KVStorageManager, "_connect_to_controller", lambda self: None) async def test_kv_backends_reject_update(mock_create, manager_cls): - """Every KV backend must refuse kv_update by name; only SimpleStorage implements it.""" mock_create.return_value = MagicMock() - # Each manager validates its own config before reaching update_data. config = { KVStorageManager: {"client_name": "YuanrongStorageClient"}, YuanrongStorageManager: {"worker_port": 31501}, @@ -260,5 +213,7 @@ async def test_kv_backends_reject_update(mock_create, manager_cls): field_schema={"x": {"dtype": torch.int64, "shape": (1,), "is_nested": False, "is_non_tensor": False}}, production_status=np.ones(1, dtype=np.int8), ) + 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, ["x"], empty=True) + await manager.update_data(meta, values, lambda _old, new: new) diff --git a/tests/test_storage_request_retry.py b/tests/test_storage_request_retry.py index 59f82edc..70be8985 100644 --- a/tests/test_storage_request_retry.py +++ b/tests/test_storage_request_retry.py @@ -29,7 +29,7 @@ from transfer_queue.storage.managers.simple_storage_manager import AsyncSimpleStorageManager, StorageUnitTimeout from transfer_queue.utils import common from transfer_queue.utils.enum_utils import Role -from transfer_queue.utils.zmq_utils import ZMQServerInfo +from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, ZMQServerInfo def _manager(with_unit: bool = True) -> AsyncSimpleStorageManager: @@ -57,6 +57,37 @@ def _single_sample_batch() -> tuple[TensorDict, BatchMeta]: return TensorDict({"input_ids": torch.zeros(1, 2, dtype=torch.int64)}, batch_size=1), metadata +class _SocketLease: + def __init__(self, socket): + self.socket = socket + + async def __aenter__(self): + return self.socket + + async def __aexit__(self, *_args): + return None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("accepted", [False, True]) +async def test_controller_notification_returns_ack_status(accepted): + manager = _manager() + manager.controller_info = MagicMock(id="controller") + socket = MagicMock() + socket.send_multipart = AsyncMock() + socket.recv_multipart = AsyncMock( + return_value=ZMQMessage.create( + request_type=ZMQRequestType.NOTIFY_DATA_UPDATE_ACK, + sender_id="controller", + body={"success": accepted}, + ).serialize() + ) + manager.notify_pool = MagicMock() + manager.notify_pool.alease.return_value = _SocketLease(socket) + + assert await manager._notify_and_wait([b"request"]) is accepted + + @pytest.mark.asyncio @pytest.mark.parametrize("data_parser, expected_attempts", [(None, 3), (lambda data: data, 1)]) async def test_put_retries_only_without_a_parser(data_parser, expected_attempts): @@ -76,6 +107,36 @@ async def test_put_retries_only_without_a_parser(data_parser, expected_attempts) manager.notify_data_update.assert_not_awaited() +@pytest.mark.asyncio +async def test_update_reports_controller_rejection_after_storage_commit(): + manager = _manager() + manager._update_to_single_storage_unit = AsyncMock( + return_value={"input_ids": {"dtype": torch.int64, "shapes": [(3,)]}} + ) + manager.notify_data_update = AsyncMock(return_value=False) + 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])) + + manager._update_to_single_storage_unit.assert_awaited_once() + manager.notify_data_update.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_update_timeout_is_not_retried_or_published(): + manager = _manager() + manager._update_to_single_storage_unit = AsyncMock(side_effect=StorageUnitTimeout("outcome unknown")) + manager.notify_data_update = AsyncMock() + 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])) + + manager._update_to_single_storage_unit.assert_awaited_once() + manager.notify_data_update.assert_not_awaited() + + @pytest.mark.asyncio async def test_lost_request_recovers_without_diagnosis(): manager = _manager() From bfc43a354c5a0d053bf65fe199cd0383e7650066 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Thu, 8 Oct 2026 22:01:26 +0800 Subject: [PATCH 8/9] [docs] Clarify kv_update semantics and tradeoffs Signed-off-by: OutstanderWang --- README.md | 9 ++++++++- docs/ssd_offload.md | 24 ++++++++++++++++++------ tutorial/02_kv_interface.py | 9 ++++----- 3 files changed, 30 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index 62610857..8482d9f2 100644 --- a/README.md +++ b/README.md @@ -116,10 +116,17 @@ 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**: Rewrite selected fields of an existing key with `parser(old, new)`, or `empty=True` to store `None`. `kv_empty()` is the same as `kv_update(..., empty=True)`. SimpleStorage only. +- **(async_)kv_update / (async_)kv_batch_update**: Merge new values into already-produced fields with `merge_fn(old, new)`. SimpleStorage only. +- **(async_)kv_empty**: Store `None` in selected fields while keeping their keys ready. SimpleStorage only. - **(async_)kv_list**: List keys and tags (metadata) in a partition. - **(async_)kv_clear**: Remove key-value pairs from storage. +`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** - **Redis-style Semantics**: Familiar KV interface (Put/Get/List) for a zero learning curve. diff --git a/docs/ssd_offload.md b/docs/ssd_offload.md index b0e5b66a..99507bfb 100644 --- a/docs/ssd_offload.md +++ b/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 emptying 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. diff --git a/tutorial/02_kv_interface.py b/tutorial/02_kv_interface.py index 601f80d8..bd182de1 100644 --- a/tutorial/02_kv_interface.py +++ b/tutorial/02_kv_interface.py @@ -159,12 +159,11 @@ def demonstrate_kv_api(): tq.kv_update( key=key, partition_id=partition_id, - fields="scratch", - values=torch.tensor([3, 4]), - parser=lambda old, new: torch.cat([old, new]), + fields={"scratch": torch.tensor([3, 4])}, + merge_fn=lambda old, new: torch.cat([old, new]), ) print(" ✓ kv_update concat: scratch of '0_0' is now [1, 2, 3, 4].") - tq.kv_empty(key=key, partition_id=partition_id, fields="scratch") + tq.kv_empty(keys=key, partition_id=partition_id, fields="scratch") print(" ✓ tq.kv_empty: scratch of '0_0' is stored as None (key remains).") # Step 6: List all keys and tags in a partition @@ -216,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 - Rewrite fields with parser(old, new), or empty=True / tq.kv_empty to store None + 3. (async_)kv_update - Merge new values into produced fields with merge_fn(old, new) 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 from storage From 369a79c9df09cf4f5d22d6e1499e2daf8ed0cc00 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 9 Oct 2026 16:12:04 +0800 Subject: [PATCH 9/9] [feat] Fold kv_empty into field-level kv_clear kv_clear(keys, partition_id, fields=...) now clears only the given fields and keeps the keys. The controller marks the cells as not produced, removes the samples from each field's FieldMeta and drops their backend metadata before storage deletes the values, so a cleared field reads as never written and later writes to the same field rebuild its schema normally. Tags, consumption status and global indexes are left untouched. Storage managers take an optional fields filter in clear_data: KV backends narrow the metadata with select_fields, and SimpleStorage units drop only those fields (unlinking their SSD files). kv_empty is removed. Signed-off-by: OutstanderWang --- README.md | 5 +- docs/ssd_offload.md | 2 +- tests/e2e/test_kv_interface_e2e.py | 50 +++++++++++---- tests/test_controller_data_partitions.py | 40 ++++++++++++ tests/test_kv_storage_manager.py | 18 ++++++ tests/test_simple_storage_unit.py | 31 +++++++++ transfer_queue/__init__.py | 4 -- transfer_queue/client.py | 23 +++++++ transfer_queue/controller.py | 42 +++++++++++++ transfer_queue/interface.py | 63 +++++-------------- transfer_queue/storage/managers/base.py | 9 ++- .../managers/simple_storage_manager.py | 9 +-- transfer_queue/storage/simple_storage.py | 20 +++--- transfer_queue/utils/zmq_utils.py | 2 + tutorial/02_kv_interface.py | 10 +-- 15 files changed, 244 insertions(+), 84 deletions(-) diff --git a/README.md b/README.md index 8482d9f2..f4b930ae 100644 --- a/README.md +++ b/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. @@ -117,9 +117,8 @@ To simplify the usage of TransferQueue, we provide a Redis-style high-level API - **(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_empty**: Store `None` in selected fields while keeping their keys ready. 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 diff --git a/docs/ssd_offload.md b/docs/ssd_offload.md index 99507bfb..a649cf8f 100644 --- a/docs/ssd_offload.md +++ b/docs/ssd_offload.md @@ -216,7 +216,7 @@ memory use while values are read or encoded. ### Dense batches can retain shared backing memory SimpleStorage slices dense batches without copying. The rows therefore share -the batch allocation, and replacing or emptying one row does not release its +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. diff --git a/tests/e2e/test_kv_interface_e2e.py b/tests/e2e/test_kv_interface_e2e.py index 875aedaf..fa5d488b 100644 --- a/tests/e2e/test_kv_interface_e2e.py +++ b/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.""" @@ -1111,14 +1139,14 @@ def test_field_expansion_across_samples(self, controller, tq_api): class TestKVUpdateE2E: - """SimpleStorage merge and empty behavior through sync and async public APIs.""" + """SimpleStorage merge behavior through sync and async public APIs.""" - def test_update_then_empty_preserves_partition_field_metadata(self, controller, tq_api, backend_name): + 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 = ["empty_0", "empty_1", "empty_2"] + keys = ["clear_field_0", "clear_field_1", "clear_field_2"] tq_api.kv_batch_put( keys=keys, partition_id=partition_id, @@ -1135,14 +1163,13 @@ def concat(old, new): merge_fn=concat, ) assert "tokens" in meta.fields - tq_api.kv_empty(keys=keys[0], partition_id=partition_id, fields="tokens") - emptied = tq_api.kv_batch_get(keys=keys[0], partition_id=partition_id, select_fields="tokens") - assert emptied["tokens"][0] is None + 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] == 1 + 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 @@ -1211,10 +1238,11 @@ def test_batch_update_multiple_fields_across_storage_units(self, controller, tq_ global_idx = partition.keys_mapping[key] assert tuple(tokens.per_sample_shapes[global_idx]) == (length + 1,) - tq_api.kv_empty(keys=keys, partition_id=partition_id, fields=["tokens", "score"]) - emptied = tq_api.kv_batch_get(keys=keys, partition_id=partition_id) - assert all(value is None for value in emptied["tokens"]) - assert all(value is None for value in emptied["score"]) + 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": diff --git a/tests/test_controller_data_partitions.py b/tests/test_controller_data_partitions.py index 2f6b874e..ef8538cc 100644 --- a/tests/test_controller_data_partitions.py +++ b/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() diff --git a/tests/test_kv_storage_manager.py b/tests/test_kv_storage_manager.py index 14980fab..dacb0be2 100644 --- a/tests/test_kv_storage_manager.py +++ b/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.""" diff --git a/tests/test_simple_storage_unit.py b/tests/test_simple_storage_unit.py index fe52f175..b8f40a99 100644 --- a/tests/test_simple_storage_unit.py +++ b/tests/test_simple_storage_unit.py @@ -472,6 +472,12 @@ def test_storage_unit_data_direct(): assert 0 not in storage_data.field_data["log_probs"] # key gone assert 1 in storage_data.field_data["log_probs"] # other key intact + # field-level clear keeps the sample and its other fields + storage_data.clear([1], fields=["rewards", "missing"]) + assert 1 not in storage_data.field_data["rewards"] + assert 1 in storage_data.field_data["log_probs"] + assert storage_data._active_keys == {1} + def test_storage_unit_data_capacity_uses_active_keys(): """Capacity check must use _active_keys, not scan field_data.""" @@ -545,6 +551,31 @@ def test_hybrid_storage_routes_each_sample_by_payload_size_and_clears(tmp_path): assert not (tmp_path / "transfer_queue_ssd_offload" / "test-run").exists() +def test_hybrid_storage_field_clear_unlinks_only_that_field(tmp_path): + storage = HybridStorageUnitData( + storage_size=10, + threshold_bytes=64, + ssd_path=str(tmp_path), + run_id="test-run", + unit_id="test-unit", + ) + try: + storage.put_data({"pixel_values": [b"p" * 100], "ids": [b"i" * 100]}, [1]) + assert storage.ssd_active_values == 2 + + storage.clear([1], fields=["pixel_values"]) + + assert storage.ssd_active_values == 1 + assert storage.ssd_active_bytes == 100 + assert len(list(tmp_path.rglob("*.bin"))) == 1 + assert storage.active_key_count == 1 + assert storage.get_data(["ids"], [1]) == {"ids": [b"i" * 100]} + with pytest.raises(StorageKeyNotFoundError): + storage.get_data(["pixel_values"], [1]) + finally: + storage.close() + + def test_hybrid_storage_cleans_completed_ssd_writes_after_batch_failure(tmp_path, monkeypatch): """A failed SSD batch removes files written by its successful tasks.""" storage = HybridStorageUnitData( diff --git a/transfer_queue/__init__.py b/transfer_queue/__init__.py index 4f3c19ad..8df00be4 100644 --- a/transfer_queue/__init__.py +++ b/transfer_queue/__init__.py @@ -23,7 +23,6 @@ async_kv_batch_put, async_kv_batch_update, async_kv_clear, - async_kv_empty, async_kv_list, async_kv_put, async_kv_update, @@ -37,7 +36,6 @@ kv_batch_put, kv_batch_update, kv_clear, - kv_empty, kv_list, kv_put, kv_update, @@ -67,7 +65,6 @@ "kv_list", "kv_clear", "kv_update", - "kv_empty", "async_kv_put", "async_kv_batch_put", "async_kv_batch_update", @@ -76,7 +73,6 @@ "async_kv_list", "async_kv_clear", "async_kv_update", - "async_kv_empty", "KVBatchMeta", ] + [ diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 41667747..41ff29d4 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -1010,6 +1010,29 @@ async def async_kv_retrieve_meta( except Exception as e: raise RuntimeError(f"[{self.client_id}] Failed in async_kv_retrieve_meta: {e}") from e + 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. + + The controller marks the fields as unproduced before storage deletes them, so no + reader sees a ready field whose value is gone. + """ + metadata = await self._kv_clear_fields_in_controller(keys, partition_id, fields) + if metadata.size > 0: + await self.storage_manager.clear_data(metadata, fields=metadata.field_names) + + @with_controller_socket + async def _kv_clear_fields_in_controller( + self, keys: list[str], partition_id: str, fields: list[str], socket=None + ) -> BatchMeta: + """Clear fields in the controller and return the cleared cells' metadata.""" + response_msg = await self._request_controller( + socket=socket, + request_type=ZMQRequestType.KV_CLEAR_FIELDS, + response_type=ZMQRequestType.KV_CLEAR_FIELDS_RESPONSE, + body={"keys": keys, "partition_id": partition_id, "fields": fields}, + ) + return response_msg.body["metadata"] + @with_controller_socket async def async_kv_retrieve_keys( self, diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index c86c0824..c2ba668f 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -941,6 +941,22 @@ def clear_data(self, indexes_to_release: list[int], clear_consumption: bool = Tr f"Attempted to clear global_indexes: {indexes_to_release}" ) + def clear_fields(self, global_indexes: list[int], field_names: list[str]): + """Mark fields as never written for these samples; keys, tags, consumption and indexes stay.""" + for field_name in field_names: + col = self.field_name_mapping.get(field_name) + if col is not None: + self.production_status[global_indexes, col] = 0 + field_meta = self.field_metadata.get(field_name) + if field_meta is not None: + field_meta.remove_samples(global_indexes) + if not field_meta.global_indexes: + self.field_metadata.pop(field_name) + for idx in global_indexes: + backend_meta = self.field_custom_backend_meta.get(idx, {}) + for field_name in field_names: + backend_meta.pop(field_name, None) + def kv_retrieve_indexes(self, keys: list[str]) -> list[int | None]: """Translate the user-specified keys to global_indexes""" global_indexes = [self.keys_mapping.get(k, None) for k in keys] @@ -1701,6 +1717,24 @@ def kv_retrieve_meta( return self.generate_batch_meta(partition_id, verified_global_indexes, data_fields, mode="force_fetch") + 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. + + Missing keys and fields are ignored. The metadata is captured before clearing so + KV backends still receive the custom backend metadata needed to delete the values. + """ + partition = self._get_partition(partition_id) + if partition is None: + return BatchMeta.empty() + global_indexes = [idx for idx in partition.kv_retrieve_indexes(keys) if idx is not None] + fields = [f for f in fields if f in partition.field_metadata] + if not global_indexes or not fields: + return BatchMeta.empty() + + metadata = self.generate_batch_meta(partition_id, global_indexes, fields, mode="force_fetch") + partition.clear_fields(global_indexes, fields) + return metadata + def kv_retrieve_keys( self, global_indexes: list[int], @@ -2010,6 +2044,7 @@ def _handle_request(self, request_msg: ZMQMessage, monitor: Any) -> ZMQMessage | ZMQRequestType.KV_LIST: self._handle_kv_list_request, ZMQRequestType.KV_DESCRIBE: self._handle_kv_describe_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, ZMQRequestType.LOAD_CONTROLLER_CHECKPOINT: self._handle_load_controller_checkpoint_request, } @@ -2212,6 +2247,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_clear_fields_request(self, request_msg: ZMQMessage) -> ZMQMessage: + params = request_msg.body + metadata = self.kv_clear_fields( + keys=params["keys"], partition_id=params["partition_id"], fields=params["fields"] + ) + return self._make_response(request_msg, ZMQRequestType.KV_CLEAR_FIELDS_RESPONSE, {"metadata": metadata}) + def _handle_kv_list_request(self, request_msg: ZMQMessage) -> ZMQMessage: requested_partition_id = request_msg.body["partition_id"] partition_ids = list(self.partitions.keys()) if requested_partition_id is None else [requested_partition_id] diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index e496f0dc..672f62d8 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -585,7 +585,7 @@ 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_clear(keys: list[str] | str, partition_id: str) -> None: +def kv_clear(keys: list[str] | str, partition_id: str, fields: list[str] | str | None = None) -> None: """Clear key-value pairs from TransferQueue. This removes the specified keys and their associated data from both @@ -594,6 +594,9 @@ def kv_clear(keys: list[str] | str, partition_id: str) -> None: Args: keys: Single key or list of keys to clear partition_id: Partition containing the keys + fields: If given, clear only these fields and keep the keys. A cleared field + reads as never written: it is not ready and can be put again. Tags and + consumption status are kept. Missing keys and fields are ignored. Example: >>> import transfer_queue as tq @@ -602,10 +605,12 @@ def kv_clear(keys: list[str] | str, partition_id: str) -> None: >>> tq.kv_clear(keys="sample_1", partition_id="train") >>> # Clear multiple keys >>> tq.kv_clear(keys=["sample_1", "sample_2"], partition_id="train") + >>> # Release one field but keep the keys + >>> tq.kv_clear(keys=["sample_1", "sample_2"], partition_id="train", fields="pixel_values") """ tq_client = _maybe_create_tq_client() - tq_client._run_coroutine(async_kv_clear(keys=keys, partition_id=partition_id)) + tq_client._run_coroutine(async_kv_clear(keys=keys, partition_id=partition_id, fields=fields)) def kv_update( @@ -663,12 +668,6 @@ def kv_batch_update( ) -def kv_empty(keys: str | list[str], partition_id: str, fields: str | list[str]) -> KVBatchMeta: - """Release produced fields by storing ``None`` while keeping the keys ready.""" - tq_client = _maybe_create_tq_client() - return tq_client._run_coroutine(async_kv_empty(keys=keys, partition_id=partition_id, fields=fields)) - - # ==================== KV Interface API ==================== async def async_kv_put( key: str, @@ -1021,7 +1020,7 @@ async def async_kv_list(partition_id: str | None = None) -> dict[str, dict[str, return partition_info -async def async_kv_clear(keys: list[str] | str, partition_id: str) -> None: +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. This removes the specified keys and their associated data from both @@ -1030,6 +1029,7 @@ async def async_kv_clear(keys: list[str] | str, partition_id: str) -> None: Args: keys: Single key or list of keys to clear partition_id: Partition containing the keys + fields: If given, clear only these fields and keep the keys. See ``kv_clear``. Example: >>> import transfer_queue as tq @@ -1038,12 +1038,20 @@ async def async_kv_clear(keys: list[str] | str, partition_id: str) -> None: >>> await tq.async_kv_clear(keys="sample_1", partition_id="train") >>> # Clear multiple keys >>> await tq.async_kv_clear(keys=["sample_1", "sample_2"], partition_id="train") + >>> # Release one field but keep the keys + >>> await tq.async_kv_clear(keys=["sample_1", "sample_2"], partition_id="train", fields="pixel_values") """ if isinstance(keys, str): keys = [keys] tq_client = _maybe_create_tq_client() + if fields is not None: + await tq_client.async_kv_clear_fields( + keys=keys, partition_id=partition_id, fields=[fields] if isinstance(fields, str) else fields + ) + return + batch_meta = await tq_client.async_kv_retrieve_meta(keys=keys, partition_id=partition_id, create=False) if batch_meta.size > 0: @@ -1117,43 +1125,6 @@ async def async_kv_batch_update( ) -async def async_kv_empty(keys: str | list[str], partition_id: str, fields: str | list[str]) -> KVBatchMeta: - """Asynchronously release fields by storing ``None`` through the ordinary put path.""" - if isinstance(keys, str): - keys = [keys] - if isinstance(fields, str): - field_names = [fields] - elif isinstance(fields, list) and fields and all(isinstance(name, str) for name in fields): - field_names = list(dict.fromkeys(fields)) - else: - raise TypeError("fields must be a field name or a non-empty list of field names") - if not keys: - raise ValueError("keys must not be empty") - - tq_client = _maybe_create_tq_client() - if not isinstance(tq_client.storage_manager, AsyncSimpleStorageManager): - raise NotImplementedError("kv_empty requires the SimpleStorage backend") - batch_meta = await tq_client.async_kv_retrieve_meta(keys=keys, partition_id=partition_id, create=False) - if batch_meta.size != len(keys): - raise ValueError("Some keys or the partition were not found") - missing_fields = sorted(set(field_names) - set(batch_meta.field_names)) - if missing_fields: - raise ValueError(f"Fields are not produced for every key: {missing_fields}") - - empty_values = TensorDict( - {field_name: NonTensorStack(*([None] * len(keys))) for field_name in field_names}, - batch_size=[len(keys)], - ) - batch_meta = await tq_client.async_put(empty_values, batch_meta) - return KVBatchMeta( - keys=keys, - tags=batch_meta.custom_meta, - partition_id=partition_id, - fields=batch_meta.field_names, - extra_info=batch_meta.extra_info, - ) - - # ==================== Low-Level Native API ==================== # For low-level API support, please refer to transfer_queue/client.py for details. def get_client(): diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index 9ea1fc30..1015172e 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -347,12 +347,13 @@ async def get_data(self, metadata: BatchMeta) -> TensorDict: raise NotImplementedError("Subclasses must implement get_data") @abstractmethod - async def clear_data(self, metadata: BatchMeta) -> None: + async def clear_data(self, metadata: BatchMeta, fields: list[str] | None = None) -> None: """ Clear data from the storage backend. Args: metadata: BatchMeta of the data to be cleared from the storage. + fields: If given, delete only these fields and keep the samples. """ raise NotImplementedError("Subclasses must implement clear_data") @@ -839,9 +840,13 @@ async def get_data(self, metadata: BatchMeta) -> TensorDict: ) return self._merge_tensors_to_tensordict(metadata, values) - async def clear_data(self, metadata: BatchMeta) -> None: + async def clear_data(self, metadata: BatchMeta, fields: list[str] | None = None) -> None: """Remove stored data associated with the given metadata.""" + if fields is not None: + metadata = metadata.select_fields(fields) + if not metadata.field_names: + return if not metadata.field_names: raise RuntimeError( "Fail to clear_data for key-value based backends due to lack of `field_names` in BatchMeta" diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py index dd482c4e..b5bf88b2 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -885,13 +885,14 @@ async def _get_from_single_storage_unit( f"Error getting data from storage unit {target_storage_unit}: {type(e).__name__}: {e}" ) from e - async def clear_data(self, metadata: BatchMeta) -> None: + async def clear_data(self, metadata: BatchMeta, fields: list[str] | None = None) -> None: """Clear data in remote StorageUnit. Routes to each SU using global_idx % num_su (hash routing). Args: metadata: BatchMeta that contains metadata for data clearing. + fields: If given, delete only these fields and keep the samples. """ logger.debug(f"[{self.storage_manager_id}]: receive clear_data request, clearing {metadata.size} samples.") @@ -902,7 +903,7 @@ async def clear_data(self, metadata: BatchMeta) -> None: routing = self._group_by_hash(metadata.global_indexes) tasks = [ - self._clear_single_storage_unit(group.global_indexes, target_storage_unit=su_id) + self._clear_single_storage_unit(group.global_indexes, fields=fields, target_storage_unit=su_id) for su_id, group in routing.items() ] @@ -913,13 +914,13 @@ async def clear_data(self, metadata: BatchMeta) -> None: logger.error(f"[{self.storage_manager_id}]: Error in clear operation task {i}: {result}") @with_storage_unit_socket - async def _clear_single_storage_unit(self, global_indexes, target_storage_unit=None, socket=None): + async def _clear_single_storage_unit(self, global_indexes, fields=None, target_storage_unit=None, socket=None): try: request_msg = ZMQMessage.create( request_type=ZMQRequestType.CLEAR_DATA, sender_id=self.storage_manager_id, receiver_id=target_storage_unit, - body={"global_indexes": global_indexes}, + body={"global_indexes": global_indexes, "fields": fields}, ) await socket.send_multipart(request_msg.serialize()) diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py index 2610792e..ca8ad853 100644 --- a/transfer_queue/storage/simple_storage.py +++ b/transfer_queue/storage/simple_storage.py @@ -268,16 +268,18 @@ def apply_update( self.put_data(computed, global_indexes) return {field: _describe_stored_values(values) for field, values in computed.items()} - def clear(self, keys: list[int]) -> None: + def clear(self, keys: list[int], fields: list[str] | None = None) -> None: """Remove data at given global index keys, immediately freeing memory. Args: keys: Global indexes to remove. + fields: If given, remove only these fields and keep the samples. """ - for f in self.field_data: + for f in self.field_data if fields is None else set(fields) & self.field_data.keys(): for key in keys: self.field_data[f].pop(key, None) - self._active_keys -= set(keys) + if fields is None: + self._active_keys -= set(keys) def save_checkpoint(self, path: str | Path, storage_unit_id: str) -> None: """Write in-memory storage state to a checkpoint.""" @@ -654,15 +656,17 @@ def get_data(self, fields: list[str], global_indexes: list) -> dict[str, list]: values[position] = self._decode_sample(raw, entry) return result - def clear(self, keys: list) -> None: + def clear(self, keys: list, fields: list[str] | None = None) -> None: """Remove values and unlink any files they reference.""" ssd_values: list[_SSDValueRef] = [] - for values in self.field_data.values(): + for field, values in self.field_data.items(): + if fields is not None and field not in fields: + continue for key in set(keys): value = values.get(key) if isinstance(value, _SSDValueRef): ssd_values.append(value) - super().clear(keys) + super().clear(keys, fields) self._ssd_active_values -= len(ssd_values) self._ssd_active_bytes -= sum(value.size_bytes for value in ssd_values) for value in ssd_values: @@ -1321,7 +1325,7 @@ def _handle_clear(self, data_parts: ZMQMessage) -> ZMQMessage: Handle clear request, clear data in storage unit according to given global_indexes. Args: - data_parts: ZMQMessage from client, including target global_indexes. + data_parts: ZMQMessage from client, including target global_indexes and optional fields. Returns: Clear data success response ZMQMessage. @@ -1332,7 +1336,7 @@ def _handle_clear(self, data_parts: ZMQMessage) -> ZMQMessage: with limit_pytorch_auto_parallel_threads( target_num_threads=TQ_NUM_THREADS, info=f"[{self.storage_unit_id}] _handle_clear" ): - self.storage_data.clear(global_indexes) + self.storage_data.clear(global_indexes, data_parts.body.get("fields")) response_msg = ZMQMessage.create( request_type=ZMQRequestType.CLEAR_DATA_RESPONSE, # type: ignore[arg-type] diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index 06ee0f8e..af30be40 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_CLEAR_FIELDS = "KV_CLEAR_FIELDS" + KV_CLEAR_FIELDS_RESPONSE = "KV_CLEAR_FIELDS_RESPONSE" # METRICS GET_METRICS = "GET_METRICS" diff --git a/tutorial/02_kv_interface.py b/tutorial/02_kv_interface.py index bd182de1..999a8193 100644 --- a/tutorial/02_kv_interface.py +++ b/tutorial/02_kv_interface.py @@ -153,8 +153,8 @@ def demonstrate_kv_api(): tq.kv_put(key=key_for_update_tags, partition_id=partition_id, fields=None, tag=tag_update) print(f" ✓ Update success: Samples '0_0' now has tag as {tag_update}.") - # Step 5: Concatenate a field with kv_update, then empty it - print("\n[Step 5] Updating a field with kv_update (concat, then empty)...") + # Step 5: Concatenate a field with kv_update, then clear only that field + print("\n[Step 5] Updating a field with kv_update (concat, then clear the field)...") tq.kv_put(key=key, partition_id=partition_id, fields={"scratch": torch.tensor([1, 2])}) tq.kv_update( key=key, @@ -163,8 +163,8 @@ def demonstrate_kv_api(): merge_fn=lambda old, new: torch.cat([old, new]), ) print(" ✓ kv_update concat: scratch of '0_0' is now [1, 2, 3, 4].") - tq.kv_empty(keys=key, partition_id=partition_id, fields="scratch") - print(" ✓ tq.kv_empty: scratch of '0_0' is stored as None (key remains).") + tq.kv_clear(keys=key, partition_id=partition_id, fields="scratch") + print(" ✓ tq.kv_clear(fields=...): scratch of '0_0' is removed (key and other fields remain).") # Step 6: List all keys and tags in a partition print("\n[Step 6] Listing all keys and tags in partition...") @@ -218,7 +218,7 @@ def main(): 3. (async_)kv_update - Merge new values into produced fields with merge_fn(old, new) 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 from storage + 6. (async_)kv_clear - Remove key-value pairs, or only selected fields, from storage Key Features: ✓ Redis-style Semantics - Familiar KV interface (Put/Get/List) for zero learning curve