From f1070b3884ee31db8a9c0082a0b2449ba196e3a8 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Thu, 1 Oct 2026 14:52:07 +0000 Subject: [PATCH 01/24] Parse coordinator-claimed Dag files with their runtime in the Dag processor The Dag processor parses a file a coordinator's Dag importer claims with LangSDKDagFileProcessorProcess, in the new lang_sdk_processor module. Its Python child finds the coordinator, reports the runtime's schema version over fd 0 and execs the runtime, which connects back to two listeners the manager owns and answers the parse request. Each returned Dag gets its unset settings from the Airflow config and must pass validate_serialized_dag. A failed start, a missing result, an invalid frame or message, or an invalid Dag is an import error. A runtime still running 5s after its result is killed, and the result is kept. Callbacks for the file are dropped. The file's Dag source is read with its Dag importer. A Dag bag cannot parse such a file yet: its importer reports an import error, and sync_bag_to_db leaves the file's import errors to the Dag processor. --- .../src/airflow/dag_processing/dagbag.py | 19 +- .../dag_processing/importer_routing.py | 59 ++ .../dag_processing/lang_sdk_processor.py | 493 ++++++++++++++++ .../src/airflow/dag_processing/manager.py | 29 +- .../unit/dag_processing/fake_lang_sdk.py | 126 +++++ .../tests/unit/dag_processing/test_dagbag.py | 20 + .../dag_processing/test_importer_routing.py | 38 ++ .../dag_processing/test_lang_sdk_processor.py | 533 ++++++++++++++++++ .../tests/unit/dag_processing/test_manager.py | 137 ++++- generated/known_sdk_imports_in_core.txt | 1 + .../airflow/sdk/coordinators/_dag_importer.py | 22 +- .../coordinators/test_dag_importer.py | 18 +- 12 files changed, 1479 insertions(+), 16 deletions(-) create mode 100644 airflow-core/src/airflow/dag_processing/importer_routing.py create mode 100644 airflow-core/src/airflow/dag_processing/lang_sdk_processor.py create mode 100644 airflow-core/tests/unit/dag_processing/fake_lang_sdk.py create mode 100644 airflow-core/tests/unit/dag_processing/test_importer_routing.py create mode 100644 airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py diff --git a/airflow-core/src/airflow/dag_processing/dagbag.py b/airflow-core/src/airflow/dag_processing/dagbag.py index 2ec6ee4afe1d5..5d97ee9031961 100644 --- a/airflow-core/src/airflow/dag_processing/dagbag.py +++ b/airflow-core/src/airflow/dag_processing/dagbag.py @@ -33,6 +33,7 @@ from airflow._shared.timezones import timezone from airflow.configuration import conf from airflow.dag_processing.bundles.local import LocalDagBundle +from airflow.dag_processing.importer_routing import get_claiming_coordinator from airflow.exceptions import ( AirflowClusterPolicyError, AirflowClusterPolicySkipDag, @@ -583,16 +584,30 @@ def sync_bag_to_db( version_data: dict[str, Any] | None = None, session: Session = NEW_SESSION, ) -> None: - """Save attributes about list of DAG to the DB.""" + """ + Save attributes about list of DAG to the DB. + + Files that a Lang-SDK runtime parses are left out, with their import errors: the Dag processor + stores those. + """ from airflow.dag_processing.collection import update_dag_parsing_results_in_db - import_errors = {(bundle_name, rel_path): error for rel_path, error in dagbag.import_errors.items()} + def is_parsed_by_runtime(rel_path: str) -> bool: + return get_claiming_coordinator(Path(dagbag.bundle_path or "", rel_path), bundle_name) is not None + + import_errors = { + (bundle_name, rel_path): error + for rel_path, error in dagbag.import_errors.items() + if not is_parsed_by_runtime(rel_path) + } # Build the set of all files that were parsed and include files with import errors # in case they are not in parsed_definitions files_parsed = set(import_errors) if dagbag.bundle_path: for rel_path in dagbag.parsed_definitions: + if is_parsed_by_runtime(rel_path): + continue files_parsed.add((bundle_name, rel_path)) # A definition nested in an archive also clears the archive's own discovery errors. if enclosing_file := find_enclosing_file(Path(dagbag.bundle_path, rel_path)): diff --git a/airflow-core/src/airflow/dag_processing/importer_routing.py b/airflow-core/src/airflow/dag_processing/importer_routing.py new file mode 100644 index 0000000000000..956723eae26be --- /dev/null +++ b/airflow-core/src/airflow/dag_processing/importer_routing.py @@ -0,0 +1,59 @@ +# +# 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. +"""Route Dag files to the process that parses them, by the bundle's Dag importer registry.""" + +from __future__ import annotations + +import logging +import os +from pathlib import Path +from typing import TYPE_CHECKING + +from airflow.sdk.coordinators._dag_importer import CoordinatorDagImporter # noqa: SDK001 +from airflow.sdk.importers import DagImporterRegistry, get_importer_registry # noqa: SDK001 + +if TYPE_CHECKING: + from airflow.sdk.coordinators._subprocess import SubprocessCoordinator # noqa: SDK001 + +log = logging.getLogger(__name__) + + +def _get_registry(bundle_name: str | None) -> DagImporterRegistry | None: + try: + return get_importer_registry(bundle_name) + except Exception: + log.exception("Cannot build the Dag importer registry for bundle %s", bundle_name) + return None + + +def get_claiming_coordinator( + path: str | os.PathLike[str], bundle_name: str | None +) -> SubprocessCoordinator | None: + """ + Return the coordinator whose runtime parses ``path``, or ``None`` when a Python child parses it. + + A runtime parses the file when its importer is a coordinator's Dag importer. + """ + if (registry := _get_registry(bundle_name)) is None: + return None + try: + importer = registry.get_importer(Path(path)) + except Exception: + log.exception("Cannot load the Dag importer for %s", path) + return None + return importer.coordinator if isinstance(importer, CoordinatorDagImporter) else None diff --git a/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py b/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py new file mode 100644 index 0000000000000..877af756ed6da --- /dev/null +++ b/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py @@ -0,0 +1,493 @@ +# 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. +"""Parse a native Lang-SDK Dag file with the runtime of the coordinator whose Dag importer claims it.""" + +from __future__ import annotations + +import functools +import os +import selectors +import signal +import time +from pathlib import Path +from socket import socket +from typing import TYPE_CHECKING, Annotated, ClassVar, Literal, cast, get_args + +import attrs +import msgspec +from pydantic import BaseModel, Field, TypeAdapter +from uuid6 import uuid7 + +from airflow.dag_processing.importer_routing import get_claiming_coordinator +from airflow.dag_processing.processor import ( + BaseDagFileProcessorProcess, + DagFileParseRequest, + DagFileParsingResult, + ToManager, +) +from airflow.exceptions import DeserializationError +from airflow.sdk.coordinators._subprocess import _is_connection_from_pid, _start_server +from airflow.sdk.execution_time import supervisor +from airflow.sdk.execution_time.comms import CommsDecoder, ErrorResponse, MaskSecret, _RequestFrame +from airflow.sdk.execution_time.supervisor import ( + ResponseSent, + length_prefixed_frame_reader, + make_buffered_socket_reader, + process_log_messages_from_subprocess, + register_request_method, +) +from airflow.sdk.importers import DagSourceCode, FilesystemDagDefinition +from airflow.serialization.serialized_objects import DagSerialization, LazyDeserializedDAG + +if TYPE_CHECKING: + from collections.abc import Generator + + from structlog.typing import FilteringBoundLogger + + from airflow.sdk.api.client import Client + from airflow.sdk.execution_time.supervisor import RequestHandler, RequestResult + from airflow.typing_compat import Self + +# How long a runtime may keep running after its parse result, as Node does while a handle stays open. +_EXIT_GRACE_PERIOD = 5.0 + + +# StartLangSDKRuntime and LangSDKRuntimeSchemaVersion pass only between the manager and its forked +# child before the exec, so they are not part of the supervisor schema the runtimes speak. + + +class StartLangSDKRuntime(BaseModel): + """Ask the parse child to exec the runtime that parses *file*.""" + + file: str + bundle_path: Path + bundle_name: str + dag_file_rel_path: str + comm_address: tuple[str, int] + logs_address: tuple[str, int] + type: Literal["StartLangSDKRuntime"] = "StartLangSDKRuntime" + + +class LangSDKRuntimeSchemaVersion(BaseModel): + """The supervisor schema version of the runtime the parse child is about to exec.""" + + schema_version: str | None + type: Literal["LangSDKRuntimeSchemaVersion"] = "LangSDKRuntimeSchemaVersion" + + +def _start_runtime_entrypoint() -> None: + """Exec the runtime that parses the file 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[StartLangSDKRuntime, LangSDKRuntimeSchemaVersion | DagFileParsingResult]( + socket=socket(fileno=os.dup(0)), + body_decoder=TypeAdapter(StartLangSDKRuntime), + ) + devnull = os.open(os.devnull, os.O_RDONLY) + os.dup2(devnull, 0) + os.close(devnull) + + msg = comms._get_response() + if not isinstance(msg, StartLangSDKRuntime): + raise RuntimeError(f"Required first message to be a StartLangSDKRuntime, it was {msg}") + + def report_schema_version(schema_version: str | None) -> None: + comms.send(LangSDKRuntimeSchemaVersion(schema_version=schema_version)) + + try: + coordinator = get_claiming_coordinator(msg.file, msg.bundle_name) + if coordinator is None: + raise RuntimeError(f"No coordinator's Dag importer claims {msg.file}") + coordinator.parse_dag( + 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( + DagFileParsingResult( + fileloc=msg.file, + serialized_dags=[], + import_errors={ + msg.dag_file_rel_path: f"Cannot start the Lang-SDK runtime: {type(e).__name__}: {e}" + }, + ) + ) + + +_Channel = Literal["comm", "logs"] + + +@attrs.define(kw_only=True) +class LangSDKDagFileProcessorProcess(BaseDagFileProcessorProcess): + """ + Parse a native Lang-SDK Dag file with its coordinator's runtime. + + 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 + ``DagFileParseRequest`` itself, so the request is sent once it has connected. + """ + + client: Client | None = None # type: ignore[assignment] + """Answers the runtime's requests; without one, as in a Dag bag, those that need it get an error.""" + + decoder = TypeAdapter( + Annotated[LangSDKRuntimeSchemaVersion | get_args(ToManager)[0], Field(discriminator="type")] + ) + + _listeners: dict[_Channel, socket] + _parse_request: DagFileParseRequest + _runtime_schema_version: str | 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) + + @classmethod + def start( # type: ignore[override] + cls, + *, + path: str | os.PathLike[str], + bundle_path: Path, + bundle_name: str, + dag_file_rel_path: str, + **kwargs, + ) -> Self: + listeners: dict[_Channel, socket] = {"comm": _start_server(), "logs": _start_server()} + try: + for listener in listeners.values(): + listener.setblocking(False) + parse_request = DagFileParseRequest( + file=os.fspath(path), bundle_path=bundle_path, bundle_name=bundle_name + ) + proc = super().start( + target=_start_runtime_entrypoint, + use_exec=supervisor._should_use_exec(), + new_process_group=True, + bundle_name=bundle_name, + dag_file_rel_path=dag_file_rel_path, + listeners=listeners, + parse_request=parse_request, + **kwargs, + ) + except BaseException: + for listener in listeners.values(): + listener.close() + raise + 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( + StartLangSDKRuntime( + file=parse_request.file, + bundle_path=bundle_path, + bundle_name=bundle_name, + dag_file_rel_path=dag_file_rel_path, + comm_address=listeners["comm"].getsockname()[:2], + logs_address=listeners["logs"].getsockname()[:2], + ), + request_id=0, + ) + return proc + + @classmethod + def run( + cls, + *, + path: str | os.PathLike[str], + bundle_path: Path, + bundle_name: str, + dag_file_rel_path: str, + timeout: float | None, + logger: FilteringBoundLogger, + ) -> DagFileParsingResult: + """ + Parse *path* outside the Dag processor and wait for the result. + + There is no API client, so each request of the runtime that needs one gets an error. + + :raises TimeoutError: if the parse does not finish within *timeout* seconds. The runtime is + killed. + """ + with selectors.DefaultSelector() as selector: + proc = cls.start( + id=uuid7(), + path=path, + bundle_path=bundle_path, + bundle_name=bundle_name, + dag_file_rel_path=dag_file_rel_path, + selector=selector, + logger=logger, + ) + try: + while not proc.is_ready: + wait = 0.1 + if timeout is not None: + if (remaining := proc.start_time + timeout - time.monotonic()) <= 0: + raise TimeoutError( + f"The Lang-SDK runtime did not parse {os.fspath(path)} within {timeout}s" + ) + wait = min(wait, remaining) + proc._service_subprocess(max_wait_time=wait) + except BaseException: + proc._kill_runtime() + raise + finally: + proc.close() + return cast("DagFileParsingResult", 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: + # A frame that does not decode would otherwise escape the Dag processor's selector loop. + try: + return read_frame(sock) + except msgspec.DecodeError as e: + 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. + """ + requests = self.handle_requests(self.process_log) + next(requests) + while True: + frame = yield + try: + self.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 = DagFileParsingResult( + fileloc=self._parse_request.file, + serialized_dags=[], + 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._schema_version_reported = True + return None, {} + + def _handle_parsing_result( + self, msg: DagFileParsingResult, 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 + import_errors = dict(msg.import_errors or {}) + serialized_dags = [] + for dag in msg.serialized_dags: + DagSerialization.fill_config_defaults(dag.data) + try: + DagSerialization.validate_serialized_dag(dag.data) + except DeserializationError as e: + message = f"Cannot load the serialized Dag: {e}" + self.process_log.warning(message) + previous = import_errors.get(self.dag_file_rel_path) + import_errors[self.dag_file_rel_path] = f"{previous}\n{message}" if previous else message + continue + serialized_dags.append(dag) + self.parsing_result = msg.model_copy( + update={ + "serialized_dags": serialized_dags, + "import_errors": import_errors or None, + "dag_source_codes": self._read_dag_source_codes(serialized_dags), + } + ) + self._parsing_result_monotonic = time.monotonic() + return None, {} + + def _read_dag_source_codes(self, serialized_dags: list[LazyDeserializedDAG]) -> dict[str, DagSourceCode]: + """ + Read the file's source with its Dag importer, for the fileloc of each Dag. + + A binary artifact cannot be read as text, so a source that cannot be read is a placeholder. + """ + if not serialized_dags: + return {} + file = self._parse_request.file + try: + coordinator = get_claiming_coordinator(file, self.bundle_name) + if coordinator is None or (importer := coordinator.get_dag_importer()) is None: + raise RuntimeError(f"No coordinator's Dag importer claims {file}") + source = importer.get_source_code(FilesystemDagDefinition(Path(file))) + except Exception as e: + self.process_log.warning("Cannot read the Dag source", fileloc=file, error=str(e)) + source = DagSourceCode(f"Cannot read the source of {self.dag_file_rel_path}: {e}", "text") + return {dag.data["dag"].get("fileloc", file): source for dag in serialized_dags} + + _request_handlers: ClassVar[dict[type[BaseModel], RequestHandler[LangSDKDagFileProcessorProcess]]] = { + **BaseDagFileProcessorProcess._common_request_handlers, + **dict([register_request_method(LangSDKRuntimeSchemaVersion, _handle_runtime_schema_version)]), + } + + def _handle_request(self, msg, log: FilteringBoundLogger, req_id: int) -> None: + if self.client is None and not isinstance( + msg, (DagFileParsingResult, LangSDKRuntimeSchemaVersion, MaskSecret) + ): + self.send_msg( + None, + request_id=req_id, + error=ErrorResponse( + detail={"message": f"{type(msg).__name__} is answered only in the Dag processor"} + ), + ) + return + super()._handle_request(msg, log, req_id) + + @property + def is_ready(self) -> bool: + self._verify_connections() + if ( + self._parsing_result_monotonic is not None + and self._exit_code is None + and time.monotonic() - self._parsing_result_monotonic > _EXIT_GRACE_PERIOD + ): + self.process_log.warning("The Lang-SDK runtime did not exit after its parse result; killing it") + self._kill_runtime() + 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 _kill_runtime(self) -> None: + """Kill the runtime and wait for it, without servicing its sockets, whose handler may have failed.""" + if self._exit_code is not None: + return + try: + self._signal_subprocess(signal.SIGKILL) + self._exit_code = self._process.wait(timeout=None) + except (self._process.ProcessNotFound, ProcessLookupError): + self._exit_code = -1 + + def close(self) -> None: + # A listener has nothing to drain, and cleanup would call its accept handler forever. + self._close_listeners() + super().close() diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index 9827cc590829c..a1cd138e672c6 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -54,7 +54,13 @@ ) from airflow.dag_processing.bundles.manager import DagBundlesManager from airflow.dag_processing.collection import update_dag_parsing_results_in_db -from airflow.dag_processing.processor import DagFileParsingResult, DagFileProcessorProcess +from airflow.dag_processing.importer_routing import get_claiming_coordinator +from airflow.dag_processing.lang_sdk_processor import LangSDKDagFileProcessorProcess +from airflow.dag_processing.processor import ( + BaseDagFileProcessorProcess, + DagFileParsingResult, + DagFileProcessorProcess, +) from airflow.models.asset import remove_references_to_deleted_dags from airflow.models.dag import DagModel from airflow.models.dagbag import DagPriorityParsingRequest @@ -262,7 +268,7 @@ class DagFileProcessorManager(LoggingMixin): _multi_team: bool = attrs.field(factory=lambda: conf.getboolean("core", "multi_team"), init=False) _bundle_name_to_team_name: dict[str, str | None] = attrs.field(factory=dict, init=False) - _processors: dict[DagFileInfo, DagFileProcessorProcess] = attrs.field(factory=dict, init=False) + _processors: dict[DagFileInfo, BaseDagFileProcessorProcess] = attrs.field(factory=dict, init=False) _parsing_start_time: float | None = attrs.field(default=None, init=False) _num_run: int = attrs.field(default=0, init=False) @@ -1277,7 +1283,7 @@ def terminate_orphan_processes(self, present: set[DagFileInfo]): def handle_parsing_result( self, file: DagFileInfo, - proc: DagFileProcessorProcess, + proc: BaseDagFileProcessorProcess, *, session: Session = NEW_SESSION, ) -> None: @@ -1461,19 +1467,17 @@ def client(self) -> Client: client.base_url = "http://in-process.invalid./" return client - def _create_process(self, dag_file: DagFileInfo) -> DagFileProcessorProcess: + def _create_process(self, dag_file: DagFileInfo) -> BaseDagFileProcessorProcess: id = uuid7() callback_to_execute_for_file = self._callback_to_execute.pop(dag_file, []) logger, logger_filehandle = self._get_logger_for_dag_file(dag_file) - - return DagFileProcessorProcess.start( + kwargs: dict[str, Any] = dict( id=id, path=dag_file.absolute_path, bundle_path=cast("Path", dag_file.bundle_path), bundle_name=dag_file.bundle_name, dag_file_rel_path=str(dag_file.rel_path), - callbacks=callback_to_execute_for_file, selector=self.selector, logger=logger, logger_filehandle=logger_filehandle, @@ -1481,6 +1485,17 @@ def _create_process(self, dag_file: DagFileInfo) -> DagFileProcessorProcess: client=self.client, ) + if get_claiming_coordinator(dag_file.absolute_path, dag_file.bundle_name) is not None: + if callback_to_execute_for_file: + self.log.warning( + "Dropping %d callbacks for %s: Lang-SDK runtimes do not run callbacks", + len(callback_to_execute_for_file), + dag_file.rel_path, + ) + return LangSDKDagFileProcessorProcess.start(**kwargs) + + return DagFileProcessorProcess.start(callbacks=callback_to_execute_for_file, **kwargs) + def _start_new_processes(self): """Start more processors if we have enough slots and files to process.""" bundle_to_team = self._get_team_names({file.bundle_name for file in self._file_queue}) diff --git a/airflow-core/tests/unit/dag_processing/fake_lang_sdk.py b/airflow-core/tests/unit/dag_processing/fake_lang_sdk.py new file mode 100644 index 0000000000000..9fc953e0ce310 --- /dev/null +++ b/airflow-core/tests/unit/dag_processing/fake_lang_sdk.py @@ -0,0 +1,126 @@ +# +# 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 that claims ``.native`` Dag files, and a runtime a test can play in the parse child.""" + +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 ( + DagFileParseRequest, + DagFileParsingResult, + ToDagProcessor, + ToManager, +) +from airflow.sdk.coordinators._dag_importer import CoordinatorDagImporter +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.importers import DagSourceCode, reset_importer_registry + +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): + """ + Claim ``.native`` files; the file'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_dag_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") + + def get_dag_importer(self) -> FakeCoordinatorDagImporter: + return FakeCoordinatorDagImporter(coordinator=self) + + +class FakeCoordinatorDagImporter(CoordinatorDagImporter): + artifact_suffix = ".native" + supported_extensions = [".native"] + + def get_source_code(self, definition) -> DagSourceCode: + return DagSourceCode(definition.read_text(), "fake") + + +@contextlib.contextmanager +def fake_coordinator(**kwargs: Any) -> Iterator[None]: + """ + Configure a ``FakeCoordinator``, with fresh coordinators and registries inside and after the block. + + The parse child is a bare fork even on macOS, so it sees the test's ``parse_dag`` patch and config. + """ + spec = {"fake": {"classpath": f"{__name__}.FakeCoordinator", "kwargs": kwargs}} + reset_importer_registry() + try: + with ( + conf_vars({("sdk", "coordinators"): json.dumps(spec)}), + mock.patch.object(supervisor, "_should_use_exec", return_value=False), + ): + yield + finally: + reset_importer_registry() + + +def write_native_file(path: Path, **spec: Any) -> Path: + path.write_text(json.dumps(spec)) + return path + + +def play_runtime( + reply: Callable[[DagFileParseRequest, CommsDecoder], DagFileParsingResult | None], + *, + schema_version: str | None = None, + log_lines: Sequence[dict[str, Any]] = (), +) -> Callable[..., None]: + """ + Return a ``parse_dag`` 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_dag(coordinator, *, comm_address, logs_address, report_schema_version, **kwargs) -> None: + report_schema_version(schema_version) + comm = socket.create_connection(comm_address) + logs = socket.create_connection(logs_address) + for line in log_lines: + logs.sendall(json.dumps(line).encode() + b"\n") + comms = CommsDecoder[ToDagProcessor, ToManager](socket=comm, body_decoder=TypeAdapter(ToDagProcessor)) + request = comms._get_response() + assert isinstance(request, DagFileParseRequest) + if (result := reply(request, comms)) is not None: + comms.send(result) + + return parse_dag diff --git a/airflow-core/tests/unit/dag_processing/test_dagbag.py b/airflow-core/tests/unit/dag_processing/test_dagbag.py index 5fd7d51c64609..0f7def4aad4af 100644 --- a/airflow-core/tests/unit/dag_processing/test_dagbag.py +++ b/airflow-core/tests/unit/dag_processing/test_dagbag.py @@ -46,6 +46,7 @@ from airflow.executors.executor_loader import ExecutorLoader from airflow.models.dag import DagModel from airflow.models.dagwarning import DagWarning, DagWarningType +from airflow.models.errors import ParseImportError from airflow.models.pool import Pool from airflow.models.serialized_dag import SerializedDagModel from airflow.sdk import DAG, BaseOperator @@ -64,6 +65,7 @@ from tests_common.test_utils import db from tests_common.test_utils.config import conf_vars from unit import cluster_policies +from unit.dag_processing.fake_lang_sdk import fake_coordinator, write_native_file from unit.models import TEST_DAGS_FOLDER pytestmark = pytest.mark.db_test @@ -1542,3 +1544,21 @@ def test_dagbag_no_bundle_path_no_syspath_modification(self, tmp_path): assert str(tmp_path) not in dag.description assert sys.path == syspath_before + + +def test_sync_bag_to_db_leaves_native_files_to_the_dag_processor(tmp_path, session, testing_dag_bundle): + db.clear_db_import_errors() + write_native_file(tmp_path / "dags.native") + session.add(ParseImportError(bundle_name="testing", filename="dags.native", stacktrace="stored")) + session.commit() + + with fake_coordinator(): + dagbag = DagBag(dag_folder=os.fspath(tmp_path), bundle_path=tmp_path, bundle_name="testing") + sync_bag_to_db(dagbag, "testing", None, session=session) + + assert dagbag.import_errors == { + "dags.native": "A native Lang-SDK Dag is parsed only by the Dag processor" + } + assert {(e.filename, e.stacktrace) for e in session.scalars(select(ParseImportError))} == { + ("dags.native", "stored") + } diff --git a/airflow-core/tests/unit/dag_processing/test_importer_routing.py b/airflow-core/tests/unit/dag_processing/test_importer_routing.py new file mode 100644 index 0000000000000..982fdc0cb680f --- /dev/null +++ b/airflow-core/tests/unit/dag_processing/test_importer_routing.py @@ -0,0 +1,38 @@ +# +# 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 + +from unittest import mock + +from airflow.dag_processing.importer_routing import get_claiming_coordinator + +from unit.dag_processing.fake_lang_sdk import FakeCoordinator, fake_coordinator + + +def test_get_claiming_coordinator_returns_the_coordinator_of_its_importer(tmp_path): + with fake_coordinator(): + coordinator = get_claiming_coordinator(tmp_path / "dags.native", "testing") + others = [get_claiming_coordinator(tmp_path / name, "testing") for name in ("dags.other", "dag.py")] + + assert isinstance(coordinator, FakeCoordinator) + assert others == [None, None] + + +@mock.patch("airflow.dag_processing.importer_routing._get_registry", autospec=True, return_value=None) +def test_get_claiming_coordinator_without_a_registry(mock_registry, tmp_path): + assert get_claiming_coordinator(tmp_path / "dags.native", "testing") is None diff --git a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py new file mode 100644 index 0000000000000..4a5fd758681d9 --- /dev/null +++ b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py @@ -0,0 +1,533 @@ +# +# 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 os +import selectors +import signal +import socket +import sys +import time +import uuid +from pathlib import Path +from unittest.mock import ANY, MagicMock, patch + +import psutil +import pytest +import structlog + +from airflow.configuration import conf +from airflow.dag_processing.lang_sdk_processor import ( + LangSDKDagFileProcessorProcess, + LangSDKRuntimeSchemaVersion, +) +from airflow.dag_processing.processor import DagFileParseRequest, DagFileParsingResult +from airflow.sdk import DAG, BaseOperator +from airflow.sdk.api.client import Client +from airflow.sdk.api.datamodels._generated import VariableResponse +from airflow.sdk.exceptions import AirflowRuntimeError +from airflow.sdk.execution_time import supervisor +from airflow.sdk.execution_time.comms import GetVariable, MaskSecret, _RequestFrame +from airflow.sdk.importers import DagSourceCode +from airflow.serialization.serialized_objects import DagSerialization, LazyDeserializedDAG + +from tests_common.test_utils.config import conf_vars +from unit.dag_processing.fake_lang_sdk import ( + FakeCoordinator, + fake_coordinator, + play_runtime, + write_native_file, +) + +# The oldest supervisor schema version, so the parse request is downgraded. +OLDEST_SCHEMA_VERSION = "2026-06-16" + + +def _serialize_dag(dag_id: str, description: str | None = None) -> LazyDeserializedDAG: + with DAG(dag_id, schedule=None, description=description) as dag: + BaseOperator(task_id="extract") + return LazyDeserializedDAG(data=DagSerialization.to_dict(dag)) + + +def _reply_with(*dags: LazyDeserializedDAG, **result): + def reply(request: DagFileParseRequest, comms) -> DagFileParsingResult: + return DagFileParsingResult(fileloc=request.file, serialized_dags=list(dags), **result) + + return reply + + +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() + + +@pytest.fixture(autouse=True) +def _coordinator(): + with fake_coordinator(): + yield + + +def _start(tmp_path, selector, *, client: Client | None = None, **spec) -> LangSDKDagFileProcessorProcess: + return LangSDKDagFileProcessorProcess.start( + id=uuid.uuid4(), + path=write_native_file(tmp_path / "dag.native", **spec), + bundle_path=tmp_path, + bundle_name="testing", + dag_file_rel_path="dag.native", + selector=selector, + logger=structlog.get_logger(), + client=client or MagicMock(spec=Client), + ) + + +@pytest.fixture +def parse(tmp_path): + """Parse ``dag.native`` as the Dag processor does, and check that nothing is left open.""" + + def _parse(**kwargs) -> LangSDKDagFileProcessorProcess: + 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 _send_an_invalid_frame(request, comms) -> None: + comms.socket.sendall(bytes.fromhex("00000003c1c1c1")) + time.sleep(60) + + +class TestLangSDKDagFileProcessorProcess: + @patch.object(FakeCoordinator, "parse_dag", autospec=True) + def test_parses_the_dags_the_runtime_returns(self, mock_parse_dag, parse, tmp_path, cap_structlog): + mock_parse_dag.side_effect = play_runtime( + _reply_with(_serialize_dag("native_dag")), + schema_version=OLDEST_SCHEMA_VERSION, + log_lines=[{"event": "Parsing the bundle", "level": "info"}], + ) + + proc = parse() + + assert proc.parsing_result.fileloc == os.fspath(tmp_path / "dag.native") + assert proc.parsing_result.import_errors is None + assert [dag.dag_id for dag in proc.parsing_result.serialized_dags] == ["native_dag"] + assert list(proc.parsing_result.dag_source_codes.values()) == [ + DagSourceCode((tmp_path / "dag.native").read_text(), "fake") + ] + assert proc._subprocess_schema_version == OLDEST_SCHEMA_VERSION + assert "Parsing the bundle" in cap_structlog + + @patch("airflow.dag_processing.lang_sdk_processor._is_connection_from_pid", autospec=True) + @patch.object(FakeCoordinator, "parse_dag", autospec=True) + def test_a_connection_is_used_once_it_is_verified(self, mock_parse_dag, mock_owned, tmp_path): + mock_parse_dag.side_effect = play_runtime(_reply_with(_serialize_dag("native_dag"))) + 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 [dag.dag_id for dag in proc.parsing_result.serialized_dags] == ["native_dag"] + + @patch.object(FakeCoordinator, "parse_dag", autospec=True) + def test_requests_are_answered_by_the_client(self, mock_parse_dag, parse): + def reply(request, comms): + variable = comms.send(GetVariable(key="native_var")) + return _reply_with(_serialize_dag("native_dag", description=variable.value))(request, comms) + + mock_parse_dag.side_effect = play_runtime(reply) + client = MagicMock(spec=Client) + client.variables = MagicMock() + client.variables.get.return_value = VariableResponse(key="native_var", value="from-db") + + proc = parse(client=client) + + [dag] = proc.parsing_result.serialized_dags + assert dag.data["dag"]["description"] == "from-db" + + @pytest.mark.parametrize( + ("change", "error"), + [ + pytest.param( + {"max_active_runs": "many"}, + "Dag 'broken_dag' does not match the schema: 'many' is not of type 'number'", + id="schema", + ), + pytest.param( + {"timetable": {"__type": "no.such.Timetable", "__var": {}}}, + "Dag 'broken_dag' cannot be deserialized: TimetableNotRegistered: " + "Timetable class 'no.such.Timetable' is not registered", + id="deserialize", + ), + ], + ) + @patch.object(FakeCoordinator, "parse_dag", autospec=True) + def test_a_dag_that_does_not_validate_is_an_import_error(self, mock_parse_dag, parse, change, error): + broken = _serialize_dag("broken_dag") + broken.data["dag"].update(change) + mock_parse_dag.side_effect = play_runtime(_reply_with(broken, _serialize_dag("good_dag"))) + + proc = parse() + + assert [dag.dag_id for dag in proc.parsing_result.serialized_dags] == ["good_dag"] + [message] = proc.parsing_result.import_errors.values() + assert message.startswith(f"Cannot load the serialized Dag: {error}") + + @conf_vars( + { + ("core", "max_active_tasks_per_dag"): "7", + ("core", "max_active_runs_per_dag"): "3", + ("scheduler", "catchup_by_default"): "True", + } + ) + @patch.object(FakeCoordinator, "parse_dag", autospec=True) + def test_a_dag_setting_left_unset_is_filled_from_the_config(self, mock_parse_dag, parse): + dag = _serialize_dag("native_dag") + del dag.data["dag"]["max_active_tasks"], dag.data["dag"]["catchup"] + dag.data["dag"]["max_active_runs"] = 16 + mock_parse_dag.side_effect = play_runtime(_reply_with(dag)) + + proc = parse() + + [stored] = proc.parsing_result.serialized_dags + assert proc.parsing_result.import_errors is None + assert stored.data["dag"]["max_active_tasks"] == 7 + assert stored.data["dag"]["max_active_runs"] == 16 + assert stored.data["dag"]["catchup"] is True + + @patch.object(FakeCoordinator, "parse_dag", autospec=True) + def test_a_dag_with_a_cycle_is_an_import_error(self, mock_parse_dag, parse): + cyclic = _serialize_dag("cyclic_dag") + [task] = cyclic.data["dag"]["tasks"] + task["__var"]["downstream_task_ids"] = ["extract"] + mock_parse_dag.side_effect = play_runtime(_reply_with(cyclic)) + + proc = parse() + + assert proc.parsing_result.serialized_dags == [] + assert proc.parsing_result.import_errors == { + "dag.native": "Cannot load the serialized Dag: Dag 'cyclic_dag' has a cycle through task 'extract'" + } + + @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, spec, reply, error): + with ( + patch.object(FakeCoordinator, "parse_dag", autospec=True, side_effect=play_runtime(reply)) + if reply + else contextlib.nullcontext() + ): + proc = parse(**spec) + + assert proc.parsing_result.serialized_dags == [] + assert proc.parsing_result.import_errors == {"dag.native": error} + + @patch.object(FakeCoordinator, "parse_dag", autospec=True) + def test_a_message_that_does_not_validate_is_an_import_error(self, mock_parse_dag, parse): + def reply(request, comms): + body = { + "type": "DagFileParsingResult", + "fileloc": request.file, + "serialized_dags": [{"data": "not a dict"}], + } + comms.socket.sendall(_RequestFrame(id=1, body=body).as_bytes()) + time.sleep(60) + + mock_parse_dag.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 "DagFileParsingResult.serialized_dags.0.data\n Input should be a valid dictionary" in message + assert proc._exit_code == -signal.SIGKILL + + @patch.object( + FakeCoordinator, "parse_dag", 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_dag, 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.lang_sdk_processor._EXIT_GRACE_PERIOD", 0.5) + @patch.object(FakeCoordinator, "parse_dag", autospec=True) + def test_a_runtime_that_runs_on_after_its_result_is_killed(self, mock_parse_dag, parse, cap_structlog): + def reply(request, comms): + comms.send(_reply_with(_serialize_dag("native_dag"))(request, comms)) + time.sleep(60) + + mock_parse_dag.side_effect = play_runtime(reply) + + proc = parse() + + assert [dag.dag_id for dag in proc.parsing_result.serialized_dags] == ["native_dag"] + assert proc._exit_code == -signal.SIGKILL + assert "The Lang-SDK runtime did not exit after its parse result; killing it" in cap_structlog + + @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" + + +class TestRun: + @staticmethod + def _run(tmp_path, *, timeout: float | None = 30) -> DagFileParsingResult: + return LangSDKDagFileProcessorProcess.run( + path=write_native_file(tmp_path / "dag.native"), + bundle_path=tmp_path, + bundle_name="testing", + dag_file_rel_path="dag.native", + timeout=timeout, + logger=structlog.get_logger(), + ) + + @patch.object(FakeCoordinator, "parse_dag", autospec=True) + def test_requests_get_an_error(self, mock_parse_dag, tmp_path): + def reply(request, comms): + with pytest.raises(AirflowRuntimeError) as ctx: + comms.send(GetVariable(key="native_var")) + description = ctx.value.error.detail["message"] + return _reply_with(_serialize_dag("native_dag", description=description))(request, comms) + + mock_parse_dag.side_effect = play_runtime(reply) + + result = self._run(tmp_path) + + assert result.serialized_dags[0].data["dag"]["description"] == ( + "GetVariable is answered only in the Dag processor" + ) + + @patch("airflow.sdk.execution_time.request_handlers.mask_secret", autospec=True) + @patch.object(FakeCoordinator, "parse_dag", autospec=True) + def test_a_secret_is_masked_without_a_client(self, mock_parse_dag, mock_mask_secret, tmp_path): + def reply(request, comms): + comms.send(MaskSecret(value="native-secret", name="native_conn")) + return _reply_with(_serialize_dag("native_dag"))(request, comms) + + mock_parse_dag.side_effect = play_runtime(reply) + + result = self._run(tmp_path) + + assert [dag.dag_id for dag in result.serialized_dags] == ["native_dag"] + mock_mask_secret.assert_called_once_with("native-secret", "native_conn") + + @pytest.mark.parametrize("connected", [False, True], ids=["before-connecting", "after-connecting"]) + def test_a_parse_past_its_timeout_is_killed(self, tmp_path, connected): + # A runtime that never connects leaves both listeners open when it is killed. + write_native_file(tmp_path / "dag.native", argv=["/bin/sh", "-c", "exec sleep 60"]) + fds_before = _get_open_fds() + + with ( + patch.object( + FakeCoordinator, + "parse_dag", + autospec=True, + side_effect=play_runtime(lambda request, comms: time.sleep(60)), + ) + if connected + else contextlib.nullcontext(), + patch.object( + LangSDKDagFileProcessorProcess, + "close", + autospec=True, + side_effect=LangSDKDagFileProcessorProcess.close, + ) as mock_close, + pytest.raises(TimeoutError, match=r"did not parse .*dag\.native within 1s"), + ): + LangSDKDagFileProcessorProcess.run( + path=tmp_path / "dag.native", + bundle_path=tmp_path, + bundle_name="testing", + dag_file_rel_path="dag.native", + timeout=1, + logger=structlog.get_logger(), + ) + + [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 + + +def _make_process(**kwargs) -> LangSDKDagFileProcessorProcess: + return LangSDKDagFileProcessorProcess( + id=uuid.uuid4(), + pid=1, + stdin=MagicMock(), + process=MagicMock(), + process_log=MagicMock(), + selector=MagicMock(), + bundle_name="testing", + dag_file_rel_path="dag.native", + listeners={}, + parse_request=DagFileParseRequest( + file="/b/dag.native", bundle_path=Path("/b"), bundle_name="testing" + ), + **kwargs, + ) + + +def test_a_dag_source_that_cannot_be_read_is_a_placeholder(): + proc = _make_process() + dag = _serialize_dag("native_dag") + + proc._handle_request(DagFileParsingResult(fileloc="/b/dag.native", serialized_dags=[dag]), MagicMock(), 1) + + source = proc.parsing_result.dag_source_codes[dag.data["dag"]["fileloc"]] + assert source.language == "text" + assert source.source_code.startswith("Cannot read the source of dag.native: [Errno 2] No such file") + + +@patch.object(LangSDKDagFileProcessorProcess, "send_msg", autospec=True) +def test_the_first_parse_result_wins(mock_send_msg): + proc = _make_process() + first = DagFileParsingResult(fileloc="/b/dag.native", serialized_dags=[_serialize_dag("first")]) + second = DagFileParsingResult(fileloc="/b/dag.native", serialized_dags=[_serialize_dag("second")]) + + proc._handle_request(first, MagicMock(), 1) + proc._handle_request(second, MagicMock(), 2) + + assert [dag.dag_id for dag in proc.parsing_result.serialized_dags] == ["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": "DagFileParsingResult", + "fileloc": "/b/dag.native", + "serialized_dags": [{"data": "not a dict"}], + }, + ).as_bytes(), + id="does-not-validate", + ), + ], +) +@patch.object(LangSDKDagFileProcessorProcess, "_kill_runtime", autospec=True) +@patch.object(LangSDKDagFileProcessorProcess, "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() + result = DagFileParsingResult(fileloc="/b/dag.native", serialized_dags=[_serialize_dag("native_dag")]) + 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=result.model_dump(mode="json")).as_bytes()) + assert read_frame(conn) + runtime.sendall(invalid_frame) + assert not read_frame(conn) + + assert [dag.dag_id for dag in proc.parsing_result.serialized_dags] == ["native_dag"] + 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(LangSDKDagFileProcessorProcess, "send_msg", autospec=True) +def test_the_schema_version_is_reported_once(mock_send_msg): + proc = _make_process() + + proc._handle_request(LangSDKRuntimeSchemaVersion(schema_version=OLDEST_SCHEMA_VERSION), MagicMock(), 1) + proc._handle_request(LangSDKRuntimeSchemaVersion(schema_version=None), MagicMock(), 2) + + assert proc._runtime_schema_version == OLDEST_SCHEMA_VERSION + assert mock_send_msg.call_args.kwargs["error"].detail["message"] == "Unhandled request" diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index fb15c2486c632..d64192631a2ca 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -51,6 +51,7 @@ from airflow.dag_processing.bundles.manager import DagBundlesManager from airflow.dag_processing.collection import update_dag_parsing_results_in_db from airflow.dag_processing.dagbag import DagBag +from airflow.dag_processing.lang_sdk_processor import LangSDKDagFileProcessorProcess from airflow.dag_processing.manager import ( BundleState, DagFileInfo, @@ -71,9 +72,9 @@ from airflow.models.serialized_dag import SerializedDagModel from airflow.models.team import Team from airflow.providers.standard.operators.empty import EmptyOperator -from airflow.sdk import DAG as SdkDAG +from airflow.sdk import DAG as SdkDAG, BaseOperator from airflow.sdk.importers import DagDefinition, DagImporterRegistry, DagImportError, DagSourceCode -from airflow.serialization.serialized_objects import LazyDeserializedDAG +from airflow.serialization.serialized_objects import DagSerialization, LazyDeserializedDAG from airflow.utils.net import get_hostname from airflow.utils.session import create_session @@ -90,6 +91,12 @@ clear_db_serialized_dags, clear_db_teams, ) +from unit.dag_processing.fake_lang_sdk import ( + FakeCoordinator, + fake_coordinator, + play_runtime, + write_native_file, +) from unit.models import TEST_DAGS_FOLDER pytestmark = pytest.mark.db_test @@ -1661,6 +1668,132 @@ def test_create_process_subprocess_logs_to_stdout( _, kwargs = mock_start.call_args assert kwargs["subprocess_logs_to_stdout"] is expected_subprocess_logs_to_stdout + @mock.patch.object(DagFileProcessorManager, "_get_logger_for_dag_file", autospec=True) + def test_create_process_parses_a_coordinator_file_with_its_runtime(self, mock_get_logger, tmp_path): + mock_get_logger.return_value = (MagicMock(), MagicMock()) + dag_file = DagFileInfo(bundle_name="testing", rel_path=Path("dags.native"), bundle_path=tmp_path) + + with ( + fake_coordinator(), + mock.patch.object(LangSDKDagFileProcessorProcess, "start", autospec=True) as mock_start, + ): + manager = DagFileProcessorManager(max_runs=1) + manager._create_process(dag_file) + + kwargs = mock_start.call_args.kwargs + assert (kwargs["path"], kwargs["dag_file_rel_path"]) == (tmp_path / "dags.native", "dags.native") + assert kwargs["client"] is manager.client + + @pytest.mark.parametrize("rel_path", ["my_dag.py", "dags.fake"]) + @mock.patch.object(DagFileProcessorManager, "_get_logger_for_dag_file", autospec=True) + def test_create_process_keeps_the_python_parse_for_other_files(self, mock_get_logger, rel_path, tmp_path): + mock_get_logger.return_value = (MagicMock(), MagicMock()) + dag_file = DagFileInfo(bundle_name="testing", rel_path=Path(rel_path), bundle_path=tmp_path) + + with ( + fake_coordinator(), + mock.patch.object(DagFileProcessorProcess, "start", autospec=True) as mock_start, + mock.patch.object(LangSDKDagFileProcessorProcess, "start", autospec=True) as mock_lang_sdk_start, + ): + DagFileProcessorManager(max_runs=1)._create_process(dag_file) + + mock_start.assert_called_once() + mock_lang_sdk_start.assert_not_called() + + @mock.patch.object(DagFileProcessorManager, "_get_logger_for_dag_file", autospec=True) + def test_create_process_drops_callbacks_for_a_coordinator_file(self, mock_get_logger, tmp_path, caplog): + mock_get_logger.return_value = (MagicMock(), MagicMock()) + dag_file = DagFileInfo(bundle_name="testing", rel_path=Path("dags.native"), bundle_path=tmp_path) + callback = DagCallbackRequest( + filepath="dags.native", + dag_id="native_dag", + run_id="run", + bundle_name="testing", + bundle_version=None, + is_failure_callback=True, + ) + + with ( + fake_coordinator(), + mock.patch.object(LangSDKDagFileProcessorProcess, "start", autospec=True) as mock_start, + caplog.at_level(logging.WARNING, logger="airflow.dag_processing.manager"), + ): + manager = DagFileProcessorManager(max_runs=1) + manager._callback_to_execute[dag_file] = [callback] + manager._create_process(dag_file) + + assert "callbacks" not in mock_start.call_args.kwargs + assert dag_file not in manager._callback_to_execute + assert [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] == [ + "Dropping 1 callbacks for dags.native: Lang-SDK runtimes do not run callbacks" + ] + + @mock.patch.object(FakeCoordinator, "parse_dag", autospec=True) + @mock.patch.object( + DagFileProcessorManager, "_find_files_in_bundle", autospec=True, return_value=[Path("good.native")] + ) + def test_coordinator_files_are_persisted( + self, mock_find_files, mock_parse_dag, tmp_path, configure_testing_dag_bundle + ): + def reply(request, comms): + with SdkDAG("native_dag", schedule=None) as dag: + BaseOperator(task_id="extract") + data = DagSerialization.to_dict(dag) + data["dag"].update(fileloc=request.file, relative_fileloc="good.native") + return DagFileParsingResult( + fileloc=request.file, serialized_dags=[LazyDeserializedDAG(data=data)] + ) + + mock_parse_dag.side_effect = play_runtime(reply) + write_native_file(tmp_path / "good.native") + + with fake_coordinator(), configure_testing_dag_bundle(tmp_path): + DagFileProcessorManager(max_runs=1, processor_timeout=60).run() + + with create_session() as session: + serialized_dag = session.scalar( + select(SerializedDagModel).where(SerializedDagModel.dag_id == "native_dag") + ) + dag_code = session.scalar(select(DagCode).where(DagCode.dag_id == "native_dag")) + + assert serialized_dag.data["dag"]["tasks"][0]["__var"]["task_id"] == "extract" + assert dag_code.source_code == (tmp_path / "good.native").read_text() + + @mock.patch.object(FakeCoordinator, "parse_dag", autospec=True) + @mock.patch.object( + DagFileProcessorManager, + "_find_files_in_bundle", + autospec=True, + return_value=[Path("garbage.native"), Path("python_dag.py")], + ) + def test_an_invalid_frame_does_not_stop_other_files_parsing( + self, mock_find_files, mock_parse_dag, tmp_path, configure_testing_dag_bundle + ): + def reply(request, comms): + comms.socket.sendall(bytes.fromhex("00000003c1c1c1")) + time.sleep(60) + + mock_parse_dag.side_effect = play_runtime(reply) + write_native_file(tmp_path / "garbage.native") + (tmp_path / "python_dag.py").write_text( + "from airflow.sdk import DAG\nfrom airflow.sdk.bases.operator import BaseOperator\n\n" + 'with DAG("python_dag", schedule=None):\n BaseOperator(task_id="task")\n' + ) + + with fake_coordinator(), configure_testing_dag_bundle(tmp_path): + manager = DagFileProcessorManager(max_runs=1, processor_timeout=60) + manager.run() + + with create_session() as session: + dag_ids = session.scalars(select(SerializedDagModel.dag_id)).all() + import_errors = session.scalars(select(ParseImportError)).all() + + assert dag_ids == ["python_dag"] + [import_error] = import_errors + assert import_error.filename == "garbage.native" + assert import_error.stacktrace.startswith("The Lang-SDK runtime sent an invalid frame: ") + assert manager.selector.get_map() == {} + def test_terminate_orphan_processes_kills_then_closes_processor(self): manager = DagFileProcessorManager(max_runs=1) processor, _ = self.mock_processor() diff --git a/generated/known_sdk_imports_in_core.txt b/generated/known_sdk_imports_in_core.txt index 8c799031129b9..7e0c840c79eed 100644 --- a/generated/known_sdk_imports_in_core.txt +++ b/generated/known_sdk_imports_in_core.txt @@ -4,6 +4,7 @@ airflow-core/src/airflow/cli/commands/task_command.py::7 airflow-core/src/airflow/cli/commands/triggerer_command.py::1 airflow-core/src/airflow/configuration.py::1 airflow-core/src/airflow/dag_processing/dagbag.py::3 +airflow-core/src/airflow/dag_processing/lang_sdk_processor.py::7 airflow-core/src/airflow/dag_processing/manager.py::5 airflow-core/src/airflow/dag_processing/processor.py::14 airflow-core/src/airflow/exceptions.py::1 diff --git a/task-sdk/src/airflow/sdk/coordinators/_dag_importer.py b/task-sdk/src/airflow/sdk/coordinators/_dag_importer.py index 124dd0143e637..54517ee38310d 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_dag_importer.py +++ b/task-sdk/src/airflow/sdk/coordinators/_dag_importer.py @@ -27,6 +27,8 @@ from airflow.sdk.execution_time.coordinator import CoordinatorManager, get_coordinator_manager from airflow.sdk.importers.base import ( AbstractDagImporter, + DagImportError, + DagImportResult, FilesystemDagDefinition, find_file_dag_definitions, get_importer_registry, @@ -56,8 +58,12 @@ class CoordinatorDagImporter(AbstractDagImporter[FilesystemDagDefinition]): needs one: the only coordinator of :attr:`coordinator_classpath`'s class, or the one ``[sdk] dag_bundle_to_coordinator`` picks for the bundle when there are several. + The Dag processor does not call :meth:`import_definition`: it runs the runtime itself and stores + the Dags the runtime serialized. A Dag bag, such as a CLI command's, reports such a file as an + import error. + Subclasses set :attr:`coordinator_classpath`, :attr:`artifact_suffix` and - :attr:`supported_extensions`, and implement :meth:`import_definition` and :meth:`get_source_code`. + :attr:`supported_extensions`, and implement :meth:`get_source_code`. """ coordinator_classpath: ClassVar[str] @@ -103,6 +109,20 @@ def list_dag_definitions( ): yield definition + def import_definition( + self, definition: FilesystemDagDefinition, bundle: BaseDagBundle + ) -> DagImportResult: + """Report that only the Dag processor parses *definition*.""" + return DagImportResult( + definition=definition, + errors=[ + DagImportError( + source_reference=repr(definition), + message="A native Lang-SDK Dag is parsed only by the Dag processor", + ) + ], + ) + def build_coordinator_dag_importers( manager: CoordinatorManager, bundle_name: str diff --git a/task-sdk/tests/task_sdk/coordinators/test_dag_importer.py b/task-sdk/tests/task_sdk/coordinators/test_dag_importer.py index e0285de5cbdf6..53442f0b24b38 100644 --- a/task-sdk/tests/task_sdk/coordinators/test_dag_importer.py +++ b/task-sdk/tests/task_sdk/coordinators/test_dag_importer.py @@ -37,7 +37,7 @@ _CoordinatorSpec, get_coordinator_manager, ) -from airflow.sdk.importers import DagImportResult, DagSourceCode, reset_importer_registry +from airflow.sdk.importers import DagSourceCode, FilesystemDagDefinition, reset_importer_registry from tests_common.test_utils.config import conf_vars @@ -55,9 +55,6 @@ class _NativeImporter(CoordinatorDagImporter): artifact_suffix = ".min.native" supported_extensions = [".native"] - def import_definition(self, definition, bundle): - return DagImportResult() - def get_source_code(self, definition) -> DagSourceCode: return DagSourceCode("", "native") @@ -109,6 +106,19 @@ def test_lists_only_its_artifacts(importer, tmp_path): assert [d.path.name for d in definitions] == ["main.min.native"] +def test_import_definition_reports_that_only_the_dag_processor_parses_it(importer, tmp_path): + bundle_file = tmp_path / "main.min.native" + bundle_file.write_text("") + definition = FilesystemDagDefinition(bundle_file) + + result = importer.import_definition(definition, SimpleNamespace(name="testing", path=tmp_path)) + + assert result.dags == [] + assert [error.message for error in result.errors] == [ + "A native Lang-SDK Dag is parsed only by the Dag processor" + ] + + class TestGetParsingCoordinator: def test_returns_the_only_coordinator_of_the_class(self, importer): with conf_vars(_coordinator_config("native")): From c8a3306646d4cd04147ce17e7f043531d30ddf7e Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Wed, 30 Sep 2026 02:58:52 +0000 Subject: [PATCH 02/24] Report a Lang-SDK parse past its import timeout as an import error The parse child resolves the get_dagbag_import_timeout policy for the file and sends the timeout with the runtime's schema version. The policy thus runs where it does for a Python Dag file, and a failure is only that file's import error. The timeout counts from the start of the parse. Past it, with no result yet, the runtime is killed and the file gets an import error. LangSDKDagFileProcessorProcess.run applies the same timeout, so it no longer takes one. Until the child reports it, dag_file_processor_timeout bounds the parse, as it does in the Dag processor. At the deadline run stops waiting even if an exited runtime left a process that holds its output open. --- .../dag_processing/lang_sdk_processor.py | 58 ++++++++---- .../dag_processing/test_lang_sdk_processor.py | 88 ++++++++++++++++--- .../tests/unit/dag_processing/test_manager.py | 57 ++++++++++++ 3 files changed, 174 insertions(+), 29 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py b/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py index 877af756ed6da..dc7d2fca7bcfc 100644 --- a/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py +++ b/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py @@ -32,6 +32,8 @@ from pydantic import BaseModel, Field, TypeAdapter from uuid6 import uuid7 +from airflow import settings +from airflow.configuration import conf from airflow.dag_processing.importer_routing import get_claiming_coordinator from airflow.dag_processing.processor import ( BaseDagFileProcessorProcess, @@ -83,12 +85,22 @@ class StartLangSDKRuntime(BaseModel): class LangSDKRuntimeSchemaVersion(BaseModel): - """The supervisor schema version of the runtime the parse child is about to exec.""" + """The schema version and the import timeout of the runtime the parse child is about to exec.""" schema_version: str | None + import_timeout: float | None = None + """Seconds from the start of the parse; ``None`` means no timeout.""" type: Literal["LangSDKRuntimeSchemaVersion"] = "LangSDKRuntimeSchemaVersion" +def _get_import_timeout(path: str) -> float | None: + """Return the ``get_dagbag_import_timeout`` policy's timeout for *path*; ``None`` means none.""" + timeout = settings.get_dagbag_import_timeout(path) + if not isinstance(timeout, (int, float)): + raise TypeError(f"Value ({timeout}) from get_dagbag_import_timeout must be int or float") + return timeout if timeout > 0 else None + + def _start_runtime_entrypoint() -> None: """Exec the runtime that parses the file named by the start request, or report why it cannot start.""" os.environ["_AIRFLOW_PROCESS_CONTEXT"] = "client" @@ -106,9 +118,11 @@ def _start_runtime_entrypoint() -> None: raise RuntimeError(f"Required first message to be a StartLangSDKRuntime, it was {msg}") def report_schema_version(schema_version: str | None) -> None: - comms.send(LangSDKRuntimeSchemaVersion(schema_version=schema_version)) + 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_claiming_coordinator(msg.file, msg.bundle_name) if coordinator is None: raise RuntimeError(f"No coordinator's Dag importer claims {msg.file}") @@ -154,6 +168,7 @@ class LangSDKDagFileProcessorProcess(BaseDagFileProcessorProcess): _listeners: dict[_Channel, socket] _parse_request: DagFileParseRequest _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) @@ -217,17 +232,16 @@ def run( bundle_path: Path, bundle_name: str, dag_file_rel_path: str, - timeout: float | None, logger: FilteringBoundLogger, ) -> DagFileParsingResult: """ Parse *path* outside the Dag processor and wait for the result. - There is no API client, so each request of the runtime that needs one gets an error. - - :raises TimeoutError: if the parse does not finish within *timeout* seconds. The runtime is - killed. + There is no API client, so each request of the runtime that needs one gets an error. The file's import + timeout bounds the parse, and ``[dag_processor] dag_file_processor_timeout`` until the parse child + reports it. """ + processor_timeout = conf.getfloat("dag_processor", "dag_file_processor_timeout") with selectors.DefaultSelector() as selector: proc = cls.start( id=uuid7(), @@ -240,14 +254,13 @@ def run( ) try: while not proc.is_ready: - wait = 0.1 - if timeout is not None: - if (remaining := proc.start_time + timeout - time.monotonic()) <= 0: - raise TimeoutError( - f"The Lang-SDK runtime did not parse {os.fspath(path)} within {timeout}s" - ) - wait = min(wait, remaining) - proc._service_subprocess(max_wait_time=wait) + timeout = proc._import_timeout if proc._schema_version_reported else processor_timeout + 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) + break + proc._service_subprocess(max_wait_time=0.1) except BaseException: proc._kill_runtime() raise @@ -381,6 +394,7 @@ def _handle_runtime_schema_version( 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, {} @@ -466,6 +480,13 @@ def is_ready(self) -> bool: ): self.process_log.warning("The Lang-SDK runtime did not exit after its parse result; killing it") 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) if self._check_subprocess_exit() is None: return False self._close_listeners() @@ -477,6 +498,13 @@ def is_ready(self) -> bool: ) return True + def _time_out(self, timeout: float) -> None: + 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" + ) + self._kill_runtime() + def _kill_runtime(self) -> None: """Kill the runtime and wait for it, without servicing its sockets, whose handler may have failed.""" if self._exit_code is not None: diff --git a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py index 4a5fd758681d9..dd3075a7c3381 100644 --- a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py @@ -36,6 +36,7 @@ from airflow.dag_processing.lang_sdk_processor import ( LangSDKDagFileProcessorProcess, LangSDKRuntimeSchemaVersion, + _get_import_timeout, ) from airflow.dag_processing.processor import DagFileParseRequest, DagFileParsingResult from airflow.sdk import DAG, BaseOperator @@ -336,6 +337,27 @@ def reply(request, comms): assert proc._exit_code == -signal.SIGKILL assert "The Lang-SDK runtime did not exit after its parse result; killing it" in cap_structlog + @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 == { + "dag.native": 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): @@ -361,13 +383,12 @@ def test_the_runtime_inherits_only_its_standard_streams(self, monkeypatch, tmp_p class TestRun: @staticmethod - def _run(tmp_path, *, timeout: float | None = 30) -> DagFileParsingResult: + def _run(tmp_path, **spec) -> DagFileParsingResult: return LangSDKDagFileProcessorProcess.run( - path=write_native_file(tmp_path / "dag.native"), + path=write_native_file(tmp_path / "dag.native", **spec), bundle_path=tmp_path, bundle_name="testing", dag_file_rel_path="dag.native", - timeout=timeout, logger=structlog.get_logger(), ) @@ -402,9 +423,9 @@ def reply(request, comms): mock_mask_secret.assert_called_once_with("native-secret", "native_conn") @pytest.mark.parametrize("connected", [False, True], ids=["before-connecting", "after-connecting"]) - def test_a_parse_past_its_timeout_is_killed(self, tmp_path, connected): + @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. - write_native_file(tmp_path / "dag.native", argv=["/bin/sh", "-c", "exec sleep 60"]) fds_before = _get_open_fds() with ( @@ -422,22 +443,61 @@ def test_a_parse_past_its_timeout_is_killed(self, tmp_path, connected): autospec=True, side_effect=LangSDKDagFileProcessorProcess.close, ) as mock_close, - pytest.raises(TimeoutError, match=r"did not parse .*dag\.native within 1s"), ): - LangSDKDagFileProcessorProcess.run( - path=tmp_path / "dag.native", - bundle_path=tmp_path, - bundle_name="testing", - dag_file_rel_path="dag.native", - timeout=1, - logger=structlog.get_logger(), - ) + result = self._run(tmp_path, argv=["/bin/sh", "-c", "exec sleep 60"]) + assert result.import_errors == { + "dag.native": f"The Lang-SDK runtime did not parse {tmp_path / 'dag.native'} within 1.0s" + } [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 + @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): + with patch.object( + LangSDKDagFileProcessorProcess, + "close", + autospec=True, + side_effect=LangSDKDagFileProcessorProcess.close, + ) as mock_close: + # The runtime exits, and the process it leaves behind keeps its output open. + result = self._run(tmp_path, argv=["/bin/sh", "-c", "sleep 30 & exit 0"]) + [proc] = [c.args[0] for c in mock_close.call_args_list] + os.killpg(proc.pid, signal.SIGKILL) + + assert result.import_errors == { + "dag.native": f"The Lang-SDK runtime did not parse {tmp_path / 'dag.native'} within 1.0s" + } + assert proc._exit_code == 0 + assert not proc._open_sockets + + @conf_vars({("dag_processor", "dag_file_processor_timeout"): "1"}) + @patch.object( + FakeCoordinator, + "_build_parse_dag_command", + autospec=True, + side_effect=lambda self, *, path: time.sleep(60), + ) + def test_the_dag_file_processor_timeout_applies_until_the_import_timeout_is_reported( + self, mock_build_parse_dag_command, tmp_path + ): + result = self._run(tmp_path) + + assert result.import_errors == { + "dag.native": f"The Lang-SDK runtime did not parse {tmp_path / 'dag.native'} within 1.0s" + } + + +@pytest.mark.parametrize(("configured", "expected"), [(30, 30), (0.5, 0.5), (0, None), (-1, None)]) +@patch("airflow.settings.get_dagbag_import_timeout", autospec=True) +def test_only_a_positive_import_timeout_applies(mock_timeout, configured, expected): + mock_timeout.return_value = configured + + assert _get_import_timeout("/b/dag.native") == expected + mock_timeout.assert_called_once_with("/b/dag.native") + def _make_process(**kwargs) -> LangSDKDagFileProcessorProcess: return LangSDKDagFileProcessorProcess( diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index d64192631a2ca..7a09829038889 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -1794,6 +1794,63 @@ def reply(request, comms): assert import_error.stacktrace.startswith("The Lang-SDK runtime sent an invalid frame: ") assert manager.selector.get_map() == {} + @mock.patch("airflow.settings.get_dagbag_import_timeout", autospec=True, return_value=1) + @mock.patch.object(FakeCoordinator, "parse_dag", autospec=True) + @mock.patch.object( + DagFileProcessorManager, "_find_files_in_bundle", autospec=True, return_value=[Path("slow.native")] + ) + def test_a_coordinator_file_past_the_import_timeout_is_an_import_error( + self, mock_find_files, mock_parse_dag, mock_timeout, tmp_path, configure_testing_dag_bundle + ): + mock_parse_dag.side_effect = play_runtime(lambda request, comms: time.sleep(60)) + slow = write_native_file(tmp_path / "slow.native") + + with fake_coordinator(), configure_testing_dag_bundle(tmp_path): + manager = DagFileProcessorManager(max_runs=1, processor_timeout=60) + manager.run() + + with create_session() as session: + [import_error] = session.scalars(select(ParseImportError)).all() + + assert import_error.filename == "slow.native" + assert import_error.stacktrace == f"The Lang-SDK runtime did not parse {slow} within 1.0s" + assert manager.selector.get_map() == {} + + @mock.patch("airflow.settings.get_dagbag_import_timeout", autospec=True) + @mock.patch.object( + DagFileProcessorManager, + "_find_files_in_bundle", + autospec=True, + return_value=[Path("broken.native"), Path("python_dag.py")], + ) + def test_a_failing_import_timeout_policy_fails_only_its_file( + self, mock_find_files, mock_timeout, tmp_path, configure_testing_dag_bundle + ): + def get_dagbag_import_timeout(dag_file_path): + if dag_file_path.endswith(".native"): + raise RuntimeError("policy bug") + return 30 + + mock_timeout.side_effect = get_dagbag_import_timeout + write_native_file(tmp_path / "broken.native") + (tmp_path / "python_dag.py").write_text( + "from airflow.sdk import DAG\nfrom airflow.sdk.bases.operator import BaseOperator\n\n" + 'with DAG("python_dag", schedule=None):\n BaseOperator(task_id="task")\n' + ) + + with fake_coordinator(), configure_testing_dag_bundle(tmp_path): + DagFileProcessorManager(max_runs=1, processor_timeout=60).run() + + with create_session() as session: + dag_ids = session.scalars(select(SerializedDagModel.dag_id)).all() + [import_error] = session.scalars(select(ParseImportError)).all() + + assert dag_ids == ["python_dag"] + assert (import_error.filename, import_error.stacktrace) == ( + "broken.native", + "Cannot start the Lang-SDK runtime: RuntimeError: policy bug", + ) + def test_terminate_orphan_processes_kills_then_closes_processor(self): manager = DagFileProcessorManager(max_runs=1) processor, _ = self.mock_processor() From bbea6488ded3385ff308e385eccbc70f2d9baca6 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Wed, 30 Sep 2026 03:03:03 +0000 Subject: [PATCH 03/24] Report a native Dag's task that reaches a Python worker A Python worker handed a task from a Dag file that a coordinator's Dag importer claims now exits with an error naming the queue to route. It checks before building a Dag bag, so it never starts the runtime. --- .../airflow/sdk/execution_time/task_runner.py | 23 +++++ .../execution_time/test_task_runner.py | 92 +++++++++++++++++++ 2 files changed, 115 insertions(+) diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py index ed622023308a2..c4bf7f9f5b29a 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -60,6 +60,7 @@ from airflow.sdk.bases.skipmixin import XCOM_SKIPMIXIN_KEY from airflow.sdk.bases.xcom import BaseXCom from airflow.sdk.configuration import conf +from airflow.sdk.coordinators._dag_importer import CoordinatorDagImporter from airflow.sdk.definitions._internal.dag_parsing_context import _airflow_parsing_context_manager from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet, is_arg_set from airflow.sdk.definitions.asset import ( @@ -150,6 +151,7 @@ from airflow.sdk.execution_time.sentry import Sentry from airflow.sdk.execution_time.tracing import detail_span from airflow.sdk.execution_time.xcom import XCom +from airflow.sdk.importers import get_importer_registry from airflow.sdk.listener import get_listener_manager from airflow.sdk.observability.metrics import stats_utils from airflow.sdk.serde import allow_class, iter_pydantic_models @@ -1018,6 +1020,16 @@ def _register_deserialization_allowed_classes(dag, log: Logger) -> None: ) +def _is_lang_sdk_dag_file(path: str, bundle_name: str) -> bool: + """Return whether a coordinator's Dag importer claims *path*, so that a Lang-SDK runtime parses it.""" + try: + importer = get_importer_registry(bundle_name).get_importer(path) + except Exception: + # Building the Dag bag reports a broken importer configuration. + return False + return isinstance(importer, CoordinatorDagImporter) + + @detail_span("parse") def parse(what: StartupDetails, log: Logger) -> RuntimeTaskInstance: # TODO: Task-SDK: @@ -1030,6 +1042,17 @@ def parse(what: StartupDetails, log: Logger) -> RuntimeTaskInstance: bundle_prepare_ms = int((time.monotonic() - bundle_prepare_start) * 1000) dag_absolute_path = os.fspath(Path(bundle_instance.path, what.dag_rel_path)) + if _is_lang_sdk_dag_file(dag_absolute_path, bundle_info.name): + log.error( + "A task of a native Lang-SDK Dag cannot run in Python. Route its queue to the coordinator " + "that parses the Dag, with [sdk] queue_to_coordinator", + dag_id=what.ti.dag_id, + task_id=what.ti.task_id, + queue=what.ti.queue, + path=what.dag_rel_path, + ) + sys.exit(1) + dag_file_parse_start = time.monotonic() bag = BundleDagBag( dag_folder=dag_absolute_path, diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py index 9cdb75446b1e9..26ae41ac01eb5 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py +++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py @@ -31,6 +31,7 @@ from unittest import mock from unittest.mock import call, patch +import attrs import pandas as pd import pytest import structlog @@ -78,6 +79,8 @@ ) from airflow.sdk.bases.operator import ExecutorSafeguard from airflow.sdk.bases.xcom import BaseXCom +from airflow.sdk.coordinators._dag_importer import CoordinatorDagImporter +from airflow.sdk.coordinators._subprocess import SubprocessCoordinator from airflow.sdk.definitions._internal.types import NOTSET, SET_DURING_EXECUTION, is_arg_set from airflow.sdk.definitions.asset import Asset, AssetAlias, AssetUniqueKey, AssetUriRef, Dataset, Model from airflow.sdk.definitions.param import DagParam @@ -176,6 +179,7 @@ TaskRunnerMarker, _defer_task, _execute_task, + _is_lang_sdk_dag_file, _make_task_span, _push_xcom_if_needed, _register_deserialization_allowed_classes, @@ -189,6 +193,7 @@ startup, ) from airflow.sdk.execution_time.xcom import XCom +from airflow.sdk.importers import DagSourceCode, reset_importer_registry from airflow.sdk.serde import deserialize from airflow.triggers.base import BaseEventTrigger, BaseTrigger, TriggerEvent from airflow.triggers.callback import CallbackTrigger @@ -902,6 +907,93 @@ def test_parse_module_in_bundle_root(tmp_path: Path, make_ti_context): assert ti.task.dag.dag_id == "dag_name" +class NativeDagImporter(CoordinatorDagImporter): + artifact_suffix = ".native" + supported_extensions = [".native"] + + def get_source_code(self, definition): + return DagSourceCode(source_code=definition.read_text(), language="native") + + +@attrs.define(kw_only=True) +class NativeCoordinator(SubprocessCoordinator): + """A coordinator whose Dag importer claims ``.native`` files in every bundle.""" + + def get_dag_importer(self): + return NativeDagImporter(coordinator=self) + + +@patch("airflow.dag_processing.dagbag.BundleDagBag", autospec=True) +def test_parse_rejects_a_task_of_a_native_dag(mock_bag, tmp_path: Path, make_ti_context): + tmp_path.joinpath("dag.native").write_text("{}") + what = StartupDetails( + ti=TaskInstance( + id=uuid7(), + task_id="a", + dag_id="native_dag", + run_id="c", + try_number=1, + dag_version_id=uuid7(), + queue="default", + ), + dag_rel_path="dag.native", + bundle_info=BundleInfo(name="my-bundle", version=None), + ti_context=make_ti_context(), + start_date=timezone.utcnow(), + sentry_integration="", + ) + bundle_config = [ + { + "name": "my-bundle", + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {"path": str(tmp_path), "refresh_interval": 1}, + } + ] + coordinators = {"native": {"classpath": f"{__name__}.NativeCoordinator", "kwargs": {}}} + log = mock.Mock() + + reset_importer_registry() + try: + with ( + patch.dict( + os.environ, + { + "AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(bundle_config), + "AIRFLOW__SDK__COORDINATORS": json.dumps(coordinators), + }, + ), + pytest.raises(SystemExit, match="1"), + ): + parse(what, log) + finally: + reset_importer_registry() + + mock_bag.assert_not_called() + log.error.assert_called_once_with( + "A task of a native Lang-SDK Dag cannot run in Python. Route its queue to the coordinator " + "that parses the Dag, with [sdk] queue_to_coordinator", + dag_id="native_dag", + task_id="a", + queue="default", + path="dag.native", + ) + + +@pytest.mark.parametrize( + ("file_name", "expected"), + [("dag.native", True), ("dag.py", False), ("dag.pyc", False), ("dags.zip", False)], +) +def test_is_lang_sdk_dag_file(file_name, expected): + coordinators = {"native": {"classpath": f"{__name__}.NativeCoordinator", "kwargs": {}}} + + reset_importer_registry() + try: + with patch.dict(os.environ, {"AIRFLOW__SDK__COORDINATORS": json.dumps(coordinators)}): + assert _is_lang_sdk_dag_file(f"/bundle/{file_name}", "my-bundle") is expected + finally: + reset_importer_registry() + + @pytest.mark.parametrize("use_queues", [False, True]) def test_run_deferred_basic(time_machine, create_runtime_ti, mock_supervisor_comms, use_queues: bool): """Test that a task can transition to a deferred state.""" From 18de8bfa32d5567994968c1c43b29fc86614dcc3 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Fri, 2 Oct 2026 02:40:19 +0000 Subject: [PATCH 04/24] Drop callbacks for Lang-SDK Dag files before queuing the file A callback run for a native file parsed the callback's pinned bundle version and stored the result as the current version, so an old run's callback could overwrite the current serialized Dag. Drop such callbacks when they arrive, with one warning naming each request. --- .../src/airflow/dag_processing/manager.py | 23 +++++-- .../tests/unit/dag_processing/test_manager.py | 68 +++++++++++++------ 2 files changed, 63 insertions(+), 28 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index a1cd138e672c6..95ff65786d5d6 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -47,6 +47,7 @@ from airflow._shared.observability.metrics import stats from airflow._shared.observability.metrics.stats import normalize_name_for_stats from airflow._shared.timezones import timezone +from airflow.callbacks.callback_requests import DagCallbackRequest from airflow.configuration import conf from airflow.dag_processing.bundles.base import ( BundleUsageTrackingManager, @@ -809,6 +810,9 @@ def prepare_callback_bundle(self, request: CallbackRequest) -> BaseDagBundle | N def _add_callback_to_queue(self, request: CallbackRequest) -> None: self.log.debug("Queuing %s CallbackRequest: %s", type(request).__name__, request) + if get_claiming_coordinator(request.filepath, request.bundle_name) is not None: + self._log_dropped_lang_sdk_callback(request) + return bundle = self.prepare_callback_bundle(request) if bundle is None: return @@ -824,6 +828,19 @@ def _add_callback_to_queue(self, request: CallbackRequest) -> None: team_name = self._get_team_name(file_info.bundle_name) stats.incr("dag_processing.other_callback_count", tags=prune_dict({"team_name": team_name})) + def _log_dropped_lang_sdk_callback(self, request: CallbackRequest) -> None: + if isinstance(request, DagCallbackRequest): + target = f"dag_id={request.dag_id} run_id={request.run_id}" + else: + ti = request.ti + target = f"dag_id={ti.dag_id} run_id={ti.run_id} task_id={ti.task_id}" + self.log.warning( + "Dropping %s for %s (%s): Lang-SDK runtimes do not run callbacks", + type(request).__name__, + request.filepath, + target, + ) + @provide_session def get_bundle_state(self, bundle_name: str, *, session: Session = NEW_SESSION) -> BundleState | None: """ @@ -1486,12 +1503,6 @@ def _create_process(self, dag_file: DagFileInfo) -> BaseDagFileProcessorProcess: ) if get_claiming_coordinator(dag_file.absolute_path, dag_file.bundle_name) is not None: - if callback_to_execute_for_file: - self.log.warning( - "Dropping %d callbacks for %s: Lang-SDK runtimes do not run callbacks", - len(callback_to_execute_for_file), - dag_file.rel_path, - ) return LangSDKDagFileProcessorProcess.start(**kwargs) return DagFileProcessorProcess.start(callbacks=callback_to_execute_for_file, **kwargs) diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 7a09829038889..07a022081d8c1 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -46,7 +46,8 @@ from uuid6 import uuid7 from airflow._shared.timezones import timezone -from airflow.callbacks.callback_requests import DagCallbackRequest +from airflow.api_fastapi.execution_api.datamodels.taskinstance import TaskInstance as TIDataModel +from airflow.callbacks.callback_requests import DagCallbackRequest, TaskCallbackRequest from airflow.dag_processing.bundles.base import BaseDagBundle, BundleVersion from airflow.dag_processing.bundles.manager import DagBundlesManager from airflow.dag_processing.collection import update_dag_parsing_results_in_db @@ -1700,33 +1701,56 @@ def test_create_process_keeps_the_python_parse_for_other_files(self, mock_get_lo mock_start.assert_called_once() mock_lang_sdk_start.assert_not_called() - @mock.patch.object(DagFileProcessorManager, "_get_logger_for_dag_file", autospec=True) - def test_create_process_drops_callbacks_for_a_coordinator_file(self, mock_get_logger, tmp_path, caplog): - mock_get_logger.return_value = (MagicMock(), MagicMock()) - dag_file = DagFileInfo(bundle_name="testing", rel_path=Path("dags.native"), bundle_path=tmp_path) - callback = DagCallbackRequest( - filepath="dags.native", - dag_id="native_dag", - run_id="run", - bundle_name="testing", - bundle_version=None, - is_failure_callback=True, - ) - + @pytest.mark.parametrize( + ("request_", "expected"), + [ + pytest.param( + DagCallbackRequest( + filepath="dags.native", + dag_id="native_dag", + run_id="run", + bundle_name="testing", + bundle_version="v1", + ), + "Dropping DagCallbackRequest for dags.native (dag_id=native_dag run_id=run): " + "Lang-SDK runtimes do not run callbacks", + id="dag", + ), + pytest.param( + TaskCallbackRequest( + filepath="dags.native", + ti=TIDataModel( + id=uuid7(), + dag_id="native_dag", + task_id="extract", + run_id="run", + try_number=1, + dag_version_id=uuid7(), + ), + bundle_name="testing", + bundle_version="v1", + ), + "Dropping TaskCallbackRequest for dags.native (dag_id=native_dag run_id=run " + "task_id=extract): Lang-SDK runtimes do not run callbacks", + id="task", + ), + ], + ) + @mock.patch.object(DagFileProcessorManager, "prepare_callback_bundle", autospec=True) + def test_a_callback_for_a_coordinator_file_is_dropped( + self, mock_prepare_callback_bundle, request_, expected, caplog + ): with ( fake_coordinator(), - mock.patch.object(LangSDKDagFileProcessorProcess, "start", autospec=True) as mock_start, caplog.at_level(logging.WARNING, logger="airflow.dag_processing.manager"), ): manager = DagFileProcessorManager(max_runs=1) - manager._callback_to_execute[dag_file] = [callback] - manager._create_process(dag_file) + manager._add_callback_to_queue(request_) - assert "callbacks" not in mock_start.call_args.kwargs - assert dag_file not in manager._callback_to_execute - assert [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] == [ - "Dropping 1 callbacks for dags.native: Lang-SDK runtimes do not run callbacks" - ] + mock_prepare_callback_bundle.assert_not_called() + assert not manager._callback_to_execute + assert not manager._file_queue + assert [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] == [expected] @mock.patch.object(FakeCoordinator, "parse_dag", autospec=True) @mock.patch.object( From a6359232b06b74b422da4ca9906caa404ce1946c Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Fri, 2 Oct 2026 02:42:21 +0000 Subject: [PATCH 05/24] Start each Dag file processor with explicit arguments Building the shared arguments as a dict[str, Any] turned off mypy's argument checks for both start() calls. --- .../src/airflow/dag_processing/manager.py | 26 ++++++++++++++----- 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index 95ff65786d5d6..6526ee8637392 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -1489,24 +1489,36 @@ def _create_process(self, dag_file: DagFileInfo) -> BaseDagFileProcessorProcess: callback_to_execute_for_file = self._callback_to_execute.pop(dag_file, []) logger, logger_filehandle = self._get_logger_for_dag_file(dag_file) - kwargs: dict[str, Any] = dict( + subprocess_logs_to_stdout = conf.get("logging", "dag_processor_log_target") == "stdout" + + if get_claiming_coordinator(dag_file.absolute_path, dag_file.bundle_name) is not None: + return LangSDKDagFileProcessorProcess.start( + id=id, + path=dag_file.absolute_path, + bundle_path=cast("Path", dag_file.bundle_path), + bundle_name=dag_file.bundle_name, + dag_file_rel_path=str(dag_file.rel_path), + selector=self.selector, + logger=logger, + logger_filehandle=logger_filehandle, + subprocess_logs_to_stdout=subprocess_logs_to_stdout, + client=self.client, + ) + + return DagFileProcessorProcess.start( id=id, path=dag_file.absolute_path, bundle_path=cast("Path", dag_file.bundle_path), bundle_name=dag_file.bundle_name, dag_file_rel_path=str(dag_file.rel_path), + callbacks=callback_to_execute_for_file, selector=self.selector, logger=logger, logger_filehandle=logger_filehandle, - subprocess_logs_to_stdout=conf.get("logging", "dag_processor_log_target") == "stdout", + subprocess_logs_to_stdout=subprocess_logs_to_stdout, client=self.client, ) - if get_claiming_coordinator(dag_file.absolute_path, dag_file.bundle_name) is not None: - return LangSDKDagFileProcessorProcess.start(**kwargs) - - return DagFileProcessorProcess.start(callbacks=callback_to_execute_for_file, **kwargs) - def _start_new_processes(self): """Start more processors if we have enough slots and files to process.""" bundle_to_team = self._get_team_names({file.bundle_name for file in self._file_queue}) From 17f65afef024161f66fdbe4c334bff4502127f8d Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Fri, 2 Oct 2026 02:48:37 +0000 Subject: [PATCH 06/24] Find the claiming coordinator in the Task SDK and wrap it in core The Task SDK now owns the lookup of the coordinator whose runtime parses a file, so the task runner can share it. Core keeps its logging wrapper. Its runtime SDK import is recorded in the SDK import inventory instead of being hidden with noqa, and the broken-config test uses a real clashing config. --- .../dag_processing/importer_routing.py | 22 +++++-------------- .../dag_processing/test_importer_routing.py | 21 ++++++++++++++---- generated/known_sdk_imports_in_core.txt | 1 + .../airflow/sdk/coordinators/_dag_importer.py | 12 ++++++++++ 4 files changed, 35 insertions(+), 21 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/importer_routing.py b/airflow-core/src/airflow/dag_processing/importer_routing.py index 956723eae26be..dafbb00f2928b 100644 --- a/airflow-core/src/airflow/dag_processing/importer_routing.py +++ b/airflow-core/src/airflow/dag_processing/importer_routing.py @@ -21,11 +21,9 @@ import logging import os -from pathlib import Path from typing import TYPE_CHECKING -from airflow.sdk.coordinators._dag_importer import CoordinatorDagImporter # noqa: SDK001 -from airflow.sdk.importers import DagImporterRegistry, get_importer_registry # noqa: SDK001 +from airflow.sdk.coordinators._dag_importer import find_claiming_coordinator if TYPE_CHECKING: from airflow.sdk.coordinators._subprocess import SubprocessCoordinator # noqa: SDK001 @@ -33,27 +31,17 @@ log = logging.getLogger(__name__) -def _get_registry(bundle_name: str | None) -> DagImporterRegistry | None: - try: - return get_importer_registry(bundle_name) - except Exception: - log.exception("Cannot build the Dag importer registry for bundle %s", bundle_name) - return None - - def get_claiming_coordinator( path: str | os.PathLike[str], bundle_name: str | None ) -> SubprocessCoordinator | None: """ Return the coordinator whose runtime parses ``path``, or ``None`` when a Python child parses it. - A runtime parses the file when its importer is a coordinator's Dag importer. + A runtime parses the file when its importer is a coordinator's Dag importer. When the bundle's + importers cannot be loaded, the error is logged and a Python child parses the file. """ - if (registry := _get_registry(bundle_name)) is None: - return None try: - importer = registry.get_importer(Path(path)) + return find_claiming_coordinator(path, bundle_name) except Exception: - log.exception("Cannot load the Dag importer for %s", path) + log.exception("Cannot load the Dag importer for %s in bundle %s", path, bundle_name) return None - return importer.coordinator if isinstance(importer, CoordinatorDagImporter) else None diff --git a/airflow-core/tests/unit/dag_processing/test_importer_routing.py b/airflow-core/tests/unit/dag_processing/test_importer_routing.py index 982fdc0cb680f..4b554f35a2221 100644 --- a/airflow-core/tests/unit/dag_processing/test_importer_routing.py +++ b/airflow-core/tests/unit/dag_processing/test_importer_routing.py @@ -17,10 +17,12 @@ # under the License. from __future__ import annotations -from unittest import mock +import json +import logging from airflow.dag_processing.importer_routing import get_claiming_coordinator +from tests_common.test_utils.config import conf_vars from unit.dag_processing.fake_lang_sdk import FakeCoordinator, fake_coordinator @@ -33,6 +35,17 @@ def test_get_claiming_coordinator_returns_the_coordinator_of_its_importer(tmp_pa assert others == [None, None] -@mock.patch("airflow.dag_processing.importer_routing._get_registry", autospec=True, return_value=None) -def test_get_claiming_coordinator_without_a_registry(mock_registry, tmp_path): - assert get_claiming_coordinator(tmp_path / "dags.native", "testing") is None +def test_get_claiming_coordinator_logs_a_broken_configuration(tmp_path, caplog): + spec = {"classpath": f"{FakeCoordinator.__module__}.FakeCoordinator", "kwargs": {}} + with ( + conf_vars({("sdk", "coordinators"): json.dumps({"first": spec, "second": spec})}), + caplog.at_level(logging.ERROR, logger="airflow.dag_processing.importer_routing"), + ): + assert get_claiming_coordinator(tmp_path / "dags.native", "testing") is None + + [record] = caplog.records + assert ( + record.getMessage() + == f"Cannot load the Dag importer for {tmp_path / 'dags.native'} in bundle testing" + ) + assert "Coordinators 'first' and 'second' both parse .native files" in str(record.exc_info[1]) diff --git a/generated/known_sdk_imports_in_core.txt b/generated/known_sdk_imports_in_core.txt index 7e0c840c79eed..8e8b4aba808a0 100644 --- a/generated/known_sdk_imports_in_core.txt +++ b/generated/known_sdk_imports_in_core.txt @@ -4,6 +4,7 @@ airflow-core/src/airflow/cli/commands/task_command.py::7 airflow-core/src/airflow/cli/commands/triggerer_command.py::1 airflow-core/src/airflow/configuration.py::1 airflow-core/src/airflow/dag_processing/dagbag.py::3 +airflow-core/src/airflow/dag_processing/importer_routing.py::1 airflow-core/src/airflow/dag_processing/lang_sdk_processor.py::7 airflow-core/src/airflow/dag_processing/manager.py::5 airflow-core/src/airflow/dag_processing/processor.py::14 diff --git a/task-sdk/src/airflow/sdk/coordinators/_dag_importer.py b/task-sdk/src/airflow/sdk/coordinators/_dag_importer.py index 54517ee38310d..3cfe6ff8fc8e4 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_dag_importer.py +++ b/task-sdk/src/airflow/sdk/coordinators/_dag_importer.py @@ -160,3 +160,15 @@ def find_claiming_importer( file = Path(path) importer = registry.get_importer(file) return importer if isinstance(importer, CoordinatorDagImporter) and importer.can_handle(file) else None + + +def find_claiming_coordinator( + path: str | os.PathLike[str], bundle_name: str | None +) -> SubprocessCoordinator | None: + """ + Return the coordinator whose runtime parses *path*, or ``None`` when a Python Dag importer parses it. + + :raises InvalidCoordinatorError: when the file's importer cannot find its coordinator. + """ + importer = find_claiming_importer(path, bundle_name) + return importer.get_parsing_coordinator() if importer is not None else None From 2d6dd214ff5f344800340c212487731c790a7d60 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Fri, 2 Oct 2026 02:49:07 +0000 Subject: [PATCH 07/24] Fail a native Dag's task on a Python worker without retries The error now names the coordinator that parses the Dag and suggests giving the Dag's tasks their own queue, since mapping a shared queue such as default would send every Python task on it to the runtime. The task reports FAILED instead of exiting 1, so it does not retry a misconfiguration that retrying cannot fix. A broken coordinators config leaves the file to the Dag bag, which reports it. --- .../airflow/sdk/execution_time/task_runner.py | 52 +++++--- .../execution_time/test_task_runner.py | 119 ++++++++++-------- 2 files changed, 103 insertions(+), 68 deletions(-) diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py index c4bf7f9f5b29a..9ee8e9427e73a 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -31,7 +31,7 @@ from datetime import datetime, timedelta, timezone from itertools import product from pathlib import Path -from typing import TYPE_CHECKING, Annotated, Any, Literal, cast +from typing import TYPE_CHECKING, Annotated, Any, Literal, NoReturn, cast from urllib.parse import quote import attrs @@ -60,7 +60,7 @@ from airflow.sdk.bases.skipmixin import XCOM_SKIPMIXIN_KEY from airflow.sdk.bases.xcom import BaseXCom from airflow.sdk.configuration import conf -from airflow.sdk.coordinators._dag_importer import CoordinatorDagImporter +from airflow.sdk.coordinators._dag_importer import find_claiming_coordinator from airflow.sdk.definitions._internal.dag_parsing_context import _airflow_parsing_context_manager from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet, is_arg_set from airflow.sdk.definitions.asset import ( @@ -143,6 +143,7 @@ get_previous_dagrun_success, set_current_context, ) +from airflow.sdk.execution_time.coordinator import get_coordinator_manager from airflow.sdk.execution_time.email_backend import ( _DEFAULT_EMAIL_BACKEND, _ErrorEmailNotifier, @@ -151,7 +152,6 @@ from airflow.sdk.execution_time.sentry import Sentry from airflow.sdk.execution_time.tracing import detail_span from airflow.sdk.execution_time.xcom import XCom -from airflow.sdk.importers import get_importer_registry from airflow.sdk.listener import get_listener_manager from airflow.sdk.observability.metrics import stats_utils from airflow.sdk.serde import allow_class, iter_pydantic_models @@ -1020,14 +1020,38 @@ def _register_deserialization_allowed_classes(dag, log: Logger) -> None: ) -def _is_lang_sdk_dag_file(path: str, bundle_name: str) -> bool: - """Return whether a coordinator's Dag importer claims *path*, so that a Lang-SDK runtime parses it.""" +def _get_claiming_coordinator_key(path: str, bundle_name: str) -> str | None: + """Return the ``[sdk] coordinators`` key of the coordinator whose runtime parses *path*, or ``None``.""" try: - importer = get_importer_registry(bundle_name).get_importer(path) + if (coordinator := find_claiming_coordinator(path, bundle_name)) is None: + return None + coordinators = get_coordinator_manager().for_bundle(bundle_name) except Exception: # Building the Dag bag reports a broken importer configuration. - return False - return isinstance(importer, CoordinatorDagImporter) + return None + return next((key for key, c in coordinators.items() if c is coordinator), None) + + +def _fail_lang_sdk_task(what: StartupDetails, coordinator_key: str, log: Logger) -> NoReturn: + """Fail a task of a native Lang-SDK Dag without retries: running it again cannot help.""" + log.error( + "This task belongs to a native Lang-SDK Dag and cannot run in Python. Give the Dag's tasks " + "their own queue and map it to coordinator %r in [sdk] queue_to_coordinator", + coordinator_key, + dag_id=what.ti.dag_id, + task_id=what.ti.task_id, + queue=what.ti.queue, + path=what.dag_rel_path, + ) + try: + SUPERVISOR_COMMS.send( + TaskState(state=TaskInstanceState.FAILED, end_date=datetime.now(tz=timezone.utc)) + ) + except Exception: + log.exception("Failed to report terminal task state to supervisor", state=TaskInstanceState.FAILED) + sys.exit(1) + # The supervisor keeps the reported state only when the process exits with 0. + sys.exit(0) @detail_span("parse") @@ -1042,16 +1066,8 @@ def parse(what: StartupDetails, log: Logger) -> RuntimeTaskInstance: bundle_prepare_ms = int((time.monotonic() - bundle_prepare_start) * 1000) dag_absolute_path = os.fspath(Path(bundle_instance.path, what.dag_rel_path)) - if _is_lang_sdk_dag_file(dag_absolute_path, bundle_info.name): - log.error( - "A task of a native Lang-SDK Dag cannot run in Python. Route its queue to the coordinator " - "that parses the Dag, with [sdk] queue_to_coordinator", - dag_id=what.ti.dag_id, - task_id=what.ti.task_id, - queue=what.ti.queue, - path=what.dag_rel_path, - ) - sys.exit(1) + if (coordinator_key := _get_claiming_coordinator_key(dag_absolute_path, bundle_info.name)) is not None: + _fail_lang_sdk_task(what, coordinator_key, log) dag_file_parse_start = time.monotonic() bag = BundleDagBag( diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py index 26ae41ac01eb5..73999b65b61cb 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py +++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py @@ -179,7 +179,7 @@ TaskRunnerMarker, _defer_task, _execute_task, - _is_lang_sdk_dag_file, + _get_claiming_coordinator_key, _make_task_span, _push_xcom_if_needed, _register_deserialization_allowed_classes, @@ -193,7 +193,7 @@ startup, ) from airflow.sdk.execution_time.xcom import XCom -from airflow.sdk.importers import DagSourceCode, reset_importer_registry +from airflow.sdk.importers import DagSourceCode from airflow.sdk.serde import deserialize from airflow.triggers.base import BaseEventTrigger, BaseTrigger, TriggerEvent from airflow.triggers.callback import CallbackTrigger @@ -919,29 +919,17 @@ def get_source_code(self, definition): class NativeCoordinator(SubprocessCoordinator): """A coordinator whose Dag importer claims ``.native`` files in every bundle.""" - def get_dag_importer(self): - return NativeDagImporter(coordinator=self) + @classmethod + def get_dag_importer_class(cls): + return NativeDagImporter -@patch("airflow.dag_processing.dagbag.BundleDagBag", autospec=True) -def test_parse_rejects_a_task_of_a_native_dag(mock_bag, tmp_path: Path, make_ti_context): +NATIVE_COORDINATOR_SPEC = {"classpath": f"{__name__}.NativeCoordinator", "kwargs": {}} + + +@pytest.fixture +def native_dag_startup(tmp_path: Path, make_ti_context): tmp_path.joinpath("dag.native").write_text("{}") - what = StartupDetails( - ti=TaskInstance( - id=uuid7(), - task_id="a", - dag_id="native_dag", - run_id="c", - try_number=1, - dag_version_id=uuid7(), - queue="default", - ), - dag_rel_path="dag.native", - bundle_info=BundleInfo(name="my-bundle", version=None), - ti_context=make_ti_context(), - start_date=timezone.utcnow(), - sentry_integration="", - ) bundle_config = [ { "name": "my-bundle", @@ -949,29 +937,50 @@ def test_parse_rejects_a_task_of_a_native_dag(mock_bag, tmp_path: Path, make_ti_ "kwargs": {"path": str(tmp_path), "refresh_interval": 1}, } ] - coordinators = {"native": {"classpath": f"{__name__}.NativeCoordinator", "kwargs": {}}} + with conf_vars( + { + ("dag_processor", "dag_bundle_config_list"): json.dumps(bundle_config), + ("sdk", "coordinators"): json.dumps({"native": NATIVE_COORDINATOR_SPEC}), + } + ): + yield StartupDetails( + ti=TaskInstance( + id=uuid7(), + task_id="a", + dag_id="native_dag", + run_id="c", + try_number=1, + dag_version_id=uuid7(), + queue="default", + ), + dag_rel_path="dag.native", + bundle_info=BundleInfo(name="my-bundle", version=None), + ti_context=make_ti_context(), + start_date=timezone.utcnow(), + sentry_integration="", + ) + + +@patch("airflow.dag_processing.dagbag.BundleDagBag", autospec=True) +def test_parse_fails_a_task_of_a_native_dag_without_retries( + mock_bag, native_dag_startup, mock_supervisor_comms, time_machine +): + instant = timezone.datetime(2024, 11, 22) + time_machine.move_to(instant, tick=False) log = mock.Mock() - reset_importer_registry() - try: - with ( - patch.dict( - os.environ, - { - "AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(bundle_config), - "AIRFLOW__SDK__COORDINATORS": json.dumps(coordinators), - }, - ), - pytest.raises(SystemExit, match="1"), - ): - parse(what, log) - finally: - reset_importer_registry() + with pytest.raises(SystemExit) as ctx: + parse(native_dag_startup, log) + assert ctx.value.code == 0 mock_bag.assert_not_called() + mock_supervisor_comms.send.assert_called_once_with( + TaskState(state=TaskInstanceState.FAILED, end_date=instant) + ) log.error.assert_called_once_with( - "A task of a native Lang-SDK Dag cannot run in Python. Route its queue to the coordinator " - "that parses the Dag, with [sdk] queue_to_coordinator", + "This task belongs to a native Lang-SDK Dag and cannot run in Python. Give the Dag's tasks " + "their own queue and map it to coordinator %r in [sdk] queue_to_coordinator", + "native", dag_id="native_dag", task_id="a", queue="default", @@ -979,19 +988,29 @@ def test_parse_rejects_a_task_of_a_native_dag(mock_bag, tmp_path: Path, make_ti_ ) +@patch("airflow.dag_processing.dagbag.BundleDagBag", autospec=True) +def test_parse_exits_non_zero_when_the_failure_cannot_be_reported( + mock_bag, native_dag_startup, mock_supervisor_comms +): + mock_supervisor_comms.send.side_effect = ConnectionError("supervisor gone") + + with pytest.raises(SystemExit, match="1"): + parse(native_dag_startup, mock.Mock()) + + @pytest.mark.parametrize( ("file_name", "expected"), - [("dag.native", True), ("dag.py", False), ("dag.pyc", False), ("dags.zip", False)], + [("dag.native", "native"), ("dag.py", None), ("dag.pyc", None), ("dags.zip", None)], ) -def test_is_lang_sdk_dag_file(file_name, expected): - coordinators = {"native": {"classpath": f"{__name__}.NativeCoordinator", "kwargs": {}}} - - reset_importer_registry() - try: - with patch.dict(os.environ, {"AIRFLOW__SDK__COORDINATORS": json.dumps(coordinators)}): - assert _is_lang_sdk_dag_file(f"/bundle/{file_name}", "my-bundle") is expected - finally: - reset_importer_registry() +def test_get_claiming_coordinator_key(file_name, expected): + with conf_vars({("sdk", "coordinators"): json.dumps({"native": NATIVE_COORDINATOR_SPEC})}): + assert _get_claiming_coordinator_key(f"/bundle/{file_name}", "my-bundle") == expected + + +def test_get_claiming_coordinator_key_with_a_clashing_config(): + coordinators = {"first": NATIVE_COORDINATOR_SPEC, "second": NATIVE_COORDINATOR_SPEC} + with conf_vars({("sdk", "coordinators"): json.dumps(coordinators)}): + assert _get_claiming_coordinator_key("/bundle/dag.native", "my-bundle") is None @pytest.mark.parametrize("use_queues", [False, True]) From 1169dbb91fab4a2bb75ef7d6e6e3a031b761f8ce Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Fri, 2 Oct 2026 02:50:38 +0000 Subject: [PATCH 08/24] Match the schema error's JSON path in a Lang-SDK parse test A schema validation error now names the path of the invalid field. --- .../tests/unit/dag_processing/test_lang_sdk_processor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py index dd3075a7c3381..a91e4811478ae 100644 --- a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py @@ -186,7 +186,7 @@ def reply(request, comms): [ pytest.param( {"max_active_runs": "many"}, - "Dag 'broken_dag' does not match the schema: 'many' is not of type 'number'", + "Dag 'broken_dag' does not match the schema at $.dag.max_active_runs: 'many' is not of type 'number'", id="schema", ), pytest.param( From 303b0be7821f968137f0f4b6a86f8f05f494bf84 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Fri, 2 Oct 2026 02:53:51 +0000 Subject: [PATCH 09/24] Kill what a Lang-SDK runtime leaves in its process group A process the runtime started could keep its output open after the runtime exited. Every kill then returned early, so the parse never finished, its result was lost and the leftover kept running. Once the runtime has exited, the grace kill, the timeout kill and close() now kill its process group. The kill also waits a bounded time, so a runtime stuck in the kernel does not stall the Dag processor loop. --- .../dag_processing/lang_sdk_processor.py | 43 ++++++++++-- .../dag_processing/test_lang_sdk_processor.py | 66 ++++++++++++++++++- 2 files changed, 102 insertions(+), 7 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py b/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py index dc7d2fca7bcfc..afa75e4c040aa 100644 --- a/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py +++ b/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py @@ -23,12 +23,14 @@ import selectors import signal import time +from contextlib import suppress from pathlib import Path from socket import 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 @@ -172,6 +174,7 @@ class LangSDKDagFileProcessorProcess(BaseDagFileProcessorProcess): _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] @@ -475,10 +478,17 @@ def is_ready(self) -> bool: self._verify_connections() if ( self._parsing_result_monotonic is not None - and self._exit_code is None and time.monotonic() - self._parsing_result_monotonic > _EXIT_GRACE_PERIOD ): - self.process_log.warning("The Lang-SDK runtime did not exit after its parse result; killing it") + 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 @@ -506,16 +516,41 @@ def _time_out(self, timeout: float) -> None: self._kill_runtime() def _kill_runtime(self) -> None: - """Kill the runtime and wait for it, without servicing its sockets, whose handler may have failed.""" + """ + 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._exit_code = self._process.wait(timeout=None) + 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/test_lang_sdk_processor.py b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py index a91e4811478ae..ca5e030f0c955 100644 --- a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py @@ -22,11 +22,12 @@ import selectors import signal import socket +import subprocess import sys import time import uuid from pathlib import Path -from unittest.mock import ANY, MagicMock, patch +from unittest.mock import ANY, MagicMock, call, patch import psutil import pytest @@ -34,6 +35,7 @@ from airflow.configuration import conf from airflow.dag_processing.lang_sdk_processor import ( + _EXIT_GRACE_PERIOD, LangSDKDagFileProcessorProcess, LangSDKRuntimeSchemaVersion, _get_import_timeout, @@ -45,6 +47,7 @@ from airflow.sdk.exceptions import AirflowRuntimeError from airflow.sdk.execution_time import supervisor from airflow.sdk.execution_time.comms import GetVariable, MaskSecret, _RequestFrame +from airflow.sdk.execution_time.supervisor import PsutilTracker from airflow.sdk.importers import DagSourceCode from airflow.serialization.serialized_objects import DagSerialization, LazyDeserializedDAG @@ -78,6 +81,14 @@ def _get_open_fds() -> set[int]: return {int(fd) for fd in os.listdir("/proc/self/fd")} if os.path.isdir("/proc/self/fd") else set() +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 + + @pytest.fixture(autouse=True) def _coordinator(): with fake_coordinator(): @@ -337,6 +348,29 @@ def reply(request, comms): 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.lang_sdk_processor._EXIT_GRACE_PERIOD", 0.5) + @patch.object(FakeCoordinator, "parse_dag", autospec=True) + def test_what_the_runtime_leaves_after_its_result_is_killed( + self, mock_parse_dag, parse, tmp_path, cap_structlog + ): + def reply(request, comms): + comms.send(_reply_with(_serialize_dag("native_dag"))(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_dag.side_effect = play_runtime(reply) + + proc = parse() + + assert [dag.dag_id for dag in proc.parsing_result.serialized_dags] == ["native_dag"] + 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 not _is_running(int((tmp_path / "leftover.pid").read_text())) + @pytest.mark.parametrize( ("policy", "error"), [ @@ -465,7 +499,6 @@ def test_the_import_timeout_holds_after_the_runtime_exits(self, mock_timeout, tm # The runtime exits, and the process it leaves behind keeps its output open. result = self._run(tmp_path, argv=["/bin/sh", "-c", "sleep 30 & exit 0"]) [proc] = [c.args[0] for c in mock_close.call_args_list] - os.killpg(proc.pid, signal.SIGKILL) assert result.import_errors == { "dag.native": f"The Lang-SDK runtime did not parse {tmp_path / 'dag.native'} within 1.0s" @@ -500,11 +533,11 @@ def test_only_a_positive_import_timeout_applies(mock_timeout, configured, expect def _make_process(**kwargs) -> LangSDKDagFileProcessorProcess: + kwargs.setdefault("process", MagicMock()) return LangSDKDagFileProcessorProcess( id=uuid.uuid4(), pid=1, stdin=MagicMock(), - process=MagicMock(), process_log=MagicMock(), selector=MagicMock(), bundle_name="testing", @@ -591,3 +624,30 @@ def test_the_schema_version_is_reported_once(mock_send_msg): assert proc._runtime_schema_version == OLDEST_SCHEMA_VERSION assert mock_send_msg.call_args.kwargs["error"].detail["message"] == "Unhandled request" + + +@pytest.mark.parametrize(("pid_reused", "killed"), [(False, True), (True, False)]) +@patch("airflow.dag_processing.lang_sdk_processor.os.killpg", autospec=True) +@patch("airflow.dag_processing.lang_sdk_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 + proc = _make_process(new_process_group=True) + proc._exit_code = 0 + + proc.close() + proc.close() + + assert mock_killpg.call_args_list == ([call(1, signal.SIGKILL)] if killed else []) + + +@patch.object(LangSDKDagFileProcessorProcess, "_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) From 19810c730b3b9b76f7a72dc0dbcae535c2d4e00d Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Fri, 2 Oct 2026 02:54:54 +0000 Subject: [PATCH 10/24] Name the setting in a Lang-SDK parse timeout error The import error is all a user sees of a timed-out parse, so it now names the limit to change: [core] dagbag_import_timeout or the get_dagbag_import_timeout policy, or [dag_processor] dag_file_processor_timeout before the runtime reports its import timeout. --- .../dag_processing/lang_sdk_processor.py | 18 +++++++++++++----- .../dag_processing/test_lang_sdk_processor.py | 9 ++++++--- .../tests/unit/dag_processing/test_manager.py | 5 ++++- 3 files changed, 23 insertions(+), 9 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py b/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py index afa75e4c040aa..2a91f9f7ecf39 100644 --- a/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py +++ b/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py @@ -69,6 +69,9 @@ # How long a runtime may keep running after its parse result, as Node does while a handle stays open. _EXIT_GRACE_PERIOD = 5.0 +_IMPORT_TIMEOUT_SETTING = "[core] dagbag_import_timeout or the get_dagbag_import_timeout policy" +_PROCESSOR_TIMEOUT_SETTING = "[dag_processor] dag_file_processor_timeout" + # StartLangSDKRuntime and LangSDKRuntimeSchemaVersion pass only between the manager and its forked # child before the exec, so they are not part of the supervisor schema the runtimes speak. @@ -257,11 +260,14 @@ def run( ) try: while not proc.is_ready: - timeout = proc._import_timeout if proc._schema_version_reported else processor_timeout + 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) + proc._time_out(timeout, setting) break proc._service_subprocess(max_wait_time=0.1) except BaseException: @@ -496,7 +502,7 @@ def is_ready(self) -> bool: and self._exit_code is None and time.monotonic() - self.start_time > self._import_timeout ): - self._time_out(self._import_timeout) + self._time_out(self._import_timeout, _IMPORT_TIMEOUT_SETTING) if self._check_subprocess_exit() is None: return False self._close_listeners() @@ -508,10 +514,12 @@ def is_ready(self) -> bool: ) return True - def _time_out(self, timeout: float) -> None: + 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 Lang-SDK runtime did not parse {self._parse_request.file} within {timeout}s, " + f"the limit set by {setting}" ) self._kill_runtime() diff --git a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py index ca5e030f0c955..ccd6cabd47608 100644 --- a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py @@ -481,7 +481,8 @@ def test_a_parse_past_the_import_timeout_is_killed(self, mock_timeout, tmp_path, result = self._run(tmp_path, argv=["/bin/sh", "-c", "exec sleep 60"]) assert result.import_errors == { - "dag.native": f"The Lang-SDK runtime did not parse {tmp_path / 'dag.native'} within 1.0s" + "dag.native": f"The Lang-SDK runtime did not parse {tmp_path / 'dag.native'} 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 @@ -501,7 +502,8 @@ def test_the_import_timeout_holds_after_the_runtime_exits(self, mock_timeout, tm [proc] = [c.args[0] for c in mock_close.call_args_list] assert result.import_errors == { - "dag.native": f"The Lang-SDK runtime did not parse {tmp_path / 'dag.native'} within 1.0s" + "dag.native": f"The Lang-SDK runtime did not parse {tmp_path / 'dag.native'} 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 @@ -519,7 +521,8 @@ def test_the_dag_file_processor_timeout_applies_until_the_import_timeout_is_repo result = self._run(tmp_path) assert result.import_errors == { - "dag.native": f"The Lang-SDK runtime did not parse {tmp_path / 'dag.native'} within 1.0s" + "dag.native": f"The Lang-SDK runtime did not parse {tmp_path / 'dag.native'} within 1.0s, " + "the limit set by [dag_processor] dag_file_processor_timeout" } diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 07a022081d8c1..0bb291b785455 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -1837,7 +1837,10 @@ def test_a_coordinator_file_past_the_import_timeout_is_an_import_error( [import_error] = session.scalars(select(ParseImportError)).all() assert import_error.filename == "slow.native" - assert import_error.stacktrace == f"The Lang-SDK runtime did not parse {slow} within 1.0s" + assert import_error.stacktrace == ( + f"The Lang-SDK runtime did not parse {slow} within 1.0s, " + "the limit set by [core] dagbag_import_timeout or the get_dagbag_import_timeout policy" + ) assert manager.selector.get_map() == {} @mock.patch("airflow.settings.get_dagbag_import_timeout", autospec=True) From b5ed50699b49dc3fa901f5a248f9ac615c4369c4 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Fri, 2 Oct 2026 02:55:49 +0000 Subject: [PATCH 11/24] Spec the mocks of a hand-built Lang-SDK processor in tests Tests assert through these mocks, so a misspelt method would pass. --- .../tests/unit/dag_processing/test_lang_sdk_processor.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py index ccd6cabd47608..091014c369dd1 100644 --- a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py @@ -32,6 +32,7 @@ import psutil import pytest import structlog +from structlog.typing import FilteringBoundLogger from airflow.configuration import conf from airflow.dag_processing.lang_sdk_processor import ( @@ -536,13 +537,13 @@ def test_only_a_positive_import_timeout_applies(mock_timeout, configured, expect def _make_process(**kwargs) -> LangSDKDagFileProcessorProcess: - kwargs.setdefault("process", MagicMock()) + kwargs.setdefault("process", MagicMock(spec=PsutilTracker)) return LangSDKDagFileProcessorProcess( id=uuid.uuid4(), pid=1, - stdin=MagicMock(), - process_log=MagicMock(), - selector=MagicMock(), + stdin=MagicMock(spec=socket.socket), + process_log=MagicMock(spec=FilteringBoundLogger), + selector=MagicMock(spec=selectors.BaseSelector), bundle_name="testing", dag_file_rel_path="dag.native", listeners={}, From 58c66651378039e95b0b10b3253fba54d9ff6003 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Fri, 2 Oct 2026 02:59:08 +0000 Subject: [PATCH 12/24] Apply the multi-team task checks to Lang-SDK Dags A Python Dag's tasks are checked for an executor available to the bundle's team, and tasks in the default pool move to the team's default pool. A Lang-SDK runtime's Dags now get the same: a Dag with an unavailable executor is the file's import error, and the pools are set in the stored payload. The Dag bag's helpers share the bundle team lookup so the processor can reuse them. --- .../src/airflow/dag_processing/dagbag.py | 43 ++++-------- .../dag_processing/lang_sdk_processor.py | 38 +++++++++-- .../dag_processing/test_lang_sdk_processor.py | 68 ++++++++++++++++++- 3 files changed, 116 insertions(+), 33 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/dagbag.py b/airflow-core/src/airflow/dag_processing/dagbag.py index 5d97ee9031961..37772d791a8bf 100644 --- a/airflow-core/src/airflow/dag_processing/dagbag.py +++ b/airflow-core/src/airflow/dag_processing/dagbag.py @@ -58,6 +58,7 @@ from airflow import DAG from airflow.models.dagwarning import DagWarning from airflow.sdk.importers import AbstractDagImporter, DagDefinition, DagImportWarning, DagSourceCode + from airflow.serialization.definitions.dag import SerializedDAG class FileLoadStat(NamedTuple): @@ -104,27 +105,24 @@ def _executor_exists(executor_name: str, team_name: str | None) -> bool: return False -def _validate_executor_fields(dag: DAG, bundle_name: str | None = None) -> None: +def _get_bundle_team_name(bundle_name: str | None) -> str | None: + """Return the team that owns *bundle_name* when multi-team is on, else ``None``.""" + if not conf.getboolean("core", "multi_team") or not bundle_name: + return None + from airflow.dag_processing.bundles.manager import DagBundlesManager + + return DagBundlesManager()._bundle_config[bundle_name].team_name + + +def _validate_executor_fields(dag: DAG | SerializedDAG, bundle_name: str | None = None) -> None: """Validate that executors specified in tasks are available and owned by the same team as the dag bundle.""" import logging log = logging.getLogger(__name__) - dag_team_name = None - # Check if multi team is available by reading the multi_team configuration (which is boolean) - if conf.getboolean("core", "multi_team"): - # Get team name from bundle configuration if available - if bundle_name: - from airflow.dag_processing.bundles.manager import DagBundlesManager - - bundle_manager = DagBundlesManager() - bundle_config = bundle_manager._bundle_config[bundle_name] - - dag_team_name = bundle_config.team_name - if dag_team_name: - log.debug( - "Found team '%s' for DAG '%s' via bundle '%s'", dag_team_name, dag.dag_id, bundle_name - ) + dag_team_name = _get_bundle_team_name(bundle_name) + if dag_team_name: + log.debug("Found team '%s' for DAG '%s' via bundle '%s'", dag_team_name, dag.dag_id, bundle_name) for task in dag.tasks: if not task.executor: @@ -151,18 +149,7 @@ def _assign_default_team_pools( bundle_name: str | None = None, ) -> None: """Assign the default team pool to tasks that do not explicitly specify a pool.""" - dag_team_name = None - - if conf.getboolean("core", "multi_team"): - if bundle_name: - from airflow.dag_processing.bundles.manager import DagBundlesManager - - bundle_manager = DagBundlesManager() - bundle_config = bundle_manager._bundle_config[bundle_name] - - dag_team_name = bundle_config.team_name - - if not dag_team_name: + if not (dag_team_name := _get_bundle_team_name(bundle_name)): return for task in dag.tasks: diff --git a/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py b/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py index 2a91f9f7ecf39..c91eb20efd3f7 100644 --- a/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py +++ b/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py @@ -18,6 +18,7 @@ from __future__ import annotations +import copy import functools import os import selectors @@ -36,6 +37,7 @@ from airflow import settings from airflow.configuration import conf +from airflow.dag_processing.dagbag import _get_bundle_team_name, _validate_executor_fields from airflow.dag_processing.importer_routing import get_claiming_coordinator from airflow.dag_processing.processor import ( BaseDagFileProcessorProcess, @@ -44,6 +46,7 @@ ToManager, ) from airflow.exceptions import DeserializationError +from airflow.models.pool import Pool from airflow.sdk.coordinators._subprocess import _is_connection_from_pid, _start_server from airflow.sdk.execution_time import supervisor from airflow.sdk.execution_time.comms import CommsDecoder, ErrorResponse, MaskSecret, _RequestFrame @@ -55,6 +58,7 @@ register_request_method, ) from airflow.sdk.importers import DagSourceCode, FilesystemDagDefinition +from airflow.serialization.enums import Encoding from airflow.serialization.serialized_objects import DagSerialization, LazyDeserializedDAG if TYPE_CHECKING: @@ -424,13 +428,17 @@ def _handle_parsing_result( DagSerialization.fill_config_defaults(dag.data) try: DagSerialization.validate_serialized_dag(dag.data) + self._apply_team_rules(dag.data) except DeserializationError as e: message = f"Cannot load the serialized Dag: {e}" - self.process_log.warning(message) - previous = import_errors.get(self.dag_file_rel_path) - import_errors[self.dag_file_rel_path] = f"{previous}\n{message}" if previous else message + except Exception as e: + message = f"{type(e).__name__}: {e}" + else: + serialized_dags.append(dag) continue - serialized_dags.append(dag) + self.process_log.warning(message) + previous = import_errors.get(self.dag_file_rel_path) + import_errors[self.dag_file_rel_path] = f"{previous}\n{message}" if previous else message self.parsing_result = msg.model_copy( update={ "serialized_dags": serialized_dags, @@ -441,6 +449,28 @@ def _handle_parsing_result( self._parsing_result_monotonic = time.monotonic() return None, {} + def _apply_team_rules(self, data: dict) -> None: + """ + Check each task's executor and move tasks in the default pool to the team's, as the Dag bag does. + + The bundle's team owns the Dag. *data* is a serialized Dag that validates, and its pools are changed + in place. + + :raises UnknownExecutorException: if a task's executor is not available to the team or globally. + """ + dag = DagSerialization.from_dict(copy.deepcopy(data)) + _validate_executor_fields(dag, self.bundle_name) + if not (team_name := _get_bundle_team_name(self.bundle_name)): + return + tasks = {task.task_id: task for task in dag.tasks} + for encoded in data["dag"]["tasks"]: + task_data = encoded[Encoding.VAR] + if tasks[task_data["task_id"]].pool != Pool.DEFAULT_POOL_NAME: + continue + # A mapped task reads its pool from its partial kwargs first. + target = task_data.setdefault("partial_kwargs", {}) if task_data.get("_is_mapped") else task_data + target["pool"] = Pool.get_default_team_pool_name(team_name) + def _read_dag_source_codes(self, serialized_dags: list[LazyDeserializedDAG]) -> dict[str, DagSourceCode]: """ Read the file's source with its Dag importer, for the fileloc of each Dag. diff --git a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py index 091014c369dd1..9e79e6ee43919 100644 --- a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py @@ -18,6 +18,7 @@ from __future__ import annotations import contextlib +import copy import os import selectors import signal @@ -42,7 +43,9 @@ _get_import_timeout, ) from airflow.dag_processing.processor import DagFileParseRequest, DagFileParsingResult -from airflow.sdk import DAG, BaseOperator +from airflow.exceptions import UnknownExecutorException +from airflow.executors.executor_loader import ExecutorLoader +from airflow.sdk import DAG, BaseOperator, task from airflow.sdk.api.client import Client from airflow.sdk.api.datamodels._generated import VariableResponse from airflow.sdk.exceptions import AirflowRuntimeError @@ -655,3 +658,66 @@ def test_killing_a_runtime_that_does_not_exit_waits_a_bounded_time(mock_signal): assert proc._exit_code is None mock_signal.assert_called_once_with(proc, signal.SIGKILL) process.wait.assert_called_once_with(_EXIT_GRACE_PERIOD) + + +@patch.object( + ExecutorLoader, + "lookup_executor_name_by_str", + autospec=True, + side_effect=UnknownExecutorException("not configured"), +) +def test_a_dag_with_an_unavailable_executor_is_an_import_error(mock_lookup): + with DAG("remote_dag", schedule=None) as remote_dag: + BaseOperator(task_id="extract", executor="no.such.Executor") + proc = _make_process() + + proc._handle_request( + DagFileParsingResult( + fileloc="/b/dag.native", + serialized_dags=[ + LazyDeserializedDAG(data=DagSerialization.to_dict(remote_dag)), + _serialize_dag("ok"), + ], + ), + MagicMock(), + 1, + ) + + assert [dag.dag_id for dag in proc.parsing_result.serialized_dags] == ["ok"] + assert proc.parsing_result.import_errors == { + "dag.native": "UnknownExecutorException: Task 'extract' specifies executor 'no.such.Executor', " + "which is not available. Make sure it is listed in your [core] executor configuration, or update " + "the task's executor to use one of the configured executors." + } + + +@conf_vars({("core", "multi_team"): "True"}) +@patch("airflow.dag_processing.bundles.manager.DagBundlesManager", autospec=True) +def test_tasks_in_the_default_pool_move_to_the_teams_pool(mock_bundles_manager): + mock_bundles_manager.return_value._bundle_config = {"testing": MagicMock(team_name="team_a")} + with DAG("team_dag", schedule=None) as dag: + BaseOperator(task_id="extract") + BaseOperator(task_id="load", pool="custom") + + @task + def fan_out(x): ... + + fan_out.expand(x=[1, 2]) + proc = _make_process() + + proc._handle_request( + DagFileParsingResult( + fileloc="/b/dag.native", serialized_dags=[LazyDeserializedDAG(data=DagSerialization.to_dict(dag))] + ), + MagicMock(), + 1, + ) + + [stored] = proc.parsing_result.serialized_dags + assert proc.parsing_result.import_errors is None + tasks = DagSerialization.from_dict(copy.deepcopy(stored.data)).tasks + assert {t.task_id: t.pool for t in tasks} == { + "extract": "default_pool_team_a", + "load": "custom", + "fan_out": "default_pool_team_a", + } From 913e53620c5e4f308075bee15e4018b16099263d Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Fri, 2 Oct 2026 03:00:14 +0000 Subject: [PATCH 13/24] Clear import errors after the native-file Dag bag sync test The test commits import error rows that the session fixture does not roll back, so they leaked into later tests. --- airflow-core/tests/unit/dag_processing/test_dagbag.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/airflow-core/tests/unit/dag_processing/test_dagbag.py b/airflow-core/tests/unit/dag_processing/test_dagbag.py index 0f7def4aad4af..00eafe90cf09b 100644 --- a/airflow-core/tests/unit/dag_processing/test_dagbag.py +++ b/airflow-core/tests/unit/dag_processing/test_dagbag.py @@ -1546,8 +1546,15 @@ def test_dagbag_no_bundle_path_no_syspath_modification(self, tmp_path): assert sys.path == syspath_before -def test_sync_bag_to_db_leaves_native_files_to_the_dag_processor(tmp_path, session, testing_dag_bundle): +@pytest.fixture +def clean_import_errors(): + db.clear_db_import_errors() + yield db.clear_db_import_errors() + + +@pytest.mark.usefixtures("clean_import_errors") +def test_sync_bag_to_db_leaves_native_files_to_the_dag_processor(tmp_path, session, testing_dag_bundle): write_native_file(tmp_path / "dags.native") session.add(ParseImportError(bundle_name="testing", filename="dags.native", stacktrace="stored")) session.commit() From 22cd5aac3cd31643568aa6328fa59d2198c2af3f Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Fri, 2 Oct 2026 03:00:54 +0000 Subject: [PATCH 14/24] Note in ADR-0010 that the manager routes claimed files for now The Dag processor manager starts the Lang-SDK processor for a coordinator-claimed file itself, which differs from the importer integration point the ADR describes. Record it as a bridge until importers own their parse process, and why. --- airflow-core/adr/lang-sdk/0010-native-dag-processing.md | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/airflow-core/adr/lang-sdk/0010-native-dag-processing.md b/airflow-core/adr/lang-sdk/0010-native-dag-processing.md index 82fe566151764..83ffb6f6c9401 100644 --- a/airflow-core/adr/lang-sdk/0010-native-dag-processing.md +++ b/airflow-core/adr/lang-sdk/0010-native-dag-processing.md @@ -152,6 +152,15 @@ DagModelOperation → PERSIST `DagImportResult.dags` is `list[DAG]`, so the importer wraps each serialized entry as a `LazyDeserializedDAG` and transforms it into an `airflow.sdk.DAG`. +### Interim: the Dag processor manager routes claimed files + +Until importers own their parse process (#73457), the Dag processor manager routes a coordinator-claimed file itself, as of #74035. When a file's importer in the bundle's +registry is a coordinator's, the manager starts `LangSDKDagFileProcessorProcess` for it instead of `DagFileProcessorProcess`. That process execs the runtime, which connects +back to it. The importer's `import_definition` only reports that the Dag processor parses the file, which is what a Dag bag outside the Dag processor sees. + +This is a bridge, not the end state. Running the runtime inside the Python parse child would need that child to relay every request and reply between the runtime and the +manager over its fd 0. Starting the runtime from the manager lets it reach the manager's request handlers directly, with no relay. + ## Consequences - The Java and TypeScript importers are never configured by hand. The runtime is declared once, in `[sdk] coordinators`, and the importer follows from it. From e434a396ff73fd78c8803e5cd57b180c34155eb6 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Fri, 2 Oct 2026 03:03:02 +0000 Subject: [PATCH 15/24] Move the client-less Lang-SDK parse out of this change Nothing here starts the Lang-SDK processor without an API client: the manager always passes one, and the first client-less caller is the Dag bag in a later change. Remove run(), the optional client and the gate that answered requests with an error without one, so the change that adds the caller can add them with it. --- .../dag_processing/lang_sdk_processor.py | 79 +----------- .../dag_processing/test_lang_sdk_processor.py | 115 +----------------- generated/known_sdk_imports_in_core.txt | 4 +- 3 files changed, 9 insertions(+), 189 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py b/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py index c91eb20efd3f7..79ed7e7e7cb9a 100644 --- a/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py +++ b/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py @@ -27,16 +27,14 @@ from contextlib import suppress from pathlib import Path from socket import socket -from typing import TYPE_CHECKING, Annotated, ClassVar, Literal, cast, get_args +from typing import TYPE_CHECKING, Annotated, ClassVar, Literal, get_args import attrs import msgspec import psutil from pydantic import BaseModel, Field, TypeAdapter -from uuid6 import uuid7 from airflow import settings -from airflow.configuration import conf from airflow.dag_processing.dagbag import _get_bundle_team_name, _validate_executor_fields from airflow.dag_processing.importer_routing import get_claiming_coordinator from airflow.dag_processing.processor import ( @@ -49,7 +47,7 @@ from airflow.models.pool import Pool from airflow.sdk.coordinators._subprocess import _is_connection_from_pid, _start_server from airflow.sdk.execution_time import supervisor -from airflow.sdk.execution_time.comms import CommsDecoder, ErrorResponse, MaskSecret, _RequestFrame +from airflow.sdk.execution_time.comms import CommsDecoder, ErrorResponse, _RequestFrame from airflow.sdk.execution_time.supervisor import ( ResponseSent, length_prefixed_frame_reader, @@ -66,7 +64,6 @@ from structlog.typing import FilteringBoundLogger - from airflow.sdk.api.client import Client from airflow.sdk.execution_time.supervisor import RequestHandler, RequestResult from airflow.typing_compat import Self @@ -74,7 +71,6 @@ _EXIT_GRACE_PERIOD = 5.0 _IMPORT_TIMEOUT_SETTING = "[core] dagbag_import_timeout or the get_dagbag_import_timeout policy" -_PROCESSOR_TIMEOUT_SETTING = "[dag_processor] dag_file_processor_timeout" # StartLangSDKRuntime and LangSDKRuntimeSchemaVersion pass only between the manager and its forked @@ -167,9 +163,6 @@ class LangSDKDagFileProcessorProcess(BaseDagFileProcessorProcess): ``DagFileParseRequest`` itself, so the request is sent once it has connected. """ - client: Client | None = None # type: ignore[assignment] - """Answers the runtime's requests; without one, as in a Dag bag, those that need it get an error.""" - decoder = TypeAdapter( Annotated[LangSDKRuntimeSchemaVersion | get_args(ToManager)[0], Field(discriminator="type")] ) @@ -234,53 +227,6 @@ def start( # type: ignore[override] ) return proc - @classmethod - def run( - cls, - *, - path: str | os.PathLike[str], - bundle_path: Path, - bundle_name: str, - dag_file_rel_path: str, - logger: FilteringBoundLogger, - ) -> DagFileParsingResult: - """ - Parse *path* outside the Dag processor and wait for the result. - - There is no API client, so each request of the runtime that needs one gets an error. The file's import - timeout bounds the parse, and ``[dag_processor] dag_file_processor_timeout`` until the parse child - reports it. - """ - processor_timeout = conf.getfloat("dag_processor", "dag_file_processor_timeout") - with selectors.DefaultSelector() as selector: - proc = cls.start( - id=uuid7(), - path=path, - bundle_path=bundle_path, - bundle_name=bundle_name, - dag_file_rel_path=dag_file_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 - proc._service_subprocess(max_wait_time=0.1) - except BaseException: - proc._kill_runtime() - raise - finally: - proc.close() - return cast("DagFileParsingResult", proc.parsing_result) - def _accept_connection(self, listener: socket, *, channel: _Channel) -> bool: try: conn, _ = listener.accept() @@ -495,20 +441,6 @@ def _read_dag_source_codes(self, serialized_dags: list[LazyDeserializedDAG]) -> **dict([register_request_method(LangSDKRuntimeSchemaVersion, _handle_runtime_schema_version)]), } - def _handle_request(self, msg, log: FilteringBoundLogger, req_id: int) -> None: - if self.client is None and not isinstance( - msg, (DagFileParsingResult, LangSDKRuntimeSchemaVersion, MaskSecret) - ): - self.send_msg( - None, - request_id=req_id, - error=ErrorResponse( - detail={"message": f"{type(msg).__name__} is answered only in the Dag processor"} - ), - ) - return - super()._handle_request(msg, log, req_id) - @property def is_ready(self) -> bool: self._verify_connections() @@ -532,7 +464,7 @@ def is_ready(self) -> bool: and self._exit_code is None and time.monotonic() - self.start_time > self._import_timeout ): - self._time_out(self._import_timeout, _IMPORT_TIMEOUT_SETTING) + self._time_out(self._import_timeout) if self._check_subprocess_exit() is None: return False self._close_listeners() @@ -544,12 +476,11 @@ def is_ready(self) -> bool: ) 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*.""" + def _time_out(self, timeout: float) -> None: 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}" + f"the limit set by {_IMPORT_TIMEOUT_SETTING}" ) self._kill_runtime() diff --git a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py index 9e79e6ee43919..89322c995568f 100644 --- a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py @@ -48,9 +48,8 @@ from airflow.sdk import DAG, BaseOperator, task from airflow.sdk.api.client import Client from airflow.sdk.api.datamodels._generated import VariableResponse -from airflow.sdk.exceptions import AirflowRuntimeError from airflow.sdk.execution_time import supervisor -from airflow.sdk.execution_time.comms import GetVariable, MaskSecret, _RequestFrame +from airflow.sdk.execution_time.comms import GetVariable, _RequestFrame from airflow.sdk.execution_time.supervisor import PsutilTracker from airflow.sdk.importers import DagSourceCode from airflow.serialization.serialized_objects import DagSerialization, LazyDeserializedDAG @@ -419,117 +418,6 @@ def test_the_runtime_inherits_only_its_standard_streams(self, monkeypatch, tmp_p assert fds["0"] == "/dev/null" -class TestRun: - @staticmethod - def _run(tmp_path, **spec) -> DagFileParsingResult: - return LangSDKDagFileProcessorProcess.run( - path=write_native_file(tmp_path / "dag.native", **spec), - bundle_path=tmp_path, - bundle_name="testing", - dag_file_rel_path="dag.native", - logger=structlog.get_logger(), - ) - - @patch.object(FakeCoordinator, "parse_dag", autospec=True) - def test_requests_get_an_error(self, mock_parse_dag, tmp_path): - def reply(request, comms): - with pytest.raises(AirflowRuntimeError) as ctx: - comms.send(GetVariable(key="native_var")) - description = ctx.value.error.detail["message"] - return _reply_with(_serialize_dag("native_dag", description=description))(request, comms) - - mock_parse_dag.side_effect = play_runtime(reply) - - result = self._run(tmp_path) - - assert result.serialized_dags[0].data["dag"]["description"] == ( - "GetVariable is answered only in the Dag processor" - ) - - @patch("airflow.sdk.execution_time.request_handlers.mask_secret", autospec=True) - @patch.object(FakeCoordinator, "parse_dag", autospec=True) - def test_a_secret_is_masked_without_a_client(self, mock_parse_dag, mock_mask_secret, tmp_path): - def reply(request, comms): - comms.send(MaskSecret(value="native-secret", name="native_conn")) - return _reply_with(_serialize_dag("native_dag"))(request, comms) - - mock_parse_dag.side_effect = play_runtime(reply) - - result = self._run(tmp_path) - - assert [dag.dag_id for dag in result.serialized_dags] == ["native_dag"] - mock_mask_secret.assert_called_once_with("native-secret", "native_conn") - - @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_dag", - autospec=True, - side_effect=play_runtime(lambda request, comms: time.sleep(60)), - ) - if connected - else contextlib.nullcontext(), - patch.object( - LangSDKDagFileProcessorProcess, - "close", - autospec=True, - side_effect=LangSDKDagFileProcessorProcess.close, - ) as mock_close, - ): - result = self._run(tmp_path, argv=["/bin/sh", "-c", "exec sleep 60"]) - - assert result.import_errors == { - "dag.native": f"The Lang-SDK runtime did not parse {tmp_path / 'dag.native'} 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 - - @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): - with patch.object( - LangSDKDagFileProcessorProcess, - "close", - autospec=True, - side_effect=LangSDKDagFileProcessorProcess.close, - ) as mock_close: - # The runtime exits, and the process it leaves behind keeps its output open. - result = self._run(tmp_path, argv=["/bin/sh", "-c", "sleep 30 & exit 0"]) - [proc] = [c.args[0] for c in mock_close.call_args_list] - - assert result.import_errors == { - "dag.native": f"The Lang-SDK runtime did not parse {tmp_path / 'dag.native'} 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 - - @conf_vars({("dag_processor", "dag_file_processor_timeout"): "1"}) - @patch.object( - FakeCoordinator, - "_build_parse_dag_command", - autospec=True, - side_effect=lambda self, *, path: time.sleep(60), - ) - def test_the_dag_file_processor_timeout_applies_until_the_import_timeout_is_reported( - self, mock_build_parse_dag_command, tmp_path - ): - result = self._run(tmp_path) - - assert result.import_errors == { - "dag.native": f"The Lang-SDK runtime did not parse {tmp_path / 'dag.native'} within 1.0s, " - "the limit set by [dag_processor] dag_file_processor_timeout" - } - - @pytest.mark.parametrize(("configured", "expected"), [(30, 30), (0.5, 0.5), (0, None), (-1, None)]) @patch("airflow.settings.get_dagbag_import_timeout", autospec=True) def test_only_a_positive_import_timeout_applies(mock_timeout, configured, expected): @@ -547,6 +435,7 @@ def _make_process(**kwargs) -> LangSDKDagFileProcessorProcess: stdin=MagicMock(spec=socket.socket), process_log=MagicMock(spec=FilteringBoundLogger), selector=MagicMock(spec=selectors.BaseSelector), + client=MagicMock(spec=Client), bundle_name="testing", dag_file_rel_path="dag.native", listeners={}, diff --git a/generated/known_sdk_imports_in_core.txt b/generated/known_sdk_imports_in_core.txt index 8e8b4aba808a0..51ae1d58804aa 100644 --- a/generated/known_sdk_imports_in_core.txt +++ b/generated/known_sdk_imports_in_core.txt @@ -4,8 +4,8 @@ airflow-core/src/airflow/cli/commands/task_command.py::7 airflow-core/src/airflow/cli/commands/triggerer_command.py::1 airflow-core/src/airflow/configuration.py::1 airflow-core/src/airflow/dag_processing/dagbag.py::3 -airflow-core/src/airflow/dag_processing/importer_routing.py::1 -airflow-core/src/airflow/dag_processing/lang_sdk_processor.py::7 +airflow-core/src/airflow/dag_processing/importer_routing.py::2 +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/exceptions.py::1 From 0e680a3b2e221c4ed8fabfe04b4edf8b1bd035cb Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Fri, 2 Oct 2026 05:39:23 +0000 Subject: [PATCH 16/24] Wait for the killed leftover in the Lang-SDK processor test SIGKILL is delivered asynchronously, so the leftover can still show as running right after the kill. --- .../unit/dag_processing/test_lang_sdk_processor.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py index 89322c995568f..21f5d7b50b53d 100644 --- a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py @@ -92,6 +92,16 @@ def _is_running(pid: int) -> bool: 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(): @@ -372,7 +382,7 @@ def reply(request, comms): "The Lang-SDK runtime left processes holding its output after its parse result; killing them" in cap_structlog ) - assert not _is_running(int((tmp_path / "leftover.pid").read_text())) + assert _stops_running(int((tmp_path / "leftover.pid").read_text())) @pytest.mark.parametrize( ("policy", "error"), From eb258495592cd9f44c5a9344332466790b7ff279 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Fri, 2 Oct 2026 11:20:42 +0000 Subject: [PATCH 17/24] Find a native Dag's coordinator key through its importer CoordinatorManager.for_bundle now returns Dag importers, so the task runner matches the claiming coordinator through each importer. --- task-sdk/src/airflow/sdk/execution_time/task_runner.py | 7 +++++-- task-sdk/tests/task_sdk/execution_time/test_task_runner.py | 5 ++--- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py index 9ee8e9427e73a..92553fe10377f 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -1025,11 +1025,14 @@ def _get_claiming_coordinator_key(path: str, bundle_name: str) -> str | None: try: if (coordinator := find_claiming_coordinator(path, bundle_name)) is None: return None - coordinators = get_coordinator_manager().for_bundle(bundle_name) + importers = get_coordinator_manager().for_bundle(bundle_name) except Exception: # Building the Dag bag reports a broken importer configuration. return None - return next((key for key, c in coordinators.items() if c is coordinator), None) + return next( + (key for key, importer in importers.items() if getattr(importer, "coordinator", None) is coordinator), + None, + ) def _fail_lang_sdk_task(what: StartupDetails, coordinator_key: str, log: Logger) -> NoReturn: diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py index 73999b65b61cb..4be6607de309b 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py +++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py @@ -919,9 +919,8 @@ def get_source_code(self, definition): class NativeCoordinator(SubprocessCoordinator): """A coordinator whose Dag importer claims ``.native`` files in every bundle.""" - @classmethod - def get_dag_importer_class(cls): - return NativeDagImporter + def get_dag_importer(self): + return NativeDagImporter(coordinator=self) NATIVE_COORDINATOR_SPEC = {"classpath": f"{__name__}.NativeCoordinator", "kwargs": {}} From 4ad06144a0e32cd277d00af9c8216fe4b179a843 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Fri, 2 Oct 2026 13:19:30 +0000 Subject: [PATCH 18/24] Pass a parse log handle to the Lang-SDK processor in tests The Dag processor always gives the Lang-SDK processor one. --- .../tests/unit/dag_processing/test_lang_sdk_processor.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py index 21f5d7b50b53d..a34a8f7f37125 100644 --- a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py @@ -28,6 +28,7 @@ import time import uuid from pathlib import Path +from typing import BinaryIO from unittest.mock import ANY, MagicMock, call, patch import psutil @@ -117,6 +118,7 @@ def _start(tmp_path, selector, *, client: Client | None = None, **spec) -> LangS dag_file_rel_path="dag.native", selector=selector, logger=structlog.get_logger(), + logger_filehandle=MagicMock(spec=BinaryIO), client=client or MagicMock(spec=Client), ) @@ -439,6 +441,7 @@ def test_only_a_positive_import_timeout_applies(mock_timeout, configured, expect def _make_process(**kwargs) -> LangSDKDagFileProcessorProcess: kwargs.setdefault("process", MagicMock(spec=PsutilTracker)) + kwargs.setdefault("logger_filehandle", MagicMock(spec=BinaryIO)) return LangSDKDagFileProcessorProcess( id=uuid.uuid4(), pid=1, @@ -537,13 +540,15 @@ def test_the_schema_version_is_reported_once(mock_send_msg): @patch("airflow.dag_processing.lang_sdk_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 - proc = _make_process(new_process_group=True) + 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() @patch.object(LangSDKDagFileProcessorProcess, "_signal_subprocess", autospec=True) From 435e639a5e9d1c045bab8e2d7515a3dfb3d30849 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Fri, 2 Oct 2026 03:11:45 +0000 Subject: [PATCH 19/24] Stub Dag file discovery with DagFileInfo in the Lang-SDK manager tests Discovery now returns DagFileInfo records instead of relative paths. --- .../tests/unit/dag_processing/test_manager.py | 26 ++++++++++++++++--- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 0bb291b785455..a5531ec57b888 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -317,6 +317,18 @@ def refresh(self) -> None: pass +def _found_files(*rel_paths: str): + """Stand in for ``_find_files_in_bundle``, listing *rel_paths* in the bundle.""" + + def find(manager, bundle): + return { + DagFileInfo(rel_path=Path(rel_path), bundle_name=bundle.name, bundle_path=bundle.path) + for rel_path in rel_paths + } + + return find + + class TestDagFileProcessorManager: @pytest.fixture(autouse=True) def _disable_examples(self): @@ -1754,7 +1766,10 @@ def test_a_callback_for_a_coordinator_file_is_dropped( @mock.patch.object(FakeCoordinator, "parse_dag", autospec=True) @mock.patch.object( - DagFileProcessorManager, "_find_files_in_bundle", autospec=True, return_value=[Path("good.native")] + DagFileProcessorManager, + "_find_files_in_bundle", + autospec=True, + side_effect=_found_files("good.native"), ) def test_coordinator_files_are_persisted( self, mock_find_files, mock_parse_dag, tmp_path, configure_testing_dag_bundle @@ -1788,7 +1803,7 @@ def reply(request, comms): DagFileProcessorManager, "_find_files_in_bundle", autospec=True, - return_value=[Path("garbage.native"), Path("python_dag.py")], + side_effect=_found_files("garbage.native", "python_dag.py"), ) def test_an_invalid_frame_does_not_stop_other_files_parsing( self, mock_find_files, mock_parse_dag, tmp_path, configure_testing_dag_bundle @@ -1821,7 +1836,10 @@ def reply(request, comms): @mock.patch("airflow.settings.get_dagbag_import_timeout", autospec=True, return_value=1) @mock.patch.object(FakeCoordinator, "parse_dag", autospec=True) @mock.patch.object( - DagFileProcessorManager, "_find_files_in_bundle", autospec=True, return_value=[Path("slow.native")] + DagFileProcessorManager, + "_find_files_in_bundle", + autospec=True, + side_effect=_found_files("slow.native"), ) def test_a_coordinator_file_past_the_import_timeout_is_an_import_error( self, mock_find_files, mock_parse_dag, mock_timeout, tmp_path, configure_testing_dag_bundle @@ -1848,7 +1866,7 @@ def test_a_coordinator_file_past_the_import_timeout_is_an_import_error( DagFileProcessorManager, "_find_files_in_bundle", autospec=True, - return_value=[Path("broken.native"), Path("python_dag.py")], + side_effect=_found_files("broken.native", "python_dag.py"), ) def test_a_failing_import_timeout_policy_fails_only_its_file( self, mock_find_files, mock_timeout, tmp_path, configure_testing_dag_bundle From e62a4815b8f5b6d7f5d9ef6fb56bc9fc5f39f6a4 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Sat, 3 Oct 2026 15:52:18 +0000 Subject: [PATCH 20/24] Name the coordinator class a native Dag task must be routed to A task of a native Dag that reaches a Python worker now fails with the coordinator class its Dag importer names, instead of the key of the coordinator that parses the Dag. The check no longer needs the parsing coordinator, so it also holds when several coordinators of that class are configured. --- .../airflow/sdk/execution_time/task_runner.py | 26 ++++---- .../execution_time/test_task_runner.py | 60 ++++++++++++------- 2 files changed, 50 insertions(+), 36 deletions(-) diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py index 92553fe10377f..27254fcb8f5e3 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -60,7 +60,7 @@ from airflow.sdk.bases.skipmixin import XCOM_SKIPMIXIN_KEY from airflow.sdk.bases.xcom import BaseXCom from airflow.sdk.configuration import conf -from airflow.sdk.coordinators._dag_importer import find_claiming_coordinator +from airflow.sdk.coordinators._dag_importer import find_claiming_importer from airflow.sdk.definitions._internal.dag_parsing_context import _airflow_parsing_context_manager from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet, is_arg_set from airflow.sdk.definitions.asset import ( @@ -143,7 +143,6 @@ get_previous_dagrun_success, set_current_context, ) -from airflow.sdk.execution_time.coordinator import get_coordinator_manager from airflow.sdk.execution_time.email_backend import ( _DEFAULT_EMAIL_BACKEND, _ErrorEmailNotifier, @@ -163,6 +162,7 @@ from pendulum.datetime import DateTime from structlog.typing import FilteringBoundLogger as Logger + from airflow.sdk.coordinators._dag_importer import CoordinatorDagImporter from airflow.sdk.definitions._internal.abstractoperator import AbstractOperator from airflow.sdk.definitions.context import Context from airflow.sdk.definitions.retry_policy import RetryDecision @@ -1020,27 +1020,21 @@ def _register_deserialization_allowed_classes(dag, log: Logger) -> None: ) -def _get_claiming_coordinator_key(path: str, bundle_name: str) -> str | None: - """Return the ``[sdk] coordinators`` key of the coordinator whose runtime parses *path*, or ``None``.""" +def _find_native_dag_importer(path: str, bundle_name: str) -> CoordinatorDagImporter | None: + """Return the coordinator Dag importer that claims *path*, so that a Lang-SDK runtime parses it.""" try: - if (coordinator := find_claiming_coordinator(path, bundle_name)) is None: - return None - importers = get_coordinator_manager().for_bundle(bundle_name) + return find_claiming_importer(path, bundle_name) except Exception: # Building the Dag bag reports a broken importer configuration. return None - return next( - (key for key, importer in importers.items() if getattr(importer, "coordinator", None) is coordinator), - None, - ) -def _fail_lang_sdk_task(what: StartupDetails, coordinator_key: str, log: Logger) -> NoReturn: +def _fail_lang_sdk_task(what: StartupDetails, coordinator_classpath: str, log: Logger) -> NoReturn: """Fail a task of a native Lang-SDK Dag without retries: running it again cannot help.""" log.error( "This task belongs to a native Lang-SDK Dag and cannot run in Python. Give the Dag's tasks " - "their own queue and map it to coordinator %r in [sdk] queue_to_coordinator", - coordinator_key, + "their own queue and route it to a %s in [sdk] queue_to_coordinator", + coordinator_classpath.rsplit(".", 1)[-1], dag_id=what.ti.dag_id, task_id=what.ti.task_id, queue=what.ti.queue, @@ -1069,8 +1063,8 @@ def parse(what: StartupDetails, log: Logger) -> RuntimeTaskInstance: bundle_prepare_ms = int((time.monotonic() - bundle_prepare_start) * 1000) dag_absolute_path = os.fspath(Path(bundle_instance.path, what.dag_rel_path)) - if (coordinator_key := _get_claiming_coordinator_key(dag_absolute_path, bundle_info.name)) is not None: - _fail_lang_sdk_task(what, coordinator_key, log) + if (importer := _find_native_dag_importer(dag_absolute_path, bundle_info.name)) is not None: + _fail_lang_sdk_task(what, importer.coordinator_classpath, log) dag_file_parse_start = time.monotonic() bag = BundleDagBag( diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py index 4be6607de309b..c1e24fabe9008 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py +++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py @@ -179,7 +179,7 @@ TaskRunnerMarker, _defer_task, _execute_task, - _get_claiming_coordinator_key, + _find_native_dag_importer, _make_task_span, _push_xcom_if_needed, _register_deserialization_allowed_classes, @@ -908,6 +908,7 @@ def test_parse_module_in_bundle_root(tmp_path: Path, make_ti_context): class NativeDagImporter(CoordinatorDagImporter): + coordinator_classpath = f"{__name__}.NativeCoordinator" artifact_suffix = ".native" supported_extensions = [".native"] @@ -917,10 +918,7 @@ def get_source_code(self, definition): @attrs.define(kw_only=True) class NativeCoordinator(SubprocessCoordinator): - """A coordinator whose Dag importer claims ``.native`` files in every bundle.""" - - def get_dag_importer(self): - return NativeDagImporter(coordinator=self) + """A coordinator of the runtime whose Dag importer claims ``.native`` files in every bundle.""" NATIVE_COORDINATOR_SPEC = {"classpath": f"{__name__}.NativeCoordinator", "kwargs": {}} @@ -936,11 +934,17 @@ def native_dag_startup(tmp_path: Path, make_ti_context): "kwargs": {"path": str(tmp_path), "refresh_interval": 1}, } ] - with conf_vars( - { - ("dag_processor", "dag_bundle_config_list"): json.dumps(bundle_config), - ("sdk", "coordinators"): json.dumps({"native": NATIVE_COORDINATOR_SPEC}), - } + with ( + mock.patch( + "airflow.sdk.coordinators._dag_importer.COORDINATOR_DAG_IMPORTERS", + (f"{__name__}.NativeDagImporter",), + ), + conf_vars( + { + ("dag_processor", "dag_bundle_config_list"): json.dumps(bundle_config), + ("sdk", "coordinators"): json.dumps({"native": NATIVE_COORDINATOR_SPEC}), + } + ), ): yield StartupDetails( ti=TaskInstance( @@ -978,8 +982,8 @@ def test_parse_fails_a_task_of_a_native_dag_without_retries( ) log.error.assert_called_once_with( "This task belongs to a native Lang-SDK Dag and cannot run in Python. Give the Dag's tasks " - "their own queue and map it to coordinator %r in [sdk] queue_to_coordinator", - "native", + "their own queue and route it to a %s in [sdk] queue_to_coordinator", + "NativeCoordinator", dag_id="native_dag", task_id="a", queue="default", @@ -997,19 +1001,35 @@ def test_parse_exits_non_zero_when_the_failure_cannot_be_reported( parse(native_dag_startup, mock.Mock()) +@patch("airflow.dag_processing.dagbag.BundleDagBag", autospec=True) +def test_parse_fails_a_task_of_a_native_dag_when_several_coordinators_could_parse_it( + mock_bag, native_dag_startup, mock_supervisor_comms +): + coordinators = {"first": NATIVE_COORDINATOR_SPEC, "second": NATIVE_COORDINATOR_SPEC} + log = mock.Mock() + + with conf_vars({("sdk", "coordinators"): json.dumps(coordinators)}), pytest.raises(SystemExit): + parse(native_dag_startup, log) + + mock_bag.assert_not_called() + assert log.error.call_args.args[1] == "NativeCoordinator" + + +@pytest.mark.usefixtures("native_dag_startup") @pytest.mark.parametrize( ("file_name", "expected"), - [("dag.native", "native"), ("dag.py", None), ("dag.pyc", None), ("dags.zip", None)], + [("dag.native", True), ("dag.py", False), ("dag.pyc", False), ("dags.zip", False)], ) -def test_get_claiming_coordinator_key(file_name, expected): - with conf_vars({("sdk", "coordinators"): json.dumps({"native": NATIVE_COORDINATOR_SPEC})}): - assert _get_claiming_coordinator_key(f"/bundle/{file_name}", "my-bundle") == expected +def test_find_native_dag_importer(file_name, expected): + importer = _find_native_dag_importer(f"/bundle/{file_name}", "my-bundle") + assert isinstance(importer, NativeDagImporter) is expected -def test_get_claiming_coordinator_key_with_a_clashing_config(): - coordinators = {"first": NATIVE_COORDINATOR_SPEC, "second": NATIVE_COORDINATOR_SPEC} - with conf_vars({("sdk", "coordinators"): json.dumps(coordinators)}): - assert _get_claiming_coordinator_key("/bundle/dag.native", "my-bundle") is None + +@pytest.mark.usefixtures("native_dag_startup") +@patch("airflow.sdk.coordinators._dag_importer.COORDINATOR_DAG_IMPORTERS", ("nonexistent.module.Importer",)) +def test_find_native_dag_importer_when_the_importers_cannot_be_built(): + assert _find_native_dag_importer("/bundle/dag.native", "my-bundle") is None @pytest.mark.parametrize("use_queues", [False, True]) From 18744e8f554f99e656609967ad0e74b85b57e5ff Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Sat, 3 Oct 2026 15:53:03 +0000 Subject: [PATCH 21/24] Find a claiming Dag importer instead of its coordinator in Dag processing get_claiming_importer replaces get_claiming_coordinator, and the Task SDK's find_claiming_coordinator is gone. The manager and the Dag bag sync only need to know whether an importer claims a file. The parse child asks the importer for its parsing coordinator, so a Dag bundle with no usable coordinator is the file's import error. The Dag source is read through the importer. The fake Lang-SDK test helpers set coordinator_classpath and register their importer in COORDINATOR_DAG_IMPORTERS. --- .../src/airflow/dag_processing/dagbag.py | 4 +-- .../dag_processing/importer_routing.py | 12 ++++----- .../dag_processing/lang_sdk_processor.py | 14 +++++----- .../src/airflow/dag_processing/manager.py | 6 ++--- .../unit/dag_processing/fake_lang_sdk.py | 17 +++++++----- .../dag_processing/test_importer_routing.py | 27 +++++++++++-------- .../dag_processing/test_lang_sdk_processor.py | 15 ++++++++++- generated/known_sdk_imports_in_core.txt | 2 +- .../airflow/sdk/coordinators/_dag_importer.py | 12 --------- 9 files changed, 58 insertions(+), 51 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/dagbag.py b/airflow-core/src/airflow/dag_processing/dagbag.py index 37772d791a8bf..b5822171de88d 100644 --- a/airflow-core/src/airflow/dag_processing/dagbag.py +++ b/airflow-core/src/airflow/dag_processing/dagbag.py @@ -33,7 +33,7 @@ from airflow._shared.timezones import timezone from airflow.configuration import conf from airflow.dag_processing.bundles.local import LocalDagBundle -from airflow.dag_processing.importer_routing import get_claiming_coordinator +from airflow.dag_processing.importer_routing import get_claiming_importer from airflow.exceptions import ( AirflowClusterPolicyError, AirflowClusterPolicySkipDag, @@ -580,7 +580,7 @@ def sync_bag_to_db( from airflow.dag_processing.collection import update_dag_parsing_results_in_db def is_parsed_by_runtime(rel_path: str) -> bool: - return get_claiming_coordinator(Path(dagbag.bundle_path or "", rel_path), bundle_name) is not None + return get_claiming_importer(Path(dagbag.bundle_path or "", rel_path), bundle_name) is not None import_errors = { (bundle_name, rel_path): error diff --git a/airflow-core/src/airflow/dag_processing/importer_routing.py b/airflow-core/src/airflow/dag_processing/importer_routing.py index dafbb00f2928b..f08f8253d6072 100644 --- a/airflow-core/src/airflow/dag_processing/importer_routing.py +++ b/airflow-core/src/airflow/dag_processing/importer_routing.py @@ -23,25 +23,25 @@ import os from typing import TYPE_CHECKING -from airflow.sdk.coordinators._dag_importer import find_claiming_coordinator +from airflow.sdk.coordinators._dag_importer import find_claiming_importer if TYPE_CHECKING: - from airflow.sdk.coordinators._subprocess import SubprocessCoordinator # noqa: SDK001 + from airflow.sdk.coordinators._dag_importer import CoordinatorDagImporter # noqa: SDK001 log = logging.getLogger(__name__) -def get_claiming_coordinator( +def get_claiming_importer( path: str | os.PathLike[str], bundle_name: str | None -) -> SubprocessCoordinator | None: +) -> CoordinatorDagImporter | None: """ - Return the coordinator whose runtime parses ``path``, or ``None`` when a Python child parses it. + Return the coordinator Dag importer that claims ``path``, or ``None`` when a Python child parses it. A runtime parses the file when its importer is a coordinator's Dag importer. When the bundle's importers cannot be loaded, the error is logged and a Python child parses the file. """ try: - return find_claiming_coordinator(path, bundle_name) + return find_claiming_importer(path, bundle_name) except Exception: log.exception("Cannot load the Dag importer for %s in bundle %s", path, bundle_name) return None diff --git a/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py b/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py index 79ed7e7e7cb9a..472fca885e574 100644 --- a/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py +++ b/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py @@ -36,7 +36,7 @@ from airflow import settings from airflow.dag_processing.dagbag import _get_bundle_team_name, _validate_executor_fields -from airflow.dag_processing.importer_routing import get_claiming_coordinator +from airflow.dag_processing.importer_routing import get_claiming_importer from airflow.dag_processing.processor import ( BaseDagFileProcessorProcess, DagFileParseRequest, @@ -128,10 +128,9 @@ def report_schema_version(schema_version: str | None) -> None: 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_claiming_coordinator(msg.file, msg.bundle_name) - if coordinator is None: - raise RuntimeError(f"No coordinator's Dag importer claims {msg.file}") - coordinator.parse_dag( + if (importer := get_claiming_importer(msg.file, msg.bundle_name)) is None: + raise RuntimeError(f"No coordinator Dag importer claims {msg.file}") + importer.get_parsing_coordinator().parse_dag( path=Path(msg.file), bundle_path=msg.bundle_path, comm_address=msg.comm_address, @@ -427,9 +426,8 @@ def _read_dag_source_codes(self, serialized_dags: list[LazyDeserializedDAG]) -> return {} file = self._parse_request.file try: - coordinator = get_claiming_coordinator(file, self.bundle_name) - if coordinator is None or (importer := coordinator.get_dag_importer()) is None: - raise RuntimeError(f"No coordinator's Dag importer claims {file}") + if (importer := get_claiming_importer(file, self.bundle_name)) is None: + raise RuntimeError(f"No coordinator Dag importer claims {file}") source = importer.get_source_code(FilesystemDagDefinition(Path(file))) except Exception as e: self.process_log.warning("Cannot read the Dag source", fileloc=file, error=str(e)) diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index 6526ee8637392..4d525a40febaa 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -55,7 +55,7 @@ ) from airflow.dag_processing.bundles.manager import DagBundlesManager from airflow.dag_processing.collection import update_dag_parsing_results_in_db -from airflow.dag_processing.importer_routing import get_claiming_coordinator +from airflow.dag_processing.importer_routing import get_claiming_importer from airflow.dag_processing.lang_sdk_processor import LangSDKDagFileProcessorProcess from airflow.dag_processing.processor import ( BaseDagFileProcessorProcess, @@ -810,7 +810,7 @@ def prepare_callback_bundle(self, request: CallbackRequest) -> BaseDagBundle | N def _add_callback_to_queue(self, request: CallbackRequest) -> None: self.log.debug("Queuing %s CallbackRequest: %s", type(request).__name__, request) - if get_claiming_coordinator(request.filepath, request.bundle_name) is not None: + if get_claiming_importer(request.filepath, request.bundle_name) is not None: self._log_dropped_lang_sdk_callback(request) return bundle = self.prepare_callback_bundle(request) @@ -1491,7 +1491,7 @@ def _create_process(self, dag_file: DagFileInfo) -> BaseDagFileProcessorProcess: logger, logger_filehandle = self._get_logger_for_dag_file(dag_file) subprocess_logs_to_stdout = conf.get("logging", "dag_processor_log_target") == "stdout" - if get_claiming_coordinator(dag_file.absolute_path, dag_file.bundle_name) is not None: + if get_claiming_importer(dag_file.absolute_path, dag_file.bundle_name) is not None: return LangSDKDagFileProcessorProcess.start( id=id, path=dag_file.absolute_path, diff --git a/airflow-core/tests/unit/dag_processing/fake_lang_sdk.py b/airflow-core/tests/unit/dag_processing/fake_lang_sdk.py index 9fc953e0ce310..456089df44409 100644 --- a/airflow-core/tests/unit/dag_processing/fake_lang_sdk.py +++ b/airflow-core/tests/unit/dag_processing/fake_lang_sdk.py @@ -15,7 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""A coordinator that claims ``.native`` Dag files, and a runtime a test can play in the parse child.""" +"""A coordinator and Dag importer for ``.native`` Dag files, and a runtime to play in the parse child.""" from __future__ import annotations @@ -50,7 +50,7 @@ @attrs.define(kw_only=True) class FakeCoordinator(SubprocessCoordinator): """ - Claim ``.native`` files; the file's JSON names the command that parses it. + Parse ``.native`` files; the file'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. @@ -62,11 +62,9 @@ def _build_parse_dag_command(self, *, path: Path) -> tuple[list[str], str | None raise FileNotFoundError(error) return spec.get("argv", ["/bin/false"]), spec.get("schema_version") - def get_dag_importer(self) -> FakeCoordinatorDagImporter: - return FakeCoordinatorDagImporter(coordinator=self) - class FakeCoordinatorDagImporter(CoordinatorDagImporter): + coordinator_classpath = f"{__name__}.FakeCoordinator" artifact_suffix = ".native" supported_extensions = [".native"] @@ -77,15 +75,20 @@ def get_source_code(self, definition) -> DagSourceCode: @contextlib.contextmanager def fake_coordinator(**kwargs: Any) -> Iterator[None]: """ - Configure a ``FakeCoordinator``, with fresh coordinators and registries inside and after the block. + Configure a ``FakeCoordinator`` and register its Dag importer. - The parse child is a bare fork even on macOS, so it sees the test's ``parse_dag`` patch and config. + The coordinators and the Dag importer registries are fresh inside and after the block. The parse child + is a bare fork even on macOS, so it sees the test's ``parse_dag`` patch and config. """ spec = {"fake": {"classpath": f"{__name__}.FakeCoordinator", "kwargs": kwargs}} reset_importer_registry() try: with ( conf_vars({("sdk", "coordinators"): json.dumps(spec)}), + mock.patch( + "airflow.sdk.coordinators._dag_importer.COORDINATOR_DAG_IMPORTERS", + (f"{__name__}.FakeCoordinatorDagImporter",), + ), mock.patch.object(supervisor, "_should_use_exec", return_value=False), ): yield diff --git a/airflow-core/tests/unit/dag_processing/test_importer_routing.py b/airflow-core/tests/unit/dag_processing/test_importer_routing.py index 4b554f35a2221..a2a956cc7466c 100644 --- a/airflow-core/tests/unit/dag_processing/test_importer_routing.py +++ b/airflow-core/tests/unit/dag_processing/test_importer_routing.py @@ -19,33 +19,38 @@ import json import logging +from unittest import mock -from airflow.dag_processing.importer_routing import get_claiming_coordinator +from airflow.dag_processing.importer_routing import get_claiming_importer from tests_common.test_utils.config import conf_vars -from unit.dag_processing.fake_lang_sdk import FakeCoordinator, fake_coordinator +from unit.dag_processing.fake_lang_sdk import FakeCoordinator, FakeCoordinatorDagImporter, fake_coordinator -def test_get_claiming_coordinator_returns_the_coordinator_of_its_importer(tmp_path): +def test_get_claiming_importer_returns_the_importer_that_claims_the_file(tmp_path): with fake_coordinator(): - coordinator = get_claiming_coordinator(tmp_path / "dags.native", "testing") - others = [get_claiming_coordinator(tmp_path / name, "testing") for name in ("dags.other", "dag.py")] + importer = get_claiming_importer(tmp_path / "dags.native", "testing") + others = [get_claiming_importer(tmp_path / name, "testing") for name in ("dags.other", "dag.py")] - assert isinstance(coordinator, FakeCoordinator) + assert isinstance(importer, FakeCoordinatorDagImporter) + assert importer.bundle_name == "testing" assert others == [None, None] -def test_get_claiming_coordinator_logs_a_broken_configuration(tmp_path, caplog): +@mock.patch( + "airflow.sdk.coordinators._dag_importer.COORDINATOR_DAG_IMPORTERS", ("nonexistent.module.Importer",) +) +def test_get_claiming_importer_logs_a_broken_configuration(tmp_path, caplog): spec = {"classpath": f"{FakeCoordinator.__module__}.FakeCoordinator", "kwargs": {}} with ( - conf_vars({("sdk", "coordinators"): json.dumps({"first": spec, "second": spec})}), + conf_vars({("sdk", "coordinators"): json.dumps({"fake": spec})}), caplog.at_level(logging.ERROR, logger="airflow.dag_processing.importer_routing"), ): - assert get_claiming_coordinator(tmp_path / "dags.native", "testing") is None + assert get_claiming_importer(tmp_path / "dags.native", "testing") is None - [record] = caplog.records + [record] = [r for r in caplog.records if r.name == "airflow.dag_processing.importer_routing"] assert ( record.getMessage() == f"Cannot load the Dag importer for {tmp_path / 'dags.native'} in bundle testing" ) - assert "Coordinators 'first' and 'second' both parse .native files" in str(record.exc_info[1]) + assert "Cannot build the coordinator Dag importers of Dag bundle 'testing'" in str(record.exc_info[1]) diff --git a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py index a34a8f7f37125..b2c529c961df2 100644 --- a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py @@ -19,6 +19,7 @@ import contextlib import copy +import json import os import selectors import signal @@ -58,6 +59,7 @@ from tests_common.test_utils.config import conf_vars from unit.dag_processing.fake_lang_sdk import ( FakeCoordinator, + FakeCoordinatorDagImporter, fake_coordinator, play_runtime, write_native_file, @@ -411,10 +413,21 @@ def test_a_failing_import_timeout_policy_is_an_import_error(self, parse, policy, @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. + # The spawned interpreter finds the coordinator and its Dag importer 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")) + monkeypatch.setenv( + "AIRFLOW__DAG_PROCESSOR__DAG_IMPORTER_CONFIGS", + json.dumps( + [ + { + "classpath": f"{FakeCoordinatorDagImporter.__module__}.FakeCoordinatorDagImporter", + "kwargs": {"bundle_name": "testing"}, + } + ] + ), + ) with selectors.DefaultSelector() as selector: proc = _start(tmp_path, selector, argv=["/bin/sh", "-c", "exec sleep 30"]) deadline = time.monotonic() + 30 diff --git a/generated/known_sdk_imports_in_core.txt b/generated/known_sdk_imports_in_core.txt index 51ae1d58804aa..e8c6aabc037f0 100644 --- a/generated/known_sdk_imports_in_core.txt +++ b/generated/known_sdk_imports_in_core.txt @@ -4,7 +4,7 @@ airflow-core/src/airflow/cli/commands/task_command.py::7 airflow-core/src/airflow/cli/commands/triggerer_command.py::1 airflow-core/src/airflow/configuration.py::1 airflow-core/src/airflow/dag_processing/dagbag.py::3 -airflow-core/src/airflow/dag_processing/importer_routing.py::2 +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 diff --git a/task-sdk/src/airflow/sdk/coordinators/_dag_importer.py b/task-sdk/src/airflow/sdk/coordinators/_dag_importer.py index 3cfe6ff8fc8e4..54517ee38310d 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_dag_importer.py +++ b/task-sdk/src/airflow/sdk/coordinators/_dag_importer.py @@ -160,15 +160,3 @@ def find_claiming_importer( file = Path(path) importer = registry.get_importer(file) return importer if isinstance(importer, CoordinatorDagImporter) and importer.can_handle(file) else None - - -def find_claiming_coordinator( - path: str | os.PathLike[str], bundle_name: str | None -) -> SubprocessCoordinator | None: - """ - Return the coordinator whose runtime parses *path*, or ``None`` when a Python Dag importer parses it. - - :raises InvalidCoordinatorError: when the file's importer cannot find its coordinator. - """ - importer = find_claiming_importer(path, bundle_name) - return importer.get_parsing_coordinator() if importer is not None else None From 302b4e0c7b53af2a5d4fe788571c17796a9b4696 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Sat, 3 Oct 2026 15:53:23 +0000 Subject: [PATCH 22/24] Test the import errors of a bundle with several coordinators of a class A Dag file that no single coordinator can parse stays claimed by its importer, and its parse fails with an import error that says why: no entry in dag_bundle_to_coordinator, an entry for another class or one that cannot be loaded, or an option that is not a JSON object of coordinator keys. --- .../unit/dag_processing/fake_lang_sdk.py | 27 ++++++-- .../dag_processing/test_importer_routing.py | 7 +++ .../dag_processing/test_lang_sdk_processor.py | 61 +++++++++++++++++++ 3 files changed, 89 insertions(+), 6 deletions(-) diff --git a/airflow-core/tests/unit/dag_processing/fake_lang_sdk.py b/airflow-core/tests/unit/dag_processing/fake_lang_sdk.py index 456089df44409..92c8e596972f6 100644 --- a/airflow-core/tests/unit/dag_processing/fake_lang_sdk.py +++ b/airflow-core/tests/unit/dag_processing/fake_lang_sdk.py @@ -73,18 +73,33 @@ def get_source_code(self, definition) -> DagSourceCode: @contextlib.contextmanager -def fake_coordinator(**kwargs: Any) -> Iterator[None]: +def fake_coordinator( + *keys: str, + dag_bundle_to_coordinator: dict[str, str] | str | None = None, + other_coordinators: dict[str, str] | None = None, + **kwargs: Any, +) -> Iterator[None]: """ - Configure a ``FakeCoordinator`` and register its Dag importer. + Configure a ``FakeCoordinator`` for each of *keys* (``fake`` by default), and register its Dag importer. - The coordinators and the Dag importer registries are fresh inside and after the block. The parse child - is a bare fork even on macOS, so it sees the test's ``parse_dag`` patch and config. + *dag_bundle_to_coordinator* is the ``[sdk] dag_bundle_to_coordinator`` option, as a dict or as the raw + text. *other_coordinators* maps keys to the classpaths of coordinators of another class. The + coordinators and the Dag importer registries are fresh inside and after the block. The parse child is + a bare fork even on macOS, so it sees the test's ``parse_dag`` patch and config. """ - spec = {"fake": {"classpath": f"{__name__}.FakeCoordinator", "kwargs": kwargs}} + spec = {key: {"classpath": f"{__name__}.FakeCoordinator", "kwargs": kwargs} for key in keys or ("fake",)} + spec.update({key: {"classpath": classpath} for key, classpath in (other_coordinators or {}).items()}) + config = {("sdk", "coordinators"): json.dumps(spec)} + if dag_bundle_to_coordinator is not None: + config[("sdk", "dag_bundle_to_coordinator")] = ( + dag_bundle_to_coordinator + if isinstance(dag_bundle_to_coordinator, str) + else json.dumps(dag_bundle_to_coordinator) + ) reset_importer_registry() try: with ( - conf_vars({("sdk", "coordinators"): json.dumps(spec)}), + conf_vars(config), mock.patch( "airflow.sdk.coordinators._dag_importer.COORDINATOR_DAG_IMPORTERS", (f"{__name__}.FakeCoordinatorDagImporter",), diff --git a/airflow-core/tests/unit/dag_processing/test_importer_routing.py b/airflow-core/tests/unit/dag_processing/test_importer_routing.py index a2a956cc7466c..25143f82fb598 100644 --- a/airflow-core/tests/unit/dag_processing/test_importer_routing.py +++ b/airflow-core/tests/unit/dag_processing/test_importer_routing.py @@ -37,6 +37,13 @@ def test_get_claiming_importer_returns_the_importer_that_claims_the_file(tmp_pat assert others == [None, None] +def test_get_claiming_importer_claims_a_file_that_several_coordinators_could_parse(tmp_path): + with fake_coordinator("first", "second"): + importer = get_claiming_importer(tmp_path / "dags.native", "testing") + + assert isinstance(importer, FakeCoordinatorDagImporter) + + @mock.patch( "airflow.sdk.coordinators._dag_importer.COORDINATOR_DAG_IMPORTERS", ("nonexistent.module.Importer",) ) diff --git a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py index b2c529c961df2..9f9bc7e9c1ce4 100644 --- a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py @@ -321,6 +321,67 @@ def test_a_failed_parse_is_an_import_error(self, parse, spec, reply, error): assert proc.parsing_result.serialized_dags == [] assert proc.parsing_result.import_errors == {"dag.native": error} + @pytest.mark.parametrize( + ("mapping", "error"), + [ + pytest.param( + None, + "Dag bundle 'testing' has 2 FakeCoordinator coordinators (first, second). " + "Map the bundle to one of them in [sdk] dag_bundle_to_coordinator.", + id="no-entry", + ), + pytest.param( + {"other-bundle": "first"}, + "Dag bundle 'testing' has 2 FakeCoordinator coordinators (first, second). " + "Map the bundle to one of them in [sdk] dag_bundle_to_coordinator.", + id="entry-for-another-bundle", + ), + pytest.param( + {"testing": "other"}, + "Dag bundle 'testing' has 2 FakeCoordinator coordinators (first, second). " + "[sdk] dag_bundle_to_coordinator maps it to 'other', a coordinator of another class. " + "Move these files to another Dag bundle, or keep one FakeCoordinator.", + id="entry-of-another-class", + ), + pytest.param( + {"testing": "missing"}, + "Dag bundle 'testing' has 2 FakeCoordinator coordinators (first, second). " + "[sdk] dag_bundle_to_coordinator maps it to 'missing', which cannot be loaded.", + id="entry-that-cannot-be-loaded", + ), + pytest.param( + "{not json", + "Unable to parse [sdk] 'dag_bundle_to_coordinator' as valid json", + id="not-json", + ), + pytest.param( + '["first"]', + "[sdk] dag_bundle_to_coordinator must be a JSON object that maps Dag bundle names to " + "coordinator keys", + id="not-an-object", + ), + pytest.param( + '{"testing": 1}', + "[sdk] dag_bundle_to_coordinator must be a JSON object that maps Dag bundle names to " + "coordinator keys", + id="not-keys", + ), + ], + ) + def test_a_file_that_several_coordinators_could_parse_is_an_import_error(self, parse, mapping, error): + with fake_coordinator( + "first", + "second", + dag_bundle_to_coordinator=mapping, + other_coordinators={"other": "airflow.sdk.execution_time.coordinator.BaseCoordinator"}, + ): + proc = parse() + + assert proc.parsing_result.serialized_dags == [] + assert proc.parsing_result.import_errors == { + "dag.native": f"Cannot start the Lang-SDK runtime: InvalidCoordinatorError: {error}" + } + @patch.object(FakeCoordinator, "parse_dag", autospec=True) def test_a_message_that_does_not_validate_is_an_import_error(self, mock_parse_dag, parse): def reply(request, comms): From 9007b91ed18c686246f8314af1126049b77f3680 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Sat, 3 Oct 2026 15:53:39 +0000 Subject: [PATCH 23/24] Apply the team rules to the Dag that validation returns validate_serialized_dag returns the Dag it deserialized, so the Lang-SDK processor uses it for the executor and pool checks instead of copying and deserializing the serialized Dag a second time. --- .../airflow/dag_processing/lang_sdk_processor.py | 13 ++++++------- .../unit/dag_processing/test_lang_sdk_processor.py | 14 ++++++++++++++ 2 files changed, 20 insertions(+), 7 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py b/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py index 472fca885e574..f79fa31e09d56 100644 --- a/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py +++ b/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py @@ -18,7 +18,6 @@ from __future__ import annotations -import copy import functools import os import selectors @@ -65,6 +64,7 @@ from structlog.typing import FilteringBoundLogger from airflow.sdk.execution_time.supervisor import RequestHandler, RequestResult + from airflow.serialization.definitions.dag import SerializedDAG from airflow.typing_compat import Self # How long a runtime may keep running after its parse result, as Node does while a handle stays open. @@ -372,8 +372,8 @@ def _handle_parsing_result( for dag in msg.serialized_dags: DagSerialization.fill_config_defaults(dag.data) try: - DagSerialization.validate_serialized_dag(dag.data) - self._apply_team_rules(dag.data) + deserialized = DagSerialization.validate_serialized_dag(dag.data) + self._apply_team_rules(dag.data, deserialized) except DeserializationError as e: message = f"Cannot load the serialized Dag: {e}" except Exception as e: @@ -394,16 +394,15 @@ def _handle_parsing_result( self._parsing_result_monotonic = time.monotonic() return None, {} - def _apply_team_rules(self, data: dict) -> None: + def _apply_team_rules(self, data: dict, dag: SerializedDAG) -> None: """ Check each task's executor and move tasks in the default pool to the team's, as the Dag bag does. - The bundle's team owns the Dag. *data* is a serialized Dag that validates, and its pools are changed - in place. + The bundle's team owns the Dag. *data* is a serialized Dag that validates, *dag* is that Dag + deserialized, and *data*'s pools are changed in place. :raises UnknownExecutorException: if a task's executor is not available to the team or globally. """ - dag = DagSerialization.from_dict(copy.deepcopy(data)) _validate_executor_fields(dag, self.bundle_name) if not (team_name := _get_bundle_team_name(self.bundle_name)): return diff --git a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py index 9f9bc7e9c1ce4..d9c851d59f103 100644 --- a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py @@ -544,6 +544,20 @@ def test_a_dag_source_that_cannot_be_read_is_a_placeholder(): assert source.source_code.startswith("Cannot read the source of dag.native: [Errno 2] No such file") +@patch.object(DagSerialization, "from_dict", autospec=True, side_effect=DagSerialization.from_dict) +def test_a_dag_is_deserialized_once_to_validate_it_and_apply_the_team_rules(mock_from_dict): + proc = _make_process() + + proc._handle_request( + DagFileParsingResult(fileloc="/b/dag.native", serialized_dags=[_serialize_dag("native_dag")]), + MagicMock(), + 1, + ) + + assert proc.parsing_result.import_errors is None + mock_from_dict.assert_called_once() + + @patch.object(LangSDKDagFileProcessorProcess, "send_msg", autospec=True) def test_the_first_parse_result_wins(mock_send_msg): proc = _make_process() From cb9d1e9ee3d8775e68a675bb6a0525117e47f5a2 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Mon, 5 Oct 2026 10:33:32 +0000 Subject: [PATCH 24/24] Catch a runtime's dead connection when its parse request is sent `_register_comm` can run after the runtime has already died, so its `send_msg` call needs the same BrokenPipeError/ConnectionResetError guard `_service_processor_sockets` uses elsewhere. Co-Authored-By: Claude --- .../airflow/dag_processing/lang_sdk_processor.py | 8 +++++++- .../unit/dag_processing/test_lang_sdk_processor.py | 13 +++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py b/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py index f79fa31e09d56..1f87abe25f1c7 100644 --- a/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py +++ b/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py @@ -296,7 +296,13 @@ def read_valid_frame(sock: socket) -> bool: # 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) + try: + self.send_msg(self._parse_request, request_id=0) + except (BrokenPipeError, ConnectionResetError): + # The runtime can die after connecting but before this send; `_check_subprocess_exit` + # reports why on the next `is_ready` poll, once the process has actually exited. + self._on_socket_closed(conn) + conn.close() def _handle_valid_requests(self) -> Generator[None, _RequestFrame, None]: """ diff --git a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py index d9c851d59f103..483ad20a9a2d8 100644 --- a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py @@ -612,6 +612,19 @@ def test_an_invalid_message_after_the_parse_result_keeps_it(mock_send_msg, mock_ mock_kill_runtime.assert_called_once_with(proc) +@patch.object(LangSDKDagFileProcessorProcess, "send_msg", autospec=True) +def test_a_runtime_that_dies_before_its_parse_request_is_sent_is_not_fatal(mock_send_msg): + mock_send_msg.side_effect = BrokenPipeError + proc = _make_process() + runtime, conn = socket.socketpair() + with runtime: + proc._register_comm(conn) + + proc.selector.unregister.assert_called_once_with(conn) + assert conn not in proc._open_sockets + assert conn.fileno() == -1 + + @patch.object(LangSDKDagFileProcessorProcess, "send_msg", autospec=True) def test_the_schema_version_is_reported_once(mock_send_msg): proc = _make_process()