Repository navigation
Conversation
ColtenOuO
left a comment
ColtenOuO
left a comment
There was a problem hiding this comment.
self review and check all 22 files.
|
Updared the PR's description (added Backward compatibility) |
|
Fixed DAG File Processor forward |
Lays the groundwork for resolving multiple XComs in one round trip instead of one per lookup. Kept separate from the caller that will actually use it so this additive, backward-compatible piece can be reviewed on its own before the behavior change that depends on it.
…call A mapped task with N plain, non-mapped-upstream XComArg kwargs made N sequential Execution API round trips to resolve them, growing linearly with the number of kwargs. Batching collapses that into one round trip for the common case, falling back to the old per-item behavior for custom XCom backends, older API servers, and anything not eligible for batching (mapped upstreams, task groups, composite XComArgs).
Bounds the batch endpoint to a fixed max item count so an authenticated task can't force an oversized bind-parameter list or unbounded work on the API server. Also replaces unspecced MagicMocks in the new batch tests with autospec'd ones so a signature or attribute drift on the mocked classes would actually fail the test, and adds coverage for the found=False path (a missing return_value XCom resolves to None; a missing non-return_value key raises, matching the per-item xcom_pull path) and for a batch item with an explicit, non-default map_index, which the original tests didn't exercise.
The regenerated Task SDK client was stale after adding the item-count cap to XComBatchRequestBody, missing the max_length constraint in the committed _generated.py. The DAG File Processor already forwards every other XCom read message (GetXCom, GetXComCount, GetXComSequenceItem, GetXComSequenceSlice) to the API server the same way the task-runner supervisor does, so GetXComBatch needs the same wiring for the message-type completeness checks to hold. The Triggerer has no equivalent need for it (deferred triggers don't resolve expand() kwargs), so it's excluded there the same way the existing sequence/count XCom messages already are.
8c8606c to
f57b7ac
Compare
potiuk
added
the
closed because of open PR limit
|
Hello @ColtenOuO - thank you for your contributions to Apache Airflow! The Airflow community has introduced a limit of 5 open pull requests at a time for contributors without write access to the repository. You currently have 19 open pull requests, so - as a one-time step of introducing the limit - we closed the ones where maintainers have not engaged yet:
These pull requests stay open because maintainers are already engaged in them - they count towards your limit:
This is not a judgement of you or of your changes. We never told contributors before that opening many pull requests at once was a problem, so there is nothing to feel bad about - and nothing is lost: your branches, commits and the review history stay where they are. What we ask you to do is to make your first prioritization decision: choose which of the pull requests above matter most to you, and reopen them (up to 5 open at a time, including the ones still open) with the "Reopen pull request" button or While your pull requests are waiting for review, the most valuable thing you can do is help in other ways - reviewing other contributors' pull requests, helping with issues, and taking part in the discussions on the devlist and Slack. Why we introduced the limit, what it means for you and how to reopen or restore a pull request is explained in https://github.com/apache/airflow/blob/main/contributing-docs/32_open_pull_request_limit.rst. Drafted-by: Claude Code (Opus 5); reviewed by @potiuk before posting |
2 participants
Summary
DictOfListsExpandInput.resolve()calls.resolve(context)on everyXComArgkwarg 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 ownti.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 inexpandinput.pyalready carried aTODOadmitting 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 asfound=Falseinstead 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 request body is capped atMAX_XCOM_BATCH_ITEMS(1000) items — each item expands to 3 SQL bind parameters, so an unbounded batch from an authenticated task could otherwise push a query past Postgres/MySQL/SQLite bind-parameter limits or spend unbounded API-server work on a single request; an oversized request is rejected with 422 before touching the database.Task SDK plumbing —
XComOperations.get_batch()on the generated client, aGetXComBatch/XComBatchResultmessage pair on the supervisor comms protocol, and the matching supervisor-side dispatch handler, following the exact same shape as the existing single-itemGetXCompath.Wiring in
DictOfListsExpandInput.resolve()— before resolving each kwarg individually, it now collects the eligible ones (plainXComArg, non-mapped upstream, not inside a mapped task group — the same condition that already made the per-item path pull withmap_indexes=None) and resolves them with a single batched call. Everything else (mapped-upstreamLazyXComSequence,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.A custom XCom backend (
core.xcom_backend) also disables batching entirely, falling back to the existing per-item path — a custom backend'sget_one/get_allmay bypass the batch endpoint's semantics, so this sidesteps that risk rather than trying to make batching custom-backend-aware. If the batch response is ever neither a valid result nor the "not supported" error (a protocol contract violation, not an expected outcome), it raises a clearTypeErrorinstead of silently resolving to wrong values.Backward compatibility
This PR touches both sides of the wire (API server and Task SDK client), so all four version-skew combinations that can show up during a rolling upgrade are handled explicitly:
AddXComBatchEndpoint,endpoint(...).didnt_existfor versions before it). An old Task SDK is pinned to an older API version, so the endpoint doesn't exist in the contract it negotiates against; it never calls it and keeps using the pre-existing per-itemGetXCompath unchanged.ErrorType.XCOM_BATCH_NOT_SUPPORTED, andDictOfListsExpandInput.resolve()catches exactly that and falls back to resolving the group per-item. The task still succeeds; it just doesn't get the speed-up until the server catches up.Everything added (the route, the
GetXComBatch/XComBatchResultcomms messages,XComOperations.get_batch()) is strictly additive — no existing endpoint, message shape, or client method changes shape — so the two sides can roll out in either order without a hard dependency.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):
The current path scales linearly with N (~17ms/call); the batched path stays flat (~25ms total) regardless of N.
Testing
found=False), empty batch, over-the-cap request rejected with 422, a batch item with an explicit non-defaultmap_index, and multi-team access-control (same-team allowed, cross-team forbidden) for the batch endpoint.ErrorType.XCOM_BATCH_NOT_SUPPORTEDfallback signal.REQUEST_TEST_CASEStable) coveringGetXComBatch→XComBatchResult.expandinput/xcom_argtests: N kwargs collapse into exactly oneGetXComBatchcall; 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); thefound=Falsepath (a missingreturn_valueXCom resolves toNone, a missing non-return_valuekey raisesXComNotFound, matching the per-itemxcom_pullpath exactly). Mocks in the new SDK-level tests areautospec'd againstBaseOperator/RuntimeTaskInstancerather than bareMagicMock, so a signature or attribute drift on those classes would fail the test instead of passing silently.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 newGetXComBatchmessage — a real regression this PR would otherwise have introduced, caught by running the full suite rather than only the new tests.task-sdkand the touchedairflow-corefiles.check-execution-api-versions,check-supervisor-schemas-versions) pass.Scope / follow-ups
Batching for a mapped-upstream
LazyXComSequence(theGetXComCountlength lookup) and forMapXComArg/ZipXComArg/ConcatXComArgcomposite 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?
Generated-by: Claude Code (Sonnet 5) following the guidelines