From f3dc7612da7bb32bb4f423ff3c4dc06c4ba47b19 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Tue, 4 Aug 2026 13:09:16 +0000 Subject: [PATCH 1/2] Extract supervisor heartbeat logic into a standalone Heartbeater MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A follow-up will let subprocess coordinators heartbeat while materializing a Dag bundle, before the task subprocess exists. That requires the heartbeat state to live outside ActivitySubprocess — especially the pid presented to the server, which must stay identical for the task instance's lifetime or the server rejects the heartbeat as "running elsewhere" and the task is killed. Reactions to fatal heartbeat outcomes (killing the process, recording SERVER_TERMINATED) remain in ActivitySubprocess, injected as callbacks. No behavior change. --- .../airflow/sdk/execution_time/supervisor.py | 204 ++++++++++++------ .../execution_time/test_supervisor.py | 8 +- 2 files changed, 141 insertions(+), 71 deletions(-) diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 9758d2b18d7c8..bbe154af44adf 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -175,7 +175,7 @@ from airflow.sdk.definitions.connection import Connection from airflow.sdk.types import RuntimeTaskInstanceProtocol as RuntimeTI -__all__ = ["ActivitySubprocess", "WatchedSubprocess", "supervise", "supervise_task"] +__all__ = ["ActivitySubprocess", "Heartbeater", "WatchedSubprocess", "supervise", "supervise_task"] log: FilteringBoundLogger = structlog.get_logger(logger_name="supervisor") @@ -1358,6 +1358,115 @@ def _remote_logging_conn(client: Client): del client +@attrs.define(kw_only=True) +class Heartbeater: + """ + Send periodic task-instance heartbeats and track their success/failure state. + + Standalone from :class:`ActivitySubprocess` so a caller can heartbeat outside a + subprocess's lifetime (e.g. a coordinator materializing a Dag bundle before the + task subprocess exists). The server rejects a heartbeat whose pid differs from + the one registered at task start as "running elsewhere", so the pid is fixed at + construction for the task instance's lifetime rather than read from a process + handle. Reactions to fatal heartbeat outcomes (killing the process, recording a + terminal state) are injected as callbacks. + """ + + client: Client + """The HTTP client to use for communication with the API server.""" + + ti_id: UUID + """ID of the task instance to heartbeat.""" + + pid: int + """The pid registered with the server at task start; presented on every heartbeat.""" + + on_server_terminated: Callable[[Any], None] + """Called with the response detail when the server says the task should no longer run.""" + + on_fatal_failures: Callable[[], None] | None = None + """Called when consecutive heartbeat failures reach ``MAX_FAILED_HEARTBEATS``. + + ``None`` disables the cap: failures are logged and retried indefinitely, for + callers that have no process to kill (e.g. heartbeating before the task + subprocess exists). + """ + + _last_successful_heartbeat: float = attrs.field(default=0, init=False) + _last_heartbeat_attempt: float = attrs.field(default=0, init=False) + + # After the failure of a heartbeat, we'll increment this counter. If it reaches `MAX_FAILED_HEARTBEATS`, we + # will kill the process. This is to handle temporary network issues etc. ensuring that the process + # does not hang around forever. + failed_heartbeats: int = attrs.field(default=0, init=False) + + def record_successful_heartbeat(self) -> None: + """Record an out-of-band beat the server already counted (e.g. the task-start API call).""" + self._last_successful_heartbeat = time.monotonic() + + def compute_max_wait_time(self) -> float: + """How long the monitor loop may block in ``select`` before the next heartbeat is due.""" + last_heartbeat_ago = time.monotonic() - self._last_successful_heartbeat + return max( + 0, # Make sure this value is never negative, + min( + # Ensure we heartbeat _at most_ 75% through the task instance heartbeat timeout time + HEARTBEAT_TIMEOUT - last_heartbeat_ago * 0.75, + MIN_HEARTBEAT_INTERVAL, + ), + ) + + def send_heartbeat_if_needed(self) -> None: + """Send a heartbeat to the client if heartbeat interval has passed.""" + # Respect the minimum interval between heartbeat attempts + if (time.monotonic() - self._last_heartbeat_attempt) < MIN_HEARTBEAT_INTERVAL: + return + + self.send_heartbeat() + + def send_heartbeat(self) -> None: + """Send a heartbeat unconditionally; pacing is the caller's responsibility.""" + self._last_heartbeat_attempt = time.monotonic() + try: + self.client.task_instances.heartbeat(self.ti_id, pid=self.pid) + # Update the last heartbeat time on success + self._last_successful_heartbeat = time.monotonic() + + # Reset the counter on success + self.failed_heartbeats = 0 + except ServerResponseError as e: + if e.response.status_code in {HTTPStatus.NOT_FOUND, HTTPStatus.GONE, HTTPStatus.CONFLICT}: + log.error( + "Server indicated the task shouldn't be running anymore", + detail=e.detail, + status_code=e.response.status_code, + ti_id=self.ti_id, + ) + self.on_server_terminated(e.detail) + else: + # If we get any other error, we'll just log it and try again next time + self._handle_heartbeat_failures(e) + except Exception as e: + self._handle_heartbeat_failures(e) + + def _handle_heartbeat_failures(self, exc: Exception) -> None: + """Increment the failed heartbeats counter and kill the process if too many failures.""" + self.failed_heartbeats += 1 + log.warning( + "Failed to send heartbeat. Will be retried", + failed_heartbeats=self.failed_heartbeats, + ti_id=self.ti_id, + max_retries=MAX_FAILED_HEARTBEATS, + exc_info=exc, + ) + # If we've failed to heartbeat too many times, kill the process + if self.on_fatal_failures is not None and self.failed_heartbeats >= MAX_FAILED_HEARTBEATS: + log.error( + "Too many failed heartbeats; terminating process", failed_heartbeats=self.failed_heartbeats + ) + self.on_fatal_failures() + + @attrs.define(kw_only=True) class ActivitySubprocess(WatchedSubprocess): client: Client @@ -1378,17 +1487,12 @@ class ActivitySubprocess(WatchedSubprocess): SucceedTask | RetryTask | DeferTask | RescheduleTask | AwaitInputTask | None ) = attrs.field(default=None, init=False) - _last_successful_heartbeat: float = attrs.field(default=0, init=False) - _last_heartbeat_attempt: float = attrs.field(default=0, init=False) + heartbeater: Heartbeater = attrs.field(init=False) + """Sends periodic heartbeats for this task instance while the subprocess runs.""" _should_retry: bool = attrs.field(default=False, init=False) """Whether the task should retry or not as decided by the API server.""" - # After the failure of a heartbeat, we'll increment this counter. If it reaches `MAX_FAILED_HEARTBEATS`, we - # will kill theprocess. This is to handle temporary network issues etc. ensuring that the process - # does not hang around forever. - failed_heartbeats: int = attrs.field(default=0, init=False) - _task_end_time_monotonic: float | None = attrs.field(default=None, init=False) _rendered_map_index: str | None = attrs.field(default=None, init=False) @@ -1396,6 +1500,27 @@ class ActivitySubprocess(WatchedSubprocess): ti: RuntimeTI | None = None + def __attrs_post_init__(self) -> None: + self.heartbeater = Heartbeater( + client=self.client, + ti_id=self.id, + pid=self._process.pid, + on_server_terminated=self._on_heartbeat_server_terminated, + on_fatal_failures=self._kill_on_heartbeat_failure, + ) + + def _on_heartbeat_server_terminated(self, detail: Any) -> None: + self.process_log.error( + "Server indicated the task shouldn't be running anymore. Terminating process", + detail=detail, + ) + self.kill(signal.SIGTERM, force=True) + self.process_log.error("Task killed!") + self._terminal_state = SERVER_TERMINATED + + def _kill_on_heartbeat_failure(self) -> None: + self.kill(signal.SIGTERM, force=True) + @classmethod def start( # type: ignore[override] cls, @@ -1448,7 +1573,8 @@ def _on_child_started( # tell us "no, stop!" for any reason) ti_context = self.client.task_instances.start(ti.id, self.pid, datetime.now(tz=timezone.utc)) self._should_retry = ti_context.should_retry - self._last_successful_heartbeat = time.monotonic() + # The start call above updated last_heartbeat_at on the server. + self.heartbeater.record_successful_heartbeat() except Exception: # On any error kill that subprocess! self.kill(signal.SIGKILL) @@ -1610,17 +1736,9 @@ def _monitor_subprocess(self): - Sends heartbeats to ensure the process is alive and checks if the subprocess has exited. """ while self._exit_code is None or self._open_sockets: - last_heartbeat_ago = time.monotonic() - self._last_successful_heartbeat # Monitor the task to see if it's done. Wait in a syscall (`select`) for as long as possible # so we notice the subprocess finishing as quick as we can. - max_wait_time = max( - 0, # Make sure this value is never negative, - min( - # Ensure we heartbeat _at most_ 75% through the task instance heartbeat timeout time - HEARTBEAT_TIMEOUT - last_heartbeat_ago * 0.75, - MIN_HEARTBEAT_INTERVAL, - ), - ) + max_wait_time = self.heartbeater.compute_max_wait_time() # Block until events are ready or the timeout is reached # This listens for activity (e.g., subprocess output) on registered file objects alive = self._service_subprocess(max_wait_time=max_wait_time) is None @@ -1666,60 +1784,12 @@ def _handle_process_overtime_if_needed(self): def _send_heartbeat_if_needed(self): """Send a heartbeat to the client if heartbeat interval has passed.""" - # Respect the minimum interval between heartbeat attempts - 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 return - self._last_heartbeat_attempt = time.monotonic() - try: - self.client.task_instances.heartbeat(self.id, pid=self._process.pid) - # Update the last heartbeat time on success - self._last_successful_heartbeat = time.monotonic() - - # Reset the counter on success - self.failed_heartbeats = 0 - except ServerResponseError as e: - if e.response.status_code in {HTTPStatus.NOT_FOUND, HTTPStatus.GONE, HTTPStatus.CONFLICT}: - log.error( - "Server indicated the task shouldn't be running anymore", - detail=e.detail, - status_code=e.response.status_code, - ti_id=self.id, - ) - self.process_log.error( - "Server indicated the task shouldn't be running anymore. Terminating process", - detail=e.detail, - ) - 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) - except Exception as e: - self._handle_heartbeat_failures(e) - - def _handle_heartbeat_failures(self, exc: Exception): - """Increment the failed heartbeats counter and kill the process if too many failures.""" - self.failed_heartbeats += 1 - log.warning( - "Failed to send heartbeat. Will be retried", - failed_heartbeats=self.failed_heartbeats, - ti_id=self.id, - max_retries=MAX_FAILED_HEARTBEATS, - exc_info=exc, - ) - # If we've failed to heartbeat too many times, kill the process - if self.failed_heartbeats >= MAX_FAILED_HEARTBEATS: - log.error( - "Too many failed heartbeats; terminating process", failed_heartbeats=self.failed_heartbeats - ) - self.kill(signal.SIGTERM, force=True) + self.heartbeater.send_heartbeat_if_needed() @property def final_state(self): 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 909f4942ef8ab..1a86b00783315 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -1041,7 +1041,7 @@ def mock_monotonic(): # Simulate sending heartbeats and ensure the process gets killed after max retries for i in range(1, max_failed_heartbeats): proc._send_heartbeat_if_needed() - assert proc.failed_heartbeats == i # Increment happens after failure + assert proc.heartbeater.failed_heartbeats == i # Increment happens after failure mock_client_heartbeat.assert_called_with(TI_ID, pid=mock_process.pid) # Ensure the retry log is present @@ -1066,7 +1066,7 @@ def mock_monotonic(): # On the final failure, the process should be killed proc._send_heartbeat_if_needed() - assert proc.failed_heartbeats == max_failed_heartbeats + assert proc.heartbeater.failed_heartbeats == max_failed_heartbeats mock_kill.assert_called_once_with(signal.SIGTERM, force=True) mock_client_heartbeat.assert_called_with(TI_ID, pid=mock_process.pid) assert { @@ -1603,7 +1603,7 @@ def test_max_wait_time_prevents_cpu_spike(self, watched_subprocess, mock_process # Set up a scenario where the last successful heartbeat was a long time ago # This will cause the heartbeat calculation to result in a negative value - mock_process._last_successful_heartbeat = time.time() - 100 # 100 seconds ago + watched_subprocess.heartbeater._last_successful_heartbeat = time.time() - 100 # 100 seconds ago # Mock process to still be alive (not exited) mock_process.wait.side_effect = psutil.TimeoutExpired(pid=12345, seconds=0) @@ -1646,7 +1646,7 @@ def test_max_wait_time_calculation_edge_cases( monkeypatch.setattr("airflow.sdk.execution_time.supervisor.HEARTBEAT_TIMEOUT", heartbeat_timeout) monkeypatch.setattr("airflow.sdk.execution_time.supervisor.MIN_HEARTBEAT_INTERVAL", min_interval) - watched_subprocess._last_successful_heartbeat = time.time() - heartbeat_ago + watched_subprocess.heartbeater._last_successful_heartbeat = time.time() - heartbeat_ago mock_process.wait.side_effect = psutil.TimeoutExpired(pid=12345, seconds=0) # Call the method and verify timeout is never less than our minimum From 03368a2d45ba11483c5ed96ba1535a36ad72dad5 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Tue, 4 Aug 2026 06:25:54 +0000 Subject: [PATCH 2/2] Report language-SDK tasks running while the runtime starts up The coordinator path left the task queued until the language runtime had connected back to the supervisor. Locating artifacts and waiting out task_startup_timeout were therefore charged to [scheduler] task_queued_timeout, whose handler revokes and requeues the task, and none of it was visible in the UI. A runtime that never started looked like a task that had never been picked up at all, and the stdout and stderr it produced while failing were discarded instead of written to the task log -- for a missing main class or a runtime that is not installed, that output is the whole explanation. The Python path reports the forked child's pid before handing it any work and materializes its Dag bundle inside that window, so this brings the two paths in line. The pid reported here is the supervisor's own because the server rejects a heartbeat whose pid differs from the one it was told at the start, and that is what lets the transition happen before the runtime exists. --- .../language-sdks/go.rst | 8 +- .../language-sdks/java.rst | 7 +- .../language-sdks/typescript.rst | 7 +- .../airflow/sdk/coordinators/_subprocess.py | 253 ++++++++++-- .../airflow/sdk/execution_time/supervisor.py | 71 +++- .../task_sdk/coordinators/test_subprocess.py | 384 +++++++++++++++++- 6 files changed, 672 insertions(+), 58 deletions(-) diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst index 158fff4ddeb2f..833a9247feee4 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst @@ -430,7 +430,9 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the * - ``task_startup_timeout`` - ``10.0`` - Seconds to wait for the bundle subprocess to connect after launch. Increase this if your - bundle startup is slow (e.g. on constrained hardware). + bundle startup is slow (e.g. on constrained hardware). The task is already ``running`` + while the coordinator waits, so exceeding this fails the task rather than leaving it + ``queued``, and whatever the bundle printed goes to the task log. .. _go-sdk/edge-worker: @@ -457,6 +459,10 @@ Limitations languages, so task names and dependencies are declared in Python with :func:`@task.stub `. This applies to both deployment modes and is a documented known limitation. +* **In coordinator mode the pid recorded on the task instance is the Airflow supervisor's, not the bundle's.** + The supervisor reports the task as running before it launches the bundle, and the server ties the run to the + pid it was given at that point, so that is the pid the API and UI show. The bundle's own pid goes to the task + log. The following are a non-exhaustive list of features the **Edge Worker** path has yet to implement. They are the main reason the coordinator path is recommended: in coordinator mode the Python supervisor handles these diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst index c8b73b4626a28..b5e705e7a8c00 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst @@ -681,7 +681,9 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the * - ``task_startup_timeout`` - ``10.0`` - Seconds to wait for the JVM subprocess to connect after launch. Increase this if your - JVM startup is slow (e.g. on constrained hardware or with a large classpath). + JVM startup is slow (e.g. on constrained hardware or with a large classpath). The task + is already ``running`` while the coordinator waits, so exceeding this fails the task + rather than leaving it ``queued``, and whatever the JVM printed goes to the task log. .. note:: @@ -724,5 +726,8 @@ Limitations * **One JVM subprocess per task instance.** Each task instance spawns a fresh JVM. Tasks that need to share in-process state between instances should use XCom or an external store instead. +* **The pid recorded on the task instance is the Airflow supervisor's, not the JVM's.** The supervisor reports + the task as running before it launches the JVM, and the server ties the run to the pid it was given at that + point, so that is the pid the API and UI show. The JVM's own pid is logged to the task log when it starts. * **Limited support for assets, deferral, and other Airflow features.** They may be implemented in the future based on user feedback and demand. diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst index 87e12cea474f4..a1e219bac0a96 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst @@ -278,7 +278,9 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the * - ``task_startup_timeout`` - ``10.0`` - Seconds to wait for the Node.js subprocess to connect after launch. Increase this if your bundle - startup is slow (e.g. on constrained hardware). + startup is slow (e.g. on constrained hardware). The task is already ``running`` while the + coordinator waits, so exceeding this fails the task rather than leaving it ``queued``, and + whatever the bundle printed goes to the task log. Limitations ----------- @@ -292,3 +294,6 @@ Limitations bundles. To serve multiple bundles, register multiple coordinators on separate queues. * **One Node.js subprocess per task instance.** Tasks that need to share in-process state between instances should use XCom or an external store instead. +* **The pid recorded on the task instance is the Airflow supervisor's, not the Node.js process's.** The + supervisor reports the task as running before it launches the bundle, and the server ties the run to the pid + it was given at that point, so that is the pid the API and UI show. The bundle's own pid goes to the task log. diff --git a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py index a0c3f518fb068..e3e0c682014c8 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py +++ b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py @@ -27,6 +27,7 @@ from __future__ import annotations +import contextlib import ipaddress import itertools import os @@ -34,31 +35,106 @@ import signal import socket import subprocess +import threading import time -from typing import TYPE_CHECKING, TypeVar, cast +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Any, NoReturn, TypeVar, cast import attrs import psutil import structlog +from airflow.sdk.api.datamodels._generated import TaskInstanceState from airflow.sdk.configuration import conf -from airflow.sdk.execution_time.coordinator import BaseCoordinator -from airflow.sdk.execution_time.supervisor import ActivitySubprocess, NeverRaised, ProcessTracker +from airflow.sdk.execution_time.coordinator import BaseCoordinator, _warm_shutdown_signals +from airflow.sdk.execution_time.supervisor import ( + MIN_HEARTBEAT_INTERVAL, + ActivitySubprocess, + Heartbeater, + NeverRaised, + ProcessTracker, +) if TYPE_CHECKING: - from collections.abc import Sequence + import uuid + from collections.abc import Generator, Sequence from structlog.typing import FilteringBoundLogger from typing_extensions import Self from airflow.sdk.api.client import Client - from airflow.sdk.api.datamodels._generated import BundleInfo, TaskInstance + from airflow.sdk.api.datamodels._generated import BundleInfo, TaskInstance, TIRunContext Tracked = TypeVar("Tracked", socket.socket, subprocess.Popen) log: FilteringBoundLogger = structlog.get_logger(logger_name="coordinators.subprocess") +class SubprocessStartupError(RuntimeError): + """ + The launched runtime never completed the startup handshake with the supervisor. + + :param exit_code: Exit code to report for the run — the runtime's own when it + managed to exit before connecting, otherwise 1. + """ + + def __init__(self, reason: str, *, exit_code: int = 1) -> None: + super().__init__(reason) + self.exit_code = exit_code + + +@contextlib.contextmanager +def _heartbeat_until_monitored( + client: Client, ti_id: uuid.UUID, pid: int, logger: FilteringBoundLogger +) -> Generator[None, None, None]: + """ + Keep the run's heartbeat fresh while the worker is getting the runtime up. + + The task is RUNNING from before the runtime exists, but nothing heartbeats until + :meth:`ActivitySubprocess.wait` starts monitoring, so a slow launch — a first-time + Dag bundle clone, a generous ``task_startup_timeout`` — would look like a zombie to + ``[scheduler] task_instance_heartbeat_timeout`` and be reaped mid-startup. The Python + path gets this for free: its supervisor is already monitoring while the child does + the same work. + + :meth:`ActivitySubprocess._monitor_subprocess` cannot cover this window — it services + and reaps a child that does not exist yet — so a dedicated thread paces a + :class:`~airflow.sdk.execution_time.supervisor.Heartbeater` instead, sharing the + send and stop policy with the monitor loop's own heartbeats. + + Nothing here aborts the launch. A disowned run stops beating and is terminated by the + monitor loop's own heartbeat moments later, which is the code that owns killing the + runtime and recording SERVER_TERMINATED; with no process to kill, + ``on_fatal_failures`` is left unset so a transient failure is simply retried. + """ + stop = threading.Event() + + def _stop_beating(detail: Any) -> None: + logger.error( + "Server disowned this run while the runtime was starting; the monitor loop will terminate it", + detail=detail, + ) + stop.set() + + heartbeater = Heartbeater(client=client, ti_id=ti_id, pid=pid, on_server_terminated=_stop_beating) + + def _beat() -> None: + # The blocking wait IS the pacing: exactly one attempt per interval, success or + # failure, with the sleep doubling as the stop signal. The monitor loop's + # if-needed gate and shrinking wait formula compensate for IO-driven select + # wake-ups, which a dedicated timer thread doesn't have. + while not stop.wait(MIN_HEARTBEAT_INTERVAL): + heartbeater.send_heartbeat() + + thread = threading.Thread(target=_beat, name=f"startup-heartbeat-{ti_id}", daemon=True) + thread.start() + try: + yield + finally: + stop.set() + thread.join(MIN_HEARTBEAT_INTERVAL) + + def _start_server() -> socket.socket: server = socket.socket() server.bind(("127.0.0.1", 0)) @@ -161,10 +237,28 @@ def _accept_connections( *, max_wait: float = 10.0, drain_size: int = 4096, + logger: FilteringBoundLogger | None = None, ) -> tuple[dict[socket.socket, socket.socket], dict[socket.socket, bytes]]: """Block until the subprocess connects to servers, draining stdout/stderr along the way.""" + task_log = logger or log accepted: dict[socket.socket, socket.socket] = {} drained: dict[socket.socket, bytes] = {s: b"" for s in drains.values()} + + def _give_up(reason: str, *, exit_code: int = 1) -> NoReturn: + for s in accepted.values(): + s.close() + # On the happy path these bytes reach the task log through + # _register_pipe_readers. Emit them here too, or the runtime's own account of + # why it never started dies with the drain buffers. + for key, soc in drains.items(): + if output := drained[soc]: + task_log.error( + "Runtime output before startup failure", + key=key, + output=output.decode(errors="replace"), + ) + raise SubprocessStartupError(reason, exit_code=exit_code) + with selectors.DefaultSelector() as sel: for key, soc in itertools.chain(servers.items(), drains.items()): sel.register(soc, selectors.EVENT_READ, data=key) @@ -172,27 +266,26 @@ def _accept_connections( while len(accepted) < len(servers): remaining = deadline - time.monotonic() if remaining <= 0: - for s in accepted.values(): - s.close() - raise TimeoutError("process did not connect within timeout") + _give_up("process did not connect within timeout") if proc.poll() is not None: - for s in accepted.values(): - s.close() - raise RuntimeError(f"process exited with {proc.returncode} before connecting") + _give_up( + f"process exited with {proc.returncode} before connecting", + exit_code=proc.returncode or 1, + ) for event, _ in sel.select(timeout=min(remaining, 1.0)): soc = cast("socket.socket", event.fileobj) if soc in drained: if incoming := soc.recv(drain_size): - log.debug("Draining child process stream", key=event.data) + task_log.debug("Draining child process stream", key=event.data) drained[soc] += incoming else: - log.warning("Child stream closed before ready!", key=event.data) + task_log.warning("Child stream closed before ready!", key=event.data) sel.unregister(soc) else: - log.debug("Accepting child process connection", key=event.data) + task_log.debug("Accepting child process connection", key=event.data) conn, _ = soc.accept() if not _is_connection_from_process(conn, proc): - log.warning( + task_log.warning( "Rejected connection not owned by child process", key=event.data, pid=proc.pid, @@ -292,6 +385,8 @@ def start( # type: ignore[override] what: TaskInstance, dag_rel_path: str | os.PathLike[str], bundle_info, + ti_context: TIRunContext, + running_since: datetime, logger: FilteringBoundLogger | None = None, sentry_integration: str = "", command: Sequence[str], @@ -299,6 +394,7 @@ def start( # type: ignore[override] startup_timeout: float = 10.0, **kwargs, ) -> Self: + task_log = logger or structlog.get_logger(logger_name="task").bind() with _ResourceTracker(timeout=startup_timeout) as tracker: comm_server, logs_server = tracker.track(_start_server(), _start_server()) stdout_r, stdout_w = tracker.track(*socket.socketpair()) @@ -326,13 +422,14 @@ def start( # type: ignore[override] tracker.track(proc) for soc in tracker.untrack(stdout_w, stderr_w): soc.close() - log.info("Starting subprocess", pid=proc.pid) + task_log.info("Starting subprocess", pid=proc.pid) socks, drained = _accept_connections( {"comm": comm_server, "logs": logs_server}, {"stdout": stdout_r, "stderr": stderr_r}, proc, max_wait=startup_timeout, + logger=task_log, ) tracker.track(*socks.values()) @@ -340,7 +437,7 @@ def start( # type: ignore[override] id=what.id, pid=proc.pid, process=PopenTracker(proc), - process_log=logger or structlog.get_logger(logger_name="task").bind(), + process_log=task_log, start_time=time.monotonic(), stdin=socks[comm_server], subprocess_schema_version=subprocess_schema_version, @@ -357,6 +454,8 @@ def start( # type: ignore[override] dag_rel_path=dag_rel_path, bundle_info=bundle_info, sentry_integration=sentry_integration, + ti_context=ti_context, + running_since=running_since, ) # Untrack everything left. 'self' keeps track of these and closes @@ -382,9 +481,15 @@ class SubprocessCoordinator(BaseCoordinator): connections, draining startup output, and tearing everything down on failure — is handled here. + The task is reported RUNNING before :meth:`_build_execute_task_command` runs, + so everything a subclass does to locate its artifacts — walking a directory, + materializing a Dag bundle — is charged to the task's runtime rather than to + its queued time, and failing it fails the task instead of leaving the run + QUEUED for ``[scheduler] task_queued_timeout`` to pick up. + :param task_startup_timeout: Maximum time the coordinator waits for the subprocess to connect to both servers, in seconds. The default is 10 - seconds. + seconds. The wait happens with the task already RUNNING. """ task_startup_timeout: float = 10.0 @@ -414,18 +519,100 @@ def execute_task( subprocess_logs_to_stdout: bool, **kwargs, ) -> BaseCoordinator.ExecutionResult: - command, subprocess_schema_version = self._build_execute_task_command(what=what) - process = _PopenActivitySubprocess.start( - what=what, - dag_rel_path=dag_rel_path, - bundle_info=bundle_info, - client=client, - logger=logger, - subprocess_logs_to_stdout=subprocess_logs_to_stdout, - sentry_integration=sentry_integration, - command=command, - subprocess_schema_version=subprocess_schema_version, - startup_timeout=self.task_startup_timeout, - ) - exit_code = process.wait() - return self.ExecutionResult(exit_code, process.final_state) + task_log = logger or structlog.get_logger(logger_name="task").bind() + # Hold the warm-shutdown handlers across the RUNNING window, as the Python + # coordinator does, so a SIGTERM here cannot orphan a task the server has + # been told is running. + with _warm_shutdown_signals(): + # Report RUNNING before any preparation work (see the class docstring), so + # `_build_execute_task_command` runs on the task's clock rather than the + # queue's. A failure here is not ours to report: there is no run in RUNNING + # yet, and the redelivery it usually means (TaskAlreadyRunningError) is for + # the executor to swallow. + # The runtime does not exist yet, so this supervisor's own pid becomes the + # run's ownership token; the runtime's pid goes to the task log instead + # ("Starting subprocess"). Every later heartbeat must present the same + # value, so it is threaded into the supervised process too (see + # ActivitySubprocess.reported_pid). + reported_pid = os.getpid() + running_since = datetime.now(tz=timezone.utc) + ti_context = client.task_instances.start(what.id, reported_pid, running_since) + try: + with _heartbeat_until_monitored(client, what.id, reported_pid, task_log): + command, subprocess_schema_version = self._build_execute_task_command(what=what) + process = _PopenActivitySubprocess.start( + what=what, + dag_rel_path=dag_rel_path, + bundle_info=bundle_info, + client=client, + ti_context=ti_context, + running_since=running_since, + reported_pid=reported_pid, + logger=task_log, + subprocess_logs_to_stdout=subprocess_logs_to_stdout, + sentry_integration=sentry_integration, + command=command, + subprocess_schema_version=subprocess_schema_version, + startup_timeout=self.task_startup_timeout, + ) + except SubprocessStartupError as error: + task_log.error("Task runtime failed to start", reason=str(error)) + return self._finish_failed_startup( + client=client, + what=what, + ti_context=ti_context, + logger=task_log, + exit_code=error.exit_code, + ) + except Exception: + task_log.exception("Failed to launch the task runtime") + return self._finish_failed_startup( + client=client, what=what, ti_context=ti_context, logger=task_log + ) + exit_code = process.wait() + return self.ExecutionResult(exit_code, process.final_state) + + def _finish_failed_startup( + self, + *, + client: Client, + what: TaskInstance, + ti_context: TIRunContext, + logger: FilteringBoundLogger, + exit_code: int = 1, + ) -> BaseCoordinator.ExecutionResult: + """ + Report the terminal state for a run that never got a runtime to monitor. + + :class:`ActivitySubprocess` only reports a terminal state out of ``wait()``, + which needs the very process this failure prevented, so without this the run + would sit RUNNING until zombie detection reaped it. + + :raises Exception: whatever the terminal-state call raised. The run is RUNNING + on the server and this is the only thing that was going to move it, so the + failure has to escape and let the executor report it for the scheduler to + reconcile — returning a state we did not manage to record would put the run + right back in the RUNNING-until-reaped hole this method exists to close. + """ + when = datetime.now(tz=timezone.utc) + state = TaskInstanceState.UP_FOR_RETRY if ti_context.should_retry else TaskInstanceState.FAILED + try: + if state is TaskInstanceState.UP_FOR_RETRY: + # UP_FOR_RETRY is not a TerminalStateNonSuccess, so it has its own endpoint. + client.task_instances.retry( + id=what.id, + end_date=when, + rendered_map_index=None, + retry_reason="Task runtime failed to start", + ) + else: + client.task_instances.finish( + id=what.id, + state=state, + when=when, + rendered_map_index=None, + ) + except Exception: + logger.exception("Failed to report the task runtime startup failure", state=state) + raise + return self.ExecutionResult(exit_code, state) diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index bbe154af44adf..e543e2e9965bc 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -56,6 +56,7 @@ ConnectionResponse, TaskInstance, TaskInstanceState, + TIRunContext, ) from airflow.sdk.configuration import conf from airflow.sdk.exceptions import ErrorType @@ -206,6 +207,7 @@ } ) + # Setting a fair buffer size here to handle most message sizes. Intention is to enforce a buffer size # that is big enough to handle small to medium messages while not enforcing hard latency issues BUFFER_SIZE = 4096 @@ -1500,11 +1502,22 @@ class ActivitySubprocess(WatchedSubprocess): ti: RuntimeTI | None = None + reported_pid: int | None = None + """ + Pid to report to the server in place of the supervised process's own. + + The server treats ``(hostname, pid)`` as the run's ownership token: it rejects a + heartbeat whose pid differs from the one ``task_instances.start`` reported, so the + RUNNING transition and every heartbeat read the one value seeded into + :attr:`Heartbeater.pid` from here. A coordinator that reports RUNNING before the + runtime exists passes its own pid; ``None`` uses the supervised process's. + """ + def __attrs_post_init__(self) -> None: self.heartbeater = Heartbeater( client=self.client, ti_id=self.id, - pid=self._process.pid, + pid=self.reported_pid if self.reported_pid is not None else self._process.pid, on_server_terminated=self._on_heartbeat_server_terminated, on_fatal_failures=self._kill_on_heartbeat_failure, ) @@ -1548,15 +1561,31 @@ def start( # type: ignore[override] new_process_group=True, **kwargs, ) + # We've forked, but the task won't start doing anything until we send it the StartupDetails + # message. But before we do that, we need to tell the server it's started (so it has the chance to + # tell us "no, stop!" for any reason) + running_since = datetime.now(tz=timezone.utc) + ti_context = proc._report_running(ti=what, when=running_since) # Tell the task process what it needs to do! proc._on_child_started( ti=what, dag_rel_path=dag_rel_path, bundle_info=bundle_info, sentry_integration=sentry_integration, + ti_context=ti_context, + running_since=running_since, ) return proc + def _report_running(self, *, ti: TaskInstance, when: datetime) -> TIRunContext: + """Transition the task instance to RUNNING, killing the child if the server refuses.""" + try: + return self.client.task_instances.start(ti.id, self.heartbeater.pid, when) + except Exception: + # On any error kill that subprocess! + self.kill(signal.SIGKILL) + raise + def _on_child_started( self, *, @@ -1564,27 +1593,33 @@ def _on_child_started( dag_rel_path: str | os.PathLike[str], bundle_info, sentry_integration: str, + ti_context: TIRunContext, + running_since: datetime, ) -> None: - """Send startup message to the subprocess.""" + """ + Send startup message to the subprocess. + + :param ti_context: Run context from the ``task_instances.start`` call that + already put the task in RUNNING — :meth:`_report_running` here, or an + earlier call by a coordinator that has to report RUNNING before it can + get this far. Re-reporting it would be rejected with a 409. + :param running_since: The timestamp that same call reported, so the task sees + the start date the server recorded. Reading the clock again here would + drift by however long the caller took to get the child talking, which for + a coordinator is a whole artifact resolution. + """ self.ti = ti # type: ignore[assignment] - try: - # We've forked, but the task won't start doing anything until we send it the StartupDetails - # message. But before we do that, we need to tell the server it's started (so it has the chance to - # tell us "no, stop!" for any reason) - ti_context = self.client.task_instances.start(ti.id, self.pid, datetime.now(tz=timezone.utc)) - self._should_retry = ti_context.should_retry - # The start call above updated last_heartbeat_at on the server. - self.heartbeater.record_successful_heartbeat() - except Exception: - # On any error kill that subprocess! - self.kill(signal.SIGKILL) - raise + # should_retry is optional in the schema; absent means "no retries left". + self._should_retry = bool(ti_context.should_retry) + # The start call that produced ti_context updated last_heartbeat_at on the server. + self.heartbeater.record_successful_heartbeat() # 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 - # for more context. Do not remove this without updating related comments and deferral handling. - start_date = ti_context.start_date or datetime.now(tz=timezone.utc) + # original start_date rather than using the resume time). We fall back to the transition's own + # timestamp otherwise. This ensures that `context["ti"].start_date` always reflects the *first* start + # time. See TIRunContext.start_date for more context. Do not remove this without updating related + # comments and deferral handling. + start_date = ti_context.start_date or running_since msg = StartupDetails.model_construct( ti=ti, diff --git a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py index 62b7fbcf39c17..9ea68ce59ee78 100644 --- a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py +++ b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py @@ -19,22 +19,27 @@ import contextlib import os +import signal import socket import subprocess import sys import threading import time +from datetime import datetime, timezone +from http import HTTPStatus from unittest.mock import ANY, MagicMock, call, patch import attrs +import httpx import psutil import pytest from uuid6 import uuid7 -from airflow.sdk.api.client import Client, TaskInstanceOperations -from airflow.sdk.api.datamodels._generated import TaskInstance +from airflow.sdk.api.client import Client, ServerResponseError, TaskInstanceOperations +from airflow.sdk.api.datamodels._generated import TaskInstance, TaskInstanceState from airflow.sdk.coordinators._subprocess import ( SubprocessCoordinator, + SubprocessStartupError, _accept_connections, _connection_owned_by_process_tree, _is_connection_from_process, @@ -42,6 +47,7 @@ _ResourceTracker, _start_server, ) +from airflow.sdk.exceptions import TaskAlreadyRunningError from airflow.sdk.execution_time.coordinator import BaseCoordinator from airflow.sdk.execution_time.supervisor import ActivitySubprocess @@ -215,10 +221,11 @@ def test_raises_timeout_when_no_connection(self): mock_proc = MagicMock(spec=subprocess.Popen) mock_proc.poll.return_value = None try: - with pytest.raises(TimeoutError, match="did not connect within timeout"): + with pytest.raises(SubprocessStartupError, match="did not connect within timeout") as exc_info: _accept_connections({"comm": server}, {}, mock_proc, max_wait=0.05) finally: server.close() + assert exc_info.value.exit_code == 1 def test_raises_runtime_error_if_process_exits_before_connecting(self): server = _start_server() @@ -226,10 +233,47 @@ def test_raises_runtime_error_if_process_exits_before_connecting(self): mock_proc.poll.return_value = 1 mock_proc.returncode = 1 try: - with pytest.raises(RuntimeError, match="process exited with 1"): + with pytest.raises(SubprocessStartupError, match="process exited with 1") as exc_info: _accept_connections({"comm": server}, {}, mock_proc) finally: server.close() + assert exc_info.value.exit_code == 1 + + def test_early_exit_carries_the_runtime_exit_code(self): + server = _start_server() + mock_proc = MagicMock(spec=subprocess.Popen) + mock_proc.poll.return_value = 127 + mock_proc.returncode = 127 + try: + with pytest.raises(SubprocessStartupError) as exc_info: + _accept_connections({"comm": server}, {}, mock_proc) + finally: + server.close() + assert exc_info.value.exit_code == 127 + + def test_drained_output_is_logged_when_startup_fails(self, cap_structlog): + """A runtime that dies before connecting usually explains itself on stderr.""" + server = _start_server() + drain_r, drain_w = socket.socketpair() + drain_w.sendall(b"Error: Could not find or load main class\n") + drain_w.shutdown(socket.SHUT_WR) + + mock_proc = MagicMock(spec=subprocess.Popen) + mock_proc.poll.return_value = None + try: + with pytest.raises(SubprocessStartupError): + _accept_connections({"comm": server}, {"stderr": drain_r}, mock_proc, max_wait=0.2) + finally: + drain_r.close() + drain_w.close() + server.close() + + assert { + "event": "Runtime output before startup failure", + "key": "stderr", + "output": "Error: Could not find or load main class\n", + "log_level": "error", + } in cap_structlog def test_returned_sockets_are_connected(self): """Accepted sockets should be real, usable connections.""" @@ -728,6 +772,296 @@ def test_returns_execution_result(self, mock_client): assert result.exit_code == 0 +@attrs.define +class _Run: + """What a driven `execute_task` call exposed for assertions.""" + + result: BaseCoordinator.ExecutionResult + events: list[str] + ti: TaskInstance + on_child_started: MagicMock + + +class TestSubprocessCoordinatorRunningTransition: + """The task must read RUNNING for every bit of work the worker does on its behalf.""" + + def _run(self, mock_client, coordinator, *, accept=None, build=None): + """ + Drive execute_task with a real (trivial) child process. + + ``subprocess.Popen`` is deliberately left unpatched: patching it replaces + the class on the stdlib module, which breaks the ``case subprocess.Popen()`` + cleanup in :class:`_ResourceTracker` that the failure paths below rely on. + """ + ti = _make_ti() + comm_sock = MagicMock(spec=socket.socket) + logs_sock = MagicMock(spec=socket.socket) + events: list[str] = [] + + def record_build(*, what): + events.append("build_command") + if build is not None: + return build(what=what) + return list(coordinator.command), None + + def default_accept(servers, drains, proc, **kw): + events.append("accept_connections") + return ( + {servers["comm"]: comm_sock, servers["logs"]: logs_sock}, + {soc: b"" for soc in drains.values()}, + ) + + mock_client.task_instances.start.side_effect = lambda *a, **kw: ( + events.append("report_running") or mock_client.task_instances.start.return_value + ) + + with ( + patch( + "airflow.sdk.coordinators._subprocess._accept_connections", + side_effect=accept or default_accept, + ), + patch.object(ActivitySubprocess, "_register_pipe_readers"), + patch.object(ActivitySubprocess, "_on_child_started") as mock_on_started, + patch.object(ActivitySubprocess, "wait", return_value=0), + patch.object(_StubSubprocessCoordinator, "_build_execute_task_command", side_effect=record_build), + ): + result = coordinator.execute_task( + what=ti, + dag_rel_path="bundle", + bundle_info=MagicMock(), + client=mock_client, + subprocess_logs_to_stdout=False, + ) + return _Run(result=result, events=events, ti=ti, on_child_started=mock_on_started) + + def test_running_reported_before_any_preparation(self, mock_client): + """Artifact resolution and Dag bundle materialization happen inside _build_execute_task_command.""" + run = self._run(mock_client, _StubSubprocessCoordinator(command=["/bin/true"])) + assert run.events == ["report_running", "build_command", "accept_connections"] + + def test_start_date_matches_the_transition_the_server_recorded(self, mock_client): + """Re-reading the clock would hand the runtime a start date later than the DB's.""" + run = self._run(mock_client, _StubSubprocessCoordinator(command=["/bin/true"])) + reported_when = mock_client.task_instances.start.call_args.args[2] + assert run.on_child_started.call_args.kwargs["running_since"] == reported_when + + def test_running_reported_once_with_the_supervisor_pid(self, mock_client): + run = self._run(mock_client, _StubSubprocessCoordinator(command=["/bin/true"])) + mock_client.task_instances.start.assert_called_once() + reported_id, reported_pid, _ = mock_client.task_instances.start.call_args.args + assert reported_id == run.ti.id + assert reported_pid == os.getpid() + + def test_heartbeat_sent_while_the_runtime_is_starting(self, mock_client, monkeypatch): + """Nothing else heartbeats until wait() monitors, so a slow launch looks like a zombie.""" + monkeypatch.setattr("airflow.sdk.coordinators._subprocess.MIN_HEARTBEAT_INTERVAL", 0.01) + beat = threading.Event() + mock_client.task_instances.heartbeat.side_effect = lambda *a, **kw: beat.set() + + def slow_build(*, what): + assert beat.wait(5), "no heartbeat arrived while preparing the launch" + return ["/bin/true"], None + + ti = _make_ti() + with ( + patch.object(_StubSubprocessCoordinator, "_build_execute_task_command", side_effect=slow_build), + patch("airflow.sdk.coordinators._subprocess.subprocess.Popen") as popen_mock, + ): + popen_mock.side_effect = ValueError("stop once the beat is observed") + _StubSubprocessCoordinator(command=["/bin/true"]).execute_task( + what=ti, + dag_rel_path="bundle", + bundle_info=MagicMock(), + client=mock_client, + subprocess_logs_to_stdout=False, + ) + + assert mock_client.task_instances.heartbeat.call_args.args[0] == ti.id + assert mock_client.task_instances.heartbeat.call_args.kwargs["pid"] == os.getpid() + + def test_startup_heartbeat_failure_does_not_abort_the_launch(self, mock_client, monkeypatch): + """Whatever broke the beat is still true once the monitor loop takes over.""" + monkeypatch.setattr("airflow.sdk.coordinators._subprocess.MIN_HEARTBEAT_INTERVAL", 0.01) + attempted = threading.Event() + + def explode(*a, **kw): + attempted.set() + raise RuntimeError("server down") + + mock_client.task_instances.heartbeat.side_effect = explode + comm_sock = MagicMock(spec=socket.socket) + logs_sock = MagicMock(spec=socket.socket) + + def slow_build(*, what): + assert attempted.wait(5), "the startup heartbeat never fired" + return ["/bin/true"], None + + with ( + patch.object(_StubSubprocessCoordinator, "_build_execute_task_command", side_effect=slow_build), + patch( + "airflow.sdk.coordinators._subprocess._accept_connections", + side_effect=lambda servers, drains, proc, **kw: ( + {servers["comm"]: comm_sock, servers["logs"]: logs_sock}, + {soc: b"" for soc in drains.values()}, + ), + ), + patch.object(ActivitySubprocess, "_register_pipe_readers"), + patch.object(ActivitySubprocess, "_on_child_started"), + patch.object(ActivitySubprocess, "wait", return_value=0), + ): + result = _StubSubprocessCoordinator(command=["/bin/true"]).execute_task( + what=_make_ti(), + dag_rel_path="bundle", + bundle_info=MagicMock(), + client=mock_client, + subprocess_logs_to_stdout=False, + ) + + assert result.exit_code == 0 + + def test_startup_heartbeat_stops_once_the_server_disowns_the_run(self, mock_client, monkeypatch): + """No point beating on: the monitor loop's own heartbeat terminates it shortly.""" + monkeypatch.setattr("airflow.sdk.coordinators._subprocess.MIN_HEARTBEAT_INTERVAL", 0.01) + disowned = threading.Event() + + def gone(*a, **kw): + disowned.set() + raise ServerResponseError( + "running elsewhere", + request=httpx.Request("PATCH", "http://server/heartbeat"), + response=httpx.Response(HTTPStatus.CONFLICT), + ) + + mock_client.task_instances.heartbeat.side_effect = gone + + def slow_build(*, what): + assert disowned.wait(5) + # Give the beat loop room to keep going if it were going to. + time.sleep(0.1) + return ["/bin/true"], None + + self._run(mock_client, _StubSubprocessCoordinator(command=["/bin/true"]), build=slow_build) + + assert mock_client.task_instances.heartbeat.call_count == 1 + + def test_handshake_failure_fails_the_task(self, mock_client): + def refuse(servers, drains, proc, **kw): + raise SubprocessStartupError("process did not connect within timeout", exit_code=3) + + run = self._run(mock_client, _StubSubprocessCoordinator(command=["/bin/true"]), accept=refuse) + + assert run.result.final_state == TaskInstanceState.FAILED + assert run.result.exit_code == 3 + mock_client.task_instances.finish.assert_called_once() + assert mock_client.task_instances.finish.call_args.kwargs["id"] == run.ti.id + assert mock_client.task_instances.finish.call_args.kwargs["state"] == TaskInstanceState.FAILED + mock_client.task_instances.retry.assert_not_called() + + def test_handshake_failure_retries_when_the_server_says_so(self, make_ti_context, mock_client): + mock_client.task_instances.start.return_value = make_ti_context(should_retry=True) + + def refuse(servers, drains, proc, **kw): + raise SubprocessStartupError("process exited with 1 before connecting") + + run = self._run(mock_client, _StubSubprocessCoordinator(command=["/bin/true"]), accept=refuse) + + assert run.result.final_state == TaskInstanceState.UP_FOR_RETRY + mock_client.task_instances.retry.assert_called_once() + assert mock_client.task_instances.retry.call_args.kwargs["id"] == run.ti.id + mock_client.task_instances.finish.assert_not_called() + + def test_command_build_failure_fails_the_task(self, mock_client): + """A misconfigured coordinator is a task failure now, not a run stuck in QUEUED.""" + ti = _make_ti() + coordinator = _StubSubprocessCoordinator(command=["/bin/true"]) + + with patch.object( + _StubSubprocessCoordinator, + "_build_execute_task_command", + side_effect=ValueError("no artifact found"), + ): + result = coordinator.execute_task( + what=ti, + dag_rel_path="bundle", + bundle_info=MagicMock(), + client=mock_client, + subprocess_logs_to_stdout=False, + ) + + assert result.final_state == TaskInstanceState.FAILED + assert result.exit_code == 1 + mock_client.task_instances.finish.assert_called_once() + + def test_terminal_report_failure_propagates(self, mock_client): + """Claiming a state the server never recorded would leave the run RUNNING forever.""" + mock_client.task_instances.finish.side_effect = RuntimeError("server down") + + with ( + patch.object( + _StubSubprocessCoordinator, + "_build_execute_task_command", + side_effect=ValueError("no artifact found"), + ), + pytest.raises(RuntimeError, match="server down"), + ): + _StubSubprocessCoordinator(command=["/bin/true"]).execute_task( + what=_make_ti(), + dag_rel_path="bundle", + bundle_info=MagicMock(), + client=mock_client, + subprocess_logs_to_stdout=False, + ) + + def test_start_conflict_propagates_without_spawning(self, mock_client): + """A redelivered workload must be rejected before a runtime is launched.""" + mock_client.task_instances.start.side_effect = TaskAlreadyRunningError("already running") + + with ( + patch("airflow.sdk.coordinators._subprocess.subprocess.Popen") as popen_mock, + patch.object(_StubSubprocessCoordinator, "_build_execute_task_command") as mock_build, + pytest.raises(TaskAlreadyRunningError), + ): + _StubSubprocessCoordinator(command=["/bin/true"]).execute_task( + what=_make_ti(), + dag_rel_path="bundle", + bundle_info=MagicMock(), + client=mock_client, + subprocess_logs_to_stdout=False, + ) + + popen_mock.assert_not_called() + mock_build.assert_not_called() + + def test_warm_shutdown_handlers_installed_for_the_running_window(self, mock_client): + """A SIGTERM in the RUNNING window must not kill this process and orphan the task.""" + handlers: list = [] + + def record_build(*, what): + handlers.append(signal.getsignal(signal.SIGTERM)) + return ["/bin/true"], None + + before = signal.getsignal(signal.SIGTERM) + with patch.object( + _StubSubprocessCoordinator, + "_build_execute_task_command", + side_effect=record_build, + ): + # Popen raises so execute_task unwinds to the ExecutionResult path without + # a real runtime, leaving `handlers` sampled from inside the window. + with patch("airflow.sdk.coordinators._subprocess.subprocess.Popen") as popen_mock: + popen_mock.side_effect = ValueError("stop here") + _StubSubprocessCoordinator(command=["/bin/true"]).execute_task( + what=_make_ti(), + dag_rel_path="bundle", + bundle_info=MagicMock(), + client=mock_client, + subprocess_logs_to_stdout=False, + ) + + assert handlers[0] is not before + assert signal.getsignal(signal.SIGTERM) is before + + class TestPopenActivitySubprocessStart: def _start_with_mocks(self, mock_client, *, command: list[str], schema_version=None): ti = _make_ti() @@ -756,6 +1090,9 @@ def _start_with_mocks(self, mock_client, *, command: list[str], schema_version=N dag_rel_path="bundle", bundle_info=MagicMock(), client=mock_client, + ti_context=mock_client.task_instances.start.return_value, + running_since=datetime(2026, 1, 1, tzinfo=timezone.utc), + reported_pid=os.getpid(), command=command, subprocess_schema_version=schema_version, subprocess_logs_to_stdout=False, @@ -790,6 +1127,8 @@ def test_on_child_started_called(self, mock_client): dag_rel_path="bundle", bundle_info=MagicMock(), client=mock_client, + ti_context=mock_client.task_instances.start.return_value, + running_since=datetime(2026, 1, 1, tzinfo=timezone.utc), command=["/bin/true"], subprocess_logs_to_stdout=False, ) @@ -799,6 +1138,41 @@ def test_on_child_started_called(self, mock_client): assert kwargs["ti"] is ti assert kwargs["dag_rel_path"] == "bundle" + def test_run_context_forwarded_to_on_child_started(self, mock_client): + """The coordinator already reported RUNNING, so start() must not do it again.""" + ti_context = mock_client.task_instances.start.return_value + with ( + patch("airflow.sdk.coordinators._subprocess.subprocess.Popen") as popen_mock, + patch( + "airflow.sdk.coordinators._subprocess._accept_connections", + side_effect=lambda servers, drains, proc, **kw: ( + {soc: MagicMock(spec=socket.socket) for soc in servers.values()}, + {soc: b"" for soc in drains.values()}, + ), + ), + patch.object(ActivitySubprocess, "_register_pipe_readers"), + patch.object(ActivitySubprocess, "_on_child_started") as mock_on_started, + ): + popen_mock.return_value.pid = 12345 + _PopenActivitySubprocess.start( + what=_make_ti(), + dag_rel_path="bundle", + bundle_info=MagicMock(), + client=mock_client, + ti_context=ti_context, + running_since=datetime(2026, 1, 1, tzinfo=timezone.utc), + command=["/bin/true"], + subprocess_logs_to_stdout=False, + ) + + assert mock_on_started.call_args.kwargs["ti_context"] is ti_context + + def test_reported_pid_is_the_supervisors(self, mock_client): + """RUNNING is reported before the runtime exists, so heartbeats must match that pid.""" + proc, _, _ = self._start_with_mocks(mock_client, command=["/bin/true"]) + assert proc.pid == 12345 + assert proc.heartbeater.pid == os.getpid() + @conf_vars({("logging", "logging_level"): "DEBUG"}) def test_resolved_log_level_passed_to_subprocess_env(self, mock_client): """A language SDK runtime gets the resolved task log level via the environment at launch.""" @@ -840,6 +1214,8 @@ def test_register_pipe_readers_called_with_four_sockets(self, mock_client): dag_rel_path="bundle", bundle_info=MagicMock(), client=mock_client, + ti_context=mock_client.task_instances.start.return_value, + running_since=datetime(2026, 1, 1, tzinfo=timezone.utc), command=["/bin/true"], subprocess_logs_to_stdout=False, )