diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index 40da0fcf8d90b..49fab9c7f5d75 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -339,10 +339,12 @@ def ti_run( }, ) - # Send the keys to the SDK so that the client requests to clear those XComs from the server. - # The reason we cannot do this here in the server is because we need to issue a purge on custom XCom backends - # too. With the current assumption, the workers ONLY have access to the custom XCom backends directly and they - # can issue the purge. + # Send the keys to the worker, which deletes those XComs. The server cannot delete them itself, + # because for a Python task the task runner also purges them from a custom XCom backend, and only + # workers can reach that backend. For a language SDK task (Java, Go, TypeScript), the supervisor + # deletes these XComs before the runtime starts the task and does not purge them, because those + # runtimes push their XComs straight to the API server, so the XCom values never pass through the + # backend. # However, do not clear it for deferral xcom_keys = [] diff --git a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py index 32af0c556b62f..424b10be91686 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py +++ b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py @@ -37,6 +37,7 @@ import stat import subprocess import time +from datetime import UTC, datetime from pathlib import PurePosixPath from typing import TYPE_CHECKING, NoReturn, TypeVar, cast @@ -46,10 +47,11 @@ from airflow.dag_processing.bundles.base import BundleVersionLock, unpack_bundle_version # noqa: SDK002 from airflow.sdk._shared.module_loading import import_string -from airflow.sdk.api.datamodels._generated import BundleInfo +from airflow.sdk.api.datamodels._generated import BundleInfo, TaskInstanceState from airflow.sdk.configuration import conf from airflow.sdk.coordinators._dag_importer import find_claiming_importer from airflow.sdk.execution_time.bundles import initialize_ti_bundle +from airflow.sdk.execution_time.comms import RetryTask, TaskState from airflow.sdk.execution_time.coordinator import BaseCoordinator, TaskLaunchError from airflow.sdk.execution_time.schema import get_schema_version_migrator from airflow.sdk.execution_time.supervisor import ActivitySubprocess, NeverRaised, ProcessTracker @@ -63,7 +65,7 @@ from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: SDK002 from airflow.sdk.api.client import Client - from airflow.sdk.api.datamodels._generated import TaskInstance + from airflow.sdk.api.datamodels._generated import TaskInstance, TIRunContext Tracked = TypeVar("Tracked", socket.socket, subprocess.Popen) @@ -403,6 +405,41 @@ def start( # type: ignore[override] return self + def _prepare_ti_context(self, ti: TaskInstance, ti_context: TIRunContext) -> TIRunContext | None: + """ + Delete the XComs in ``ti_context.xcom_keys_to_clear`` so the launched runtime does not have to. + + Unlike ``XCom.delete``, this method does not purge the XCom backend, for two reasons. XComs that a + runtime pushes go straight to the API server and never pass through the backend. Also, a purge may + need a connection that a hook running in the supervisor process cannot look up, such as a + connection that exists only in Airflow's metadata database. So if Python code stored an XCom for + this task instance through the backend, deleting that XCom here leaves the stored object in the + backend's storage. + + If an XCom cannot be deleted, this try fails and the task does not run. + """ + for key in ti_context.xcom_keys_to_clear or (): + self.process_log.debug("Clearing XCom with key", key=key) + try: + self.client.xcoms.delete(ti.dag_id, ti.run_id, ti.task_id, key, ti.map_index) + except Exception as e: + self.process_log.exception("Failed to clear XCom; not running the task", key=key) + end_date = datetime.now(tz=UTC) + reason = f"Failed to clear XCom {key!r}: {e}" + # wait() sends this outcome to the server after the runtime process exits, the same + # way it sends an outcome that the runtime itself reported. + self._pending_terminal_state_msg = ( + RetryTask(end_date=end_date, retry_reason=reason) + if ti_context.should_retry + else TaskState(state=TaskInstanceState.FAILED, end_date=end_date, retry_reason=reason) + ) + self.kill(signal.SIGKILL) + return None + # The supervisor has already deleted the listed XComs, so it sends the runtime an empty list. + # If a runtime also deletes these XComs itself, the empty list keeps the runtime from deleting + # them a second time. + return ti_context.model_copy(update={"xcom_keys_to_clear": []}) + def wait(self) -> int: code = super().wait() self._close_unused_sockets(self._comm_server, self._logs_server) diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 960b2a01d1394..2d16870d21bdc 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -191,6 +191,7 @@ from structlog.typing import FilteringBoundLogger, WrappedLogger from airflow.executors.workloads import BundleInfo + from airflow.sdk.api.datamodels._generated import TIRunContext from airflow.sdk.bases.secrets_backend import BaseSecretsBackend from airflow.sdk.definitions.connection import Connection from airflow.sdk.types import RuntimeTaskInstanceProtocol as RuntimeTI @@ -1735,6 +1736,10 @@ def _on_child_started( return raise + if (prepared := self._prepare_ti_context(ti, ti_context)) is None: + return + ti_context = prepared + # ti_context.start_date is only populated by the server when resuming from a deferral (to preserve the # original start_date rather than using the resume time). We fall back to now() otherwise. This ensures # that `context["ti"].start_date` always reflects the *first* start time. See TIRunContext.start_date @@ -1759,6 +1764,17 @@ def _on_child_started( # Debug is fine, the process will have shown _something_ in it's last_chance exception handler log.debug("Couldn't send startup message to Subprocess - it died very early", pid=self.pid) + def _prepare_ti_context(self, ti: TaskInstance, ti_context: TIRunContext) -> TIRunContext | None: + """ + Return the context the supervisor sends to the task process, or None when the task must not run. + + An implementation that returns None must first stop the task process and record the outcome that + will be reported. This base version returns the context unchanged, because the Python task runner + itself deletes the XComs in ``ti_context.xcom_keys_to_clear`` with ``XCom.delete``, which also + purges them from the XCom backend. + """ + return ti_context + def wait(self) -> int: if self._wait_completed and self._exit_code is not None: return self._exit_code diff --git a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py index e150a0122d577..2b7615ea42c17 100644 --- a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py +++ b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py @@ -19,6 +19,7 @@ import contextlib import json +import logging import os import pathlib import signal @@ -27,16 +28,20 @@ import sys import threading import time +from datetime import UTC, datetime from unittest.mock import ANY, MagicMock, call, patch import attrs +import httpx2 +import msgspec import psutil import pytest +from task_sdk import FAKE_BUNDLE, make_client from uuid6 import uuid7 from airflow.dag_processing.bundles.base import BaseDagBundle, BundleVersion from airflow.sdk.api.client import Client, TaskInstanceOperations -from airflow.sdk.api.datamodels._generated import BundleInfo, TaskInstance +from airflow.sdk.api.datamodels._generated import BundleInfo, TaskInstance, TIRunContext from airflow.sdk.coordinators._dag_importer import CoordinatorDagImporter from airflow.sdk.coordinators._subprocess import ( SubprocessCoordinator, @@ -61,7 +66,7 @@ pytest.skip("Coordinator is only compatible with Airflow >= 3.3.0", allow_module_level=True) -def _make_ti(dag_id: str = "tutorial_dag", queue: str = "socket") -> TaskInstance: +def _make_ti(dag_id: str = "tutorial_dag", queue: str = "socket", map_index: int = -1) -> TaskInstance: return TaskInstance( id=uuid7(), dag_version_id=uuid7(), @@ -69,7 +74,7 @@ def _make_ti(dag_id: str = "tutorial_dag", queue: str = "socket") -> TaskInstanc dag_id=dag_id, run_id="run_1", try_number=1, - map_index=-1, + map_index=map_index, queue=queue, ) @@ -933,6 +938,174 @@ def test_register_pipe_readers_called_with_four_sockets(self, mock_client): assert mock_register.mock_calls == [call(ANY, ANY, ANY, ANY, data=ANY)] +# This script connects back to the supervisor the way a language SDK runtime does, then waits for +# StartupDetails. It exits with code 0 as soon as the supervisor sends anything or closes the +# connection. It exits with code 3 if 10 seconds pass and neither has happened. +_RUNTIME_AWAITING_STARTUP = """ +import socket +import sys + +def connect(flag): + host, _, port = next(arg for arg in sys.argv if arg.startswith(flag)).split("=", 1)[1].rpartition(":") + return socket.create_connection((host, int(port))) + +comm = connect("--comm=") +logs = connect("--logs=") +comm.settimeout(10) +try: + comm.recv(1) +except TimeoutError: + sys.exit(3) +""" + + +class TestPopenActivitySubprocessClearXComs: + @pytest.mark.parametrize( + ("map_index", "query"), + [ + pytest.param(-1, "", id="unmapped"), + pytest.param(0, "map_index=0", id="map-index-0"), + pytest.param(2, "map_index=2", id="map-index-2"), + ], + ) + @pytest.mark.log_level(logging.DEBUG) + @patch.object(ActivitySubprocess, "_register_pipe_readers", autospec=True) + @patch("airflow.sdk.coordinators._subprocess._accept_connections", autospec=True) + @patch("airflow.sdk.coordinators._subprocess.subprocess.Popen", autospec=True) + def test_deletes_listed_xcoms_before_the_runtime_gets_the_task( + self, mock_popen, mock_accept, mock_register, make_ti_context_dict, captured_logs, map_index, query + ): + ti = _make_ti(map_index=map_index) + ti_context = { + **make_ti_context_dict(), + "xcom_keys_to_clear": ["return_value", "skipmixin_key"], + "start_date": "2024-12-01T00:30:00Z", + "arg_bindings": [{"name": "country", "kind": "literal", "value": "uk"}], + } + events: list = [] + + def logged_keys() -> list[tuple[str, str, str]]: + return [ + (log["logger"], log["level"], log["key"]) + for log in captured_logs + if log.get("event") == "Clearing XCom with key" + ] + + def handle_request(request: httpx2.Request) -> httpx2.Response: + events.append((request.method, request.url.path, request.url.query.decode(), logged_keys())) + if request.url.path.endswith("/run"): + return httpx2.Response(200, json=ti_context) + return httpx2.Response(200, json={}) + + def record_frame(frame: bytearray) -> None: + body = msgspec.msgpack.decode(frame[4:])[1] + events.append((body["type"], TIRunContext.model_validate(body["ti_context"]))) + + comm_sock = MagicMock(spec=socket.socket) + comm_sock.sendall.side_effect = record_frame + mock_popen.return_value.pid = 12345 + mock_accept.side_effect = lambda servers, drains, proc, **kw: ( + {servers["comm"]: comm_sock, servers["logs"]: MagicMock(spec=socket.socket)}, + {soc: b"" for soc in drains.values()}, + ) + + _PopenActivitySubprocess.start( + what=ti, + dag_rel_path="bundle", + bundle_info=FAKE_BUNDLE, + client=make_client(transport=httpx2.MockTransport(handle_request)), + command=["/bin/true"], + subprocess_logs_to_stdout=False, + ) + + first = ("task", "debug", "return_value") + second = ("task", "debug", "skipmixin_key") + assert events == [ + ("PATCH", f"/task-instances/{ti.id}/run", "", []), + ("DELETE", "/xcoms/tutorial_dag/run_1/task_1/return_value", query, [first]), + ("DELETE", "/xcoms/tutorial_dag/run_1/task_1/skipmixin_key", query, [first, second]), + ("StartupDetails", TIRunContext.model_validate({**ti_context, "xcom_keys_to_clear": []})), + ] + + @pytest.mark.parametrize( + ("delete_error", "should_retry", "state", "error_text"), + [ + pytest.param("forbidden", True, "up_for_retry", "Forbidden", id="forbidden-retry"), + pytest.param("forbidden", False, "failed", "Forbidden", id="forbidden-no-retry"), + pytest.param("unreachable", True, "up_for_retry", "Connection refused", id="unreachable-retry"), + ], + ) + # The patch turns off the API client's retry policy, so the case where the API server is + # unreachable fails at once instead of after the client's retries and their waits. + @patch.object(Client, "_request_with_retry", Client._request_without_retry) + @patch.object(ActivitySubprocess, "send_msg", autospec=True) + def test_fails_the_try_without_starting_the_task_when_an_xcom_cannot_be_deleted( + self, + mock_send_msg, + make_ti_context_dict, + captured_logs, + time_machine, + delete_error, + should_retry, + state, + error_text, + ): + time_machine.move_to(datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC), tick=False) + ti = _make_ti() + ti_context = { + **make_ti_context_dict(), + "xcom_keys_to_clear": ["return_value", "skipmixin_key", "summary"], + "should_retry": should_retry, + } + requests: list[tuple[str, str]] = [] + reports: list[tuple[str, str, str]] = [] + + def handle_request(request: httpx2.Request) -> httpx2.Response: + requests.append((request.method, request.url.path)) + if request.url.path.endswith("/run"): + return httpx2.Response(200, json=ti_context) + if request.method == "DELETE": + if not request.url.path.endswith("/skipmixin_key"): + return httpx2.Response(200, json={}) + if delete_error == "unreachable": + raise httpx2.ConnectError("Connection refused", request=request) + return httpx2.Response(403, json={"detail": "Forbidden"}) + report = json.loads(request.content) + reports.append((report["state"], report["end_date"], report["retry_reason"])) + return httpx2.Response(204) + + proc = _PopenActivitySubprocess.start( + what=ti, + dag_rel_path="bundle", + bundle_info=FAKE_BUNDLE, + client=make_client(transport=httpx2.MockTransport(handle_request)), + command=[sys.executable, "-c", _RUNTIME_AWAITING_STARTUP], + subprocess_logs_to_stdout=False, + ) + + assert proc.wait() == -signal.SIGKILL + mock_send_msg.assert_not_called() + assert requests == [ + ("PATCH", f"/task-instances/{ti.id}/run"), + ("DELETE", "/xcoms/tutorial_dag/run_1/task_1/return_value"), + ("DELETE", "/xcoms/tutorial_dag/run_1/task_1/skipmixin_key"), + ("PATCH", f"/task-instances/{ti.id}/state"), + ] + assert reports == [ + (state, "2026-01-02T03:04:05Z", f"Failed to clear XCom 'skipmixin_key': {error_text}"), + ] + assert proc.final_state == state + assert { + "event": "Failed to clear XCom; not running the task", + "key": "skipmixin_key", + "level": "error", + "logger": "task", + "timestamp": ANY, + "exc_info": ANY, + "loc": ANY, + } in captured_logs + + class TestResolveArtifactBundle: """Execute-time resolution of the Dag bundle holding the artifacts.""" 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 433d6d11f3685..a80c10133d728 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -144,6 +144,7 @@ SetTaskStateStore, SetXCom, SkipDownstreamTasks, + StartupDetails, SucceedTask, TaskBreadcrumbsResult, TaskRescheduleStartDate, @@ -1447,6 +1448,73 @@ def subprocess_main(): target=subprocess_main, ) + def test_prepare_ti_context_leaves_listed_xcoms_to_python_task_runner(self, make_ti_context): + """ + For Python tasks, the supervisor leaves the listed XComs to the Python task runner. + + The task runner deletes them with XCom.delete, a call that also purges them from the XCom backend. + """ + requests = [] + + def handle_request(request: httpx2.Request) -> httpx2.Response: + requests.append((request.method, request.url.path)) + return httpx2.Response(200, json={}) + + proc = ActivitySubprocess( + process_log=MagicMock(spec=FilteringBoundLogger), + id=TI_ID, + pid=12345, + stdin=MagicMock(spec=socket.socket), + client=make_client(transport=httpx2.MockTransport(handle_request)), + process=MagicMock(spec=psutil.Process), + ) + ti_context = make_ti_context().model_copy(update={"xcom_keys_to_clear": ["return_value"]}) + ti = TaskInstance( + id=TI_ID, + task_id="b", + dag_id="c", + run_id="d", + try_number=1, + dag_version_id=uuid7(), + queue="default", + ) + + try: + prepared = proc._prepare_ti_context(ti, ti_context) + finally: + proc.selector.close() + assert requests == [] + assert prepared is not None + assert prepared.xcom_keys_to_clear == ["return_value"] + + def test_python_task_process_gets_the_listed_xcom_keys(self, mocker, make_ti_context): + mock_client = MagicMock(spec=sdk_client.Client) + mock_client.task_instances.start.return_value = make_ti_context().model_copy( + update={"xcom_keys_to_clear": ["return_value"]} + ) + mock_send = mocker.patch.object(ActivitySubprocess, "send_msg", autospec=True) + + proc = ActivitySubprocess.start( + dag_rel_path=os.devnull, + bundle_info=FAKE_BUNDLE, + what=TaskInstance( + id=uuid7(), + task_id="b", + dag_id="c", + run_id="d", + try_number=1, + dag_version_id=uuid7(), + queue="default", + ), + client=mock_client, + target=lambda: None, + ) + assert proc.wait() == 0 + + (startup,) = [c.args[1] for c in mock_send.call_args_list if isinstance(c.args[1], StartupDetails)] + assert startup.ti_context.xcom_keys_to_clear == ["return_value"] + mock_client.xcoms.delete.assert_not_called() + @pytest.mark.parametrize("captured_logs", [logging.WARNING], indirect=True) def test_heartbeat_failures_handling(self, monkeypatch, mocker, captured_logs, time_machine): """