From d8c78379e77219de9ca893abf18a8b5c2813ef82 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 29 Jul 2026 09:51:42 +0000 Subject: [PATCH 01/11] Add dag_bundle_name option to subprocess coordinators Language-SDK coordinators previously required a hand-managed filesystem root (jars_root/executables_root/bundles_root) for compiled artifacts, kept outside Airflow's Dag-bundle machinery and prone to drifting from the Dags they serve. Allowing artifacts to be located through a configured Dag bundle -- or the task's own bundle when none is named -- lets them be delivered and versioned the same way as everything else. --- .../airflow/dag_processing/bundles/manager.py | 5 + .../bundles/test_dag_bundle_manager.py | 9 + .../airflow/sdk/coordinators/_subprocess.py | 171 +++++++++++++++--- .../coordinators/executable/coordinator.py | 14 +- .../sdk/coordinators/java/coordinator.py | 16 +- .../sdk/coordinators/node/coordinator.py | 14 +- .../airflow/sdk/execution_time/task_runner.py | 28 ++- .../executable/test_coordinator.py | 49 ++++- .../coordinators/java/test_coordinator.py | 25 +++ .../coordinators/node/test_coordinator.py | 47 ++++- .../task_sdk/coordinators/test_subprocess.py | 132 +++++++++++++- .../execution_time/test_task_runner.py | 24 +++ 12 files changed, 475 insertions(+), 59 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/bundles/manager.py b/airflow-core/src/airflow/dag_processing/bundles/manager.py index 2c20f923739bd..0245143a0cdda 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/manager.py +++ b/airflow-core/src/airflow/dag_processing/bundles/manager.py @@ -673,6 +673,11 @@ def get_bundle( name=name, version=version, version_data=version_data, **cfg_bundle.kwargs ) + @classmethod + def is_bundle_configured(cls, name: str) -> bool: + """Return whether *name* is a configured Dag bundle, without constructing the bundle.""" + return name in cls()._bundle_config + def get_all_dag_bundles(self) -> Iterable[BaseDagBundle]: """ Get all DAG bundles. diff --git a/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py b/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py index 7d23af9afefc0..f1ce2a92ca237 100644 --- a/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py +++ b/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py @@ -205,6 +205,15 @@ def test_get_bundle(): assert bundle.version is None +def test_is_bundle_configured(): + """is_bundle_configured reports membership without constructing the bundle.""" + with patch.dict( + os.environ, {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(BASIC_BUNDLE_CONFIG)} + ): + assert DagBundlesManager.is_bundle_configured("my-test-bundle") is True + assert DagBundlesManager.is_bundle_configured("bundle-that-doesn't-exist") is False + + @pytest.fixture def clear_db(): clear_db_dag_bundles() diff --git a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py index a0c3f518fb068..f157c82cfef98 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py +++ b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py @@ -27,6 +27,8 @@ from __future__ import annotations +import contextlib +import enum import ipaddress import itertools import os @@ -46,6 +48,7 @@ from airflow.sdk.execution_time.supervisor import ActivitySubprocess, NeverRaised, ProcessTracker if TYPE_CHECKING: + import pathlib from collections.abc import Sequence from structlog.typing import FilteringBoundLogger @@ -371,6 +374,17 @@ def wait(self) -> int: return code +class _ArtifactSource(enum.Enum): + """How a subprocess coordinator locates the compiled task artifacts.""" + + EXPLICIT_ROOT = enum.auto() + """An explicit filesystem root (``jars_root`` / ``executables_root`` / ``bundles_root``).""" + NAMED_BUNDLE = enum.auto() + """``dag_bundle_name`` names a configured Dag bundle; its latest version is used.""" + TASK_BUNDLE = enum.auto() + """Neither is set: the task's own (co-located) bundle, pinned to the run's version.""" + + @attrs.define(kw_only=True) class SubprocessCoordinator(BaseCoordinator): """ @@ -385,23 +399,134 @@ class SubprocessCoordinator(BaseCoordinator): :param task_startup_timeout: Maximum time the coordinator waits for the subprocess to connect to both servers, in seconds. The default is 10 seconds. + :param dag_bundle_name: Locate artifacts through a configured Dag bundle rather + than an explicit root. Mutually exclusive with the subclass's explicit root; + if neither is set, the task's own bundle is used. A named bundle resolves to + its latest version; the task's own bundle is pinned to the run's version. """ task_startup_timeout: float = 10.0 + dag_bundle_name: str | None = None + + # Classified once at construction by :meth:`_classify_artifact_source` and + # dispatched on by :meth:`_init_root_source` at execute time. + _artifact_source: _ArtifactSource | None = attrs.field(init=False, default=None) + # The subclass's explicit root, recorded at construction so the base can + # resolve roots without knowing the subclass field name. + _configured_roots: list[pathlib.Path] = attrs.field(init=False, factory=list) + # The task's own bundle, bound for the duration of a single :meth:`execute_task` + # call by :meth:`_set_current_bundle` so :meth:`_init_root_source` can resolve + # co-located artifacts. + _active_bundle_info: BundleInfo | None = attrs.field(init=False, default=None) + + def _classify_artifact_source(self, configured: Sequence[pathlib.Path], *, root_kwarg: str) -> None: + """ + Classify and validate how this coordinator locates artifacts (construction time). + + Subclasses call this from ``__attrs_post_init__`` with their own root + field. It rejects setting both an explicit root and ``dag_bundle_name``, + fails fast when ``dag_bundle_name`` names a bundle that is not configured, + and records the resulting :class:`_ArtifactSource` and explicit root. + """ + if configured and self.dag_bundle_name is not None: + raise ValueError( + f"Set at most one of {root_kwarg!r} or 'dag_bundle_name': {root_kwarg!r} for an " + f"explicit path, 'dag_bundle_name' for a configured Dag bundle, or leave both " + f"unset to scan the task's own bundle." + ) + if configured: + source = _ArtifactSource.EXPLICIT_ROOT + self._configured_roots = list(configured) + elif self.dag_bundle_name is not None: + source = _ArtifactSource.NAMED_BUNDLE + from airflow.dag_processing.bundles.manager import DagBundlesManager # noqa: SDK002 + + if not DagBundlesManager.is_bundle_configured(self.dag_bundle_name): + raise ValueError( + f"Coordinator 'dag_bundle_name' references unconfigured Dag bundle " + f"{self.dag_bundle_name!r}." + ) + else: + source = _ArtifactSource.TASK_BUNDLE + + self._artifact_source = source + details: dict[str, str | list[str]] = {"mode": source.name} + if self.dag_bundle_name is not None: + details["dag_bundle_name"] = self.dag_bundle_name + if self._configured_roots: + details["configured_roots"] = [str(root) for root in self._configured_roots] + log.debug("Coordinator artifact source selected", **details) + + def _init_root_source(self) -> list[pathlib.Path]: + """ + Resolve the directories to scan for artifacts for the current task. + + Dispatches on the :class:`_ArtifactSource` classified at construction. + An explicit root is returned as-is (no Dag bundle is resolved); otherwise + the root is a Dag bundle's materialized path — the named + ``dag_bundle_name`` bundle at its latest version, or the task's own + bundle pinned to the run's version. Called by :meth:`execute_task`, which + forwards the result to :meth:`_build_execute_task_command`. + """ + if self._artifact_source is _ArtifactSource.EXPLICIT_ROOT: + return self._configured_roots + + if self._artifact_source is _ArtifactSource.NAMED_BUNDLE: + from airflow.sdk.api.datamodels._generated import BundleInfo + + # NAMED_BUNDLE implies dag_bundle_name is set. + target = BundleInfo(name=cast("str", self.dag_bundle_name)) + elif self._artifact_source is _ArtifactSource.TASK_BUNDLE: + if self._active_bundle_info is None: + raise RuntimeError("_init_root_source requires an active task; call it during execute_task.") + target = self._active_bundle_info + else: + raise RuntimeError( + "Coordinator artifact source was not classified; call _classify_artifact_source first." + ) + + # Lazy import: task_runner is a heavy module and importing it at module + # load would risk an import cycle through the supervisor. + from airflow.sdk.execution_time.task_runner import initialize_ti_bundle - def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], str | None]: + bundle = initialize_ti_bundle(target, log) + path = bundle.path + if not path.exists(): + raise FileNotFoundError(f"Dag bundle {target.name!r} resolved to {path}, which does not exist.") + return [path] + + def _build_execute_task_command( + self, *, what: TaskInstance, roots: list[pathlib.Path] + ) -> tuple[list[str], str | None]: """ Build the subprocess command and resolve its supervisor wire-schema version for *what*. - Returns a ``(command, subprocess_schema_version)`` pair. *command* - MUST NOT include the ``--comm`` / ``--logs`` flags — those are - appended by :class:`_PopenActivitySubprocess` once the listening - sockets have been bound. A ``None`` schema version disables schema - migration; messages are then exchanged at the runtime's native wire - format. + *roots* are the directories to scan for artifacts, already resolved by + :meth:`_init_root_source` from the coordinator's configured source. + Returns a ``(command, subprocess_schema_version)`` pair. *command* MUST + NOT include the ``--comm`` / ``--logs`` flags — those are appended by + :class:`_PopenActivitySubprocess` once the listening sockets have been + bound. A ``None`` schema version disables schema migration; messages are + then exchanged at the runtime's native wire format. """ raise NotImplementedError + @contextlib.contextmanager + def _set_current_bundle(self, bundle_info: BundleInfo): + """ + Bind *bundle_info* as the active task for the duration of the block, clearing it on exit. + + Rejects a second concurrent bind: this coordinator runs one blocking task + per process and is not re-entrant. + """ + if self._active_bundle_info is not None: + raise RuntimeError("SubprocessCoordinator.execute_task is not re-entrant.") + self._active_bundle_info = bundle_info + try: + yield + finally: + self._active_bundle_info = None + def execute_task( self, *, @@ -414,18 +539,20 @@ def execute_task( subprocess_logs_to_stdout: bool, **kwargs, ) -> BaseCoordinator.ExecutionResult: - command, subprocess_schema_version = self._build_execute_task_command(what=what) - process = _PopenActivitySubprocess.start( - what=what, - dag_rel_path=dag_rel_path, - bundle_info=bundle_info, - client=client, - logger=logger, - subprocess_logs_to_stdout=subprocess_logs_to_stdout, - sentry_integration=sentry_integration, - command=command, - subprocess_schema_version=subprocess_schema_version, - startup_timeout=self.task_startup_timeout, - ) - exit_code = process.wait() - return self.ExecutionResult(exit_code, process.final_state) + with self._set_current_bundle(bundle_info): + roots = self._init_root_source() + command, subprocess_schema_version = self._build_execute_task_command(what=what, roots=roots) + process = _PopenActivitySubprocess.start( + what=what, + dag_rel_path=dag_rel_path, + bundle_info=bundle_info, + client=client, + logger=logger, + subprocess_logs_to_stdout=subprocess_logs_to_stdout, + sentry_integration=sentry_integration, + command=command, + subprocess_schema_version=subprocess_schema_version, + startup_timeout=self.task_startup_timeout, + ) + exit_code = process.wait() + return self.ExecutionResult(exit_code, process.final_state) diff --git a/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py index 70ba5490450c3..0bc1e9e72f8d0 100644 --- a/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py @@ -339,16 +339,22 @@ class ExecutableCoordinator(SubprocessCoordinator): :param executables_root: A list of directories scanned for executable bundles when a Python stub DAG delegates task execution to a native - runtime. + runtime. See :class:`SubprocessCoordinator` for its interaction with + ``dag_bundle_name``. :param task_startup_timeout: Maximum time the coordinator waits for a task process to start, in seconds. The default is 10 seconds. """ executables_root: list[pathlib.Path] = attrs.field( converter=convert_roots, - validator=attrs.validators.min_len(1), + factory=list, ) - def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], str | None]: - bundle = _Bundle.find(self.executables_root, what.dag_id) + def __attrs_post_init__(self) -> None: + self._classify_artifact_source(self.executables_root, root_kwarg="executables_root") + + def _build_execute_task_command( + self, *, what: TaskInstance, roots: list[pathlib.Path] + ) -> tuple[list[str], str | None]: + bundle = _Bundle.find(roots, what.dag_id) return [str(bundle.path)], bundle.schema_version diff --git a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py index 0ebe089162ba1..5f4ad7fd70aaa 100644 --- a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py @@ -172,7 +172,8 @@ class JavaCoordinator(SubprocessCoordinator): :param java_executable: Path to the ``java`` command (defaults to ``"java"``, which relies on ``$PATH``). :param jvm_args: Extra arguments passed to the JVM (e.g. ``["-Xmx512m"]``). - :param jars_root: A list of directories scanned for JAR bundles. + :param jars_root: A list of directories scanned for JAR bundles. See + :class:`SubprocessCoordinator` for its interaction with ``dag_bundle_name``. :param main_class: Explicit entry point to execute with *java_executable*. :param task_startup_timeout: Maximum time the coordinator waits for a task process to start, in seconds. The default is 10 seconds. @@ -200,16 +201,21 @@ class JavaCoordinator(SubprocessCoordinator): jvm_args: list[str] = attrs.field(factory=list) jars_root: list[pathlib.Path] = attrs.field( converter=convert_roots, - validator=attrs.validators.min_len(1), + factory=list, ) main_class: str = "" - def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], str | None]: - jar = _JarInfo.find(self.jars_root, self.main_class) + def __attrs_post_init__(self) -> None: + self._classify_artifact_source(self.jars_root, root_kwarg="jars_root") + + def _build_execute_task_command( + self, *, what: TaskInstance, roots: list[pathlib.Path] + ) -> tuple[list[str], str | None]: + jar = _JarInfo.find(roots, self.main_class) command = [ self.java_executable, "-classpath", - _calculate_classpath(self.jars_root), + _calculate_classpath(roots), *self.jvm_args, jar.main_class, ] diff --git a/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py index cd352f3871bdc..a4c181388daec 100644 --- a/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py @@ -111,7 +111,8 @@ class NodeCoordinator(SubprocessCoordinator): :param node_executable: Path to the ``node`` binary (defaults to ``"node"``, which relies on ``$PATH``). :param bundles_root: Directories searched recursively, in order, for the first verified - ``*.min.mjs`` bundle declaring the task instance's Dag. + ``*.min.mjs`` bundle declaring the task instance's Dag. See :class:`SubprocessCoordinator` + for its interaction with ``dag_bundle_name``. :param task_startup_timeout: Maximum time the coordinator waits for a task process to start, in seconds. The default is 10 seconds. """ @@ -119,9 +120,14 @@ class NodeCoordinator(SubprocessCoordinator): node_executable: str = "node" bundles_root: list[pathlib.Path] = attrs.field( converter=convert_roots, - validator=attrs.validators.min_len(1), + factory=list, ) - def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], str | None]: - bundle = _Bundle.find(self.bundles_root, what.dag_id) + def __attrs_post_init__(self) -> None: + self._classify_artifact_source(self.bundles_root, root_kwarg="bundles_root") + + def _build_execute_task_command( + self, *, what: TaskInstance, roots: list[pathlib.Path] + ) -> tuple[list[str], str | None]: + bundle = _Bundle.find(roots, what.dag_id) return [self.node_executable, os.fspath(bundle.path)], bundle.schema_version 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 46f2fdae41128..415f24b0a2093 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -159,6 +159,7 @@ from pendulum.datetime import DateTime from structlog.typing import FilteringBoundLogger as Logger + from airflow.sdk.api.datamodels._generated import BundleInfo from airflow.sdk.definitions._internal.abstractoperator import AbstractOperator from airflow.sdk.definitions.context import Context from airflow.sdk.definitions.retry_policy import RetryDecision @@ -1025,14 +1026,15 @@ def _register_deserialization_allowed_classes(dag, log: Logger) -> None: ) -@detail_span("parse") -def parse(what: StartupDetails, log: Logger) -> RuntimeTaskInstance: - # TODO: Task-SDK: - # Using BundleDagBag here is about 98% wrong, but it'll do for now - from airflow.dag_processing.dagbag import BundleDagBag +def initialize_ti_bundle(bundle_info: BundleInfo, log: Logger) -> BaseDagBundle: + """ + Resolve, initialize, and access-check the Dag bundle for a task instance. - bundle_info = what.bundle_info - bundle_prepare_start = time.monotonic() + Shared by :func:`parse` (Python task path) and the subprocess coordinators + (language-SDK path), which both need a task instance's bundle materialized + on disk before use. Returns the initialized bundle so callers can read + ``bundle.path``. + """ bundle_instance = DagBundlesManager().get_bundle( name=bundle_info.name, version=bundle_info.version, @@ -1040,6 +1042,18 @@ def parse(what: StartupDetails, log: Logger) -> RuntimeTaskInstance: ) bundle_instance.initialize() _verify_bundle_access(bundle_instance, log) + return bundle_instance + + +@detail_span("parse") +def parse(what: StartupDetails, log: Logger) -> RuntimeTaskInstance: + # TODO: Task-SDK: + # Using BundleDagBag here is about 98% wrong, but it'll do for now + from airflow.dag_processing.dagbag import BundleDagBag + + bundle_info = what.bundle_info + bundle_prepare_start = time.monotonic() + bundle_instance = initialize_ti_bundle(bundle_info, log) bundle_prepare_ms = int((time.monotonic() - bundle_prepare_start) * 1000) dag_absolute_path = os.fspath(Path(bundle_instance.path, what.dag_rel_path)) diff --git a/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py index 5963e5f4892a7..4aad178df823f 100644 --- a/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py @@ -370,13 +370,42 @@ def test_executables_root_accepts_list(self, tmp_path): coordinator = ExecutableCoordinator(executables_root=[str(tmp_path), other]) assert coordinator.executables_root == [tmp_path, other] - def test_executables_root_required(self): - with pytest.raises(TypeError, match="executables_root"): - ExecutableCoordinator() - - def test_executables_root_must_be_non_empty(self): - with pytest.raises(ValueError, match="executables_root"): - ExecutableCoordinator(executables_root=None) + def test_executables_root_optional_defaults_to_empty(self): + # Neither an explicit root nor dag_bundle_name: co-located mode, valid. + coordinator = ExecutableCoordinator() + assert coordinator.executables_root == [] + assert coordinator.dag_bundle_name is None + + def test_none_executables_root_normalized_to_empty(self): + coordinator = ExecutableCoordinator(executables_root=None) + assert coordinator.executables_root == [] + + def test_root_and_dag_bundle_name_are_mutually_exclusive(self, tmp_path): + with pytest.raises(ValueError, match="at most one of 'executables_root' or 'dag_bundle_name'"): + ExecutableCoordinator(executables_root=[tmp_path], dag_bundle_name="artifacts") + + @patch("airflow.dag_processing.bundles.manager.DagBundlesManager") + def test_unconfigured_dag_bundle_name_raises(self, mock_manager): + mock_manager.is_bundle_configured.return_value = False + with pytest.raises(ValueError, match="unconfigured Dag bundle 'ghost'"): + ExecutableCoordinator(dag_bundle_name="ghost") + + @patch("airflow.dag_processing.bundles.manager.DagBundlesManager") + def test_configured_dag_bundle_name_accepted(self, mock_manager): + mock_manager.is_bundle_configured.return_value = True + coordinator = ExecutableCoordinator(dag_bundle_name="artifacts") + assert coordinator.dag_bundle_name == "artifacts" + mock_manager.is_bundle_configured.assert_called_once_with("artifacts") + + def test_build_command_scans_given_roots(self, tmp_path): + # The base resolves roots and hands them in; the subclass just scans them. + binary = _build_bundle(tmp_path / "my_bundle", dag_ids=["tutorial_dag"]) + coordinator = ExecutableCoordinator(executables_root=[tmp_path]) + command, schema_version = coordinator._build_execute_task_command( + what=_make_ti(dag_id="tutorial_dag"), roots=[tmp_path] + ) + assert command == [str(binary.resolve())] + assert schema_version == "2026-06-16" class TestBuildExecuteTaskCommand: @@ -385,7 +414,7 @@ def test_returns_resolved_executable_and_schema_version(self, tmp_path): ti = _make_ti(dag_id="tutorial_dag") coordinator = ExecutableCoordinator(executables_root=[tmp_path]) - command, schema_version = coordinator._build_execute_task_command(what=ti) + command, schema_version = coordinator._build_execute_task_command(what=ti, roots=[tmp_path]) assert command == [str(binary.resolve())] assert schema_version == "2026-06-16" @@ -397,7 +426,7 @@ def test_raises_when_bundle_omits_schema_version(self, tmp_path): coordinator = ExecutableCoordinator(executables_root=[tmp_path]) with pytest.raises(FileNotFoundError, match="matching bundles were rejected"): - coordinator._build_execute_task_command(what=ti) + coordinator._build_execute_task_command(what=ti, roots=[tmp_path]) def test_raises_when_dag_id_not_found(self, tmp_path): _build_bundle(tmp_path / "my_bundle", dag_ids=["other_dag"]) @@ -405,7 +434,7 @@ def test_raises_when_dag_id_not_found(self, tmp_path): coordinator = ExecutableCoordinator(executables_root=[tmp_path]) with pytest.raises(FileNotFoundError, match="cannot find executable bundle"): - coordinator._build_execute_task_command(what=ti) + coordinator._build_execute_task_command(what=ti, roots=[tmp_path]) @pytest.fixture diff --git a/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py index c30ecff35b9d2..9979197f1a00e 100644 --- a/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py @@ -237,6 +237,31 @@ def test_custom_kwargs(self): assert coordinator.jvm_args == ["-Xmx512m", "-Xms256m"] assert coordinator.jars_root == [pathlib.Path("/airflow/java-bundles")] + def test_jars_root_optional_defaults_to_empty(self): + # Neither an explicit root nor dag_bundle_name: co-located mode, valid. + coordinator = JavaCoordinator() + assert coordinator.jars_root == [] + assert coordinator.dag_bundle_name is None + + def test_root_and_dag_bundle_name_are_mutually_exclusive(self): + with pytest.raises(ValueError, match="at most one of 'jars_root' or 'dag_bundle_name'"): + JavaCoordinator(jars_root="/airflow/java-bundles", dag_bundle_name="artifacts") + + @patch("airflow.dag_processing.bundles.manager.DagBundlesManager") + def test_unconfigured_dag_bundle_name_raises(self, mock_manager): + mock_manager.is_bundle_configured.return_value = False + with pytest.raises(ValueError, match="unconfigured Dag bundle 'ghost'"): + JavaCoordinator(dag_bundle_name="ghost") + + def test_build_command_scans_given_roots(self, tmp_path): + # The base resolves roots and hands them in; the subclass just scans them. + _make_jar(tmp_path / "app.jar", main_class="com.example.TaskRunner", schema_version="2026-06-16") + coordinator = JavaCoordinator(jars_root=tmp_path) + command, schema_version = coordinator._build_execute_task_command(what=_make_ti(), roots=[tmp_path]) + assert command[0] == "java" + assert command[-1] == "com.example.TaskRunner" + assert schema_version == "2026-06-16" + @pytest.fixture def jars_root(tmp_path): diff --git a/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py index 492c410c8aab6..b984478c5385c 100644 --- a/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py @@ -21,6 +21,7 @@ import json import pathlib from unittest import mock +from unittest.mock import patch import pytest from task_sdk.coordinators.node._bundle_test_utils import ( @@ -79,9 +80,33 @@ def test_custom_kwargs(self): ] assert coordinator.task_startup_timeout == 30.0 - def test_bundles_root_is_required(self): - with pytest.raises(ValueError, match="Length of 'bundles_root' must be >= 1"): - NodeCoordinator(bundles_root=None) + def test_bundles_root_optional_defaults_to_empty(self): + # Neither an explicit root nor dag_bundle_name: co-located mode, valid. + coordinator = NodeCoordinator() + assert coordinator.bundles_root == [] + assert coordinator.dag_bundle_name is None + + def test_none_bundles_root_normalized_to_empty(self): + coordinator = NodeCoordinator(bundles_root=None) + assert coordinator.bundles_root == [] + + def test_root_and_dag_bundle_name_are_mutually_exclusive(self): + with pytest.raises(ValueError, match="at most one of 'bundles_root' or 'dag_bundle_name'"): + NodeCoordinator(bundles_root="/airflow/ts-bundles", dag_bundle_name="artifacts") + + @patch("airflow.dag_processing.bundles.manager.DagBundlesManager") + def test_unconfigured_dag_bundle_name_raises(self, mock_manager): + mock_manager.is_bundle_configured.return_value = False + with pytest.raises(ValueError, match="unconfigured Dag bundle 'ghost'"): + NodeCoordinator(dag_bundle_name="ghost") + + def test_build_command_scans_given_roots(self, tmp_path): + # The base resolves roots and hands them in; the subclass just scans them. + bundle = write_bundle(tmp_path, "test_dag") + coordinator = NodeCoordinator(bundles_root=tmp_path) + command, schema_version = coordinator._build_execute_task_command(what=_make_ti(), roots=[tmp_path]) + assert command == ["node", str(bundle)] + assert schema_version == SCHEMA_VERSION class TestNodeCoordinatorExecuteTaskCommand: @@ -92,11 +117,25 @@ def test_selects_bundle_by_dag_id(self, tmp_path): bundles_root=tmp_path, ) - command, schema_version = coordinator._build_execute_task_command(what=_make_ti(dag_id="sales")) + command, schema_version = coordinator._build_execute_task_command( + what=_make_ti(dag_id="sales"), roots=[tmp_path] + ) assert command == ["/opt/node/bin/node", str(selected)] assert schema_version == SCHEMA_VERSION + def test_build_execute_task_command_returns_node_bundle_and_schema_version(self, tmp_path): + bundle = write_bundle(tmp_path, "test_dag") + coordinator = NodeCoordinator( + node_executable="/opt/node/bin/node", + bundles_root=tmp_path, + ) + + command, schema_version = coordinator._build_execute_task_command(what=_make_ti(), roots=[tmp_path]) + + assert command == ["/opt/node/bin/node", str(bundle)] + assert schema_version == SCHEMA_VERSION + class TestBundleFind: @pytest.mark.parametrize( diff --git a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py index 62b7fbcf39c17..760e44f745507 100644 --- a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py +++ b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py @@ -19,6 +19,7 @@ import contextlib import os +import pathlib import socket import subprocess import sys @@ -32,15 +33,17 @@ from uuid6 import uuid7 from airflow.sdk.api.client import Client, TaskInstanceOperations -from airflow.sdk.api.datamodels._generated import TaskInstance +from airflow.sdk.api.datamodels._generated import BundleInfo, TaskInstance from airflow.sdk.coordinators._subprocess import ( SubprocessCoordinator, _accept_connections, + _ArtifactSource, _connection_owned_by_process_tree, _is_connection_from_process, _PopenActivitySubprocess, _ResourceTracker, _start_server, + log, ) from airflow.sdk.execution_time.coordinator import BaseCoordinator from airflow.sdk.execution_time.supervisor import ActivitySubprocess @@ -572,7 +575,11 @@ class _StubSubprocessCoordinator(SubprocessCoordinator): command: list[str] schema_version: str | None = None - def _build_execute_task_command(self, *, what): + def __attrs_post_init__(self) -> None: + # Explicit root so execute_task resolves without touching a Dag bundle. + self._classify_artifact_source([pathlib.Path(".")], root_kwarg="command") + + def _build_execute_task_command(self, *, what, roots): return list(self.command), self.schema_version @@ -598,7 +605,7 @@ class _Plain(SubprocessCoordinator): pass with pytest.raises(NotImplementedError): - _Plain()._build_execute_task_command(what=_make_ti()) + _Plain()._build_execute_task_command(what=_make_ti(), roots=[]) class TestSubprocessCoordinatorExecuteTask: @@ -844,3 +851,122 @@ def test_register_pipe_readers_called_with_four_sockets(self, mock_client): subprocess_logs_to_stdout=False, ) assert mock_register.mock_calls == [call(ANY, ANY, ANY, ANY, data=ANY)] + + +class TestClassifyArtifactSource: + """Construction-time classification and per-mode validation.""" + + def test_explicit_root_classified_and_recorded(self): + coordinator = SubprocessCoordinator() + configured = [pathlib.Path("/artifacts")] + coordinator._classify_artifact_source(configured, root_kwarg="jars_root") + assert coordinator._artifact_source is _ArtifactSource.EXPLICIT_ROOT + assert coordinator._configured_roots == configured + + def test_task_bundle_classified_when_neither_set(self): + coordinator = SubprocessCoordinator() + coordinator._classify_artifact_source([], root_kwarg="jars_root") + assert coordinator._artifact_source is _ArtifactSource.TASK_BUNDLE + + def test_rejects_explicit_root_and_dag_bundle_name_together(self): + coordinator = SubprocessCoordinator(dag_bundle_name="artifacts") + with pytest.raises(ValueError, match="at most one of 'jars_root' or 'dag_bundle_name'"): + coordinator._classify_artifact_source([pathlib.Path("/artifacts")], root_kwarg="jars_root") + + @patch("airflow.dag_processing.bundles.manager.DagBundlesManager") + def test_rejects_unconfigured_dag_bundle_name(self, mock_manager): + mock_manager.is_bundle_configured.return_value = False + coordinator = SubprocessCoordinator(dag_bundle_name="ghost") + with pytest.raises(ValueError, match="unconfigured Dag bundle 'ghost'"): + coordinator._classify_artifact_source([], root_kwarg="jars_root") + + @patch("airflow.dag_processing.bundles.manager.DagBundlesManager") + def test_accepts_configured_dag_bundle_name(self, mock_manager): + mock_manager.is_bundle_configured.return_value = True + coordinator = SubprocessCoordinator(dag_bundle_name="artifacts") + coordinator._classify_artifact_source([], root_kwarg="jars_root") + assert coordinator._artifact_source is _ArtifactSource.NAMED_BUNDLE + mock_manager.is_bundle_configured.assert_called_once_with("artifacts") + + +class TestInitRootSource: + """Execute-time root resolution, dispatched on the classified mode.""" + + def test_returns_configured_root_in_explicit_mode(self): + coordinator = SubprocessCoordinator() + configured = [pathlib.Path("/a"), pathlib.Path("/b")] + coordinator._artifact_source = _ArtifactSource.EXPLICIT_ROOT + coordinator._configured_roots = configured + assert coordinator._init_root_source() == configured + + @patch("airflow.sdk.execution_time.task_runner.initialize_ti_bundle") + def test_task_bundle_mode_uses_active_task_bundle_verbatim(self, mock_initialize, tmp_path): + mock_initialize.return_value = MagicMock(path=tmp_path) + coordinator = SubprocessCoordinator() + coordinator._artifact_source = _ArtifactSource.TASK_BUNDLE + bundle_info = BundleInfo(name="dags", version="v3", version_data={"k": "v"}) + coordinator._active_bundle_info = bundle_info + + roots = coordinator._init_root_source() + + assert roots == [tmp_path] + mock_initialize.assert_called_once_with(bundle_info, log) + + def test_task_bundle_mode_without_active_task_raises(self): + coordinator = SubprocessCoordinator() + coordinator._artifact_source = _ArtifactSource.TASK_BUNDLE + with pytest.raises(RuntimeError, match="requires an active task"): + coordinator._init_root_source() + + @patch("airflow.sdk.execution_time.task_runner.initialize_ti_bundle") + def test_named_bundle_mode_resolves_latest_version(self, mock_initialize, tmp_path): + mock_initialize.return_value = MagicMock(path=tmp_path) + coordinator = SubprocessCoordinator(dag_bundle_name="artifacts") + coordinator._artifact_source = _ArtifactSource.NAMED_BUNDLE + + roots = coordinator._init_root_source() + + assert roots == [tmp_path] + mock_initialize.assert_called_once_with(BundleInfo(name="artifacts"), log) + + @patch("airflow.sdk.execution_time.task_runner.initialize_ti_bundle") + def test_missing_resolved_path_raises(self, mock_initialize, tmp_path): + missing = tmp_path / "nope" + mock_initialize.return_value = MagicMock(path=missing) + coordinator = SubprocessCoordinator(dag_bundle_name="artifacts") + coordinator._artifact_source = _ArtifactSource.NAMED_BUNDLE + + with pytest.raises(FileNotFoundError, match="does not exist"): + coordinator._init_root_source() + + +class TestSetCurrentBundle: + def test_binds_for_the_block_and_clears_on_exit_allowing_reuse(self): + coordinator = _StubSubprocessCoordinator(command=["/bin/true"]) + bundle = MagicMock() + with coordinator._set_current_bundle(bundle): + assert coordinator._active_bundle_info is bundle + assert coordinator._active_bundle_info is None + # A subsequent task on the same (reused) coordinator binds cleanly. + with coordinator._set_current_bundle(MagicMock()): + assert coordinator._active_bundle_info is not None + assert coordinator._active_bundle_info is None + + def test_rejects_nested_reentry(self): + coordinator = _StubSubprocessCoordinator(command=["/bin/true"]) + with coordinator._set_current_bundle(MagicMock()): + with pytest.raises(RuntimeError, match="not re-entrant"): + with coordinator._set_current_bundle(MagicMock()): + pass + + def test_execute_task_is_not_reentrant(self, mock_client): + coordinator = _StubSubprocessCoordinator(command=["/bin/true"]) + coordinator._active_bundle_info = MagicMock() # simulate an in-flight task + with pytest.raises(RuntimeError, match="not re-entrant"): + coordinator.execute_task( + what=_make_ti(), + dag_rel_path="dag.py", + bundle_info=MagicMock(), + client=mock_client, + subprocess_logs_to_stdout=False, + ) 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 198b6ec070b9c..bb152fef8187b 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 @@ -887,6 +887,30 @@ def test_verify_bundle_access_skips_nonexistent_path(tmp_path: Path): _verify_bundle_access(mock_bundle, mock.Mock()) +def test_initialize_ti_bundle_resolves_initializes_and_verifies(tmp_path: Path): + """initialize_ti_bundle resolves the bundle from BundleInfo, initializes it, and access-checks it.""" + from airflow.sdk.api.datamodels._generated import BundleInfo + from airflow.sdk.execution_time.task_runner import initialize_ti_bundle + + bundle_path = tmp_path / "bundle" + bundle_path.mkdir() + mock_bundle = mock.Mock() + mock_bundle.path = bundle_path + mock_bundle.name = "my-bundle" + + bundle_info = BundleInfo(name="my-bundle", version="v2", version_data={"k": "v"}) + + with patch("airflow.sdk.execution_time.task_runner.DagBundlesManager") as mock_manager: + mock_manager.return_value.get_bundle.return_value = mock_bundle + result = initialize_ti_bundle(bundle_info, mock.Mock()) + + assert result is mock_bundle + mock_manager.return_value.get_bundle.assert_called_once_with( + name="my-bundle", version="v2", version_data={"k": "v"} + ) + mock_bundle.initialize.assert_called_once_with() + + @pytest.mark.parametrize("use_queues", [False, True]) def test_run_deferred_basic(time_machine, create_runtime_ti, mock_supervisor_comms, use_queues: bool): """Test that a task can transition to a deferred state.""" From 9257c50cf1668cb985ceb9a6948ebe60185692f6 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Fri, 31 Jul 2026 15:30:13 +0000 Subject: [PATCH 02/11] Document dag_bundle_name option for language-SDK coordinators --- .../language-sdks/go.rst | 20 +++++++++++++++++-- .../language-sdks/java.rst | 20 +++++++++++++++++-- .../language-sdks/typescript.rst | 20 +++++++++++++++++-- contributing-docs/30_new_language_sdk.rst | 9 ++++++++- 4 files changed, 62 insertions(+), 7 deletions(-) diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst index 72ea37bf66583..2c3bb1c58b0ef 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst @@ -507,14 +507,30 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the - Default - Description * - ``executables_root`` - - *(required)* + - *(optional)* - One or more directories scanned recursively for executable bundles. Accepts a string, - a path, or a list of strings/paths. + a path, or a list of strings/paths. When omitted, bundles are located through a Dag + bundle instead (see the note below). + * - ``dag_bundle_name`` + - *(auto: task's own bundle)* + - Name of a configured Dag bundle to load executable bundles from. Mutually exclusive + with ``executables_root``. * - ``task_startup_timeout`` - ``10.0`` - Seconds to wait for the bundle subprocess to connect after launch. Increase this if your bundle startup is slow (e.g. on constrained hardware). +.. note:: + + **Locating bundles.** ``executables_root`` and ``dag_bundle_name`` are mutually exclusive, + and both are optional: + + * Set ``executables_root`` to scan explicit filesystem directories you manage yourself. + * Set ``dag_bundle_name`` to load bundles from a configured Dag bundle (its latest version), + so they are delivered and versioned through the same bundle machinery as your Dags. + * Leave both unset (the default) to load bundles from the **task's own** Dag bundle, pinned + to the version the run was created with. + .. _go-sdk/limitations: Limitations diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst index 1561ecfffcfab..fa27d715924d0 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst @@ -1022,9 +1022,14 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the - Default - Description * - ``jars_root`` - - *(required)* + - *(optional)* - One or more directories scanned recursively for ``.jar`` files. Accepts a string, - a path, or a list of strings/paths. + a path, or a list of strings/paths. When omitted, JARs are located through a Dag + bundle instead (see the note below). + * - ``dag_bundle_name`` + - *(auto: task's own bundle)* + - Name of a configured Dag bundle to load JARs from. Mutually exclusive with + ``jars_root``. * - ``java_executable`` - ``"java"`` - Path to the ``java`` binary. Defaults to ``java`` on ``$PATH``. @@ -1042,6 +1047,17 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the - Seconds to wait for the JVM subprocess to connect after launch. Increase this if your JVM startup is slow (e.g. on constrained hardware or with a large classpath). +.. note:: + + **Locating JARs.** ``jars_root`` and ``dag_bundle_name`` are mutually exclusive, and both + are optional: + + * Set ``jars_root`` to scan explicit filesystem directories you manage yourself. + * Set ``dag_bundle_name`` to load JARs from a configured Dag bundle (its latest version), so + they are delivered and versioned through the same bundle machinery as your Dags. + * Leave both unset (the default) to load JARs from the **task's own** Dag bundle, pinned to + the version the run was created with. + .. note:: The ``[sdk]`` configuration is read at startup, so changes to ``coordinators`` or diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst index d3ddca0c9b49b..d59893b3359e6 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst @@ -446,9 +446,14 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the - Default - Description * - ``bundles_root`` - - *(required)* + - *(optional)* - One or more directories searched recursively, in order, for an integrity-verified ``*.min.mjs`` - bundle that declares the requested Dag. Accepts a string, a path, or a list of strings/paths. + bundle that declares the requested Dag. Accepts a string, a path, or a list of strings/paths. When + omitted, the bundle is located through a Dag bundle instead (see the note below). + * - ``dag_bundle_name`` + - *(auto: task's own bundle)* + - Name of a configured Dag bundle to load the ``*.min.mjs`` bundle from. Mutually exclusive with + ``bundles_root``. * - ``node_executable`` - ``"node"`` - Path to the ``node`` binary. Defaults to ``node`` on ``$PATH``. @@ -457,6 +462,17 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the - Seconds to wait for the Node.js subprocess to connect after launch. Increase this if your bundle startup is slow (e.g. on constrained hardware). +.. note:: + + **Locating the bundle.** ``bundles_root`` and ``dag_bundle_name`` are mutually exclusive, and both + are optional: + + * Set ``bundles_root`` to scan explicit filesystem directories you manage yourself. + * Set ``dag_bundle_name`` to load the bundle from a configured Dag bundle (its latest version), so it + is delivered and versioned through the same bundle machinery as your Dags. + * Leave both unset (the default) to load the bundle from the **task's own** Dag bundle, pinned to the + version the run was created with. + Limitations ----------- diff --git a/contributing-docs/30_new_language_sdk.rst b/contributing-docs/30_new_language_sdk.rst index 80fbe9e4235cb..8983be1681de4 100644 --- a/contributing-docs/30_new_language_sdk.rst +++ b/contributing-docs/30_new_language_sdk.rst @@ -110,7 +110,9 @@ the subprocess understands: .. code-block:: python - def _build_execute_task_command(self, *, what: TaskInstanceDTO) -> tuple[list[str], str]: ... + def _build_execute_task_command( + self, *, what: TaskInstanceDTO, roots: list[pathlib.Path] + ) -> tuple[list[str], str]: ... The method returns a ``(command, subprocess_schema_version)`` pair: @@ -120,6 +122,11 @@ The method returns a ``(command, subprocess_schema_version)`` pair: subprocess understands, used by the supervisor to negotiate message formats across SDK versions. See `Supervisor Schema`_ below. +``roots`` are the artifact directories the base class has already resolved from +the coordinator's configured source — an explicit filesystem root, a named Dag +bundle (``dag_bundle_name``), or the task's own bundle — so subclasses scan +``roots`` rather than reading the configured root directly. + Supervisor Schema ~~~~~~~~~~~~~~~~~ From 0e857247f41286884ff606a561c95c022c98392c Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Sat, 1 Aug 2026 14:11:22 +0000 Subject: [PATCH 03/11] Harden subprocess coordinator artifact bundle resolution Language-SDK subprocess coordinators materialized artifact bundles without the version-lock protection the Python task path holds, so bundle cleanup could delete a version a running JVM was still reading from. Resolution now holds the version lock for the subprocess lifetime, mirroring the Python task path, and materialization failures surface in the task log rather than the worker log. Classification moves to the base coordinator so a subclass that adds no explicit-root kwarg falls back to the task-bundle default instead of failing every task, the Java coordinator now requires an explicit entrypoint when it has to scan a whole bundle, and configured-bundle lookup checks names only so an unrelated unimportable bundle no longer breaks coordinator construction. --- .../airflow/dag_processing/bundles/manager.py | 19 +- .../bundles/test_dag_bundle_manager.py | 26 +++ .../airflow/sdk/coordinators/_subprocess.py | 71 ++++++-- .../coordinators/executable/coordinator.py | 8 +- .../sdk/coordinators/java/coordinator.py | 20 ++- .../sdk/coordinators/node/coordinator.py | 8 +- .../executable/test_coordinator.py | 7 +- .../coordinators/java/test_coordinator.py | 22 ++- .../coordinators/node/test_coordinator.py | 7 +- .../task_sdk/coordinators/test_subprocess.py | 162 ++++++++++++++---- 10 files changed, 271 insertions(+), 79 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/bundles/manager.py b/airflow-core/src/airflow/dag_processing/bundles/manager.py index 0245143a0cdda..4021ee79d6538 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/manager.py +++ b/airflow-core/src/airflow/dag_processing/bundles/manager.py @@ -16,6 +16,7 @@ # under the License. from __future__ import annotations +import functools import importlib import logging import os @@ -107,6 +108,20 @@ def _bundle_item_exc(msg): ) +@functools.cache +def _configured_bundle_names() -> frozenset[str]: + """Return configured Dag bundle names, parsed by name only (no class import).""" + config_list = conf.getjson("dag_processor", "dag_bundle_config_list") + if not config_list: + return frozenset() + if not isinstance(config_list, list): + raise AirflowConfigException( + "Section `dag_processor` key `dag_bundle_config_list` " + f"must be list but got {config_list.__class__}" + ) + return frozenset(cfg.name for cfg in _parse_bundle_config(config_list)) + + def _parse_bundle_config(config_list) -> list[_ExternalBundleConfig]: bundles = {} for item in config_list: @@ -675,8 +690,8 @@ def get_bundle( @classmethod def is_bundle_configured(cls, name: str) -> bool: - """Return whether *name* is a configured Dag bundle, without constructing the bundle.""" - return name in cls()._bundle_config + """Return whether *name* is a configured Dag bundle, by config name only (no class import).""" + return name in _configured_bundle_names() def get_all_dag_bundles(self) -> Iterable[BaseDagBundle]: """ diff --git a/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py b/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py index f1ce2a92ca237..d4fa9f3209262 100644 --- a/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py +++ b/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py @@ -31,6 +31,7 @@ from airflow.dag_processing.bundles.base import BaseDagBundle from airflow.dag_processing.bundles.manager import ( DagBundlesManager, + _configured_bundle_names, _get_configured_bundle_team_names, _guess_best_bundle_for_fileloc, ) @@ -205,6 +206,15 @@ def test_get_bundle(): assert bundle.version is None +@pytest.fixture(autouse=True) +def _clear_configured_bundle_names_cache(): + # _configured_bundle_names is a process-global cache; clear it so tests that + # set different bundle configs don't observe each other's cached result. + _configured_bundle_names.cache_clear() + yield + _configured_bundle_names.cache_clear() + + def test_is_bundle_configured(): """is_bundle_configured reports membership without constructing the bundle.""" with patch.dict( @@ -214,6 +224,22 @@ def test_is_bundle_configured(): assert DagBundlesManager.is_bundle_configured("bundle-that-doesn't-exist") is False +def test_is_bundle_configured_does_not_import_bundle_class(): + """A validly-named bundle whose class is unimportable is still reported as configured. + + The check reads names only; the class is imported lazily at materialization, + so an unimportable classpath must not make a coordinator fail at construction. + """ + config = [{"name": "artifacts", "classpath": "does.not.exist.Bundle", "kwargs": {}}] + with patch.dict(os.environ, {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(config)}): + assert DagBundlesManager.is_bundle_configured("artifacts") is True + + +def test_is_bundle_configured_empty_config(): + with patch.dict(os.environ, {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": "[]"}): + assert DagBundlesManager.is_bundle_configured("anything") is False + + @pytest.fixture def clear_db(): clear_db_dag_bundles() diff --git a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py index f157c82cfef98..6da048fe156f7 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py +++ b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py @@ -37,7 +37,7 @@ import socket import subprocess import time -from typing import TYPE_CHECKING, TypeVar, cast +from typing import TYPE_CHECKING, ClassVar, TypeVar, cast import attrs import psutil @@ -54,6 +54,7 @@ from structlog.typing import FilteringBoundLogger from typing_extensions import Self + from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: SDK002 from airflow.sdk.api.client import Client from airflow.sdk.api.datamodels._generated import BundleInfo, TaskInstance @@ -408,6 +409,12 @@ class SubprocessCoordinator(BaseCoordinator): task_startup_timeout: float = 10.0 dag_bundle_name: str | None = None + # Name of the subclass's explicit-root kwarg, used only in error messages. + # Subclasses that expose an explicit root set this and override + # :meth:`_explicit_artifact_roots`; a subclass that does neither lands on the + # task-bundle default instead of failing at execute time. + _root_kwarg: ClassVar[str] = "root" + # Classified once at construction by :meth:`_classify_artifact_source` and # dispatched on by :meth:`_init_root_source` at execute time. _artifact_source: _ArtifactSource | None = attrs.field(init=False, default=None) @@ -419,12 +426,24 @@ class SubprocessCoordinator(BaseCoordinator): # co-located artifacts. _active_bundle_info: BundleInfo | None = attrs.field(init=False, default=None) + @property + def _explicit_artifact_roots(self) -> Sequence[pathlib.Path]: + """Subclass's explicit artifact roots; empty (the default) selects task-bundle mode.""" + return [] + + def _validate_artifact_source(self) -> None: + """Run subclass-specific validation after classification; default no-op.""" + + def __attrs_post_init__(self) -> None: + self._classify_artifact_source(self._explicit_artifact_roots, root_kwarg=self._root_kwarg) + self._validate_artifact_source() + def _classify_artifact_source(self, configured: Sequence[pathlib.Path], *, root_kwarg: str) -> None: """ Classify and validate how this coordinator locates artifacts (construction time). - Subclasses call this from ``__attrs_post_init__`` with their own root - field. It rejects setting both an explicit root and ``dag_bundle_name``, + Called from :meth:`__attrs_post_init__` with the subclass's explicit + roots. It rejects setting both an explicit root and ``dag_bundle_name``, fails fast when ``dag_bundle_name`` names a bundle that is not configured, and records the resulting :class:`_ArtifactSource` and explicit root. """ @@ -457,23 +476,23 @@ def _classify_artifact_source(self, configured: Sequence[pathlib.Path], *, root_ details["configured_roots"] = [str(root) for root in self._configured_roots] log.debug("Coordinator artifact source selected", **details) - def _init_root_source(self) -> list[pathlib.Path]: + def _init_root_source( + self, logger: FilteringBoundLogger + ) -> tuple[list[pathlib.Path], BaseDagBundle | None]: """ - Resolve the directories to scan for artifacts for the current task. - - Dispatches on the :class:`_ArtifactSource` classified at construction. - An explicit root is returned as-is (no Dag bundle is resolved); otherwise - the root is a Dag bundle's materialized path — the named - ``dag_bundle_name`` bundle at its latest version, or the task's own - bundle pinned to the run's version. Called by :meth:`execute_task`, which - forwards the result to :meth:`_build_execute_task_command`. + Resolve the directories to scan for artifacts, dispatched on the classified mode. + + Returns ``(roots, bundle)``: an explicit root yields no bundle (``None``); + a Dag-bundle mode returns the materialized path and the resolved bundle so + :meth:`execute_task` can hold a version lock over it. *logger* is the task + logger, so materialization failures surface in the task log. """ if self._artifact_source is _ArtifactSource.EXPLICIT_ROOT: - return self._configured_roots + return self._configured_roots, None - if self._artifact_source is _ArtifactSource.NAMED_BUNDLE: - from airflow.sdk.api.datamodels._generated import BundleInfo + from airflow.sdk.api.datamodels._generated import BundleInfo + if self._artifact_source is _ArtifactSource.NAMED_BUNDLE: # NAMED_BUNDLE implies dag_bundle_name is set. target = BundleInfo(name=cast("str", self.dag_bundle_name)) elif self._artifact_source is _ArtifactSource.TASK_BUNDLE: @@ -489,11 +508,11 @@ def _init_root_source(self) -> list[pathlib.Path]: # load would risk an import cycle through the supervisor. from airflow.sdk.execution_time.task_runner import initialize_ti_bundle - bundle = initialize_ti_bundle(target, log) + bundle = initialize_ti_bundle(target, logger) path = bundle.path if not path.exists(): raise FileNotFoundError(f"Dag bundle {target.name!r} resolved to {path}, which does not exist.") - return [path] + return [path], bundle def _build_execute_task_command( self, *, what: TaskInstance, roots: list[pathlib.Path] @@ -539,8 +558,22 @@ def execute_task( subprocess_logs_to_stdout: bool, **kwargs, ) -> BaseCoordinator.ExecutionResult: - with self._set_current_bundle(bundle_info): - roots = self._init_root_source() + task_logger = logger or log + with contextlib.ExitStack() as stack: + stack.enter_context(self._set_current_bundle(bundle_info)) + roots, resolved_bundle = self._init_root_source(task_logger) + if resolved_bundle is not None: + # Hold the version lock across start()/wait() so bundle cleanup + # cannot rmtree a version this task is still reading from, + # mirroring task_runner.main() for the Python task path. + from airflow.dag_processing.bundles.base import BundleVersionLock # noqa: SDK002 + + stack.enter_context( + BundleVersionLock( + bundle_name=resolved_bundle.name, + bundle_version=resolved_bundle.version, + ) + ) command, subprocess_schema_version = self._build_execute_task_command(what=what, roots=roots) process = _PopenActivitySubprocess.start( what=what, diff --git a/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py index 0bc1e9e72f8d0..bc08ea6d6b7cc 100644 --- a/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py @@ -25,7 +25,7 @@ import stat import struct from collections import OrderedDict -from typing import TYPE_CHECKING, Any, BinaryIO +from typing import TYPE_CHECKING, Any, BinaryIO, ClassVar import attrs import structlog @@ -349,9 +349,11 @@ class ExecutableCoordinator(SubprocessCoordinator): converter=convert_roots, factory=list, ) + _root_kwarg: ClassVar[str] = "executables_root" - def __attrs_post_init__(self) -> None: - self._classify_artifact_source(self.executables_root, root_kwarg="executables_root") + @property + def _explicit_artifact_roots(self) -> list[pathlib.Path]: + return self.executables_root def _build_execute_task_command( self, *, what: TaskInstance, roots: list[pathlib.Path] diff --git a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py index 5f4ad7fd70aaa..d4f44ffbfcc11 100644 --- a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py @@ -24,7 +24,7 @@ import pathlib import stat import zipfile -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar import attrs import structlog @@ -204,9 +204,21 @@ class JavaCoordinator(SubprocessCoordinator): factory=list, ) main_class: str = "" - - def __attrs_post_init__(self) -> None: - self._classify_artifact_source(self.jars_root, root_kwarg="jars_root") + _root_kwarg: ClassVar[str] = "jars_root" + + @property + def _explicit_artifact_roots(self) -> list[pathlib.Path]: + return self.jars_root + + def _validate_artifact_source(self) -> None: + # Without an explicit root the scan falls back to a whole Dag bundle, where + # _JarInfo.find would pick the first executable JAR in walk order — a + # non-deterministic entrypoint. Require an explicit main_class in that case. + if not self.jars_root and not self.main_class: + raise ValueError( + "JavaCoordinator requires 'main_class' when 'jars_root' is not set, so the " + "entrypoint is unambiguous when scanning a Dag bundle for JARs." + ) def _build_execute_task_command( self, *, what: TaskInstance, roots: list[pathlib.Path] diff --git a/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py index a4c181388daec..d983983cd8c3f 100644 --- a/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py @@ -21,7 +21,7 @@ import os import pathlib -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar import attrs import structlog @@ -122,9 +122,11 @@ class NodeCoordinator(SubprocessCoordinator): converter=convert_roots, factory=list, ) + _root_kwarg: ClassVar[str] = "bundles_root" - def __attrs_post_init__(self) -> None: - self._classify_artifact_source(self.bundles_root, root_kwarg="bundles_root") + @property + def _explicit_artifact_roots(self) -> list[pathlib.Path]: + return self.bundles_root def _build_execute_task_command( self, *, what: TaskInstance, roots: list[pathlib.Path] diff --git a/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py index 4aad178df823f..be5b434e2d5d2 100644 --- a/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py @@ -397,10 +397,11 @@ def test_configured_dag_bundle_name_accepted(self, mock_manager): assert coordinator.dag_bundle_name == "artifacts" mock_manager.is_bundle_configured.assert_called_once_with("artifacts") - def test_build_command_scans_given_roots(self, tmp_path): - # The base resolves roots and hands them in; the subclass just scans them. + def test_build_command_scans_passed_roots_in_colocated_mode(self, tmp_path): + # Co-located mode: no configured root, so the subclass must scan the roots + # the base hands in rather than a configured executables_root field. binary = _build_bundle(tmp_path / "my_bundle", dag_ids=["tutorial_dag"]) - coordinator = ExecutableCoordinator(executables_root=[tmp_path]) + coordinator = ExecutableCoordinator() command, schema_version = coordinator._build_execute_task_command( what=_make_ti(dag_id="tutorial_dag"), roots=[tmp_path] ) diff --git a/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py index 9979197f1a00e..2ce073f147c69 100644 --- a/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py @@ -238,11 +238,20 @@ def test_custom_kwargs(self): assert coordinator.jars_root == [pathlib.Path("/airflow/java-bundles")] def test_jars_root_optional_defaults_to_empty(self): - # Neither an explicit root nor dag_bundle_name: co-located mode, valid. - coordinator = JavaCoordinator() + # Neither an explicit root nor dag_bundle_name: co-located mode, valid + # once an explicit main_class disambiguates the entrypoint. + coordinator = JavaCoordinator(main_class="com.example.Main") assert coordinator.jars_root == [] assert coordinator.dag_bundle_name is None + def test_colocated_mode_requires_main_class(self): + with pytest.raises(ValueError, match="requires 'main_class' when 'jars_root' is not set"): + JavaCoordinator() + + def test_explicit_root_does_not_require_main_class(self, tmp_path): + coordinator = JavaCoordinator(jars_root=tmp_path) + assert coordinator.main_class == "" + def test_root_and_dag_bundle_name_are_mutually_exclusive(self): with pytest.raises(ValueError, match="at most one of 'jars_root' or 'dag_bundle_name'"): JavaCoordinator(jars_root="/airflow/java-bundles", dag_bundle_name="artifacts") @@ -251,12 +260,13 @@ def test_root_and_dag_bundle_name_are_mutually_exclusive(self): def test_unconfigured_dag_bundle_name_raises(self, mock_manager): mock_manager.is_bundle_configured.return_value = False with pytest.raises(ValueError, match="unconfigured Dag bundle 'ghost'"): - JavaCoordinator(dag_bundle_name="ghost") + JavaCoordinator(main_class="com.example.Main", dag_bundle_name="ghost") - def test_build_command_scans_given_roots(self, tmp_path): - # The base resolves roots and hands them in; the subclass just scans them. + def test_build_command_scans_passed_roots_in_colocated_mode(self, tmp_path): + # Co-located mode: no configured root, so the subclass must scan the roots + # the base hands in rather than a configured jars_root field. _make_jar(tmp_path / "app.jar", main_class="com.example.TaskRunner", schema_version="2026-06-16") - coordinator = JavaCoordinator(jars_root=tmp_path) + coordinator = JavaCoordinator(main_class="com.example.TaskRunner") command, schema_version = coordinator._build_execute_task_command(what=_make_ti(), roots=[tmp_path]) assert command[0] == "java" assert command[-1] == "com.example.TaskRunner" diff --git a/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py index b984478c5385c..cde7b46a790d0 100644 --- a/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py @@ -100,10 +100,11 @@ def test_unconfigured_dag_bundle_name_raises(self, mock_manager): with pytest.raises(ValueError, match="unconfigured Dag bundle 'ghost'"): NodeCoordinator(dag_bundle_name="ghost") - def test_build_command_scans_given_roots(self, tmp_path): - # The base resolves roots and hands them in; the subclass just scans them. + def test_build_command_scans_passed_roots_in_colocated_mode(self, tmp_path): + # Co-located mode: no configured root, so the subclass must scan the roots + # the base hands in rather than a configured bundles_root field. bundle = write_bundle(tmp_path, "test_dag") - coordinator = NodeCoordinator(bundles_root=tmp_path) + coordinator = NodeCoordinator() command, schema_version = coordinator._build_execute_task_command(what=_make_ti(), roots=[tmp_path]) assert command == ["node", str(bundle)] assert schema_version == SCHEMA_VERSION diff --git a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py index 760e44f745507..bb2d1265e154b 100644 --- a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py +++ b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py @@ -25,6 +25,7 @@ import sys import threading import time +from typing import ClassVar from unittest.mock import ANY, MagicMock, call, patch import attrs @@ -570,16 +571,26 @@ def test_untrack_unknown_object_does_not_raise(self): @attrs.define(kw_only=True) class _StubSubprocessCoordinator(SubprocessCoordinator): - """Minimal SubprocessCoordinator subclass used to exercise the base machinery.""" + """Minimal SubprocessCoordinator subclass used to exercise the base machinery. + + ``explicit_roots`` defaults to a real path so the coordinator classifies as + EXPLICIT_ROOT and execute_task resolves without touching a Dag bundle; pass + ``explicit_roots=[]`` to exercise TASK_BUNDLE mode. Roots handed to the + command builder are recorded in ``recorded_roots`` so wiring can be asserted. + """ command: list[str] schema_version: str | None = None + explicit_roots: list[pathlib.Path] = attrs.field(factory=lambda: [pathlib.Path(".")]) + recorded_roots: list[list[pathlib.Path]] = attrs.field(init=False, factory=list) + _root_kwarg: ClassVar[str] = "explicit_roots" - def __attrs_post_init__(self) -> None: - # Explicit root so execute_task resolves without touching a Dag bundle. - self._classify_artifact_source([pathlib.Path(".")], root_kwarg="command") + @property + def _explicit_artifact_roots(self) -> list[pathlib.Path]: + return self.explicit_roots def _build_execute_task_command(self, *, what, roots): + self.recorded_roots.append(roots) return list(self.command), self.schema_version @@ -854,37 +865,50 @@ def test_register_pipe_readers_called_with_four_sockets(self, mock_client): class TestClassifyArtifactSource: - """Construction-time classification and per-mode validation.""" + """Construction-time classification and per-mode validation. + + Classification runs from the base ``__attrs_post_init__``, so it is exercised + through construction rather than by calling the helper directly. + """ def test_explicit_root_classified_and_recorded(self): - coordinator = SubprocessCoordinator() configured = [pathlib.Path("/artifacts")] - coordinator._classify_artifact_source(configured, root_kwarg="jars_root") + coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=configured) assert coordinator._artifact_source is _ArtifactSource.EXPLICIT_ROOT assert coordinator._configured_roots == configured def test_task_bundle_classified_when_neither_set(self): - coordinator = SubprocessCoordinator() - coordinator._classify_artifact_source([], root_kwarg="jars_root") + coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=[]) assert coordinator._artifact_source is _ArtifactSource.TASK_BUNDLE + def test_subclass_without_overrides_defaults_to_task_bundle(self): + """A subclass that overrides neither hook is classified, not left to error at execute time.""" + + @attrs.define(kw_only=True) + class _Bare(SubprocessCoordinator): + def _build_execute_task_command(self, *, what, roots): + return [], None + + assert _Bare()._artifact_source is _ArtifactSource.TASK_BUNDLE + def test_rejects_explicit_root_and_dag_bundle_name_together(self): - coordinator = SubprocessCoordinator(dag_bundle_name="artifacts") - with pytest.raises(ValueError, match="at most one of 'jars_root' or 'dag_bundle_name'"): - coordinator._classify_artifact_source([pathlib.Path("/artifacts")], root_kwarg="jars_root") + with pytest.raises(ValueError, match="at most one of 'explicit_roots' or 'dag_bundle_name'"): + _StubSubprocessCoordinator( + command=["x"], explicit_roots=[pathlib.Path("/artifacts")], dag_bundle_name="artifacts" + ) @patch("airflow.dag_processing.bundles.manager.DagBundlesManager") def test_rejects_unconfigured_dag_bundle_name(self, mock_manager): mock_manager.is_bundle_configured.return_value = False - coordinator = SubprocessCoordinator(dag_bundle_name="ghost") with pytest.raises(ValueError, match="unconfigured Dag bundle 'ghost'"): - coordinator._classify_artifact_source([], root_kwarg="jars_root") + _StubSubprocessCoordinator(command=["x"], explicit_roots=[], dag_bundle_name="ghost") @patch("airflow.dag_processing.bundles.manager.DagBundlesManager") def test_accepts_configured_dag_bundle_name(self, mock_manager): mock_manager.is_bundle_configured.return_value = True - coordinator = SubprocessCoordinator(dag_bundle_name="artifacts") - coordinator._classify_artifact_source([], root_kwarg="jars_root") + coordinator = _StubSubprocessCoordinator( + command=["x"], explicit_roots=[], dag_bundle_name="artifacts" + ) assert coordinator._artifact_source is _ArtifactSource.NAMED_BUNDLE mock_manager.is_bundle_configured.assert_called_once_with("artifacts") @@ -892,52 +916,118 @@ def test_accepts_configured_dag_bundle_name(self, mock_manager): class TestInitRootSource: """Execute-time root resolution, dispatched on the classified mode.""" - def test_returns_configured_root_in_explicit_mode(self): - coordinator = SubprocessCoordinator() + def test_returns_configured_root_and_no_bundle_in_explicit_mode(self): configured = [pathlib.Path("/a"), pathlib.Path("/b")] - coordinator._artifact_source = _ArtifactSource.EXPLICIT_ROOT - coordinator._configured_roots = configured - assert coordinator._init_root_source() == configured + coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=configured) + roots, bundle = coordinator._init_root_source(log) + assert roots == configured + assert bundle is None @patch("airflow.sdk.execution_time.task_runner.initialize_ti_bundle") - def test_task_bundle_mode_uses_active_task_bundle_verbatim(self, mock_initialize, tmp_path): - mock_initialize.return_value = MagicMock(path=tmp_path) - coordinator = SubprocessCoordinator() - coordinator._artifact_source = _ArtifactSource.TASK_BUNDLE + def test_task_bundle_mode_uses_active_task_bundle(self, mock_initialize, tmp_path): + resolved = MagicMock(supports_versioning=False, path=tmp_path) + mock_initialize.return_value = resolved + coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=[]) bundle_info = BundleInfo(name="dags", version="v3", version_data={"k": "v"}) coordinator._active_bundle_info = bundle_info - roots = coordinator._init_root_source() + roots, bundle = coordinator._init_root_source(log) assert roots == [tmp_path] + assert bundle is resolved mock_initialize.assert_called_once_with(bundle_info, log) def test_task_bundle_mode_without_active_task_raises(self): - coordinator = SubprocessCoordinator() - coordinator._artifact_source = _ArtifactSource.TASK_BUNDLE + coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=[]) with pytest.raises(RuntimeError, match="requires an active task"): - coordinator._init_root_source() + coordinator._init_root_source(log) + + def test_unclassified_source_raises(self): + coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=[]) + coordinator._artifact_source = None + with pytest.raises(RuntimeError, match="was not classified"): + coordinator._init_root_source(log) @patch("airflow.sdk.execution_time.task_runner.initialize_ti_bundle") - def test_named_bundle_mode_resolves_latest_version(self, mock_initialize, tmp_path): - mock_initialize.return_value = MagicMock(path=tmp_path) - coordinator = SubprocessCoordinator(dag_bundle_name="artifacts") + def test_named_bundle_mode_resolves_bundle(self, mock_initialize, tmp_path): + resolved = MagicMock(supports_versioning=False, path=tmp_path) + mock_initialize.return_value = resolved + coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=[]) + coordinator.dag_bundle_name = "artifacts" coordinator._artifact_source = _ArtifactSource.NAMED_BUNDLE - roots = coordinator._init_root_source() + roots, bundle = coordinator._init_root_source(log) assert roots == [tmp_path] + assert bundle is resolved mock_initialize.assert_called_once_with(BundleInfo(name="artifacts"), log) @patch("airflow.sdk.execution_time.task_runner.initialize_ti_bundle") def test_missing_resolved_path_raises(self, mock_initialize, tmp_path): missing = tmp_path / "nope" - mock_initialize.return_value = MagicMock(path=missing) - coordinator = SubprocessCoordinator(dag_bundle_name="artifacts") + mock_initialize.return_value = MagicMock(supports_versioning=False, path=missing) + coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=[]) + coordinator.dag_bundle_name = "artifacts" coordinator._artifact_source = _ArtifactSource.NAMED_BUNDLE with pytest.raises(FileNotFoundError, match="does not exist"): - coordinator._init_root_source() + coordinator._init_root_source(log) + + +class TestExecuteTaskBundleWiring: + """execute_task binds the task bundle, forwards resolved roots, holds the version lock, uses the task log.""" + + @patch("airflow.dag_processing.bundles.base.BundleVersionLock") + @patch("airflow.sdk.execution_time.task_runner.initialize_ti_bundle") + @patch.object(_PopenActivitySubprocess, "start") + def test_task_bundle_mode_binds_forwards_roots_and_locks( + self, mock_start, mock_initialize, mock_lock, mock_client, tmp_path + ): + resolved = MagicMock(supports_versioning=False, path=tmp_path, version="v9") + resolved.name = "dags" + mock_initialize.return_value = resolved + mock_start.return_value.wait.return_value = 0 + task_logger = MagicMock() + + coordinator = _StubSubprocessCoordinator(command=["/runtime"], explicit_roots=[]) + bundle_info = BundleInfo(name="dags", version="v9") + + coordinator.execute_task( + what=_make_ti(), + dag_rel_path="dag.py", + bundle_info=bundle_info, + client=mock_client, + logger=task_logger, + subprocess_logs_to_stdout=False, + ) + + # Resolution runs against the bound task bundle, using the task logger (not the module log). + mock_initialize.assert_called_once_with(bundle_info, task_logger) + # The resolved root is forwarded to the subclass command builder. + assert coordinator.recorded_roots == [[tmp_path]] + # The version lock is held around the subprocess for the resolved bundle. + mock_lock.assert_called_once_with(bundle_name="dags", bundle_version="v9") + mock_lock.return_value.__enter__.assert_called_once() + mock_lock.return_value.__exit__.assert_called_once() + + @patch("airflow.dag_processing.bundles.base.BundleVersionLock") + @patch.object(_PopenActivitySubprocess, "start") + def test_explicit_root_mode_forwards_roots_without_locking( + self, mock_start, mock_lock, mock_client, tmp_path + ): + mock_start.return_value.wait.return_value = 0 + coordinator = _StubSubprocessCoordinator(command=["/runtime"], explicit_roots=[tmp_path]) + + coordinator.execute_task( + what=_make_ti(), + dag_rel_path="dag.py", + bundle_info=MagicMock(), + client=mock_client, + subprocess_logs_to_stdout=False, + ) + + assert coordinator.recorded_roots == [[tmp_path]] + mock_lock.assert_not_called() class TestSetCurrentBundle: From e6c1deca9d9bc1f4f27102648138a6b5fd593c2b Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Mon, 3 Aug 2026 06:39:34 +0000 Subject: [PATCH 04/11] Make JavaCoordinator main_class optional in co-located mode When a subprocess coordinator locates artifacts through a Dag bundle rather than an explicit root, JavaCoordinator previously rejected a missing main_class outright. That made the co-located default unusable without extra configuration and left the entrypoint no more determinate than an auto-detected scan would be. Auto-detect over the resolved bundle now handles that case, with the remaining cross-bundle ambiguity deferred to an IMPORT_ERROR-stage check under AIP-85. A named Dag bundle resolves to its latest version, so its version was unset and the version lock held around the subprocess silently did nothing, letting bundle cleanup delete artifacts a running task was still reading. The concrete current version is now pinned before locking. --- .../language-sdks/java.rst | 8 ++--- .../airflow/sdk/coordinators/_subprocess.py | 19 ++++++++---- .../sdk/coordinators/java/coordinator.py | 15 +++------ .../coordinators/java/test_coordinator.py | 13 +++----- .../task_sdk/coordinators/test_subprocess.py | 31 +++++++++++++++++++ 5 files changed, 58 insertions(+), 28 deletions(-) diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst index fa27d715924d0..5d38a29262827 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst @@ -1038,10 +1038,10 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the - Extra JVM arguments such as ``["-Xmx1g", "-Dsome.property=value"]``. * - ``main_class`` - *(auto-detect)* - - Explicit entry-point class. If omitted, - :class:`~airflow.sdk.coordinators.java.JavaCoordinator` scans ``jars_root`` for a - JAR whose manifest sets ``Main-Class``. If multiple executable JARs are found the - result is non-deterministic; set ``main_class`` explicitly in that case. + - Explicit entry-point class. If omitted, the coordinator scans for a JAR whose + manifest sets ``Main-Class`` — in ``jars_root`` when set, otherwise across the + resolved Dag bundle. If multiple executable JARs match the result is + non-deterministic; set ``main_class`` explicitly in that case. * - ``task_startup_timeout`` - ``10.0`` - Seconds to wait for the JVM subprocess to connect after launch. Increase this if your diff --git a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py index 6da048fe156f7..9baf491b77048 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py +++ b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py @@ -431,12 +431,8 @@ def _explicit_artifact_roots(self) -> Sequence[pathlib.Path]: """Subclass's explicit artifact roots; empty (the default) selects task-bundle mode.""" return [] - def _validate_artifact_source(self) -> None: - """Run subclass-specific validation after classification; default no-op.""" - def __attrs_post_init__(self) -> None: self._classify_artifact_source(self._explicit_artifact_roots, root_kwarg=self._root_kwarg) - self._validate_artifact_source() def _classify_artifact_source(self, configured: Sequence[pathlib.Path], *, root_kwarg: str) -> None: """ @@ -566,12 +562,23 @@ def execute_task( # Hold the version lock across start()/wait() so bundle cleanup # cannot rmtree a version this task is still reading from, # mirroring task_runner.main() for the Python task path. - from airflow.dag_processing.bundles.base import BundleVersionLock # noqa: SDK002 + from airflow.dag_processing.bundles.base import ( # noqa: SDK002 + BundleVersionLock, + unpack_bundle_version, + ) + # NAMED_BUNDLE resolves "latest" with no pinned version, so + # resolved_bundle.version is None and the lock would be a silent + # no-op. Pin the concrete current version so the lock protects the read. + lock_version = resolved_bundle.version + if lock_version is None: + lock_version, _ = unpack_bundle_version( + resolved_bundle.get_current_version(), resolved_bundle + ) stack.enter_context( BundleVersionLock( bundle_name=resolved_bundle.name, - bundle_version=resolved_bundle.version, + bundle_version=lock_version, ) ) command, subprocess_schema_version = self._build_execute_task_command(what=what, roots=roots) diff --git a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py index d4f44ffbfcc11..3127ca7e6d73f 100644 --- a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py @@ -210,19 +210,14 @@ class JavaCoordinator(SubprocessCoordinator): def _explicit_artifact_roots(self) -> list[pathlib.Path]: return self.jars_root - def _validate_artifact_source(self) -> None: - # Without an explicit root the scan falls back to a whole Dag bundle, where - # _JarInfo.find would pick the first executable JAR in walk order — a - # non-deterministic entrypoint. Require an explicit main_class in that case. - if not self.jars_root and not self.main_class: - raise ValueError( - "JavaCoordinator requires 'main_class' when 'jars_root' is not set, so the " - "entrypoint is unambiguous when scanning a Dag bundle for JARs." - ) - def _build_execute_task_command( self, *, what: TaskInstance, roots: list[pathlib.Path] ) -> tuple[list[str], str | None]: + # TODO: Scanning a whole Dag bundle without an explicit main_class lets + # _JarInfo.find pick the first executable JAR in walk order, so duplicate + # entrypoints across bundles resolve non-deterministically — the same + # duplicate-dag_id ambiguity the Python Dag path has. Reject it at the + # IMPORT_ERROR stage once AIP-85 exposes an interface to raise there. jar = _JarInfo.find(roots, self.main_class) command = [ self.java_executable, diff --git a/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py index 2ce073f147c69..ed5f26a995f8e 100644 --- a/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py @@ -238,15 +238,12 @@ def test_custom_kwargs(self): assert coordinator.jars_root == [pathlib.Path("/airflow/java-bundles")] def test_jars_root_optional_defaults_to_empty(self): - # Neither an explicit root nor dag_bundle_name: co-located mode, valid - # once an explicit main_class disambiguates the entrypoint. - coordinator = JavaCoordinator(main_class="com.example.Main") + # Neither an explicit root nor dag_bundle_name: co-located mode. main_class + # is optional here — the entrypoint is auto-detected from the bundle scan. + coordinator = JavaCoordinator() assert coordinator.jars_root == [] assert coordinator.dag_bundle_name is None - - def test_colocated_mode_requires_main_class(self): - with pytest.raises(ValueError, match="requires 'main_class' when 'jars_root' is not set"): - JavaCoordinator() + assert coordinator.main_class == "" def test_explicit_root_does_not_require_main_class(self, tmp_path): coordinator = JavaCoordinator(jars_root=tmp_path) @@ -260,7 +257,7 @@ def test_root_and_dag_bundle_name_are_mutually_exclusive(self): def test_unconfigured_dag_bundle_name_raises(self, mock_manager): mock_manager.is_bundle_configured.return_value = False with pytest.raises(ValueError, match="unconfigured Dag bundle 'ghost'"): - JavaCoordinator(main_class="com.example.Main", dag_bundle_name="ghost") + JavaCoordinator(dag_bundle_name="ghost") def test_build_command_scans_passed_roots_in_colocated_mode(self, tmp_path): # Co-located mode: no configured root, so the subclass must scan the roots diff --git a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py index bb2d1265e154b..b3e6130d261ab 100644 --- a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py +++ b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py @@ -1029,6 +1029,37 @@ def test_explicit_root_mode_forwards_roots_without_locking( assert coordinator.recorded_roots == [[tmp_path]] mock_lock.assert_not_called() + @patch("airflow.dag_processing.bundles.base.BundleVersionLock") + @patch("airflow.sdk.execution_time.task_runner.initialize_ti_bundle") + @patch.object(_PopenActivitySubprocess, "start") + def test_named_bundle_mode_locks_resolved_current_version( + self, mock_start, mock_initialize, mock_lock, mock_client, tmp_path + ): + # NAMED_BUNDLE resolves "latest" with version=None; the lock must pin the + # concrete current version so it does not silently become a no-op. + from airflow.dag_processing.bundles.base import BundleVersion + + resolved = MagicMock(path=tmp_path, version=None) + resolved.name = "artifacts" + resolved.get_current_version.return_value = BundleVersion(version="sha-abc", data=None) + mock_initialize.return_value = resolved + mock_start.return_value.wait.return_value = 0 + + coordinator = _StubSubprocessCoordinator(command=["/runtime"], explicit_roots=[]) + coordinator.dag_bundle_name = "artifacts" + coordinator._artifact_source = _ArtifactSource.NAMED_BUNDLE + + coordinator.execute_task( + what=_make_ti(), + dag_rel_path="dag.py", + bundle_info=MagicMock(), + client=mock_client, + subprocess_logs_to_stdout=False, + ) + + resolved.get_current_version.assert_called_once_with() + mock_lock.assert_called_once_with(bundle_name="artifacts", bundle_version="sha-abc") + class TestSetCurrentBundle: def test_binds_for_the_block_and_clears_on_exit_allowing_reuse(self): From 8c674c638d5cd3cb8e798df0329d433eab98a3c8 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Mon, 10 Aug 2026 01:59:52 +0000 Subject: [PATCH 05/11] Resolve coordinator Dag bundles without importing the task runner A language-SDK task start needs a bundle materialized on disk before the subprocess can be launched, but reaching that through the task runner made every such start pay for the operator tree, pydantic, and OTel plumbing that the supervisor side never uses. The bundle lifecycle now lives on its own, so the coordinator depends on the Dag-bundle machinery alone. The bundle-name check a coordinator runs at construction and the class-importing parse the manager runs at startup previously read the configuration independently, so the two could disagree about what is configured. Keying the parsed configuration on the configuration itself lets both share one snapshot while a configuration change is still picked up. Subclasses had to declare the name of their explicit-root kwarg separately from its value, which is two things to keep in step for one piece of information. --- .../language-sdks/go.rst | 5 +- .../language-sdks/java.rst | 5 +- .../language-sdks/typescript.rst | 5 +- .../airflow/dag_processing/bundles/manager.py | 61 +++++- .../tests/integration/otel/test_otel.py | 2 +- .../bundles/test_dag_bundle_manager.py | 49 ++++- generated/known_airflow_exceptions.txt | 2 +- .../airflow/sdk/coordinators/_subprocess.py | 123 +++++------ .../coordinators/executable/coordinator.py | 7 +- .../sdk/coordinators/java/coordinator.py | 14 +- .../sdk/coordinators/node/coordinator.py | 7 +- .../src/airflow/sdk/execution_time/bundles.py | 92 +++++++++ .../airflow/sdk/execution_time/task_runner.py | 97 +-------- .../src/airflow/sdk/execution_time/tracing.py | 60 ++++++ .../executable/test_coordinator.py | 7 +- .../coordinators/java/test_coordinator.py | 7 +- .../coordinators/node/test_coordinator.py | 5 +- .../task_sdk/coordinators/test_subprocess.py | 104 ++++++---- .../task_sdk/execution_time/test_bundles.py | 86 ++++++++ .../execution_time/test_task_runner.py | 195 +----------------- .../task_sdk/execution_time/test_tracing.py | 144 +++++++++++++ 21 files changed, 630 insertions(+), 447 deletions(-) create mode 100644 task-sdk/src/airflow/sdk/execution_time/bundles.py create mode 100644 task-sdk/src/airflow/sdk/execution_time/tracing.py create mode 100644 task-sdk/tests/task_sdk/execution_time/test_bundles.py create mode 100644 task-sdk/tests/task_sdk/execution_time/test_tracing.py diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst index 2c3bb1c58b0ef..965fabbc292f6 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst @@ -526,8 +526,9 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the and both are optional: * Set ``executables_root`` to scan explicit filesystem directories you manage yourself. - * Set ``dag_bundle_name`` to load bundles from a configured Dag bundle (its latest version), - so they are delivered and versioned through the same bundle machinery as your Dags. + * Set ``dag_bundle_name`` to load bundles from a configured Dag bundle, so they are delivered + and versioned through the same bundle machinery as your Dags. The task uses the version that + bundle is on when it starts, pinned for the whole task. * Leave both unset (the default) to load bundles from the **task's own** Dag bundle, pinned to the version the run was created with. diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst index 5d38a29262827..2587bd7950a68 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst @@ -1053,8 +1053,9 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the are optional: * Set ``jars_root`` to scan explicit filesystem directories you manage yourself. - * Set ``dag_bundle_name`` to load JARs from a configured Dag bundle (its latest version), so - they are delivered and versioned through the same bundle machinery as your Dags. + * Set ``dag_bundle_name`` to load JARs from a configured Dag bundle, so they are delivered and + versioned through the same bundle machinery as your Dags. The task uses the version that bundle + is on when it starts, pinned for the whole task. * Leave both unset (the default) to load JARs from the **task's own** Dag bundle, pinned to the version the run was created with. diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst index d59893b3359e6..376771ff4cddd 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst @@ -468,8 +468,9 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the are optional: * Set ``bundles_root`` to scan explicit filesystem directories you manage yourself. - * Set ``dag_bundle_name`` to load the bundle from a configured Dag bundle (its latest version), so it - is delivered and versioned through the same bundle machinery as your Dags. + * Set ``dag_bundle_name`` to load the bundle from a configured Dag bundle, so it is delivered + and versioned through the same bundle machinery as your Dags. The task uses the version that + bundle is on when it starts, pinned for the whole task. * Leave both unset (the default) to load the bundle from the **task's own** Dag bundle, pinned to the version the run was created with. diff --git a/airflow-core/src/airflow/dag_processing/bundles/manager.py b/airflow-core/src/airflow/dag_processing/bundles/manager.py index 4021ee79d6538..e19d702d71fca 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/manager.py +++ b/airflow-core/src/airflow/dag_processing/bundles/manager.py @@ -18,12 +18,13 @@ import functools import importlib +import json import logging import os import warnings from collections import defaultdict from pathlib import Path -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, NamedTuple, cast from itsdangerous import URLSafeSerializer from pydantic import BaseModel, ValidationError @@ -108,18 +109,54 @@ def _bundle_item_exc(msg): ) +class _BundleConfigSnapshot(NamedTuple): + """The configured Dag bundles, as names only and as full configs.""" + + configs: tuple[_ExternalBundleConfig, ...] + names: frozenset[str] + + +_EMPTY_BUNDLE_CONFIG_SNAPSHOT = _BundleConfigSnapshot(configs=(), names=frozenset()) + + @functools.cache -def _configured_bundle_names() -> frozenset[str]: - """Return configured Dag bundle names, parsed by name only (no class import).""" +def _parse_bundle_config_snapshot(config_json: str, load_examples: bool) -> _BundleConfigSnapshot: + """ + Build the snapshot for one configuration, without importing any bundle class. + + Keyed on the configuration rather than cached outright, so every reader of the + same configuration shares one snapshot while a configuration change is still + picked up. + """ + bundle_config_list = _parse_bundle_config(json.loads(config_json)) + if load_examples: + _add_example_dag_bundle(bundle_config_list) + _add_provider_example_dags_to_bundle(bundle_config_list) + return _BundleConfigSnapshot( + configs=tuple(bundle_config_list), + names=frozenset(cfg.name for cfg in bundle_config_list), + ) + + +def _load_bundle_config_snapshot() -> _BundleConfigSnapshot: + """ + Read and validate the configured Dag bundles. + + Shared by :meth:`DagBundlesManager.parse_config`, which goes on to import each + bundle class, and by :meth:`DagBundlesManager.is_bundle_configured`, which only + needs the names, so both resolve to the same snapshot. + """ config_list = conf.getjson("dag_processor", "dag_bundle_config_list") if not config_list: - return frozenset() + return _EMPTY_BUNDLE_CONFIG_SNAPSHOT if not isinstance(config_list, list): raise AirflowConfigException( "Section `dag_processor` key `dag_bundle_config_list` " f"must be list but got {config_list.__class__}" ) - return frozenset(cfg.name for cfg in _parse_bundle_config(config_list)) + return _parse_bundle_config_snapshot( + json.dumps(config_list, sort_keys=True), conf.getboolean("core", "LOAD_EXAMPLES") + ) def _parse_bundle_config(config_list) -> list[_ExternalBundleConfig]: @@ -315,12 +352,9 @@ def parse_config(self) -> None: if self._bundle_config: return - bundle_config_list = _read_bundle_config_list() + bundle_config_list = _load_bundle_config_snapshot().configs if not bundle_config_list: return - if conf.getboolean("core", "LOAD_EXAMPLES"): - _add_example_dag_bundle(bundle_config_list) - _add_provider_example_dags_to_bundle(bundle_config_list) for bundle_config in bundle_config_list: if bundle_config.team_name and not conf.getboolean("core", "multi_team"): @@ -690,8 +724,13 @@ def get_bundle( @classmethod def is_bundle_configured(cls, name: str) -> bool: - """Return whether *name* is a configured Dag bundle, by config name only (no class import).""" - return name in _configured_bundle_names() + """ + Return whether *name* is a configured Dag bundle. + + Deliberately reads configured names only: a caller validating a bundle name + must not depend on every *other* configured bundle's class being importable. + """ + return name in _load_bundle_config_snapshot().names def get_all_dag_bundles(self) -> Iterable[BaseDagBundle]: """ diff --git a/airflow-core/tests/integration/otel/test_otel.py b/airflow-core/tests/integration/otel/test_otel.py index 39d7ee6bdab61..5be887bbdfa7d 100644 --- a/airflow-core/tests/integration/otel/test_otel.py +++ b/airflow-core/tests/integration/otel/test_otel.py @@ -330,7 +330,7 @@ def test_export_metrics_during_process_shutdown(self, capfd): # Additional detail spans are deferred to follow-up PRs; tracked # at https://linear.app/astronomer/issue/ACD-157. { - "_verify_bundle_access": "parse", + "verify_bundle_access": "parse", "parse": "startup", "get_template_context": "startup", "startup": "worker.task1", diff --git a/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py b/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py index d4fa9f3209262..f080fe137c6ce 100644 --- a/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py +++ b/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py @@ -31,9 +31,10 @@ from airflow.dag_processing.bundles.base import BaseDagBundle from airflow.dag_processing.bundles.manager import ( DagBundlesManager, - _configured_bundle_names, _get_configured_bundle_team_names, _guess_best_bundle_for_fileloc, + _load_bundle_config_snapshot, + _parse_bundle_config_snapshot, ) from airflow.exceptions import AirflowConfigException from airflow.models.dag import DagModel @@ -207,19 +208,22 @@ def test_get_bundle(): @pytest.fixture(autouse=True) -def _clear_configured_bundle_names_cache(): - # _configured_bundle_names is a process-global cache; clear it so tests that - # set different bundle configs don't observe each other's cached result. - _configured_bundle_names.cache_clear() +def _clear_bundle_config_snapshot_cache(): + _parse_bundle_config_snapshot.cache_clear() yield - _configured_bundle_names.cache_clear() + _parse_bundle_config_snapshot.cache_clear() + + +def _bundle_config_env(config) -> dict[str, str]: + return { + "AIRFLOW__CORE__LOAD_EXAMPLES": "False", + "AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(config), + } def test_is_bundle_configured(): """is_bundle_configured reports membership without constructing the bundle.""" - with patch.dict( - os.environ, {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(BASIC_BUNDLE_CONFIG)} - ): + with patch.dict(os.environ, _bundle_config_env(BASIC_BUNDLE_CONFIG)): assert DagBundlesManager.is_bundle_configured("my-test-bundle") is True assert DagBundlesManager.is_bundle_configured("bundle-that-doesn't-exist") is False @@ -231,15 +235,38 @@ def test_is_bundle_configured_does_not_import_bundle_class(): so an unimportable classpath must not make a coordinator fail at construction. """ config = [{"name": "artifacts", "classpath": "does.not.exist.Bundle", "kwargs": {}}] - with patch.dict(os.environ, {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(config)}): + with patch.dict(os.environ, _bundle_config_env(config)): assert DagBundlesManager.is_bundle_configured("artifacts") is True def test_is_bundle_configured_empty_config(): - with patch.dict(os.environ, {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": "[]"}): + with patch.dict(os.environ, _bundle_config_env([])): assert DagBundlesManager.is_bundle_configured("anything") is False +def test_is_bundle_configured_sees_implicitly_added_bundles(): + """Bundles parse_config adds implicitly are configured for the name check too.""" + with patch.dict( + os.environ, + { + "AIRFLOW__CORE__LOAD_EXAMPLES": "True", + "AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(BASIC_BUNDLE_CONFIG), + }, + ): + assert DagBundlesManager.is_bundle_configured("example_dags") is True + + +def test_bundle_config_snapshot_is_shared_and_config_scoped(): + """Readers of one configuration share a snapshot; a config change builds a new one.""" + with patch.dict(os.environ, _bundle_config_env(BASIC_BUNDLE_CONFIG)): + first = _load_bundle_config_snapshot() + assert _load_bundle_config_snapshot() is first + + other_config = [{"name": "other", "classpath": BASIC_BUNDLE_CONFIG[0]["classpath"], "kwargs": {}}] + with patch.dict(os.environ, _bundle_config_env(other_config)): + assert _load_bundle_config_snapshot().names == {"other"} + + @pytest.fixture def clear_db(): clear_db_dag_bundles() diff --git a/generated/known_airflow_exceptions.txt b/generated/known_airflow_exceptions.txt index cb860b5cd71a4..ca2b46365cbb1 100644 --- a/generated/known_airflow_exceptions.txt +++ b/generated/known_airflow_exceptions.txt @@ -407,6 +407,6 @@ task-sdk/src/airflow/sdk/definitions/_internal/setup_teardown.py::1 task-sdk/src/airflow/sdk/definitions/connection.py::4 task-sdk/src/airflow/sdk/definitions/decorators/setup_teardown.py::4 task-sdk/src/airflow/sdk/definitions/xcom_arg.py::1 -task-sdk/src/airflow/sdk/execution_time/task_runner.py::1 +task-sdk/src/airflow/sdk/execution_time/bundles.py::1 task-sdk/tests/task_sdk/bases/test_sensor.py::1 task-sdk/tests/task_sdk/execution_time/test_task_runner.py::2 diff --git a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py index 9baf491b77048..b34aaf77fb1db 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py +++ b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py @@ -37,13 +37,17 @@ import socket import subprocess import time -from typing import TYPE_CHECKING, ClassVar, TypeVar, cast +from typing import TYPE_CHECKING, TypeVar, cast import attrs import psutil import structlog +from airflow.dag_processing.bundles.base import BundleVersionLock, unpack_bundle_version # noqa: SDK002 +from airflow.dag_processing.bundles.manager import DagBundlesManager # noqa: SDK002 +from airflow.sdk.api.datamodels._generated import BundleInfo from airflow.sdk.configuration import conf +from airflow.sdk.execution_time.bundles import initialize_ti_bundle from airflow.sdk.execution_time.coordinator import BaseCoordinator from airflow.sdk.execution_time.supervisor import ActivitySubprocess, NeverRaised, ProcessTracker @@ -56,7 +60,7 @@ from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: SDK002 from airflow.sdk.api.client import Client - from airflow.sdk.api.datamodels._generated import BundleInfo, TaskInstance + from airflow.sdk.api.datamodels._generated import TaskInstance Tracked = TypeVar("Tracked", socket.socket, subprocess.Popen) @@ -375,15 +379,43 @@ def wait(self) -> int: return code +def _initialize_pinned_bundle(target: BundleInfo, logger: FilteringBoundLogger) -> BaseDagBundle: + """ + Materialize *target* at a concrete version, so the tree handed to the subprocess is lockable. + + A bundle resolved without a version points at the bundle's shared, mutable + checkout: another task refreshing the same bundle resets it underneath a running + subprocess, and ``BundleVersionLock`` cannot protect it because a version-less + lock is a no-op. Re-resolving at the version current now yields a private + ``versions/`` tree that the lock does cover. + + Bundles that do not track versions have nothing to pin and keep their single path. + """ + bundle = initialize_ti_bundle(target) + if bundle.version is not None: + return bundle + + version, version_data = unpack_bundle_version(bundle.get_current_version(), bundle) + if version is None: + return bundle + logger.debug("Pinning Dag bundle to its current version", bundle=target.name, version=version) + return initialize_ti_bundle(BundleInfo(name=target.name, version=version, version_data=version_data)) + + class _ArtifactSource(enum.Enum): """How a subprocess coordinator locates the compiled task artifacts.""" EXPLICIT_ROOT = enum.auto() """An explicit filesystem root (``jars_root`` / ``executables_root`` / ``bundles_root``).""" NAMED_BUNDLE = enum.auto() - """``dag_bundle_name`` names a configured Dag bundle; its latest version is used.""" + """``dag_bundle_name`` names a configured Dag bundle; its version current at task start is used.""" TASK_BUNDLE = enum.auto() - """Neither is set: the task's own (co-located) bundle, pinned to the run's version.""" + """ + Neither is set: artifacts are *co-located* with the Python stub Dag. + + The task's own bundle is scanned, so the compiled artifacts ship in the same + bundle, at the same version, as the Dag that delegates to them. + """ @attrs.define(kw_only=True) @@ -403,46 +435,42 @@ class SubprocessCoordinator(BaseCoordinator): :param dag_bundle_name: Locate artifacts through a configured Dag bundle rather than an explicit root. Mutually exclusive with the subclass's explicit root; if neither is set, the task's own bundle is used. A named bundle resolves to - its latest version; the task's own bundle is pinned to the run's version. + the version current when the task starts; the task's own bundle uses the + run's version. Either way the resolved version is pinned for the whole task. """ task_startup_timeout: float = 10.0 dag_bundle_name: str | None = None - # Name of the subclass's explicit-root kwarg, used only in error messages. - # Subclasses that expose an explicit root set this and override - # :meth:`_explicit_artifact_roots`; a subclass that does neither lands on the - # task-bundle default instead of failing at execute time. - _root_kwarg: ClassVar[str] = "root" - - # Classified once at construction by :meth:`_classify_artifact_source` and - # dispatched on by :meth:`_init_root_source` at execute time. - _artifact_source: _ArtifactSource | None = attrs.field(init=False, default=None) + _artifact_source: _ArtifactSource = attrs.field(init=False) # The subclass's explicit root, recorded at construction so the base can # resolve roots without knowing the subclass field name. _configured_roots: list[pathlib.Path] = attrs.field(init=False, factory=list) - # The task's own bundle, bound for the duration of a single :meth:`execute_task` - # call by :meth:`_set_current_bundle` so :meth:`_init_root_source` can resolve - # co-located artifacts. _active_bundle_info: BundleInfo | None = attrs.field(init=False, default=None) @property - def _explicit_artifact_roots(self) -> Sequence[pathlib.Path]: - """Subclass's explicit artifact roots; empty (the default) selects task-bundle mode.""" - return [] + def _explicit_artifact_roots(self) -> tuple[str, Sequence[pathlib.Path]]: + """ + The subclass's explicit-root kwarg name and its configured value. + + The name is only used in error messages. An empty value — the default, for a + subclass that does not override this — selects task-bundle mode rather than + failing at execute time. + """ + return "root", () def __attrs_post_init__(self) -> None: - self._classify_artifact_source(self._explicit_artifact_roots, root_kwarg=self._root_kwarg) + self._classify_artifact_source() - def _classify_artifact_source(self, configured: Sequence[pathlib.Path], *, root_kwarg: str) -> None: + def _classify_artifact_source(self) -> None: """ Classify and validate how this coordinator locates artifacts (construction time). - Called from :meth:`__attrs_post_init__` with the subclass's explicit - roots. It rejects setting both an explicit root and ``dag_bundle_name``, - fails fast when ``dag_bundle_name`` names a bundle that is not configured, - and records the resulting :class:`_ArtifactSource` and explicit root. + Rejects setting both an explicit root and ``dag_bundle_name``, fails fast + when ``dag_bundle_name`` names a bundle that is not configured, and records + the resulting :class:`_ArtifactSource` and explicit root. """ + root_kwarg, configured = self._explicit_artifact_roots if configured and self.dag_bundle_name is not None: raise ValueError( f"Set at most one of {root_kwarg!r} or 'dag_bundle_name': {root_kwarg!r} for an " @@ -454,8 +482,6 @@ def _classify_artifact_source(self, configured: Sequence[pathlib.Path], *, root_ self._configured_roots = list(configured) elif self.dag_bundle_name is not None: source = _ArtifactSource.NAMED_BUNDLE - from airflow.dag_processing.bundles.manager import DagBundlesManager # noqa: SDK002 - if not DagBundlesManager.is_bundle_configured(self.dag_bundle_name): raise ValueError( f"Coordinator 'dag_bundle_name' references unconfigured Dag bundle " @@ -465,12 +491,12 @@ def _classify_artifact_source(self, configured: Sequence[pathlib.Path], *, root_ source = _ArtifactSource.TASK_BUNDLE self._artifact_source = source - details: dict[str, str | list[str]] = {"mode": source.name} - if self.dag_bundle_name is not None: - details["dag_bundle_name"] = self.dag_bundle_name - if self._configured_roots: - details["configured_roots"] = [str(root) for root in self._configured_roots] - log.debug("Coordinator artifact source selected", **details) + log.debug( + "Coordinator artifact source selected", + mode=source.name, + dag_bundle_name=self.dag_bundle_name, + configured_roots=[str(root) for root in self._configured_roots], + ) def _init_root_source( self, logger: FilteringBoundLogger @@ -486,25 +512,15 @@ def _init_root_source( if self._artifact_source is _ArtifactSource.EXPLICIT_ROOT: return self._configured_roots, None - from airflow.sdk.api.datamodels._generated import BundleInfo - if self._artifact_source is _ArtifactSource.NAMED_BUNDLE: # NAMED_BUNDLE implies dag_bundle_name is set. target = BundleInfo(name=cast("str", self.dag_bundle_name)) - elif self._artifact_source is _ArtifactSource.TASK_BUNDLE: + else: if self._active_bundle_info is None: raise RuntimeError("_init_root_source requires an active task; call it during execute_task.") target = self._active_bundle_info - else: - raise RuntimeError( - "Coordinator artifact source was not classified; call _classify_artifact_source first." - ) - - # Lazy import: task_runner is a heavy module and importing it at module - # load would risk an import cycle through the supervisor. - from airflow.sdk.execution_time.task_runner import initialize_ti_bundle - bundle = initialize_ti_bundle(target, logger) + bundle = _initialize_pinned_bundle(target, logger) path = bundle.path if not path.exists(): raise FileNotFoundError(f"Dag bundle {target.name!r} resolved to {path}, which does not exist.") @@ -562,23 +578,10 @@ def execute_task( # Hold the version lock across start()/wait() so bundle cleanup # cannot rmtree a version this task is still reading from, # mirroring task_runner.main() for the Python task path. - from airflow.dag_processing.bundles.base import ( # noqa: SDK002 - BundleVersionLock, - unpack_bundle_version, - ) - - # NAMED_BUNDLE resolves "latest" with no pinned version, so - # resolved_bundle.version is None and the lock would be a silent - # no-op. Pin the concrete current version so the lock protects the read. - lock_version = resolved_bundle.version - if lock_version is None: - lock_version, _ = unpack_bundle_version( - resolved_bundle.get_current_version(), resolved_bundle - ) stack.enter_context( BundleVersionLock( bundle_name=resolved_bundle.name, - bundle_version=lock_version, + bundle_version=resolved_bundle.version, ) ) command, subprocess_schema_version = self._build_execute_task_command(what=what, roots=roots) diff --git a/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py index bc08ea6d6b7cc..16d57f39b5832 100644 --- a/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py @@ -25,7 +25,7 @@ import stat import struct from collections import OrderedDict -from typing import TYPE_CHECKING, Any, BinaryIO, ClassVar +from typing import TYPE_CHECKING, Any, BinaryIO import attrs import structlog @@ -349,11 +349,10 @@ class ExecutableCoordinator(SubprocessCoordinator): converter=convert_roots, factory=list, ) - _root_kwarg: ClassVar[str] = "executables_root" @property - def _explicit_artifact_roots(self) -> list[pathlib.Path]: - return self.executables_root + def _explicit_artifact_roots(self) -> tuple[str, list[pathlib.Path]]: + return "executables_root", self.executables_root def _build_execute_task_command( self, *, what: TaskInstance, roots: list[pathlib.Path] diff --git a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py index 3127ca7e6d73f..299ea5e07d826 100644 --- a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py @@ -24,7 +24,7 @@ import pathlib import stat import zipfile -from typing import TYPE_CHECKING, ClassVar +from typing import TYPE_CHECKING import attrs import structlog @@ -204,20 +204,16 @@ class JavaCoordinator(SubprocessCoordinator): factory=list, ) main_class: str = "" - _root_kwarg: ClassVar[str] = "jars_root" @property - def _explicit_artifact_roots(self) -> list[pathlib.Path]: - return self.jars_root + def _explicit_artifact_roots(self) -> tuple[str, list[pathlib.Path]]: + return "jars_root", self.jars_root def _build_execute_task_command( self, *, what: TaskInstance, roots: list[pathlib.Path] ) -> tuple[list[str], str | None]: - # TODO: Scanning a whole Dag bundle without an explicit main_class lets - # _JarInfo.find pick the first executable JAR in walk order, so duplicate - # entrypoints across bundles resolve non-deterministically — the same - # duplicate-dag_id ambiguity the Python Dag path has. Reject it at the - # IMPORT_ERROR stage once AIP-85 exposes an interface to raise there. + # Without main_class, the first executable JAR in walk order wins; tracked at + # https://github.com/apache/airflow/issues/71134 jar = _JarInfo.find(roots, self.main_class) command = [ self.java_executable, diff --git a/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py index d983983cd8c3f..f9b23c179ccc3 100644 --- a/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py @@ -21,7 +21,7 @@ import os import pathlib -from typing import TYPE_CHECKING, ClassVar +from typing import TYPE_CHECKING import attrs import structlog @@ -122,11 +122,10 @@ class NodeCoordinator(SubprocessCoordinator): converter=convert_roots, factory=list, ) - _root_kwarg: ClassVar[str] = "bundles_root" @property - def _explicit_artifact_roots(self) -> list[pathlib.Path]: - return self.bundles_root + def _explicit_artifact_roots(self) -> tuple[str, list[pathlib.Path]]: + return "bundles_root", self.bundles_root def _build_execute_task_command( self, *, what: TaskInstance, roots: list[pathlib.Path] diff --git a/task-sdk/src/airflow/sdk/execution_time/bundles.py b/task-sdk/src/airflow/sdk/execution_time/bundles.py new file mode 100644 index 0000000000000..cfc1fe66cfb53 --- /dev/null +++ b/task-sdk/src/airflow/sdk/execution_time/bundles.py @@ -0,0 +1,92 @@ +# 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. +""" +Materialize a task instance's Dag bundle on the worker. + +Kept apart from :mod:`airflow.sdk.execution_time.task_runner` so the supervisor +side — which needs a bundle on disk before it can launch a language-SDK +subprocess — does not have to import the task runner and everything it pulls in. +""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING + +from airflow.dag_processing.bundles.manager import DagBundlesManager # noqa: SDK002 +from airflow.sdk.execution_time.tracing import detail_span + +if TYPE_CHECKING: + from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: SDK002 + from airflow.sdk.api.datamodels._generated import BundleInfo + +__all__ = ["initialize_ti_bundle", "verify_bundle_access"] + + +def initialize_ti_bundle(bundle_info: BundleInfo) -> BaseDagBundle: + """ + Resolve, initialize, and access-check the Dag bundle for a task instance. + + Shared by :func:`~airflow.sdk.execution_time.task_runner.parse` (Python task + path) and the subprocess coordinators (language-SDK path), which both need a + task instance's bundle materialized on disk before use. Returns the + initialized bundle so callers can read ``bundle.path``. + """ + bundle_instance = DagBundlesManager().get_bundle( + name=bundle_info.name, + version=bundle_info.version, + version_data=bundle_info.version_data, + ) + bundle_instance.initialize() + verify_bundle_access(bundle_instance) + return bundle_instance + + +@detail_span("verify_bundle_access") +def verify_bundle_access(bundle_instance: BaseDagBundle) -> None: + """ + Verify bundle is accessible by the current user. + + This is called after user impersonation (if any) to ensure the bundle + is actually accessible. Uses os.access() which works with any permission + scheme (standard Unix permissions, ACLs, SELinux, etc.). + + :param bundle_instance: The bundle instance to check + :raises AirflowException: if bundle is not accessible + """ + from getpass import getuser + + from airflow.sdk.exceptions import AirflowException + + bundle_path = bundle_instance.path + + if not bundle_path.exists(): + # Already handled by initialize() with a warning + return + + # Check read permission (and execute for directories to list contents) + access_mode = os.R_OK + if bundle_path.is_dir(): + access_mode |= os.X_OK + + if not os.access(bundle_path, access_mode): + raise AirflowException( + f"Bundle '{bundle_instance.name}' path '{bundle_path}' is not accessible " + f"by user '{getuser()}'. When using run_as_user, ensure bundle directories " + f"are readable by the impersonated user. " + f"See: https://airflow.apache.org/docs/apache-airflow/stable/administration-and-deployment/dag-bundles.html" + ) 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 415f24b0a2093..fe68a5c19dc58 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -43,10 +43,8 @@ from structlog.contextvars import bind_contextvars from airflow.dag_processing.bundles.base import BaseDagBundle, BundleVersionLock -from airflow.dag_processing.bundles.manager import DagBundlesManager from airflow.sdk._shared.observability.metrics import stats from airflow.sdk._shared.observability.metrics.stats import build_dag_metric_tags -from airflow.sdk._shared.observability.traces import get_task_span_detail_level from airflow.sdk._shared.template_rendering import truncate_rendered_value from airflow.sdk.api.client import get_hostname, getuser from airflow.sdk.api.datamodels._generated import ( @@ -85,6 +83,7 @@ TaskAwaitingInput, TaskDeferred, ) +from airflow.sdk.execution_time.bundles import initialize_ti_bundle from airflow.sdk.execution_time.callback_runner import create_executable_runner from airflow.sdk.execution_time.comms import ( AssetEventDagRunReferenceResult, @@ -147,6 +146,7 @@ _LegacyEmailBackendNotifier, ) from airflow.sdk.execution_time.sentry import Sentry +from airflow.sdk.execution_time.tracing import detail_span from airflow.sdk.execution_time.xcom import XCom from airflow.sdk.listener import get_listener_manager from airflow.sdk.observability.metrics import stats_utils @@ -159,7 +159,6 @@ from pendulum.datetime import DateTime from structlog.typing import FilteringBoundLogger as Logger - from airflow.sdk.api.datamodels._generated import BundleInfo from airflow.sdk.definitions._internal.abstractoperator import AbstractOperator from airflow.sdk.definitions.context import Context from airflow.sdk.definitions.retry_policy import RetryDecision @@ -171,38 +170,6 @@ tracer = trace.get_tracer(__name__) -class detail_span: - """Context manager and decorator that creates a child span when detail level > 1.""" - - def __init__(self, *args, **kwargs): - self._args = args - self._kwargs = kwargs - self._ctx = None - - def _make_ctx(self): - parent_span = trace.get_current_span() - config_level = get_task_span_detail_level(span=parent_span) - if config_level > 1: - return tracer.start_as_current_span(*self._args, **self._kwargs) - return trace.INVALID_SPAN - - def __enter__(self): - self._ctx = self._make_ctx() - return self._ctx.__enter__() - - def __exit__(self, *exc_info): - return self._ctx.__exit__(*exc_info) - - def __call__(self, f): - @functools.wraps(f) - def wrapper(*inner_args, **inner_kwargs): - with self._make_ctx(): - return f(*inner_args, **inner_kwargs) - - wrapper.__signature__ = inspect.signature(f) - return wrapper - - @contextmanager def _make_task_span(msg: StartupDetails): parent_context = ( @@ -1026,25 +993,6 @@ def _register_deserialization_allowed_classes(dag, log: Logger) -> None: ) -def initialize_ti_bundle(bundle_info: BundleInfo, log: Logger) -> BaseDagBundle: - """ - Resolve, initialize, and access-check the Dag bundle for a task instance. - - Shared by :func:`parse` (Python task path) and the subprocess coordinators - (language-SDK path), which both need a task instance's bundle materialized - on disk before use. Returns the initialized bundle so callers can read - ``bundle.path``. - """ - bundle_instance = DagBundlesManager().get_bundle( - name=bundle_info.name, - version=bundle_info.version, - version_data=bundle_info.version_data, - ) - bundle_instance.initialize() - _verify_bundle_access(bundle_instance, log) - return bundle_instance - - @detail_span("parse") def parse(what: StartupDetails, log: Logger) -> RuntimeTaskInstance: # TODO: Task-SDK: @@ -1053,7 +1001,7 @@ def parse(what: StartupDetails, log: Logger) -> RuntimeTaskInstance: bundle_info = what.bundle_info bundle_prepare_start = time.monotonic() - bundle_instance = initialize_ti_bundle(bundle_info, log) + bundle_instance = initialize_ti_bundle(bundle_info) bundle_prepare_ms = int((time.monotonic() - bundle_prepare_start) * 1000) dag_absolute_path = os.fspath(Path(bundle_instance.path, what.dag_rel_path)) @@ -1149,43 +1097,6 @@ def parse(what: StartupDetails, log: Logger) -> RuntimeTaskInstance: # 3. Shutdown and report status -@detail_span("_verify_bundle_access") -def _verify_bundle_access(bundle_instance: BaseDagBundle, log: Logger) -> None: - """ - Verify bundle is accessible by the current user. - - This is called after user impersonation (if any) to ensure the bundle - is actually accessible. Uses os.access() which works with any permission - scheme (standard Unix permissions, ACLs, SELinux, etc.). - - :param bundle_instance: The bundle instance to check - :param log: Logger instance - :raises AirflowException: if bundle is not accessible - """ - from getpass import getuser - - from airflow.sdk.exceptions import AirflowException - - bundle_path = bundle_instance.path - - if not bundle_path.exists(): - # Already handled by initialize() with a warning - return - - # Check read permission (and execute for directories to list contents) - access_mode = os.R_OK - if bundle_path.is_dir(): - access_mode |= os.X_OK - - if not os.access(bundle_path, access_mode): - raise AirflowException( - f"Bundle '{bundle_instance.name}' path '{bundle_path}' is not accessible " - f"by user '{getuser()}'. When using run_as_user, ensure bundle directories " - f"are readable by the impersonated user. " - f"See: https://airflow.apache.org/docs/apache-airflow/stable/administration-and-deployment/dag-bundles.html" - ) - - def get_startup_details() -> StartupDetails: # The parent sends us a StartupDetails message un-prompted. After this, every single message is only sent # in response to us sending a request. @@ -1293,8 +1204,6 @@ def _serialize_template_field( Uses the SDK secrets masker to redact secrets in the serialized output. """ - import inspect - from airflow.sdk._shared.module_loading import qualname from airflow.sdk._shared.secrets_masker import redact diff --git a/task-sdk/src/airflow/sdk/execution_time/tracing.py b/task-sdk/src/airflow/sdk/execution_time/tracing.py new file mode 100644 index 0000000000000..86a28199c61da --- /dev/null +++ b/task-sdk/src/airflow/sdk/execution_time/tracing.py @@ -0,0 +1,60 @@ +# 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. +"""Span helpers shared by the modules that make up a task run.""" + +from __future__ import annotations + +import functools +import inspect + +from opentelemetry import trace + +from airflow.sdk._shared.observability.traces import get_task_span_detail_level + +tracer = trace.get_tracer(__name__) + + +class detail_span: + """Context manager and decorator that creates a child span when detail level > 1.""" + + def __init__(self, *args, **kwargs): + self._args = args + self._kwargs = kwargs + self._ctx = None + + def _make_ctx(self): + parent_span = trace.get_current_span() + config_level = get_task_span_detail_level(span=parent_span) + if config_level > 1: + return tracer.start_as_current_span(*self._args, **self._kwargs) + return trace.INVALID_SPAN + + def __enter__(self): + self._ctx = self._make_ctx() + return self._ctx.__enter__() + + def __exit__(self, *exc_info): + return self._ctx.__exit__(*exc_info) + + def __call__(self, f): + @functools.wraps(f) + def wrapper(*inner_args, **inner_kwargs): + with self._make_ctx(): + return f(*inner_args, **inner_kwargs) + + wrapper.__signature__ = inspect.signature(f) + return wrapper diff --git a/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py index be5b434e2d5d2..d62e420e857e8 100644 --- a/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py @@ -371,7 +371,6 @@ def test_executables_root_accepts_list(self, tmp_path): assert coordinator.executables_root == [tmp_path, other] def test_executables_root_optional_defaults_to_empty(self): - # Neither an explicit root nor dag_bundle_name: co-located mode, valid. coordinator = ExecutableCoordinator() assert coordinator.executables_root == [] assert coordinator.dag_bundle_name is None @@ -384,13 +383,13 @@ def test_root_and_dag_bundle_name_are_mutually_exclusive(self, tmp_path): with pytest.raises(ValueError, match="at most one of 'executables_root' or 'dag_bundle_name'"): ExecutableCoordinator(executables_root=[tmp_path], dag_bundle_name="artifacts") - @patch("airflow.dag_processing.bundles.manager.DagBundlesManager") + @patch("airflow.sdk.coordinators._subprocess.DagBundlesManager") def test_unconfigured_dag_bundle_name_raises(self, mock_manager): mock_manager.is_bundle_configured.return_value = False with pytest.raises(ValueError, match="unconfigured Dag bundle 'ghost'"): ExecutableCoordinator(dag_bundle_name="ghost") - @patch("airflow.dag_processing.bundles.manager.DagBundlesManager") + @patch("airflow.sdk.coordinators._subprocess.DagBundlesManager") def test_configured_dag_bundle_name_accepted(self, mock_manager): mock_manager.is_bundle_configured.return_value = True coordinator = ExecutableCoordinator(dag_bundle_name="artifacts") @@ -398,8 +397,6 @@ def test_configured_dag_bundle_name_accepted(self, mock_manager): mock_manager.is_bundle_configured.assert_called_once_with("artifacts") def test_build_command_scans_passed_roots_in_colocated_mode(self, tmp_path): - # Co-located mode: no configured root, so the subclass must scan the roots - # the base hands in rather than a configured executables_root field. binary = _build_bundle(tmp_path / "my_bundle", dag_ids=["tutorial_dag"]) coordinator = ExecutableCoordinator() command, schema_version = coordinator._build_execute_task_command( diff --git a/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py index ed5f26a995f8e..e84486150d96d 100644 --- a/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py @@ -238,8 +238,7 @@ def test_custom_kwargs(self): assert coordinator.jars_root == [pathlib.Path("/airflow/java-bundles")] def test_jars_root_optional_defaults_to_empty(self): - # Neither an explicit root nor dag_bundle_name: co-located mode. main_class - # is optional here — the entrypoint is auto-detected from the bundle scan. + # main_class stays optional: the entrypoint is auto-detected from the bundle scan. coordinator = JavaCoordinator() assert coordinator.jars_root == [] assert coordinator.dag_bundle_name is None @@ -253,15 +252,13 @@ def test_root_and_dag_bundle_name_are_mutually_exclusive(self): with pytest.raises(ValueError, match="at most one of 'jars_root' or 'dag_bundle_name'"): JavaCoordinator(jars_root="/airflow/java-bundles", dag_bundle_name="artifacts") - @patch("airflow.dag_processing.bundles.manager.DagBundlesManager") + @patch("airflow.sdk.coordinators._subprocess.DagBundlesManager") def test_unconfigured_dag_bundle_name_raises(self, mock_manager): mock_manager.is_bundle_configured.return_value = False with pytest.raises(ValueError, match="unconfigured Dag bundle 'ghost'"): JavaCoordinator(dag_bundle_name="ghost") def test_build_command_scans_passed_roots_in_colocated_mode(self, tmp_path): - # Co-located mode: no configured root, so the subclass must scan the roots - # the base hands in rather than a configured jars_root field. _make_jar(tmp_path / "app.jar", main_class="com.example.TaskRunner", schema_version="2026-06-16") coordinator = JavaCoordinator(main_class="com.example.TaskRunner") command, schema_version = coordinator._build_execute_task_command(what=_make_ti(), roots=[tmp_path]) diff --git a/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py index cde7b46a790d0..2e8ea8eed9001 100644 --- a/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py @@ -81,7 +81,6 @@ def test_custom_kwargs(self): assert coordinator.task_startup_timeout == 30.0 def test_bundles_root_optional_defaults_to_empty(self): - # Neither an explicit root nor dag_bundle_name: co-located mode, valid. coordinator = NodeCoordinator() assert coordinator.bundles_root == [] assert coordinator.dag_bundle_name is None @@ -94,15 +93,13 @@ def test_root_and_dag_bundle_name_are_mutually_exclusive(self): with pytest.raises(ValueError, match="at most one of 'bundles_root' or 'dag_bundle_name'"): NodeCoordinator(bundles_root="/airflow/ts-bundles", dag_bundle_name="artifacts") - @patch("airflow.dag_processing.bundles.manager.DagBundlesManager") + @patch("airflow.sdk.coordinators._subprocess.DagBundlesManager") def test_unconfigured_dag_bundle_name_raises(self, mock_manager): mock_manager.is_bundle_configured.return_value = False with pytest.raises(ValueError, match="unconfigured Dag bundle 'ghost'"): NodeCoordinator(dag_bundle_name="ghost") def test_build_command_scans_passed_roots_in_colocated_mode(self, tmp_path): - # Co-located mode: no configured root, so the subclass must scan the roots - # the base hands in rather than a configured bundles_root field. bundle = write_bundle(tmp_path, "test_dag") coordinator = NodeCoordinator() command, schema_version = coordinator._build_execute_task_command(what=_make_ti(), roots=[tmp_path]) diff --git a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py index b3e6130d261ab..bf1220237bc16 100644 --- a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py +++ b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py @@ -25,7 +25,6 @@ import sys import threading import time -from typing import ClassVar from unittest.mock import ANY, MagicMock, call, patch import attrs @@ -33,6 +32,7 @@ import pytest from uuid6 import uuid7 +from airflow.dag_processing.bundles.base import BundleVersion from airflow.sdk.api.client import Client, TaskInstanceOperations from airflow.sdk.api.datamodels._generated import BundleInfo, TaskInstance from airflow.sdk.coordinators._subprocess import ( @@ -583,11 +583,10 @@ class _StubSubprocessCoordinator(SubprocessCoordinator): schema_version: str | None = None explicit_roots: list[pathlib.Path] = attrs.field(factory=lambda: [pathlib.Path(".")]) recorded_roots: list[list[pathlib.Path]] = attrs.field(init=False, factory=list) - _root_kwarg: ClassVar[str] = "explicit_roots" @property - def _explicit_artifact_roots(self) -> list[pathlib.Path]: - return self.explicit_roots + def _explicit_artifact_roots(self) -> tuple[str, list[pathlib.Path]]: + return "explicit_roots", self.explicit_roots def _build_execute_task_command(self, *, what, roots): self.recorded_roots.append(roots) @@ -897,13 +896,13 @@ def test_rejects_explicit_root_and_dag_bundle_name_together(self): command=["x"], explicit_roots=[pathlib.Path("/artifacts")], dag_bundle_name="artifacts" ) - @patch("airflow.dag_processing.bundles.manager.DagBundlesManager") + @patch("airflow.sdk.coordinators._subprocess.DagBundlesManager") def test_rejects_unconfigured_dag_bundle_name(self, mock_manager): mock_manager.is_bundle_configured.return_value = False with pytest.raises(ValueError, match="unconfigured Dag bundle 'ghost'"): _StubSubprocessCoordinator(command=["x"], explicit_roots=[], dag_bundle_name="ghost") - @patch("airflow.dag_processing.bundles.manager.DagBundlesManager") + @patch("airflow.sdk.coordinators._subprocess.DagBundlesManager") def test_accepts_configured_dag_bundle_name(self, mock_manager): mock_manager.is_bundle_configured.return_value = True coordinator = _StubSubprocessCoordinator( @@ -923,9 +922,9 @@ def test_returns_configured_root_and_no_bundle_in_explicit_mode(self): assert roots == configured assert bundle is None - @patch("airflow.sdk.execution_time.task_runner.initialize_ti_bundle") + @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle") def test_task_bundle_mode_uses_active_task_bundle(self, mock_initialize, tmp_path): - resolved = MagicMock(supports_versioning=False, path=tmp_path) + resolved = MagicMock(path=tmp_path, version="v3") mock_initialize.return_value = resolved coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=[]) bundle_info = BundleInfo(name="dags", version="v3", version_data={"k": "v"}) @@ -935,23 +934,49 @@ def test_task_bundle_mode_uses_active_task_bundle(self, mock_initialize, tmp_pat assert roots == [tmp_path] assert bundle is resolved - mock_initialize.assert_called_once_with(bundle_info, log) + mock_initialize.assert_called_once_with(bundle_info) def test_task_bundle_mode_without_active_task_raises(self): coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=[]) with pytest.raises(RuntimeError, match="requires an active task"): coordinator._init_root_source(log) - def test_unclassified_source_raises(self): + @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle") + def test_version_less_bundle_is_rematerialized_at_its_current_version(self, mock_initialize, tmp_path): + """A bundle resolved without a version is re-resolved at the version current now. + + Otherwise the subprocess reads the bundle's shared mutable checkout, which no + version lock can protect and which another task's refresh can reset underneath it. + """ + shared_checkout = tmp_path / "tracking_repo" + shared_checkout.mkdir() + pinned_tree = tmp_path / "versions" / "sha-abc" + pinned_tree.mkdir(parents=True) + + unpinned = MagicMock(path=shared_checkout, version=None) + unpinned.get_current_version.return_value = BundleVersion(version="sha-abc", data={"k": "v"}) + pinned = MagicMock(path=pinned_tree, version="sha-abc") + mock_initialize.side_effect = [unpinned, pinned] + coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=[]) - coordinator._artifact_source = None - with pytest.raises(RuntimeError, match="was not classified"): - coordinator._init_root_source(log) + coordinator.dag_bundle_name = "artifacts" + coordinator._artifact_source = _ArtifactSource.NAMED_BUNDLE - @patch("airflow.sdk.execution_time.task_runner.initialize_ti_bundle") - def test_named_bundle_mode_resolves_bundle(self, mock_initialize, tmp_path): - resolved = MagicMock(supports_versioning=False, path=tmp_path) + roots, bundle = coordinator._init_root_source(log) + + assert roots == [pinned_tree] + assert bundle is pinned + assert mock_initialize.call_args_list == [ + call(BundleInfo(name="artifacts")), + call(BundleInfo(name="artifacts", version="sha-abc", version_data={"k": "v"})), + ] + + @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle") + def test_unversioned_bundle_keeps_its_single_path(self, mock_initialize, tmp_path): + resolved = MagicMock(path=tmp_path, version=None) + resolved.get_current_version.return_value = None mock_initialize.return_value = resolved + coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=[]) coordinator.dag_bundle_name = "artifacts" coordinator._artifact_source = _ArtifactSource.NAMED_BUNDLE @@ -960,12 +985,12 @@ def test_named_bundle_mode_resolves_bundle(self, mock_initialize, tmp_path): assert roots == [tmp_path] assert bundle is resolved - mock_initialize.assert_called_once_with(BundleInfo(name="artifacts"), log) + mock_initialize.assert_called_once_with(BundleInfo(name="artifacts")) - @patch("airflow.sdk.execution_time.task_runner.initialize_ti_bundle") + @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle") def test_missing_resolved_path_raises(self, mock_initialize, tmp_path): missing = tmp_path / "nope" - mock_initialize.return_value = MagicMock(supports_versioning=False, path=missing) + mock_initialize.return_value = MagicMock(path=missing, version="v1") coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=[]) coordinator.dag_bundle_name = "artifacts" coordinator._artifact_source = _ArtifactSource.NAMED_BUNDLE @@ -975,19 +1000,18 @@ def test_missing_resolved_path_raises(self, mock_initialize, tmp_path): class TestExecuteTaskBundleWiring: - """execute_task binds the task bundle, forwards resolved roots, holds the version lock, uses the task log.""" + """execute_task binds the task bundle, forwards resolved roots, and holds the version lock.""" - @patch("airflow.dag_processing.bundles.base.BundleVersionLock") - @patch("airflow.sdk.execution_time.task_runner.initialize_ti_bundle") + @patch("airflow.sdk.coordinators._subprocess.BundleVersionLock") + @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle") @patch.object(_PopenActivitySubprocess, "start") def test_task_bundle_mode_binds_forwards_roots_and_locks( self, mock_start, mock_initialize, mock_lock, mock_client, tmp_path ): - resolved = MagicMock(supports_versioning=False, path=tmp_path, version="v9") + resolved = MagicMock(path=tmp_path, version="v9") resolved.name = "dags" mock_initialize.return_value = resolved mock_start.return_value.wait.return_value = 0 - task_logger = MagicMock() coordinator = _StubSubprocessCoordinator(command=["/runtime"], explicit_roots=[]) bundle_info = BundleInfo(name="dags", version="v9") @@ -997,12 +1021,11 @@ def test_task_bundle_mode_binds_forwards_roots_and_locks( dag_rel_path="dag.py", bundle_info=bundle_info, client=mock_client, - logger=task_logger, subprocess_logs_to_stdout=False, ) - # Resolution runs against the bound task bundle, using the task logger (not the module log). - mock_initialize.assert_called_once_with(bundle_info, task_logger) + # Resolution runs against the bound task bundle. + mock_initialize.assert_called_once_with(bundle_info) # The resolved root is forwarded to the subclass command builder. assert coordinator.recorded_roots == [[tmp_path]] # The version lock is held around the subprocess for the resolved bundle. @@ -1010,7 +1033,7 @@ def test_task_bundle_mode_binds_forwards_roots_and_locks( mock_lock.return_value.__enter__.assert_called_once() mock_lock.return_value.__exit__.assert_called_once() - @patch("airflow.dag_processing.bundles.base.BundleVersionLock") + @patch("airflow.sdk.coordinators._subprocess.BundleVersionLock") @patch.object(_PopenActivitySubprocess, "start") def test_explicit_root_mode_forwards_roots_without_locking( self, mock_start, mock_lock, mock_client, tmp_path @@ -1029,20 +1052,21 @@ def test_explicit_root_mode_forwards_roots_without_locking( assert coordinator.recorded_roots == [[tmp_path]] mock_lock.assert_not_called() - @patch("airflow.dag_processing.bundles.base.BundleVersionLock") - @patch("airflow.sdk.execution_time.task_runner.initialize_ti_bundle") + @patch("airflow.sdk.coordinators._subprocess.BundleVersionLock") + @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle") @patch.object(_PopenActivitySubprocess, "start") - def test_named_bundle_mode_locks_resolved_current_version( + def test_named_bundle_mode_locks_the_tree_it_scans( self, mock_start, mock_initialize, mock_lock, mock_client, tmp_path ): - # NAMED_BUNDLE resolves "latest" with version=None; the lock must pin the - # concrete current version so it does not silently become a no-op. - from airflow.dag_processing.bundles.base import BundleVersion - - resolved = MagicMock(path=tmp_path, version=None) - resolved.name = "artifacts" - resolved.get_current_version.return_value = BundleVersion(version="sha-abc", data=None) - mock_initialize.return_value = resolved + """The lock names the pinned version whose tree is handed to the subprocess.""" + pinned_tree = tmp_path / "versions" / "sha-abc" + pinned_tree.mkdir(parents=True) + + unpinned = MagicMock(path=tmp_path, version=None) + unpinned.get_current_version.return_value = BundleVersion(version="sha-abc", data=None) + pinned = MagicMock(path=pinned_tree, version="sha-abc") + pinned.name = "artifacts" + mock_initialize.side_effect = [unpinned, pinned] mock_start.return_value.wait.return_value = 0 coordinator = _StubSubprocessCoordinator(command=["/runtime"], explicit_roots=[]) @@ -1057,7 +1081,7 @@ def test_named_bundle_mode_locks_resolved_current_version( subprocess_logs_to_stdout=False, ) - resolved.get_current_version.assert_called_once_with() + assert coordinator.recorded_roots == [[pinned_tree]] mock_lock.assert_called_once_with(bundle_name="artifacts", bundle_version="sha-abc") diff --git a/task-sdk/tests/task_sdk/execution_time/test_bundles.py b/task-sdk/tests/task_sdk/execution_time/test_bundles.py new file mode 100644 index 0000000000000..0a58b6484dc64 --- /dev/null +++ b/task-sdk/tests/task_sdk/execution_time/test_bundles.py @@ -0,0 +1,86 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +from pathlib import Path +from unittest import mock +from unittest.mock import patch + +import pytest + +from airflow.sdk.api.datamodels._generated import BundleInfo +from airflow.sdk.exceptions import AirflowException +from airflow.sdk.execution_time.bundles import initialize_ti_bundle, verify_bundle_access + + +def test_verify_bundle_access_raises_when_not_accessible(tmp_path: Path): + bundle_path = tmp_path / "test_bundle" + bundle_path.mkdir() + + mock_bundle = mock.Mock() + mock_bundle.path = bundle_path + mock_bundle.name = "test-bundle" + + # Mock os.access to simulate permission denied (avoids root user issues in CI) + with patch("airflow.sdk.execution_time.bundles.os.access", return_value=False): + with pytest.raises(AirflowException) as exc_info: + verify_bundle_access(mock_bundle) + + assert "not accessible" in str(exc_info.value) + assert "test-bundle" in str(exc_info.value) + + +def test_verify_bundle_access_succeeds_when_readable(tmp_path: Path): + bundle_path = tmp_path / "accessible_bundle" + bundle_path.mkdir() + + mock_bundle = mock.Mock() + mock_bundle.path = bundle_path + mock_bundle.name = "test-bundle" + + verify_bundle_access(mock_bundle) + + +def test_verify_bundle_access_skips_nonexistent_path(tmp_path: Path): + mock_bundle = mock.Mock() + mock_bundle.path = tmp_path / "nonexistent" + mock_bundle.name = "test-bundle" + + # Should not raise - nonexistent paths are handled by initialize() + verify_bundle_access(mock_bundle) + + +def test_initialize_ti_bundle_resolves_initializes_and_verifies(tmp_path: Path): + """initialize_ti_bundle resolves the bundle from BundleInfo, initializes it, and access-checks it.""" + bundle_path = tmp_path / "bundle" + bundle_path.mkdir() + mock_bundle = mock.Mock() + mock_bundle.path = bundle_path + mock_bundle.name = "my-bundle" + + bundle_info = BundleInfo(name="my-bundle", version="v2", version_data={"k": "v"}) + + with patch("airflow.sdk.execution_time.bundles.DagBundlesManager") as mock_manager: + mock_manager.return_value.get_bundle.return_value = mock_bundle + result = initialize_ti_bundle(bundle_info) + + assert result is mock_bundle + mock_manager.return_value.get_bundle.assert_called_once_with( + name="my-bundle", version="v2", version_data={"k": "v"} + ) + mock_bundle.initialize.assert_called_once_with() 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 bb152fef8187b..da4cef6ff6c20 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 @@ -180,7 +180,6 @@ _run_execute_callable, _serialize_outlet_events, _xcom_push, - detail_span, finalize, get_startup_details, parse, @@ -837,80 +836,6 @@ def test_parse_module_in_bundle_root(tmp_path: Path, make_ti_context): assert ti.task.dag.dag_id == "dag_name" -def test_verify_bundle_access_raises_when_not_accessible(tmp_path: Path, make_ti_context): - """Test that _verify_bundle_access raises AirflowException when bundle path is not accessible.""" - from airflow.sdk.execution_time.task_runner import _verify_bundle_access - - # Create a directory that exists - bundle_path = tmp_path / "test_bundle" - bundle_path.mkdir() - - # Create a mock bundle instance - mock_bundle = mock.Mock() - mock_bundle.path = bundle_path - mock_bundle.name = "test-bundle" - - # Mock os.access to simulate permission denied (avoids root user issues in CI) - with patch("airflow.sdk.execution_time.task_runner.os.access", return_value=False): - with pytest.raises(AirflowException) as exc_info: - _verify_bundle_access(mock_bundle, mock.Mock()) - - assert "not accessible" in str(exc_info.value) - assert "test-bundle" in str(exc_info.value) - - -def test_verify_bundle_access_succeeds_when_readable(tmp_path: Path, make_ti_context): - """Test that _verify_bundle_access succeeds when bundle path is accessible.""" - from airflow.sdk.execution_time.task_runner import _verify_bundle_access - - # Create a directory with read permissions - bundle_path = tmp_path / "accessible_bundle" - bundle_path.mkdir() - - mock_bundle = mock.Mock() - mock_bundle.path = bundle_path - mock_bundle.name = "test-bundle" - - # Should not raise - _verify_bundle_access(mock_bundle, mock.Mock()) - - -def test_verify_bundle_access_skips_nonexistent_path(tmp_path: Path): - """Test that _verify_bundle_access does nothing when bundle path doesn't exist.""" - from airflow.sdk.execution_time.task_runner import _verify_bundle_access - - mock_bundle = mock.Mock() - mock_bundle.path = tmp_path / "nonexistent" - mock_bundle.name = "test-bundle" - - # Should not raise - nonexistent paths are handled by initialize() - _verify_bundle_access(mock_bundle, mock.Mock()) - - -def test_initialize_ti_bundle_resolves_initializes_and_verifies(tmp_path: Path): - """initialize_ti_bundle resolves the bundle from BundleInfo, initializes it, and access-checks it.""" - from airflow.sdk.api.datamodels._generated import BundleInfo - from airflow.sdk.execution_time.task_runner import initialize_ti_bundle - - bundle_path = tmp_path / "bundle" - bundle_path.mkdir() - mock_bundle = mock.Mock() - mock_bundle.path = bundle_path - mock_bundle.name = "my-bundle" - - bundle_info = BundleInfo(name="my-bundle", version="v2", version_data={"k": "v"}) - - with patch("airflow.sdk.execution_time.task_runner.DagBundlesManager") as mock_manager: - mock_manager.return_value.get_bundle.return_value = mock_bundle - result = initialize_ti_bundle(bundle_info, mock.Mock()) - - assert result is mock_bundle - mock_manager.return_value.get_bundle.assert_called_once_with( - name="my-bundle", version="v2", version_data={"k": "v"} - ) - mock_bundle.initialize.assert_called_once_with() - - @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.""" @@ -6188,120 +6113,6 @@ def test_operator_metrics_respect_team_name( backend.incr.assert_any_call(ti_metric, tags=stats_tags) -class TestDetailSpan: - """Tests for the detail_span decorator / context manager.""" - - @pytest.fixture(autouse=True) - def _sampled_carrier_provider(self): - """Make new_dagrun_trace_carrier produce a SAMPLED carrier. - - new_dagrun_trace_carrier consults the global tracer provider's sampler to - decide the carrier's SAMPLED flag. In the test process the global provider - is a no-op ProxyTracerProvider (no sampler) -> unsampled carrier, which - would make the parent span (and its detail children) non-recording. Patch - the lookup to a real SDK provider whose default sampler - (parentbased_always_on) samples the root, mirroring "otel on" in production. - """ - provider = TracerProvider() - with mock.patch( - "airflow._shared.observability.traces.trace.get_tracer_provider", - return_value=provider, - ): - yield - - def test_level_1_no_child_span_as_context_manager(self): - """At detail level 1, entering detail_span should not create a real recorded span.""" - exporter = InMemorySpanExporter() - provider = TracerProvider() - provider.add_span_processor(SimpleSpanProcessor(exporter)) - t = provider.get_tracer("test") - carrier = new_dagrun_trace_carrier(task_span_detail_level=1) - parent_ctx = TraceContextTextMapPropagator().extract(carrier) - - with mock.patch("airflow.sdk.execution_time.task_runner.tracer", t): - with t.start_as_current_span("parent", context=parent_ctx): - with detail_span("child") as span: - assert span is trace.INVALID_SPAN - - # Only the "parent" span should be recorded; no "child". - names = [s.name for s in exporter.get_finished_spans()] - assert "child" not in names - - def test_level_2_creates_child_span_as_context_manager(self): - """At detail level 2, detail_span should create a real recorded child span.""" - exporter = InMemorySpanExporter() - provider = TracerProvider() - provider.add_span_processor(SimpleSpanProcessor(exporter)) - t = provider.get_tracer("test") - carrier = new_dagrun_trace_carrier(task_span_detail_level=2) - parent_ctx = TraceContextTextMapPropagator().extract(carrier) - - with mock.patch("airflow.sdk.execution_time.task_runner.tracer", t): - with t.start_as_current_span("parent", context=parent_ctx): - with detail_span("child"): - pass - - names = [s.name for s in exporter.get_finished_spans()] - assert "child" in names - - def test_decorator_at_level_1_does_not_create_span(self): - """@detail_span at level 1 should not produce a recorded span.""" - exporter = InMemorySpanExporter() - provider = TracerProvider() - provider.add_span_processor(SimpleSpanProcessor(exporter)) - t = provider.get_tracer("test") - carrier = new_dagrun_trace_carrier(task_span_detail_level=1) - parent_ctx = TraceContextTextMapPropagator().extract(carrier) - - @detail_span("decorated") - def my_func(): - return 42 - - with mock.patch("airflow.sdk.execution_time.task_runner.tracer", t): - with t.start_as_current_span("parent", context=parent_ctx): - result = my_func() - - assert result == 42 - names = [s.name for s in exporter.get_finished_spans()] - assert "decorated" not in names - - def test_decorator_at_level_2_creates_span_and_preserves_return_value(self): - """@detail_span at level 2 creates a span and the wrapped function's return value is preserved.""" - exporter = InMemorySpanExporter() - provider = TracerProvider() - provider.add_span_processor(SimpleSpanProcessor(exporter)) - t = provider.get_tracer("test") - carrier = new_dagrun_trace_carrier(task_span_detail_level=2) - parent_ctx = TraceContextTextMapPropagator().extract(carrier) - - @detail_span("decorated") - def my_func(x): - return x * 2 - - with mock.patch("airflow.sdk.execution_time.task_runner.tracer", t): - with t.start_as_current_span("parent", context=parent_ctx): - result = my_func(7) - - assert result == 14 - names = [s.name for s in exporter.get_finished_spans()] - assert "decorated" in names - - def test_exception_in_context_manager_propagates(self): - """Exceptions inside `with detail_span(...)` propagate normally.""" - exporter = InMemorySpanExporter() - provider = TracerProvider() - provider.add_span_processor(SimpleSpanProcessor(exporter)) - t = provider.get_tracer("test") - carrier = new_dagrun_trace_carrier(task_span_detail_level=2) - parent_ctx = TraceContextTextMapPropagator().extract(carrier) - - with mock.patch("airflow.sdk.execution_time.task_runner.tracer", t): - with t.start_as_current_span("parent", context=parent_ctx): - with pytest.raises(ValueError, match="boom"): - with detail_span("child"): - raise ValueError("boom") - - class TestRunExecuteCallable: """Tests for ``_run_execute_callable``. @@ -6383,7 +6194,7 @@ def test_emits_task_execute_span_at_detail_level_2(self): task = self._make_task() - with mock.patch("airflow.sdk.execution_time.task_runner.tracer", t): + with mock.patch("airflow.sdk.execution_time.tracing.tracer", t): with t.start_as_current_span("parent", context=parent_ctx): result = _run_execute_callable(context={}, execute=lambda context: "ok", task=task) @@ -6411,7 +6222,7 @@ def execute(context): with t.start_as_current_span("operator_child"): return "ok" - with mock.patch("airflow.sdk.execution_time.task_runner.tracer", t): + with mock.patch("airflow.sdk.execution_time.tracing.tracer", t): with t.start_as_current_span("parent", context=parent_ctx): result = _run_execute_callable(context={}, execute=execute, task=task) @@ -6430,7 +6241,7 @@ def test_no_task_execute_span_at_detail_level_1(self): task = self._make_task() - with mock.patch("airflow.sdk.execution_time.task_runner.tracer", t): + with mock.patch("airflow.sdk.execution_time.tracing.tracer", t): with t.start_as_current_span("parent", context=parent_ctx): result = _run_execute_callable(context={}, execute=lambda context: "ok", task=task) diff --git a/task-sdk/tests/task_sdk/execution_time/test_tracing.py b/task-sdk/tests/task_sdk/execution_time/test_tracing.py new file mode 100644 index 0000000000000..fcd467ba6b8e7 --- /dev/null +++ b/task-sdk/tests/task_sdk/execution_time/test_tracing.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. + +from __future__ import annotations + +from unittest import mock + +import pytest +from opentelemetry import trace +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator + +from airflow.sdk._shared.observability.traces import new_dagrun_trace_carrier +from airflow.sdk.execution_time.tracing import detail_span + + +class TestDetailSpan: + """Tests for the detail_span decorator / context manager.""" + + @pytest.fixture(autouse=True) + def _sampled_carrier_provider(self): + """Make new_dagrun_trace_carrier produce a SAMPLED carrier. + + new_dagrun_trace_carrier consults the global tracer provider's sampler to + decide the carrier's SAMPLED flag. In the test process the global provider + is a no-op ProxyTracerProvider (no sampler) -> unsampled carrier, which + would make the parent span (and its detail children) non-recording. Patch + the lookup to a real SDK provider whose default sampler + (parentbased_always_on) samples the root, mirroring "otel on" in production. + """ + provider = TracerProvider() + with mock.patch( + "airflow._shared.observability.traces.trace.get_tracer_provider", + return_value=provider, + ): + yield + + def test_level_1_no_child_span_as_context_manager(self): + """At detail level 1, entering detail_span should not create a real recorded span.""" + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + t = provider.get_tracer("test") + carrier = new_dagrun_trace_carrier(task_span_detail_level=1) + parent_ctx = TraceContextTextMapPropagator().extract(carrier) + + with mock.patch("airflow.sdk.execution_time.tracing.tracer", t): + with t.start_as_current_span("parent", context=parent_ctx): + with detail_span("child") as span: + assert span is trace.INVALID_SPAN + + # Only the "parent" span should be recorded; no "child". + names = [s.name for s in exporter.get_finished_spans()] + assert "child" not in names + + def test_level_2_creates_child_span_as_context_manager(self): + """At detail level 2, detail_span should create a real recorded child span.""" + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + t = provider.get_tracer("test") + carrier = new_dagrun_trace_carrier(task_span_detail_level=2) + parent_ctx = TraceContextTextMapPropagator().extract(carrier) + + with mock.patch("airflow.sdk.execution_time.tracing.tracer", t): + with t.start_as_current_span("parent", context=parent_ctx): + with detail_span("child"): + pass + + names = [s.name for s in exporter.get_finished_spans()] + assert "child" in names + + def test_decorator_at_level_1_does_not_create_span(self): + """@detail_span at level 1 should not produce a recorded span.""" + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + t = provider.get_tracer("test") + carrier = new_dagrun_trace_carrier(task_span_detail_level=1) + parent_ctx = TraceContextTextMapPropagator().extract(carrier) + + @detail_span("decorated") + def my_func(): + return 42 + + with mock.patch("airflow.sdk.execution_time.tracing.tracer", t): + with t.start_as_current_span("parent", context=parent_ctx): + result = my_func() + + assert result == 42 + names = [s.name for s in exporter.get_finished_spans()] + assert "decorated" not in names + + def test_decorator_at_level_2_creates_span_and_preserves_return_value(self): + """@detail_span at level 2 creates a span and the wrapped function's return value is preserved.""" + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + t = provider.get_tracer("test") + carrier = new_dagrun_trace_carrier(task_span_detail_level=2) + parent_ctx = TraceContextTextMapPropagator().extract(carrier) + + @detail_span("decorated") + def my_func(x): + return x * 2 + + with mock.patch("airflow.sdk.execution_time.tracing.tracer", t): + with t.start_as_current_span("parent", context=parent_ctx): + result = my_func(7) + + assert result == 14 + names = [s.name for s in exporter.get_finished_spans()] + assert "decorated" in names + + def test_exception_in_context_manager_propagates(self): + """Exceptions inside `with detail_span(...)` propagate normally.""" + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + t = provider.get_tracer("test") + carrier = new_dagrun_trace_carrier(task_span_detail_level=2) + parent_ctx = TraceContextTextMapPropagator().extract(carrier) + + with mock.patch("airflow.sdk.execution_time.tracing.tracer", t): + with t.start_as_current_span("parent", context=parent_ctx): + with pytest.raises(ValueError, match="boom"): + with detail_span("child"): + raise ValueError("boom") From e9da1565f5b411603a1314bf6d10ea933480bc0f Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Mon, 10 Aug 2026 02:44:31 +0000 Subject: [PATCH 06/11] Read the Dag bundle configuration through a single entry point Splitting the read across a keyed builder and a wrapper meant two names for one idea, and the empty case leaked out as a module-level constant that existed only to be returned. One function now answers "what is configured", parsed once per process so both readers of it see the same answer. Parsing once per process means a later configuration change is not observed. That only arises in tests, which switch configuration freely, so overriding configuration there drops the memo rather than silently serving a stale one. --- .../airflow/dag_processing/bundles/manager.py | 41 +++------------ .../bundles/test_dag_bundle_manager.py | 52 ++++--------------- .../src/tests_common/pytest_plugin.py | 26 ++++++++++ .../src/tests_common/test_utils/config.py | 13 +++++ 4 files changed, 56 insertions(+), 76 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/bundles/manager.py b/airflow-core/src/airflow/dag_processing/bundles/manager.py index e19d702d71fca..12b3b50fbca21 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/manager.py +++ b/airflow-core/src/airflow/dag_processing/bundles/manager.py @@ -18,7 +18,6 @@ import functools import importlib -import json import logging import os import warnings @@ -116,20 +115,13 @@ class _BundleConfigSnapshot(NamedTuple): names: frozenset[str] -_EMPTY_BUNDLE_CONFIG_SNAPSHOT = _BundleConfigSnapshot(configs=(), names=frozenset()) - - @functools.cache -def _parse_bundle_config_snapshot(config_json: str, load_examples: bool) -> _BundleConfigSnapshot: - """ - Build the snapshot for one configuration, without importing any bundle class. - - Keyed on the configuration rather than cached outright, so every reader of the - same configuration shares one snapshot while a configuration change is still - picked up. - """ - bundle_config_list = _parse_bundle_config(json.loads(config_json)) - if load_examples: +def _load_bundle_config_snapshot() -> _BundleConfigSnapshot: + """Read and validate the configured Dag bundles, without importing their classes.""" + bundle_config_list = _read_bundle_config_list() + if not bundle_config_list: + return _BundleConfigSnapshot(configs=(), names=frozenset()) + if conf.getboolean("core", "LOAD_EXAMPLES"): _add_example_dag_bundle(bundle_config_list) _add_provider_example_dags_to_bundle(bundle_config_list) return _BundleConfigSnapshot( @@ -138,27 +130,6 @@ def _parse_bundle_config_snapshot(config_json: str, load_examples: bool) -> _Bun ) -def _load_bundle_config_snapshot() -> _BundleConfigSnapshot: - """ - Read and validate the configured Dag bundles. - - Shared by :meth:`DagBundlesManager.parse_config`, which goes on to import each - bundle class, and by :meth:`DagBundlesManager.is_bundle_configured`, which only - needs the names, so both resolve to the same snapshot. - """ - config_list = conf.getjson("dag_processor", "dag_bundle_config_list") - if not config_list: - return _EMPTY_BUNDLE_CONFIG_SNAPSHOT - if not isinstance(config_list, list): - raise AirflowConfigException( - "Section `dag_processor` key `dag_bundle_config_list` " - f"must be list but got {config_list.__class__}" - ) - return _parse_bundle_config_snapshot( - json.dumps(config_list, sort_keys=True), conf.getboolean("core", "LOAD_EXAMPLES") - ) - - def _parse_bundle_config(config_list) -> list[_ExternalBundleConfig]: bundles = {} for item in config_list: diff --git a/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py b/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py index f080fe137c6ce..6eac0c575fbb0 100644 --- a/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py +++ b/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py @@ -34,7 +34,6 @@ _get_configured_bundle_team_names, _guess_best_bundle_for_fileloc, _load_bundle_config_snapshot, - _parse_bundle_config_snapshot, ) from airflow.exceptions import AirflowConfigException from airflow.models.dag import DagModel @@ -209,9 +208,9 @@ def test_get_bundle(): @pytest.fixture(autouse=True) def _clear_bundle_config_snapshot_cache(): - _parse_bundle_config_snapshot.cache_clear() + _load_bundle_config_snapshot.cache_clear() yield - _parse_bundle_config_snapshot.cache_clear() + _load_bundle_config_snapshot.cache_clear() def _bundle_config_env(config) -> dict[str, str]: @@ -256,17 +255,6 @@ def test_is_bundle_configured_sees_implicitly_added_bundles(): assert DagBundlesManager.is_bundle_configured("example_dags") is True -def test_bundle_config_snapshot_is_shared_and_config_scoped(): - """Readers of one configuration share a snapshot; a config change builds a new one.""" - with patch.dict(os.environ, _bundle_config_env(BASIC_BUNDLE_CONFIG)): - first = _load_bundle_config_snapshot() - assert _load_bundle_config_snapshot() is first - - other_config = [{"name": "other", "classpath": BASIC_BUNDLE_CONFIG[0]["classpath"], "kwargs": {}}] - with patch.dict(os.environ, _bundle_config_env(other_config)): - assert _load_bundle_config_snapshot().names == {"other"} - - @pytest.fixture def clear_db(): clear_db_dag_bundles() @@ -283,9 +271,7 @@ def _get_bundle_names_and_active(): ).all() # Initial add - with patch.dict( - os.environ, {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(BASIC_BUNDLE_CONFIG)} - ): + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(BASIC_BUNDLE_CONFIG)}): manager = DagBundlesManager() manager.sync_bundles_to_db() assert _get_bundle_names_and_active() == [("my-test-bundle", True)] @@ -310,9 +296,7 @@ def _get_bundle_names_and_active(): assert session.scalar(select(func.count(ParseImportError.id))) == 0 # Re-enable one that reappears in config - with patch.dict( - os.environ, {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(BASIC_BUNDLE_CONFIG)} - ): + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(BASIC_BUNDLE_CONFIG)}): manager = DagBundlesManager() manager.sync_bundles_to_db() assert _get_bundle_names_and_active() == [ @@ -360,16 +344,12 @@ def _get_bundle_names_and_active(): ).all() # Processor A: only knows "my-test-bundle". - with patch.dict( - os.environ, {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(BASIC_BUNDLE_CONFIG)} - ): + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(BASIC_BUNDLE_CONFIG)}): DagBundlesManager().sync_bundles_to_db(deactivate_missing=False) assert _get_bundle_names_and_active() == [("my-test-bundle", True)] # Processor B: only knows "other-test-bundle". It must not touch A's bundle. - with patch.dict( - os.environ, {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(OTHER_BUNDLE_CONFIG)} - ): + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(OTHER_BUNDLE_CONFIG)}): DagBundlesManager().sync_bundles_to_db(deactivate_missing=False) assert _get_bundle_names_and_active() == [ ("my-test-bundle", True), @@ -377,9 +357,7 @@ def _get_bundle_names_and_active(): ] # Processor A runs again: still must not disable B's bundle. - with patch.dict( - os.environ, {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(BASIC_BUNDLE_CONFIG)} - ): + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(BASIC_BUNDLE_CONFIG)}): DagBundlesManager().sync_bundles_to_db(deactivate_missing=False) assert _get_bundle_names_and_active() == [ ("my-test-bundle", True), @@ -478,9 +456,7 @@ def test_bundle_model_render_url(clear_db, session): @conf_vars({("core", "LOAD_EXAMPLES"): "False"}) def test_template_params_update_on_sync(clear_db, session): """Test that template parameters are updated when bundle configuration changes.""" - with patch.dict( - os.environ, {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(TEMPLATE_BUNDLE_CONFIG)} - ): + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(TEMPLATE_BUNDLE_CONFIG)}): manager = DagBundlesManager() manager.sync_bundles_to_db() @@ -503,9 +479,7 @@ def test_template_params_update_on_sync(clear_db, session): } ] - with patch.dict( - os.environ, {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(updated_config)} - ): + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(updated_config)}): manager = DagBundlesManager() manager.sync_bundles_to_db() @@ -522,9 +496,7 @@ def test_template_params_update_on_sync(clear_db, session): def test_template_update_on_sync(clear_db, session): """Test that templates are updated when bundle configuration changes.""" # First, sync with initial template - with patch.dict( - os.environ, {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(TEMPLATE_BUNDLE_CONFIG)} - ): + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(TEMPLATE_BUNDLE_CONFIG)}): manager = DagBundlesManager() manager.sync_bundles_to_db() @@ -547,9 +519,7 @@ def test_template_update_on_sync(clear_db, session): } ] - with patch.dict( - os.environ, {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(updated_config)} - ): + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(updated_config)}): manager = DagBundlesManager() manager.sync_bundles_to_db() diff --git a/devel-common/src/tests_common/pytest_plugin.py b/devel-common/src/tests_common/pytest_plugin.py index 445581cd2b8dc..bb85fe60c2d2d 100644 --- a/devel-common/src/tests_common/pytest_plugin.py +++ b/devel-common/src/tests_common/pytest_plugin.py @@ -2186,6 +2186,32 @@ def clear_current_task_instance_session(): task_instance_session.__current_task_instance_session = None +@pytest.fixture(autouse=True) +def reset_dag_bundle_config_cache(): + """Reset the per-process Dag bundle configuration cache between tests. + + The configuration is parsed once per process, so a test that sets a different + ``[dag_processor] dag_bundle_config_list`` would otherwise be served the previous + test's bundles. ``conf_vars`` clears it too, for tests that switch config midway. + """ + if importlib.util.find_spec("airflow") is None: + yield + return + + try: + from airflow.dag_processing.bundles.manager import _load_bundle_config_snapshot + except ImportError: + # compat for airflow versions without the snapshot cache + yield + return + + _load_bundle_config_snapshot.cache_clear() + try: + yield + finally: + _load_bundle_config_snapshot.cache_clear() + + @pytest.fixture(autouse=True) def refuse_to_run_test_from_wrongly_named_files(request: pytest.FixtureRequest): filepath = request.node.path diff --git a/devel-common/src/tests_common/test_utils/config.py b/devel-common/src/tests_common/test_utils/config.py index 9893da665a4a7..619a88a6930f3 100644 --- a/devel-common/src/tests_common/test_utils/config.py +++ b/devel-common/src/tests_common/test_utils/config.py @@ -98,6 +98,7 @@ def conf_vars(overrides): if "airflow.configuration" in sys.modules: settings.configure_vars() + _clear_dag_bundle_config_cache() try: yield @@ -116,6 +117,18 @@ def conf_vars(overrides): if "airflow.configuration" in sys.modules: settings.configure_vars() + _clear_dag_bundle_config_cache() + + +def _clear_dag_bundle_config_cache() -> None: + """Drop the per-process Dag bundle configuration cache so the new config is read.""" + import sys + + manager = sys.modules.get("airflow.dag_processing.bundles.manager") + # compat for airflow versions without the snapshot cache + cache = getattr(manager, "_load_bundle_config_snapshot", None) + if cache is not None: + cache.cache_clear() @overload From 68cca25a3e13866f6c37f62a10fbfeced50f1bb7 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Mon, 17 Aug 2026 08:46:03 +0000 Subject: [PATCH 07/11] Preserve custom subprocess coordinator compatibility External language SDK coordinators can implement the documented command hook. Keeping that contract stable avoids unexpected-keyword failures when upgrading while allowing Dag-bundle-backed artifact discovery. --- contributing-docs/30_new_language_sdk.rst | 13 +++++---- .../airflow/sdk/coordinators/_subprocess.py | 27 ++++++++++++++----- .../coordinators/executable/coordinator.py | 5 ++-- .../sdk/coordinators/java/coordinator.py | 5 ++-- .../sdk/coordinators/node/coordinator.py | 5 ++-- .../executable/test_coordinator.py | 24 +++++++++++------ .../coordinators/java/test_coordinator.py | 3 ++- .../coordinators/node/test_coordinator.py | 6 +++-- .../task_sdk/coordinators/test_subprocess.py | 19 ++++++++++--- 9 files changed, 70 insertions(+), 37 deletions(-) diff --git a/contributing-docs/30_new_language_sdk.rst b/contributing-docs/30_new_language_sdk.rst index 8983be1681de4..34c4cdbb4a27a 100644 --- a/contributing-docs/30_new_language_sdk.rst +++ b/contributing-docs/30_new_language_sdk.rst @@ -110,9 +110,7 @@ the subprocess understands: .. code-block:: python - def _build_execute_task_command( - self, *, what: TaskInstanceDTO, roots: list[pathlib.Path] - ) -> tuple[list[str], str]: ... + def _build_execute_task_command(self, *, what: TaskInstanceDTO) -> tuple[list[str], str]: ... The method returns a ``(command, subprocess_schema_version)`` pair: @@ -122,10 +120,11 @@ The method returns a ``(command, subprocess_schema_version)`` pair: subprocess understands, used by the supervisor to negotiate message formats across SDK versions. See `Supervisor Schema`_ below. -``roots`` are the artifact directories the base class has already resolved from -the coordinator's configured source — an explicit filesystem root, a named Dag -bundle (``dag_bundle_name``), or the task's own bundle — so subclasses scan -``roots`` rather than reading the configured root directly. +Call ``self._get_scan_roots()`` to retrieve the artifact directories the base +class has already resolved from the coordinator's configured source — an +explicit filesystem root, a named Dag bundle (``dag_bundle_name``), or the +task's own bundle. Subclasses should scan those roots rather than reading the +configured root directly. Supervisor Schema ~~~~~~~~~~~~~~~~~ diff --git a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py index b34aaf77fb1db..ffed5ad2e356c 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py +++ b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py @@ -447,6 +447,7 @@ class SubprocessCoordinator(BaseCoordinator): # resolve roots without knowing the subclass field name. _configured_roots: list[pathlib.Path] = attrs.field(init=False, factory=list) _active_bundle_info: BundleInfo | None = attrs.field(init=False, default=None) + _active_scan_roots: tuple[pathlib.Path, ...] | None = attrs.field(init=False, default=None) @property def _explicit_artifact_roots(self) -> tuple[str, Sequence[pathlib.Path]]: @@ -526,14 +527,18 @@ def _init_root_source( raise FileNotFoundError(f"Dag bundle {target.name!r} resolved to {path}, which does not exist.") return [path], bundle - def _build_execute_task_command( - self, *, what: TaskInstance, roots: list[pathlib.Path] - ) -> tuple[list[str], str | None]: + def _get_scan_roots(self) -> tuple[pathlib.Path, ...]: + """Return the artifact roots resolved for the active task.""" + if self._active_scan_roots is None: + raise RuntimeError("_get_scan_roots requires an active task; call it during execute_task.") + return self._active_scan_roots + + def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], str | None]: """ Build the subprocess command and resolve its supervisor wire-schema version for *what*. - *roots* are the directories to scan for artifacts, already resolved by - :meth:`_init_root_source` from the coordinator's configured source. + Subclasses can retrieve the directories to scan for artifacts with + :meth:`_get_scan_roots`. Returns a ``(command, subprocess_schema_version)`` pair. *command* MUST NOT include the ``--comm`` / ``--logs`` flags — those are appended by :class:`_PopenActivitySubprocess` once the listening sockets have been @@ -558,6 +563,15 @@ def _set_current_bundle(self, bundle_info: BundleInfo): finally: self._active_bundle_info = None + @contextlib.contextmanager + def _set_scan_roots(self, roots: Sequence[pathlib.Path]): + """Expose *roots* to the command builder for the duration of the task.""" + self._active_scan_roots = tuple(roots) + try: + yield + finally: + self._active_scan_roots = None + def execute_task( self, *, @@ -584,7 +598,8 @@ def execute_task( bundle_version=resolved_bundle.version, ) ) - command, subprocess_schema_version = self._build_execute_task_command(what=what, roots=roots) + stack.enter_context(self._set_scan_roots(roots)) + command, subprocess_schema_version = self._build_execute_task_command(what=what) process = _PopenActivitySubprocess.start( what=what, dag_rel_path=dag_rel_path, diff --git a/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py index 16d57f39b5832..59e05300b8568 100644 --- a/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py @@ -354,8 +354,7 @@ class ExecutableCoordinator(SubprocessCoordinator): def _explicit_artifact_roots(self) -> tuple[str, list[pathlib.Path]]: return "executables_root", self.executables_root - def _build_execute_task_command( - self, *, what: TaskInstance, roots: list[pathlib.Path] - ) -> tuple[list[str], str | None]: + def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], str | None]: + roots = self._get_scan_roots() bundle = _Bundle.find(roots, what.dag_id) return [str(bundle.path)], bundle.schema_version diff --git a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py index 299ea5e07d826..c5db8557cc19d 100644 --- a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py @@ -209,11 +209,10 @@ class JavaCoordinator(SubprocessCoordinator): def _explicit_artifact_roots(self) -> tuple[str, list[pathlib.Path]]: return "jars_root", self.jars_root - def _build_execute_task_command( - self, *, what: TaskInstance, roots: list[pathlib.Path] - ) -> tuple[list[str], str | None]: + def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], str | None]: # Without main_class, the first executable JAR in walk order wins; tracked at # https://github.com/apache/airflow/issues/71134 + roots = self._get_scan_roots() jar = _JarInfo.find(roots, self.main_class) command = [ self.java_executable, diff --git a/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py index f9b23c179ccc3..35aff8181d3e5 100644 --- a/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py @@ -127,8 +127,7 @@ class NodeCoordinator(SubprocessCoordinator): def _explicit_artifact_roots(self) -> tuple[str, list[pathlib.Path]]: return "bundles_root", self.bundles_root - def _build_execute_task_command( - self, *, what: TaskInstance, roots: list[pathlib.Path] - ) -> tuple[list[str], str | None]: + def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], str | None]: + roots = self._get_scan_roots() bundle = _Bundle.find(roots, what.dag_id) return [self.node_executable, os.fspath(bundle.path)], bundle.schema_version diff --git a/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py index d62e420e857e8..502456d9e1a9b 100644 --- a/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py @@ -399,9 +399,10 @@ def test_configured_dag_bundle_name_accepted(self, mock_manager): def test_build_command_scans_passed_roots_in_colocated_mode(self, tmp_path): binary = _build_bundle(tmp_path / "my_bundle", dag_ids=["tutorial_dag"]) coordinator = ExecutableCoordinator() - command, schema_version = coordinator._build_execute_task_command( - what=_make_ti(dag_id="tutorial_dag"), roots=[tmp_path] - ) + with coordinator._set_scan_roots([tmp_path]): + command, schema_version = coordinator._build_execute_task_command( + what=_make_ti(dag_id="tutorial_dag") + ) assert command == [str(binary.resolve())] assert schema_version == "2026-06-16" @@ -412,7 +413,8 @@ def test_returns_resolved_executable_and_schema_version(self, tmp_path): ti = _make_ti(dag_id="tutorial_dag") coordinator = ExecutableCoordinator(executables_root=[tmp_path]) - command, schema_version = coordinator._build_execute_task_command(what=ti, roots=[tmp_path]) + with coordinator._set_scan_roots([tmp_path]): + command, schema_version = coordinator._build_execute_task_command(what=ti) assert command == [str(binary.resolve())] assert schema_version == "2026-06-16" @@ -423,16 +425,22 @@ def test_raises_when_bundle_omits_schema_version(self, tmp_path): ti = _make_ti(dag_id="tutorial_dag") coordinator = ExecutableCoordinator(executables_root=[tmp_path]) - with pytest.raises(FileNotFoundError, match="matching bundles were rejected"): - coordinator._build_execute_task_command(what=ti, roots=[tmp_path]) + with ( + coordinator._set_scan_roots([tmp_path]), + pytest.raises(FileNotFoundError, match="matching bundles were rejected"), + ): + coordinator._build_execute_task_command(what=ti) def test_raises_when_dag_id_not_found(self, tmp_path): _build_bundle(tmp_path / "my_bundle", dag_ids=["other_dag"]) ti = _make_ti(dag_id="tutorial_dag") coordinator = ExecutableCoordinator(executables_root=[tmp_path]) - with pytest.raises(FileNotFoundError, match="cannot find executable bundle"): - coordinator._build_execute_task_command(what=ti, roots=[tmp_path]) + with ( + coordinator._set_scan_roots([tmp_path]), + pytest.raises(FileNotFoundError, match="cannot find executable bundle"), + ): + coordinator._build_execute_task_command(what=ti) @pytest.fixture diff --git a/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py index e84486150d96d..8aeb27abc88e6 100644 --- a/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py @@ -261,7 +261,8 @@ def test_unconfigured_dag_bundle_name_raises(self, mock_manager): def test_build_command_scans_passed_roots_in_colocated_mode(self, tmp_path): _make_jar(tmp_path / "app.jar", main_class="com.example.TaskRunner", schema_version="2026-06-16") coordinator = JavaCoordinator(main_class="com.example.TaskRunner") - command, schema_version = coordinator._build_execute_task_command(what=_make_ti(), roots=[tmp_path]) + with coordinator._set_scan_roots([tmp_path]): + command, schema_version = coordinator._build_execute_task_command(what=_make_ti()) assert command[0] == "java" assert command[-1] == "com.example.TaskRunner" assert schema_version == "2026-06-16" diff --git a/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py index 2e8ea8eed9001..57e25798b5aff 100644 --- a/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py @@ -102,7 +102,8 @@ def test_unconfigured_dag_bundle_name_raises(self, mock_manager): def test_build_command_scans_passed_roots_in_colocated_mode(self, tmp_path): bundle = write_bundle(tmp_path, "test_dag") coordinator = NodeCoordinator() - command, schema_version = coordinator._build_execute_task_command(what=_make_ti(), roots=[tmp_path]) + with coordinator._set_scan_roots([tmp_path]): + command, schema_version = coordinator._build_execute_task_command(what=_make_ti()) assert command == ["node", str(bundle)] assert schema_version == SCHEMA_VERSION @@ -129,7 +130,8 @@ def test_build_execute_task_command_returns_node_bundle_and_schema_version(self, bundles_root=tmp_path, ) - command, schema_version = coordinator._build_execute_task_command(what=_make_ti(), roots=[tmp_path]) + with coordinator._set_scan_roots([tmp_path]): + command, schema_version = coordinator._build_execute_task_command(what=_make_ti()) assert command == ["/opt/node/bin/node", str(bundle)] assert schema_version == SCHEMA_VERSION diff --git a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py index bf1220237bc16..06bcca9fa7d4a 100644 --- a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py +++ b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py @@ -588,8 +588,8 @@ class _StubSubprocessCoordinator(SubprocessCoordinator): def _explicit_artifact_roots(self) -> tuple[str, list[pathlib.Path]]: return "explicit_roots", self.explicit_roots - def _build_execute_task_command(self, *, what, roots): - self.recorded_roots.append(roots) + def _build_execute_task_command(self, *, what): + self.recorded_roots.append(list(self._get_scan_roots())) return list(self.command), self.schema_version @@ -615,7 +615,7 @@ class _Plain(SubprocessCoordinator): pass with pytest.raises(NotImplementedError): - _Plain()._build_execute_task_command(what=_make_ti(), roots=[]) + _Plain()._build_execute_task_command(what=_make_ti()) class TestSubprocessCoordinatorExecuteTask: @@ -885,7 +885,7 @@ def test_subclass_without_overrides_defaults_to_task_bundle(self): @attrs.define(kw_only=True) class _Bare(SubprocessCoordinator): - def _build_execute_task_command(self, *, what, roots): + def _build_execute_task_command(self, *, what): return [], None assert _Bare()._artifact_source is _ArtifactSource.TASK_BUNDLE @@ -1085,6 +1085,17 @@ def test_named_bundle_mode_locks_the_tree_it_scans( mock_lock.assert_called_once_with(bundle_name="artifacts", bundle_version="sha-abc") +class TestGetScanRoots: + def test_returns_bound_roots_and_clears_them_on_exit(self, tmp_path): + coordinator = _StubSubprocessCoordinator(command=["x"]) + + with coordinator._set_scan_roots([tmp_path]): + assert coordinator._get_scan_roots() == (tmp_path,) + + with pytest.raises(RuntimeError, match="requires an active task"): + coordinator._get_scan_roots() + + class TestSetCurrentBundle: def test_binds_for_the_block_and_clears_on_exit_allowing_reuse(self): coordinator = _StubSubprocessCoordinator(command=["/bin/true"]) From a6557be73bd3f48dc879a01ea43122c20a06c39f Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Mon, 17 Aug 2026 12:42:27 +0000 Subject: [PATCH 08/11] Make task bundle access dependencies explicit Module-level imports keep dependency analysis reliable and follow Airflow's import-placement contract without changing bundle access behavior. --- task-sdk/src/airflow/sdk/execution_time/bundles.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/task-sdk/src/airflow/sdk/execution_time/bundles.py b/task-sdk/src/airflow/sdk/execution_time/bundles.py index cfc1fe66cfb53..8ab131beb64f3 100644 --- a/task-sdk/src/airflow/sdk/execution_time/bundles.py +++ b/task-sdk/src/airflow/sdk/execution_time/bundles.py @@ -25,9 +25,11 @@ from __future__ import annotations import os +from getpass import getuser from typing import TYPE_CHECKING from airflow.dag_processing.bundles.manager import DagBundlesManager # noqa: SDK002 +from airflow.sdk.exceptions import AirflowException from airflow.sdk.execution_time.tracing import detail_span if TYPE_CHECKING: @@ -68,10 +70,6 @@ def verify_bundle_access(bundle_instance: BaseDagBundle) -> None: :param bundle_instance: The bundle instance to check :raises AirflowException: if bundle is not accessible """ - from getpass import getuser - - from airflow.sdk.exceptions import AirflowException - bundle_path = bundle_instance.path if not bundle_path.exists(): From bada0a79306b8e5db7daf62e8808d35b02cf2509 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Fri, 21 Aug 2026 07:05:30 +0000 Subject: [PATCH 09/11] Reject concurrent subprocess coordinator execution consistently Bundle resolution does not need mutable instance state, while overwriting scan roots can make a command build against another task's artifacts. --- .../airflow/sdk/coordinators/_subprocess.py | 28 ++------- .../task_sdk/coordinators/test_subprocess.py | 62 ++++++------------- 2 files changed, 25 insertions(+), 65 deletions(-) diff --git a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py index ffed5ad2e356c..9e69255ba3d07 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py +++ b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py @@ -446,7 +446,6 @@ class SubprocessCoordinator(BaseCoordinator): # The subclass's explicit root, recorded at construction so the base can # resolve roots without knowing the subclass field name. _configured_roots: list[pathlib.Path] = attrs.field(init=False, factory=list) - _active_bundle_info: BundleInfo | None = attrs.field(init=False, default=None) _active_scan_roots: tuple[pathlib.Path, ...] | None = attrs.field(init=False, default=None) @property @@ -500,7 +499,7 @@ def _classify_artifact_source(self) -> None: ) def _init_root_source( - self, logger: FilteringBoundLogger + self, bundle_info: BundleInfo, logger: FilteringBoundLogger ) -> tuple[list[pathlib.Path], BaseDagBundle | None]: """ Resolve the directories to scan for artifacts, dispatched on the classified mode. @@ -517,9 +516,7 @@ def _init_root_source( # NAMED_BUNDLE implies dag_bundle_name is set. target = BundleInfo(name=cast("str", self.dag_bundle_name)) else: - if self._active_bundle_info is None: - raise RuntimeError("_init_root_source requires an active task; call it during execute_task.") - target = self._active_bundle_info + target = bundle_info bundle = _initialize_pinned_bundle(target, logger) path = bundle.path @@ -547,25 +544,11 @@ def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], """ raise NotImplementedError - @contextlib.contextmanager - def _set_current_bundle(self, bundle_info: BundleInfo): - """ - Bind *bundle_info* as the active task for the duration of the block, clearing it on exit. - - Rejects a second concurrent bind: this coordinator runs one blocking task - per process and is not re-entrant. - """ - if self._active_bundle_info is not None: - raise RuntimeError("SubprocessCoordinator.execute_task is not re-entrant.") - self._active_bundle_info = bundle_info - try: - yield - finally: - self._active_bundle_info = None - @contextlib.contextmanager def _set_scan_roots(self, roots: Sequence[pathlib.Path]): """Expose *roots* to the command builder for the duration of the task.""" + if self._active_scan_roots is not None: + raise RuntimeError("SubprocessCoordinator.execute_task is not re-entrant.") self._active_scan_roots = tuple(roots) try: yield @@ -586,8 +569,7 @@ def execute_task( ) -> BaseCoordinator.ExecutionResult: task_logger = logger or log with contextlib.ExitStack() as stack: - stack.enter_context(self._set_current_bundle(bundle_info)) - roots, resolved_bundle = self._init_root_source(task_logger) + roots, resolved_bundle = self._init_root_source(bundle_info, task_logger) if resolved_bundle is not None: # Hold the version lock across start()/wait() so bundle cleanup # cannot rmtree a version this task is still reading from, diff --git a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py index 06bcca9fa7d4a..5297f56dead4f 100644 --- a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py +++ b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py @@ -918,29 +918,23 @@ class TestInitRootSource: def test_returns_configured_root_and_no_bundle_in_explicit_mode(self): configured = [pathlib.Path("/a"), pathlib.Path("/b")] coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=configured) - roots, bundle = coordinator._init_root_source(log) + roots, bundle = coordinator._init_root_source(MagicMock(), log) assert roots == configured assert bundle is None @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle") - def test_task_bundle_mode_uses_active_task_bundle(self, mock_initialize, tmp_path): + def test_task_bundle_mode_uses_passed_task_bundle(self, mock_initialize, tmp_path): resolved = MagicMock(path=tmp_path, version="v3") mock_initialize.return_value = resolved coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=[]) bundle_info = BundleInfo(name="dags", version="v3", version_data={"k": "v"}) - coordinator._active_bundle_info = bundle_info - roots, bundle = coordinator._init_root_source(log) + roots, bundle = coordinator._init_root_source(bundle_info, log) assert roots == [tmp_path] assert bundle is resolved mock_initialize.assert_called_once_with(bundle_info) - def test_task_bundle_mode_without_active_task_raises(self): - coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=[]) - with pytest.raises(RuntimeError, match="requires an active task"): - coordinator._init_root_source(log) - @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle") def test_version_less_bundle_is_rematerialized_at_its_current_version(self, mock_initialize, tmp_path): """A bundle resolved without a version is re-resolved at the version current now. @@ -962,7 +956,7 @@ def test_version_less_bundle_is_rematerialized_at_its_current_version(self, mock coordinator.dag_bundle_name = "artifacts" coordinator._artifact_source = _ArtifactSource.NAMED_BUNDLE - roots, bundle = coordinator._init_root_source(log) + roots, bundle = coordinator._init_root_source(MagicMock(), log) assert roots == [pinned_tree] assert bundle is pinned @@ -981,7 +975,7 @@ def test_unversioned_bundle_keeps_its_single_path(self, mock_initialize, tmp_pat coordinator.dag_bundle_name = "artifacts" coordinator._artifact_source = _ArtifactSource.NAMED_BUNDLE - roots, bundle = coordinator._init_root_source(log) + roots, bundle = coordinator._init_root_source(MagicMock(), log) assert roots == [tmp_path] assert bundle is resolved @@ -996,11 +990,11 @@ def test_missing_resolved_path_raises(self, mock_initialize, tmp_path): coordinator._artifact_source = _ArtifactSource.NAMED_BUNDLE with pytest.raises(FileNotFoundError, match="does not exist"): - coordinator._init_root_source(log) + coordinator._init_root_source(MagicMock(), log) class TestExecuteTaskBundleWiring: - """execute_task binds the task bundle, forwards resolved roots, and holds the version lock.""" + """execute_task passes the task bundle, forwards resolved roots, and holds the version lock.""" @patch("airflow.sdk.coordinators._subprocess.BundleVersionLock") @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle") @@ -1024,11 +1018,8 @@ def test_task_bundle_mode_binds_forwards_roots_and_locks( subprocess_logs_to_stdout=False, ) - # Resolution runs against the bound task bundle. mock_initialize.assert_called_once_with(bundle_info) - # The resolved root is forwarded to the subclass command builder. assert coordinator.recorded_roots == [[tmp_path]] - # The version lock is held around the subprocess for the resolved bundle. mock_lock.assert_called_once_with(bundle_name="dags", bundle_version="v9") mock_lock.return_value.__enter__.assert_called_once() mock_lock.return_value.__exit__.assert_called_once() @@ -1095,34 +1086,21 @@ def test_returns_bound_roots_and_clears_them_on_exit(self, tmp_path): with pytest.raises(RuntimeError, match="requires an active task"): coordinator._get_scan_roots() - -class TestSetCurrentBundle: - def test_binds_for_the_block_and_clears_on_exit_allowing_reuse(self): + def test_rejects_nested_reentry(self, tmp_path): coordinator = _StubSubprocessCoordinator(command=["/bin/true"]) - bundle = MagicMock() - with coordinator._set_current_bundle(bundle): - assert coordinator._active_bundle_info is bundle - assert coordinator._active_bundle_info is None - # A subsequent task on the same (reused) coordinator binds cleanly. - with coordinator._set_current_bundle(MagicMock()): - assert coordinator._active_bundle_info is not None - assert coordinator._active_bundle_info is None - - def test_rejects_nested_reentry(self): - coordinator = _StubSubprocessCoordinator(command=["/bin/true"]) - with coordinator._set_current_bundle(MagicMock()): + with coordinator._set_scan_roots([tmp_path]): with pytest.raises(RuntimeError, match="not re-entrant"): - with coordinator._set_current_bundle(MagicMock()): + with coordinator._set_scan_roots([tmp_path]): pass - def test_execute_task_is_not_reentrant(self, mock_client): + def test_execute_task_is_not_reentrant(self, mock_client, tmp_path): coordinator = _StubSubprocessCoordinator(command=["/bin/true"]) - coordinator._active_bundle_info = MagicMock() # simulate an in-flight task - with pytest.raises(RuntimeError, match="not re-entrant"): - coordinator.execute_task( - what=_make_ti(), - dag_rel_path="dag.py", - bundle_info=MagicMock(), - client=mock_client, - subprocess_logs_to_stdout=False, - ) + with coordinator._set_scan_roots([tmp_path]): + with pytest.raises(RuntimeError, match="not re-entrant"): + coordinator.execute_task( + what=_make_ti(), + dag_rel_path="dag.py", + bundle_info=MagicMock(), + client=mock_client, + subprocess_logs_to_stdout=False, + ) From e8e1a8ba50605f971cf85ab954b0a9dd67a7491d Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Fri, 21 Aug 2026 07:15:11 +0000 Subject: [PATCH 10/11] Reject empty subprocess coordinator artifact roots An explicitly empty root is likely a configuration or templating error. Treating it as omitted can silently scan the task's Dag bundle and run unintended artifacts. --- .../language-sdks/go.rst | 3 +- .../language-sdks/java.rst | 3 +- .../language-sdks/typescript.rst | 3 +- .../sdk/coordinators/_bundle_metadata.py | 28 ++++++++++++++++++- .../coordinators/executable/coordinator.py | 7 +++-- .../sdk/coordinators/java/coordinator.py | 10 +++++-- .../sdk/coordinators/node/coordinator.py | 11 ++++++-- .../executable/test_coordinator.py | 7 +++-- .../coordinators/java/test_coordinator.py | 5 ++++ .../coordinators/node/test_coordinator.py | 12 ++++---- 10 files changed, 67 insertions(+), 22 deletions(-) diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst index 965fabbc292f6..bd2a3e4ac761a 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst @@ -510,7 +510,8 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the - *(optional)* - One or more directories scanned recursively for executable bundles. Accepts a string, a path, or a list of strings/paths. When omitted, bundles are located through a Dag - bundle instead (see the note below). + bundle instead (see the note below). Explicitly setting this option to ``null`` or + an empty list is invalid. * - ``dag_bundle_name`` - *(auto: task's own bundle)* - Name of a configured Dag bundle to load executable bundles from. Mutually exclusive diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst index 2587bd7950a68..0ad2850896db3 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst @@ -1025,7 +1025,8 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the - *(optional)* - One or more directories scanned recursively for ``.jar`` files. Accepts a string, a path, or a list of strings/paths. When omitted, JARs are located through a Dag - bundle instead (see the note below). + bundle instead (see the note below). Explicitly setting this option to ``null`` or + an empty list is invalid. * - ``dag_bundle_name`` - *(auto: task's own bundle)* - Name of a configured Dag bundle to load JARs from. Mutually exclusive with diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst index 376771ff4cddd..344229f4f2d95 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst @@ -449,7 +449,8 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the - *(optional)* - One or more directories searched recursively, in order, for an integrity-verified ``*.min.mjs`` bundle that declares the requested Dag. Accepts a string, a path, or a list of strings/paths. When - omitted, the bundle is located through a Dag bundle instead (see the note below). + omitted, the bundle is located through a Dag bundle instead (see the note below). Explicitly setting + this option to ``null`` or an empty list is invalid. * - ``dag_bundle_name`` - *(auto: task's own bundle)* - Name of a configured Dag bundle to load the ``*.min.mjs`` bundle from. Mutually exclusive with diff --git a/task-sdk/src/airflow/sdk/coordinators/_bundle_metadata.py b/task-sdk/src/airflow/sdk/coordinators/_bundle_metadata.py index 9a0ee5e5b32ae..b542483262385 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_bundle_metadata.py +++ b/task-sdk/src/airflow/sdk/coordinators/_bundle_metadata.py @@ -22,7 +22,7 @@ import os import pathlib import stat -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Final import attrs import structlog @@ -38,6 +38,13 @@ log: FilteringBoundLogger = structlog.get_logger(logger_name="coordinators") +class _UnsetArtifactRoots: + """Sentinel distinguishing an omitted artifact-roots option from an empty value.""" + + +ARTIFACT_ROOTS_NOT_CONFIGURED: Final = _UnsetArtifactRoots() + + def convert_roots( value: None | os.PathLike[str] | pathlib.Path | list[os.PathLike[str] | pathlib.Path], ) -> list[pathlib.Path]: @@ -96,6 +103,25 @@ def _sorted_children(directory: pathlib.Path) -> list[pathlib.Path]: return [] +def convert_configured_roots( + value: _UnsetArtifactRoots + | None + | os.PathLike[str] + | pathlib.Path + | list[os.PathLike[str] | pathlib.Path], +) -> list[pathlib.Path]: + """Normalize configured roots while rejecting explicitly empty values.""" + if isinstance(value, _UnsetArtifactRoots): + return [] + roots = convert_roots(value) + if not roots: + raise ValueError( + "Artifact roots must contain at least one path when provided; " + "omit the option to use the task's Dag bundle." + ) + return roots + + def validate_schema_version(instance, _, value) -> str: """Attrs validator resolving a bundle's supervisor schema version to a known one.""" return get_schema_version_migrator().resolve_version(str(value)) diff --git a/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py index 59e05300b8568..9d6b80b249f81 100644 --- a/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py @@ -31,8 +31,9 @@ import structlog from airflow.sdk.coordinators._bundle_metadata import ( + ARTIFACT_ROOTS_NOT_CONFIGURED, ResolvedBundle, - convert_roots, + convert_configured_roots, extract_supervisor_schema_version, parse_metadata_mapping, ) @@ -346,8 +347,8 @@ class ExecutableCoordinator(SubprocessCoordinator): """ executables_root: list[pathlib.Path] = attrs.field( - converter=convert_roots, - factory=list, + default=ARTIFACT_ROOTS_NOT_CONFIGURED, + converter=convert_configured_roots, ) @property diff --git a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py index c5db8557cc19d..1f0f6e6efd5f9 100644 --- a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py @@ -29,7 +29,11 @@ import attrs import structlog -from airflow.sdk.coordinators._bundle_metadata import convert_roots, validate_schema_version +from airflow.sdk.coordinators._bundle_metadata import ( + ARTIFACT_ROOTS_NOT_CONFIGURED, + convert_configured_roots, + validate_schema_version, +) from airflow.sdk.coordinators._subprocess import SubprocessCoordinator if TYPE_CHECKING: @@ -200,8 +204,8 @@ class JavaCoordinator(SubprocessCoordinator): java_executable: str = "java" jvm_args: list[str] = attrs.field(factory=list) jars_root: list[pathlib.Path] = attrs.field( - converter=convert_roots, - factory=list, + default=ARTIFACT_ROOTS_NOT_CONFIGURED, + converter=convert_configured_roots, ) main_class: str = "" diff --git a/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py index 35aff8181d3e5..dd2a3f6f0416f 100644 --- a/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py @@ -26,7 +26,12 @@ import attrs import structlog -from airflow.sdk.coordinators._bundle_metadata import ResolvedBundle, convert_roots, walk_files +from airflow.sdk.coordinators._bundle_metadata import ( + ARTIFACT_ROOTS_NOT_CONFIGURED, + ResolvedBundle, + convert_configured_roots, + walk_files, +) from airflow.sdk.coordinators._subprocess import SubprocessCoordinator from airflow.sdk.coordinators.node._bundle_reader import read_bundle @@ -119,8 +124,8 @@ class NodeCoordinator(SubprocessCoordinator): node_executable: str = "node" bundles_root: list[pathlib.Path] = attrs.field( - converter=convert_roots, - factory=list, + default=ARTIFACT_ROOTS_NOT_CONFIGURED, + converter=convert_configured_roots, ) @property diff --git a/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py index 502456d9e1a9b..2fb1ea96d097d 100644 --- a/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py @@ -375,9 +375,10 @@ def test_executables_root_optional_defaults_to_empty(self): assert coordinator.executables_root == [] assert coordinator.dag_bundle_name is None - def test_none_executables_root_normalized_to_empty(self): - coordinator = ExecutableCoordinator(executables_root=None) - assert coordinator.executables_root == [] + @pytest.mark.parametrize("executables_root", [None, []], ids=["none", "empty-list"]) + def test_explicit_empty_executables_root_raises(self, executables_root): + with pytest.raises(ValueError, match="must contain at least one path when provided"): + ExecutableCoordinator(executables_root=executables_root) def test_root_and_dag_bundle_name_are_mutually_exclusive(self, tmp_path): with pytest.raises(ValueError, match="at most one of 'executables_root' or 'dag_bundle_name'"): diff --git a/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py index 8aeb27abc88e6..5d4dab326c5fe 100644 --- a/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py @@ -244,6 +244,11 @@ def test_jars_root_optional_defaults_to_empty(self): assert coordinator.dag_bundle_name is None assert coordinator.main_class == "" + @pytest.mark.parametrize("jars_root", [None, []], ids=["none", "empty-list"]) + def test_explicit_empty_jars_root_raises(self, jars_root): + with pytest.raises(ValueError, match="must contain at least one path when provided"): + JavaCoordinator(jars_root=jars_root) + def test_explicit_root_does_not_require_main_class(self, tmp_path): coordinator = JavaCoordinator(jars_root=tmp_path) assert coordinator.main_class == "" diff --git a/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py index 57e25798b5aff..4aa1c4608c711 100644 --- a/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py @@ -85,9 +85,10 @@ def test_bundles_root_optional_defaults_to_empty(self): assert coordinator.bundles_root == [] assert coordinator.dag_bundle_name is None - def test_none_bundles_root_normalized_to_empty(self): - coordinator = NodeCoordinator(bundles_root=None) - assert coordinator.bundles_root == [] + @pytest.mark.parametrize("bundles_root", [None, []], ids=["none", "empty-list"]) + def test_explicit_empty_bundles_root_raises(self, bundles_root): + with pytest.raises(ValueError, match="must contain at least one path when provided"): + NodeCoordinator(bundles_root=bundles_root) def test_root_and_dag_bundle_name_are_mutually_exclusive(self): with pytest.raises(ValueError, match="at most one of 'bundles_root' or 'dag_bundle_name'"): @@ -116,9 +117,8 @@ def test_selects_bundle_by_dag_id(self, tmp_path): bundles_root=tmp_path, ) - command, schema_version = coordinator._build_execute_task_command( - what=_make_ti(dag_id="sales"), roots=[tmp_path] - ) + with coordinator._set_scan_roots([tmp_path]): + command, schema_version = coordinator._build_execute_task_command(what=_make_ti(dag_id="sales")) assert command == ["/opt/node/bin/node", str(selected)] assert schema_version == SCHEMA_VERSION From 18334707dec6270b3ddd418f48fee87b1a57b223 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Mon, 28 Sep 2026 14:53:55 +0000 Subject: [PATCH 11/11] Reject blank subprocess coordinator artifact roots An empty string normalizes to the current directory, so it bypassed the empty-roots check and scanned the worker's working directory. --- task-sdk/src/airflow/sdk/coordinators/_bundle_metadata.py | 8 ++++---- .../task_sdk/coordinators/executable/test_coordinator.py | 8 ++++++-- .../tests/task_sdk/coordinators/java/test_coordinator.py | 8 ++++++-- .../tests/task_sdk/coordinators/node/test_coordinator.py | 8 ++++++-- 4 files changed, 22 insertions(+), 10 deletions(-) diff --git a/task-sdk/src/airflow/sdk/coordinators/_bundle_metadata.py b/task-sdk/src/airflow/sdk/coordinators/_bundle_metadata.py index b542483262385..b60aa94144bbb 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_bundle_metadata.py +++ b/task-sdk/src/airflow/sdk/coordinators/_bundle_metadata.py @@ -113,13 +113,13 @@ def convert_configured_roots( """Normalize configured roots while rejecting explicitly empty values.""" if isinstance(value, _UnsetArtifactRoots): return [] - roots = convert_roots(value) - if not roots: + values = [value] if isinstance(value, (str, os.PathLike)) else (value or []) + if not values or any(not os.fspath(v).strip() for v in values): raise ValueError( - "Artifact roots must contain at least one path when provided; " + "Artifact roots must contain at least one path when provided, and each path must be non-empty; " "omit the option to use the task's Dag bundle." ) - return roots + return convert_roots(value) def validate_schema_version(instance, _, value) -> str: diff --git a/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py index 2fb1ea96d097d..c20913e26eae3 100644 --- a/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py @@ -375,9 +375,13 @@ def test_executables_root_optional_defaults_to_empty(self): assert coordinator.executables_root == [] assert coordinator.dag_bundle_name is None - @pytest.mark.parametrize("executables_root", [None, []], ids=["none", "empty-list"]) + @pytest.mark.parametrize( + "executables_root", + [None, [], "", " ", [""]], + ids=["none", "empty-list", "empty-str", "blank-str", "list-of-empty-str"], + ) def test_explicit_empty_executables_root_raises(self, executables_root): - with pytest.raises(ValueError, match="must contain at least one path when provided"): + with pytest.raises(ValueError, match="and each path must be non-empty"): ExecutableCoordinator(executables_root=executables_root) def test_root_and_dag_bundle_name_are_mutually_exclusive(self, tmp_path): diff --git a/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py index 5d4dab326c5fe..c55a2cc891253 100644 --- a/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py @@ -244,9 +244,13 @@ def test_jars_root_optional_defaults_to_empty(self): assert coordinator.dag_bundle_name is None assert coordinator.main_class == "" - @pytest.mark.parametrize("jars_root", [None, []], ids=["none", "empty-list"]) + @pytest.mark.parametrize( + "jars_root", + [None, [], "", " ", [""]], + ids=["none", "empty-list", "empty-str", "blank-str", "list-of-empty-str"], + ) def test_explicit_empty_jars_root_raises(self, jars_root): - with pytest.raises(ValueError, match="must contain at least one path when provided"): + with pytest.raises(ValueError, match="and each path must be non-empty"): JavaCoordinator(jars_root=jars_root) def test_explicit_root_does_not_require_main_class(self, tmp_path): diff --git a/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py index 4aa1c4608c711..146cb91c1d94c 100644 --- a/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py @@ -85,9 +85,13 @@ def test_bundles_root_optional_defaults_to_empty(self): assert coordinator.bundles_root == [] assert coordinator.dag_bundle_name is None - @pytest.mark.parametrize("bundles_root", [None, []], ids=["none", "empty-list"]) + @pytest.mark.parametrize( + "bundles_root", + [None, [], "", " ", [""]], + ids=["none", "empty-list", "empty-str", "blank-str", "list-of-empty-str"], + ) def test_explicit_empty_bundles_root_raises(self, bundles_root): - with pytest.raises(ValueError, match="must contain at least one path when provided"): + with pytest.raises(ValueError, match="and each path must be non-empty"): NodeCoordinator(bundles_root=bundles_root) def test_root_and_dag_bundle_name_are_mutually_exclusive(self):