kaxil commented on code in PR #73249:
URL: https://github.com/apache/airflow/pull/73249#discussion_r4074603490
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -1611,22 +1597,27 @@ def update_task_state_if_needed(self):
)
def _send_terminal_state_msg(
- self, msg: SucceedTask | RetryTask | DeferTask | RescheduleTask |
AwaitInputTask
+ self, msg: TaskState | SucceedTask | RetryTask | DeferTask |
RescheduleTask | AwaitInputTask
) -> None:
- # Capture the message BEFORE the API call so the recovery dispatcher
- # in `update_task_state_if_needed` can re-issue it if the call raises
- # (network blip, transient server 5xx). Clear the pending slot and
- # record the resulting state only after the call returns successfully.
+ if self._terminal_state == SERVER_TERMINATED:
+ return
+ self._terminal_state = msg.state
self._pending_terminal_state_msg = msg
- if isinstance(msg, SucceedTask):
+ if isinstance(msg, TaskState):
+ self.client.task_instances.finish(
+ id=self.id,
+ state=msg.state,
+ when=datetime.now(tz=timezone.utc),
Review Comment:
Now that the whole `TaskState` is in hand at delivery, `when=msg.end_date or
datetime.now(tz=timezone.utc)` would match the `SucceedTask` branch below and
make the server-side end_date the task's end rather than process exit, which
now includes callback overtime. Pre-existing on main, so fine to ride with the
follow-on.
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -1585,23 +1576,18 @@ def wait(self) -> int:
return self._exit_code
def update_task_state_if_needed(self):
- # If a direct-state API call (succeed / retry / defer / reschedule)
- # was attempted but raised, `_pending_terminal_state_msg` still holds
- # the original request. Re-issue the matching dedicated API call so
- # the server learns the terminal state we couldn't deliver earlier.
- # Without this recovery, a transient API failure during the direct
- # call would leave the TI stuck RUNNING on the server — `finish()`
- # cannot substitute because the server-side `finish` endpoint does
- # not accept SUCCESS / DEFERRED / SERVER_TERMINATED transitions.
+ if self._terminal_state == SERVER_TERMINATED:
+ self._pending_terminal_state_msg = None
+ return
+
if self._pending_terminal_state_msg is not None:
- self._replay_pending_terminal_state_msg()
+ if isinstance(self._pending_terminal_state_msg, TaskState):
+ self._send_terminal_state_msg(self._pending_terminal_state_msg)
+ else:
+ self._replay_pending_terminal_state_msg()
return
- # If the process has finished in a non-directly-patched state (e.g.
- # FAILED, or SKIPPED reported via a TaskState message), `finish()` is
- # the dedicated endpoint for those transitions. For states already in
- # STATES_SENT_DIRECTLY whose direct API call succeeded, no further
- # action is needed.
+ # A process can exit before sending an outcome; report the inferred
state.
Review Comment:
Not quite true for the retry case. When a worker dies without sending
anything (OOM kill, segfault) on a TI with retries left, `final_state` infers
`UP_FOR_RETRY`, which is in `STATES_SENT_DIRECTLY`, so the guard below skips
and nothing is reported; the TI sits until zombie detection picks it up. Same
as main, so not for this PR, but the comment reads as if the inference is
always delivered. Either narrow it to "when `finish()` can carry it", or route
the retry inference through `retry()` in the follow-on.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
