Support DagBundle based deployment for Coordinator (#70805)

* 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.

* Document dag_bundle_name option for language-SDK coordinators

* 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.

* 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.

* 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.

* 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.

* 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.

* 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.

* 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.

* 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.

* 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.
This commit is contained in:
Jason(Zhe-You) Liu authored and GitHub committed 2026-09-30 16:10:25 +08:00
1 parent 76bf88a487
commit 1a67bf46b5
25 files changed
+1250 -346

No files matched your search

@@ -507,14 +507,32 @@ 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). 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
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, 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.
.. _go-sdk/limitations:
Limitations
@@ -1045,9 +1045,15 @@ 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). 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
``jars_root``.
* - ``java_executable``
- ``"java"``
- Path to the ``java`` binary. Defaults to ``java`` on ``$PATH``.
@@ -1056,15 +1062,27 @@ 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
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, 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.
.. note::
The ``[sdk]`` configuration is read at startup, so changes to ``coordinators`` or
@@ -482,9 +482,15 @@ 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). 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
``bundles_root``.
* - ``node_executable``
- ``"node"``
- Path to the ``node`` binary. Defaults to ``node`` on ``$PATH``.
@@ -493,6 +499,18 @@ 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, 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.
Limitations
-----------
@@ -16,13 +16,14 @@
# under the License.
from __future__ import annotations
import functools
import importlib
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
@@ -107,6 +108,28 @@ 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]
@functools.cache
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(
configs=tuple(bundle_config_list),
names=frozenset(cfg.name for cfg in bundle_config_list),
)
def _parse_bundle_config(config_list) -> list[_ExternalBundleConfig]:
bundles = {}
for item in config_list:
@@ -300,12 +323,9 @@ class DagBundlesManager(LoggingMixin):
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"):
@@ -673,6 +693,16 @@ class DagBundlesManager(LoggingMixin):
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.
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]:
"""
Get all DAG bundles.
@@ -330,7 +330,7 @@ class TestOtelIntegration:
# 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",
@@ -33,6 +33,7 @@ from airflow.dag_processing.bundles.manager import (
DagBundlesManager,
_get_configured_bundle_team_names,
_guess_best_bundle_for_fileloc,
_load_bundle_config_snapshot,
)
from airflow.exceptions import AirflowConfigException
from airflow.models.dag import DagModel
@@ -205,6 +206,55 @@ def test_get_bundle():
assert bundle.version is None
@pytest.fixture(autouse=True)
def _clear_bundle_config_snapshot_cache():
_load_bundle_config_snapshot.cache_clear()
yield
_load_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, _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
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, _bundle_config_env(config)):
assert DagBundlesManager.is_bundle_configured("artifacts") is True
def test_is_bundle_configured_empty_config():
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
@pytest.fixture
def clear_db():
clear_db_dag_bundles()
@@ -221,9 +271,7 @@ def test_sync_bundles_to_db(clear_db, session):
).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)]
@@ -248,9 +296,7 @@ def test_sync_bundles_to_db(clear_db, session):
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() == [
@@ -298,16 +344,12 @@ def test_sync_bundles_to_db_partial_config_does_not_disable_other_bundles(clear_
).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),
@@ -315,9 +357,7 @@ def test_sync_bundles_to_db_partial_config_does_not_disable_other_bundles(clear_
]
# 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),
@@ -416,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()
@@ -441,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()
@@ -460,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()
@@ -485,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()
@@ -120,6 +120,12 @@ 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.
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
~~~~~~~~~~~~~~~~~
@@ -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
@@ -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
+1 -1
View File
@@ -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
@@ -22,7 +22,7 @@ from __future__ import annotations
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 @@ if TYPE_CHECKING:
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 []
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, and each path must be non-empty; "
"omit the option to use the task's Dag bundle."
)
return convert_roots(value)
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))
@@ -27,6 +27,8 @@ draining machinery in this module rather than re-implementing it.
from __future__ import annotations
import contextlib
import enum
import ipaddress
import itertools
import os
@@ -41,18 +43,24 @@ 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
if TYPE_CHECKING:
import pathlib
from collections.abc import Sequence
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
from airflow.sdk.api.datamodels._generated import TaskInstance
Tracked = TypeVar("Tracked", socket.socket, subprocess.Popen)
@@ -371,6 +379,45 @@ class _PopenActivitySubprocess(ActivitySubprocess):
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/<version>`` 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 version current at task start is used."""
TASK_BUNDLE = enum.auto()
"""
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)
class SubprocessCoordinator(BaseCoordinator):
"""
@@ -385,23 +432,129 @@ 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
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
_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)
_active_scan_roots: tuple[pathlib.Path, ...] | None = attrs.field(init=False, default=None)
@property
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()
def _classify_artifact_source(self) -> None:
"""
Classify and validate how this coordinator locates artifacts (construction time).
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 "
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
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
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, bundle_info: BundleInfo, logger: FilteringBoundLogger
) -> tuple[list[pathlib.Path], BaseDagBundle | None]:
"""
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, None
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))
else:
target = bundle_info
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.")
return [path], bundle
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*.
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.
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
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_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
finally:
self._active_scan_roots = None
def execute_task(
self,
*,
@@ -414,18 +567,32 @@ class SubprocessCoordinator(BaseCoordinator):
subprocess_logs_to_stdout: bool,
**kwargs,
) -> BaseCoordinator.ExecutionResult:
command, subprocess_schema_version = self._build_execute_task_command(what=what)
process = _PopenActivitySubprocess.start(
what=what,
dag_rel_path=dag_rel_path,
bundle_info=bundle_info,
client=client,
logger=logger,
subprocess_logs_to_stdout=subprocess_logs_to_stdout,
sentry_integration=sentry_integration,
command=command,
subprocess_schema_version=subprocess_schema_version,
startup_timeout=self.task_startup_timeout,
)
exit_code = process.wait()
return self.ExecutionResult(exit_code, process.final_state)
task_logger = logger or log
with contextlib.ExitStack() as stack:
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,
# mirroring task_runner.main() for the Python task path.
stack.enter_context(
BundleVersionLock(
bundle_name=resolved_bundle.name,
bundle_version=resolved_bundle.version,
)
)
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,
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)
@@ -31,8 +31,9 @@ import attrs
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,
)
@@ -339,16 +340,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),
default=ARTIFACT_ROOTS_NOT_CONFIGURED,
converter=convert_configured_roots,
)
@property
def _explicit_artifact_roots(self) -> tuple[str, list[pathlib.Path]]:
return "executables_root", self.executables_root
def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], str | None]:
bundle = _Bundle.find(self.executables_root, what.dag_id)
roots = self._get_scan_roots()
bundle = _Bundle.find(roots, what.dag_id)
return [str(bundle.path)], bundle.schema_version
@@ -29,7 +29,11 @@ from typing import TYPE_CHECKING
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:
@@ -172,7 +176,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.
@@ -199,17 +204,24 @@ 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,
validator=attrs.validators.min_len(1),
default=ARTIFACT_ROOTS_NOT_CONFIGURED,
converter=convert_configured_roots,
)
main_class: str = ""
@property
def _explicit_artifact_roots(self) -> tuple[str, list[pathlib.Path]]:
return "jars_root", self.jars_root
def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], str | None]:
jar = _JarInfo.find(self.jars_root, self.main_class)
# 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,
"-classpath",
_calculate_classpath(self.jars_root),
_calculate_classpath(roots),
*self.jvm_args,
jar.main_class,
]
@@ -26,7 +26,12 @@ from typing import TYPE_CHECKING
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
@@ -111,17 +116,23 @@ 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.
"""
node_executable: str = "node"
bundles_root: list[pathlib.Path] = attrs.field(
converter=convert_roots,
validator=attrs.validators.min_len(1),
default=ARTIFACT_ROOTS_NOT_CONFIGURED,
converter=convert_configured_roots,
)
@property
def _explicit_artifact_roots(self) -> tuple[str, list[pathlib.Path]]:
return "bundles_root", self.bundles_root
def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], str | None]:
bundle = _Bundle.find(self.bundles_root, what.dag_id)
roots = self._get_scan_roots()
bundle = _Bundle.find(roots, what.dag_id)
return [self.node_executable, os.fspath(bundle.path)], bundle.schema_version
@@ -0,0 +1,90 @@
# 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 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:
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
"""
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"
)
@@ -44,10 +44,8 @@ from pydantic import AwareDatetime, ConfigDict, Field, JsonValue, TypeAdapter
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 (
@@ -86,6 +84,7 @@ from airflow.sdk.exceptions import (
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,
@@ -148,6 +147,7 @@ from airflow.sdk.execution_time.email_backend import (
_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
@@ -171,38 +171,6 @@ log = structlog.get_logger("task")
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 = (
@@ -1034,13 +1002,7 @@ def parse(what: StartupDetails, log: Logger) -> RuntimeTaskInstance:
bundle_info = what.bundle_info
bundle_prepare_start = time.monotonic()
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)
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))
@@ -1136,43 +1098,6 @@ SUPERVISOR_COMMS: CommsDecoder[ToTask, ToSupervisor]
# 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.
@@ -1280,8 +1205,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
@@ -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
@@ -370,13 +370,46 @@ class TestExecutableCoordinatorAttributes:
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_optional_defaults_to_empty(self):
coordinator = ExecutableCoordinator()
assert coordinator.executables_root == []
assert coordinator.dag_bundle_name is None
def test_executables_root_must_be_non_empty(self):
with pytest.raises(ValueError, match="executables_root"):
ExecutableCoordinator(executables_root=None)
@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="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):
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.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.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")
assert coordinator.dag_bundle_name == "artifacts"
mock_manager.is_bundle_configured.assert_called_once_with("artifacts")
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()
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"
class TestBuildExecuteTaskCommand:
@@ -385,7 +418,8 @@ class TestBuildExecuteTaskCommand:
ti = _make_ti(dag_id="tutorial_dag")
coordinator = ExecutableCoordinator(executables_root=[tmp_path])
command, schema_version = coordinator._build_execute_task_command(what=ti)
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"
@@ -396,7 +430,10 @@ class TestBuildExecuteTaskCommand:
ti = _make_ti(dag_id="tutorial_dag")
coordinator = ExecutableCoordinator(executables_root=[tmp_path])
with pytest.raises(FileNotFoundError, match="matching bundles were rejected"):
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):
@@ -404,7 +441,10 @@ class TestBuildExecuteTaskCommand:
ti = _make_ti(dag_id="tutorial_dag")
coordinator = ExecutableCoordinator(executables_root=[tmp_path])
with pytest.raises(FileNotFoundError, match="cannot find executable bundle"):
with (
coordinator._set_scan_roots([tmp_path]),
pytest.raises(FileNotFoundError, match="cannot find executable bundle"),
):
coordinator._build_execute_task_command(what=ti)
@@ -237,6 +237,45 @@ class TestJavaCoordinatorAttributes:
assert coordinator.jvm_args == ["-Xmx512m", "-Xms256m"]
assert coordinator.jars_root == [pathlib.Path("/airflow/java-bundles")]
def test_jars_root_optional_defaults_to_empty(self):
# 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
assert coordinator.main_class == ""
@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="and each path must be non-empty"):
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 == ""
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.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):
_make_jar(tmp_path / "app.jar", main_class="com.example.TaskRunner", schema_version="2026-06-16")
coordinator = JavaCoordinator(main_class="com.example.TaskRunner")
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"
@pytest.fixture
def jars_root(tmp_path):
@@ -21,6 +21,7 @@ from __future__ import annotations
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,37 @@ class TestNodeCoordinatorAttributes:
]
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):
coordinator = NodeCoordinator()
assert coordinator.bundles_root == []
assert coordinator.dag_bundle_name is None
@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="and each path must be non-empty"):
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'"):
NodeCoordinator(bundles_root="/airflow/ts-bundles", dag_bundle_name="artifacts")
@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):
bundle = write_bundle(tmp_path, "test_dag")
coordinator = NodeCoordinator()
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
class TestNodeCoordinatorExecuteTaskCommand:
@@ -92,11 +121,25 @@ class TestNodeCoordinatorExecuteTaskCommand:
bundles_root=tmp_path,
)
command, schema_version = coordinator._build_execute_task_command(what=_make_ti(dag_id="sales"))
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
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,
)
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
class TestBundleFind:
@pytest.mark.parametrize(
@@ -19,6 +19,7 @@ from __future__ import annotations
import contextlib
import os
import pathlib
import socket
import subprocess
import sys
@@ -31,16 +32,19 @@ import psutil
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 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
@@ -567,12 +571,25 @@ class TestResourceTracker:
@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)
@property
def _explicit_artifact_roots(self) -> tuple[str, list[pathlib.Path]]:
return "explicit_roots", self.explicit_roots
def _build_execute_task_command(self, *, what):
self.recorded_roots.append(list(self._get_scan_roots()))
return list(self.command), self.schema_version
@@ -844,3 +861,246 @@ class TestPopenActivitySubprocessStart:
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.
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):
configured = [pathlib.Path("/artifacts")]
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 = _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):
return [], None
assert _Bare()._artifact_source is _ArtifactSource.TASK_BUNDLE
def test_rejects_explicit_root_and_dag_bundle_name_together(self):
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.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.sdk.coordinators._subprocess.DagBundlesManager")
def test_accepts_configured_dag_bundle_name(self, mock_manager):
mock_manager.is_bundle_configured.return_value = True
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")
class TestInitRootSource:
"""Execute-time root resolution, dispatched on the classified mode."""
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(MagicMock(), log)
assert roots == configured
assert bundle is None
@patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle")
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"})
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)
@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.dag_bundle_name = "artifacts"
coordinator._artifact_source = _ArtifactSource.NAMED_BUNDLE
roots, bundle = coordinator._init_root_source(MagicMock(), 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
roots, bundle = coordinator._init_root_source(MagicMock(), log)
assert roots == [tmp_path]
assert bundle is resolved
mock_initialize.assert_called_once_with(BundleInfo(name="artifacts"))
@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(path=missing, version="v1")
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(MagicMock(), log)
class TestExecuteTaskBundleWiring:
"""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")
@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(path=tmp_path, version="v9")
resolved.name = "dags"
mock_initialize.return_value = resolved
mock_start.return_value.wait.return_value = 0
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,
subprocess_logs_to_stdout=False,
)
mock_initialize.assert_called_once_with(bundle_info)
assert coordinator.recorded_roots == [[tmp_path]]
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.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
):
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()
@patch("airflow.sdk.coordinators._subprocess.BundleVersionLock")
@patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle")
@patch.object(_PopenActivitySubprocess, "start")
def test_named_bundle_mode_locks_the_tree_it_scans(
self, mock_start, mock_initialize, mock_lock, mock_client, tmp_path
):
"""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=[])
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,
)
assert coordinator.recorded_roots == [[pinned_tree]]
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()
def test_rejects_nested_reentry(self, tmp_path):
coordinator = _StubSubprocessCoordinator(command=["/bin/true"])
with coordinator._set_scan_roots([tmp_path]):
with pytest.raises(RuntimeError, match="not re-entrant"):
with coordinator._set_scan_roots([tmp_path]):
pass
def test_execute_task_is_not_reentrant(self, mock_client, tmp_path):
coordinator = _StubSubprocessCoordinator(command=["/bin/true"])
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,
)
@@ -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()
@@ -181,7 +181,6 @@ from airflow.sdk.execution_time.task_runner import (
_run_execute_callable,
_serialize_outlet_events,
_xcom_push,
detail_span,
finalize,
get_startup_details,
parse,
@@ -838,56 +837,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())
@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 +6137,6 @@ class TestTaskInstanceMetrics:
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 +6218,7 @@ class TestRunExecuteCallable:
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 +6246,7 @@ class TestRunExecuteCallable:
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 +6265,7 @@ class TestRunExecuteCallable:
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)
@@ -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")