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. diff --git a/airflow-core/src/airflow/dag_processing/dagbag.py b/airflow-core/src/airflow/dag_processing/dagbag.py index 2ec6ee4afe1d5..b5822171de88d 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_importer from airflow.exceptions import ( AirflowClusterPolicyError, AirflowClusterPolicySkipDag, @@ -57,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): @@ -103,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: @@ -150,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: @@ -583,16 +571,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_importer(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..f08f8253d6072 --- /dev/null +++ b/airflow-core/src/airflow/dag_processing/importer_routing.py @@ -0,0 +1,47 @@ +# +# 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 typing import TYPE_CHECKING + +from airflow.sdk.coordinators._dag_importer import find_claiming_importer + +if TYPE_CHECKING: + from airflow.sdk.coordinators._dag_importer import CoordinatorDagImporter # noqa: SDK001 + +log = logging.getLogger(__name__) + + +def get_claiming_importer( + path: str | os.PathLike[str], bundle_name: str | None +) -> CoordinatorDagImporter | None: + """ + 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_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 new file mode 100644 index 0000000000000..1f87abe25f1c7 --- /dev/null +++ b/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py @@ -0,0 +1,528 @@ +# 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 contextlib import suppress +from pathlib import Path +from socket import socket +from typing import TYPE_CHECKING, Annotated, ClassVar, Literal, get_args + +import attrs +import msgspec +import psutil +from pydantic import BaseModel, Field, TypeAdapter + +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_importer +from airflow.dag_processing.processor import ( + BaseDagFileProcessorProcess, + DagFileParseRequest, + DagFileParsingResult, + 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, _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.enums import Encoding +from airflow.serialization.serialized_objects import DagSerialization, LazyDeserializedDAG + +if TYPE_CHECKING: + from collections.abc import Generator + + 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. +_EXIT_GRACE_PERIOD = 5.0 + +_IMPORT_TIMEOUT_SETTING = "[core] dagbag_import_timeout or the get_dagbag_import_timeout policy" + + +# 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 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" + # 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, 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) + 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, + 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. + """ + + 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) + _import_timeout: float | None = attrs.field(default=None, init=False) + _schema_version_reported: bool = attrs.field(default=False, init=False) + _parsing_result_monotonic: float | None = attrs.field(default=None, init=False) + _unverified_connections: list[tuple[socket, _Channel]] = attrs.field(factory=list, init=False) + _group_killed: bool = attrs.field(default=False, init=False) + + @classmethod + def start( # type: ignore[override] + cls, + *, + 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 + + 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 + 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]: + """ + 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._import_timeout = msg.import_timeout + 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: + 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: + message = f"{type(e).__name__}: {e}" + else: + serialized_dags.append(dag) + continue + 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, + "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 _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, *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. + """ + _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. + + 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: + 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)) + 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)]), + } + + @property + def is_ready(self) -> bool: + self._verify_connections() + if ( + self._parsing_result_monotonic is not None + and time.monotonic() - self._parsing_result_monotonic > _EXIT_GRACE_PERIOD + ): + if self._exit_code is None: + self.process_log.warning( + "The Lang-SDK runtime did not exit after its parse result; killing it" + ) + elif not self._group_killed: + self.process_log.warning( + "The Lang-SDK runtime left processes holding its output after its parse result; " + "killing them" + ) + self._kill_runtime() + if ( + self._import_timeout is not None + and self.parsing_result is None + and self._exit_code is None + and time.monotonic() - self.start_time > self._import_timeout + ): + self._time_out(self._import_timeout) + if self._check_subprocess_exit() is None: + return False + self._close_listeners() + if not super().is_ready: + return False + if self.parsing_result is None: + self._set_import_error( + f"The Lang-SDK runtime exited with code {self._exit_code} without a parse result" + ) + return True + + def _time_out(self, timeout: float) -> 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 {_IMPORT_TIMEOUT_SETTING}" + ) + self._kill_runtime() + + def _kill_runtime(self) -> None: + """ + Kill the runtime's process group and wait for the runtime, without servicing its sockets. + + Their handler may have failed. The wait is bounded so that a runtime stuck in the kernel does not + stall the caller's loop; ``is_ready`` sees it exit later. + """ + if self._exit_code is not None: + self._kill_leftovers() + return + try: + self._signal_subprocess(signal.SIGKILL) + self._group_killed = True + self._exit_code = self._process.wait(timeout=_EXIT_GRACE_PERIOD) + except (self._process.ProcessNotFound, ProcessLookupError): + self._exit_code = -1 + except self._process.TimeoutExpired: + self.process_log.warning("The Lang-SDK runtime did not exit after SIGKILL", pid=self.pid) + + def _kill_leftovers(self) -> None: + """ + Kill the processes an exited runtime left in its process group, which can keep its sockets open. + + The group is killed once, and not when another process has reused the runtime's pid, since the + group could then be that process's. + """ + if self._exit_code is None or self._group_killed or not self._new_process_group: + return + self._group_killed = True + if psutil.pid_exists(self.pid): + return + with suppress(ProcessLookupError, PermissionError): + os.killpg(self.pid, signal.SIGKILL) + + def close(self) -> None: + # A listener has nothing to drain, and cleanup would call its accept handler forever. + self._close_listeners() + self._kill_leftovers() + super().close() diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index 9827cc590829c..4d525a40febaa 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, @@ -54,7 +55,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_importer +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 +269,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) @@ -803,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_importer(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 @@ -818,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: """ @@ -1277,7 +1300,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,11 +1484,26 @@ 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) + subprocess_logs_to_stdout = conf.get("logging", "dag_processor_log_target") == "stdout" + + 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, + 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, @@ -1477,7 +1515,7 @@ def _create_process(self, dag_file: DagFileInfo) -> DagFileProcessorProcess: 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, ) 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..92c8e596972f6 --- /dev/null +++ b/airflow-core/tests/unit/dag_processing/fake_lang_sdk.py @@ -0,0 +1,144 @@ +# +# 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 and Dag importer for ``.native`` Dag files, and a runtime to 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): + """ + 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. + """ + + 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") + + +class FakeCoordinatorDagImporter(CoordinatorDagImporter): + coordinator_classpath = f"{__name__}.FakeCoordinator" + artifact_suffix = ".native" + supported_extensions = [".native"] + + def get_source_code(self, definition) -> DagSourceCode: + return DagSourceCode(definition.read_text(), "fake") + + +@contextlib.contextmanager +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`` for each of *keys* (``fake`` by default), and register its Dag importer. + + *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 = {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(config), + mock.patch( + "airflow.sdk.coordinators._dag_importer.COORDINATOR_DAG_IMPORTERS", + (f"{__name__}.FakeCoordinatorDagImporter",), + ), + 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..00eafe90cf09b 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,28 @@ 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 + + +@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() + + 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..25143f82fb598 --- /dev/null +++ b/airflow-core/tests/unit/dag_processing/test_importer_routing.py @@ -0,0 +1,63 @@ +# +# 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 json +import logging +from unittest import mock + +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, FakeCoordinatorDagImporter, fake_coordinator + + +def test_get_claiming_importer_returns_the_importer_that_claims_the_file(tmp_path): + with fake_coordinator(): + 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(importer, FakeCoordinatorDagImporter) + assert importer.bundle_name == "testing" + 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",) +) +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({"fake": spec})}), + caplog.at_level(logging.ERROR, logger="airflow.dag_processing.importer_routing"), + ): + assert get_claiming_importer(tmp_path / "dags.native", "testing") is None + + [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 "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 new file mode 100644 index 0000000000000..483ad20a9a2d8 --- /dev/null +++ b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py @@ -0,0 +1,728 @@ +# +# 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 copy +import json +import os +import selectors +import signal +import socket +import subprocess +import sys +import time +import uuid +from pathlib import Path +from typing import BinaryIO +from unittest.mock import ANY, MagicMock, call, patch + +import psutil +import pytest +import structlog +from structlog.typing import FilteringBoundLogger + +from airflow.configuration import conf +from airflow.dag_processing.lang_sdk_processor import ( + _EXIT_GRACE_PERIOD, + LangSDKDagFileProcessorProcess, + LangSDKRuntimeSchemaVersion, + _get_import_timeout, +) +from airflow.dag_processing.processor import DagFileParseRequest, DagFileParsingResult +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.execution_time import supervisor +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 + +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, +) + +# 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() + + +def _is_running(pid: int) -> bool: + try: + # Init may not have reaped the killed process yet. + return psutil.Process(pid).status() != psutil.STATUS_ZOMBIE + except psutil.NoSuchProcess: + return False + + +def _stops_running(pid: int, timeout: float = 10) -> bool: + # SIGKILL is delivered asynchronously, so the process may still run right after the kill. + deadline = time.monotonic() + timeout + while _is_running(pid): + if time.monotonic() > deadline: + return False + time.sleep(0.05) + return True + + +@pytest.fixture(autouse=True) +def _coordinator(): + with fake_coordinator(): + yield + + +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(), + logger_filehandle=MagicMock(spec=BinaryIO), + 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 at $.dag.max_active_runs: '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} + + @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): + 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 + + @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 _stops_running(int((tmp_path / "leftover.pid").read_text())) + + @pytest.mark.parametrize( + ("policy", "error"), + [ + pytest.param( + {"side_effect": RuntimeError("policy bug")}, "RuntimeError: policy bug", id="raises" + ), + pytest.param( + {"return_value": "30"}, + "TypeError: Value (30) from get_dagbag_import_timeout must be int or float", + id="not-a-number", + ), + ], + ) + def test_a_failing_import_timeout_policy_is_an_import_error(self, parse, policy, error): + with patch("airflow.settings.get_dagbag_import_timeout", autospec=True, **policy): + proc = parse() + + assert proc.parsing_result.import_errors == { + "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): + if use_exec: + # 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 + while psutil.Process(proc.pid).name() != "sleep": + assert time.monotonic() < deadline, "the runtime did not start" + proc._service_subprocess(max_wait_time=0.1) + fd_dir = Path(f"/proc/{proc.pid}/fd") + fds = {fd.name: os.readlink(fd) for fd in fd_dir.iterdir()} + proc.kill(signal.SIGKILL) + proc.close() + + assert sorted(fds) == ["0", "1", "2"] + assert fds["0"] == "/dev/null" + + +@pytest.mark.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: + kwargs.setdefault("process", MagicMock(spec=PsutilTracker)) + kwargs.setdefault("logger_filehandle", MagicMock(spec=BinaryIO)) + return LangSDKDagFileProcessorProcess( + id=uuid.uuid4(), + pid=1, + 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={}, + 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(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() + 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_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() + + 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" + + +@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 + 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) +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) + + +@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", + } diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index fb15c2486c632..a5531ec57b888 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -46,11 +46,13 @@ 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 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 +73,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 +92,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 @@ -309,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): @@ -1661,6 +1681,221 @@ 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() + + @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(), + caplog.at_level(logging.WARNING, logger="airflow.dag_processing.manager"), + ): + manager = DagFileProcessorManager(max_runs=1) + manager._add_callback_to_queue(request_) + + 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( + 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 + ): + 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, + 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 + ): + 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() == {} + + @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, + 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 + ): + 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, " + "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) + @mock.patch.object( + DagFileProcessorManager, + "_find_files_in_bundle", + autospec=True, + 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 + ): + 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() diff --git a/generated/known_sdk_imports_in_core.txt b/generated/known_sdk_imports_in_core.txt index 8c799031129b9..e8c6aabc037f0 100644 --- a/generated/known_sdk_imports_in_core.txt +++ b/generated/known_sdk_imports_in_core.txt @@ -4,6 +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::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 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/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py index ed622023308a2..27254fcb8f5e3 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,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 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 ( @@ -161,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 @@ -1018,6 +1020,37 @@ def _register_deserialization_allowed_classes(dag, log: Logger) -> 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: + return find_claiming_importer(path, bundle_name) + except Exception: + # Building the Dag bag reports a broken importer configuration. + return None + + +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 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, + 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") def parse(what: StartupDetails, log: Logger) -> RuntimeTaskInstance: # TODO: Task-SDK: @@ -1030,6 +1063,9 @@ 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 (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( dag_folder=dag_absolute_path, 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")): 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..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 @@ -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, + _find_native_dag_importer, _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 from airflow.sdk.serde import deserialize from airflow.triggers.base import BaseEventTrigger, BaseTrigger, TriggerEvent from airflow.triggers.callback import CallbackTrigger @@ -902,6 +907,131 @@ def test_parse_module_in_bundle_root(tmp_path: Path, make_ti_context): assert ti.task.dag.dag_id == "dag_name" +class NativeDagImporter(CoordinatorDagImporter): + coordinator_classpath = f"{__name__}.NativeCoordinator" + 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 of the runtime whose Dag importer claims ``.native`` files in every bundle.""" + + +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("{}") + bundle_config = [ + { + "name": "my-bundle", + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {"path": str(tmp_path), "refresh_interval": 1}, + } + ] + 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( + 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() + + 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( + "This task belongs to a native Lang-SDK Dag and cannot run in Python. Give the Dag's tasks " + "their own queue and route it to a %s in [sdk] queue_to_coordinator", + "NativeCoordinator", + dag_id="native_dag", + task_id="a", + queue="default", + path="dag.native", + ) + + +@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()) + + +@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", True), ("dag.py", False), ("dag.pyc", False), ("dags.zip", False)], +) +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 + + +@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]) 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."""