Skip to content

Fetch several XCom values of a Taskinstance in one Execution API request - #70223

Draft
dabla wants to merge 53 commits into
apache:mainfrom
dabla:feature/add-xcoms-keys-route
Draft

dabla wants to merge 53 commits into
apache:mainfrom
dabla:feature/add-xcoms-keys-route

Conversation

@dabla

@dabla dabla commented Jul 22, 2026 •

Copy link
Copy Markdown
Contributor

Review only this change: the PR's own commit · diff against #62922's branch. Everything else in this PR's commit list and file list is #62922, which this PR is stacked on.

Stacked on #62922 (Task Iteration), which owns XComIterable: only the last commit, 76625b4, belongs to this PR. Squashed from this branch's earlier history and re-based on that PR, with the batched reads moved to XComIterable's current home in bases/xcom.py.

Add a POST /xcoms/{dag_id}/{run_id}/{task_id}/keys execution API endpoint that
accepts a list of XCom keys and returns all matching values in a single database
query. Use it in XComIterable to reduce iteration from N round-trips to one.

Motivation

XComIterable (introduced in AIP-104 Task Iteration) stores per-index results
under distinct keys (return_value_0, return_value_1, …) with the same
map_index. The existing GET …/slice endpoint cannot be reused because it
ranges over map_index for a single key — the inverse structure. Iterating or
slicing a 1000-item result previously issued 1000 separate XCom.get_one calls.

Changes

  • New endpoint POST /xcoms/{dag_id}/{run_id}/{task_id}/keys — accepts
    {"keys": [...]}, returns values in request-key order (null for missing
    keys), filtered by map_index query param (default -1)
  • has_xcom_access simplified: reads the optional key from
    request.path_params instead of requiring a {key} path segment, so a
    separate router/dependency is no longer needed
  • XComKeysRequest Pydantic model added to execution API data models
  • GetXComByKeys comms message, XComOperations.get_by_keys client method,
    handle_get_xcom_by_keys supervisor handler, and processor dispatch wiring
  • BaseXCom.get_by_keys sends the message and deserializes the values;
    XComIterable.__iter__ and __getitem__ slice read through it, one
    round-trip for a full iteration or slice. A single index stays one XCom read.
  • AddXComKeysEndpoint execution API version change and AddGetXComByKeys
    supervisor schema version entry, with the regenerated schema snapshot.
  • Tests for the new endpoint (TestGetXComByKeys) and for XComIterable
    (in test_xcom.py), plus supervisor message coverage.

Was generative AI tooling used to co-author this PR?
  • Yes (please specify the tool below)

Generated-by: Claude Code following the guidelines

🤖 Generated with Claude Code

@amoghrajesh amoghrajesh left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

I haven't followed up closely with AIP 104, but I have some qns from an initial look.

Comment thread task-sdk/src/airflow/sdk/execution_time/schema/schema.json
Comment thread task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py
Comment thread airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py
Comment thread airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py
@dabla dabla added this to the Airflow 3.4.0 milestone Jul 22, 2026
@dabla
dabla force-pushed the feature/add-xcoms-keys-route branch from 5258e09 to dbf25d5 Compare July 22, 2026 14:10
@dabla
dabla marked this pull request as draft August 17, 2026 08:52
@dabla

dabla commented Aug 17, 2026

Copy link
Copy Markdown
Contributor Author

This PR should only be merged once AIP-104 is merged.

@dabla

dabla commented Sep 27, 2026

Copy link
Copy Markdown
Contributor Author

Interaction with #73807 to handle when both land: FlattenedXComIterable there inherits XComIterable.__iter__, which this PR turns into a single batched fetch of the raw pages. Whichever of the two merges second must give FlattenedXComIterable its own __iter__ over flattened positions; see the matching note on #73807.

@dabla
dabla force-pushed the feature/add-xcoms-keys-route branch 12 times, most recently from 04cc40e to b74f42c Compare October 1, 2026 10:19
dabla added 27 commits October 14, 2026 22:43
…ts type

