Add Execution API endpoints for callback lifecycle and token swap#66141
Add Execution API endpoints for callback lifecycle and token swap#66141wjddn279 wants to merge 2 commits into
Conversation
|
@wjddn279 Hey! Confirmed the regression is real and on me, generate_token flipped to workload scope at the base class but supervise_callback never got the paired swap. Thanks for catching it. Quick thought before you go deeper on the PR: this is now the third place we'd be hand-rolling the workload→execution swap (tasks /run, your callback /run, and #62343's connection-test PATCH). Wonder if it's worth lifting it into a generic POST /auth/exchange + a transparent swap in the SDK client, so resource endpoints stay execution-only and future workload types inherit the fix automatically. Happy to chat / draft if it sounds reasonable. Would love to hear @ashb @amoghrajesh @kaxil thoughts on this as well |
|
So you're suggesting we separate out only the API that exchanges the token into its own? Considering that workload types will likely grow, that sounds like a reasonable idea. But right now, we're tightening the security of the token swap by extremely restricting when the API can succeed, right? (e.g., only allowing it to succeed on the queued → running state transition.) Would that still be possible if we generalize it? |
Yeah, thinking about it again, you're right, the QUEUED → RUNNING gate gives us single-use structurally, and that's worth keeping, I think, maybe abstraction is not worth it as of now.Thanks for pointing out |
7709b17 to
4af312d
Compare
|
@wjddn279 — Your unresolved review thread(s) from @amoghrajesh appear to have been addressed (post-review commits and/or in-thread replies on every thread, with the latest commit pushed after the most recent thread). I've added the @amoghrajesh — could you take another look when you have a chance? If you agree the feedback was addressed, please mark the threads as resolved so the queue signal stays accurate. If a thread still needs work, please reply in-line — @wjddn279 will follow up. Note: This comment was drafted by an AI-assisted triage tool and may contain mistakes. Once you have addressed the points above, an Apache Airflow maintainer — a real person — will take the next look at your PR. We use this two-stage triage process so that our maintainers' limited time is spent where it matters most: the conversation with you. |
1325cbb to
137b659
Compare
|
conflicts :( |
137b659 to
a95e678
Compare
potiuk
left a comment
There was a problem hiding this comment.
Solid architectural change — moving callback state transitions onto the Execution API closes the EdgeExecutor (and any future remote-execution executor) gap, and the workload-token swap fixes the 403s on Connection / Variable lookups inside callbacks. The _process_executor_events fallback is correctly forward-only — terminal states won't regress on stale executor events. Tests are decent: 264 lines for the new endpoints, parametrized forward-only matrix in test_scheduler_job.py, and the supervise_callback bookend tests cover all three subprocess outcomes (success / non-zero exit / raise). CI is green.
@amoghrajesh's earlier review was substantially addressed in 7a0e3523 / a95e678d — the race-condition concern (callbacks.py:88), the callback.output-before-assignment regression, and the spec-less mocks all look fixed. One thread is still open: the dead-code branch on RUNNING events in scheduler_job_runner.py (see inline #1).
My own pass added a few more observations inline. Disposition: COMMENT — none are blocking, but I'd want #1, #2, and #3 addressed before APPROVE. The architecture is sound; this is about tightening edges.
This review was drafted by an AI-assisted tool and
confirmed by an Apache Airflow maintainer. The findings
below are observations, not blockers; an Apache Airflow
maintainer — a real person — will take the next look at the
PR. If you think a finding is mis-applied, please reply on
the PR and a maintainer will weigh in.More on how Apache Airflow handles maintainer review:
contributing-docs/05_pull_requests.rst.
potiuk
left a comment
There was a problem hiding this comment.
Follow-up from a second-read pass through Codex (adversarial reviewer). It found a real stranded-state path I missed in the primary review — [high] needs-attention, see inline. The fix interacts with my finding #1 (dead-code RUNNING branch): both are about the fallback being too narrow, and the right resolution probably addresses both at once.
This review was drafted by an AI-assisted tool and
confirmed by an Apache Airflow maintainer. The findings
below are observations, not blockers; an Apache Airflow
maintainer — a real person — will take the next look at the
PR. If you think a finding is mis-applied, please reply on
the PR and a maintainer will weigh in.More on how Apache Airflow handles maintainer review:
contributing-docs/05_pull_requests.rst.
potiuk
left a comment
There was a problem hiding this comment.
Requesting changes until the adversarial review findings are fixed. It looks pretty plausible.
9fefce2 to
9afb7c9
Compare
|
resolved it. thanks! |
ferruzzi
left a comment
There was a problem hiding this comment.
Thanks for sorting this out, I guess that was a gap that I missed. Cheers. I've left two small comments. Combined with Jarek's comments, and assuming the CI passes, I think it looks good.
b0ed76b to
e3f645b
Compare
|
Welcome back @ferruzzi and I hope you had a nice vacation! all is resolved and ci is green! thanks for checking |
|
fix test failure in #67542 |
d7258e2 to
0eb4daa
Compare
19571d3 to
0bad822
Compare
d0f31d4 to
a7a1b9f
Compare
|
@wjddn279 — A reviewer (@potiuk) has requested changes on this PR, so I've removed the Note: This comment was drafted by an AI-assisted triage tool and may contain mistakes. Once you have addressed the points above, an Apache Airflow maintainer — a real person — will take the next look at your PR. We use this two-stage triage process so that our maintainers' limited time is spent where it matters most: the conversation with you. |
e76d570 to
1ab613b
Compare
@potiuk maybe false-positive AI checker.. and I think it is ready to review |
ashb
left a comment
There was a problem hiding this comment.
Please follow our documented commit message rules
6ec6ac7 to
81c0410
Compare
Workers and the callback supervisor previously had no API-driven way to transition a callback through its lifecycle, forcing reliance on paths outside the Execution API. Add dedicated Execution API routes to mark a callback running and to record its terminal state, plus a token-swap endpoint so a callback can obtain a scoped JWT for its own callback id. The scheduler and the Task SDK client/callback supervisor are updated to use the new endpoints, keeping callback execution behind the Execution API instead of direct database access.
81c0410 to
c74431d
Compare
potiuk
left a comment
There was a problem hiding this comment.
Everything from the prior pass is addressed — including the one that actually mattered. The executor-event fallback now advances QUEUED → terminal (scheduler_job_runner.py:316-331), so a callback whose start() fails before the row flips to RUNNING no longer strands in QUEUED forever, and the forward-only invariant still holds (terminal rows aren't in the (QUEUED, RUNNING) set). Regression test test_process_executor_events_failed_event_captures_output locks it in. The dead RUNNING branch is gone, CallbackTerminalState is a real enum (no more generic State in _generated.py), the idempotent-retry + workload-token-swap test is in, and client.finish() is typed to the enum. Really solid follow-up.
Approving the substance — but please hold the merge until the one inline point below (the version filename vs label mismatch) is reconciled.
This review was drafted by an AI-assisted tool and confirmed by an Airflow maintainer. The maintainer approving this PR has read the findings and signed off. If something feels off, please reply on the PR and a maintainer will follow up.
More on how Airflow handles maintainer review: contributing-docs/05_pull_requests.rst.
| AddPartitionDateField, | ||
| ), | ||
| Version( | ||
| "2026-06-01", |
There was a problem hiding this comment.
Please reconcile before this merges: the version module is named v2026_07_01.py, but this label is "2026-06-01". They should agree. Two ways:
- If
2026-06-01is correct, rename the file tov2026_06_01.py(the import above and this entry's position between2026-06-30and2026-04-06are already right for that date). - If
2026-07-01was intended, change this label to"2026-07-01"and move the wholeVersion(...)block above the2026-06-30entry so the descending order stays correct.
It works today (Cadwyn keys off the string, and 2026-06-01 is correctly ordered), so this isn't a runtime bug — but the filename/label disagreement is a trap for the next person editing the version bundle.
|
I guess that one is too late for 3.3.0 ? But if the file name mismatch is solved and @ashb withdraws change request, maybe it can still be aded if there is an rc3 or maybe in 3.3.1 -> this one seems to solve Edge Executor issue. |
There was a problem hiding this comment.
Since this has missed 3.3.0 we will want to future-date this version to ~2026-07-31 (and the release manager will fix the version when it's included in a release)
| Mirrors ``PATCH /task-instances/{id}/run``: this is the single endpoint that | ||
| accepts a workload-scoped token and atomically (a) transitions the callback | ||
| from QUEUED to RUNNING and (b) issues a fresh execution-scoped token via the | ||
| ``Refreshed-API-Token`` response header. All subsequent supervisor calls hit | ||
| execution-only routes. |
There was a problem hiding this comment.
I believe this doc string is what shows up in the OpenAPI generated spec, and thus as inline docs for any clients generated off this.
As a result we should re-word this to make it be more user/client focused. (It's pretty close, I just don't think the "Mirrors PATCH" part is relevant to users. That can be a comment inside the fn itself if you think that is relevant.
There was a problem hiding this comment.
That's good to know, I'll try to remember that in the future.
| async def _token(request: Request) -> TIToken: | ||
| return TIToken(id=uuid4(), claims=TIClaims(scope="execution")) |
There was a problem hiding this comment.
I think we also need to rename TIToken -- given it's. now used for Callbacks (which is the right thing to do) calling this TI isn't correct.
That or we need to make a new type, CallbackToken and then have the security decorator return TIToken | CallbackToken (and then something inside TI routes ensure it's of the right type)
Removing blocking request changes, but see existing comments I've left
|
Not sure if or how this might interact with @seanghaeli 's |
Problems
1. Workload token leaks into callback supervisor's runtime API calls
related #60108, When the executor enqueues an ExecuteCallback, the supervisor receives a long-lived
workload-scoped token (intended only as a one-time admission ticket, like for ExecuteTask). However, unlike ExecuteTask — which immediately swaps it for anexecution-scoped token via PATCH /task-instances/{id}/run — supervise_callback never performs the swap. It carries theworkloadtoken straight through into runtime API calls (e.g. Connection.get(), Variable.get()).These shared endpoints only accept execution-scoped tokens, so any callback that touches Connections or Variables hits 403 from require_auth's scope check.
2. process_executor_events is the only writer of callback state, and it's executor-specific
Today, callback state transitions (QUEUED → RUNNING → SUCCESS/FAILED) are written only by the scheduler's _process_executor_events, which reads from each executor's in-process event_buffer. It does not work for executors that don't use the scheduler's in-process event channel — most notably EdgeExecutor, where workers run on remote machines and report back via HTTP. Their callbacks stay forever in QUEUED.
Solution
Move callback state transitions to the Execution API, mirroring the ExecuteTask pattern:
POST /callbacks/{id}/run— accepts both workload and execution tokens; transitions QUEUED → RUNNING and returns a fresh execution token via the Refreshed-API-Token header (same swap mechanism PR Two-token mechanism for task execution to prevent token expiration while tasks wait in executor queues #60108 introduced for TI). All subsequent callback API calls use the execution token.PATCH /callbacks/{id}/state— execution-only; records the terminal SUCCESS / FAILED outcome with optional output.supervise_callbackcallsclient.callbacks.start()before forking the subprocess andclient.callbacks.finish()afterward, so the API is the single source of truth for callbackstate — works uniformly across all executors (Local, Celery, Kubernetes, Edge).
_process_executor_eventskeeps its callback handling as a forward-only fallback for cases where the supervisor crashes before reporting a terminal state; it never regresses a state already advanced by the API.Was generative AI tooling used to co-author this PR?
claude code (opus 4.7) to make test codes and ci green
{pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.