From d4a9f49aa5c9d45341673fbbab22524d168ab93b Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Thu, 5 Nov 2026 17:38:01 +0800 Subject: [PATCH] [fix] Omit missing fields in kv_batch_get instead of raising #190 made `kv_batch_get` and `kv_batch_get_by_meta` raise `ValueError` when a `select_fields` entry is not ready in every requested key. Before, a field the existing keys do not hold was silently left out of the returned TensorDict. Restore that behavior: drop the check, so missing and unfinalized fields are omitted again. Missing keys still raise, and so does a `kv_batch_get_by_meta` `select_fields` entry absent from the given meta. Signed-off-by: OutstanderWang --- tests/e2e/test_kv_interface_e2e.py | 10 ++++++---- transfer_queue/interface.py | 15 ++++----------- 2 files changed, 10 insertions(+), 15 deletions(-) diff --git a/tests/e2e/test_kv_interface_e2e.py b/tests/e2e/test_kv_interface_e2e.py index 03a03298..c3ec0c13 100644 --- a/tests/e2e/test_kv_interface_e2e.py +++ b/tests/e2e/test_kv_interface_e2e.py @@ -1390,10 +1390,12 @@ def tokens_ready(): meta = tq_api.kv_update(key=key, partition_id=partition_id, fields=fields, merge_fn=concat, finalize=False) assert not tokens_ready() assert set(tq_api.kv_batch_get(keys=key, partition_id=partition_id).keys()) == {"prompt"} - with pytest.raises(ValueError, match="not finalized"): - tq_api.kv_batch_get(keys=key, partition_id=partition_id, select_fields="tokens") - with pytest.raises(ValueError, match="not finalized"): - tq_api.kv_batch_get_by_meta(meta) + assert "tokens" not in tq_api.kv_batch_get(keys=key, partition_id=partition_id, select_fields="tokens").keys() + partial = tq_api.kv_batch_get(keys=key, partition_id=partition_id, select_fields=["prompt", "tokens"]) + assert set(partial.keys()) == {"prompt"} + absent = tq_api.kv_batch_get(keys=key, partition_id=partition_id, select_fields=["prompt", "never_written"]) + assert set(absent.keys()) == {"prompt"} + assert "tokens" not in tq_api.kv_batch_get_by_meta(meta).keys() row = tq_api.kv_describe(keys=key, partition_id=partition_id)[key] assert row == {"fields": ["prompt", "tokens"], "unfinalized": ["tokens"], "links": {}} with pytest.raises(RuntimeError, match="not found"): diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 5aeadc42..a9615c9b 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -544,7 +544,8 @@ def kv_batch_get( Args: keys: Single key or list of keys to retrieve partition_id: Partition containing the keys - select_fields: Optional field(s) to retrieve. If None, retrieves all fields + select_fields: Optional field(s) to retrieve. If None, retrieves all fields. Fields + not ready in every key (missing, or not finalized) are omitted from the result. data_parser: Advanced. Optional ``data_parser(data, args) -> dict`` that each storage unit runs on its share of the batch before replying, so only the parsed values cross the network (the get-side counterpart of ``kv_batch_put``'s ``data_parser``). @@ -576,7 +577,6 @@ def kv_batch_get( Raises: ValueError: If the partition or any of the keys is not found ValueError: If empty fields exist in any key (sample) - ValueError: If a field in ``select_fields`` is not ready in every found key 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 @@ -1136,7 +1136,8 @@ async def async_kv_batch_get( Args: keys: Single key or list of keys to retrieve partition_id: Partition containing the keys - select_fields: Optional field(s) to retrieve. If None, retrieves all fields + select_fields: Optional field(s) to retrieve. If None, retrieves all fields. Fields + not ready in every key (missing, or not finalized) are omitted from the result. data_parser: Advanced. Optional parser the storage units run before replying; see ``kv_batch_get`` for its contract and constraints. parser_args: Optional per-key arguments for ``data_parser``, aligned with ``keys``. @@ -1147,7 +1148,6 @@ async def async_kv_batch_get( Raises: ValueError: If the partition or any of the keys is not found ValueError: If empty fields exist in any key (sample) - ValueError: If a field in ``select_fields`` is not ready in every found key 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 @@ -1184,13 +1184,6 @@ async def async_kv_batch_get( fields_to_fetch = [select_fields] else: fields_to_fetch = select_fields - # The meta keeps only fields ready in every key; select_fields would drop the rest silently. - not_ready = [f for f in fields_to_fetch if f not in batch_meta.field_names] - if not_ready: - raise ValueError( - f"Fields {not_ready} are not ready in all the requested keys: missing, or written by " - "kv_update(finalize=False) and not finalized yet" - ) batch_meta = batch_meta.select_fields(fields_to_fetch) if not batch_meta.is_ready: