airflow.providers.common.ai.batch.results¶
Streaming result merge, JSONL landing, and XCom manifest assembly for @task.llm_batch.
Memory-bounded streaming is a hard requirement, not an optimization: a batch
can have up to 100,000 results, each up to a few KB, so materializing the
full result set (list(), sorted(), or otherwise) risks running the
worker out of memory. The only structure this module keeps resident for the
whole batch is seen, a set[int] of indexes already written (~4MB for
100k ints) – everything else is processed and written one item at a time, in
whatever order the provider streams it in (ordered: false in the
manifest is an honest declaration of this, not an apology for it).
Attributes¶
Classes¶
Anomalies encountered while merging, surfaced separately from |
Functions¶
|
Build one JSONL row from a provider-native result. |
|
Stream every result for |
|
|
|
Assemble the XCom manifest (the sole XCom payload of |
|
Return the sorted indexes in |
Module Contents¶
- airflow.providers.common.ai.batch.results.build_result_row(raw, adapter, spec)[source]¶
Build one JSONL row from a provider-native result.
Non-success outcomes (
raw.provider_status in ("errored", "expired", "cancelled")) never go through output validation – there is nothing to validate, the request itself never produced usable output ("expired"/"cancelled"are their own per-item statuses, not folded into"errored"– a provider that reports them at item granularity is telling you why that item has no output, and collapsing that into a generic provider error would make it indistinguishable from an actual API failure like rate limiting). Only a successful provider response is extracted and validated againstspec, and validation failure producesstatus: "invalid_output", never an exception: a model returning schema-non-conforming JSON is expected batch data, not a reason to abort the merge.
- class airflow.providers.common.ai.batch.results.MergeDiagnostics[source]¶
Anomalies encountered while merging, surfaced separately from
counts.Never inflates the officially reconciled
counts– visible instead of silently dropped or silently double-counted.
- airflow.providers.common.ai.batch.results.stream_results_to_jsonl(*, adapter, batch_id, output_spec, request_count, custom_id_prefix, destination)[source]¶
Stream every result for
batch_idtodestinationas JSONL, one row per input index.Every index in
range(request_count)gets exactly one row: results that never appeared in the provider stream (dropped by the provider, or genuinely never processed) are filled in asstatus: "missing"once the stream is exhausted. Validation happens inline, per item, as results arrive – not after collecting them all – so the memory bound holds for the invalid-output case too.Two defensive checks keep a single anomalous item from corrupting the whole batch’s accounting or making the manifest impossible to ever produce:
An index the adapter yields twice is written once (the first occurrence); the repeat is dropped and counted in
MergeDiagnostics.duplicate_result_count, never double-counted intocounts.An index outside
range(request_count)(a corrupt/foreigncustom_id) is dropped entirely – it cannot be rejoined to any input – and counted inMergeDiagnostics.out_of_range_result_count.
Without this, either anomaly would inflate the total past
request_countand makeassemble_manifest()’s reconciliation check raise on every attempt to finalize this batch, including every retry – the dirty data never goes away on its own.- Returns:
a tuple of (per-status row counts for
"success"/"error"/"invalid_output"/"expired"/"cancelled"/"missing", diagnostics for anomalies excluded from those counts).- Return type:
tuple[dict[str, int], MergeDiagnostics]
- airflow.providers.common.ai.batch.results.assemble_manifest(*, batch_id, adapter_name, llm_conn_id, model_id, output_spec, result_uri, request_count, merge_counts, extra_counts=None, merge_diagnostics=None, custom_id_prefix, submitted_at, completed_at)[source]¶
Assemble the XCom manifest (the sole XCom payload of
@task.llm_batch– results never are).merge_countsis the per-item breakdown fromstream_results_to_jsonl(), including"expired"/"cancelled"rows the provider reported per item.extra_countscarries the job-level{"expired", "cancelled"}counts from the batch’s terminal event; it is the only source of that information for requests the provider never wrote a result line for.None/absent keys default to0.An
expiredorcancelledterminal event still reaches this function (unlike a plaintimeout, which fails the task before any results are fetched) so that partial, already-billed results are not discarded.The two sources are combined without double-counting: per-item counts are taken as-is, and
extra_countsonly re-labels whatevermerge_counts["missing"]still has left after subtracting what the job-level figure already agrees was accounted for per item. Any part ofextra_countsthatmissingcannot cover is dropped rather than allowed to inflate the total pastrequest_count.terminal_reasonis"expired"if any request expired, else"cancelled"if any was cancelled, else"succeeded"when every request succeeded and"partial"otherwise.- Raises:
ValueError – the counts do not reconcile against
request_count– every request must land in exactly one bucket.