kaxil commented on code in PR #73142:
URL: https://github.com/apache/airflow/pull/73142#discussion_r4029077533
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -1804,23 +1804,32 @@ def final_state(self):
"""
The final state of the TaskInstance.
- By default, this will be derived from the exit code of the task
- (0=success, failed otherwise) but can be changed by the subprocess
- sending a TaskState message, as long as the process exits with 0
+ If the subprocess reported a terminal state via message (TaskState,
SucceedTask,
Review Comment:
Your correction comment settles the mechanism, but the retracted framing is
still in three places that outlive this thread: the title, the body's first
paragraph, and commit `855019e0ef`. Airflow squash-merges with every commit
message concatenated into the merge body (see `4c3ec9cbae`), so that paragraph
and `2adc8b28ae`'s "Add newsfragment" for a fragment that no longer exists both
become permanent history. Worth a retitle plus a reword of the body and the
branch commits.
The body's claimed harm also can't happen as written: `UP_FOR_RETRY` is
inside `STATES_SENT_DIRECTLY`, so `finish()` is never reached, and
`client.finish()` would raise `ValueError` on it anyway since
`TerminalStateNonSuccess` has no `up_for_retry`. The overtime SIGTERM example
here isn't wrong, it's reachable and does produce the bug, it just isn't the
case you captured, so naming the captured one is clearer. The strongest case
for this fix is `TaskState(SKIPPED)`: that arm (1844) makes no API call, so the
row is still `running` when a non-zero exit lands, and base either writes
`finish(FAILED)` over a skipped task or, with retries configured, skips the
write and leaves the row `running`.
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -4178,6 +4178,41 @@ def
test_non_signal_exit_code_without_retry_goes_to_failed(self, mocker):
assert mock_watched_subprocess.final_state == TaskInstanceState.FAILED
+ @pytest.mark.parametrize("should_retry", [True, False])
+ def
test_confirmed_terminal_state_takes_precedence_over_later_nonzero_exit_code(
+ self, mocker, should_retry
+ ):
+ """
+ A terminal state reported via message (e.g. SucceedTask) is
authoritative even if the
+ subprocess is later killed with a genuinely non-zero exit code -- e.g.
+ `_handle_process_overtime_if_needed()` sending SIGTERM once
`_terminal_state` is already
+ set. Regression test for
https://github.com/apache/airflow/issues/65708.
+
+ Only `should_retry=False` actually reproduces the original crash: with
+ `should_retry=True` the pre-fix code already returned UP_FOR_RETRY,
which is in
+ STATES_SENT_DIRECTLY, so `update_task_state_if_needed()` would already
skip `.finish()`.
+ The crash needs `should_retry=False`, where the pre-fix code fell
through to FAILED,
+ which is *not* in that set -- triggering a spurious `.finish()` call
and a 409 against
+ the already-correct DB row. Both values are asserted here so the fix
is pinned for
+ either configuration, not just the value that happens to match the new
state.
+ """
+ mock_watched_subprocess = ActivitySubprocess(
+ process_log=mocker.MagicMock(),
+ id=TI_ID,
+ pid=12345,
+ stdin=mocker.Mock(),
+ process=mocker.Mock(),
+ client=mocker.Mock(),
+ )
+ mock_watched_subprocess._terminal_state = TaskInstanceState.SUCCESS
+ mock_watched_subprocess._exit_code = 1 # genuinely observed, e.g.
from a SIGTERM kill
+ mock_watched_subprocess._should_retry = should_retry
+
+ assert mock_watched_subprocess.final_state == TaskInstanceState.SUCCESS
Review Comment:
This assertion fails on base for both parametrized values, so pytest stops
here and the `finish.assert_not_called()` two lines down never runs against the
unfixed code. Moving `update_task_state_if_needed()` and that assertion above
this line costs nothing and makes `should_retry=False` actually pin the
spurious `finish()` call, which was the point of adding it.
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -4178,6 +4178,41 @@ def
test_non_signal_exit_code_without_retry_goes_to_failed(self, mocker):
assert mock_watched_subprocess.final_state == TaskInstanceState.FAILED
+ @pytest.mark.parametrize("should_retry", [True, False])
+ def
test_confirmed_terminal_state_takes_precedence_over_later_nonzero_exit_code(
+ self, mocker, should_retry
+ ):
+ """
+ A terminal state reported via message (e.g. SucceedTask) is
authoritative even if the
+ subprocess is later killed with a genuinely non-zero exit code -- e.g.
+ `_handle_process_overtime_if_needed()` sending SIGTERM once
`_terminal_state` is already
+ set. Regression test for
https://github.com/apache/airflow/issues/65708.
+
+ Only `should_retry=False` actually reproduces the original crash: with
+ `should_retry=True` the pre-fix code already returned UP_FOR_RETRY,
which is in
+ STATES_SENT_DIRECTLY, so `update_task_state_if_needed()` would already
skip `.finish()`.
+ The crash needs `should_retry=False`, where the pre-fix code fell
through to FAILED,
+ which is *not* in that set -- triggering a spurious `.finish()` call
and a 409 against
+ the already-correct DB row. Both values are asserted here so the fix
is pinned for
+ either configuration, not just the value that happens to match the new
state.
+ """
+ mock_watched_subprocess = ActivitySubprocess(
+ process_log=mocker.MagicMock(),
+ id=TI_ID,
+ pid=12345,
+ stdin=mocker.Mock(),
+ process=mocker.Mock(),
+ client=mocker.Mock(),
+ )
+ mock_watched_subprocess._terminal_state = TaskInstanceState.SUCCESS
Review Comment:
This pins the 409 for `should_retry=False`, which is worth having, but
`SUCCESS` is the one `_terminal_state` where `succeed()` has already written
the row, so nothing here covers the half of the change that alters what gets
persisted. `FAILED`, `SKIPPED` and `REMOVED` are the only `_terminal_state`
values outside `STATES_SENT_DIRECTLY`, and their arm at 1844 makes no API call,
so the row is still `running` when the overtime SIGTERM lands: base either
writes `finish(FAILED)` over a skipped task or, with retries configured,
returns `UP_FOR_RETRY` and skips the write entirely, leaving the row `running`.
Neither is tested -- `final_state` and `update_task_state_if_needed` only ever
see `SUCCESS` (702, 3701, 4207) and `SERVER_TERMINATED` (4162), and all three
`finish` assertions in this file are `assert_not_called`. A `SKIPPED` case
asserting `finish` is called with `state=skipped` would cover it, and
`test_overtime_handling` already parametrizes `SKIPPED` with
`expected_kill=True` (1
097), so the combination is reachable.
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -1804,23 +1804,32 @@ def final_state(self):
"""
The final state of the TaskInstance.
- By default, this will be derived from the exit code of the task
- (0=success, failed otherwise) but can be changed by the subprocess
- sending a TaskState message, as long as the process exits with 0
+ If the subprocess reported a terminal state via message (TaskState,
SucceedTask,
+ RetryTask, etc.) before exiting, that message is authoritative and
takes precedence
+ over the exit code -- even over a genuinely non-zero one. This matters
when the
+ subprocess is killed *after* it already reported success or another
terminal state,
+ e.g. `_handle_process_overtime_if_needed()` sending SIGTERM once
`_terminal_state` is
+ set: without this precedence, a real (non-defaulted) non-zero exit
code would override
+ an already-confirmed terminal state and re-derive a stale one from the
exit code alone,
+ triggering a redundant `update_task_state_if_needed()` -> `.finish()`
call that 409s
+ against the row the earlier message-driven update already wrote
correctly. Only fall
Review Comment:
`TaskState` is the first message the paragraph names, and its arm at 1844
makes no API call, so there is no row that "the earlier message-driven update
already wrote correctly": the row is still `running` and `finish()` is the
first and intended writer, as the comment on that arm says. "Reported via
message" is also narrower than the `is not None` the code actually checks,
since `_send_heartbeat_if_needed` sets `SERVER_TERMINATED` with no message from
the subprocess at all.
--
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]
