Repository navigation
Conversation
dabla
requested review from
amoghrajesh,
ashb,
ephraimbuddy,
jedcunningham and
kaxil
as code owners
July 22, 2026 09:24
amoghrajesh
reviewed
Jul 22, 2026
amoghrajesh
left a comment
Contributor
There was a problem hiding this comment.
I haven't followed up closely with AIP 104, but I have some qns from an initial look.
dabla
force-pushed
the
feature/add-xcoms-keys-route
branch
from
July 22, 2026 14:10
5258e09 to
dbf25d5
Compare
dabla
marked this pull request as draft
August 17, 2026 08:52
Contributor
Author
|
This PR should only be merged once AIP-104 is merged. |
dabla
force-pushed
the
feature/add-xcoms-keys-route
branch
from
September 27, 2026 20:36
a6f73fd to
f688623
Compare
1 task done
Contributor
Author
|
Interaction with #73807 to handle when both land: |
dabla
force-pushed
the
feature/add-xcoms-keys-route
branch
12 times, most recently
from
October 1, 2026 10:19
04cc40e to
b74f42c
Compare
…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
force-pushed
the
feature/add-xcoms-keys-route
branch
from
October 14, 2026 20:50
e4e3d40 to
56343f2
Compare
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 toXComIterable's current home inbases/xcom.py.Add a
POST /xcoms/{dag_id}/{run_id}/{task_id}/keysexecution API endpoint thataccepts a list of XCom keys and returns all matching values in a single database
query. Use it in
XComIterableto reduce iteration from N round-trips to one.Motivation
XComIterable(introduced in AIP-104 Task Iteration) stores per-index resultsunder distinct keys (
return_value_0,return_value_1, …) with the samemap_index. The existingGET …/sliceendpoint cannot be reused because itranges over
map_indexfor a single key — the inverse structure. Iterating orslicing a 1000-item result previously issued 1000 separate
XCom.get_onecalls.Changes
POST /xcoms/{dag_id}/{run_id}/{task_id}/keys— accepts{"keys": [...]}, returns values in request-key order (nullfor missingkeys), filtered by
map_indexquery param (default-1)has_xcom_accesssimplified: reads the optionalkeyfromrequest.path_paramsinstead of requiring a{key}path segment, so aseparate router/dependency is no longer needed
XComKeysRequestPydantic model added to execution API data modelsGetXComByKeyscomms message,XComOperations.get_by_keysclient method,handle_get_xcom_by_keyssupervisor handler, and processor dispatch wiringBaseXCom.get_by_keyssends the message and deserializes the values;XComIterable.__iter__and__getitem__slice read through it, oneround-trip for a full iteration or slice. A single index stays one XCom read.
AddXComKeysEndpointexecution API version change andAddGetXComByKeyssupervisor schema version entry, with the regenerated schema snapshot.
TestGetXComByKeys) and forXComIterable(in
test_xcom.py), plus supervisor message coverage.Was generative AI tooling used to co-author this PR?
Generated-by: Claude Code following the guidelines
🤖 Generated with Claude Code