From eb88e9377afaf9c432c24065ad81308f76fd4401 Mon Sep 17 00:00:00 2001 From: 0oshowero0 Date: Thu, 15 Oct 2026 00:57:45 +0800 Subject: [PATCH 1/2] [fix] kv_batch_get: raise when some keys are missing The controller drops missing keys from the retrieved meta, so a partial miss used to return fewer rows than keys, with nothing telling the caller which rows were dropped. kv_batch_get (and the _by_meta variants, which call it) now raises ValueError naming the missing keys. The check used to apply only when parser_args was given; it now covers every get. Signed-off-by: 0oshowero0 Co-authored-by: Cursor --- tests/e2e/test_kv_interface_e2e.py | 14 ++++++++++++- transfer_queue/interface.py | 32 +++++++++++++++++------------- 2 files changed, 31 insertions(+), 15 deletions(-) diff --git a/tests/e2e/test_kv_interface_e2e.py b/tests/e2e/test_kv_interface_e2e.py index 228f22f6..bb4ce2f2 100644 --- a/tests/e2e/test_kv_interface_e2e.py +++ b/tests/e2e/test_kv_interface_e2e.py @@ -727,6 +727,18 @@ def test_kv_batch_get_nonexistent_key(self, controller, tq_api): # Or it might raise an error about keys not found assert "not found" in str(e).lower() or "empty" in str(e).lower() + def test_kv_batch_get_raises_when_some_keys_are_missing(self, controller, tq_api): + """A partial miss must not return fewer rows than keys.""" + partition_id = "test_partition" + keys = ["partial_0", "partial_1"] + fields = TensorDict({"x": torch.tensor([[1], [2]])}, batch_size=2) + tq_api.kv_batch_put(keys=keys, partition_id=partition_id, fields=fields) + + with pytest.raises(ValueError, match=r"1 of 3 keys were not found in partition 'test_partition': \['absent'\]"): + tq_api.kv_batch_get(keys=["partial_0", "absent", "partial_1"], partition_id=partition_id) + + tq_api.kv_clear(keys=keys, partition_id=partition_id) + def test_kv_batch_get_data_parser_parses_each_field_with_per_key_args(self, tq_api, backend_name): if backend_name != "SimpleStorage": pytest.skip("kv_batch_get data_parser is implemented only for SimpleStorage") @@ -756,7 +768,7 @@ def parse(data, args): tq_api.kv_batch_get(keys=keys, partition_id=partition_id, data_parser=parse, parser_args=[2, 3]) with pytest.raises(ValueError, match="requires data_parser"): tq_api.kv_batch_get(keys=keys, partition_id=partition_id, parser_args=[2, 3, 1]) - with pytest.raises(ValueError, match="parser_args needs every key"): + with pytest.raises(ValueError, match="1 of 2 keys were not found"): tq_api.kv_batch_get( keys=["parse_0", "absent"], partition_id=partition_id, data_parser=parse, parser_args=[1, 1] ) diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 3d6be893..91f35713 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -502,7 +502,7 @@ def kv_batch_get_by_meta(meta: KVBatchMeta, select_fields: list[str] | str | Non TensorDict with the requested data Raises: - ValueError: If keys or partition are not found + ValueError: If the partition or any of the keys is not found ValueError: If empty fields exist in any key (sample) ValueError: If any field in select_fields doesn't exist in KVBatchMeta.fields @@ -557,17 +557,17 @@ def kv_batch_get( Only supported by SimpleStorage; KV backends raise ``NotImplementedError``. parser_args: Optional per-key arguments for ``data_parser``, aligned with ``keys``. - Requires ``data_parser``, and every key must exist. They are serialized like - request data, so tuples arrive as lists. + Requires ``data_parser``. They are serialized like request data, so tuples + arrive as lists. Returns: TensorDict with the requested data Raises: - ValueError: If keys or partition are not found + ValueError: If the partition or any of the keys is not found ValueError: If empty fields exist in any key (sample) - ValueError: If ``parser_args`` is given without ``data_parser``, does not match the - number of keys, or some keys are missing + ValueError: If ``parser_args`` is given without ``data_parser`` or does not match the + number of keys RuntimeError: If ``data_parser`` fails or breaks its contract on a storage unit Example: @@ -947,7 +947,7 @@ async def async_kv_batch_get_by_meta(meta: KVBatchMeta, select_fields: list[str] TensorDict with the requested data Raises: - ValueError: If keys or partition are not found + ValueError: If the partition or any of the keys is not found ValueError: If empty fields exist in any key (sample) ValueError: If any field in select_fields doesn't exist in KVBatchMeta.fields @@ -1008,10 +1008,10 @@ async def async_kv_batch_get( TensorDict with the requested data Raises: - ValueError: If keys or partition are not found + ValueError: If the partition or any of the keys is not found ValueError: If empty fields exist in any key (sample) - ValueError: If ``parser_args`` is given without ``data_parser``, does not match the - number of keys, or some keys are missing + ValueError: If ``parser_args`` is given without ``data_parser`` or does not match the + number of keys RuntimeError: If ``data_parser`` fails or breaks its contract on a storage unit Example: @@ -1032,10 +1032,14 @@ async def async_kv_batch_get( if batch_meta.size == 0: raise ValueError("keys or partition were not found!") - # Missing keys are dropped from the meta, which would shift parser_args onto other keys. - num_keys = 1 if isinstance(keys, str) else len(keys) - if parser_args is not None and batch_meta.size != num_keys: - raise ValueError(f"only {batch_meta.size} of {num_keys} keys were found; parser_args needs every key") + # The controller drops missing keys from the meta, so the rows would no longer line up with keys. + requested = [keys] if isinstance(keys, str) else keys + if batch_meta.size != len(requested): + found = set(await tq_client.async_kv_retrieve_keys(batch_meta.global_indexes, partition_id)) + missing = [key for key in requested if key not in found] + raise ValueError( + f"{len(missing)} of {len(requested)} keys were not found in partition {partition_id!r}: {missing[:20]}" + ) if select_fields is not None: if isinstance(select_fields, str): From 47d9cba91cff3c0792ec14fc18c7c00edf90580f Mon Sep 17 00:00:00 2001 From: 0oshowero0 Date: Thu, 15 Oct 2026 00:59:42 +0800 Subject: [PATCH 2/2] [fix] kv_batch_put / kv_batch_update: reject duplicate keys before writing KVBatchMeta rejects duplicate keys, but only when it is built for the return value, after the data was already written. A kv_batch_put with duplicate keys stored the data and then raised; a kv_batch_update applied the merge and then raised. Both now raise ValueError naming the duplicates before any request is sent. Signed-off-by: 0oshowero0 Co-authored-by: Cursor --- tests/e2e/test_kv_interface_e2e.py | 28 ++++++++++++++++++++++++++++ transfer_queue/interface.py | 12 ++++++++++++ 2 files changed, 40 insertions(+) diff --git a/tests/e2e/test_kv_interface_e2e.py b/tests/e2e/test_kv_interface_e2e.py index bb4ce2f2..10ba4b3f 100644 --- a/tests/e2e/test_kv_interface_e2e.py +++ b/tests/e2e/test_kv_interface_e2e.py @@ -602,6 +602,15 @@ def test_kv_batch_put_returns_cumulative_fields(self, controller, tq_api): tq_api.kv_clear(keys=keys, partition_id=partition_id) + def test_kv_batch_put_rejects_duplicate_keys_before_writing(self, controller, tq_api): + partition_id = "test_partition" + fields = TensorDict({"x": torch.tensor([[1], [2]])}, batch_size=2) + + with pytest.raises(ValueError, match=r"keys must be unique, got duplicates: \['dup_put'\]"): + tq_api.kv_batch_put(keys=["dup_put", "dup_put"], partition_id=partition_id, fields=fields) + + assert "dup_put" not in tq_api.kv_list(partition_id=partition_id).get(partition_id, {}) + class TestKVGetE2E: """End-to-end tests for kv_batch_get functionality.""" @@ -1189,6 +1198,25 @@ def test_field_expansion_across_samples(self, controller, tq_api): class TestKVUpdateE2E: """SimpleStorage merge behavior through sync and async public APIs.""" + def test_batch_update_rejects_duplicate_keys_before_merging(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" + key = "dup_update" + tq_api.kv_batch_put(keys=[key], partition_id=partition_id, fields=TensorDict({"x": torch.tensor([[1]])}, 1)) + + with pytest.raises(ValueError, match=r"keys must be unique, got duplicates: \['dup_update'\]"): + tq_api.kv_batch_update( + keys=[key, key], + partition_id=partition_id, + fields=TensorDict({"x": torch.tensor([[10], [20]])}, batch_size=2), + merge_fn=lambda old, new: old + new, + ) + + assert_tensor_equal(tq_api.kv_batch_get(keys=key, partition_id=partition_id)["x"], torch.tensor([[1]])) + tq_api.kv_clear(keys=key, partition_id=partition_id) + 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") diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 91f35713..f307666d 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -33,6 +33,7 @@ import shutil import subprocess import time +from collections import Counter from importlib import resources from pathlib import Path from typing import Any, Callable @@ -456,6 +457,7 @@ def kv_batch_put( ValueError: When both ``fields`` and ``tags`` are empty. ValueError: When ``fields`` batch size mismatches key count. ValueError: When ``tags`` length mismatches key count. + ValueError: When ``keys`` contains duplicates. RuntimeError: When retrieved metadata size mismatches input key count. Example: @@ -832,6 +834,13 @@ async def async_kv_put( ) +def _check_unique_keys(keys: list[str]) -> None: + # KVBatchMeta rejects duplicates only after the write has landed, so reject them before any RPC. + duplicates = sorted(key for key, count in Counter(keys).items() if count > 1) + if duplicates: + raise ValueError(f"keys must be unique, got duplicates: {duplicates[:20]}") + + async def async_kv_batch_put( keys: list[str], partition_id: str, @@ -869,6 +878,7 @@ async def async_kv_batch_put( Raises: ValueError: If neither `fields` nor `tags` is provided ValueError: If length of `keys` doesn't match length of `tags` or the batch_size of `fields` TensorDict + ValueError: If `keys` contains duplicates RuntimeError: If retrieved BatchMeta size doesn't match length of `keys` Example: @@ -892,6 +902,7 @@ async def async_kv_batch_put( f"`keys` with length {len(keys)} does not match the `fields` TensorDict with " f"batch_size {fields.batch_size[0]}" ) + _check_unique_keys(keys) tq_client = _maybe_create_tq_client() @@ -1182,6 +1193,7 @@ async def async_kv_batch_update( 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") + _check_unique_keys(keys) tq_client = _maybe_create_tq_client() batch_meta = await tq_client.async_kv_retrieve_meta(keys=keys, partition_id=partition_id, create=False)