This is an automated email from the ASF dual-hosted git repository.
ashb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 61340ea9a73 Separate task outcomes from report delivery in TaskSDK
Supervisor (#73249)
61340ea9a73 is described below
commit 61340ea9a73eb82228ac8c24fc366cd265de6e7d
Author: Ash Berlin-Taylor <[email protected]>
AuthorDate: Wed Sep 23 21:25:19 2026 +0100
Separate task outcomes from report delivery in TaskSDK Supervisor (#73249)
The supervisor uses _terminal_state inconsistently: failed and skipped
outcomes are recorded when the worker reports them, while other outcomes
are recorded only after the API call succeeds. This couples the known
task outcome to API availability and makes deferred reporting harder to
reason about.
This changes things so we record the worker’s outcome immediately on an
instance field and use the existing pending message to track delivery.
Failed API calls retain that message and keep heartbeats running. Upon
task process exit pending reports take precedence over exit-code
inference. Server-directed termination (zombie detection, external
clearing, etc.) overrides pending reporting and any worker outcomes
received during shutdown.
The previous tests required _terminal_state to remain unset after an API
failure because setting it stopped heartbeats. That assertion reflected
the old coupling between outcome and delivery. The updated tests check
the underlying guarantees: retain the report, continue heartbeats, and
recover through the correct endpoint.
Post-task finalization timeouts now depend on the recorded completion
timestamp, independently of whether reporting succeeds. This also bounds
finalization when the initial API call fails.
---
.../src/airflow/sdk/execution_time/supervisor.py | 97 +++++-------
.../task_sdk/execution_time/test_supervisor.py | 171 +++++++++++++++++++--
2 files changed, 199 insertions(+), 69 deletions(-)
diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
index 7c28f0aeaca..d722f7d5a55 100644
--- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
+++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
@@ -1598,18 +1598,9 @@ class ActivitySubprocess(WatchedSubprocess):
"""The HTTP client to use for communication with the API server."""
_terminal_state: str | None = attrs.field(default=None, init=False)
- _final_state: str | None = attrs.field(default=None, init=False)
- # The terminal-state message currently being processed by
`_handle_request`,
- # captured BEFORE the dedicated API call (succeed / retry / defer /
- # reschedule). If the API call raises (network blip, server 5xx, etc.),
- # this attribute stays set and the dispatcher in
- # `update_task_state_if_needed` re-issues the matching API call on
- # subprocess exit — re-attempting the original transition rather than
- # falling back to `finish()`, which doesn't accept SUCCESS / DEFERRED /
- # SERVER_TERMINATED on the server side. Cleared (and `_terminal_state`
- # set) only after the API call returns successfully.
+ # Retain the full report until delivery succeeds, including reports
deferred until process exit.
_pending_terminal_state_msg: (
- SucceedTask | RetryTask | DeferTask | RescheduleTask | AwaitInputTask
| None
+ TaskState | SucceedTask | RetryTask | DeferTask | RescheduleTask |
AwaitInputTask | None
) = attrs.field(default=None, init=False)
_last_successful_heartbeat: float = attrs.field(default=0, init=False)
@@ -1736,23 +1727,18 @@ class ActivitySubprocess(WatchedSubprocess):
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.
+ # Without a worker outcome, only report inferred states that finish()
accepts.
if self.final_state not in STATES_SENT_DIRECTLY:
self.client.task_instances.finish(
id=self.id,
@@ -1762,14 +1748,20 @@ class ActivitySubprocess(WatchedSubprocess):
)
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=msg.end_date or datetime.now(tz=timezone.utc),
+ rendered_map_index=self._rendered_map_index,
+ )
+ elif isinstance(msg, SucceedTask):
self.client.task_instances.succeed(
id=self.id,
when=msg.end_date,
@@ -1777,7 +1769,6 @@ class ActivitySubprocess(WatchedSubprocess):
outlet_events=msg.outlet_events,
rendered_map_index=self._rendered_map_index,
)
- self._terminal_state = msg.state
elif isinstance(msg, RetryTask):
self.client.task_instances.retry(
id=self.id,
@@ -1786,16 +1777,12 @@ class ActivitySubprocess(WatchedSubprocess):
retry_delay_seconds=getattr(msg, "retry_delay_seconds", None),
retry_reason=getattr(msg, "retry_reason", None),
)
- self._terminal_state = msg.state
elif isinstance(msg, DeferTask):
self.client.task_instances.defer(self.id, msg)
- self._terminal_state = TaskInstanceState.DEFERRED
elif isinstance(msg, RescheduleTask):
self.client.task_instances.reschedule(self.id, msg)
- self._terminal_state = TaskInstanceState.UP_FOR_RESCHEDULE
elif isinstance(msg, AwaitInputTask):
self.client.task_instances.await_input(self.id, msg)
- self._terminal_state = TaskInstanceState.AWAITING_INPUT
self._pending_terminal_state_msg = None
def _replay_pending_terminal_state_msg(self) -> None:
@@ -1883,13 +1870,9 @@ class ActivitySubprocess(WatchedSubprocess):
def _handle_process_overtime_if_needed(self):
"""Handle termination of auxiliary processes if the task exceeds the
configured overtime."""
- # If the task has reached a terminal state, we can start monitoring
the overtime
- if not self._terminal_state:
+ if self._task_end_time_monotonic is None:
return
- if (
- self._task_end_time_monotonic
- and (time.monotonic() - self._task_end_time_monotonic) >
TASK_OVERTIME_THRESHOLD
- ):
+ if (time.monotonic() - self._task_end_time_monotonic) >
TASK_OVERTIME_THRESHOLD:
log.warning(
"Task success overtime reached; terminating process. "
"Modify `task_success_overtime` setting in [core] section of "
@@ -1904,9 +1887,9 @@ class ActivitySubprocess(WatchedSubprocess):
if (time.monotonic() - self._last_heartbeat_attempt) <
MIN_HEARTBEAT_INTERVAL:
return
- if self._terminal_state:
- # If the task has finished, and we are in "overtime" (running OL
listeners etc) we shouldn't
- # heartbeat
+ if self._terminal_state == SERVER_TERMINATED:
+ return
+ if self._terminal_state and self._pending_terminal_state_msg is None:
return
self._last_heartbeat_attempt = time.monotonic()
@@ -1929,9 +1912,11 @@ class ActivitySubprocess(WatchedSubprocess):
"Server indicated the task shouldn't be running anymore.
Terminating process",
detail=e.detail,
)
+ # kill() drains worker messages while waiting for the process
to exit.
+ self._terminal_state = SERVER_TERMINATED
+ self._pending_terminal_state_msg = None
self.kill(signal.SIGTERM, force=True)
self.process_log.error("Task killed!")
- self._terminal_state = SERVER_TERMINATED
else:
# If we get any other error, we'll just log it and try again
next time
self._handle_heartbeat_failures(e)
@@ -1962,14 +1947,17 @@ class ActivitySubprocess(WatchedSubprocess):
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
+ sending a TaskState message, as long as the process exits with 0.
+
+ A nonzero exit code can override the worker outcome here. Deliver
reports
+ using the message state, not this inferred state.
Not valid before the process has finished.
"""
+ if self._terminal_state == SERVER_TERMINATED:
+ return self._terminal_state
if self._exit_code == 0:
return self._terminal_state or TaskInstanceState.SUCCESS
- if self._exit_code != 0 and self._terminal_state == SERVER_TERMINATED:
- return SERVER_TERMINATED
# Any non zero exit code indicates a failure
# If retries are configured, mark as UP_FOR_RETRY
@@ -1989,12 +1977,11 @@ class ActivitySubprocess(WatchedSubprocess):
super()._handle_request(msg, log, req_id)
def _handle_task_state(self, msg: TaskState, log: FilteringBoundLogger,
req_id: int) -> RequestResult:
- # No direct API call here — the recovery path in
- # `update_task_state_if_needed` will call `finish()` for
- # non-direct states (FAILED, etc.) once the subprocess exits.
- self._terminal_state = msg.state
- self._task_end_time_monotonic = time.monotonic()
- self._rendered_map_index = msg.rendered_map_index
+ if self._terminal_state != SERVER_TERMINATED:
+ self._terminal_state = msg.state
+ self._pending_terminal_state_msg = msg
+ self._task_end_time_monotonic = time.monotonic()
+ self._rendered_map_index = msg.rendered_map_index
return None, {}
def _handle_finished_task(
diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
index e94b901179d..a55cef8aa58 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
@@ -35,7 +35,7 @@ from operator import attrgetter
from random import randint
from textwrap import dedent
from time import sleep
-from typing import TYPE_CHECKING, Any, get_args
+from typing import TYPE_CHECKING, Any, get_args, get_type_hints
from unittest import mock
from unittest.mock import MagicMock, patch
@@ -1091,6 +1091,9 @@ class TestWatchedSubprocess:
False,
id="no_terminal_state",
),
+ pytest.param(None, None, 10, False, id="execution_still_running"),
+ pytest.param(None, 9.0, 10, True,
id="finalization_without_reported_state"),
+ pytest.param(None, 0.0, 10, True,
id="finalization_started_at_zero"),
pytest.param(TaskInstanceState.SUCCESS, 15.0, 10, False,
id="below_threshold"),
pytest.param(TaskInstanceState.SUCCESS, 9.0, 10, True,
id="above_threshold"),
pytest.param(TaskInstanceState.FAILED, 9.0, 10, True,
id="above_threshold_failed_state"),
@@ -3330,6 +3333,119 @@ REQUEST_TEST_CASES = [
class TestHandleRequest:
+ @pytest.mark.parametrize("arrival", ["during_kill", "after_kill"])
+ @pytest.mark.parametrize(
+ ("msg", "api_method"),
+ [
+ (TaskState(state=TaskInstanceState.FAILED), "finish"),
+ (SucceedTask(end_date=timezone.parse("2024-10-31T12:00:00Z")),
"succeed"),
+ (RetryTask(end_date=timezone.parse("2024-10-31T12:00:00Z")),
"retry"),
+ (DeferTask(next_method="execute_complete",
classpath="test.Trigger", trigger_kwargs={}), "defer"),
+ (
+ RescheduleTask(
+ end_date=timezone.parse("2024-10-31T12:00:00Z"),
+ reschedule_date=timezone.parse("2024-10-31T12:01:00Z"),
+ ),
+ "reschedule",
+ ),
+ (AwaitInputTask(next_method="execute_complete"), "await_input"),
+ ],
+ )
+ def test_server_termination_ignores_late_outcome(
+ self, watched_subprocess, mocker, arrival, msg, api_method
+ ):
+ process, _ = watched_subprocess
+ process._handle_request(TaskState(state=TaskInstanceState.FAILED),
structlog.get_logger(), req_id=1)
+ process.client.task_instances.heartbeat.side_effect =
ServerResponseError.from_response(
+ httpx.Response(
+ 409,
+ request=httpx.Request("PUT", "http://server/heartbeat";),
+ json={"detail": "already stopped"},
+ )
+ )
+ observed_at_kill = []
+
+ def terminate(self, signal_to_send, force):
+ observed_at_kill.append((self._terminal_state,
self._pending_terminal_state_msg))
+ if arrival == "during_kill":
+ self._handle_request(msg, structlog.get_logger(), req_id=2)
+ self._exit_code = -signal.SIGTERM
+
+ mocker.patch.object(ActivitySubprocess, "kill", autospec=True,
side_effect=terminate)
+ process._should_retry = True
+
+ process._send_heartbeat_if_needed()
+ if arrival == "after_kill":
+ process._handle_request(msg, structlog.get_logger(), req_id=2)
+ process.update_task_state_if_needed()
+
+ assert observed_at_kill == [(supervisor.SERVER_TERMINATED, None)]
+ assert process.final_state == supervisor.SERVER_TERMINATED
+ assert process._pending_terminal_state_msg is None
+ getattr(process.client.task_instances, api_method).assert_not_called()
+ process.client.task_instances.finish.assert_not_called()
+
+ @pytest.mark.parametrize("exit_code", [0, -signal.SIGTERM])
+ def test_server_termination_cancels_pending_report(self,
watched_subprocess, exit_code):
+ process, _ = watched_subprocess
+ process._handle_request(TaskState(state=TaskInstanceState.FAILED),
structlog.get_logger(), req_id=1)
+ process._terminal_state = supervisor.SERVER_TERMINATED
+ process._exit_code = exit_code
+
+ process._send_heartbeat_if_needed()
+ process.update_task_state_if_needed()
+
+ process.client.task_instances.heartbeat.assert_not_called()
+ process.client.task_instances.finish.assert_not_called()
+ assert process._pending_terminal_state_msg is None
+ assert process.final_state == supervisor.SERVER_TERMINATED
+
+ @pytest.mark.parametrize(
+ "state", [TaskInstanceState.FAILED, TaskInstanceState.SKIPPED,
TaskInstanceState.REMOVED]
+ )
+ @pytest.mark.parametrize("exit_code", [0, 1, -signal.SIGTERM])
+ @pytest.mark.parametrize("should_retry", [False, True])
+ @pytest.mark.parametrize("end_date", [None,
timezone.parse("2024-10-31T12:00:00Z")])
+ def test_task_state_waits_for_exit_and_keeps_heartbeating(
+ self, watched_subprocess, time_machine, state, exit_code,
should_retry, end_date
+ ):
+ process, _ = watched_subprocess
+ process._should_retry = should_retry
+ msg = TaskState(state=state, rendered_map_index="label",
end_date=end_date)
+
+ process._handle_request(msg, structlog.get_logger(), req_id=1)
+
+ assert process._terminal_state == state
+ assert process._pending_terminal_state_msg is msg
+ process.client.task_instances.finish.assert_not_called()
+ process._send_heartbeat_if_needed()
+ process.client.task_instances.heartbeat.assert_called_once()
+ process._exit_code = exit_code
+ report_time = timezone.parse("2024-10-31T12:01:00Z")
+ time_machine.move_to(report_time, tick=False)
+ process.update_task_state_if_needed()
+
+ assert process.client.task_instances.finish.call_args.kwargs["state"]
== state
+ assert process.client.task_instances.finish.call_args.kwargs["when"]
== (end_date or report_time)
+ assert
process.client.task_instances.finish.call_args.kwargs["rendered_map_index"] ==
"label"
+ assert process._pending_terminal_state_msg is None
+
+ def test_failed_finish_retains_pending_report(self, watched_subprocess,
mocker):
+ process, _ = watched_subprocess
+ msg = TaskState(state=TaskInstanceState.FAILED)
+ process._handle_request(msg, structlog.get_logger(), req_id=1)
+ process._exit_code = 0
+ process.client.task_instances.finish = mocker.Mock(
+ spec=sdk_client.TaskInstanceOperations.finish,
+ side_effect=httpx.ConnectError("connection refused"),
+ )
+
+ with pytest.raises(httpx.ConnectError):
+ process.update_task_state_if_needed()
+
+ assert process._terminal_state == TaskInstanceState.FAILED
+ assert process._pending_terminal_state_msg is msg
+
class _OverrideActivitySubprocess(ActivitySubprocess):
def _handle_set_rendered_map_index(
self, msg: SetRenderedMapIndex, log: FilteringBoundLogger, req_id:
int
@@ -3711,32 +3827,35 @@ class TestHandleRequest:
TaskInstanceState.UP_FOR_RESCHEDULE,
id="reschedule",
),
+ pytest.param(
+ AwaitInputTask(next_method="execute_complete"),
+ "await_input",
+ TaskInstanceState.AWAITING_INPUT,
+ id="await_input",
+ ),
],
)
- def test_terminal_state_not_set_when_direct_api_fails(
+ def test_worker_outcome_retained_when_direct_api_fails(
self, watched_subprocess, mocker, msg, api_method, expected_state
):
- """`_terminal_state` must NOT be set when the dedicated terminal-state
- API raises.
-
- The original message is captured in `_pending_terminal_state_msg`
- BEFORE the API call so the recovery dispatcher in
- `update_task_state_if_needed` can re-issue it on subprocess exit.
- Covers all four terminal-state message types.
- """
watched_subprocess, _ = watched_subprocess
setattr(
watched_subprocess.client.task_instances,
api_method,
- mocker.Mock(side_effect=httpx.ConnectError("connection refused")),
+ mocker.Mock(
+ spec=getattr(sdk_client.TaskInstanceOperations, api_method),
+ side_effect=httpx.ConnectError("connection refused"),
+ ),
)
with pytest.raises(httpx.ConnectError):
- watched_subprocess._handle_request(msg, mocker.Mock(), req_id=1)
+ watched_subprocess._handle_request(msg, structlog.get_logger(),
req_id=1)
- assert watched_subprocess._terminal_state is None
+ assert watched_subprocess._terminal_state == expected_state
# Pending msg preserved so the recovery dispatcher can re-issue.
assert watched_subprocess._pending_terminal_state_msg is msg
+ watched_subprocess._send_heartbeat_if_needed()
+ watched_subprocess.client.task_instances.heartbeat.assert_called_once()
@pytest.mark.parametrize(
("msg", "api_method", "expected_state"),
@@ -3772,6 +3891,12 @@ class TestHandleRequest:
TaskInstanceState.UP_FOR_RESCHEDULE,
id="reschedule",
),
+ pytest.param(
+ AwaitInputTask(next_method="execute_complete"),
+ "await_input",
+ TaskInstanceState.AWAITING_INPUT,
+ id="await_input",
+ ),
],
)
def test_update_task_state_replays_pending_terminal_state_call(
@@ -3780,7 +3905,7 @@ class TestHandleRequest:
"""If a direct terminal-state API call was attempted and raised, the
recovery dispatcher must re-issue the dedicated endpoint (not
`finish()`, which the server-side endpoint refuses for SUCCESS /
- DEFERRED / SERVER_TERMINATED). Covers all four message types.
+ DEFERRED / SERVER_TERMINATED).
"""
watched_subprocess, _ = watched_subprocess
watched_subprocess._exit_code = 0
@@ -3809,6 +3934,24 @@ class TestHandleRequest:
watched_subprocess.client.task_instances.succeed.assert_not_called()
[email protected](
+ "test_name",
+ [
+ "test_worker_outcome_retained_when_direct_api_fails",
+ "test_update_task_state_replays_pending_terminal_state_call",
+ ],
+)
+def test_terminal_message_parameter_coverage(test_name):
+ marks = getattr(TestHandleRequest, test_name).pytestmark
+ cases = next(mark.args[1] for mark in marks if mark.name == "parametrize")
+ message_types =
set(get_args(get_type_hints(ActivitySubprocess._send_terminal_state_msg)["msg"]))
+ assert {type(case.values[0]) for case in cases} == message_types -
{TaskState}
+
+ marks =
TestHandleRequest.test_task_state_waits_for_exit_and_keeps_heartbeating.pytestmark
+ states = next(mark.args[1] for mark in marks if mark.name == "parametrize"
and mark.args[0] == "state")
+ assert set(states) ==
set(get_args(TaskState.model_fields["state"].annotation))
+
+
class TestSetSupervisorComms:
class DummyComms:
pass
