ColtenOuO opened a new pull request, #72107:
URL: https://github.com/apache/airflow/pull/72107
## Summary
`DictOfListsExpandInput.resolve()` calls `.resolve(context)` on every
`XComArg` kwarg of a mapped task's `.expand()` one at a time, and for the
common case (a non-mapped upstream) each of those turns into its own
`ti.xcom_pull()` call -- one Execution API round trip per kwarg. A mapped task
with N XComArg kwargs makes N sequential round trips today. Both call sites in
`expandinput.py` already carried a `TODO` admitting this. This PR adds a batch
XCom lookup endpoint and wires the common case through it, so N kwargs collapse
into one round trip instead of N.
## Change
**New batch endpoint** (`POST /execution/xcoms/{dag_id}/{run_id}/batch`) --
looks up multiple `(task_id, key, map_index)` triples in a single request and a
single DB query (`tuple_(...).in_(...)`), scoped to one dag run. A missing item
is reported as `found=False` instead of failing the whole batch, since a
partial miss (an upstream that hasn't pushed yet) is an expected batch outcome,
not an error. Access control reuses the existing multi-team dag-level check.
The endpoint is registered on a Cadwyn migration (`AddXComBatchEndpoint`) so
older Task SDK clients aren't affected.
**Task SDK plumbing** -- `XComOperations.get_batch()` on the generated
client, a `GetXComBatch`/`XComBatchResult` message pair on the supervisor comms
protocol, and the matching supervisor-side dispatch handler, following the
exact same shape as the existing single-item `GetXCom` path.
**Wiring in `DictOfListsExpandInput.resolve()`** -- before resolving each
kwarg individually, it now collects the eligible ones (plain `XComArg`,
non-mapped upstream, not inside a mapped task group -- the same condition that
already made the per-item path pull with `map_indexes=None`) and resolves them
with a single batched call. Everything else (mapped-upstream
`LazyXComSequence`, `MapXComArg`/`ZipXComArg`/`ConcatXComArg`) is unchanged and
still resolves per-item; batching those is left for a follow-up, since they
need additional design work (lazy count/slice semantics, composite-tree
walking) with smaller marginal payoff.
Three explicit fallbacks keep this safe rather than merely fast:
- A custom XCom backend is configured (`core.xcom_backend`) → batching is
skipped entirely, since a custom backend's `get_one`/`get_all` may bypass the
batch endpoint's semantics. Falls back to the existing per-item path.
- The API server doesn't support the batch endpoint yet (old server, newer
Task SDK; a 404) → falls back to the existing per-item path for that group.
- The batch response is neither a valid result nor the "not supported" error
→ raises a clear `TypeError` instead of silently resolving to wrong values.
## Before / After
Before: N XComArg kwargs on one `.expand()` call = N sequential Execution
API round trips (worker → supervisor → API server → DB), one per kwarg.
After: the eligible kwargs resolve in a single round trip regardless of N.
Measured locally (verification script, not included in this PR) against a
real FastAPI execution-API app and a real breeze Postgres database (in-process
ASGI transport — this excludes only the raw TCP hop a separate worker process
would add on top, so these numbers are a conservative floor, not an inflated
one):
| N kwargs | current (s) | batched (s) | speedup |
|---|---|---|---|
| 2 | 0.054 | 0.023 | 2.3x |
| 5 | 0.086 | 0.025 | 3.4x |
| 10 | 0.176 | 0.026 | 6.7x |
| 20 | 0.337 | 0.026 | 12.8x |
| 50 | 0.890 | 0.028 | 32.2x |
The current path scales linearly with N (~17ms/call); the batched path stays
flat (~25ms total) regardless of N.
## Testing
- Full task-sdk suite: 2796 passed, 7 skipped (pre-existing/environment,
unrelated), 0 failed.
- New execution-API route tests: multiple items found, partial miss
(`found=False`), empty batch, and multi-team access-control (same-team allowed,
cross-team forbidden) for the batch endpoint.
- New Task SDK client tests: successful batch call, and the 404 →
`ErrorType.XCOM_BATCH_NOT_SUPPORTED` fallback signal.
- New supervisor dispatch test (added to the existing `REQUEST_TEST_CASES`
table) covering `GetXComBatch` → `XComBatchResult`.
- New `expandinput`/`xcom_arg` tests: N kwargs collapse into exactly one
`GetXComBatch` call; correct per-kwarg value distribution; old-server fallback
(404 → per-item calls, correct values); custom-backend skip (batching never
attempted); mixed eligibility (a mapped-task-group XComArg alongside plain ones
— only the plain ones batch).
- Fixed 3 pre-existing tests (`test_map_cross_product`,
`test_map_product_same`,
`test_mapped_render_template_fields_validating_operator`) whose hand-rolled
comms mocks didn't know about the new `GetXComBatch` message -- a real
regression this PR would otherwise have introduced, caught by running the full
suite rather than only the new tests.
- mypy clean on both `task-sdk` and the touched `airflow-core` files.
- Both Cadwyn version-check prek hooks (`check-execution-api-versions`,
`check-supervisor-schemas-versions`) pass.
## Scope / follow-ups
Batching for a mapped-upstream `LazyXComSequence` (the `GetXComCount` length
lookup) and for `MapXComArg`/`ZipXComArg`/`ConcatXComArg` composite resolution
is intentionally out of scope here -- those need their own design work and have
smaller marginal payoff for typical DAG shapes (most `.expand()` calls use a
handful of plain kwargs, not many mapped-upstream ones). Tracked as a natural
next step, not filed as a separate issue since no workaround is being shipped
here.
---
##### Was generative AI tooling used to co-author this PR?
- [X] Yes — Claude Code (Sonnet 5)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