_execute_task asked twice whether the task instance was an
IndexedTaskInstance: to skip exporting AIRFLOW_CTX_* to os.environ and to
leave the execution timeout to the iterated task. Each class now answers
for itself: RuntimeTaskInstance.export_context_to_env() and
enforces_execution_timeout(), which IndexedTaskInstance overrides, so the
runner no longer needs to know the subclass. Boy scout rule on code this
change added.

Measured with a CodSpeed simulation (valgrind), base vs head, median of
three rounds, instructions:

  control (checkpoint read)            57,690 ->     57,813  +0.2%
  plain task through _execute_task    310,483 ->    313,530  +1.0% (rounds spread 3%)
  ten sync items                   74,464,963 -> 74,650,754  +0.25%

In isolation the method call costs the same as a failing isinstance on a
plain task instance (76 vs 80 ns, pydantic models go through
ABCMeta.__instancecheck__) and about 60 ns more per call on an indexed one.
…its neighbours

SubOperatorRegister was the last Sub* name next to IndexedTaskInstance,
IndexedTaskRunner, IndexedTaskState and IndexedTaskOutcomes; it becomes
IndexedOperatorRegister. The docstrings of IndexedTaskRunner say indexed task
instead of sub-task as well.
…or an item

The classes say IndexedTaskInstance, IndexedTaskRunner, IndexedTaskState and
IndexedOperatorRegister, while the functions, messages, docstrings, comments,
docs page and test names around them called the same thing a sub-task, a
sub-operator or an item. An item now only means the input value an index runs
with; iteration stays, as the feature's own word.

- keeps_failure_to_its_item, _run_item, _run_task, _report_item and
  _create_task become keeps_failure_to_its_indexed_task, _run_index,
  _run_indexed_task, _report_indexed_task and _create_indexed_task
