From ff728dc0408b45649850aa18ff94f042552c47e9 Mon Sep 17 00:00:00 2001 From: Vincent Hsiao <124506982+fat-catTW@users.noreply.github.com> Date: Mon, 6 Jul 2026 15:09:02 +0000 Subject: [PATCH] Clarify Task Instance terminal state reporting --- .../test_task_instance_operations.py | 2 +- task-sdk/src/airflow/sdk/api/client.py | 7 +++--- .../airflow/sdk/execution_time/supervisor.py | 25 +++++++++---------- task-sdk/tests/task_sdk/api/test_client.py | 6 ++--- .../execution_time/test_supervisor.py | 14 +++++------ 5 files changed, 27 insertions(+), 27 deletions(-) diff --git a/task-sdk-integration-tests/tests/task_sdk_tests/test_task_instance_operations.py b/task-sdk-integration-tests/tests/task_sdk_tests/test_task_instance_operations.py index 947a435501301..a1f0def8556ae 100644 --- a/task-sdk-integration-tests/tests/task_sdk_tests/test_task_instance_operations.py +++ b/task-sdk-integration-tests/tests/task_sdk_tests/test_task_instance_operations.py @@ -177,7 +177,7 @@ def test_ti_state_transitions(sdk_client, task_instance_id): Test task instance state transition to terminal state. """ console.print("[yellow]Testing state transition: RUNNING → FAILED...") - sdk_client.task_instances.finish( + sdk_client.task_instances.set_terminal_state( id=task_instance_id, state=TerminalStateNonSuccess.FAILED, when=utcnow(), rendered_map_index="-1" ) diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index 55668c9f3bdec..26378c0c9d934 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -265,11 +265,12 @@ def start(self, id: uuid.UUID, pid: int, when: datetime) -> TIRunContext: raise return TIRunContext.model_validate_json(resp.read()) - def finish(self, id: uuid.UUID, state: TerminalStateNonSuccess, when: datetime, rendered_map_index): - """Tell the API server that this TI has reached a terminal state.""" + def set_terminal_state( + self, id: uuid.UUID, state: TerminalStateNonSuccess, when: datetime, rendered_map_index + ): + """Tell the API server that this TI has reached a non-success terminal state.""" if state == TaskInstanceState.SUCCESS: raise ValueError("Logic error. SUCCESS state should call the `succeed` function instead") - # TODO: handle the naming better. finish sounds wrong as "even" deferred is essentially finishing. body = TITerminalStatePayload( end_date=when, state=TerminalStateNonSuccess(state), rendered_map_index=rendered_map_index ) diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 5a5e3852fd735..aee0810d3750f 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -1316,9 +1316,9 @@ class ActivitySubprocess(WatchedSubprocess): # 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. + # falling back to the generic non-success terminal-state update, which + # does not accept SUCCESS / DEFERRED / SERVER_TERMINATED. Cleared + # (and `_terminal_state` set) only after the API call returns successfully. _pending_terminal_state_msg: ( SucceedTask | RetryTask | DeferTask | RescheduleTask | AwaitInputTask | None ) = attrs.field(default=None, init=False) @@ -1446,20 +1446,19 @@ def update_task_state_if_needed(self): # 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. + # call would leave the TI stuck RUNNING on the server. The generic + # non-success terminal-state update cannot substitute because its + # payload does not accept SUCCESS / DEFERRED / SERVER_TERMINATED. if self._pending_terminal_state_msg is not None: self._replay_pending_terminal_state_msg() return - # If the process has finished a non-directly-patched state (e.g. - # FAILED, UP_FOR_RETRY without RetryTask), `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. + # If the process has finished in a non-success terminal state that was + # not sent directly, report it with `set_terminal_state()`. For states + # already in STATES_SENT_DIRECTLY whose direct API call succeeded, no + # further action is needed. if self.final_state not in STATES_SENT_DIRECTLY: - self.client.task_instances.finish( + self.client.task_instances.set_terminal_state( id=self.id, state=self.final_state, when=datetime.now(tz=timezone.utc), @@ -1695,7 +1694,7 @@ def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger, req_id: dump_opts: dict[str, bool] = {} if isinstance(msg, TaskState): # No direct API call here — the recovery path in - # `update_task_state_if_needed` will call `finish()` for + # `update_task_state_if_needed` will call `set_terminal_state()` for # non-direct states (FAILED, etc.) once the subprocess exits. self._terminal_state = msg.state self._task_end_time_monotonic = time.monotonic() diff --git a/task-sdk/tests/task_sdk/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py index 92b970f7ee51a..3cb24382615e1 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -405,8 +405,8 @@ def handle_request(request: httpx.Request) -> httpx.Response: @pytest.mark.parametrize( "state", [state for state in TerminalTIState if state != TerminalTIState.SUCCESS] ) - def test_task_instance_finish(self, state): - # Simulate a successful response from the server that finishes (moved to terminal state) a task + def test_task_instance_set_terminal_state(self, state): + # Simulate a successful response from the server that sets a non-success terminal state for a task ti_id = uuid6.uuid7() def handle_request(request: httpx.Request) -> httpx.Response: @@ -421,7 +421,7 @@ def handle_request(request: httpx.Request) -> httpx.Response: return httpx.Response(status_code=400, json={"detail": "Bad Request"}) client = make_client(transport=httpx.MockTransport(handle_request)) - client.task_instances.finish( + client.task_instances.set_terminal_state( ti_id, state=state, when="2024-10-31T12:00:00Z", rendered_map_index="test" ) 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 491a1b0e3fd95..aaf1114ae2d71 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -3233,9 +3233,9 @@ def test_update_task_state_replays_pending_terminal_state_call( self, watched_subprocess, mocker, msg, api_method, expected_state ): """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. + recovery dispatcher must re-issue the dedicated method, not the + generic non-success terminal-state update, which refuses SUCCESS / + DEFERRED / SERVER_TERMINATED. Covers all four message types. """ watched_subprocess, _ = watched_subprocess watched_subprocess._exit_code = 0 @@ -3244,9 +3244,9 @@ def test_update_task_state_replays_pending_terminal_state_call( watched_subprocess.update_task_state_if_needed() - # Recovery re-issues the dedicated endpoint, NOT finish(). + # Recovery re-issues the dedicated method, not the generic terminal-state update. getattr(watched_subprocess.client.task_instances, api_method).assert_called_once() - watched_subprocess.client.task_instances.finish.assert_not_called() + watched_subprocess.client.task_instances.set_terminal_state.assert_not_called() assert watched_subprocess._terminal_state == expected_state assert watched_subprocess._pending_terminal_state_msg is None @@ -3260,7 +3260,7 @@ def test_update_task_state_no_recovery_without_pending_msg(self, watched_subproc watched_subprocess.update_task_state_if_needed() - watched_subprocess.client.task_instances.finish.assert_not_called() + watched_subprocess.client.task_instances.set_terminal_state.assert_not_called() watched_subprocess.client.task_instances.succeed.assert_not_called() @@ -3378,7 +3378,7 @@ def execute(self, context: Context): ) # Patch the API client used by InProcessTestSupervisor to return a predictable TI context - fake_task_instances = mock.MagicMock(spec_set=["start", "finish"]) + fake_task_instances = mock.MagicMock(spec_set=["start", "set_terminal_state"]) fake_task_instances.start.return_value = make_ti_context() fake_client = mock.MagicMock(spec_set=["task_instances"]) fake_client.task_instances = fake_task_instances