Repository navigation
Conversation
boring-cyborg
Bot
added
area:API
a4c315a to
3dd249f
Compare
fat-catTW
left a comment
fat-catTW
left a comment
There was a problem hiding this comment.
I think the scheduler fallback could race with the new state endpoint here.
The scheduler reads the callback without a lock, checks that it is still QUEUED/RUNNING, and later writes the executor event's terminal state. In between those operations, the callback supervisor can successfully report SUCCESS (including its output) through PATCH /callbacks/{callback_id}/state.
For example, the callback subprocess can succeed and report SUCCESS, while the outer supervisor subsequently exits abnormally and the executor emits a late FAILED event. The scheduler could then overwrite the supervisor's already-recorded result, despite the intended first-writer-wins behaviour.
Maybe we could lock the callback row before rechecking its state, or use a conditional update that only succeeds while the persisted state is still QUEUED or RUNNING and add a regression test for this interleaving?
fat-catTW
left a comment
fat-catTW
left a comment
There was a problem hiding this comment.
lgtm
Thanks for the PR :D
seanghaeli
left a comment
seanghaeli
left a comment
There was a problem hiding this comment.
Also a conflict to deal with, but otherwise looks ready!
ramitkataria
left a comment
ramitkataria
left a comment
There was a problem hiding this comment.
The conditional UPDATE looks like a nice fix for the race. A few smaller things my agent found on top of the inline comments:
- +1 to the suggestion on `routes/callbacks.py` about returning 204 when the row is already in the requested state, so retries don't log a traceback.
- Which executors is this mainly for? Edge seems to already pass callback events back to the scheduler since #74784, so maybe worth naming the ones that don't in the description. For those, a failed `/run` or an OOM-killed supervisor would still leave the row stuck, so the comment at `callback_supervisor.py:469` might want a tweak, and maybe a follow-up issue for a timeout-based cleanup of stuck `QUEUED`/`RUNNING` callbacks?
- `jwt_token_authentication.rst` still only lists `/run` for callbacks, so maybe worth adding `/state` there too.
- After rebasing onto #75689, it might be worth double-checking the 404 case in `test_state_validates_token_scope_and_sub`, since it depends on the order the auth dependencies resolve in.
| raise | ||
| finally: | ||
| try: | ||
| client.callbacks.finish(UUID(id), state=terminal_state, output=terminal_output) |
There was a problem hiding this comment.
The execution token from /run seems to only live for [execution_api] jwt_expiration_time (600s by default), and since the callback supervisor doesn't heartbeat, I don't think anything refreshes it while the callback runs. So a callback that runs longer than ~10 min without reading a Connection or Variable would get a 403 here, which isn't retried and gets swallowed by the except. On executors with no executor-event fallback, that would leave the row RUNNING. Maybe the supervisor could refresh the token periodically?
There was a problem hiding this comment.
Let's also handle this in the same follow-up
| update(Callback) | ||
| .where( | ||
| Callback.id == callback_uuid, | ||
| Callback.state.in_((CallbackState.QUEUED, CallbackState.RUNNING)), |
There was a problem hiding this comment.
Could a duplicate delivery (e.g. Celery acks_late redelivery) end up overriding a real success here? The second supervisor's /run would get a 409 and fail, this UPDATE would write FAILED since the row is RUNNING, and then the original supervisor's finish(SUCCESS) gets a 409 and its SUCCESS event is ignored. Before this change the later SUCCESS would have overwritten it. Not sure what the cleanest fix is, since this path still needs to fail a RUNNING row whose supervisor died, so mostly flagging it in case it's worth handling or noting.
There was a problem hiding this comment.
Yeah, I see what you are saying. I think I can see a fix for it, but it would be a bigger change worth a followup. I'll think on it more, but the best answer I have right now would solve both of your points: We could maybe add a new column to the callback table for the heartbeat timestamp the same way TaskInstance has TaskInstance.last_heartbeat_at.
That would need a new column with a migration, a new PATCH endpoint to update it with a cadwyn migration, a new heartbeat loop in the supervisor to generate it, (three?) new config options to set the new heartbeat's timeout, a new reaper to clean up anything that goes stale... then once that's all in we can add the liveness check that makes sure that a FAILED event can only move a RUNNING row if the heartbeat is state. I think that would work but it feels heavy-handed... I'm sure there's a better way. Let me think on it some more.
There was a problem hiding this comment.
I can't think of a better solution and that's a relatively big bite of work to close a race condition that already exists. Are you alright with throwing that in the backlog and getting this merged? I don't believe (I could be wrong) this is making the existing state worse than main or introducing a new bug, right?
There was a problem hiding this comment.
Yes, that makes sense - I haven't fully thought through solutions yet either
Callback state was written only by the scheduler, reading each executor's in-process event buffer. Executors whose workers report back over HTTP never reach that channel, so their callbacks sat in QUEUED forever. The executor-event path stays as a fallback, now narrowed to rows still in flight so a late event cannot overwrite an outcome the supervisor already reported, or the output captured with it. Replaces apache#66141, whose /run half landed separately in apache#71192. Co-authored-by: wjddn279 <ehwjddn279@naver.com>
8486c4e to
726dc55
Compare
o-nikolas
left a comment
o-nikolas
left a comment
There was a problem hiding this comment.
Agreed about a separate PR for the callouts from @ramitkataria
| terminal_output = f"Callback subprocess exited with code {exit_code}" | ||
| except Exception as e: | ||
| terminal_output = f"Callback supervisor error: {type(e).__name__}: {e}" | ||
| log.exception("Callback supervisor error", workload_id=id) |
There was a problem hiding this comment.
nit: this log.exception plus raise gives two tracebacks for the same error, because the caller logs it too (LocalExecutor logs Workload execution failed. with the exception). Could you drop the log here and just record terminal_output?
There was a problem hiding this comment.
Ah, fair catch. I'll make this change and open the Issue for Ramit's follow-up on Monday.
Callback state was written only by the scheduler, reading each executor's in-process event buffer. Executors whose workers report back over HTTP never reach that channel, so their callbacks sat in QUEUED forever.
The executor-event path stays as a fallback, now narrowed to rows still in flight so a late event cannot overwrite an outcome the supervisor already reported, or the output captured with it.
Replaces #66141, whose /run half landed separately in #71192. Most of this is @wjddn279's code from that PR ported on top of the current main, co-author credit should carry over.
Was generative AI tooling used to co-author this PR?
Inheriting the disclaimer from the parent 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.