- the errors name the indexed task ("Indexed task x[i] ...", "N of M
  indexed tasks ran")
- the docs page defines the iterated task and the indexed task
…pstream provides

The partial kwargs holding an XComArg are the same for every index, so they
are found once per iterated task and only they are walked for the
fingerprint; a task without one walks nothing per index. A success or skip
report without a callback only sets the state and end_date, which makes no
SDK call, so it runs on the loop instead of in a worker thread.

CodSpeed simulation over 20 indexed tasks: -4.1% instructions for sync
operators, -3.6% for async ones.
…hen they differ

Unmapping an index that iterates params resolved its values into the Param
objects a shallow copy shares with the mapped operator: iterate() unmaps every
index in one process, so one index's params became the next one's. They are
copied deeply when the index brings its own.

An index's own params also never reached execute: the run context was cloned
from the parent's again. They are processed once, when the built operator's
params differ from the iterated task's, and both the render and the run
context read them; any other index keeps the parent's processed params, which
the clone copies, instead of processing them again.

CodSpeed simulation over 20 indexed tasks: -2.2% instructions for sync
operators, -2.0% for async ones.
Only ImportError and TypeError were taken as a refusal: a pydantic result whose
model gained a required field since the checkpoint was written raises a
ValidationError, a ValueError, and the indexed task that had succeeded failed
every retry and every clear until the checkpoint expired. The checkpoint is
already read, so nothing transient can fail while restoring it, and any
exception is now taken as a refusal.
A partial kwarg from an iterated upstream renders to an XComIterable, whose
serialized form keeps a handle and no item: clearing that upstream with this
task and getting other items of the same length left the fingerprint as it
was, and the indexed tasks that had succeeded replayed results computed from
the old input. A mapped upstream's LazyXComSequence was fetched whole by its
serializer for every index.

Both are now digested by their items, read one at a time, once per run: the
first indexed task to ask takes the digest and the others reuse it. CodSpeed
simulation over 20 indexed tasks with a partial kwarg from an upstream of 100
items: -34.6% instructions and 2000 -> 100 item reads for a mapped upstream,
+3.4% and 0 -> 100 reads for an iterated one; flat without an upstream.
Only an XComIterable or a LazyXComSequence in a partial kwarg was digested
by its items. The sequences .map(), .zip() and .concat() wrap around an
upstream took the plain serde path: a _MapResult could not be serialized,
which left the whole index without a fingerprint, so its checkpoint was
replayed whatever had changed, the iterated input included; a _ZipResult or
_ConcatResult around an iterated upstream was digested by its handle, which
stays the same when the upstream is cleared and returns other items.

Every lazy sequence, anything but a str, bytes, list or tuple, is now
digested by its items. A value whose items cannot all be read gets no digest
rather than failing every indexed task: a .map() callable may raise for an
item the indexed tasks never read.
The digest of a lazy upstream in a partial kwarg reads every item, one
supervisor round trip each, before the first indexed task starts. A kill
landing meanwhile stopped the indexed tasks waiting for the digest, but not
the read itself, which went on to the last item and held the process until
the supervisor escalated.

The digest now asks the iteration's stop flag before every read, as the
executor does before every pull, and leaves the value without a digest once
it is set: no indexed task runs after the kill to use it. CodSpeed
simulation of the hot-path benchmarks: within 0.4% of the previous head.
serde writes a set's members in iteration order, which for strings follows the
hash seed, and sort_keys only orders dict keys: the same input fingerprinted
differently on another worker, and a retry there ran every succeeded index
again. The members of every set are now put in a fixed order before the
digest, in the input fingerprint and in the digest of a lazy upstream's items.
A value without a set is serialized exactly as before, so the fingerprints of
earlier checkpoints still match.
…ndex

Every indexed task serialized the partial kwargs an upstream provides again for
its fingerprint, although they are the same for every index: a lookup list or a
config from an upstream behind .partial() cost its whole serialization once per
index. They are now made canonical once per run, by the first indexed task that
asks, and each index only adds its own iterated input to the digest. Without a
partial kwarg from an upstream the fingerprint is the same as before.

CodSpeed simulation, 20 indexed tasks with a partial kwarg from a plain upstream
list of 1,000 rows: 1,258M -> 188M instructions (-85%), peak memory 2.1-2.6 MiB
-> 0.95 MiB; the cases without such a kwarg are unchanged within noise.
…d loop is not kept

An asyncio.Lock that has had a waiter keeps a reference to its loop, so as the
value of the decoder's per-loop WeakKeyDictionary it kept its own key alive:
a sync indexed task that runs asyncio.run() per index with overlapping calls
kept one closed loop per index until the task ended. A lock no call waited for
holds no such reference and is still reused across calls; one that a call
waited for is dropped when the loop's last call leaves.

The per-loop bookkeeping moves out of asend() into _LoopQueues, which owns the
weak-keyed map, and _LoopQueue, an async context manager per loop that counts
its calls, notes contention and leaves the map when it must.

CodSpeed simulation of the comms benchmark: one asend at a time +1.0%, 50
concurrent +0.4%, send unchanged; the wall-clock wait of asyncio.to_thread work
next to 50 calls in flight is unchanged (0.24 -> 0.26 ms median).
…gest

One item of a lazy upstream that could not be read or serialized, such as a
.map() callable raising for an element no indexed task reads, left the whole
value without a digest on every attempt, so the fingerprints still matched
after the upstream was cleared and produced other values. The items are now
read by position and an unreadable one is digested as a fixed marker naming its
exception's type, so the other items still count; a warning names the path, as
the indexed tasks then run again, or not, for a reason found in that upstream.
…ays its own

_get_unmap_kwargs deep-copied params only when an index iterated them. Otherwise
copy.copy kept the same Param objects, which BaseOperator.__init__ wraps again
without copying and ParamsDict.__setitem__ resolves a new value into, so an
operator setting a param in its __init__ changed the iterated task's params too
and every index read the parent's value, and a param one index's execute set
reached its siblings. The params are now always copied deeply.

With its own copy, an index whose operator set a param differs from the
iterated task again, so the shortcut that processes only such an index's params
stays: processing every index's params measured +1.7% (sync) and +2.9% (async)
instructions in CodSpeed simulation over 20 indexed tasks.
…s own

IterableOperator had grown to 33 methods doing six jobs; eight of them, and two
caches with their locks in IterationState, only served the fingerprint that
tells whether an indexed task's checkpoint still applies. InputFingerprint now
owns that question for one run: the partial kwargs that come from an upstream,
the walk of declared against rendered values, the digests of lazy upstream
values, the canonical JSON and both per-run caches. IterationState hands out the
run's instance, and _create_indexed_task asks it for the fingerprint.

Every test case is kept, with its call target moved; the one branch the old
code left untested (an upstream input serde cannot serialize) now has a test.
CodSpeed simulation shows every hot-path case within noise.
…vel helpers

InputFingerprint held two caches with their locks, the digests of lazy upstream
values and the canonical text of the partial kwargs from an upstream, and a
static digest() that took that text as an argument although it was the
instance's own state. _LazyDigests and _UpstreamInputs now own one cache each,
the canonical JSON is a module function the three share, and the digest is part
of InputFingerprint.of().

Every test case is kept with its call target moved; the module stays fully
covered, and CodSpeed simulation shows every hot-path case within noise.
…erator stays a Dag node

The review rounds of Task Iteration each fixed a real problem (kill paths,
checkpoint ordering, thread affinity, params, fingerprints), and each fix landed
where the code already was: as another method, parameter or flag on
IterableOperator. By the approval round the class was 1,033 lines with 33
methods. Every method was sound on its own; the class no longer had one
responsibility. The input fingerprint moved out first (two commits before this
one). This commit finishes the split, with no change in behaviour.

What was wrong, and the principle each fix applies:

- One run's executor, context, template environment, checkpoint flags and
  outcomes were passed from method to method through six methods
  (_run_index took seven parameters). Values that travel together are an
  object that is missing (Fowler, Introduce Parameter Object): IterationRun,
  one per execute(), holds them.
- One indexed task's task and runner travelled the same way through its
  steps (replay, publish, record skip, record failure, report). They are an
  object as well (Fowler, Extract Class): IndexedTaskLifecycle, one per index.
- Building an indexed task (unmap, params, render, fingerprint) used fields
  nothing else in the run touched: IndexedTaskUnmapper, as unmapping and
  rendering go together in MappedOperator.render_template_fields.
- IterationState held two jobs behind two locks: what a kill must reach,
  which exists before the run, and the run's own data. It keeps the kill
  bookkeeping and kills the indexed operators itself; the run's data moves to
  IterationRun (single responsibility, Martin's SRP).
- Static and class methods that used no state of their class were moved to
  where their data lives (Fowler, Move Function): _report_indexed_task to
  IndexedTaskLifecycle, whose executor it took; _retryable and _fail_fast_for,
  the runner's rules for an exception, to module functions;
  _refuse_operators_that_skip_downstream inlined next to the guard of the same
  kind in __init__; axcom_push, a one-line wrapper, to the lifecycle.
- Checkpoints also answers whether an index's checkpoint may be trusted, so
  trust_checkpoints and since stop travelling with each indexed task.

The new classes are collaborators IterableOperator holds or makes, never bases
it inherits from (composition over inheritance): it subclasses BaseOperator
because the Dag needs one, and delegates the run. It comes last in the module,
after everything it uses.

IterableOperator: 1,033 -> 399 lines (about 160 of them the class docstring),
33 -> 12 methods. No class in the module has more than 12 methods. ruff's
opt-in design rules (PLR0904, PLR0913, PLR0917, PLR6301), which flagged
_run_index, _run_indexed_task, _publish and axcom_push before, pass.

Tests: every test function of test_iterableoperator.py is kept under its name
(one moved from TestIterationState to TestIterationRun with the length it
tests); only call targets changed, through a run_of() helper that builds a run
as execute() does. Two tests are new: IterationState.kill reaching every
operator whatever one raises, and the unmapper's do_xcom_push. Line and branch
coverage of the module are unchanged (17 and 4 missed, base 18 and 4).

CodSpeed simulation, base vs head, two rounds each, plus peak memory:
| Benchmark      | base 9726e07   | head              | mean   |
|----------------|-------------------|-------------------|--------|
| sync items     | 156.42M / 154.50M | 152.44M / 153.99M | -1.4%  |
| async items    | 152.87M / 152.56M | 152.38M / 152.14M | -0.3%  |
| iterated upstr | 106.81M / 107.18M | 106.94M / 106.61M | -0.2%  |
| mapped upstr   | 107.53M / 105.84M | 106.51M / 105.40M | -0.7%  |
| plain upstr    | 189.08M / 187.85M | 188.34M / 186.83M | -0.5%  |
| control        |   0.517M / 0.519M |   0.518M / 0.519M |  0     |
| peak memory    | 1.02 / 0.86 MiB   | 0.95 / 0.86 MiB   |        |

Every case is within its round-to-round spread: the split costs nothing
measurable, in time or memory.
…ide the loop

Nothing tested .iterate() together with loops at parse time, nor what an
iterated task's input holds when it reads a loop task:

- inside the loop, over a task of the same loop, it validates and its input is
  the current pass's value;
- outside the loop (or in another loop), it validates too, waits for the loop's
  gate, and its input is a lazy sequence of every pass's value, as for any task
  that reads a loop task from outside it. .expand() over the same task is
  refused at parse time; these tests pin today's .iterate() behaviour until
  that is decided;
- each pass is a task instance of its own and writes its indexed tasks'
  checkpoints under it, so a later pass never replays an earlier one.
…) is

The loop check only looked at mapped tasks. An iterated task is not mapped,
so .iterate() over a task of a loop passed from outside that loop, where
.expand() over the same task is refused. It then read every pass's value
and ran one index per pass instead of one per item, without any warning.

Every operator now says which upstream tasks it runs once per item of:
iter_expansion_sources(), empty by default. A mapped task returns its mapped
dependencies; an iterated task the upstream tasks of its input, which are not
reported as mapped dependencies, as the iterated task is not mapped and the
scheduler must not count those upstream items.

The check starts from each loop's gate: a task outside a loop that reads
one of its tasks depends on the gate instead (get_edge_source), so only the
gate's downstream tasks can fan out over a task of the loop. It no longer
walks every mapped task of the Dag. CodSpeed simulation, the check and
DAG.validate() on three Dag shapes, two rounds each:

| Benchmark                  | before              | after             |
|----------------------------|---------------------|-------------------|
| check, 1000 plain, no loop |   4.15M /   4.17M   | 2.97M / 2.97M -29% |
| check, 1000 plain + a loop |   4.30M /   4.32M   | 3.13M / 3.12M -28% |
| check, 999 mixed + a loop  |  12.43M /  12.37M   | 3.10M / 3.12M -75% |
| validate(), 999 mixed      |  13.22M /  13.14M   | 3.87M / 3.86M -71% |

The .iterate() cases from outside a loop and from another loop move into
the rejecting test; the two tests that pinned the old behaviour go.
…oint

XComIterable, the result of an iterated task, stores one XCom per iteration
under distinct keys (return_value_0, return_value_1, ...) of the same task
instance. The existing slice endpoint for a mapped task's XComs ranges over
map_index for a single key, the inverse shape, so iterating or slicing an
XComIterable cost one GET per value.

The new endpoint takes a list of keys and returns their values in that order,
None for a key without an XCom, filtered by map_index (-1 by default), in one
database query. It resolves the producer as the other XCom reads do, so inside
a loop the keys come from the caller's iteration, or from the previous one with
previous_iteration, the flag an XComIterable read through loop.previous
carries. Around it:

* XComKeysRequest body model, and an AddXComKeysEndpoint execution API version
  change so older clients keep their contract.
* has_xcom_access reads the optional key from the request's path parameters,
  so the endpoint needs no separate router or dependency.
* GetXComByKeys supervisor message, XComOperations.get_by_keys client method,
  handle_get_xcom_by_keys handler registered with the task supervisor and the
  DAG processor, an AddGetXComByKeys supervisor schema version entry and the
  regenerated schema snapshot.
* BaseXCom.get_by_keys sends the message and deserializes the values, and
  XComIterable iterates and slices through it: one request instead of one per
  value. A single index stays one XCom read.

Stacked on the task iteration PR, which owns XComIterable; only this commit
belongs to the endpoint. Squashed from the earlier history of this branch and
re-based on that PR, with XComIterable's batched reads moved to its current
home in bases/xcom.py.
@dabla
dabla force-pushed the feature/add-xcoms-keys-route branch from e4e3d40 to 56343f2 Compare October 14, 2026 20:50

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:API Airflow's REST/HTTP API area:dag-processor area:task-sdk skip newsfragment check Skip the newsfragment PR number check

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants