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 9758d2b18d7c8..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 @@ -175,7 +176,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") @@ -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 @@ -1358,6 +1360,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 +1489,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 +1502,38 @@ 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.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, + ) + + 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, @@ -1423,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, *, @@ -1439,26 +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 - self._last_successful_heartbeat = time.monotonic() - 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, @@ -1610,17 +1771,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 +1819,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/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, ) 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