This is an automated email from the ASF dual-hosted git repository.
vatsrahul1001 pushed a commit to branch v3-4-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-4-test by this push:
new 801d9a456a2 [v3-4-test] Clear language SDK task XComs in the
coordinator when a try starts (#74221) (#74936)
801d9a456a2 is described below
commit 801d9a456a2485f8ecb68de5ee25caa321c0a14f
Author: github-actions[bot]
<41898282+github-actions[bot]@users.noreply.github.com>
AuthorDate: Sun Oct 18 11:28:16 2026 +0530
[v3-4-test] Clear language SDK task XComs in the coordinator when a try
starts (#74221) (#74936)
* Clear language SDK task XComs in the coordinator when a try starts
* Keep XCom deletion in the Go runtime until backports are merged
* Use httpx2 in the new XCom clearing tests
---------
(cherry picked from commit 4b464ee29bae582a09ffda951b783fe2031de1ce)
Signed-off-by: PoAn Yang <[email protected]>
Co-authored-by: PoAn Yang <[email protected]>
---
.../execution_api/routes/task_instances.py | 10 +-
.../src/airflow/sdk/coordinators/_subprocess.py | 41 ++++-
.../src/airflow/sdk/execution_time/supervisor.py | 16 ++
.../tests/task_sdk/coordinators/test_subprocess.py | 179 ++++++++++++++++++++-
.../task_sdk/execution_time/test_supervisor.py | 68 ++++++++
5 files changed, 305 insertions(+), 9 deletions(-)
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 3e204488fd3..eb6686cfe97 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
@@ -310,10 +310,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 32af0c556b6..424b10be916 100644
--- a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py
+++ b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py
@@ -37,6 +37,7 @@ import socket
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 @@ import structlog
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 @@ if TYPE_CHECKING:
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 @@ class _PopenActivitySubprocess(ActivitySubprocess):
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 5f870d7be78..78bba4c3642 100644
--- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
+++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
@@ -190,6 +190,7 @@ if TYPE_CHECKING:
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
@@ -1719,6 +1720,10 @@ class ActivitySubprocess(WatchedSubprocess):
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
@@ -1743,6 +1748,17 @@ class ActivitySubprocess(WatchedSubprocess):
# 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 e150a0122d5..2b7615ea42c 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 @@ from __future__ import annotations
import contextlib
import json
+import logging
import os
import pathlib
import signal
@@ -27,16 +28,20 @@ import subprocess
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 @@ if not AIRFLOW_V_3_3_PLUS:
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 @@ class TestPopenActivitySubprocessStart:
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 7c3fee25cce..df7a8ba9a89 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
@@ -145,6 +145,7 @@ from airflow.sdk.execution_time.comms import (
SetTaskStateStore,
SetXCom,
SkipDownstreamTasks,
+ StartupDetails,
SucceedTask,
TaskBreadcrumbsResult,
TaskRescheduleStartDate,
@@ -1448,6 +1449,73 @@ class TestWatchedSubprocess:
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):
"""
