Skip to content

[fix] Omit missing fields in kv_batch_get instead of raising #209

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
0oshowero0 merged 1 commit into from
Nov 5, 2026
Merged

[fix] Omit missing fields in kv_batch_get instead of raising #209

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 6 additions & 4 deletions tests/e2e/test_kv_interface_e2e.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"):
Expand Down
15 changes: 4 additions & 11 deletions transfer_queue/interface.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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``).
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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.

@chatgpt-codex-connector chatgpt-codex-connector Bot Nov 5, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Handle batches with no common ready field

When the requested keys have no common ready field—for example, two existing keys with disjoint field sets, or a new key whose only kv_update(..., finalize=False) field is hidden—kv_retrieve_meta returns a nonempty BatchMeta with no fields and is_ready == False. The deleted not_ready check is therefore not the only error path: the unchanged readiness check at lines 1189-1190 still raises ValueError, rather than omitting the unavailable selection as promised here. The new test masks this because its unrelated prompt field keeps the retrieved metadata ready, so behavior now depends on whether some unselected common field happens to exist.

Useful? React with 👍 / 👎.

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``.
Expand All @@ -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
Expand Down Expand Up @@ -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:
Expand Down
Loading