diff --git a/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md b/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md index 6ef154a0f6360..7a6e91e1567d7 100644 --- a/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md +++ b/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md @@ -43,7 +43,7 @@ parse_dag parse_task_handler parent process parent process │ DagFileParseRequest │ TaskHandlerParseRequest ▼ (ToDagProcessor) ▼ (ToSDKTaskHandlerProcessor) - coordinator (raw byte forward) coordinator (raw byte forward) + coordinator (execs the runtime) coordinator (execs the runtime) ▼ ▼ runtime runtime │ DagFileParsingResult │ TaskHandlerParsingResult @@ -51,8 +51,9 @@ parse_dag parse_task_handler parent process parent process ``` -Both verbs are byte forwarders: the coordinator spawns the runtime, wires `fd 0` to the comm socket, and never decodes the payload. The process that spawned the parse decodes the -reply. They stay two methods, not one `parse(request)`, because a coordinator can serve handlers without serving native Dag parsing. Appendix A has the longer argument. +Neither verb forwards bytes: in a child of the process that asked, the coordinator builds the command and execs the runtime, which connects back to that process. +The coordinator never decodes the payload, and the process that spawned the parse decodes the reply. They stay two methods, not one `parse(request)`, +because a coordinator can serve handlers without serving native Dag parsing. Appendix A has the longer argument. ### The reply travels on `ToManager` @@ -74,8 +75,9 @@ ToManager = DagFileParsingResult | TaskHandlerParsingResult connections, variables and XComs do not depend on which parse it was asked for. `ToManager` is named for the process that usually holds the other end, but the role it describes is "whoever spawned this parse". The Dag processor manager fills it for -`DagFileProcessorProcess`; a Dag-parsing child fills it for the two processes below, relaying anything that is not a parsing result up its own `ToManager` channel unchanged. That -relay is only type-safe because both hops speak the same pair, which is the reason not to mint a separate `ToCoordinator`. +`DagFileProcessorProcess` and `LangSDKDagFileProcessorProcess`; a Dag-parsing child fills it for `LangSDKTaskHandlerProcessorProcess`, +relaying each request that needs a client up its own `ToManager` channel unchanged. +That relay is only type-safe because both hops speak the same pair, which is the reason not to mint a separate `ToCoordinator`. ### Message shapes @@ -118,54 +120,59 @@ schema snapshot. A handler declaration carries no Dag, so it code-generates and ``` WatchedSubprocess - └── BaseParsingProcess socket lifecycle · ToManager decoding · Get* handling + └── BaseDagFileProcessorProcess socket lifecycle · ToManager decoding · Get* handling │ · log forwarding under dag_processor.* - ├── DagFileProcessorProcess (shipped, now a subclass) - │ │ target = _parse_file_entrypoint - │ │ DagFileParseRequest → DagFileParsingResult - │ │ - │ └── LangSDKDagFileProcessorProcess (new — ADR-0010) - │ target = _parse_lang_sdk_dag_entrypoint - │ └── coordinator.parse_dag() — spawn runtime, forward fd 0 ⇄ comm socket - │ same request and result types as its base class + ├── DagFileProcessorProcess (shipped) + │ target = _parse_file_entrypoint + │ DagFileParseRequest → DagFileParsingResult │ - └── SDKTaskHandlerProcessorProcess (new — ADR-0011) - target = _parse_task_handler_entrypoint - └── coordinator.parse_task_handler() — same forwarding + ├── LangSDKDagFileProcessorProcess (new, ADR-0010) + │ target = _start_runtime_entrypoint + │ └── coordinator.parse_dag(): exec the runtime, which connects back + │ DagFileParseRequest → DagFileParsingResult + │ + └── LangSDKTaskHandlerProcessorProcess (new, ADR-0011) + target = _start_task_handler_runtime_entrypoint + └── coordinator.parse_task_handler(): the same exec TaskHandlerParseRequest → TaskHandlerParsingResult ``` -`BaseParsingProcess` is `DagFileProcessorProcess` minus the Dag-specific request and result: the comm socket, the `ToManager` decode, the `Get*` dispatch, and the -`task.` → `dag_processor.` log-forwarder rename. The subclasses supply the first message they send, the result they collect, and the target the child runs. +`BaseDagFileProcessorProcess` holds what every parse process shares: the comm socket, the `ToManager` decode, the `Get*` dispatch, +and the `task.` → `dag_processor.` log-forwarder rename. Each subclass supplies the first message it sends, the result it collects, and the target its child runs. -`LangSDKDagFileProcessorProcess` differs from its base in the target callable alone. Everything else — the request, the result, the socket, the logging — is inherited, because a -native Lang-SDK Dag answers the same question a Python file does. +The two Lang-SDK processes have one shape. +Their child finds the coordinator, reports the runtime's schema version and execs the runtime, which connects back to two sockets the process owns and answers the request itself. -Answering `Get*` needs a `Client`, which only the manager holds. `BaseParsingProcess` therefore resolves a request one of two ways: directly against `self.client` when the manager -is the parent, or by relaying it up `SUPERVISOR_COMMS` when a Dag-parsing child is. +Answering `Get*` needs a `Client`, which only the manager holds. A process with a client answers directly, and one without answers with an error. +The exception is `LangSDKTaskHandlerProcessorProcess`: it runs in a Dag-parsing child, so it relays each such request up `SUPERVISOR_COMMS`. ### Coordinator interface ``` BaseCoordinator execution_time/coordinator.py - ├── execute_task (shipped) - ├── parse_dag (new — native Dags, ADR-0010) - └── parse_task_handler (new — handlers, ADR-0011) + └── execute_task (shipped) │ SubprocessCoordinator coordinators/_subprocess.py - implements all three; each resolves (command, subprocess_schema_version) - from a hook and owns the socket lifecycle: + implements the three verbs; each resolves (command, subprocess_schema_version) + from a command builder, and execute_task also owns the socket lifecycle: + ├── execute_task (shipped) + ├── parse_dag (new, native Dags, ADR-0010) + ├── parse_task_handler (new, handlers, ADR-0011) ├── _build_execute_task_command (shipped) ├── _build_parse_dag_command (new) - └── _build_parse_task_handler_command (new) + ├── _build_parse_task_handler_command (new) + └── _find_task_handler_artifact (new, called by the Dag processor: the artifact a stub task of a Dag runs, found as execute_task finds it) │ JavaCoordinator · ExecutableCoordinator · NodeCoordinator - supply the three commands; no socket or protocol code + supply the hooks; no socket or protocol code ``` -Names follow the shipped `execute_task` / `_build_execute_task_command` pair and supersede ADR-0004's `run_dag_parsing` / `dag_parsing_cmd`. Each hook returns its own +Names follow the shipped `execute_task` / `_build_execute_task_command` pair and supersede ADR-0004's `run_dag_parsing` / `dag_parsing_cmd`. Each command builder returns its own `subprocess_schema_version`, so handler parsing negotiates the schema the same way task execution does. +Only a `SubprocessCoordinator` can parse, so `BaseCoordinator` gains no parse verb. +`parse_task_handler` refuses a runtime whose schema version is older than `TASK_HANDLER_PARSING_SCHEMA_VERSION`, since it cannot answer the request. + ## Consequences - The handler channel is typed; the Dag channel is opaque. Every SDK gets generated models for the declaration shape. @@ -176,9 +183,14 @@ Names follow the shipped `execute_task` / `_build_execute_task_command` pair and already holds. - Nothing in the protocol distinguishes a coordinator-backed parse from a Python one. A runtime's `Get*` request is answered by the same handlers that answer a Python parser's, through however many relay hops lie between it and the manager. -- `DagFileProcessorProcess` becomes a subclass. Its public surface does not move, but the shipped `_handle_request` and socket code shifts to `BaseParsingProcess`. +- `DagFileProcessorProcess` becomes a subclass. Its public surface does not move, but the shipped `_handle_request` and socket code shifts to `BaseDagFileProcessorProcess`. - Neither verb is reached through ADR-0004's `can_handle_dag_file` scan. `parse_dag` is reached through the importer registered for the artifact's extension ([ADR-0010](0010-native-dag-processing.md)); `parse_task_handler` through `queue → coordinator` ([ADR-0011](0011-mixed-language-dag-processing.md)). +- Where the Dag processor starts its children with exec instead of fork (the default on macOS), + a probe started from a Dag-parsing child inherits that child's ORM-blocking environment and dies at `import airflow`, + so the check logs a warning and its stub tasks stay unchecked. +- The parent death signal reaches only the exec'd runtime, so a runtime must exec and leave no children: what it starts survives an abrupt kill of the Dag-parsing child, + though normal exits and timeouts kill what it leaves in its process group. - Terms track Language SDK spec `1.0`. A spec rename of `TaskHandler` lands here too. ## References @@ -201,9 +213,9 @@ The two verbs differ in what they ask for and in which command starts the runtim distinction: a coordinator can serve mixed-language handlers with no interest in native Dag parsing, and an absent method states that better than a runtime rejection does. The reply direction does not need the same split. An earlier draft gave handler parsing its own `ToRuntime` / `ToCoordinator` pair on the theory that the runtime was answering the -coordinator rather than the manager. It is not: the coordinator forwards bytes in both directions and decodes nothing, so the peer at the far end of the socket is whichever process -spawned the parse. Giving that peer two unions to decode would mean two `CommsDecoder` configurations, two relay paths for the identical `Get*` traffic, and two registry entries for -bodies that never differ. Reusing `ToManager` leaves one reply union with one new member. +coordinator rather than the manager. It is not: the runtime connects back to whichever process spawned the parse, and the coordinator decodes nothing. +Giving that peer two unions to decode would mean two `CommsDecoder` configurations, two relay paths for the identical `Get*` traffic, +and two registry entries for bodies that never differ. Reusing `ToManager` leaves one reply union with one new member. A boolean on `DagFileParseRequest` was the other alternative. It cannot work: a `DagRef` and a `TaskHandlerRef` are different payloads, not two subsets of one, so the flag would select between shapes the result type cannot both hold. @@ -238,6 +250,6 @@ under one name, so a member reached twice is not a clash. Two prek hooks guard the generated snapshot. The rule for a new body is documented in `task-sdk/src/airflow/sdk/execution_time/schema/AGENTS.md`. -Artifact roots are resolved per mode by `_init_root_source`, but published through `_get_scan_roots()`, which is scoped to an active task and raises outside one. Both parse-side -commands need those roots with no `TaskInstance` in hand, so the scope that publishes them has to open for a parse as well as for a task. [ADR-0010](0010-native-dag-processing.md) -covers the modes themselves. +Command builders read artifact roots through `_get_scan_roots()`, which raises outside the scope `_set_scan_roots` opens. +`execute_task` opens it with the roots it resolves for its mode, and `parse_dag` and `parse_task_handler` with the root of the Dag bundle that holds the file, +so both parse-side commands get roots with no `TaskInstance` in hand. [ADR-0010](0010-native-dag-processing.md) covers the modes themselves. diff --git a/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md b/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md index 65275c36c3912..d043ea927941b 100644 --- a/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md +++ b/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md @@ -386,8 +386,8 @@ DagFileProcessorProcess(etl.py) [no DB — client con │ anything differs ──▶ probe │ ├─6─ PROBE, per differing candidate — one subprocess - │ SDKTaskHandlerProcessorProcess.start( - │ target=_parse_task_handler_entrypoint, + │ LangSDKTaskHandlerProcessorProcess.start( + │ target=_start_task_handler_runtime_entrypoint, │ coordinator=JavaCoordinator("jdk-17"), │ path=) │ ──TaskHandlerParseRequest(file=…)──▶ runtime diff --git a/airflow-core/src/airflow/dag_processing/task_handler_processor.py b/airflow-core/src/airflow/dag_processing/task_handler_processor.py new file mode 100644 index 0000000000000..83b6e176159d4 --- /dev/null +++ b/airflow-core/src/airflow/dag_processing/task_handler_processor.py @@ -0,0 +1,592 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Ask the Lang-SDK runtime of a coordinator which task handlers an artifact registers.""" + +from __future__ import annotations + +import functools +import os +import selectors +import signal +import time +from contextlib import suppress +from pathlib import Path +from socket import MSG_DONTWAIT, socket +from typing import TYPE_CHECKING, Annotated, ClassVar, Literal, cast, get_args + +import attrs +import msgspec +import psutil +from pydantic import BaseModel, Field, TypeAdapter +from uuid6 import uuid7 + +from airflow.configuration import conf +from airflow.dag_processing.lang_sdk_processor import ( + _EXIT_GRACE_PERIOD, + _IMPORT_TIMEOUT_SETTING, + _PROCESSOR_TIMEOUT_SETTING, + LangSDKRuntimeSchemaVersion, + _Channel, + _get_import_timeout, +) +from airflow.dag_processing.processor import ( + BaseDagFileProcessorProcess, + DagFileParsingResult, + TaskHandlerParseRequest, + TaskHandlerParsingResult, + ToManager, +) +from airflow.sdk.coordinators._subprocess import SubprocessCoordinator, _is_connection_from_pid, _start_server +from airflow.sdk.exceptions import AirflowRuntimeError +from airflow.sdk.execution_time import supervisor, task_runner +from airflow.sdk.execution_time.comms import CommsDecoder, ErrorResponse, _RequestFrame +from airflow.sdk.execution_time.coordinator import get_coordinator_manager +from airflow.sdk.execution_time.supervisor import ( + ResponseSent, + length_prefixed_frame_reader, + make_buffered_socket_reader, + process_log_messages_from_subprocess, + register_request_method, +) + +if TYPE_CHECKING: + from collections.abc import Generator + + from structlog.typing import FilteringBoundLogger + + from airflow.sdk.execution_time.supervisor import RequestHandler, RequestResult + from airflow.typing_compat import Self + + +# StartTaskHandlerRuntime and LangSDKRuntimeStartFailed, like LangSDKRuntimeSchemaVersion, pass only +# between the probe and its forked child before the exec, so they are not part of the supervisor schema +# the runtimes speak. + + +class StartTaskHandlerRuntime(BaseModel): + """Ask the parse child to exec the runtime of the coordinator configured under *coordinator*.""" + + file: str + bundle_path: Path + coordinator: str + """The coordinator's key in ``[sdk] coordinators``.""" + comm_address: tuple[str, int] + logs_address: tuple[str, int] + type: Literal["StartTaskHandlerRuntime"] = "StartTaskHandlerRuntime" + + +class LangSDKRuntimeStartFailed(BaseModel): + """Why the parse child could not start the runtime.""" + + error: str + type: Literal["LangSDKRuntimeStartFailed"] = "LangSDKRuntimeStartFailed" + + +def _start_task_handler_runtime_entrypoint() -> None: + """Exec the runtime that probes the artifact named by the start request, or report why it cannot start.""" + os.environ["_AIRFLOW_PROCESS_CONTEXT"] = "client" + # fd 0 becomes the runtime's stdin, so the request channel moves to a close-on-exec copy. + comms = CommsDecoder[StartTaskHandlerRuntime, LangSDKRuntimeSchemaVersion | LangSDKRuntimeStartFailed]( + socket=socket(fileno=os.dup(0)), + body_decoder=TypeAdapter(StartTaskHandlerRuntime), + ) + devnull = os.open(os.devnull, os.O_RDONLY) + os.dup2(devnull, 0) + os.close(devnull) + + msg = comms._get_response() + if not isinstance(msg, StartTaskHandlerRuntime): + raise RuntimeError(f"Required first message to be a StartTaskHandlerRuntime, it was {msg}") + + def report_schema_version(schema_version: str | None) -> None: + comms.send(LangSDKRuntimeSchemaVersion(schema_version=schema_version, import_timeout=import_timeout)) + + try: + # The policy is user code: it runs in this child, where a failure is only this file's import error. + import_timeout = _get_import_timeout(msg.file) + coordinator = get_coordinator_manager().get_coordinator(msg.coordinator) + if not isinstance(coordinator, SubprocessCoordinator): + raise NotImplementedError(f"{type(coordinator).__name__} does not parse task handlers") + coordinator.parse_task_handler( + path=Path(msg.file), + bundle_path=msg.bundle_path, + comm_address=msg.comm_address, + logs_address=msg.logs_address, + report_schema_version=report_schema_version, + ) + except Exception as e: + comms.send(LangSDKRuntimeStartFailed(error=f"{type(e).__name__}: {e}")) + + +class _ReadsWithoutWaiting: + """ + The runtime's comm socket, with reads that return at once instead of waiting for data. + + A runtime that stops in the middle of a frame then cannot block the caller's loop, which keeps + checking the import timeout. Replies are sent on the socket itself, which stays blocking. + """ + + def __init__(self, sock: socket) -> None: + self._sock = sock + + def recv(self, bufsize: int) -> bytes: + return self._sock.recv(bufsize, MSG_DONTWAIT) + + def recv_into(self, buffer: memoryview) -> int: + return self._sock.recv_into(buffer, 0, MSG_DONTWAIT) + + +@attrs.define(kw_only=True) +class LangSDKTaskHandlerProcessorProcess(BaseDagFileProcessorProcess): + """ + Ask a coordinator's Lang-SDK runtime for every task handler an artifact registers. + + The forked parse child finds the coordinator, reports the runtime's schema version and execs the + runtime. The runtime connects back to two listeners this process owns and answers the + ``TaskHandlerParseRequest`` itself, so the request is sent once it has connected. A failed start, + a missing result, an invalid frame or message, or a timeout is an import error on the result, + keyed by the artifact's path in its Dag bundle. Processes the runtime leaves in its process group + are killed when it exits or the probe times out, but not when the Dag-parsing child is killed abruptly. + + The probe has no API client: each request of the runtime that needs one is relayed up the supervisor + channel of the process this runs in, as in a Dag-parsing child, or gets an error when there is none. + On Linux the kernel kills only the runtime when the thread that started this process exits, + so start it from a thread that outlives the probe. A runtime must exec and leave no children, + since the ones it starts survive an abrupt kill of the Dag-parsing child. + + Where the Dag processor starts its children with exec instead of fork (the default on macOS), + a probe started from a Dag-parsing child inherits that child's ORM-blocking environment, + and the probe's own child dies at ``import airflow``, so the probe fails. + """ + + parsing_result: TaskHandlerParsingResult | None = None # type: ignore[assignment] + + decoder = TypeAdapter( + Annotated[ + LangSDKRuntimeSchemaVersion | LangSDKRuntimeStartFailed | get_args(ToManager)[0], + Field(discriminator="type"), + ] + ) + + coordinator: str + """The coordinator's key in ``[sdk] coordinators``.""" + + _listeners: dict[_Channel, socket] + _parse_request: TaskHandlerParseRequest + _runtime_schema_version: str | None = attrs.field(default=None, init=False) + _import_timeout: float | None = attrs.field(default=None, init=False) + _schema_version_reported: bool = attrs.field(default=False, init=False) + _parsing_result_monotonic: float | None = attrs.field(default=None, init=False) + _unverified_connections: list[tuple[socket, _Channel]] = attrs.field(factory=list, init=False) + _group_killed: bool = attrs.field(default=False, init=False) + + @classmethod + def start( # type: ignore[override] + cls, + *, + coordinator: str, + path: str | os.PathLike[str], + bundle_path: Path, + bundle_name: str, + artifact_rel_path: str, + **kwargs, + ) -> Self: + """ + Start probing the artifact at *path* for every task handler it registers. + + *bundle_path* and *bundle_name* are those of the Dag bundle holding the artifact, and + *artifact_rel_path* is the artifact's path in it. + """ + listeners: dict[_Channel, socket] = {"comm": _start_server(), "logs": _start_server()} + try: + for listener in listeners.values(): + listener.setblocking(False) + parse_request = TaskHandlerParseRequest( + file=os.fspath(path), bundle_path=bundle_path, bundle_name=bundle_name + ) + proc = super().start( + target=_start_task_handler_runtime_entrypoint, + use_exec=supervisor._should_use_exec(), + new_process_group=True, + coordinator=coordinator, + bundle_name=bundle_name, + dag_file_rel_path=artifact_rel_path, + listeners=listeners, + parse_request=parse_request, + **kwargs, + ) + except BaseException: + for listener in listeners.values(): + listener.close() + raise + try: + for channel, listener in listeners.items(): + proc._open_sockets[listener] = f"{channel}-listener" + proc.selector.register( + listener, + selectors.EVENT_READ, + (functools.partial(proc._accept_connection, channel=channel), proc._on_socket_closed), + ) + proc.send_msg( + StartTaskHandlerRuntime( + file=parse_request.file, + bundle_path=bundle_path, + coordinator=coordinator, + comm_address=listeners["comm"].getsockname()[:2], + logs_address=listeners["logs"].getsockname()[:2], + ), + request_id=0, + ) + except BaseException: + proc._kill_runtime() + proc.close() + raise + return proc + + @classmethod + def run( + cls, + *, + coordinator: str, + path: str | os.PathLike[str], + bundle_path: Path, + bundle_name: str, + artifact_rel_path: str, + logger: FilteringBoundLogger, + deadline: float | None = None, + ) -> TaskHandlerParsingResult: + """ + Probe the artifact at *path* as :meth:`start` does, and wait for the result. + + The artifact's import timeout bounds the probe, and ``[dag_processor] dag_file_processor_timeout`` + until the parse child reports it. *deadline*, a :func:`time.monotonic` value, also bounds it: a probe + still running then is killed, and its result is an import error unless the runtime had already + answered. + """ + processor_timeout = conf.getfloat("dag_processor", "dag_file_processor_timeout") + with selectors.DefaultSelector() as selector: + proc = cls.start( + id=uuid7(), + coordinator=coordinator, + path=path, + bundle_path=bundle_path, + bundle_name=bundle_name, + artifact_rel_path=artifact_rel_path, + selector=selector, + logger=logger, + ) + try: + while not proc.is_ready: + if proc._schema_version_reported: + timeout, setting = proc._import_timeout, _IMPORT_TIMEOUT_SETTING + else: + timeout, setting = processor_timeout, _PROCESSOR_TIMEOUT_SETTING + if timeout is not None and time.monotonic() - proc.start_time > timeout: + # Unlike is_ready, this does not wait for an exited runtime's leftover processes, + # which can hold its sockets open. close() closes them. + proc._time_out(timeout, setting) + break + if deadline is not None and time.monotonic() >= deadline: + if proc.parsing_result is None: + proc._set_import_error( + f"The Lang-SDK runtime did not parse {proc._parse_request.file} " + "by its deadline" + ) + proc._kill_runtime() + break + proc._service_subprocess(max_wait_time=0.1) + except BaseException: + proc._kill_runtime() + raise + finally: + proc.close() + return cast("TaskHandlerParsingResult", proc.parsing_result) + + def _accept_connection(self, listener: socket, *, channel: _Channel) -> bool: + try: + conn, _ = listener.accept() + except (BlockingIOError, InterruptedError): + return True + conn.setblocking(True) + self._unverified_connections.append((conn, channel)) + self._verify_connections() + return True + + def _verify_connections(self) -> None: + """ + Use each accepted connection once it is confirmed to come from the runtime. + + A connection that is not visible yet stays pending and is checked again on the next + ``is_ready`` poll, so the caller's loop never waits here. + """ + pending = [] + for conn, channel in self._unverified_connections: + if channel not in self._listeners: + # The runtime already connected this channel. + conn.close() + continue + try: + owned = _is_connection_from_pid(conn, self.pid) + except OSError: + conn.close() + continue + if not owned: + pending.append((conn, channel)) + continue + self._close_listener(channel) + if channel == "comm": + self._register_comm(conn) + else: + self._register_logs(conn) + self._unverified_connections = pending + + def _close_listener(self, channel: _Channel) -> None: + if (listener := self._listeners.pop(channel, None)) is not None: + self._on_socket_closed(listener) + listener.close() + + def _close_listeners(self) -> None: + """Close the listeners of a runtime that did not connect, and connections never verified.""" + for channel in list(self._listeners): + self._close_listener(channel) + for conn, _ in self._unverified_connections: + conn.close() + self._unverified_connections = [] + + def _register_comm(self, conn: socket) -> None: + self.stdin = conn + self._open_sockets[conn] = "requests" + read_frame, on_close = length_prefixed_frame_reader( + self._handle_valid_requests(), on_close=self._on_socket_closed + ) + + def read_valid_frame(sock: socket) -> bool: + try: + return read_frame(cast("socket", _ReadsWithoutWaiting(sock))) + except BlockingIOError: + # The rest of the frame has not arrived; the reader keeps what it has read so far. + return True + except msgspec.DecodeError as e: + # A frame that does not decode would otherwise escape the Dag processor's selector loop. + self._fail_on_invalid_message(f"The Lang-SDK runtime sent an invalid frame: {e}") + return False + + self.selector.register(conn, selectors.EVENT_READ, (read_valid_frame, on_close)) + # The parse child reports the version and waits for the reply before it execs the runtime, + # so the version is known here. It is set only now, so the child's messages are not migrated. + self._subprocess_schema_version = self._runtime_schema_version + self.send_msg(self._parse_request, request_id=0) + + def _handle_valid_requests(self) -> Generator[None, _RequestFrame, None]: + """ + Pass each request on to ``handle_requests``, or kill the runtime at one that does not validate. + + ``handle_requests`` would only log such a request, and the runtime would wait for a reply. The + runtime speaks ``ToManager`` only; the start messages come from the parse child. + """ + requests = self.handle_requests(self.process_log) + next(requests) + while True: + frame = yield + try: + BaseDagFileProcessorProcess.decoder.validate_python(self._deserialize_request(frame.body)) + except ValueError as e: + self._fail_on_invalid_message( + f"The Lang-SDK runtime sent a message that does not validate: {e}" + ) + return + requests.send(frame) + + def _fail_on_invalid_message(self, message: str) -> None: + """Kill the runtime; *message* is the import error unless a parse result was already received.""" + if self.parsing_result is None: + self._set_import_error(message) + else: + self.process_log.warning( + "Ignoring an invalid message from the Lang-SDK runtime after its parse result", error=message + ) + self._kill_runtime() + + def _register_logs(self, conn: socket) -> None: + self._open_sockets[conn] = "logs" + self.selector.register( + conn, + selectors.EVENT_READ, + make_buffered_socket_reader( + process_log_messages_from_subprocess(self._get_target_loggers()), + on_close=self._on_socket_closed, + ), + ) + + def _set_import_error(self, message: str) -> None: + self.parsing_result = TaskHandlerParsingResult( + fileloc=self._parse_request.file, + task_handlers={}, + import_errors={self.dag_file_rel_path: message}, + ) + + def _handle_runtime_schema_version( + self, msg: LangSDKRuntimeSchemaVersion, log: FilteringBoundLogger, req_id: int + ) -> RequestResult | ResponseSent: + if self._schema_version_reported: + self._reject_request(msg, log, req_id) + return ResponseSent.ALREADY_SENT + self._runtime_schema_version = msg.schema_version + self._import_timeout = msg.import_timeout + self._schema_version_reported = True + return None, {} + + def _handle_start_failed( + self, msg: LangSDKRuntimeStartFailed, log: FilteringBoundLogger, req_id: int + ) -> RequestResult | ResponseSent: + self._set_import_error(f"Cannot start the Lang-SDK runtime: {msg.error}") + return None, {} + + def _handle_task_handler_parsing_result( + self, msg: TaskHandlerParsingResult, log: FilteringBoundLogger, req_id: int + ) -> RequestResult | ResponseSent: + if self.parsing_result is not None: + log.warning("Ignoring another parse result from the Lang-SDK runtime", fileloc=msg.fileloc) + self.send_msg( + None, + request_id=req_id, + error=ErrorResponse(detail={"message": "A parse result was already received"}), + ) + return ResponseSent.ALREADY_SENT + self.parsing_result = msg + self._parsing_result_monotonic = time.monotonic() + return None, {} + + # A runtime's Dag parse result is not a task handler result, so DagFileParsingResult stays unhandled. + _request_handlers: ClassVar[dict[type[BaseModel], RequestHandler[LangSDKTaskHandlerProcessorProcess]]] = { + **{ + message_type: handler + for message_type, handler in BaseDagFileProcessorProcess._common_request_handlers.items() + if message_type is not DagFileParsingResult + }, + **dict( + [ + register_request_method(LangSDKRuntimeSchemaVersion, _handle_runtime_schema_version), + register_request_method(LangSDKRuntimeStartFailed, _handle_start_failed), + register_request_method(TaskHandlerParsingResult, _handle_task_handler_parsing_result), + ] + ), + } + + def _handle_request(self, msg, log: FilteringBoundLogger, req_id: int) -> None: + if ( + self.client is None + and isinstance(msg, self._client_request_types) + and (comms := getattr(task_runner, "SUPERVISOR_COMMS", None)) is not None + ): + self._relay_request(comms, msg, req_id) + return + super()._handle_request(msg, log, req_id) + + def _relay_request(self, comms: CommsDecoder, msg: BaseModel, req_id: int) -> None: + """Answer the runtime's request through the supervisor channel of this process, errors included.""" + try: + response = comms.send(msg) + except AirflowRuntimeError as e: + self.send_msg(None, request_id=req_id, error=e.error) + return + # Only the fields the parent sent, under their wire names, so the runtime gets the same body. + self.send_msg(response, request_id=req_id, exclude_unset=True, by_alias=True) + + @property + def is_ready(self) -> bool: + self._verify_connections() + if ( + self._parsing_result_monotonic is not None + and time.monotonic() - self._parsing_result_monotonic > _EXIT_GRACE_PERIOD + ): + if self._exit_code is None: + self.process_log.warning( + "The Lang-SDK runtime did not exit after its parse result; killing it" + ) + elif not self._group_killed: + self.process_log.warning( + "The Lang-SDK runtime left processes holding its output after its parse result; " + "killing them" + ) + self._kill_runtime() + if ( + self._import_timeout is not None + and self.parsing_result is None + and self._exit_code is None + and time.monotonic() - self.start_time > self._import_timeout + ): + self._time_out(self._import_timeout, _IMPORT_TIMEOUT_SETTING) + if self._check_subprocess_exit() is None: + return False + self._close_listeners() + if not super().is_ready: + return False + if self.parsing_result is None: + self._set_import_error( + f"The Lang-SDK runtime exited with code {self._exit_code} without a parse result" + ) + return True + + def _time_out(self, timeout: float, setting: str) -> None: + """Kill the runtime; unless a parse result was received, the import error names *setting*.""" + if self.parsing_result is None: + self._set_import_error( + f"The Lang-SDK runtime did not parse {self._parse_request.file} within {timeout}s, " + f"the limit set by {setting}" + ) + self._kill_runtime() + + def _kill_runtime(self) -> None: + """ + Kill the runtime's process group and wait for the runtime, without servicing its sockets. + + Their handler may have failed. The wait is bounded so that a runtime stuck in the kernel does not + stall the caller's loop; ``is_ready`` sees it exit later. + """ + if self._exit_code is not None: + self._kill_leftovers() + return + try: + self._signal_subprocess(signal.SIGKILL) + self._group_killed = True + self._exit_code = self._process.wait(timeout=_EXIT_GRACE_PERIOD) + except (self._process.ProcessNotFound, ProcessLookupError): + self._exit_code = -1 + except self._process.TimeoutExpired: + self.process_log.warning("The Lang-SDK runtime did not exit after SIGKILL", pid=self.pid) + + def _kill_leftovers(self) -> None: + """ + Kill the processes an exited runtime left in its process group, which can keep its sockets open. + + The group is killed once, and not when another process has reused the runtime's pid, since the + group could then be that process's. + """ + if self._exit_code is None or self._group_killed or not self._new_process_group: + return + self._group_killed = True + if psutil.pid_exists(self.pid): + return + with suppress(ProcessLookupError, PermissionError): + os.killpg(self.pid, signal.SIGKILL) + + def close(self) -> None: + # A listener has nothing to drain, and cleanup would call its accept handler forever. + self._close_listeners() + self._kill_leftovers() + super().close() diff --git a/airflow-core/tests/unit/dag_processing/fake_task_handler_runtime.py b/airflow-core/tests/unit/dag_processing/fake_task_handler_runtime.py new file mode 100644 index 0000000000000..103c6a1894d34 --- /dev/null +++ b/airflow-core/tests/unit/dag_processing/fake_task_handler_runtime.py @@ -0,0 +1,122 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""A coordinator whose artifacts name the command that parses them, and a runtime a test can play.""" + +from __future__ import annotations + +import contextlib +import json +import socket +from pathlib import Path +from typing import TYPE_CHECKING, Any +from unittest import mock + +import attrs +from pydantic import TypeAdapter + +from airflow.dag_processing.processor import ( + TaskHandlerParseRequest, + TaskHandlerParsingResult, + ToManager, + ToSDKTaskHandlerProcessor, +) +from airflow.sdk.coordinators._subprocess import SubprocessCoordinator +from airflow.sdk.execution_time import supervisor +from airflow.sdk.execution_time.comms import CommsDecoder +from airflow.sdk.execution_time.coordinator import reset_coordinator_manager + +from tests_common.test_utils.config import conf_vars + +if TYPE_CHECKING: + from collections.abc import Callable, Iterator, Sequence + + +@attrs.define(kw_only=True) +class FakeCoordinator(SubprocessCoordinator): + """ + The artifact's JSON names the command that parses it. + + ``argv`` is the command, ``schema_version`` its schema version, and ``command_error`` an error to + raise instead. + """ + + def _build_parse_task_handler_command(self, *, path: Path) -> tuple[list[str], str | None]: + spec = json.loads(path.read_text()) + if error := spec.get("command_error"): + raise FileNotFoundError(error) + return spec.get("argv", ["/bin/false"]), spec.get("schema_version") + + +@contextlib.contextmanager +def fake_coordinator(*, other_coordinators: dict[str, str] | None = None, **kwargs: Any) -> Iterator[None]: + """ + Configure a ``FakeCoordinator`` as ``fake``, with fresh coordinators inside and after the block. + + *other_coordinators* maps keys to the classpaths of coordinators of another class. The parse child is + a bare fork even on macOS, so it sees the test's patches and config. + """ + spec = {"fake": {"classpath": f"{__name__}.FakeCoordinator", "kwargs": kwargs}} + spec.update({key: {"classpath": classpath} for key, classpath in (other_coordinators or {}).items()}) + reset_coordinator_manager() + try: + with ( + conf_vars({("sdk", "coordinators"): json.dumps(spec)}), + mock.patch.object(supervisor, "_should_use_exec", return_value=False), + ): + yield + finally: + reset_coordinator_manager() + + +def write_artifact(path: Path, **spec: Any) -> Path: + path.write_text(json.dumps(spec)) + return path + + +def play_runtime( + reply: Callable[[TaskHandlerParseRequest, CommsDecoder], TaskHandlerParsingResult | None], + *, + schema_version: str | None = None, + log_lines: Sequence[dict[str, Any]] = (), +) -> Callable[..., None]: + """ + Return a ``parse_task_handler`` that plays the runtime in the parse child instead of exec'ing one. + + It connects back as a runtime does, writes *log_lines* to its logs channel and answers the parse + request with what *reply* returns; ``None`` sends no result. + """ + + def parse_task_handler( + coordinator, *, comm_address, logs_address, report_schema_version, **kwargs + ) -> None: + report_schema_version(schema_version) + with ( + socket.create_connection(comm_address) as comm, + socket.create_connection(logs_address) as logs, + ): + for line in log_lines: + logs.sendall(json.dumps(line).encode() + b"\n") + comms = CommsDecoder[ToSDKTaskHandlerProcessor, ToManager]( + socket=comm, body_decoder=TypeAdapter(ToSDKTaskHandlerProcessor) + ) + request = comms._get_response() + assert isinstance(request, TaskHandlerParseRequest) + if (result := reply(request, comms)) is not None: + comms.send(result) + + return parse_task_handler diff --git a/airflow-core/tests/unit/dag_processing/test_task_handler_processor.py b/airflow-core/tests/unit/dag_processing/test_task_handler_processor.py new file mode 100644 index 0000000000000..a1792f98c4612 --- /dev/null +++ b/airflow-core/tests/unit/dag_processing/test_task_handler_processor.py @@ -0,0 +1,1090 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import contextlib +import io +import os +import selectors +import signal +import socket +import subprocess +import sys +import threading +import time +import typing +import uuid +from pathlib import Path +from typing import BinaryIO +from unittest.mock import ANY, MagicMock, call, patch + +import psutil +import pytest +import structlog +from pydantic import TypeAdapter +from structlog.typing import FilteringBoundLogger + +from airflow.configuration import conf +from airflow.dag_processing.lang_sdk_processor import _EXIT_GRACE_PERIOD, LangSDKRuntimeSchemaVersion +from airflow.dag_processing.processor import ( + BaseDagFileProcessorProcess, + DagFileParseRequest, + DagFileParsingResult, + DagFileProcessorProcess, + TaskHandlerDeclaration, + TaskHandlerParam, + TaskHandlerParseRequest, + TaskHandlerParsingResult, + ToDagProcessor, + ToManager, +) +from airflow.dag_processing.task_handler_processor import ( + LangSDKRuntimeStartFailed, + LangSDKTaskHandlerProcessorProcess, +) +from airflow.sdk.api.client import Client, VariableOperations +from airflow.sdk.api.datamodels._generated import VariableResponse +from airflow.sdk.coordinators._subprocess import ( + TASK_HANDLER_PARSING_SCHEMA_VERSION, + SubprocessCoordinator, +) +from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType +from airflow.sdk.execution_time import supervisor, task_runner +from airflow.sdk.execution_time.comms import ( + CommsDecoder, + ErrorResponse, + GetVariable, + MaskSecret, + VariableResult, + _RequestFrame, +) +from airflow.sdk.execution_time.supervisor import PsutilTracker + +from tests_common.test_utils.config import conf_vars +from unit.dag_processing.fake_task_handler_runtime import ( + FakeCoordinator, + fake_coordinator, + play_runtime, + write_artifact, +) + + +def _reply_with(*task_ids: str, **result): + """Reply with a handler for each of *task_ids* under the Dag ``etl``.""" + + def reply(request: TaskHandlerParseRequest, comms) -> TaskHandlerParsingResult: + declarations = [ + TaskHandlerDeclaration(task_id=task_id, binding="positional", params=[]) for task_id in task_ids + ] + return TaskHandlerParsingResult(fileloc=request.file, task_handlers={"etl": declarations}, **result) + + return reply + + +def _get_task_ids(result: TaskHandlerParsingResult) -> list[str]: + return [declaration.task_id for declaration in result.task_handlers["etl"]] + + +def _get_open_fds() -> set[int]: + # Without /proc, as on macOS, this is empty, so the fd leak checks pass trivially. + return {int(fd) for fd in os.listdir("/proc/self/fd")} if os.path.isdir("/proc/self/fd") else set() + + +# A runtime that exits once a child has left its process group; the child keeps the runtime's +# stdout and stderr open and writes its pid to argv[1]. +_LEAVE_A_CHILD_OUTSIDE_THE_GROUP = """ +import os, sys, time +read_fd, write_fd = os.pipe() +if os.fork(): + os.read(read_fd, 1) + os._exit(0) +os.setsid() +with open(sys.argv[1], "w") as f: + f.write(str(os.getpid())) +os.write(write_fd, b"x") +time.sleep(30) +""" + + +def _is_running(pid: int) -> bool: + try: + # Init may not have reaped the killed process yet. + return psutil.Process(pid).status() != psutil.STATUS_ZOMBIE + except psutil.NoSuchProcess: + return False + + +def _stops_running(pid: int, timeout: float = 10) -> bool: + # SIGKILL is delivered asynchronously, so the process may still run right after the kill. + deadline = time.monotonic() + timeout + while _is_running(pid): + if time.monotonic() > deadline: + return False + time.sleep(0.05) + return True + + +@pytest.fixture(autouse=True) +def _coordinator(): + with fake_coordinator(): + yield + + +@pytest.fixture +def supervisor_comms(monkeypatch): + """Give this process a supervisor channel, as a Dag-parsing child has.""" + comms = MagicMock(spec=CommsDecoder) + monkeypatch.setattr(task_runner, "SUPERVISOR_COMMS", comms, raising=False) + return comms + + +def _start(tmp_path, selector, **spec) -> LangSDKTaskHandlerProcessorProcess: + return LangSDKTaskHandlerProcessorProcess.start( + id=uuid.uuid4(), + coordinator="fake", + path=write_artifact(tmp_path / "etl.artifact", **spec), + bundle_path=tmp_path, + bundle_name="task-handlers", + artifact_rel_path="etl.artifact", + selector=selector, + logger=structlog.get_logger(), + ) + + +@pytest.fixture +def parse(tmp_path): + """Probe ``etl.artifact`` under a caller's selector loop, and check that nothing is left open.""" + + def _parse(**kwargs) -> LangSDKTaskHandlerProcessorProcess: + fds_before = _get_open_fds() + with selectors.DefaultSelector() as selector: + proc = _start(tmp_path, selector, **kwargs) + deadline = time.monotonic() + 30 + while not proc.is_ready: + assert time.monotonic() < deadline, "the Lang-SDK parse did not finish" + proc._service_subprocess(max_wait_time=0.1) + assert selector.get_map() == {} + proc.close() + assert _get_open_fds() <= fds_before + return proc + + return _parse + + +def _block_until_killed(comms) -> None: + """Hang as a stuck runtime does, until it is killed or the parent closes the comm socket.""" + comms.socket.recv(1) + + +def _send_an_invalid_frame(request, comms) -> None: + comms.socket.sendall(bytes.fromhex("00000003c1c1c1")) + _block_until_killed(comms) + + +def _send_a_frame(comms, body: dict) -> None: + comms.socket.sendall(_RequestFrame(id=1, body=body).as_bytes()) + _block_until_killed(comms) + + +def _send_an_invalid_result(request, comms) -> None: + _send_a_frame( + comms, {"type": "TaskHandlerParsingResult", "fileloc": request.file, "task_handlers": "none"} + ) + + +def _send_a_start_message(request, comms) -> None: + _send_a_frame(comms, LangSDKRuntimeSchemaVersion(schema_version=None).model_dump()) + + +def _run( + tmp_path, *, coordinator: str = "fake", deadline: float | None = None, **spec +) -> TaskHandlerParsingResult: + return LangSDKTaskHandlerProcessorProcess.run( + coordinator=coordinator, + path=write_artifact(tmp_path / "etl.artifact", **spec), + bundle_path=tmp_path, + bundle_name="task-handlers", + artifact_rel_path="etl.artifact", + logger=structlog.get_logger(), + deadline=deadline, + ) + + +class TestLangSDKTaskHandlerProcessorProcess: + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_collects_the_result_the_runtime_returns( + self, mock_parse_task_handler, parse, tmp_path, cap_structlog + ): + mock_parse_task_handler.side_effect = play_runtime( + _reply_with("extract"), + schema_version=TASK_HANDLER_PARSING_SCHEMA_VERSION, + log_lines=[{"event": "Registering handlers", "level": "info"}], + ) + + proc = parse() + + assert proc.parsing_result.fileloc == os.fspath(tmp_path / "etl.artifact") + assert proc.parsing_result.import_errors is None + assert _get_task_ids(proc.parsing_result) == ["extract"] + assert proc.bundle_name == "task-handlers" + assert proc._subprocess_schema_version == TASK_HANDLER_PARSING_SCHEMA_VERSION + assert "Registering handlers" in cap_structlog + + @patch("airflow.dag_processing.task_handler_processor._is_connection_from_pid", autospec=True) + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_connection_is_used_once_it_is_verified(self, mock_parse_task_handler, mock_owned, tmp_path): + mock_parse_task_handler.side_effect = play_runtime(_reply_with("extract")) + mock_owned.return_value = False + with selectors.DefaultSelector() as selector: + proc = _start(tmp_path, selector) + child_stdin = proc.stdin + deadline = time.monotonic() + 30 + while len(proc._unverified_connections) < 2: + assert proc.stdin is child_stdin, "an unverified connection was used" + assert time.monotonic() < deadline, "the runtime did not connect" + proc._service_subprocess(max_wait_time=0.1) + assert proc.stdin is child_stdin + + mock_owned.return_value = True + while not proc.is_ready: + assert time.monotonic() < deadline, "the Lang-SDK parse did not finish" + proc._service_subprocess(max_wait_time=0.1) + proc.close() + + assert _get_task_ids(proc.parsing_result) == ["extract"] + + @patch.object( + LangSDKTaskHandlerProcessorProcess, + "send_msg", + autospec=True, + side_effect=OSError("cannot send the start request"), + ) + def test_a_start_that_fails_after_the_fork_leaves_nothing_behind(self, mock_send_msg, tmp_path): + fds_before = _get_open_fds() + children_before = {child.pid for child in psutil.Process().children()} + + with selectors.DefaultSelector() as selector: + with pytest.raises(OSError, match="cannot send the start request"): + _start(tmp_path, selector) + assert selector.get_map() == {} + + assert {child.pid for child in psutil.Process().children()} == children_before + assert _get_open_fds() <= fds_before + + @patch.object(BaseDagFileProcessorProcess, "start", autospec=True, side_effect=OSError("fork failed")) + def test_a_start_that_fails_before_the_fork_closes_its_listeners(self, mock_start, tmp_path): + fds_before = _get_open_fds() + + with selectors.DefaultSelector() as selector: + with pytest.raises(OSError, match="fork failed"): + _start(tmp_path, selector) + + listeners = mock_start.call_args.kwargs["listeners"] + assert [listener.fileno() for listener in listeners.values()] == [-1, -1] + assert _get_open_fds() <= fds_before + + @pytest.mark.parametrize( + ("spec", "reply", "error"), + [ + pytest.param( + {"command_error": "no runtime"}, + None, + "Cannot start the Lang-SDK runtime: FileNotFoundError: no runtime", + id="command-not-resolved", + ), + pytest.param( + {"argv": ["/no/such/runtime"]}, + None, + "Cannot start the Lang-SDK runtime: FileNotFoundError: " + "[Errno 2] No such file or directory: '/no/such/runtime'", + id="exec-failed", + ), + pytest.param( + {"argv": ["/bin/sh", "-c", "exit 3"]}, + None, + "The Lang-SDK runtime exited with code 3 without a parse result", + id="exits-before-connecting", + ), + pytest.param( + {}, + lambda request, comms: None, + "The Lang-SDK runtime exited with code 0 without a parse result", + id="exits-without-a-result", + ), + pytest.param( + {}, + _send_an_invalid_frame, + "The Lang-SDK runtime sent an invalid frame: MessagePack data is malformed: " + "invalid opcode '\\xc1' (byte 0)", + id="invalid-frame", + ), + ], + ) + def test_a_failed_parse_is_an_import_error(self, parse, tmp_path, spec, reply, error): + with ( + patch.object( + FakeCoordinator, "parse_task_handler", autospec=True, side_effect=play_runtime(reply) + ) + if reply + else contextlib.nullcontext() + ): + proc = parse(**spec) + + assert proc.parsing_result == TaskHandlerParsingResult( + fileloc=os.fspath(tmp_path / "etl.artifact"), + task_handlers={}, + import_errors={"etl.artifact": error}, + ) + + def test_a_coordinator_that_is_not_configured_is_an_import_error(self, tmp_path): + result = _run(tmp_path, coordinator="missing") + + assert result == TaskHandlerParsingResult( + fileloc=os.fspath(tmp_path / "etl.artifact"), + task_handlers={}, + import_errors={ + "etl.artifact": "Cannot start the Lang-SDK runtime: " + "InvalidCoordinatorError: No coordinator 'missing' in [sdk] coordinators" + }, + ) + + @patch.object( + FakeCoordinator, + "_build_parse_task_handler_command", + SubprocessCoordinator._build_parse_task_handler_command, + ) + def test_a_coordinator_that_does_not_parse_task_handlers_is_an_import_error(self, parse): + proc = parse() + + assert proc.parsing_result.import_errors == { + "etl.artifact": "Cannot start the Lang-SDK runtime: " + "NotImplementedError: FakeCoordinator does not parse task handlers" + } + + def test_a_coordinator_that_is_not_a_subprocess_coordinator_is_an_import_error(self, tmp_path): + with fake_coordinator( + other_coordinators={"plain": "airflow.sdk.execution_time.coordinator.BaseCoordinator"} + ): + result = _run(tmp_path, coordinator="plain") + + assert result.import_errors == { + "etl.artifact": "Cannot start the Lang-SDK runtime: " + "NotImplementedError: BaseCoordinator does not parse task handlers" + } + + @pytest.mark.parametrize( + ("reply", "error"), + [ + pytest.param( + _send_an_invalid_result, + "TaskHandlerParsingResult.task_handlers\n Input should be a valid dictionary", + id="result", + ), + pytest.param( + _send_a_start_message, "does not match any of the expected tags", id="start-message" + ), + ], + ) + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_message_that_does_not_validate_is_an_import_error( + self, mock_parse_task_handler, parse, reply, error + ): + mock_parse_task_handler.side_effect = play_runtime(reply) + + proc = parse() + + [message] = proc.parsing_result.import_errors.values() + assert message.startswith("The Lang-SDK runtime sent a message that does not validate: ") + assert error in message + assert proc._exit_code == -signal.SIGKILL + + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_dag_file_parsing_result_is_not_a_task_handler_result(self, mock_parse_task_handler, tmp_path): + def reply(request, comms): + with pytest.raises(AirflowRuntimeError, match="Unhandled request"): + comms.send(DagFileParsingResult(fileloc=request.file, serialized_dags=[])) + + mock_parse_task_handler.side_effect = play_runtime(reply) + + result = _run(tmp_path) + + assert result.import_errors == { + "etl.artifact": "The Lang-SDK runtime exited with code 0 without a parse result" + } + + @patch.object( + FakeCoordinator, "parse_task_handler", autospec=True, side_effect=play_runtime(_send_an_invalid_frame) + ) + def test_killing_the_runtime_is_not_reported_as_out_of_memory( + self, mock_parse_task_handler, parse, cap_structlog + ): + proc = parse() + + assert proc._exit_code == -signal.SIGKILL + assert not any("Likely out of memory" in str(entry.get("event")) for entry in cap_structlog.entries) + + @patch("airflow.dag_processing.task_handler_processor._EXIT_GRACE_PERIOD", 0.5) + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_runtime_that_runs_on_after_its_result_is_killed( + self, mock_parse_task_handler, parse, cap_structlog + ): + def reply(request, comms): + comms.send(_reply_with("extract")(request, comms)) + _block_until_killed(comms) + + mock_parse_task_handler.side_effect = play_runtime(reply) + + proc = parse() + + assert _get_task_ids(proc.parsing_result) == ["extract"] + assert proc._exit_code == -signal.SIGKILL + assert "The Lang-SDK runtime did not exit after its parse result; killing it" in cap_structlog + + @patch("airflow.dag_processing.task_handler_processor._EXIT_GRACE_PERIOD", 0.5) + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_what_the_runtime_leaves_after_its_result_is_killed( + self, mock_parse_task_handler, parse, tmp_path, cap_structlog + ): + def reply(request, comms): + comms.send(_reply_with("extract")(request, comms)) + # The leftover inherits the runtime's stdout, so the parse is not done when the runtime exits. + leftover = subprocess.Popen(["sleep", "60"]) + (tmp_path / "leftover.pid").write_text(str(leftover.pid)) + + mock_parse_task_handler.side_effect = play_runtime(reply) + + proc = parse() + + assert _get_task_ids(proc.parsing_result) == ["extract"] + assert proc._exit_code == 0 + assert ( + "The Lang-SDK runtime left processes holding its output after its parse result; killing them" + in cap_structlog + ) + assert _stops_running(int((tmp_path / "leftover.pid").read_text())) + + @pytest.mark.parametrize( + ("policy", "error"), + [ + pytest.param( + {"side_effect": RuntimeError("policy bug")}, "RuntimeError: policy bug", id="raises" + ), + pytest.param( + {"return_value": "30"}, + "TypeError: Value (30) from get_dagbag_import_timeout must be int or float", + id="not-a-number", + ), + ], + ) + def test_a_failing_import_timeout_policy_is_an_import_error(self, parse, policy, error): + with patch("airflow.settings.get_dagbag_import_timeout", autospec=True, **policy): + proc = parse() + + assert proc.parsing_result.import_errors == { + "etl.artifact": f"Cannot start the Lang-SDK runtime: {error}" + } + + @pytest.mark.skipif(not Path("/proc/self/fd").is_dir(), reason="reads /proc") + @pytest.mark.parametrize("use_exec", [False, True], ids=["fork", "spawn"]) + def test_the_runtime_inherits_only_its_standard_streams(self, monkeypatch, tmp_path, use_exec): + if use_exec: + # The spawned interpreter finds the coordinator again from its environment. + monkeypatch.setattr(supervisor, "_should_use_exec", lambda: True) + monkeypatch.setenv("PYTHONPATH", os.pathsep.join(sys.path)) + monkeypatch.setenv("AIRFLOW__SDK__COORDINATORS", conf.get("sdk", "coordinators")) + with selectors.DefaultSelector() as selector: + proc = _start(tmp_path, selector, argv=["/bin/sh", "-c", "exec sleep 30"]) + deadline = time.monotonic() + 30 + while psutil.Process(proc.pid).name() != "sleep": + assert time.monotonic() < deadline, "the runtime did not start" + proc._service_subprocess(max_wait_time=0.1) + fd_dir = Path(f"/proc/{proc.pid}/fd") + fds = {fd.name: os.readlink(fd) for fd in fd_dir.iterdir()} + proc.kill(signal.SIGKILL) + proc.close() + + assert sorted(fds) == ["0", "1", "2"] + assert fds["0"] == "/dev/null" + + @pytest.mark.skipif(sys.platform != "linux", reason="PR_SET_PDEATHSIG is Linux-only") + @pytest.mark.execution_timeout(60) + def test_the_runtime_dies_with_the_dag_parsing_child(self, tmp_path): + artifact = write_artifact(tmp_path / "etl.artifact", argv=["/bin/sh", "-c", "exec sleep 60"]) + proc = DagFileProcessorProcess.start( + id=1, + path=artifact, + bundle_path=tmp_path, + bundle_name="task-handlers", + dag_file_rel_path="etl.artifact", + callbacks=[], + target=_probe_from_a_dag_parsing_child, + logger=structlog.get_logger(), + logger_filehandle=io.BytesIO(), + client=MagicMock(spec=Client), + ) + runtime = None + try: + deadline = time.monotonic() + 30 + while runtime is None: + assert time.monotonic() < deadline, "the runtime did not start" + proc._service_subprocess(max_wait_time=0.1) + with contextlib.suppress(psutil.NoSuchProcess): + runtime = next( + (c for c in psutil.Process(proc.pid).children(recursive=True) if c.name() == "sleep"), + None, + ) + + # The Dag processor manager kills a Dag-parsing child by pid, with no chance to kill its runtime. + os.kill(proc.pid, signal.SIGKILL) + + assert _stops_running(runtime.pid), "the runtime outlived the Dag-parsing child" + finally: + if runtime is not None: + with contextlib.suppress(psutil.NoSuchProcess): + runtime.kill() + proc.kill(signal.SIGKILL) + deadline = time.monotonic() + 30 + while not proc.is_ready: + assert time.monotonic() < deadline, "the Dag-parsing child did not finish" + proc._service_subprocess(max_wait_time=0.1) + proc.close() + + def test_registered_message_types(self): + to_manager = set(typing.get_args(typing.get_args(ToManager)[0])) + handlers = LangSDKTaskHandlerProcessorProcess._request_handlers + + assert set(handlers) == (to_manager - {DagFileParsingResult}) | { + LangSDKRuntimeSchemaVersion, + LangSDKRuntimeStartFailed, + } + for message_type in to_manager - {DagFileParsingResult, TaskHandlerParsingResult}: + assert ( + handlers[message_type] is BaseDagFileProcessorProcess._common_request_handlers[message_type] + ) + assert TaskHandlerParsingResult not in BaseDagFileProcessorProcess._common_request_handlers + + +def _probe_from_a_dag_parsing_child() -> None: + """Stand in for ``_parse_file_entrypoint``: probe the file it is asked to parse, and return the result.""" + comms_decoder = CommsDecoder[ToDagProcessor, ToManager](body_decoder=TypeAdapter(ToDagProcessor)) + request = comms_decoder._get_response() + assert isinstance(request, DagFileParseRequest) + task_runner.SUPERVISOR_COMMS = comms_decoder # type: ignore[assignment] + + result = LangSDKTaskHandlerProcessorProcess.run( + coordinator="fake", + path=request.file, + bundle_path=request.bundle_path, + bundle_name=request.bundle_name, + artifact_rel_path="etl.artifact", + logger=structlog.get_logger(logger_name="task"), + ) + comms_decoder.send( + DagFileParsingResult( + fileloc=request.file, serialized_dags=[], warnings=[result.model_dump(mode="json")] + ) + ) + + +class TestRuntimeRequests: + @pytest.mark.parametrize( + "request_type", BaseDagFileProcessorProcess._client_request_types, ids=lambda t: t.__name__ + ) + @patch.object(LangSDKTaskHandlerProcessorProcess, "send_msg", autospec=True) + def test_a_client_request_is_relayed_to_the_parent(self, mock_send_msg, request_type, supervisor_comms): + proc = _make_process() + msg = request_type.model_construct() + + proc._handle_request(msg, structlog.get_logger(), 7) + + supervisor_comms.send.assert_called_once_with(msg) + mock_send_msg.assert_called_once_with( + proc, supervisor_comms.send.return_value, request_id=7, exclude_unset=True, by_alias=True + ) + + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_request_is_relayed_to_the_parent(self, mock_parse_task_handler, parse, supervisor_comms): + def reply(request, comms): + variable = comms.send(GetVariable(key="probe_var")) + return _reply_with(variable.value)(request, comms) + + mock_parse_task_handler.side_effect = play_runtime( + reply, schema_version=TASK_HANDLER_PARSING_SCHEMA_VERSION + ) + # As decoded from the parent's frame, which always carries its type. + supervisor_comms.send.return_value = VariableResult.model_validate( + {"key": "probe_var", "value": "from-the-parent", "type": "VariableResult"} + ) + + proc = parse() + + assert _get_task_ids(proc.parsing_result) == ["from-the-parent"] + supervisor_comms.send.assert_called_once_with(GetVariable(key="probe_var")) + + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_the_parent_s_error_reaches_the_runtime(self, mock_parse_task_handler, parse, supervisor_comms): + def reply(request, comms): + with pytest.raises(AirflowRuntimeError) as ctx: + comms.send(GetVariable(key="probe_var")) + error = ctx.value.error + return _reply_with(f"{error.error.value}:{error.detail['key']}")(request, comms) + + mock_parse_task_handler.side_effect = play_runtime(reply) + supervisor_comms.send.side_effect = AirflowRuntimeError( + ErrorResponse(error=ErrorType.VARIABLE_NOT_FOUND, detail={"key": "probe_var"}) + ) + + proc = parse() + + assert _get_task_ids(proc.parsing_result) == [f"{ErrorType.VARIABLE_NOT_FOUND.value}:probe_var"] + + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_request_gets_an_error_without_a_parent(self, mock_parse_task_handler, monkeypatch, tmp_path): + monkeypatch.delattr(task_runner, "SUPERVISOR_COMMS", raising=False) + + def reply(request, comms): + with pytest.raises(AirflowRuntimeError) as ctx: + comms.send(GetVariable(key="probe_var")) + return _reply_with(ctx.value.error.detail["message"])(request, comms) + + mock_parse_task_handler.side_effect = play_runtime(reply) + + result = _run(tmp_path) + + assert _get_task_ids(result) == ["GetVariable is answered only in the Dag processor"] + + @patch("airflow.sdk.log._secrets_masker", autospec=True) + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_secret_is_masked_here_and_by_the_parent( + self, mock_parse_task_handler, mock_secrets_masker, parse, supervisor_comms + ): + def reply(request, comms): + comms.send(MaskSecret(value="probe-secret", name="probe_conn")) + return _reply_with("extract")(request, comms) + + mock_parse_task_handler.side_effect = play_runtime(reply) + + proc = parse() + + assert _get_task_ids(proc.parsing_result) == ["extract"] + mock_secrets_masker.return_value.add_mask.assert_called_once_with("probe-secret", "probe_conn") + supervisor_comms.send.assert_called_once_with(MaskSecret(value="probe-secret", name="probe_conn")) + + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_runtime_request_is_answered_by_the_dag_processor(self, mock_parse_task_handler, tmp_path): + def reply(request, comms): + variable = comms.send(GetVariable(key="probe_var")) + return _reply_with(variable.value)(request, comms) + + mock_parse_task_handler.side_effect = play_runtime(reply) + client = MagicMock(spec=Client) + client.variables = MagicMock(spec=VariableOperations) + client.variables.get.return_value = VariableResponse(key="probe_var", value="from-the-dag-processor") + artifact = write_artifact(tmp_path / "etl.artifact") + + proc = DagFileProcessorProcess.start( + id=1, + path=artifact, + bundle_path=tmp_path, + bundle_name="task-handlers", + dag_file_rel_path="etl.artifact", + callbacks=[], + target=_probe_from_a_dag_parsing_child, + logger=structlog.get_logger(), + logger_filehandle=io.BytesIO(), + client=client, + ) + deadline = time.monotonic() + 30 + while not proc.is_ready: + assert time.monotonic() < deadline, "the Dag-parsing child did not finish" + proc._service_subprocess(max_wait_time=0.1) + proc.close() + + client.variables.get.assert_called_once_with("probe_var") + [probe_result] = proc.parsing_result.warnings + assert TaskHandlerParsingResult.model_validate(probe_result) == TaskHandlerParsingResult( + fileloc=os.fspath(artifact), + task_handlers={ + "etl": [ + TaskHandlerDeclaration(task_id="from-the-dag-processor", binding="positional", params=[]) + ] + }, + ) + + +class TestRun: + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_returns_every_handler_the_runtime_registers(self, mock_parse_task_handler, tmp_path): + def reply(request, comms): + # Echo what the request carries, so the test can check it. + param = TaskHandlerParam(name=request.bundle_name, value_schema={"type": "string"}) + return TaskHandlerParsingResult( + fileloc=request.file, + task_handlers={ + dag_id: [ + TaskHandlerDeclaration( + task_id=os.fspath(request.bundle_path), binding="positional", params=[param] + ) + ] + for dag_id in ("etl", "report") + }, + ) + + mock_parse_task_handler.side_effect = play_runtime(reply) + + result = _run(tmp_path) + + param = TaskHandlerParam(name="task-handlers", value_schema={"type": "string"}) + declaration = TaskHandlerDeclaration( + task_id=os.fspath(tmp_path), binding="positional", params=[param] + ) + assert result == TaskHandlerParsingResult( + fileloc=os.fspath(tmp_path / "etl.artifact"), + task_handlers={"etl": [declaration], "report": [declaration]}, + ) + + @pytest.mark.execution_timeout(30) + @pytest.mark.parametrize("connected", [False, True], ids=["before-connecting", "after-connecting"]) + @patch("airflow.settings.get_dagbag_import_timeout", autospec=True, return_value=1) + def test_a_parse_past_the_import_timeout_is_killed(self, mock_timeout, tmp_path, connected): + # A runtime that never connects leaves both listeners open when it is killed. + fds_before = _get_open_fds() + + with ( + patch.object( + FakeCoordinator, + "parse_task_handler", + autospec=True, + side_effect=play_runtime(lambda request, comms: _block_until_killed(comms)), + ) + if connected + else contextlib.nullcontext(), + patch.object( + LangSDKTaskHandlerProcessorProcess, + "close", + autospec=True, + side_effect=LangSDKTaskHandlerProcessorProcess.close, + ) as mock_close, + ): + result = _run(tmp_path, argv=["/bin/sh", "-c", "exec sleep 60"]) + + assert result.import_errors == { + "etl.artifact": f"The Lang-SDK runtime did not parse {tmp_path / 'etl.artifact'} within 1.0s, " + "the limit set by [core] dagbag_import_timeout or the get_dagbag_import_timeout policy" + } + [proc] = [c.args[0] for c in mock_close.call_args_list] + assert proc._exit_code == -9 + assert not proc._open_sockets + assert _get_open_fds() <= fds_before + + @pytest.mark.execution_timeout(30) + @patch("airflow.settings.get_dagbag_import_timeout", autospec=True, return_value=1) + def test_what_the_runtime_leaves_is_killed_at_the_import_timeout(self, mock_timeout, tmp_path): + pid_file = tmp_path / "leftover.pid" + with patch.object( + LangSDKTaskHandlerProcessorProcess, + "close", + autospec=True, + side_effect=LangSDKTaskHandlerProcessorProcess.close, + ) as mock_close: + # The runtime exits, and the process it leaves in its process group keeps its output open. + result = _run(tmp_path, argv=["/bin/sh", "-c", f"sleep 30 & echo $! > {pid_file}; exit 0"]) + [proc] = [c.args[0] for c in mock_close.call_args_list] + leftover = int(pid_file.read_text()) + try: + assert result.import_errors == { + "etl.artifact": f"The Lang-SDK runtime did not parse {tmp_path / 'etl.artifact'} within 1.0s, " + "the limit set by [core] dagbag_import_timeout or the get_dagbag_import_timeout policy" + } + assert proc._exit_code == 0 + assert _stops_running(leftover) + finally: + with contextlib.suppress(ProcessLookupError): + os.kill(leftover, signal.SIGKILL) + + @pytest.mark.execution_timeout(30) + @patch("airflow.settings.get_dagbag_import_timeout", autospec=True, return_value=1) + def test_the_import_timeout_holds_after_the_runtime_exits(self, mock_timeout, tmp_path): + pid_file = tmp_path / "leftover.pid" + # The runtime exits, and the process it leaves outside its process group keeps its output open. + try: + with patch.object( + LangSDKTaskHandlerProcessorProcess, + "close", + autospec=True, + side_effect=LangSDKTaskHandlerProcessorProcess.close, + ) as mock_close: + result = _run( + tmp_path, + argv=[sys.executable, "-c", _LEAVE_A_CHILD_OUTSIDE_THE_GROUP, os.fspath(pid_file)], + ) + finally: + if pid_file.exists(): + with contextlib.suppress(ProcessLookupError): + os.kill(int(pid_file.read_text()), signal.SIGKILL) + [proc] = [c.args[0] for c in mock_close.call_args_list] + + assert result.import_errors == { + "etl.artifact": f"The Lang-SDK runtime did not parse {tmp_path / 'etl.artifact'} within 1.0s, " + "the limit set by [core] dagbag_import_timeout or the get_dagbag_import_timeout policy" + } + assert proc._exit_code == 0 + assert not proc._open_sockets + + @pytest.mark.execution_timeout(30) + @patch("airflow.settings.get_dagbag_import_timeout", autospec=True, return_value=1) + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_runtime_that_stops_mid_frame_is_timed_out( + self, mock_parse_task_handler, mock_timeout, tmp_path + ): + def reply(request, comms): + comms.socket.sendall((100).to_bytes(4, byteorder="big")) + _block_until_killed(comms) + + mock_parse_task_handler.side_effect = play_runtime(reply) + + result = _run(tmp_path) + + assert result.import_errors == { + "etl.artifact": f"The Lang-SDK runtime did not parse {tmp_path / 'etl.artifact'} within 1.0s, " + "the limit set by [core] dagbag_import_timeout or the get_dagbag_import_timeout policy" + } + + @pytest.mark.execution_timeout(30) + @pytest.mark.parametrize("reported", [False, True], ids=["before-the-schema-version", "after-it"]) + @patch.object( + LangSDKTaskHandlerProcessorProcess, + "close", + autospec=True, + side_effect=LangSDKTaskHandlerProcessorProcess.close, + ) + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_probe_past_its_deadline_is_killed( + self, mock_parse_task_handler, mock_close, tmp_path, reported + ): + if reported: + mock_parse_task_handler.side_effect = play_runtime( + lambda request, comms: _block_until_killed(comms) + ) + else: + mock_parse_task_handler.side_effect = lambda self, **kwargs: threading.Event().wait() + + result = _run(tmp_path, deadline=time.monotonic() + 1) + + assert result.import_errors == { + "etl.artifact": f"The Lang-SDK runtime did not parse {tmp_path / 'etl.artifact'} by its deadline" + } + [proc] = [c.args[0] for c in mock_close.call_args_list] + assert proc._exit_code == -signal.SIGKILL + + @pytest.mark.execution_timeout(30) + @patch("airflow.settings.get_dagbag_import_timeout", autospec=True, return_value=60) + def test_the_deadline_holds_after_the_runtime_exits(self, mock_timeout, tmp_path): + pid_file = tmp_path / "leftover.pid" + # The runtime exits, and the process it leaves outside its process group keeps its output open. + try: + with patch.object( + LangSDKTaskHandlerProcessorProcess, + "close", + autospec=True, + side_effect=LangSDKTaskHandlerProcessorProcess.close, + ) as mock_close: + result = _run( + tmp_path, + deadline=time.monotonic() + 1, + argv=[sys.executable, "-c", _LEAVE_A_CHILD_OUTSIDE_THE_GROUP, os.fspath(pid_file)], + ) + finally: + if pid_file.exists(): + with contextlib.suppress(ProcessLookupError): + os.kill(int(pid_file.read_text()), signal.SIGKILL) + [proc] = [c.args[0] for c in mock_close.call_args_list] + + assert result.import_errors == { + "etl.artifact": f"The Lang-SDK runtime did not parse {tmp_path / 'etl.artifact'} by its deadline" + } + assert proc._exit_code == 0 + assert not proc._open_sockets + + @pytest.mark.execution_timeout(30) + @patch("airflow.settings.get_dagbag_import_timeout", autospec=True, return_value=1) + def test_the_import_timeout_applies_when_it_ends_before_the_deadline(self, mock_timeout, tmp_path): + result = _run(tmp_path, deadline=time.monotonic() + 30, argv=["/bin/sh", "-c", "exec sleep 60"]) + + assert result.import_errors == { + "etl.artifact": f"The Lang-SDK runtime did not parse {tmp_path / 'etl.artifact'} within 1.0s, " + "the limit set by [core] dagbag_import_timeout or the get_dagbag_import_timeout policy" + } + + @pytest.mark.execution_timeout(30) + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_an_answer_received_before_the_deadline_is_kept(self, mock_parse_task_handler, tmp_path): + def reply(request, comms): + comms.send(_reply_with(import_errors={"etl.artifact": "handler registry failed"})(request, comms)) + _block_until_killed(comms) + + mock_parse_task_handler.side_effect = play_runtime(reply) + + result = _run(tmp_path, deadline=time.monotonic() + 1) + + assert result.import_errors == {"etl.artifact": "handler registry failed"} + + @pytest.mark.execution_timeout(30) + @conf_vars({("dag_processor", "dag_file_processor_timeout"): "1"}) + @patch.object( + FakeCoordinator, + "_build_parse_task_handler_command", + autospec=True, + side_effect=lambda self, *, path: threading.Event().wait(), + ) + def test_the_dag_file_processor_timeout_applies_until_the_import_timeout_is_reported( + self, mock_build_command, tmp_path + ): + result = _run(tmp_path) + + assert result.import_errors == { + "etl.artifact": f"The Lang-SDK runtime did not parse {tmp_path / 'etl.artifact'} within 1.0s, " + "the limit set by [dag_processor] dag_file_processor_timeout" + } + + +def _make_process(**kwargs) -> LangSDKTaskHandlerProcessorProcess: + kwargs.setdefault("process", MagicMock(spec=PsutilTracker)) + kwargs.setdefault("logger_filehandle", MagicMock(spec=BinaryIO)) + return LangSDKTaskHandlerProcessorProcess( + id=uuid.uuid4(), + pid=1, + stdin=MagicMock(spec=socket.socket), + process_log=MagicMock(spec=FilteringBoundLogger), + selector=MagicMock(spec=selectors.BaseSelector), + bundle_name="task-handlers", + dag_file_rel_path="etl.artifact", + coordinator="fake", + listeners={}, + parse_request=TaskHandlerParseRequest( + file="/b/etl.artifact", bundle_path=Path("/b"), bundle_name="task-handlers" + ), + **kwargs, + ) + + +def _build_result(*task_ids: str) -> TaskHandlerParsingResult: + return TaskHandlerParsingResult( + fileloc="/b/etl.artifact", + task_handlers={ + "etl": [ + TaskHandlerDeclaration(task_id=task_id, binding="positional", params=[]) + for task_id in task_ids + ] + }, + ) + + +@patch.object(LangSDKTaskHandlerProcessorProcess, "send_msg", autospec=True) +def test_the_first_parse_result_wins(mock_send_msg): + proc = _make_process() + + proc._handle_request(_build_result("first"), structlog.get_logger(), 1) + proc._handle_request(_build_result("second"), structlog.get_logger(), 2) + + assert _get_task_ids(proc.parsing_result) == ["first"] + assert mock_send_msg.call_args.kwargs["error"].detail == { + "message": "A parse result was already received" + } + + +@pytest.mark.parametrize( + "invalid_frame", + [ + pytest.param(bytes.fromhex("00000003c1c1c1"), id="does-not-decode"), + pytest.param( + _RequestFrame( + id=2, + body={ + "type": "TaskHandlerParsingResult", + "fileloc": "/b/etl.artifact", + "task_handlers": "none", + }, + ).as_bytes(), + id="does-not-validate", + ), + ], +) +@patch.object(LangSDKTaskHandlerProcessorProcess, "_kill_runtime", autospec=True) +@patch.object(LangSDKTaskHandlerProcessorProcess, "send_msg", autospec=True) +def test_an_invalid_message_after_the_parse_result_keeps_it(mock_send_msg, mock_kill_runtime, invalid_frame): + proc = _make_process() + runtime, conn = socket.socketpair() + with runtime, conn: + proc._register_comm(conn) + read_frame, _ = proc.selector.register.call_args.args[2] + runtime.sendall(_RequestFrame(id=1, body=_build_result("extract").model_dump(mode="json")).as_bytes()) + assert read_frame(conn) + runtime.sendall(invalid_frame) + assert not read_frame(conn) + + assert _get_task_ids(proc.parsing_result) == ["extract"] + assert proc.parsing_result.import_errors is None + proc.process_log.warning.assert_called_with( + "Ignoring an invalid message from the Lang-SDK runtime after its parse result", error=ANY + ) + mock_kill_runtime.assert_called_once_with(proc) + + +@patch.object(LangSDKTaskHandlerProcessorProcess, "send_msg", autospec=True) +def test_the_schema_version_is_reported_once(mock_send_msg): + proc = _make_process() + + proc._handle_request( + LangSDKRuntimeSchemaVersion(schema_version=TASK_HANDLER_PARSING_SCHEMA_VERSION), + structlog.get_logger(), + 1, + ) + proc._handle_request(LangSDKRuntimeSchemaVersion(schema_version=None), structlog.get_logger(), 2) + + assert proc._runtime_schema_version == TASK_HANDLER_PARSING_SCHEMA_VERSION + assert mock_send_msg.call_args.kwargs["error"].detail["message"] == "Unhandled request" + + +@patch.object(LangSDKTaskHandlerProcessorProcess, "_signal_subprocess", autospec=True) +def test_killing_a_runtime_that_does_not_exit_waits_a_bounded_time(mock_signal): + process = MagicMock(spec=psutil.Process) + process.wait.side_effect = psutil.TimeoutExpired(_EXIT_GRACE_PERIOD) + proc = _make_process(process=PsutilTracker(process)) + + proc._kill_runtime() + + assert proc._exit_code is None + mock_signal.assert_called_once_with(proc, signal.SIGKILL) + process.wait.assert_called_once_with(_EXIT_GRACE_PERIOD) + + +@pytest.mark.parametrize(("pid_reused", "killed"), [(False, True), (True, False)]) +@patch("airflow.dag_processing.task_handler_processor.os.killpg", autospec=True) +@patch("airflow.dag_processing.task_handler_processor.psutil.pid_exists", autospec=True) +def test_close_kills_what_an_exited_runtime_left_once(mock_pid_exists, mock_killpg, pid_reused, killed): + mock_pid_exists.return_value = pid_reused + logger_filehandle = MagicMock(spec=BinaryIO) + proc = _make_process(new_process_group=True, logger_filehandle=logger_filehandle) + proc._exit_code = 0 + + proc.close() + proc.close() + + assert mock_killpg.call_args_list == ([call(1, signal.SIGKILL)] if killed else []) + logger_filehandle.close.assert_called() diff --git a/generated/known_sdk_imports_in_core.txt b/generated/known_sdk_imports_in_core.txt index e8c6aabc037f0..91362f4aabb8e 100644 --- a/generated/known_sdk_imports_in_core.txt +++ b/generated/known_sdk_imports_in_core.txt @@ -8,6 +8,7 @@ airflow-core/src/airflow/dag_processing/importer_routing.py::1 airflow-core/src/airflow/dag_processing/lang_sdk_processor.py::6 airflow-core/src/airflow/dag_processing/manager.py::5 airflow-core/src/airflow/dag_processing/processor.py::14 +airflow-core/src/airflow/dag_processing/task_handler_processor.py::7 airflow-core/src/airflow/exceptions.py::1 airflow-core/src/airflow/executors/base_executor.py::3 airflow-core/src/airflow/jobs/triggerer_job_runner.py::18 diff --git a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py index feed0da740b04..b6b26eceaf2b4 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py +++ b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py @@ -28,6 +28,7 @@ from __future__ import annotations import contextlib +import ctypes import ipaddress import itertools import os @@ -36,9 +37,10 @@ import socket import stat import subprocess +import sys import time from pathlib import PurePosixPath -from typing import TYPE_CHECKING, NoReturn, TypeVar, cast +from typing import TYPE_CHECKING, Final, NoReturn, TypeVar, cast import attrs import psutil @@ -64,11 +66,27 @@ from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: SDK002 from airflow.sdk.api.client import Client from airflow.sdk.api.datamodels._generated import TaskInstance + from airflow.sdk.coordinators._bundle_metadata import ResolvedBundle Tracked = TypeVar("Tracked", socket.socket, subprocess.Popen) log: FilteringBoundLogger = structlog.get_logger(logger_name="coordinators.subprocess") +# The first supervisor schema version whose runtime can answer a task handler parse request. +TASK_HANDLER_PARSING_SCHEMA_VERSION: Final = "2026-10-30" + + +def supports_task_handler_parsing(schema_version: str) -> bool: + """ + Return whether a runtime of supervisor schema *schema_version* can answer a task handler parse request. + + :raises ValueError: when this Task SDK does not know *schema_version*. + """ + # Dated versions order as strings; SchemaVersionMigrator compares them the same way. + return ( + get_schema_version_migrator().resolve_version(schema_version) >= TASK_HANDLER_PARSING_SCHEMA_VERSION + ) + def _start_server() -> socket.socket: server = socket.socket() @@ -200,6 +218,28 @@ def _build_runtime_env() -> dict[str, str]: } +# From +_PR_SET_PDEATHSIG = 1 + + +def _set_parent_death_signal() -> None: + """ + On Linux, have the kernel SIGKILL this process when the thread that forked it exits. + + The setting survives the exec, so the runtime dies with the process that reads its sockets, + even when that process is killed without the chance to kill the runtime first. Processes the + runtime starts do not inherit the setting. + """ + if sys.platform != "linux": + return + try: + libc = ctypes.CDLL(None, use_errno=True) + if libc.prctl(_PR_SET_PDEATHSIG, signal.SIGKILL, 0, 0, 0) != 0: + log.warning("Failed to set PR_SET_PDEATHSIG", errno=ctypes.get_errno()) + except Exception: + log.warning("Unable to set PR_SET_PDEATHSIG", exc_info=True) + + def _accept_connections( servers: dict[str, socket.socket], drains: dict[str, socket.socket], @@ -490,6 +530,9 @@ class SubprocessCoordinator(BaseCoordinator): connections, draining startup output, and tearing everything down on failure — is handled here. + A subclass that also answers task handler parse requests implements + :meth:`_find_task_handler_artifact` and :meth:`_build_parse_task_handler_command`. + :param task_startup_timeout: Maximum time the coordinator waits for the subprocess to connect to both servers, in seconds. The default is 10 seconds. @@ -528,10 +571,11 @@ def _resolve_artifact_bundle( return bundle def _get_scan_roots(self) -> tuple[pathlib.Path, ...]: - """Return the artifact roots resolved for the active task or Dag parse.""" + """Return the artifact roots resolved for the active task, Dag parse or task handler parse.""" if self._active_scan_roots is None: raise RuntimeError( - "_get_scan_roots requires an active task or Dag parse; call it during execute_task or parse_dag." + "_get_scan_roots requires an active task, Dag parse or task handler parse; " + "call it during execute_task, parse_dag or parse_task_handler." ) return self._active_scan_roots @@ -611,11 +655,89 @@ def parse_dag( _set_close_on_exec_above_stderr() os.execvpe(argv[0], argv, _build_runtime_env()) + def _find_task_handler_artifact(self, *, bundle_path: pathlib.Path, dag_id: str) -> ResolvedBundle: + """ + Return the artifact that a stub task of *dag_id* runs, found the way :meth:`execute_task` finds it. + + *bundle_path* is the root of the Dag bundle that holds the artifacts. + + :raises FileNotFoundError: when a task of *dag_id* would find no artifact it can run. + :raises ValueError: when the artifact found uses a supervisor schema version this Task SDK + does not know. + """ + raise NotImplementedError(f"{type(self).__name__} does not parse task handlers") + + def _build_parse_task_handler_command(self, *, path: pathlib.Path) -> tuple[list[str], str | None]: + """ + Build the command that reports the task handlers at *path* and resolve its wire-schema version. + + Subclasses can retrieve the root of the Dag bundle holding *path* with :meth:`_get_scan_roots`. + The contract is that of :meth:`_build_execute_task_command`: *command* MUST NOT include the + ``--comm`` / ``--logs`` flags. + """ + raise NotImplementedError(f"{type(self).__name__} does not parse task handlers") + + def parse_task_handler( + self, + *, + path: pathlib.Path, + bundle_path: pathlib.Path, + comm_address: tuple[str, int], + logs_address: tuple[str, int], + report_schema_version: Callable[[str | None], None], + ) -> NoReturn: + """ + Replace the current process with the runtime that reports the task handlers of the artifact at *path*. + + Call this in a child process whose standard streams are already set up; the runtime + inherits them and no other file descriptor. The command and its supervisor wire-schema + version are resolved against *bundle_path*, and the version is passed to + *report_schema_version* just before the exec. The runtime connects back to + *comm_address* and *logs_address*. On Linux the runtime dies with the thread that forked + the calling process, and it is not started once that parent has exited. A failure after the + version is reported leaves this process set to die with its parent and with the default + SIGPIPE and SIGXFSZ handling, so the caller should exit. + + :raises Exception: when the command cannot be resolved or the runtime's supervisor schema + version is unknown or older than :data:`TASK_HANDLER_PARSING_SCHEMA_VERSION`; the process + is then unchanged. + :raises RuntimeError: when the parent has exited. + :raises OSError: when the exec fails. + """ + with self._set_scan_roots([bundle_path]): + command, schema_version = self._build_parse_task_handler_command(path=path) + if schema_version is not None and not supports_task_handler_parsing(schema_version): + raise ValueError( + f"{path} uses supervisor schema {schema_version}, " + "which cannot answer a task handler parse request" + ) + argv = [ + *command, + f"--comm={comm_address[0]}:{comm_address[1]}", + f"--logs={logs_address[0]}:{logs_address[1]}", + ] + # Taken before the report, which fails if the parent has already exited. A different parent once the + # death signal is set means the parent exited in between, and the kernel does not send it for that. + parent_pid = os.getppid() + report_schema_version(schema_version) + # Python ignores these at startup and exec keeps ignored signals; subprocess.Popen resets + # them the same way for the task runtime. + for name in ("SIGPIPE", "SIGXFSZ"): + if (sig := getattr(signal, name, None)) is not None: + signal.signal(sig, signal.SIG_DFL) + _set_parent_death_signal() + if os.getppid() != parent_pid: + raise RuntimeError("The process that started the runtime has exited") + _set_close_on_exec_above_stderr() + os.execvpe(argv[0], argv, _build_runtime_env()) + @contextlib.contextmanager def _set_scan_roots(self, roots: Sequence[pathlib.Path]): - """Expose *roots* to the command builder for the duration of the task or Dag parse.""" + """Expose *roots* to the command builder for the task, Dag parse or task handler parse.""" if self._active_scan_roots is not None: - raise RuntimeError("SubprocessCoordinator.execute_task and parse_dag are not re-entrant.") + raise RuntimeError( + "SubprocessCoordinator.execute_task, parse_dag and parse_task_handler are not re-entrant." + ) self._active_scan_roots = tuple(roots) try: yield diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 40a43208269c5..5152abc697815 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -390,8 +390,9 @@ def _reopen_std_io_handles(child_stdin, child_stdout, child_stderr): del sock # We open the socket/fd as binary, and then pass it to a TextIOWrapper so that it looks more like a - # normal sys.stdout etc. - binary = os.fdopen(fd, mode + "b") + # normal sys.stdout etc. The handle does not own the fd: a child this process supervises replaces + # it after dup2 onto the same fd, and closing it there would close the child's new stream. + binary = os.fdopen(fd, mode + "b", closefd=False) handle = io.TextIOWrapper(binary, line_buffering=True) setattr(sys, handle_name, handle) diff --git a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py index 497212f0ef717..c61214b10ca4f 100644 --- a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py +++ b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py @@ -49,6 +49,7 @@ _ResourceTracker, _start_server, log, + supports_task_handler_parsing, ) from airflow.sdk.execution_time.coordinator import BaseCoordinator, InvalidCoordinatorError, TaskLaunchError from airflow.sdk.execution_time.supervisor import ActivitySubprocess @@ -1467,6 +1468,13 @@ def _build_parse_dag_command(self, *, path): return [*self.command, os.fspath(path)], self.schema_version +@attrs.define(kw_only=True) +class _TaskHandlerParsingCoordinator(_StubSubprocessCoordinator): + def _build_parse_task_handler_command(self, *, path): + self.recorded_roots.append(list(self._get_scan_roots())) + return [*self.command, os.fspath(path)], self.schema_version + + class TestParseDag: def test_build_parse_dag_command_default_raises(self, tmp_path): with pytest.raises(NotImplementedError): @@ -1522,3 +1530,182 @@ def test_rejects_an_unknown_schema_version_before_reporting(self, mock_execvpe, assert reported == [] mock_execvpe.assert_not_called() + + +def _parse_task_handler(coordinator: SubprocessCoordinator, tmp_path: pathlib.Path, reported: list) -> None: + coordinator.parse_task_handler( + path=tmp_path / "handlers.artifact", + bundle_path=tmp_path, + comm_address=("127.0.0.1", 1001), + logs_address=("127.0.0.1", 1002), + report_schema_version=reported.append, + ) + + +def _is_running(process: psutil.Process) -> bool: + try: + return process.status() != psutil.STATUS_ZOMBIE + except psutil.NoSuchProcess: + return False + + +class TestParseTaskHandler: + def test_find_task_handler_artifact_default_raises(self, tmp_path): + coordinator = _StubSubprocessCoordinator(command=["x"]) + + with pytest.raises( + NotImplementedError, match="_StubSubprocessCoordinator does not parse task handlers" + ): + coordinator._find_task_handler_artifact(bundle_path=tmp_path, dag_id="dag") + + def test_build_parse_task_handler_command_default_raises(self, tmp_path): + coordinator = _StubSubprocessCoordinator(command=["x"]) + + with pytest.raises( + NotImplementedError, match="_StubSubprocessCoordinator does not parse task handlers" + ): + coordinator._build_parse_task_handler_command(path=tmp_path / "handlers.artifact") + + @patch("airflow.sdk.coordinators._subprocess._set_close_on_exec_above_stderr", autospec=True) + @patch("airflow.sdk.coordinators._subprocess._set_parent_death_signal", autospec=True) + @patch("airflow.sdk.coordinators._subprocess.signal.signal", autospec=True) + @patch( + "airflow.sdk.coordinators._subprocess.os.execvpe", autospec=True, side_effect=OSError("exec failed") + ) + def test_resolves_the_command_under_the_bundle_root_then_execs_it( + self, mock_execvpe, mock_signal, mock_death_signal, mock_close_on_exec, tmp_path + ): + coordinator = _TaskHandlerParsingCoordinator(command=["runtime"], schema_version="2026-10-30") + reported: list[str | None] = [] + + with pytest.raises(OSError, match="exec failed"): + _parse_task_handler(coordinator, tmp_path, reported) + + assert coordinator.recorded_roots == [[tmp_path]] + assert coordinator._active_scan_roots is None + assert reported == ["2026-10-30"] + argv = [ + "runtime", + os.fspath(tmp_path / "handlers.artifact"), + "--comm=127.0.0.1:1001", + "--logs=127.0.0.1:1002", + ] + mock_execvpe.assert_called_once_with("runtime", argv, ANY) + assert "AIRFLOW__LOGGING__LOGGING_LEVEL" in mock_execvpe.call_args.args[2] + restored = {c.args[0] for c in mock_signal.call_args_list if c.args[1] == signal.SIG_DFL} + assert {signal.SIGPIPE, signal.SIGXFSZ} <= restored + mock_death_signal.assert_called_once_with() + mock_close_on_exec.assert_called_once_with() + + @patch("airflow.sdk.coordinators._subprocess._set_close_on_exec_above_stderr", autospec=True) + @patch("airflow.sdk.coordinators._subprocess._set_parent_death_signal", autospec=True) + @patch("airflow.sdk.coordinators._subprocess.signal.signal", autospec=True) + @patch("airflow.sdk.coordinators._subprocess.os.getppid", autospec=True, side_effect=[4242, 1]) + @patch("airflow.sdk.coordinators._subprocess.os.execvpe", autospec=True) + def test_does_not_exec_once_the_parent_has_exited( + self, mock_execvpe, mock_getppid, mock_signal, mock_death_signal, mock_close_on_exec, tmp_path + ): + coordinator = _TaskHandlerParsingCoordinator(command=["runtime"]) + steps = MagicMock() + steps.attach_mock(mock_getppid, "getppid") + steps.attach_mock(mock_death_signal, "death_signal") + + with pytest.raises(RuntimeError, match="The process that started the runtime has exited"): + coordinator.parse_task_handler( + path=tmp_path / "handlers.artifact", + bundle_path=tmp_path, + comm_address=("127.0.0.1", 1001), + logs_address=("127.0.0.1", 1002), + report_schema_version=steps.report, + ) + + # The parent is read before the death signal is set, so an exit in between is seen. + assert steps.mock_calls == [call.getppid(), call.report(None), call.death_signal(), call.getppid()] + mock_execvpe.assert_not_called() + + @patch("airflow.sdk.coordinators._subprocess._set_close_on_exec_above_stderr", autospec=True) + @patch("airflow.sdk.coordinators._subprocess._set_parent_death_signal", autospec=True) + @patch("airflow.sdk.coordinators._subprocess.signal.signal", autospec=True) + @patch("airflow.sdk.coordinators._subprocess.os.execvpe", autospec=True) + def test_does_not_start_a_runtime_that_cannot_answer( + self, mock_execvpe, mock_signal, mock_death_signal, mock_close_on_exec, tmp_path + ): + coordinator = _TaskHandlerParsingCoordinator(command=["runtime"], schema_version="2026-06-16") + reported: list[str | None] = [] + + with pytest.raises(ValueError, match="cannot answer a task handler parse request"): + _parse_task_handler(coordinator, tmp_path, reported) + + assert reported == [] + mock_execvpe.assert_not_called() + + @patch("airflow.sdk.coordinators._subprocess._set_close_on_exec_above_stderr", autospec=True) + @patch("airflow.sdk.coordinators._subprocess._set_parent_death_signal", autospec=True) + @patch("airflow.sdk.coordinators._subprocess.signal.signal", autospec=True) + @patch("airflow.sdk.coordinators._subprocess.os.execvpe", autospec=True) + def test_rejects_an_unknown_schema_version_before_reporting( + self, mock_execvpe, mock_signal, mock_death_signal, mock_close_on_exec, tmp_path + ): + coordinator = _TaskHandlerParsingCoordinator(command=["runtime"], schema_version="1999-01-01") + reported: list[str | None] = [] + + with pytest.raises(ValueError, match="'1999-01-01' not found"): + _parse_task_handler(coordinator, tmp_path, reported) + + assert reported == [] + mock_execvpe.assert_not_called() + + @pytest.mark.skipif(sys.platform != "linux", reason="PR_SET_PDEATHSIG is Linux-only") + def test_the_runtime_dies_with_the_process_that_started_it(self, tmp_path): + coordinator = _TaskHandlerParsingCoordinator(command=["/bin/sh", "-c", "exec sleep 60"]) + read_fd, write_fd = os.pipe() + parent_pid = os.fork() + if parent_pid == 0: + # The process reading the runtime's sockets: it starts the runtime, reports its pid and waits. + try: + os.close(read_fd) + runtime_pid = os.fork() + if runtime_pid == 0: + try: + _parse_task_handler(coordinator, tmp_path, []) + finally: + os._exit(1) + os.write(write_fd, str(runtime_pid).encode()) + os.close(write_fd) + os.waitpid(runtime_pid, 0) + finally: + os._exit(0) + + os.close(write_fd) + with os.fdopen(read_fd) as reader: + runtime = psutil.Process(int(reader.read())) + try: + # Polls: the runtime is not a child to wait for, and time_machine cannot reach another process. + deadline = time.monotonic() + 30 + while runtime.name() != "sleep": + assert time.monotonic() < deadline, "the runtime did not start" + time.sleep(0.05) + + os.kill(parent_pid, signal.SIGKILL) + os.waitpid(parent_pid, 0) + + deadline = time.monotonic() + 10 + while _is_running(runtime): + assert time.monotonic() < deadline, "the runtime outlived the process that started it" + time.sleep(0.05) + finally: + with contextlib.suppress(psutil.NoSuchProcess): + runtime.kill() + + +class TestSupportsTaskHandlerParsing: + @pytest.mark.parametrize( + ("schema_version", "expected"), + [("2026-06-16", False), ("2026-10-30", True)], + ) + def test_supports_task_handler_parsing(self, schema_version, expected): + assert supports_task_handler_parsing(schema_version) is expected + + def test_supports_task_handler_parsing_rejects_an_unknown_version(self): + with pytest.raises(ValueError, match="not found"): + supports_task_handler_parsing("1999-01-01") 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 4fbdb76b4a68a..3b0fdbc72e0f2 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -173,6 +173,7 @@ WatchedSubprocess, _make_process_nondumpable, _remote_logging_conn, + _reopen_std_io_handles, forward_to_log, in_process_api_server, make_buffered_socket_reader, @@ -5317,6 +5318,25 @@ def on_close(sock): log_write.close() +def test_a_nested_child_keeps_its_standard_streams(): + """A child of a supervised child reopens handles its parent already reopened, and keeps them open.""" + pid = os.fork() + if pid == 0: + code = 1 + try: + for _ in ("supervised child", "its child"): + pairs = [socket.socketpair() for _ in range(3)] + _reopen_std_io_handles(*(child_end for child_end, _ in pairs)) + for fd in (0, 1, 2): + os.fstat(fd) + code = 0 + finally: + os._exit(code) + _, status = os.waitpid(pid, 0) + + assert os.waitstatus_to_exitcode(status) == 0 + + def test_reinit_supervisor_comms(monkeypatch, client_with_ti_start, caplog): def subprocess_main(): # This is run in the subprocess!