[fix] Raise on missing keys in kv_batch_get and reject duplicate keys before batch writes

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."""
@@ -727,6 +736,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 +777,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]
             )
@@ -1177,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")

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:
@@ -502,7 +504,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 +559,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:
@@ -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()
 
@@ -947,7 +958,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 +1019,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 +1043,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):
@@ -1178,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)