From bd9be97f70e84d3d7d46a037d248ce9d84c82f80 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 08:53:13 +0000 Subject: [PATCH 01/94] Replace coordinator artifact roots with task_handler_bundle_name ADR-0013 locates Lang-SDK task handlers through a named Dag bundle, so the unversioned jars_root, executables_root and bundles_root directories go away instead of living beside it. dag_bundle_name becomes task_handler_bundle_name because the native-Dag work will use the old name with a different meaning. It has not shipped in a release, so there is no shim. A coordinator reads handlers from the named bundle at its current version, or from the task's own Dag bundle at the run's version, pinned for the task either way. CoordinatorManager.from_config validates the name of every spec without constructing coordinators, so a typo fails at config load instead of on the first task. --- .../src/airflow/config_templates/config.yml | 10 +- .../sdk/coordinators/_bundle_metadata.py | 45 +---- .../airflow/sdk/coordinators/_subprocess.py | 131 +++--------- .../coordinators/executable/coordinator.py | 30 +-- .../sdk/coordinators/java/coordinator.py | 34 ++-- .../sdk/coordinators/node/coordinator.py | 31 +-- .../airflow/sdk/execution_time/coordinator.py | 18 +- .../executable/test_coordinator.py | 73 +++---- .../coordinators/java/test_coordinator.py | 92 ++++----- .../coordinators/node/test_coordinator.py | 45 +---- .../task_sdk/coordinators/test_subprocess.py | 190 +++++------------- .../execution_time/test_coordinator.py | 31 ++- 12 files changed, 232 insertions(+), 498 deletions(-) diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index 34849d5eb5e13..0d93997abe387 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -2191,6 +2191,12 @@ sdk: example, two ``JavaCoordinator`` instances pinned to different JDK versions). + The language SDK coordinators read their compiled task handlers from the + Dag bundle named by the ``task_handler_bundle_name`` kwarg. It must name a + bundle in ``[dag_processor] dag_bundle_config_list`` and is checked when + this option is loaded. When it is unset, a task reads its handlers from its + own Dag bundle. + An entry may also carry an optional ``extra`` object for additional information associated with the coordinator that the coordinator itself does not receive; other components read it as needed. For example, @@ -2205,7 +2211,7 @@ sdk: "jdk-17": { "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", "kwargs": { - "jars_root": ["/opt/airflow/java-bundles"], + "task_handler_bundle_name": "java-task-handlers", "java_executable": "/usr/lib/jvm/java-17-openjdk/bin/java", "jvm_args": ["-Xmx1024m"] }, @@ -2218,7 +2224,7 @@ sdk: "go-sdk": { "classpath": "airflow.sdk.coordinators.executable.ExecutableCoordinator", "kwargs": { - "executables_root": ["/opt/airflow/executable-bundles"] + "task_handler_bundle_name": "go-task-handlers" } } } diff --git a/task-sdk/src/airflow/sdk/coordinators/_bundle_metadata.py b/task-sdk/src/airflow/sdk/coordinators/_bundle_metadata.py index ec6e6419cbab2..b97ed80e8d82e 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_bundle_metadata.py +++ b/task-sdk/src/airflow/sdk/coordinators/_bundle_metadata.py @@ -19,10 +19,9 @@ from __future__ import annotations -import os import pathlib import stat -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any import attrs import structlog @@ -38,24 +37,6 @@ 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]: - """Normalize a coordinator's root-directories kwarg into a list of expanded paths.""" - if value is None: - return [] - if isinstance(value, (str, os.PathLike, pathlib.Path)): - return [pathlib.Path(value).expanduser()] - return [pathlib.Path(v).expanduser() for v in value] - - def walk_files( roots: Iterable[pathlib.Path], *, match: Callable[[pathlib.Path], bool] ) -> Iterator[pathlib.Path]: @@ -103,30 +84,6 @@ 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 [] - if isinstance(value, (str, os.PathLike)): - values = [value] - elif value: - values = value - else: - values = [] - 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(values) - - 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/_subprocess.py b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py index 9e69255ba3d07..7ddddc9eeca4e 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py +++ b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py @@ -28,7 +28,6 @@ from __future__ import annotations import contextlib -import enum import ipaddress import itertools import os @@ -44,7 +43,6 @@ 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 @@ -402,22 +400,6 @@ def _initialize_pinned_bundle(target: BundleInfo, logger: FilteringBoundLogger) 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): """ @@ -432,97 +414,38 @@ 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. + :param task_handler_bundle_name: Name of the Dag bundle that holds the compiled + task handlers. It must be registered in ``[dag_processor] dag_bundle_config_list``. + If unset, the task's own Dag 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 + task_handler_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( + def _resolve_artifact_bundle( self, bundle_info: BundleInfo, logger: FilteringBoundLogger - ) -> tuple[list[pathlib.Path], BaseDagBundle | None]: + ) -> BaseDagBundle: """ - Resolve the directories to scan for artifacts, dispatched on the classified mode. + Materialize the Dag bundle holding the artifacts for a task of *bundle_info*. - 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. + That is the bundle named by ``task_handler_bundle_name``, or the task's own + bundle when it is unset. *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: + if self.task_handler_bundle_name is None: target = bundle_info + else: + target = BundleInfo(name=self.task_handler_bundle_name) 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 + return bundle def _get_scan_roots(self) -> tuple[pathlib.Path, ...]: """Return the artifact roots resolved for the active task.""" @@ -567,20 +490,14 @@ def execute_task( subprocess_logs_to_stdout: bool, **kwargs, ) -> BaseCoordinator.ExecutionResult: - 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)) + bundle = self._resolve_artifact_bundle(bundle_info, logger or log) + # 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. + with ( + BundleVersionLock(bundle_name=bundle.name, bundle_version=bundle.version), + self._set_scan_roots([bundle.path]), + ): command, subprocess_schema_version = self._build_execute_task_command(what=what) 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 9d6b80b249f81..237cc8dbbdf90 100644 --- a/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py @@ -31,9 +31,7 @@ import structlog from airflow.sdk.coordinators._bundle_metadata import ( - ARTIFACT_ROOTS_NOT_CONFIGURED, ResolvedBundle, - convert_configured_roots, extract_supervisor_schema_version, parse_metadata_mapping, ) @@ -294,10 +292,10 @@ def _walk_executables( @attrs.define class _Bundle(ResolvedBundle): @classmethod - def find(cls, executables_root: Sequence[pathlib.Path], dag_id: str) -> Self: - log.debug("Finding executable bundles recursively", roots=executables_root) + def find(cls, roots: Sequence[pathlib.Path], dag_id: str) -> Self: + log.debug("Finding executable bundles recursively", roots=roots) rejected: list[tuple[pathlib.Path, str]] = [] - for p in _find_executables(executables_root): + for p in _find_executables(roots): if (metadata := _read_bundle_metadata(p)) is None: continue if dag_id not in _dag_ids(metadata): @@ -310,7 +308,7 @@ def find(cls, executables_root: Sequence[pathlib.Path], dag_id: str) -> Self: rejected.append((p.resolve(), str(exc))) continue - resolved_paths = os.pathsep.join(str(r.resolve()) for r in executables_root) + resolved_paths = os.pathsep.join(str(r.resolve()) for r in roots) if rejected: details = "; ".join(f"{path}: {reason}" for path, reason in rejected) tp = ( @@ -334,27 +332,19 @@ class ExecutableCoordinator(SubprocessCoordinator): "go": { "classpath": "airflow.sdk.coordinators.executable.ExecutableCoordinator", "kwargs": { - "executables_root": ["~/airflow/executable-bundles"] + "task_handler_bundle_name": "go-task-handlers" } } - :param executables_root: A list of directories scanned for executable - bundles when a Python stub DAG delegates task execution to a native - runtime. See :class:`SubprocessCoordinator` for its interaction with - ``dag_bundle_name``. + :param task_handler_bundle_name: Name of the Dag bundle holding the + executable bundles a Python stub Dag delegates task execution to. It must + be registered in ``[dag_processor] dag_bundle_config_list``. If unset, + the task's own Dag bundle is used. Only files with the executable bit set + are considered. :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( - 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]: roots = self._get_scan_roots() bundle = _Bundle.find(roots, what.dag_id) diff --git a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py index 1f0f6e6efd5f9..5d483110580f2 100644 --- a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py @@ -29,11 +29,7 @@ import attrs import structlog -from airflow.sdk.coordinators._bundle_metadata import ( - ARTIFACT_ROOTS_NOT_CONFIGURED, - convert_configured_roots, - validate_schema_version, -) +from airflow.sdk.coordinators._bundle_metadata import validate_schema_version from airflow.sdk.coordinators._subprocess import SubprocessCoordinator if TYPE_CHECKING: @@ -86,8 +82,8 @@ def _iter_dir(directory: pathlib.Path) -> Iterator[pathlib.Path]: return -def _calculate_classpath(jars_root: Sequence[pathlib.Path]) -> str: - jars = (p.as_posix() for p in _find_jars(jars_root)) +def _calculate_classpath(roots: Sequence[pathlib.Path]) -> str: + jars = (p.as_posix() for p in _find_jars(roots)) return os.pathsep.join(sorted(jars)) # Keep output deterministic. @@ -167,7 +163,7 @@ class JavaCoordinator(SubprocessCoordinator): "jdk-17": { "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", "kwargs": { - "jars_root": ["~/airflow/jars"], + "task_handler_bundle_name": "java-task-handlers", "java_executable": "/usr/lib/jvm/java-17-openjdk/bin/java", "jvm_args": ["-Xmx1024m"] } @@ -176,16 +172,22 @@ 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. See - :class:`SubprocessCoordinator` for its interaction with ``dag_bundle_name``. + :param task_handler_bundle_name: Name of the Dag bundle holding the JARs. It + must be registered in ``[dag_processor] dag_bundle_config_list``. If + unset, the task's own Dag bundle is used. :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. - If *main_class* is not explicitly set, JavaCoordinator scans *jars_root* to + Every JAR in the bundle goes on the classpath, so one bundle is one + classpath. Handlers that need conflicting dependency versions belong in + separate bundles, each served by its own coordinator and queue. + + If *main_class* is not explicitly set, JavaCoordinator scans the bundle to find an executable JAR (one with Main-Class set in its metadata). If more than one executable JAR is found, it may be nondeterministic which one ends - up being executed. + up being executed, so set *main_class* when more than one JAR in the bundle + declares Main-Class. A JAR containing metadata *Airflow-Supervisor-Schema-Version* should also be available to specify the wire schema version. The JAR containing the Java @@ -203,16 +205,8 @@ class JavaCoordinator(SubprocessCoordinator): java_executable: str = "java" jvm_args: list[str] = attrs.field(factory=list) - jars_root: list[pathlib.Path] = attrs.field( - 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]: # Without main_class, the first executable JAR in walk order wins; tracked at # https://github.com/apache/airflow/issues/71134 diff --git a/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py index dd2a3f6f0416f..27fc142a3036c 100644 --- a/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py @@ -26,12 +26,7 @@ import attrs import structlog -from airflow.sdk.coordinators._bundle_metadata import ( - ARTIFACT_ROOTS_NOT_CONFIGURED, - ResolvedBundle, - convert_configured_roots, - walk_files, -) +from airflow.sdk.coordinators._bundle_metadata import ResolvedBundle, walk_files from airflow.sdk.coordinators._subprocess import SubprocessCoordinator from airflow.sdk.coordinators.node._bundle_reader import read_bundle @@ -55,11 +50,11 @@ def _is_bundle(path: pathlib.Path) -> bool: @attrs.define class _Bundle(ResolvedBundle): @classmethod - def find(cls, bundles_root: Sequence[pathlib.Path], dag_id: str) -> Self: + def find(cls, roots: Sequence[pathlib.Path], dag_id: str) -> Self: """Return the first verified configured bundle that declares *dag_id*.""" - log.debug("Finding TypeScript bundles recursively", roots=bundles_root, dag_id=dag_id) + log.debug("Finding TypeScript bundles recursively", roots=roots, dag_id=dag_id) rejected: list[tuple[pathlib.Path, str]] = [] - for candidate in walk_files(bundles_root, match=_is_bundle): + for candidate in walk_files(roots, match=_is_bundle): try: metadata = read_bundle(candidate) if dag_id not in metadata.dag_ids: @@ -85,7 +80,7 @@ def find(cls, bundles_root: Sequence[pathlib.Path], dag_id: str) -> Self: log.debug("Selected TypeScript bundle", path=candidate, dag_id=dag_id) return bundle - searched = os.pathsep.join(os.fspath(root) for root in bundles_root) + searched = os.pathsep.join(os.fspath(root) for root in roots) if rejected: details = "; ".join(f"{path}: {reason}" for path, reason in rejected) raise FileNotFoundError( @@ -108,29 +103,21 @@ class NodeCoordinator(SubprocessCoordinator): "classpath": "airflow.sdk.coordinators.node.NodeCoordinator", "kwargs": { "node_executable": "node", - "bundles_root": ["/opt/airflow/ts-bundles"], + "task_handler_bundle_name": "ts-task-handlers", }, } } :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. See :class:`SubprocessCoordinator` - for its interaction with ``dag_bundle_name``. + :param task_handler_bundle_name: Name of the Dag bundle searched recursively for the first + verified ``*.min.mjs`` bundle declaring the task instance's Dag. It must be registered in + ``[dag_processor] dag_bundle_config_list``. If unset, the task's own Dag bundle is used. :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( - 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]: roots = self._get_scan_roots() diff --git a/task-sdk/src/airflow/sdk/execution_time/coordinator.py b/task-sdk/src/airflow/sdk/execution_time/coordinator.py index 1469b27888c51..8aad5c6899dbc 100644 --- a/task-sdk/src/airflow/sdk/execution_time/coordinator.py +++ b/task-sdk/src/airflow/sdk/execution_time/coordinator.py @@ -48,6 +48,7 @@ import pydantic import structlog +from airflow.dag_processing.bundles.manager import DagBundlesManager # noqa: SDK002 from airflow.sdk._shared.module_loading import import_string from airflow.sdk.configuration import conf @@ -253,7 +254,13 @@ class CoordinatorManager: @classmethod def from_config(cls) -> Self: - """Load coordinator specs from configuration without initialization.""" + """ + Load coordinator specs from configuration without initialization. + + Every ``queue_to_coordinator`` key and every ``task_handler_bundle_name`` + kwarg is validated here, so a typo fails at config load rather than on the + first task routed to the coordinator. + """ coordinator_specs = { k: _CoordinatorSpec.model_validate(v) for k, v in conf.getjson("sdk", "coordinators", fallback={}).items() @@ -262,6 +269,15 @@ def from_config(cls) -> Self: for key in queue_to_coordinator.values(): if key not in coordinator_specs: raise ValueError(f"[sdk] queue_to_coordinator references invalid coordinator key: {key!r}") + for key, spec in coordinator_specs.items(): + bundle_name = spec.kwargs.get("task_handler_bundle_name") + if bundle_name is None: + continue + if not isinstance(bundle_name, str) or not DagBundlesManager.is_bundle_configured(bundle_name): + raise InvalidCoordinatorError( + f"[sdk] coordinators {key!r} sets task_handler_bundle_name={bundle_name!r}, " + f"which is not a bundle in [dag_processor] dag_bundle_config_list" + ) return cls(coordinator_specs=coordinator_specs, queue_to_coordinator=queue_to_coordinator) def _find_queue(self, key: str) -> BaseCoordinator: 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 c20913e26eae3..7071c9134bfb9 100644 --- a/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py @@ -18,7 +18,7 @@ from __future__ import annotations import hashlib -import pathlib +import json import socket import stat import struct @@ -43,6 +43,7 @@ from airflow.sdk.execution_time.coordinator import BaseCoordinator from airflow.sdk.execution_time.supervisor import ActivitySubprocess +from tests_common.test_utils.config import conf_vars from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS if not AIRFLOW_V_3_3_PLUS: @@ -361,46 +362,6 @@ def test_logs_when_metadata_is_not_a_mapping(self, tmp_path): class TestExecutableCoordinatorAttributes: - def test_executables_root_accepts_single_path(self, tmp_path): - coordinator = ExecutableCoordinator(executables_root=str(tmp_path)) - assert coordinator.executables_root == [tmp_path] - - def test_executables_root_accepts_list(self, tmp_path): - other = tmp_path / "other" - coordinator = ExecutableCoordinator(executables_root=[str(tmp_path), other]) - assert coordinator.executables_root == [tmp_path, other] - - def test_executables_root_optional_defaults_to_empty(self): - coordinator = ExecutableCoordinator() - assert coordinator.executables_root == [] - assert coordinator.dag_bundle_name is 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() @@ -417,7 +378,7 @@ def test_returns_resolved_executable_and_schema_version(self, tmp_path): binary = _build_bundle(tmp_path / "my_bundle", dag_ids=["tutorial_dag"]) ti = _make_ti(dag_id="tutorial_dag") - coordinator = ExecutableCoordinator(executables_root=[tmp_path]) + coordinator = ExecutableCoordinator() with coordinator._set_scan_roots([tmp_path]): command, schema_version = coordinator._build_execute_task_command(what=ti) assert command == [str(binary.resolve())] @@ -429,7 +390,7 @@ def test_raises_when_bundle_omits_schema_version(self, tmp_path): _build_bundle(tmp_path / "my_bundle", dag_ids=["tutorial_dag"], metadata=metadata) ti = _make_ti(dag_id="tutorial_dag") - coordinator = ExecutableCoordinator(executables_root=[tmp_path]) + coordinator = ExecutableCoordinator() with ( coordinator._set_scan_roots([tmp_path]), pytest.raises(FileNotFoundError, match="matching bundles were rejected"), @@ -440,7 +401,7 @@ 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]) + coordinator = ExecutableCoordinator() with ( coordinator._set_scan_roots([tmp_path]), pytest.raises(FileNotFoundError, match="cannot find executable bundle"), @@ -454,6 +415,18 @@ def bundles_dir(tmp_path): return tmp_path +@pytest.fixture +def go_task_handlers(bundles_dir): + """Register *bundles_dir* as the ``go-task-handlers`` Dag bundle and return that name.""" + bundle = { + "name": "go-task-handlers", + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {"path": str(bundles_dir)}, + } + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps([bundle])}): + yield bundle["name"] + + @pytest.fixture def mock_client(make_ti_context): client = MagicMock() @@ -462,10 +435,10 @@ def mock_client(make_ti_context): class TestExecutableCoordinatorExecuteTask: - def _captured_popen_cmd(self, bundles_dir: pathlib.Path, mock_client) -> list[str]: + def _captured_popen_cmd(self, bundle_name: str, mock_client) -> list[str]: """Run execute_task with mocked subprocess and return the command list.""" ti = _make_ti(dag_id="tutorial_dag") - coordinator = ExecutableCoordinator(executables_root=[bundles_dir]) + coordinator = ExecutableCoordinator(task_handler_bundle_name=bundle_name) mock_proc = MagicMock(spec=subprocess.Popen) mock_proc.pid = 12345 @@ -505,14 +478,14 @@ def capture_popen(cmd, **kwargs): assert popen_calls, "subprocess.Popen was not called" return popen_calls[0] - def test_executable_path_is_first_arg(self, bundles_dir, mock_client): - cmd = self._captured_popen_cmd(bundles_dir, mock_client) + def test_executable_path_is_first_arg(self, bundles_dir, go_task_handlers, mock_client): + cmd = self._captured_popen_cmd(go_task_handlers, mock_client) expected = str((bundles_dir / "my_bundle").resolve()) assert cmd[0] == expected - def test_returns_execution_result(self, bundles_dir, mock_client): + def test_returns_execution_result(self, go_task_handlers, mock_client): ti = _make_ti(dag_id="tutorial_dag") - coordinator = ExecutableCoordinator(executables_root=[bundles_dir]) + coordinator = ExecutableCoordinator(task_handler_bundle_name=go_task_handlers) mock_proc = MagicMock(spec=subprocess.Popen) mock_proc.pid = 99999 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 c55a2cc891253..eda377e6251fd 100644 --- a/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py @@ -18,6 +18,7 @@ from __future__ import annotations +import json import os import pathlib import re @@ -39,6 +40,7 @@ from airflow.sdk.execution_time.coordinator import BaseCoordinator from airflow.sdk.execution_time.supervisor import ActivitySubprocess +from tests_common.test_utils.config import conf_vars from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS if not AIRFLOW_V_3_3_PLUS: @@ -222,50 +224,22 @@ def test_skip_logged_when_directory_revisited(self, tmp_path): class TestJavaCoordinatorAttributes: def test_default_kwargs(self): - coordinator = JavaCoordinator(jars_root="/airflow/java-bundles") + coordinator = JavaCoordinator() assert coordinator.java_executable == "java" assert coordinator.jvm_args == [] - assert coordinator.jars_root == [pathlib.Path("/airflow/java-bundles")] + assert coordinator.task_handler_bundle_name is None + # main_class stays optional: the entrypoint is auto-detected from the bundle scan. + assert coordinator.main_class == "" def test_custom_kwargs(self): coordinator = JavaCoordinator( java_executable="/opt/java/bin/java", jvm_args=["-Xmx512m", "-Xms256m"], - jars_root=["/airflow/java-bundles"], + task_handler_bundle_name="java-task-handlers", ) assert coordinator.java_executable == "/opt/java/bin/java" 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") + assert coordinator.task_handler_bundle_name == "java-task-handlers" 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") @@ -283,6 +257,18 @@ def jars_root(tmp_path): return tmp_path +@pytest.fixture +def java_task_handlers(jars_root): + """Register *jars_root* as the ``java-task-handlers`` Dag bundle and return that name.""" + bundle = { + "name": "java-task-handlers", + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {"path": str(jars_root)}, + } + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps([bundle])}): + yield bundle["name"] + + @pytest.fixture def mock_client(make_ti_context): client = MagicMock() @@ -293,7 +279,7 @@ def mock_client(make_ti_context): class TestJavaCoordinatorExecuteTask: def _captured_popen_cmd( self, - jars_root: pathlib.Path, + bundle_name: str, mock_client, *, java_executable: str = "java", @@ -304,7 +290,7 @@ def _captured_popen_cmd( coordinator = JavaCoordinator( java_executable=java_executable, jvm_args=jvm_args or [], - jars_root=jars_root, + task_handler_bundle_name=bundle_name, ) mock_proc = MagicMock(spec=subprocess.Popen) @@ -345,56 +331,58 @@ def capture_popen(cmd, **kwargs): assert popen_calls, "subprocess.Popen was not called" return popen_calls[0] - def test_java_executable_is_first_arg(self, jars_root, mock_client): + def test_java_executable_is_first_arg(self, java_task_handlers, mock_client): cmd = self._captured_popen_cmd( - jars_root, mock_client, java_executable="/usr/lib/jvm/java-17/bin/java" + java_task_handlers, mock_client, java_executable="/usr/lib/jvm/java-17/bin/java" ) assert cmd[0] == "/usr/lib/jvm/java-17/bin/java" - def test_classpath_flag_and_value_present(self, jars_root, mock_client): - cmd = self._captured_popen_cmd(jars_root, mock_client) + def test_classpath_flag_and_value_present(self, jars_root, java_task_handlers, mock_client): + cmd = self._captured_popen_cmd(java_task_handlers, mock_client) assert "-classpath" in cmd cp_idx = cmd.index("-classpath") classpath = cmd[cp_idx + 1] assert jars_root.joinpath("app.jar").as_posix() in classpath - def test_main_class_present(self, jars_root, mock_client): - cmd = self._captured_popen_cmd(jars_root, mock_client) + def test_main_class_present(self, java_task_handlers, mock_client): + cmd = self._captured_popen_cmd(java_task_handlers, mock_client) assert "com.example.TaskRunner" in cmd - def test_comm_and_logs_args_present(self, jars_root, mock_client): - cmd = self._captured_popen_cmd(jars_root, mock_client) + def test_comm_and_logs_args_present(self, java_task_handlers, mock_client): + cmd = self._captured_popen_cmd(java_task_handlers, mock_client) comm_args = [a for a in cmd if a.startswith("--comm=")] logs_args = [a for a in cmd if a.startswith("--logs=")] assert len(comm_args) == 1 assert len(logs_args) == 1 - def test_comm_and_logs_contain_port(self, jars_root, mock_client): - cmd = self._captured_popen_cmd(jars_root, mock_client) + def test_comm_and_logs_contain_port(self, java_task_handlers, mock_client): + cmd = self._captured_popen_cmd(java_task_handlers, mock_client) comm_arg = next(a for a in cmd if a.startswith("--comm=")) logs_arg = next(a for a in cmd if a.startswith("--logs=")) # format is host:port assert ":" in comm_arg.split("=", 1)[1] assert ":" in logs_arg.split("=", 1)[1] - def test_jvm_args_inserted_before_main_class(self, jars_root, mock_client): - cmd = self._captured_popen_cmd(jars_root, mock_client, jvm_args=["-Xmx512m", "-Dsome.prop=value"]) + def test_jvm_args_inserted_before_main_class(self, java_task_handlers, mock_client): + cmd = self._captured_popen_cmd( + java_task_handlers, mock_client, jvm_args=["-Xmx512m", "-Dsome.prop=value"] + ) main_idx = cmd.index("com.example.TaskRunner") for jvm_arg in ["-Xmx512m", "-Dsome.prop=value"]: assert jvm_arg in cmd assert cmd.index(jvm_arg) < main_idx - def test_comm_and_logs_after_main_class(self, jars_root, mock_client): - cmd = self._captured_popen_cmd(jars_root, mock_client) + def test_comm_and_logs_after_main_class(self, java_task_handlers, mock_client): + cmd = self._captured_popen_cmd(java_task_handlers, mock_client) main_idx = cmd.index("com.example.TaskRunner") comm_idx = next(i for i, a in enumerate(cmd) if a.startswith("--comm=")) logs_idx = next(i for i, a in enumerate(cmd) if a.startswith("--logs=")) assert comm_idx > main_idx assert logs_idx > main_idx - def test_returns_execution_result(self, jars_root, mock_client): + def test_returns_execution_result(self, java_task_handlers, mock_client): ti = _make_ti() - coordinator = JavaCoordinator(jars_root=jars_root) + coordinator = JavaCoordinator(task_handler_bundle_name=java_task_handlers) mock_proc = MagicMock(spec=subprocess.Popen) mock_proc.pid = 99999 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 146cb91c1d94c..3dc01be33f297 100644 --- a/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py @@ -21,7 +21,6 @@ import json import pathlib from unittest import mock -from unittest.mock import patch import pytest from task_sdk.coordinators.node._bundle_test_utils import ( @@ -60,50 +59,22 @@ def _make_ti(dag_id: str = "test_dag", queue: str = "ts") -> TaskInstance: class TestNodeCoordinatorAttributes: def test_default_kwargs(self): - coordinator = NodeCoordinator(bundles_root="/airflow/ts-bundles") + coordinator = NodeCoordinator() assert coordinator.node_executable == "node" - assert coordinator.bundles_root == [pathlib.Path("/airflow/ts-bundles")] assert coordinator.task_startup_timeout == 10.0 def test_custom_kwargs(self): coordinator = NodeCoordinator( node_executable="/opt/node/bin/node", - bundles_root=["/airflow/ts-bundles", "~/extra-bundles"], + task_handler_bundle_name="ts-task-handlers", task_startup_timeout=30.0, ) assert coordinator.node_executable == "/opt/node/bin/node" - assert coordinator.bundles_root == [ - pathlib.Path("/airflow/ts-bundles"), - pathlib.Path("~/extra-bundles").expanduser(), - ] + assert coordinator.task_handler_bundle_name == "ts-task-handlers" assert coordinator.task_startup_timeout == 30.0 - 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() @@ -116,10 +87,7 @@ def test_build_command_scans_passed_roots_in_colocated_mode(self, tmp_path): class TestNodeCoordinatorExecuteTaskCommand: def test_selects_bundle_by_dag_id(self, tmp_path): selected = write_bundle(tmp_path, "sales") - coordinator = NodeCoordinator( - node_executable="/opt/node/bin/node", - bundles_root=tmp_path, - ) + coordinator = NodeCoordinator(node_executable="/opt/node/bin/node") with coordinator._set_scan_roots([tmp_path]): command, schema_version = coordinator._build_execute_task_command(what=_make_ti(dag_id="sales")) @@ -129,10 +97,7 @@ def test_selects_bundle_by_dag_id(self, tmp_path): 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, - ) + coordinator = NodeCoordinator(node_executable="/opt/node/bin/node") with coordinator._set_scan_roots([tmp_path]): command, schema_version = coordinator._build_execute_task_command(what=_make_ti()) diff --git a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py index 5297f56dead4f..909a2c605f0ba 100644 --- a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py +++ b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py @@ -32,13 +32,12 @@ import pytest from uuid6 import uuid7 -from airflow.dag_processing.bundles.base import BundleVersion +from airflow.dag_processing.bundles.base import BaseDagBundle, BundleVersion from airflow.sdk.api.client import Client, TaskInstanceOperations 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, @@ -573,26 +572,37 @@ def test_untrack_unknown_object_does_not_raise(self): class _StubSubprocessCoordinator(SubprocessCoordinator): """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. + 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 +def _make_bundle(path: pathlib.Path, *, version: str | None = None, name: str = "dags") -> MagicMock: + bundle = MagicMock(spec=BaseDagBundle, path=path, version=version) + bundle.name = name + return bundle + + +@pytest.fixture +def artifact_bundle(tmp_path): + """Resolve every task's artifacts to an unversioned bundle at *tmp_path*, without the bundle machinery.""" + with patch.object( + SubprocessCoordinator, + "_resolve_artifact_bundle", + autospec=True, + return_value=_make_bundle(tmp_path), + ) as mock_resolve: + yield mock_resolve + + @pytest.fixture def mock_client(make_ti_context): client = MagicMock(spec=Client) @@ -618,6 +628,7 @@ class _Plain(SubprocessCoordinator): _Plain()._build_execute_task_command(what=_make_ti()) +@pytest.mark.usefixtures("artifact_bundle") class TestSubprocessCoordinatorExecuteTask: def _captured_popen_cmd( self, @@ -863,79 +874,20 @@ def test_register_pipe_readers_called_with_four_sockets(self, mock_client): 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.""" +class TestResolveArtifactBundle: + """Execute-time resolution of the Dag bundle holding the artifacts.""" - 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") + @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle", autospec=True) + def test_unset_name_uses_the_task_bundle(self, mock_initialize, tmp_path): + resolved = _make_bundle(tmp_path, version="v3") mock_initialize.return_value = resolved - coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=[]) + coordinator = _StubSubprocessCoordinator(command=["x"]) 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 + assert coordinator._resolve_artifact_bundle(bundle_info, log) is resolved mock_initialize.assert_called_once_with(bundle_info) - @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle") + @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle", autospec=True) 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. @@ -947,67 +899,52 @@ def test_version_less_bundle_is_rematerialized_at_its_current_version(self, mock pinned_tree = tmp_path / "versions" / "sha-abc" pinned_tree.mkdir(parents=True) - unpinned = MagicMock(path=shared_checkout, version=None) + unpinned = _make_bundle(shared_checkout) unpinned.get_current_version.return_value = BundleVersion(version="sha-abc", data={"k": "v"}) - pinned = MagicMock(path=pinned_tree, version="sha-abc") + pinned = _make_bundle(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 + coordinator = _StubSubprocessCoordinator(command=["x"], task_handler_bundle_name="artifacts") - roots, bundle = coordinator._init_root_source(MagicMock(), log) - - assert roots == [pinned_tree] - assert bundle is pinned + assert coordinator._resolve_artifact_bundle(MagicMock(spec=BundleInfo), log) 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") + @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle", autospec=True) def test_unversioned_bundle_keeps_its_single_path(self, mock_initialize, tmp_path): - resolved = MagicMock(path=tmp_path, version=None) + resolved = _make_bundle(tmp_path) 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) + coordinator = _StubSubprocessCoordinator(command=["x"], task_handler_bundle_name="artifacts") - assert roots == [tmp_path] - assert bundle is resolved + assert coordinator._resolve_artifact_bundle(MagicMock(spec=BundleInfo), log) is resolved mock_initialize.assert_called_once_with(BundleInfo(name="artifacts")) - @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle") + @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle", autospec=True) 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 + mock_initialize.return_value = _make_bundle(tmp_path / "nope", version="v1") + coordinator = _StubSubprocessCoordinator(command=["x"], task_handler_bundle_name="artifacts") with pytest.raises(FileNotFoundError, match="does not exist"): - coordinator._init_root_source(MagicMock(), log) + coordinator._resolve_artifact_bundle(MagicMock(spec=BundleInfo), 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("airflow.sdk.coordinators._subprocess.BundleVersionLock", autospec=True) + @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle", autospec=True) @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_initialize.return_value = _make_bundle(tmp_path, version="v9") mock_start.return_value.wait.return_value = 0 - coordinator = _StubSubprocessCoordinator(command=["/runtime"], explicit_roots=[]) + coordinator = _StubSubprocessCoordinator(command=["/runtime"]) bundle_info = BundleInfo(name="dags", version="v9") coordinator.execute_task( @@ -1024,27 +961,8 @@ 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.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("airflow.sdk.coordinators._subprocess.BundleVersionLock", autospec=True) + @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle", autospec=True) @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 @@ -1053,21 +971,18 @@ def test_named_bundle_mode_locks_the_tree_it_scans( pinned_tree = tmp_path / "versions" / "sha-abc" pinned_tree.mkdir(parents=True) - unpinned = MagicMock(path=tmp_path, version=None) + unpinned = _make_bundle(tmp_path, name="artifacts") unpinned.get_current_version.return_value = BundleVersion(version="sha-abc", data=None) - pinned = MagicMock(path=pinned_tree, version="sha-abc") - pinned.name = "artifacts" + pinned = _make_bundle(pinned_tree, version="sha-abc", 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 = _StubSubprocessCoordinator(command=["/runtime"], task_handler_bundle_name="artifacts") coordinator.execute_task( what=_make_ti(), dag_rel_path="dag.py", - bundle_info=MagicMock(), + bundle_info=BundleInfo(name="dags", version="v9"), client=mock_client, subprocess_logs_to_stdout=False, ) @@ -1093,6 +1008,7 @@ def test_rejects_nested_reentry(self, tmp_path): with coordinator._set_scan_roots([tmp_path]): pass + @pytest.mark.usefixtures("artifact_bundle") def test_execute_task_is_not_reentrant(self, mock_client, tmp_path): coordinator = _StubSubprocessCoordinator(command=["/bin/true"]) with coordinator._set_scan_roots([tmp_path]): diff --git a/task-sdk/tests/task_sdk/execution_time/test_coordinator.py b/task-sdk/tests/task_sdk/execution_time/test_coordinator.py index c4185c1939092..3659bce784dbf 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_coordinator.py +++ b/task-sdk/tests/task_sdk/execution_time/test_coordinator.py @@ -37,6 +37,8 @@ reset_coordinator_manager, ) +from tests_common.test_utils.config import conf_vars + class _CoordinatorA(BaseCoordinator): def __init__(self, *, label: str = "a"): @@ -134,6 +136,22 @@ def test_from_config_rejects_invalid_queue_mapping(self, sdk_config): ): CoordinatorManager.from_config() + @pytest.mark.parametrize("bundle_name", ["ghost", ["ghost"]], ids=["unconfigured", "not-a-string"]) + def test_from_config_rejects_invalid_task_handler_bundle_name(self, sdk_config, bundle_name): + """Validated for every spec, routed or not, without importing or constructing the coordinator.""" + sdk_config( + coordinators=json.dumps( + { + "boom": { + "classpath": f"{_ExplodingCoordinator.__module__}._ExplodingCoordinator", + "kwargs": {"task_handler_bundle_name": bundle_name}, + } + } + ), + ) + with pytest.raises(InvalidCoordinatorError, match=r"'boom' sets task_handler_bundle_name="): + CoordinatorManager.from_config() + @pytest.mark.parametrize( ("coordinator_spec", "expected_match"), [ @@ -261,8 +279,8 @@ def test_extra_for_queue_does_not_instantiate_coordinator(self, sdk_config): class TestConfigYamlCoordinatorsExample: """Guard the ``[sdk] coordinators`` example in ``config.yml`` against drift. - Nothing else exercises the example, so a broken one (e.g. dropping the - required ``jars_root`` kwarg) can ship unnoticed. Loading it through + Nothing else exercises the example, so a broken one (e.g. a kwarg the + coordinator does not accept) can ship unnoticed. Loading it through CoordinatorManager and constructing every entry keeps the example honest. """ @@ -279,7 +297,14 @@ def test_every_example_coordinator_constructs(self, sdk_config): coordinators=coordinators_example, queue_to_coordinator=json.dumps(queue_to_coordinator), ) - manager = CoordinatorManager.from_config() + # The deployment registers the Dag bundles the example names. + artifact_bundles = [ + {"name": name, "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", "kwargs": {}} + for spec in specs.values() + if (name := spec.get("kwargs", {}).get("task_handler_bundle_name")) + ] + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(artifact_bundles)}): + manager = CoordinatorManager.from_config() assert set(manager._coordinator_specs) == set(specs) for queue, key in queue_to_coordinator.items(): From ace4b2eeaa9855b389ca57741a9e94482a5601a6 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 08:57:00 +0000 Subject: [PATCH 02/94] Serve lang-SDK harness artifacts from Dag bundles The compose e2e and the k8s lang-SDK test configured coordinators with the explicit artifact roots, which no longer exist. The Java e2e needs three bundles because one bundle is one classpath, and its three coordinators each need their own dependencies. Artifacts stay mounted on the workers only, because only the workers run them; the other components need just the bundle registration for the name to resolve. The k8s init container keeps staging from S3 because the S3 download drops the Go binary's execute bit, which the coordinator requires. --- .../tests/airflow_e2e_tests/conftest.py | 56 ++++++++++++++----- kubernetes-tests/lang_sdk/README.md | 14 ++--- kubernetes-tests/lang_sdk/config/values.yaml | 14 ++++- .../pod_templates/lang_sdk_golang.yaml | 3 +- .../lang_sdk/pod_templates/lang_sdk_java.yaml | 3 +- kubernetes-tests/lang_sdk/stage_artifacts.py | 8 +-- 6 files changed, 70 insertions(+), 28 deletions(-) diff --git a/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py b/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py index 39ff3fd0a64b8..db08b27b06397 100644 --- a/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py +++ b/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py @@ -293,6 +293,22 @@ def _setup_xcom_object_storage_integration(dot_env_file, tmp_dir): ] +def _dag_bundle_config_with_artifacts(artifact_bundles: dict[str, str]) -> str: + """Return a ``dag_bundle_config_list`` of the Dags folder plus a ``LocalDagBundle`` per artifact path. + + Registration is what makes a coordinator's ``task_handler_bundle_name`` resolvable on the + worker. The artifact directories are mounted only there; the Dag processor finds no files + in them. + """ + local_bundle = "airflow.dag_processing.bundles.local.LocalDagBundle" + bundles = [{"name": "dags-folder", "classpath": local_bundle, "kwargs": {}}] + bundles.extend( + {"name": name, "classpath": local_bundle, "kwargs": {"path": path}} + for name, path in artifact_bundles.items() + ) + return json.dumps(bundles) + + def _run_java_sdk_gradle(workdir, *gradle_argv, capture_output=False, native=False): """Run the Java SDK Gradle wrapper natively or inside the pinned JDK container. @@ -416,7 +432,7 @@ def _setup_java_sdk_integration(dot_env_file, tmp_dir): copyfile(JAVA_DOCKERFILE_PATH, tmp_dir / "Dockerfile.java") # Copy each bundle's JARs into its own directory; the compose bind-mounts - # expose them to the worker, and each JavaCoordinator globs its own dir. + # expose them to the worker, where each is registered as its own Dag bundle. copytree(JAVA_SDK_EXAMPLE_LIBS_PATH, tmp_dir / "java-jars") copytree(SCALA_SPARK_EXAMPLE_LIBS_PATH, tmp_dir / "scala-jars") copytree(JAVA_TEST_BUNDLE_LIBS_PATH, tmp_dir / "java-test-jars") @@ -456,20 +472,27 @@ def _setup_java_sdk_integration(dot_env_file, tmp_dir): ) # One JavaCoordinator per queue on the same worker image, each serving its - # own bundle. The scala-jdk entry pins main_class (Spark's large classpath - # makes Main-Class discovery ambiguous) and carries Spark's Java 17 module - # openings, a small driver heap, and a longer startup timeout for its large - # dependency classpath. + # own artifact bundle (one bundle is one classpath). The scala-jdk entry pins + # main_class (Spark's large classpath makes Main-Class discovery ambiguous) + # and carries Spark's Java 17 module openings, a small driver heap, and a + # longer startup timeout for its large dependency classpath. + dag_bundle_config = _dag_bundle_config_with_artifacts( + { + "java-task-handlers": "/opt/airflow/java-jars", + "scala-task-handlers": "/opt/airflow/scala-jars", + "java-test-task-handlers": "/opt/airflow/java-test-jars", + } + ) coordinator_config = json.dumps( { "java-jdk": { "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", - "kwargs": {"jars_root": ["/opt/airflow/java-jars"]}, + "kwargs": {"task_handler_bundle_name": "java-task-handlers"}, }, "scala-jdk": { "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", "kwargs": { - "jars_root": ["/opt/airflow/scala-jars"], + "task_handler_bundle_name": "scala-task-handlers", "main_class": "org.apache.airflow.example.ScalaSparkBundleBuilder", "jvm_args": ["-Xmx512m", *_SPARK_JAVA_MODULE_OPTIONS], "task_startup_timeout": 60.0, @@ -477,7 +500,7 @@ def _setup_java_sdk_integration(dot_env_file, tmp_dir): }, "java-test-jdk": { "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", - "kwargs": {"jars_root": ["/opt/airflow/java-test-jars"]}, + "kwargs": {"task_handler_bundle_name": "java-test-task-handlers"}, }, } ) @@ -502,6 +525,7 @@ def _setup_java_sdk_integration(dot_env_file, tmp_dir): dot_env_file.write_text( f"AIRFLOW_UID={os.getuid()}\n" # Single-quote the JSON values so Docker Compose reads them literally. + f"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST='{dag_bundle_config}'\n" f"AIRFLOW__SDK__COORDINATORS='{coordinator_config}'\n" f"AIRFLOW__SDK__QUEUE_TO_COORDINATOR='{queue_to_coordinator}'\n" f"AIRFLOW_CONN_TEST_HTTP='{test_http_conn}'\n" @@ -598,9 +622,9 @@ def _setup_go_sdk_integration(dot_env_file, tmp_dir): """Set up the go_sdk E2E test mode. Compiles the Go SDK example bundle into a self-contained executable bundle - via the ``airflow-go-pack`` tooling, drops it into the directory the - ``ExecutableCoordinator`` scans, copies the Python stub Dag, and writes the - coordinator configuration. + via the ``airflow-go-pack`` tooling, drops it into the directory registered + as the ``go-task-handlers`` Dag bundle, copies the Python stub Dag, and + writes the coordinator configuration. The packed bundle is a statically linked native executable (built with ``CGO_ENABLED=0``), so the stock Airflow worker image can exec it directly @@ -624,13 +648,14 @@ def _setup_go_sdk_integration(dot_env_file, tmp_dir): copyfile(GO_SDK_DAGS_PATH / "go_examples.py", tmp_dir / "dags" / "go_examples.py") # Coordinator registry: maps the logical name "go-sdk" to ExecutableCoordinator, - # which scans executables_root for the packed bundle by dag_id. + # which scans the go-task-handlers Dag bundle for the packed bundle by dag_id. # Queue mapping: routes tasks on the "golang" queue to "go-sdk". + dag_bundle_config = _dag_bundle_config_with_artifacts({"go-task-handlers": "/opt/airflow/go-bundles"}) coordinator_config = json.dumps( { "go-sdk": { "classpath": "airflow.sdk.coordinators.executable.ExecutableCoordinator", - "kwargs": {"executables_root": ["/opt/airflow/go-bundles"]}, + "kwargs": {"task_handler_bundle_name": "go-task-handlers"}, } } ) @@ -639,6 +664,7 @@ def _setup_go_sdk_integration(dot_env_file, tmp_dir): dot_env_file.write_text( f"AIRFLOW_UID={os.getuid()}\n" # Single-quote the JSON values so Docker Compose reads them literally. + f"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST='{dag_bundle_config}'\n" f"AIRFLOW__SDK__COORDINATORS='{coordinator_config}'\n" f"AIRFLOW__SDK__QUEUE_TO_COORDINATOR='{queue_to_coordinator}'\n" # Connection and variable read by the Go example bundle tasks. @@ -730,12 +756,13 @@ def _setup_ts_sdk_integration(dot_env_file, tmp_dir): for dag_file in ("typescript_example.py", "typescript_taskflow_example.py"): copyfile(TS_SDK_EXAMPLE_PATH / "dags" / dag_file, tmp_dir / "dags" / dag_file) + dag_bundle_config = _dag_bundle_config_with_artifacts({"ts-task-handlers": "/opt/airflow/ts-bundles"}) coordinator_config = json.dumps( { "ts": { "classpath": "airflow.sdk.coordinators.node.NodeCoordinator", "kwargs": { - "bundles_root": ["/opt/airflow/ts-bundles"], + "task_handler_bundle_name": "ts-task-handlers", "node_executable": "/opt/nodejs/node", }, } @@ -747,6 +774,7 @@ def _setup_ts_sdk_integration(dot_env_file, tmp_dir): f"AIRFLOW_UID={os.getuid()}\n" f"NODE_IMAGE={NODE_IMAGE}\n" # single-quoted so Docker Compose reads the JSON literally + f"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST='{dag_bundle_config}'\n" f"AIRFLOW__SDK__COORDINATORS='{coordinator_config}'\n" f"AIRFLOW__SDK__QUEUE_TO_COORDINATOR='{queue_to_coordinator}'\n" "AIRFLOW_CONN_TYPESCRIPT_EXAMPLE_HTTP=http://user:pass@example.com/\n" diff --git a/kubernetes-tests/lang_sdk/README.md b/kubernetes-tests/lang_sdk/README.md index 93583ecbd8b44..f19ae247529eb 100644 --- a/kubernetes-tests/lang_sdk/README.md +++ b/kubernetes-tests/lang_sdk/README.md @@ -36,16 +36,16 @@ End-to-end test that one Dag mixing **Python + Go + Java** tasks runs to success worker pod (from that pod template): initContainer stage_artifacts.py ── S3DagBundle.initialize() ──► pulls go-artifacts / java-artifacts bucket into the shared - emptyDir = executables_root / jars_root + emptyDir = go-task-handlers / java-task-handlers LocalDagBundle base container supervisor → coordinator forks the Go binary / Java jar ``` -Key point: in coordinator mode the worker pod does **not** run the Python task-runner, -so the artifact is not downloaded by the normal Dag-bundle path. The init container runs -`stage_artifacts.py`, which reuses the **DagBundle interface** +Key point: each coordinator's `task_handler_bundle_name` names a `LocalDagBundle` over the +shared `emptyDir`, registered in `dagProcessor.dagBundleConfigList`. The init container fills it +by running `stage_artifacts.py`, which reuses the **DagBundle interface** (`DagBundlesManager().get_bundle(name).initialize()` — the download half of -`task_runner.parse`) to pull the artifact from its S3 bucket into the path the -coordinator scans. +`task_runner.parse`) to pull the artifact from its S3 bucket, then restores the Go binary's +execute bit, which the S3 download drops and the coordinator requires. ## Components @@ -58,7 +58,7 @@ coordinator scans. | `pod_templates/lang_sdk_golang.yaml` | `golang` queue worker pod: prod image + go-artifacts init container. | | `pod_templates/lang_sdk_java.yaml` | `java` queue worker pod: JVM image + java-artifacts init container. | | `manifests/localstack.yaml` | In-cluster S3 (localstack). | -| `config/values.yaml` | Helm overrides: KubernetesExecutor, coordinators (+extra.pod_template_file), queue routing, stub-Dag S3 bundle, AWS conn, scheduler pod-template mount. | +| `config/values.yaml` | Helm overrides: KubernetesExecutor, coordinators (+extra.pod_template_file), queue routing, stub-Dag S3 bundle, artifact Dag bundles, AWS conn, scheduler pod-template mount. | The Go binary, Java jar, and stub Dag share one object store (localstack) but live in **separate buckets** (`go-artifacts`, `java-artifacts`, `dags`). diff --git a/kubernetes-tests/lang_sdk/config/values.yaml b/kubernetes-tests/lang_sdk/config/values.yaml index a4b5b4004b403..6f4d9671ad6a5 100644 --- a/kubernetes-tests/lang_sdk/config/values.yaml +++ b/kubernetes-tests/lang_sdk/config/values.yaml @@ -50,6 +50,10 @@ scheduler: # Configure it here (dagProcessor.dagBundleConfigList) rather than overriding # config.dag_processor.dag_bundle_config_list directly -- the chart renders the config from this and # it is the recommended, drift-safe path. +# +# go-task-handlers / java-task-handlers are the coordinators' task_handler_bundle_name bundles: the +# worker pod's init container stages each artifact bucket into these paths (stage_artifacts.py). +# They exist only in the worker pods, so the dag-processor finds no files in them. dagProcessor: dagBundleConfigList: - name: dags-folder @@ -60,6 +64,14 @@ dagProcessor: kwargs: bucket_name: "dags" aws_conn_id: "aws_localstack" + - name: go-task-handlers + classpath: "airflow.dag_processing.bundles.local.LocalDagBundle" + kwargs: + path: "/opt/airflow/artifacts/go-artifacts" + - name: java-task-handlers + classpath: "airflow.dag_processing.bundles.local.LocalDagBundle" + kwargs: + path: "/opt/airflow/artifacts/java-artifacts" config: # Route the golang/java queues to their coordinators, each carrying an @@ -70,5 +82,5 @@ config: # Single-line JSON: airflow.cfg is an ini file, so a value with embedded # newlines would break the rendered configmap. # yamllint disable-line rule:line-length - coordinators: '{"go-sdk": {"classpath": "airflow.sdk.coordinators.executable.ExecutableCoordinator", "kwargs": {"executables_root": ["/opt/airflow/artifacts/go-artifacts"]}, "extra": {"pod_template_file": "/opt/airflow/lang_sdk/pod_templates/lang_sdk_golang.yaml"}}, "java-sdk": {"classpath": "airflow.sdk.coordinators.java.JavaCoordinator", "kwargs": {"jars_root": ["/opt/airflow/artifacts/java-artifacts"]}, "extra": {"pod_template_file": "/opt/airflow/lang_sdk/pod_templates/lang_sdk_java.yaml", "worker_container_repository": "lang-sdk-java-worker", "worker_container_tag": "latest"}}}' + coordinators: '{"go-sdk": {"classpath": "airflow.sdk.coordinators.executable.ExecutableCoordinator", "kwargs": {"task_handler_bundle_name": "go-task-handlers"}, "extra": {"pod_template_file": "/opt/airflow/lang_sdk/pod_templates/lang_sdk_golang.yaml"}}, "java-sdk": {"classpath": "airflow.sdk.coordinators.java.JavaCoordinator", "kwargs": {"task_handler_bundle_name": "java-task-handlers"}, "extra": {"pod_template_file": "/opt/airflow/lang_sdk/pod_templates/lang_sdk_java.yaml", "worker_container_repository": "lang-sdk-java-worker", "worker_container_tag": "latest"}}}' queue_to_coordinator: '{"golang": "go-sdk", "java": "java-sdk"}' diff --git a/kubernetes-tests/lang_sdk/pod_templates/lang_sdk_golang.yaml b/kubernetes-tests/lang_sdk/pod_templates/lang_sdk_golang.yaml index e3e700d700012..00d6b990690ae 100644 --- a/kubernetes-tests/lang_sdk/pod_templates/lang_sdk_golang.yaml +++ b/kubernetes-tests/lang_sdk/pod_templates/lang_sdk_golang.yaml @@ -19,7 +19,8 @@ # # An init container stages the packed Go bundle from the localstack ``go-artifacts`` # bucket into a shared emptyDir at /opt/airflow/artifacts via the DagBundle -# interface (stage_artifacts.py). That path is the coordinator's executables_root. +# interface (stage_artifacts.py). That path is the go-task-handlers Dag bundle the +# coordinator's task_handler_bundle_name names. # The base image is the plain prod image: a Go bundle is a self-contained static # binary, so no extra runtime is needed. Image placeholders are rendered by the # ``breeze k8s setup-lang-sdk-test`` command. diff --git a/kubernetes-tests/lang_sdk/pod_templates/lang_sdk_java.yaml b/kubernetes-tests/lang_sdk/pod_templates/lang_sdk_java.yaml index eaccd63f4d2ad..ad4456610fb73 100644 --- a/kubernetes-tests/lang_sdk/pod_templates/lang_sdk_java.yaml +++ b/kubernetes-tests/lang_sdk/pod_templates/lang_sdk_java.yaml @@ -19,7 +19,8 @@ # # An init container stages the Java jar from the localstack ``java-artifacts`` # bucket into a shared emptyDir at /opt/airflow/artifacts via the DagBundle -# interface (stage_artifacts.py). That path is the coordinator's jars_root. The +# interface (stage_artifacts.py). That path is the java-task-handlers Dag bundle the +# coordinator's task_handler_bundle_name names. The # base image is the prod image plus a headless JRE (built from Dockerfile.java) # so the JavaCoordinator can exec ``java`` -- distinct from the Go queue's plain # prod image. Image placeholders are rendered by ``breeze k8s setup-lang-sdk-test``. diff --git a/kubernetes-tests/lang_sdk/stage_artifacts.py b/kubernetes-tests/lang_sdk/stage_artifacts.py index 3da9ac6fb5a72..45a9d69882b03 100644 --- a/kubernetes-tests/lang_sdk/stage_artifacts.py +++ b/kubernetes-tests/lang_sdk/stage_artifacts.py @@ -23,10 +23,10 @@ parse), so the Go binary / Java jar stored in an S3 (localstack) bucket is pulled into ``bundle.path`` = ``{dag_bundle_storage_path}/{name}``. -That path is the ``emptyDir`` shared with the worker container and is exactly the -``executables_root`` / ``jars_root`` the coordinator scans. Coordinator mode never -runs the Python task-runner in the worker pod, so without this step the artifact -would never reach the pod. +That path is the ``emptyDir`` shared with the worker container, where it is registered +as the ``LocalDagBundle`` the coordinator's ``task_handler_bundle_name`` names. The +coordinator does not name the S3 bundle directly because its download drops the Go +binary's execute bit, which the coordinator requires; Java is staged the same way. Configuration is read from the environment (set by the pod template): From 4bfc18153e83066b031a9150634f21825bf9ccc0 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 09:04:57 +0000 Subject: [PATCH 03/94] Document artifact Dag bundles for language SDK coordinators The Go, Java and TypeScript guides, the bundle specs, the READMEs and the contributor guide still pointed users at the removed jars_root, executables_root and bundles_root kwargs. A Dag bundle now decides what a coordinator sees, so two of its limits reach users. For Java, every JAR in the bundle goes on one classpath, so conflicting dependency versions cannot share a bundle, and a second Main-Class in it makes the entry point ambiguous. For Go, the coordinator skips files without the execute bit, so an object-store Dag bundle that drops it leaves the coordinator with no executable to find. --- .agents/skills/airflow-java-sdk/SKILL.md | 11 ++- .../language-sdks/go.rst | 82 +++++++++-------- .../language-sdks/java.rst | 92 ++++++++++--------- .../language-sdks/typescript.rst | 79 ++++++++-------- contributing-docs/30_new_language_sdk.rst | 8 +- go-sdk/README.md | 35 ++++--- go-sdk/example/bundle/Justfile | 5 +- java-sdk/README.md | 19 +++- .../airflow/sdk/plugin/AirflowSdkPlugin.kt | 4 +- java-sdk/scala_spark_example/README.md | 6 +- task-sdk/docs/executable-bundle-spec.rst | 27 +++--- task-sdk/docs/ts-bundle-spec.rst | 26 +++--- ts-sdk/README.md | 17 +++- ts-sdk/example/README.md | 14 ++- 14 files changed, 246 insertions(+), 179 deletions(-) diff --git a/.agents/skills/airflow-java-sdk/SKILL.md b/.agents/skills/airflow-java-sdk/SKILL.md index 8edb742679535..9921832eafee5 100644 --- a/.agents/skills/airflow-java-sdk/SKILL.md +++ b/.agents/skills/airflow-java-sdk/SKILL.md @@ -50,8 +50,9 @@ subclasses must only import from `org.apache.airflow.sdk`; any import of ## Bundle composition and coordinator discovery -A **bundle** is a directory of JAR files (typically `build/bundle/`) placed on the coordinator's -`jars_root`. The coordinator scans the directory at task-dispatch time to find: +A **bundle** is a directory of JAR files (typically `build/bundle/`) placed in the Dag bundle named +by the coordinator's `task_handler_bundle_name` (the task's own Dag bundle when unset). The +coordinator scans that Dag bundle at task-dispatch time to find: 1. **`Main-Class`** (standard JAR manifest attribute) — the fully-qualified class name of the entry point that the coordinator invokes with `java -classpath … --comm … --logs …`. @@ -65,7 +66,7 @@ A **bundle** is a directory of JAR files (typically `build/bundle/`) placed on t `runtimeClasspath` and copies it into the shadow JAR manifest. In thin-JAR mode (`fatJar = false`), the value stays in the `airflow-sdk` JAR deployed alongside the bundle JAR. -The Python coordinator (`JavaCoordinator`) scans every JAR under `jars_root` with +The Python coordinator (`JavaCoordinator`) scans every JAR in that Dag bundle with `_JarInfo.find()`, reads `META-INF/MANIFEST.MF` out of each ZIP, and collects `Main-Class` and `Airflow-Supervisor-Schema-Version` from whichever JARs carry them. The resolved schema version is then passed as the `schema_version` return value from `_build_execute_task_command`, which @@ -74,7 +75,7 @@ the base `SubprocessCoordinator` uses to negotiate the supervisor wire protocol. If `main_class` is set explicitly on the `JavaCoordinator` instance (via `[sdk] coordinators` kwargs), the scan uses it as a filter; otherwise the first JAR with a `Main-Class` attribute wins. Either way, `Airflow-Supervisor-Schema-Version` must be present in at least one JAR in -`jars_root` or startup fails. +the Dag bundle or startup fails. Every JAR in the Dag bundle goes on one classpath. --- @@ -117,7 +118,7 @@ E2E_TEST_MODE=java_sdk uv run --project airflow-e2e-tests pytest \ `coordinator.py` extends `SubprocessCoordinator`. The only method subclasses must implement is `_build_execute_task_command`, which returns `(argv, schema_version)`. Look at the existing -implementation for how `jars_root`, `java_executable`, `jvm_args`, and `main_class` are +implementation for how the scanned Dag bundle, `java_executable`, `jvm_args`, and `main_class` are assembled into the command. Do not reach into the JVM process from Python beyond what this method provides. 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 bd2a3e4ac761a..b0599035f9601 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst @@ -48,7 +48,7 @@ Prerequisites * Go 1.24 or later to build and pack bundles. This is a build-time requirement only; the worker that runs a packed bundle needs no Go toolchain, because the bundle is a self-contained native executable. -* The packed bundle must be accessible from the Airflow worker, under a directory the coordinator scans. +* The packed bundle must be accessible from the Airflow worker, in the Dag bundle the coordinator scans. * The ``apache-airflow-task-sdk`` package (installed with Airflow) provides the coordinator; no additional Python packages are needed. @@ -178,33 +178,44 @@ to pass on with ``bundle.Register(reports.Handlers()...)``. Coordinator configuration ~~~~~~~~~~~~~~~~~~~~~~~~~~~ -Register the coordinator and route the queue to it under ``[sdk]`` in ``airflow.cfg`` (or the equivalent -``AIRFLOW__SDK__*`` environment variables): +Register a Dag bundle for the packed bundles, register the coordinator, and route the queue to it in +``airflow.cfg`` (or the equivalent ``AIRFLOW__*`` environment variables): .. code-block:: ini + [dag_processor] + dag_bundle_config_list = [ + {"name": "dags-folder", "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", "kwargs": {}}, + { + "name": "go-task-handlers", + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {"path": "/opt/airflow/go-task-handlers"} + } + ] + [sdk] coordinators = { "go": { "classpath": "airflow.sdk.coordinators.executable.ExecutableCoordinator", - "kwargs": {"executables_root": ["~/airflow/executable-bundles"]} + "kwargs": {"task_handler_bundle_name": "go-task-handlers"} } } queue_to_coordinator = {"golang": "go"} -``executables_root`` is one or more directories the coordinator scans for bundles; ``queue_to_coordinator`` -routes stub tasks with ``queue="golang"`` to this Go coordinator. See :ref:`go-sdk/coordinator-config` for -the full list of accepted ``kwargs``. +``task_handler_bundle_name`` names the Dag bundle the coordinator scans for packed bundles; +``queue_to_coordinator`` routes stub tasks with ``queue="golang"`` to this Go coordinator. See +:ref:`go-sdk/coordinator-config` for the full list of accepted ``kwargs`` and how bundles are located. There is no separate Go worker to run: the Airflow worker forks the bundle binary once per task instance. .. note:: - The coordinator is part of the Airflow worker, so the ``[sdk]`` config (and the bundle files in - ``executables_root``) only need to be present wherever tasks actually execute. With ``CeleryExecutor``, - setting it on the Celery workers is sufficient. With ``LocalExecutor``, tasks run inside the scheduler - process, so it must be set where the scheduler can read it. The API server and Dag processor do not need - it. + The coordinator is part of the Airflow worker, so the ``[sdk]`` config and the packed bundle files only + need to be present wherever tasks actually execute. With ``CeleryExecutor``, setting it on the Celery + workers is sufficient. With ``LocalExecutor``, tasks run inside the scheduler process, so it must be set + where the scheduler can read it. The API server and Dag processor do not need it. The worker resolves + ``task_handler_bundle_name`` through ``[dag_processor] dag_bundle_config_list``, so the workers need that + setting too. Writing tasks ------------- @@ -444,12 +455,12 @@ Build and pack in one step; any flags after ``--`` are forwarded verbatim to ``g go tool airflow-go-pack ./example/bundle -- -trimpath -tags=prod -Use ``--output `` to write the packed bundle straight into a directory the coordinator scans -(``executables_root``): +Use ``--output `` to write the packed bundle straight into the directory of the Dag bundle the +coordinator scans (see `Deploying`_): .. code-block:: bash - go tool airflow-go-pack --output ~/airflow/executable-bundles/sample-dag-bundle ./example/bundle + go tool airflow-go-pack --output /opt/airflow/go-task-handlers/sample-dag-bundle ./example/bundle Cross-platform builds ~~~~~~~~~~~~~~~~~~~~~~~ @@ -461,7 +472,7 @@ machine (for example, deploying to a Linux host from an Apple-silicon ``darwin/a .. code-block:: bash go tool airflow-go-pack --goos linux --goarch amd64 \ - --output ~/airflow/executable-bundles/sample-dag-bundle \ + --output /opt/airflow/go-task-handlers/sample-dag-bundle \ ./example/bundle Alternatively, pack a pre-built binary with ``--executable`` / ``--source``. The packer normally execs the @@ -485,12 +496,16 @@ with ``--airflow-metadata``: Deploying ~~~~~~~~~ -Copy or mount the packed bundle into a directory listed in the coordinator's ``executables_root``. The -:class:`~airflow.sdk.coordinators.executable.ExecutableCoordinator` scans those directories recursively, +Copy or mount the packed bundle into the Dag bundle named by the coordinator's ``task_handler_bundle_name``. +The :class:`~airflow.sdk.coordinators.executable.ExecutableCoordinator` scans that Dag bundle recursively, matches the incoming ``dag_id`` against each bundle's manifest, verifies the bundle's integrity hash, and launches the matching bundle. Bundles are identified by the trailer magic, not by filename (no extension on Linux/macOS, ``.exe`` on Windows), so the file name on the worker is irrelevant. +Only files with the executable bit set are considered, so the packed bundles need a Dag bundle that keeps +it. A ``LocalDagBundle`` over a directory you manage does; object-store Dag bundles such as ``S3DagBundle`` +do not. + .. _go-sdk/coordinator-config: :class:`~airflow.sdk.coordinators.executable.ExecutableCoordinator` configuration @@ -506,16 +521,11 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the * - Parameter - Default - Description - * - ``executables_root`` - - *(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). 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_handler_bundle_name`` + - *(task's own Dag bundle)* + - Name of the Dag bundle scanned recursively for executable bundles. It must be registered in + ``[dag_processor] dag_bundle_config_list``. It is checked when the ``[sdk]`` configuration is + loaded, so a typo fails there rather than on the first task. * - ``task_startup_timeout`` - ``10.0`` - Seconds to wait for the bundle subprocess to connect after launch. Increase this if your @@ -523,15 +533,15 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the .. note:: - **Locating bundles.** ``executables_root`` and ``dag_bundle_name`` are mutually exclusive, - and both are optional: + **Locating bundles.** Packed bundles are read from a Dag bundle, so they are delivered, refreshed and + versioned by the same machinery as your Dags. - * 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. + * The expected layout is a separate Dag bundle for the packed bundles, named by + ``task_handler_bundle_name``, rather than the Dag bundle that holds your ``.py`` files. The task + uses the version that Dag bundle is on when it starts, pinned for the whole task. + * If ``task_handler_bundle_name`` is unset, packed bundles are read from the **task's own** Dag + bundle, pinned to the version the run was created with. + * The Dag bundle must keep the executable bit (see `Deploying`_). .. _go-sdk/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 3f69772e4d92d..5f875617e57e9 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst @@ -56,7 +56,8 @@ deployment process are the same. See :ref:`java-sdk/interface-api` for the inter The Python Dag source and the Java Gradle project are independent. They do not need to be in the same repository or have any particular relative filesystem layout. The Dag follows the deployment's normal Dag -delivery process; only the compiled Java bundle is deployed from the Gradle project to ``jars_root``. +delivery process; only the compiled Java bundle is deployed from the Gradle project, into a separate Dag +bundle that holds the JARs. Define the Python Dag ~~~~~~~~~~~~~~~~~~~~~ @@ -245,29 +246,42 @@ Deploy ``sales_pipeline.py`` separately through the deployment's normal Dag deli that process might sync it to ``${AIRFLOW_HOME}/dags/`` or package it in a Dag bundle; neither location is inside or relative to ``sales-pipeline-java/``. -Configure Airflow so the coordinator scans the parent JAR directory recursively and routes the ``java`` queue -to it. Add the following ``[sdk]`` section to the file selected by ``AIRFLOW_CONFIG`` (by default, -``${AIRFLOW_HOME}/airflow.cfg``), or set the equivalent ``AIRFLOW__SDK__*`` environment variables: +Configure Airflow to register the JAR directory as a Dag bundle, point the coordinator at it, and route the +``java`` queue to the coordinator. Add the following sections to the file selected by ``AIRFLOW_CONFIG`` (by +default, ``${AIRFLOW_HOME}/airflow.cfg``), or set the equivalent ``AIRFLOW__*`` environment variables: .. code-block:: ini + [dag_processor] + dag_bundle_config_list = [ + {"name": "dags-folder", "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", "kwargs": {}}, + { + "name": "java-task-handlers", + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {"path": "/opt/airflow/jars/sales-pipeline"} + } + ] + [sdk] coordinators = { "java": { "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", - "kwargs": {"jars_root": ["/opt/airflow/jars"]} + "kwargs": {"task_handler_bundle_name": "java-task-handlers"} } } queue_to_coordinator = {"java": "java"} ``java`` is a user-chosen coordinator name, not a reserved value. The value assigned to the queue in -``queue_to_coordinator`` must match a key in ``coordinators``. +``queue_to_coordinator`` must match a key in ``coordinators``, and ``task_handler_bundle_name`` must match a +Dag bundle name in ``dag_bundle_config_list``. See :ref:`java-sdk/coordinator-config` for how JARs are +located. Restart the affected Airflow components after changing this configuration. The coordinator config and JARs must be available wherever tasks execute. With ``CeleryExecutor``, that means the Celery workers; with -``LocalExecutor``, tasks run in subprocesses on the scheduler's host. The API server and Dag processor do not -need the JARs, while the Dag processor must receive ``sales_pipeline.py`` through the separate Dag delivery -process. +``LocalExecutor``, tasks run in subprocesses on the scheduler's host. The worker resolves +``task_handler_bundle_name`` through ``[dag_processor] dag_bundle_config_list``, so the workers need that +setting too. The API server and Dag processor do not need the JARs, while the Dag processor must receive +``sales_pipeline.py`` through the separate Dag delivery process. After Airflow has parsed the Dag, trigger it from the UI or command line: @@ -726,7 +740,7 @@ configuration): "java-jdk17": { "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", "kwargs": { - "jars_root": ["/opt/airflow/jars"], + "task_handler_bundle_name": "java-task-handlers", "jvm_args": ["-Djava.util.logging.config.file=/opt/airflow/logging.properties"] } } @@ -851,9 +865,10 @@ Then run: ./gradlew bundle -The ``build/bundle/`` directory contains all required JAR(s). Copy or mount it into the directory pointed to -by ``jars_root`` in the coordinator configuration. :class:`~airflow.sdk.coordinators.java.JavaCoordinator` -scans ``jars_root`` recursively and builds the classpath automatically. +The ``build/bundle/`` directory contains all required JAR(s). Copy or mount it into the Dag bundle named by +``task_handler_bundle_name`` in the coordinator configuration. +:class:`~airflow.sdk.coordinators.java.JavaCoordinator` scans that Dag bundle recursively and builds the +classpath automatically. .. note:: @@ -865,8 +880,7 @@ scans ``jars_root`` recursively and builds the classpath automatically. The plugin generates a fat JAR with the `Shadow `__ plugin by default. This is generally a good idea since you only deploy one JAR file to avoid dependency issues between projects. If this does not suit you, set ``fatJar = false`` in ``airflowBundle`` to produce thin JARs instead. The rest of the - process stays the same, but you will need to put all dependency JARs somewhere Airflow can find with - ``jars_root``. + process stays the same, but you will need to put all dependency JARs in the same Dag bundle. .. _java-sdk/build/maven: @@ -956,8 +970,8 @@ Then run: mvn package -The fat JAR is written to ``target/-.jar``. Copy it to the directory configured as -``jars_root`` in your coordinator. +The fat JAR is written to ``target/-.jar``. Copy it into the Dag bundle named by +``task_handler_bundle_name`` in your coordinator. **Option 2: thin JAR with separate dependencies** @@ -1017,8 +1031,8 @@ Then run: mvn package -``target/bundle/`` will contain the thin JAR and all runtime dependency JARs. Point ``jars_root`` at -this directory. +``target/bundle/`` will contain the thin JAR and all runtime dependency JARs. Copy or mount this +directory into the Dag bundle named by ``task_handler_bundle_name``. .. note:: @@ -1044,16 +1058,11 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the * - Parameter - Default - Description - * - ``jars_root`` - - *(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). 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``. + * - ``task_handler_bundle_name`` + - *(task's own Dag bundle)* + - Name of the Dag bundle scanned recursively for ``.jar`` files. It must be registered in + ``[dag_processor] dag_bundle_config_list``. It is checked when the ``[sdk]`` configuration is + loaded, so a typo fails there rather than on the first task. * - ``java_executable`` - ``"java"`` - Path to the ``java`` binary. Defaults to ``java`` on ``$PATH``. @@ -1062,10 +1071,9 @@ 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, 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. + - Explicit entry-point class. If omitted, the coordinator scans the Dag bundle for a JAR + whose manifest sets ``Main-Class``. If more than one JAR in that Dag bundle sets it, which + one runs is non-deterministic, so 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 @@ -1073,15 +1081,17 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the .. note:: - **Locating JARs.** ``jars_root`` and ``dag_bundle_name`` are mutually exclusive, and both - are optional: + **Locating JARs.** JARs are read from a Dag bundle, so they are delivered, refreshed and versioned + by the same machinery as your Dags. - * 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 + * The expected layout is a separate Dag bundle for the JARs, named by ``task_handler_bundle_name``, + rather than the Dag bundle that holds your ``.py`` files. The task uses the version that Dag 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. + * If ``task_handler_bundle_name`` is unset, JARs are read from the **task's own** Dag bundle, pinned + to the version the run was created with. + * Every JAR in the Dag bundle goes on one classpath, so all handlers in it share one set of + dependencies. To isolate conflicting dependency versions, put the handlers in a second Dag bundle + served by a second coordinator on its own queue. .. note:: @@ -1110,7 +1120,7 @@ point ``java_executable`` at it explicitly: "java-jdk17": { "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", "kwargs": { - "jars_root": ["/opt/airflow/jars"], + "task_handler_bundle_name": "java-task-handlers", "java_executable": "/opt/homebrew/opt/openjdk@17/bin/java" } } 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 47070c581da15..452a9dca2c68c 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst @@ -50,7 +50,7 @@ Prerequisites * Node.js 22 or later must be available on the Airflow worker nodes. * The packed bundle (a single ``bundle.min.mjs`` file, see :ref:`typescript-sdk/build`) must be accessible - from the worker, under a directory the coordinator scans. + from the worker, in the Dag bundle the coordinator scans. * The ``apache-airflow-task-sdk`` package (installed with Airflow) provides the coordinator; no additional Python packages are needed. * In the TypeScript project, install the ``apache-airflow-ts-sdk`` npm package to author task handlers: @@ -209,34 +209,46 @@ check. Coordinator configuration ~~~~~~~~~~~~~~~~~~~~~~~~~ -Register the coordinator and route the queue to it under ``[sdk]`` in ``airflow.cfg`` (or the equivalent -``AIRFLOW__SDK__*`` environment variables): +Register a Dag bundle for the packed bundles, register the coordinator, and route the queue to it in +``airflow.cfg`` (or the equivalent ``AIRFLOW__*`` environment variables): .. code-block:: ini + [dag_processor] + dag_bundle_config_list = [ + {"name": "dags-folder", "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", "kwargs": {}}, + { + "name": "ts-task-handlers", + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {"path": "/opt/airflow/ts-bundles"} + } + ] + [sdk] coordinators = { "ts": { "classpath": "airflow.sdk.coordinators.node.NodeCoordinator", - "kwargs": {"bundles_root": ["/opt/airflow/ts-bundles"]} + "kwargs": {"task_handler_bundle_name": "ts-task-handlers"} } } queue_to_coordinator = {"typescript": "ts"} -``bundles_root`` is one or more directories the coordinator scans for bundles; ``queue_to_coordinator`` -routes stub tasks with ``queue="typescript"`` to this coordinator. See -:ref:`typescript-sdk/coordinator-config` for the full list of accepted ``kwargs``. +``task_handler_bundle_name`` names the Dag bundle the coordinator scans for packed bundles; +``queue_to_coordinator`` routes stub tasks with ``queue="typescript"`` to this coordinator. See +:ref:`typescript-sdk/coordinator-config` for the full list of accepted ``kwargs`` and how the bundle is +located. There is no separate Node.js worker to run: the Airflow worker launches the bundle with ``node`` once per task instance. .. note:: - The coordinator runs inside the Airflow worker, so the ``[sdk]`` config (and the packed ``*.min.mjs`` - bundles in ``bundles_root``) only need to be present wherever tasks actually execute. With - ``CeleryExecutor``, setting them on the Celery workers is sufficient. With ``LocalExecutor``, tasks run - inside the scheduler process, so they must be present where the scheduler can read them. The API server - and Dag processor do not need them. + The coordinator runs inside the Airflow worker, so the ``[sdk]`` config and the packed ``*.min.mjs`` + bundles only need to be present wherever tasks actually execute. With ``CeleryExecutor``, setting them on + the Celery workers is sufficient. With ``LocalExecutor``, tasks run inside the scheduler process, so they + must be present where the scheduler can read them. The API server and Dag processor do not need them. The + worker resolves ``task_handler_bundle_name`` through ``[dag_processor] dag_bundle_config_list``, so the + workers need that setting too. .. _typescript-sdk/native-dag: @@ -499,7 +511,7 @@ that only supply utilities or types are not embedded. npx airflow-ts-pack src/main.ts --outdir dist Use ``--outdir `` to choose the output directory (default ``dist``), ``--outfile `` to name the -artifact exactly, which helps when one ``bundles_root`` holds several bundles, and ``--source `` to set +artifact exactly, which helps when one Dag bundle holds several bundles, and ``--source `` to set the source name displayed in the Airflow UI (default: the entry file's basename). ``--outdir`` and ``--outfile`` are mutually exclusive, and an ``--outfile`` name must end in ``.min.mjs`` so the coordinator can find it. @@ -507,12 +519,11 @@ can find it. Deploying ~~~~~~~~~ -Copy or mount the bundle into a directory listed in the coordinator's ``bundles_root``. -:class:`~airflow.sdk.coordinators.node.NodeCoordinator` searches the configured directories in order, -recursively, and launches the first integrity-verified ``*.min.mjs`` bundle whose metadata declares the task -instance's Dag. The artifact's name does not matter beyond that suffix, so one root can hold several bundles -and a Dag is routed to whichever declares it. If multiple bundles declare the same Dag, the first configured -root wins, and within a root the first in sorted path order. +Copy or mount the bundle into the Dag bundle named by the coordinator's ``task_handler_bundle_name``. +:class:`~airflow.sdk.coordinators.node.NodeCoordinator` searches that Dag bundle recursively and launches the +first integrity-verified ``*.min.mjs`` bundle whose metadata declares the task instance's Dag. The artifact's +name does not matter beyond that suffix, so one Dag bundle can hold several bundles and a Dag is routed to +whichever declares it. If multiple bundles declare the same Dag, the first in sorted path order wins. .. _typescript-sdk/coordinator-config: @@ -529,16 +540,11 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the * - Parameter - Default - Description - * - ``bundles_root`` - - *(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). 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``. + * - ``task_handler_bundle_name`` + - *(task's own Dag bundle)* + - Name of the Dag bundle searched recursively for an integrity-verified ``*.min.mjs`` bundle that + declares the requested Dag. It must be registered in ``[dag_processor] dag_bundle_config_list``. It is + checked when the ``[sdk]`` configuration is loaded, so a typo fails there rather than on the first task. * - ``node_executable`` - ``"node"`` - Path to the ``node`` binary. Defaults to ``node`` on ``$PATH``. @@ -549,15 +555,14 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the .. note:: - **Locating the bundle.** ``bundles_root`` and ``dag_bundle_name`` are mutually exclusive, and both - are optional: + **Locating the bundle.** Packed bundles are read from a Dag bundle, so they are delivered, refreshed and + versioned by the same machinery as your Dags. - * 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. + * The expected layout is a separate Dag bundle for the packed bundles, named by + ``task_handler_bundle_name``, rather than the Dag bundle that holds your ``.py`` files. The task uses + the version that Dag bundle is on when it starts, pinned for the whole task. + * If ``task_handler_bundle_name`` is unset, the bundle is read 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 34c4cdbb4a27a..86c65efc716ea 100644 --- a/contributing-docs/30_new_language_sdk.rst +++ b/contributing-docs/30_new_language_sdk.rst @@ -121,10 +121,10 @@ The method returns a ``(command, subprocess_schema_version)`` pair: 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. +class has already resolved: the Dag bundle named by ``task_handler_bundle_name``, +or the task's own Dag bundle when it is unset, pinned for the whole task. +Subclasses should scan those directories rather than locating artifacts +themselves. Supervisor Schema ~~~~~~~~~~~~~~~~~ diff --git a/go-sdk/README.md b/go-sdk/README.md index a6dedde08049b..5e8dba63ea414 100644 --- a/go-sdk/README.md +++ b/go-sdk/README.md @@ -250,15 +250,15 @@ the full range of task states, and alternate XCom backends without implementing go tool airflow-go-pack ./example/bundle -- -trimpath -tags=prod ``` - Use `--output ` to write the packed bundle straight into a directory the coordinator scans - (`executables_root`), and pass extra `go build` flags after `--`. + Use `--output ` to write the packed bundle straight into the directory of the Dag bundle the + coordinator scans (`task_handler_bundle_name`), and pass extra `go build` flags after `--`. For cross-compiling (e.g. deploy to a Linux host from an Apple-silicon (darwin/arm64) machine), pass `--goos`/`--goarch` and the packer cross-builds for you: ```bash go tool airflow-go-pack --goos linux --goarch amd64 \ - --output ~/airflow/executable-bundles/sample-dag-bundle \ + --output /opt/airflow/go-task-handlers/sample-dag-bundle \ ./example/bundle ``` @@ -280,24 +280,33 @@ the full range of task states, and alternate XCom backends without implementing > `tool github.com/apache/airflow/go-sdk/cmd/airflow-go-pack` to your bundle module's `go.mod` and run > it with `go tool airflow-go-pack`. This pins the packer version per project. -- Register the coordinator and route the queue to it, under `[sdk]` in `airflow.cfg` (or the equivalent - `AIRFLOW__SDK__*` env vars): +- Register a Dag bundle for the packed bundles, register the coordinator, and route the queue to it, in + `airflow.cfg` (or the equivalent `AIRFLOW__*` env vars): ```ini + [dag_processor] + dag_bundle_config_list = [ + {"name": "dags-folder", "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", "kwargs": {}}, + {"name": "go-task-handlers", "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", "kwargs": {"path": "/opt/airflow/go-task-handlers"}} + ] + [sdk] - coordinators = {"go": {"classpath": "airflow.sdk.coordinators.executable.ExecutableCoordinator", "kwargs": {"executables_root": ["~/airflow/executable-bundles"]}}} + coordinators = {"go": {"classpath": "airflow.sdk.coordinators.executable.ExecutableCoordinator", "kwargs": {"task_handler_bundle_name": "go-task-handlers"}}} queue_to_coordinator = {"golang": "go"} ``` - `executables_root` is one or more directories the coordinator scans for bundles; `queue_to_coordinator` - routes stub tasks with `queue="golang"` to this Go coordinator. + `task_handler_bundle_name` names the Dag bundle the coordinator scans for packed bundles (the task's own + Dag bundle when unset); `queue_to_coordinator` routes stub tasks with `queue="golang"` to this Go + coordinator. Only files with the executable bit are considered, so use a Dag bundle that keeps it: a + `LocalDagBundle` does, object-store Dag bundles such as `S3DagBundle` do not. > [!IMPORTANT] - > The coordinator is part of the Airflow worker, so the `[sdk]` config (and the bundle files in - > `executables_root`) only need to be present wherever tasks actually execute. With `CeleryExecutor`, - > setting it on the Celery workers is sufficient. With `LocalExecutor`, tasks run inside the scheduler - > process, so it must be set where the scheduler can read it. The API server and Dag processor do not - > need it. + > The coordinator is part of the Airflow worker, so the `[sdk]` config and the packed bundle files only + > need to be present wherever tasks actually execute. With `CeleryExecutor`, setting it on the Celery + > workers is sufficient. With `LocalExecutor`, tasks run inside the scheduler process, so it must be set + > where the scheduler can read it. The API server and Dag processor do not need it. The worker resolves + > `task_handler_bundle_name` through `[dag_processor] dag_bundle_config_list`, so the workers need that + > setting too. - Deploy the matching Python stub Dag (above) into Airflow. There is no separate Go worker to run: the Airflow worker forks the bundle binary once per task instance. diff --git a/go-sdk/example/bundle/Justfile b/go-sdk/example/bundle/Justfile index 5a12c4407aed1..07df82a2280a6 100644 --- a/go-sdk/example/bundle/Justfile +++ b/go-sdk/example/bundle/Justfile @@ -29,8 +29,9 @@ build: pack # DAG/task identity via --airflow-metadata, and appends the source plus # airflow-metadata.yaml plus AFBNDL01 trailer. The output is a single # self-contained executable bundle, named after the bundle's package -# directory and written to the current directory. Drop it into -# ExecutableCoordinator's executables_root to deploy. +# directory and written to the current directory. Drop it into the +# Dag bundle named by ExecutableCoordinator's task_handler_bundle_name +# to deploy. pack: @echo "Packing example DAG bundle..." go tool airflow-go-pack --output ../../bin/example_dags diff --git a/java-sdk/README.md b/java-sdk/README.md index b8600ab5d44fa..5ec0995708546 100644 --- a/java-sdk/README.md +++ b/java-sdk/README.md @@ -96,13 +96,22 @@ Now `cd example` into the example project, and * Ensure the `java` command is available in the same environment the Airflow task worker is in. -* Configure Airflow to route tasks in the *java* queue to be run with Java: +* Register the packaged example as a Dag bundle, and configure Airflow to route tasks in the *java* + queue to be run with Java from it: ```bash + export AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST='[ + {"name": "dags-folder", "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", "kwargs": {}}, + { + "name": "java-task-handlers", + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {"path": "/opt/airflow/java-sdk/example/build/bundle"} + } + ]' export AIRFLOW__SDK__COORDINATORS='{ "java": { "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", - "kwargs": {"jars_root": ["/opt/airflow/java-sdk/example/build/bundle"]} + "kwargs": {"task_handler_bundle_name": "java-task-handlers"} } }' export AIRFLOW__SDK__QUEUE_TO_COORDINATOR='{"java": "java"}' @@ -644,9 +653,9 @@ where Airflow can find it. When the Airflow supervisor identifies that a task should run with Java, it launches the JVM application as a subprocess. The flow is: -1. `JavaCoordinator.execute_task()` (Python) scans `jars_root`, builds the - classpath, and spawns `java -cp --comm=: - --logs=:`. +1. `JavaCoordinator.execute_task()` (Python) scans the Dag bundle named by + `task_handler_bundle_name`, builds the classpath, and spawns + `java -cp --comm=: --logs=:`. 2. `Server.kt` connects to both sockets immediately on startup. 3. The supervisor sends a `StartupDetails` MessagePack message; the JVM reads it, looks up the matching task by `dag_id` + `task_id`, and calls the diff --git a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt index 95836c12fc5fe..3756ef688e174 100644 --- a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt +++ b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt @@ -83,8 +83,8 @@ abstract class AirflowBundleExtension { * * The plugin automatically sets the `Main-Class` metadata, and provides a new * task `bundle` to create a Dag bundle in one command. This builds deploy-ready - * artifacts to `build/bundle/` that can be copied directly into an Airflow Java - * coordinator's `jars_root`. + * artifacts to `build/bundle/` that can be copied directly into the Dag bundle + * named by an Airflow Java coordinator's `task_handler_bundle_name`. * * By default, plugin `com.github.johnrengelman.shadow` is applied automatically * to enable fat JAR build. In this mode, one single JAR with user code and all diff --git a/java-sdk/scala_spark_example/README.md b/java-sdk/scala_spark_example/README.md index f59c01e0e52c6..f736ecd08ee2b 100644 --- a/java-sdk/scala_spark_example/README.md +++ b/java-sdk/scala_spark_example/README.md @@ -41,7 +41,9 @@ cd scala_spark_example ``` `fatJar` is disabled, so `build/bundle/` holds the bundle JAR plus every runtime -JAR (Spark included) — copy it into a Java coordinator's `jars_root`. +JAR (Spark included) — copy it into the Dag bundle named by a Java coordinator's +`task_handler_bundle_name`. Every JAR in that Dag bundle goes on one classpath, +so give Spark its own Dag bundle, coordinator and queue. ## Running Spark under the Java SDK @@ -57,7 +59,7 @@ and bypasses that launcher, so the coordinator has to pass them itself via "scala-jdk": { "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", "kwargs": { - "jars_root": ["/path/to/scala-jars"], + "task_handler_bundle_name": "scala-task-handlers", "main_class": "org.apache.airflow.example.ScalaSparkBundleBuilder", "jvm_args": [ "-Xmx512m", diff --git a/task-sdk/docs/executable-bundle-spec.rst b/task-sdk/docs/executable-bundle-spec.rst index 856aeca94e922..f83443741b101 100644 --- a/task-sdk/docs/executable-bundle-spec.rst +++ b/task-sdk/docs/executable-bundle-spec.rst @@ -283,8 +283,8 @@ Code Signing The bundle format itself does not require OS-level code signing. ``binary_sha256`` provides integrity against truncation, in-flight -corruption, and naive tampering, and Airflow's threat model treats -``executables_root`` as Deployment-Manager-controlled — *authenticity* +corruption, and naive tampering, and Airflow's threat model treats the Dag +bundle holding the bundles as Deployment-Manager-controlled — *authenticity* (signed by a trusted identity) is a deployment-time concern rather than a bundle-format one. @@ -299,22 +299,23 @@ yields a byte-identical bundle file (and therefore a stable Deployment Layout ----------------- -Bundle files are placed **as-is** in any of the directories configured as the -``executables_root`` kwarg on the +Bundle files are placed **as-is** in the Dag bundle named by the +``task_handler_bundle_name`` kwarg on the :class:`~airflow.sdk.coordinators.executable.ExecutableCoordinator` entry -under ``[sdk] coordinators``. The scanner walks each directory **recursively** -and considers only regular files whose **executable bit is set** for the -invoking user; files without the executable bit are skipped without reading -their trailer. For each candidate it reads the last 64 bytes and treats files -whose magic matches ``"AFBNDL01"`` as bundles. Matched files are then -SHA-256-verified per the reader algorithm; a mismatch demotes the file back -to "ignored, with an error log." Files without the magic are silently -ignored, so non-bundle files (READMEs, dotfiles) MAY share the directory +under ``[sdk] coordinators`` (or, when it is unset, in the task's own Dag +bundle). The scanner walks the Dag bundle **recursively** and considers only +regular files whose **executable bit is set** for the invoking user; files +without the executable bit are skipped without reading their trailer, so the +Dag bundle must preserve it. For each candidate it reads the last 64 bytes and +treats files whose magic matches ``"AFBNDL01"`` as bundles. Matched files are +then SHA-256-verified per the reader algorithm; a mismatch demotes the file +back to "ignored, with an error log." Files without the magic are silently +ignored, so non-bundle files (READMEs, dotfiles) MAY share the Dag bundle without interfering with the scan. :: - /opt/airflow/executable-bundles/ + /opt/airflow/go-task-handlers/ ├── example ├── team-a/ │ └── pipeline diff --git a/task-sdk/docs/ts-bundle-spec.rst b/task-sdk/docs/ts-bundle-spec.rst index 705227efa55b2..03ee63ceea0a6 100644 --- a/task-sdk/docs/ts-bundle-spec.rst +++ b/task-sdk/docs/ts-bundle-spec.rst @@ -24,8 +24,8 @@ This document specifies the bundle format produced by ``airflow-ts-pack`` and co Artifact Name ------------- -A bundle's name must end in ``.min.mjs``. Nothing else about it is significant: the coordinator searches each -configured root recursively and routes on embedded metadata, so one root may hold several differently named bundles. +A bundle's name must end in ``.min.mjs``. Nothing else about it is significant: the coordinator searches its Dag +bundle recursively and routes on embedded metadata, so one Dag bundle may hold several differently named bundles. ``airflow-ts-pack`` writes ``bundle.min.mjs`` by default and accepts ``--outfile`` for any other name ending in that suffix. @@ -153,11 +153,12 @@ documents, and neither is used to execute the bundle. Reader and Selection Algorithm ------------------------------ -For each candidate in ``bundles_root``, the coordinator: +For each candidate in the Dag bundle named by ``task_handler_bundle_name`` (or, when it is unset, the task's own Dag +bundle), the coordinator: -1. Opens it once. Candidates are the files whose name ends in ``.min.mjs``, found by walking each root recursively, - roots in configured order and each directory's entries in sorted order, so selection does not depend on the order - a filesystem returns entries in. Directories are deduplicated by ``(st_dev, st_ino)``, so a symlink loop +1. Opens it once. Candidates are the files whose name ends in ``.min.mjs``, found by walking the Dag bundle + recursively, each directory's entries in sorted order, so selection does not depend on the order a filesystem + returns entries in. Directories are deduplicated by ``(st_dev, st_ino)``, so a symlink loop terminates the walk instead of exhausting the interpreter stack. 2. Reads a bounded first line and decodes the named metadata and code ranges. 3. Reads the bounded metadata line and checks that the declared metadata range matches its physical location. @@ -171,15 +172,14 @@ For each candidate in ``bundles_root``, the coordinator: 9. Resolves the supervisor schema version and selects the first usable match. A missing, unrelated, unreadable, malformed, corrupt, or incompatible earlier candidate does not prevent selection -of a later usable match. When more than one usable bundle declares the same Dag, the first configured match wins. If -none matches, the error identifies the requested Dag, searched roots, and rejected candidates. +of a later usable match. When more than one usable bundle declares the same Dag, the first match in walk order wins. +If none matches, the error identifies the requested Dag, the searched directory, and rejected candidates. -Every ``.min.mjs`` file under a root is therefore opened, and one that is not a usable bundle is named among those -rejected candidates. ``bundles_root`` names directories of deployed Airflow bundles, so unrelated minified modules do -not belong there. +Every ``.min.mjs`` file in the Dag bundle is therefore opened, and one that is not a usable bundle is named among +those rejected candidates. That Dag bundle holds deployed Airflow bundles, so unrelated minified modules do not belong +there. -The coordinator does not cache Dag-to-path routing. It checks root ordering and the current deployed files for each -task selection. It may reuse section digests from a bounded process-local cache when the open file identity, +The coordinator does not cache Dag-to-path routing. It checks the current deployed files for each task selection. It may reuse section digests from a bounded process-local cache when the open file identity, timestamps, size, layout ranges, and declared digests have not changed. Integrity, Authenticity, and Provenance diff --git a/ts-sdk/README.md b/ts-sdk/README.md index 361e987fa3116..75c5859b5ca8f 100644 --- a/ts-sdk/README.md +++ b/ts-sdk/README.md @@ -88,18 +88,29 @@ sales_pipeline() Airflow coordinator config: ```ini +[dag_processor] +dag_bundle_config_list = [ + {"name": "dags-folder", "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", "kwargs": {}}, + { + "name": "ts-task-handlers", + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {"path": "/opt/airflow/ts-bundles"} + } + ] + [sdk] coordinators = { "ts": { "classpath": "airflow.sdk.coordinators.node.NodeCoordinator", - "kwargs": {"bundles_root": ["/opt/airflow/ts-bundles"]} + "kwargs": {"task_handler_bundle_name": "ts-task-handlers"} } } queue_to_coordinator = {"typescript": "ts"} ``` -Each configured bundle directory is searched recursively for `*.min.mjs` bundles built with `airflow-ts-pack` -(see [Packing bundles](#packing-bundles)), which embeds the Airflow metadata in the bundle itself. +The Dag bundle named by `task_handler_bundle_name` (the task's own Dag bundle when unset) is searched +recursively for `*.min.mjs` bundles built with `airflow-ts-pack` (see [Packing bundles](#packing-bundles)), +which embeds the Airflow metadata in the bundle itself. TypeScript entrypoint: diff --git a/ts-sdk/example/README.md b/ts-sdk/example/README.md index 94c2ead395110..85b5c74099f80 100644 --- a/ts-sdk/example/README.md +++ b/ts-sdk/example/README.md @@ -58,14 +58,22 @@ ts-sdk/example/dist/ ## Airflow Configuration -Configure Airflow to route the `typescript` queue to the Node coordinator and -point it at the example bundle directory: +Register the example bundle directory as a Dag bundle, and configure Airflow to +route the `typescript` queue to the Node coordinator that reads it: ```bash +export AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST='[ + {"name": "dags-folder", "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", "kwargs": {}}, + { + "name": "ts-task-handlers", + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {"path": "/absolute/path/to/airflow/ts-sdk/example/dist"} + } +]' export AIRFLOW__SDK__COORDINATORS='{ "ts": { "classpath": "airflow.sdk.coordinators.node.NodeCoordinator", - "kwargs": {"bundles_root": ["/absolute/path/to/airflow/ts-sdk/example/dist"]} + "kwargs": {"task_handler_bundle_name": "ts-task-handlers"} } }' export AIRFLOW__SDK__QUEUE_TO_COORDINATOR='{"typescript": "ts"}' From b9cde7fe612e1d6a84a301cc63fedaa256840a3c Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 09:25:19 +0000 Subject: [PATCH 04/94] Tidy names and mocks in the artifact bundle tests Review flagged a Java test fixture still named after the removed jars_root kwarg, an e2e helper without a verb, and two execute_task patches that were not autospecced like the ones beside them. --- .../tests/airflow_e2e_tests/conftest.py | 8 ++++---- .../task_sdk/coordinators/java/test_coordinator.py | 12 ++++++------ .../tests/task_sdk/coordinators/test_subprocess.py | 4 ++-- 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py b/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py index db08b27b06397..1282326cc9a1c 100644 --- a/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py +++ b/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py @@ -293,7 +293,7 @@ def _setup_xcom_object_storage_integration(dot_env_file, tmp_dir): ] -def _dag_bundle_config_with_artifacts(artifact_bundles: dict[str, str]) -> str: +def _build_dag_bundle_config(artifact_bundles: dict[str, str]) -> str: """Return a ``dag_bundle_config_list`` of the Dags folder plus a ``LocalDagBundle`` per artifact path. Registration is what makes a coordinator's ``task_handler_bundle_name`` resolvable on the @@ -476,7 +476,7 @@ def _setup_java_sdk_integration(dot_env_file, tmp_dir): # main_class (Spark's large classpath makes Main-Class discovery ambiguous) # and carries Spark's Java 17 module openings, a small driver heap, and a # longer startup timeout for its large dependency classpath. - dag_bundle_config = _dag_bundle_config_with_artifacts( + dag_bundle_config = _build_dag_bundle_config( { "java-task-handlers": "/opt/airflow/java-jars", "scala-task-handlers": "/opt/airflow/scala-jars", @@ -650,7 +650,7 @@ def _setup_go_sdk_integration(dot_env_file, tmp_dir): # Coordinator registry: maps the logical name "go-sdk" to ExecutableCoordinator, # which scans the go-task-handlers Dag bundle for the packed bundle by dag_id. # Queue mapping: routes tasks on the "golang" queue to "go-sdk". - dag_bundle_config = _dag_bundle_config_with_artifacts({"go-task-handlers": "/opt/airflow/go-bundles"}) + dag_bundle_config = _build_dag_bundle_config({"go-task-handlers": "/opt/airflow/go-bundles"}) coordinator_config = json.dumps( { "go-sdk": { @@ -756,7 +756,7 @@ def _setup_ts_sdk_integration(dot_env_file, tmp_dir): for dag_file in ("typescript_example.py", "typescript_taskflow_example.py"): copyfile(TS_SDK_EXAMPLE_PATH / "dags" / dag_file, tmp_dir / "dags" / dag_file) - dag_bundle_config = _dag_bundle_config_with_artifacts({"ts-task-handlers": "/opt/airflow/ts-bundles"}) + dag_bundle_config = _build_dag_bundle_config({"ts-task-handlers": "/opt/airflow/ts-bundles"}) coordinator_config = json.dumps( { "ts": { 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 eda377e6251fd..56e21d970e612 100644 --- a/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py @@ -252,18 +252,18 @@ def test_build_command_scans_passed_roots_in_colocated_mode(self, tmp_path): @pytest.fixture -def jars_root(tmp_path): +def jars_dir(tmp_path): _make_jar(tmp_path.joinpath("app.jar"), main_class="com.example.TaskRunner", schema_version="2026-06-16") return tmp_path @pytest.fixture -def java_task_handlers(jars_root): - """Register *jars_root* as the ``java-task-handlers`` Dag bundle and return that name.""" +def java_task_handlers(jars_dir): + """Register *jars_dir* as the ``java-task-handlers`` Dag bundle and return that name.""" bundle = { "name": "java-task-handlers", "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", - "kwargs": {"path": str(jars_root)}, + "kwargs": {"path": str(jars_dir)}, } with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps([bundle])}): yield bundle["name"] @@ -337,12 +337,12 @@ def test_java_executable_is_first_arg(self, java_task_handlers, mock_client): ) assert cmd[0] == "/usr/lib/jvm/java-17/bin/java" - def test_classpath_flag_and_value_present(self, jars_root, java_task_handlers, mock_client): + def test_classpath_flag_and_value_present(self, jars_dir, java_task_handlers, mock_client): cmd = self._captured_popen_cmd(java_task_handlers, mock_client) assert "-classpath" in cmd cp_idx = cmd.index("-classpath") classpath = cmd[cp_idx + 1] - assert jars_root.joinpath("app.jar").as_posix() in classpath + assert jars_dir.joinpath("app.jar").as_posix() in classpath def test_main_class_present(self, java_task_handlers, mock_client): cmd = self._captured_popen_cmd(java_task_handlers, mock_client) diff --git a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py index 909a2c605f0ba..f836c6797350c 100644 --- a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py +++ b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py @@ -937,7 +937,7 @@ class TestExecuteTaskBundleWiring: @patch("airflow.sdk.coordinators._subprocess.BundleVersionLock", autospec=True) @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle", autospec=True) - @patch.object(_PopenActivitySubprocess, "start") + @patch.object(_PopenActivitySubprocess, "start", autospec=True) def test_task_bundle_mode_binds_forwards_roots_and_locks( self, mock_start, mock_initialize, mock_lock, mock_client, tmp_path ): @@ -963,7 +963,7 @@ def test_task_bundle_mode_binds_forwards_roots_and_locks( @patch("airflow.sdk.coordinators._subprocess.BundleVersionLock", autospec=True) @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle", autospec=True) - @patch.object(_PopenActivitySubprocess, "start") + @patch.object(_PopenActivitySubprocess, "start", autospec=True) def test_named_bundle_mode_locks_the_tree_it_scans( self, mock_start, mock_initialize, mock_lock, mock_client, tmp_path ): From 23898f0586c10e539981a881187ebad9a0880d9f Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 09:26:23 +0000 Subject: [PATCH 05/94] Say artifact Dag bundles are registered on every component The guides said only the workers need the artifact Dag bundle registered, which reads as permission to register it there alone. The [sdk] configuration is also loaded by the scheduler under KubernetesExecutor, and it is rejected wherever the name is missing, so the bundle belongs in the shared bundle list. --- .../docs/authoring-and-scheduling/language-sdks/go.rst | 7 ++++--- .../docs/authoring-and-scheduling/language-sdks/java.rst | 9 +++++---- .../language-sdks/typescript.rst | 7 ++++--- go-sdk/README.md | 7 ++++--- 4 files changed, 17 insertions(+), 13 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 b0599035f9601..d3b1f7743df7e 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst @@ -213,9 +213,10 @@ There is no separate Go worker to run: the Airflow worker forks the bundle binar The coordinator is part of the Airflow worker, so the ``[sdk]`` config and the packed bundle files only need to be present wherever tasks actually execute. With ``CeleryExecutor``, setting it on the Celery workers is sufficient. With ``LocalExecutor``, tasks run inside the scheduler process, so it must be set - where the scheduler can read it. The API server and Dag processor do not need it. The worker resolves - ``task_handler_bundle_name`` through ``[dag_processor] dag_bundle_config_list``, so the workers need that - setting too. + where the scheduler can read it. The API server and Dag processor do not need it. Register the Dag + bundle in ``[dag_processor] dag_bundle_config_list`` on every component, like your other Dag bundles: + the worker resolves ``task_handler_bundle_name`` through it, and wherever the ``[sdk]`` config is read + it is rejected if the name is missing there. Writing tasks ------------- 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 5f875617e57e9..36c5f12243674 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst @@ -278,10 +278,11 @@ located. Restart the affected Airflow components after changing this configuration. The coordinator config and JARs must be available wherever tasks execute. With ``CeleryExecutor``, that means the Celery workers; with -``LocalExecutor``, tasks run in subprocesses on the scheduler's host. The worker resolves -``task_handler_bundle_name`` through ``[dag_processor] dag_bundle_config_list``, so the workers need that -setting too. The API server and Dag processor do not need the JARs, while the Dag processor must receive -``sales_pipeline.py`` through the separate Dag delivery process. +``LocalExecutor``, tasks run in subprocesses on the scheduler's host. Register the Dag bundle in +``[dag_processor] dag_bundle_config_list`` on every component, like your other Dag bundles: the worker +resolves ``task_handler_bundle_name`` through it, and wherever the ``[sdk]`` config is read it is rejected if +the name is missing there. The API server and Dag processor do not need the JARs, while the Dag processor must +receive ``sales_pipeline.py`` through the separate Dag delivery process. After Airflow has parsed the Dag, trigger it from the UI or command line: 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 452a9dca2c68c..dfb7ed7d636e7 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst @@ -246,9 +246,10 @@ task instance. The coordinator runs inside the Airflow worker, so the ``[sdk]`` config and the packed ``*.min.mjs`` bundles only need to be present wherever tasks actually execute. With ``CeleryExecutor``, setting them on the Celery workers is sufficient. With ``LocalExecutor``, tasks run inside the scheduler process, so they - must be present where the scheduler can read them. The API server and Dag processor do not need them. The - worker resolves ``task_handler_bundle_name`` through ``[dag_processor] dag_bundle_config_list``, so the - workers need that setting too. + must be present where the scheduler can read them. The API server and Dag processor do not need them. + Register the Dag bundle in ``[dag_processor] dag_bundle_config_list`` on every component, like your other + Dag bundles: the worker resolves ``task_handler_bundle_name`` through it, and wherever the ``[sdk]`` config + is read it is rejected if the name is missing there. .. _typescript-sdk/native-dag: diff --git a/go-sdk/README.md b/go-sdk/README.md index 5e8dba63ea414..da477119b0c38 100644 --- a/go-sdk/README.md +++ b/go-sdk/README.md @@ -304,9 +304,10 @@ the full range of task states, and alternate XCom backends without implementing > The coordinator is part of the Airflow worker, so the `[sdk]` config and the packed bundle files only > need to be present wherever tasks actually execute. With `CeleryExecutor`, setting it on the Celery > workers is sufficient. With `LocalExecutor`, tasks run inside the scheduler process, so it must be set - > where the scheduler can read it. The API server and Dag processor do not need it. The worker resolves - > `task_handler_bundle_name` through `[dag_processor] dag_bundle_config_list`, so the workers need that - > setting too. + > where the scheduler can read it. The API server and Dag processor do not need it. Register the Dag + > bundle in `[dag_processor] dag_bundle_config_list` on every component, like your other Dag bundles: + > the worker resolves `task_handler_bundle_name` through it, and wherever the `[sdk]` config is read it + > is rejected if the name is missing there. - Deploy the matching Python stub Dag (above) into Airflow. There is no separate Go worker to run: the Airflow worker forks the bundle binary once per task instance. From 8997a357c46d7e90f43cf19ef700d57ec84e2da0 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 10:16:08 +0000 Subject: [PATCH 06/94] Drop em-dashes from the rewritten doc sentences Two sentences rewritten for artifact Dag bundles kept their em-dashes, which the docs style for this work avoids. --- java-sdk/scala_spark_example/README.md | 2 +- task-sdk/docs/executable-bundle-spec.rst | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/java-sdk/scala_spark_example/README.md b/java-sdk/scala_spark_example/README.md index f736ecd08ee2b..34f248a017cf1 100644 --- a/java-sdk/scala_spark_example/README.md +++ b/java-sdk/scala_spark_example/README.md @@ -41,7 +41,7 @@ cd scala_spark_example ``` `fatJar` is disabled, so `build/bundle/` holds the bundle JAR plus every runtime -JAR (Spark included) — copy it into the Dag bundle named by a Java coordinator's +JAR (Spark included). Copy it into the Dag bundle named by a Java coordinator's `task_handler_bundle_name`. Every JAR in that Dag bundle goes on one classpath, so give Spark its own Dag bundle, coordinator and queue. diff --git a/task-sdk/docs/executable-bundle-spec.rst b/task-sdk/docs/executable-bundle-spec.rst index f83443741b101..27e48a0d04fde 100644 --- a/task-sdk/docs/executable-bundle-spec.rst +++ b/task-sdk/docs/executable-bundle-spec.rst @@ -284,9 +284,9 @@ Code Signing The bundle format itself does not require OS-level code signing. ``binary_sha256`` provides integrity against truncation, in-flight corruption, and naive tampering, and Airflow's threat model treats the Dag -bundle holding the bundles as Deployment-Manager-controlled — *authenticity* -(signed by a trusted identity) is a deployment-time concern rather than a -bundle-format one. +bundle holding the bundles as Deployment-Manager-controlled, so +*authenticity* (signed by a trusted identity) is a deployment-time concern +rather than a bundle-format one. **Compressors** such as UPX are NOT supported. They rewrite the file end-to-end, destroying both the trailer and the hash invariant. From 8983f6dc234ff14015dd8791da498036fdf1d46f Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 1 Oct 2026 13:45:21 +0000 Subject: [PATCH 07/94] Say the Dag processor needs the language SDK setup too In Airflow 3.4 the Dag processor checks a Python Dag's stub tasks against the task handlers the language SDK artifacts register, so a deployment that follows the old worker-only advice fails that check. task_handler_bundle_name only serves those stub tasks, which the docs did not say. --- .../language-sdks/go.rst | 32 ++++++++++-------- .../language-sdks/index.rst | 6 ++-- .../language-sdks/java.rst | 33 +++++++++++-------- .../language-sdks/typescript.rst | 29 ++++++++-------- .../src/airflow/config_templates/config.yml | 15 ++++++--- go-sdk/README.md | 23 +++++++------ java-sdk/README.md | 11 +++++-- ts-sdk/README.md | 8 ++++- 8 files changed, 96 insertions(+), 61 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 d3b1f7743df7e..f35012832d6e2 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst @@ -48,7 +48,8 @@ Prerequisites * Go 1.24 or later to build and pack bundles. This is a build-time requirement only; the worker that runs a packed bundle needs no Go toolchain, because the bundle is a self-contained native executable. -* The packed bundle must be accessible from the Airflow worker, in the Dag bundle the coordinator scans. +* The packed bundle must be accessible from the Airflow worker and the Dag processor, in the Dag bundle the + coordinator scans, and built for the operating system and CPU architecture of both. * The ``apache-airflow-task-sdk`` package (installed with Airflow) provides the coordinator; no additional Python packages are needed. @@ -210,13 +211,14 @@ There is no separate Go worker to run: the Airflow worker forks the bundle binar .. note:: - The coordinator is part of the Airflow worker, so the ``[sdk]`` config and the packed bundle files only - need to be present wherever tasks actually execute. With ``CeleryExecutor``, setting it on the Celery - workers is sufficient. With ``LocalExecutor``, tasks run inside the scheduler process, so it must be set - where the scheduler can read it. The API server and Dag processor do not need it. Register the Dag - bundle in ``[dag_processor] dag_bundle_config_list`` on every component, like your other Dag bundles: - the worker resolves ``task_handler_bundle_name`` through it, and wherever the ``[sdk]`` config is read - it is rejected if the name is missing there. + The ``[sdk]`` config and the packed bundle files must be present wherever tasks execute and on the Dag + processor. With ``CeleryExecutor``, tasks execute on the Celery workers; with ``LocalExecutor``, they run + inside the scheduler process. The Dag processor checks the stub tasks of each Python Dag against the task + handlers the packed bundles register, so it runs them too and needs bundles built for its operating + system and CPU architecture. The API server does not need any of it. Register the Dag bundle in + ``[dag_processor] dag_bundle_config_list`` on every component, like your other Dag bundles: the worker + and the Dag processor resolve ``task_handler_bundle_name`` through it, and wherever the ``[sdk]`` config + is read it is rejected if the name is missing there. Writing tasks ------------- @@ -466,9 +468,9 @@ coordinator scans (see `Deploying`_): Cross-platform builds ~~~~~~~~~~~~~~~~~~~~~~~ -The worker that runs a bundle often uses a different operating system or CPU architecture than your build -machine (for example, deploying to a Linux host from an Apple-silicon ``darwin/arm64`` laptop). Pass -``--goos`` / ``--goarch`` and the packer cross-builds for you: +The worker and the Dag processor that run a bundle often use a different operating system or CPU +architecture than your build machine (for example, deploying to a Linux host from an Apple-silicon +``darwin/arm64`` laptop). Pass ``--goos`` / ``--goarch`` and the packer cross-builds for you: .. code-block:: bash @@ -524,7 +526,9 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the - Description * - ``task_handler_bundle_name`` - *(task's own Dag bundle)* - - Name of the Dag bundle scanned recursively for executable bundles. It must be registered in + - Name of the Dag bundle scanned recursively for executable bundles. It is used only by + mixed-language Dags, to locate the task handlers for the ``@task.stub`` tasks of a Python Dag; + Dags defined natively in a language SDK do not use it. It must be registered in ``[dag_processor] dag_bundle_config_list``. It is checked when the ``[sdk]`` configuration is loaded, so a typo fails there rather than on the first task. * - ``task_startup_timeout`` @@ -534,8 +538,8 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the .. note:: - **Locating bundles.** Packed bundles are read from a Dag bundle, so they are delivered, refreshed and - versioned by the same machinery as your Dags. + **Locating bundles.** The packed bundles for the ``@task.stub`` tasks of a Python Dag are read from a Dag + bundle, so they are delivered, refreshed and versioned by the same machinery as your Dags. * The expected layout is a separate Dag bundle for the packed bundles, named by ``task_handler_bundle_name``, rather than the Dag bundle that holds your ``.py`` files. The task diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/index.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/index.rst index 79a0cb3df1f4d..159911d84c9f3 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/index.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/index.rst @@ -155,9 +155,9 @@ Coordinators are registered in ``airflow.cfg`` (or via environment variables) un } } - The ``classpath`` value must be importable by the worker. The ``kwargs`` are passed directly - to the coordinator's constructor. See the language-specific guide for the accepted kwargs - of each coordinator (e.g. :ref:`java-sdk/coordinator-config` for + The ``classpath`` value must be importable by the worker and the Dag processor. The ``kwargs`` + are passed directly to the coordinator's constructor. See the language-specific guide for the + accepted kwargs of each coordinator (e.g. :ref:`java-sdk/coordinator-config` for :class:`~airflow.sdk.coordinators.java.JavaCoordinator`). ``extra`` is an optional object for any additional information you want to associate with a 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 36c5f12243674..eafd542df6955 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst @@ -41,8 +41,8 @@ Prerequisites * JDK 11 or later is required on the machine that builds the Java project. A local Gradle installation is only needed to generate the Gradle Wrapper for a new project. -* JRE 11 or later must be available on the Airflow worker nodes. -* The compiled task JAR(s) and JVM dependencies must be accessible from the worker. +* JRE 11 or later must be available on the Airflow worker nodes and the Dag processor. +* The compiled task JAR(s) and JVM dependencies must be accessible from the worker and the Dag processor. * The ``apache-airflow-task-sdk`` package (installed with Airflow) provides the coordinator; no additional Python packages are needed. @@ -276,13 +276,15 @@ default, ``${AIRFLOW_HOME}/airflow.cfg``), or set the equivalent ``AIRFLOW__*`` Dag bundle name in ``dag_bundle_config_list``. See :ref:`java-sdk/coordinator-config` for how JARs are located. -Restart the affected Airflow components after changing this configuration. The coordinator config and JARs -must be available wherever tasks execute. With ``CeleryExecutor``, that means the Celery workers; with -``LocalExecutor``, tasks run in subprocesses on the scheduler's host. Register the Dag bundle in -``[dag_processor] dag_bundle_config_list`` on every component, like your other Dag bundles: the worker -resolves ``task_handler_bundle_name`` through it, and wherever the ``[sdk]`` config is read it is rejected if -the name is missing there. The API server and Dag processor do not need the JARs, while the Dag processor must -receive ``sales_pipeline.py`` through the separate Dag delivery process. +Restart the affected Airflow components after changing this configuration. The coordinator config, the JARs +and a JRE must be available wherever tasks execute and on the Dag processor. With ``CeleryExecutor``, tasks +execute on the Celery workers; with ``LocalExecutor``, they run in subprocesses on the scheduler's host. The +Dag processor checks the stub tasks of ``sales_pipeline.py`` against the task handlers the JARs register, so +it runs them too. The API server does not need any of it. Register the Dag bundle in +``[dag_processor] dag_bundle_config_list`` on every component, like your other Dag bundles: the worker and +the Dag processor resolve ``task_handler_bundle_name`` through it, and wherever the ``[sdk]`` config is read +it is rejected if the name is missing there. The Dag processor still receives ``sales_pipeline.py`` through +the separate Dag delivery process. After Airflow has parsed the Dag, trigger it from the UI or command line: @@ -1061,7 +1063,9 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the - Description * - ``task_handler_bundle_name`` - *(task's own Dag bundle)* - - Name of the Dag bundle scanned recursively for ``.jar`` files. It must be registered in + - Name of the Dag bundle scanned recursively for ``.jar`` files. It is used only by + mixed-language Dags, to locate the task handlers for the ``@task.stub`` tasks of a Python Dag; + Dags defined natively in a language SDK do not use it. It must be registered in ``[dag_processor] dag_bundle_config_list``. It is checked when the ``[sdk]`` configuration is loaded, so a typo fails there rather than on the first task. * - ``java_executable`` @@ -1082,8 +1086,8 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the .. note:: - **Locating JARs.** JARs are read from a Dag bundle, so they are delivered, refreshed and versioned - by the same machinery as your Dags. + **Locating JARs.** The JARs for the ``@task.stub`` tasks of a Python Dag are read from a Dag bundle, so + they are delivered, refreshed and versioned by the same machinery as your Dags. * The expected layout is a separate Dag bundle for the JARs, named by ``task_handler_bundle_name``, rather than the Dag bundle that holds your ``.py`` files. The task uses the version that Dag bundle @@ -1098,8 +1102,9 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the The ``[sdk]`` configuration is read at startup, so changes to ``coordinators`` or ``queue_to_coordinator`` (for example adding ``jvm_args``) only take effect after you restart the - scheduler (or ``airflow standalone``). A rebuilt bundle JAR, by contrast, is picked up on the next - task launch without a restart, because a fresh JVM is spawned per task instance. + components that read it: the workers (the scheduler with ``LocalExecutor``), the Dag processor, or + ``airflow standalone``. A rebuilt bundle JAR, by contrast, is picked up on the next task launch without a + restart, because a fresh JVM is spawned per task instance. .. _java-sdk/java-executable: 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 dfb7ed7d636e7..9364305ee2ddd 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst @@ -48,9 +48,9 @@ The SDK is the ``apache-airflow-ts-sdk`` package (ESM-only). It is currently in Prerequisites ------------- -* Node.js 22 or later must be available on the Airflow worker nodes. +* Node.js 22 or later must be available on the Airflow worker nodes and the Dag processor. * The packed bundle (a single ``bundle.min.mjs`` file, see :ref:`typescript-sdk/build`) must be accessible - from the worker, in the Dag bundle the coordinator scans. + from the worker and the Dag processor, in the Dag bundle the coordinator scans. * The ``apache-airflow-task-sdk`` package (installed with Airflow) provides the coordinator; no additional Python packages are needed. * In the TypeScript project, install the ``apache-airflow-ts-sdk`` npm package to author task handlers: @@ -243,13 +243,14 @@ task instance. .. note:: - The coordinator runs inside the Airflow worker, so the ``[sdk]`` config and the packed ``*.min.mjs`` - bundles only need to be present wherever tasks actually execute. With ``CeleryExecutor``, setting them on - the Celery workers is sufficient. With ``LocalExecutor``, tasks run inside the scheduler process, so they - must be present where the scheduler can read them. The API server and Dag processor do not need them. - Register the Dag bundle in ``[dag_processor] dag_bundle_config_list`` on every component, like your other - Dag bundles: the worker resolves ``task_handler_bundle_name`` through it, and wherever the ``[sdk]`` config - is read it is rejected if the name is missing there. + The ``[sdk]`` config, the packed ``*.min.mjs`` bundles and Node.js must be present wherever tasks execute + and on the Dag processor. With ``CeleryExecutor``, tasks execute on the Celery workers; with + ``LocalExecutor``, they run inside the scheduler process. The Dag processor checks the stub tasks of each + Python Dag against the task handlers the packed bundles register, so it runs them too. The API server + does not need any of it. Register the Dag bundle in ``[dag_processor] dag_bundle_config_list`` on every + component, like your other Dag bundles: the worker and the Dag processor resolve + ``task_handler_bundle_name`` through it, and wherever the ``[sdk]`` config is read it is rejected if the + name is missing there. .. _typescript-sdk/native-dag: @@ -544,8 +545,10 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the * - ``task_handler_bundle_name`` - *(task's own Dag bundle)* - Name of the Dag bundle searched recursively for an integrity-verified ``*.min.mjs`` bundle that - declares the requested Dag. It must be registered in ``[dag_processor] dag_bundle_config_list``. It is - checked when the ``[sdk]`` configuration is loaded, so a typo fails there rather than on the first task. + declares the requested Dag. It is used only by mixed-language Dags, to locate the task handlers for + the ``@task.stub`` tasks of a Python Dag; Dags defined natively in a language SDK do not use it. It + must be registered in ``[dag_processor] dag_bundle_config_list``. It is checked when the ``[sdk]`` + configuration is loaded, so a typo fails there rather than on the first task. * - ``node_executable`` - ``"node"`` - Path to the ``node`` binary. Defaults to ``node`` on ``$PATH``. @@ -556,8 +559,8 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the .. note:: - **Locating the bundle.** Packed bundles are read from a Dag bundle, so they are delivered, refreshed and - versioned by the same machinery as your Dags. + **Locating the bundle.** The packed bundles for the ``@task.stub`` tasks of a Python Dag are read from a + Dag bundle, so they are delivered, refreshed and versioned by the same machinery as your Dags. * The expected layout is a separate Dag bundle for the packed bundles, named by ``task_handler_bundle_name``, rather than the Dag bundle that holds your ``.py`` files. The task uses diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index 0d93997abe387..eafaecdda739b 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -2191,12 +2191,19 @@ sdk: example, two ``JavaCoordinator`` instances pinned to different JDK versions). - The language SDK coordinators read their compiled task handlers from the - Dag bundle named by the ``task_handler_bundle_name`` kwarg. It must name a - bundle in ``[dag_processor] dag_bundle_config_list`` and is checked when - this option is loaded. When it is unset, a task reads its handlers from its + The language SDK coordinators read the task handlers for the + ``@task.stub`` tasks of a Python Dag from the Dag bundle named by the + ``task_handler_bundle_name`` kwarg. Dags defined natively in a language + SDK do not use it. It must name a bundle in + ``[dag_processor] dag_bundle_config_list`` and is checked when this + option is loaded. When it is unset, a task reads its handlers from its own Dag bundle. + The Dag processor needs this option too, with the language runtime the + coordinator starts and the files of that Dag bundle, because it checks + the stub tasks of each Python Dag against the task handlers those files + register. The API server does not need it. + An entry may also carry an optional ``extra`` object for additional information associated with the coordinator that the coordinator itself does not receive; other components read it as needed. For example, diff --git a/go-sdk/README.md b/go-sdk/README.md index da477119b0c38..400726c1acb47 100644 --- a/go-sdk/README.md +++ b/go-sdk/README.md @@ -296,18 +296,21 @@ the full range of task states, and alternate XCom backends without implementing ``` `task_handler_bundle_name` names the Dag bundle the coordinator scans for packed bundles (the task's own - Dag bundle when unset); `queue_to_coordinator` routes stub tasks with `queue="golang"` to this Go - coordinator. Only files with the executable bit are considered, so use a Dag bundle that keeps it: a - `LocalDagBundle` does, object-store Dag bundles such as `S3DagBundle` do not. + Dag bundle when unset). It is used only by mixed-language Dags, to locate the task handlers for the + `@task.stub` tasks of a Python Dag; Dags defined natively in a language SDK do not use it. + `queue_to_coordinator` routes stub tasks with `queue="golang"` to this Go coordinator. Only files with the + executable bit are considered, so use a Dag bundle that keeps it: a `LocalDagBundle` does, object-store + Dag bundles such as `S3DagBundle` do not. > [!IMPORTANT] - > The coordinator is part of the Airflow worker, so the `[sdk]` config and the packed bundle files only - > need to be present wherever tasks actually execute. With `CeleryExecutor`, setting it on the Celery - > workers is sufficient. With `LocalExecutor`, tasks run inside the scheduler process, so it must be set - > where the scheduler can read it. The API server and Dag processor do not need it. Register the Dag - > bundle in `[dag_processor] dag_bundle_config_list` on every component, like your other Dag bundles: - > the worker resolves `task_handler_bundle_name` through it, and wherever the `[sdk]` config is read it - > is rejected if the name is missing there. + > The `[sdk]` config and the packed bundle files must be present wherever tasks execute and on the Dag + > processor. With `CeleryExecutor`, tasks execute on the Celery workers; with `LocalExecutor`, they run + > inside the scheduler process. The Dag processor checks the stub tasks of each Python Dag against the + > task handlers the packed bundles register, so it runs them too and needs bundles built for its + > operating system and CPU architecture. The API server does not need any of it. Register the Dag bundle + > in `[dag_processor] dag_bundle_config_list` on every component, like your other Dag bundles: the worker + > and the Dag processor resolve `task_handler_bundle_name` through it, and wherever the `[sdk]` config is + > read it is rejected if the name is missing there. - Deploy the matching Python stub Dag (above) into Airflow. There is no separate Go worker to run: the Airflow worker forks the bundle binary once per task instance. diff --git a/java-sdk/README.md b/java-sdk/README.md index 5ec0995708546..752253b6ab826 100644 --- a/java-sdk/README.md +++ b/java-sdk/README.md @@ -93,8 +93,8 @@ Now `cd example` into the example project, and * Put the [DAG with stub tasks](./example/src/resources/dags) to somewhere Airflow can find. -* Ensure the `java` command is available in the same environment the Airflow - task worker is in. +* Ensure the `java` command is available in the same environments the Airflow + task worker and the Dag processor are in. * Register the packaged example as a Dag bundle, and configure Airflow to route tasks in the *java* queue to be run with Java from it: @@ -117,6 +117,13 @@ Now `cd example` into the example project, and export AIRFLOW__SDK__QUEUE_TO_COORDINATOR='{"java": "java"}' ``` + Set these, and make the bundle directory available, for the Dag processor + as well as the worker: the Dag processor checks the stub tasks against the + task handlers the JARs register. The API server does not need them. + `task_handler_bundle_name` is used only by mixed-language Dags, to locate + the task handlers for the `@task.stub` tasks of a Python Dag; Dags defined + natively in a language SDK do not use it. + * Ensure the Connection and Variable needed by the example DAG are available: ```bash diff --git a/ts-sdk/README.md b/ts-sdk/README.md index 75c5859b5ca8f..b105f768c6a41 100644 --- a/ts-sdk/README.md +++ b/ts-sdk/README.md @@ -110,7 +110,13 @@ queue_to_coordinator = {"typescript": "ts"} The Dag bundle named by `task_handler_bundle_name` (the task's own Dag bundle when unset) is searched recursively for `*.min.mjs` bundles built with `airflow-ts-pack` (see [Packing bundles](#packing-bundles)), -which embeds the Airflow metadata in the bundle itself. +which embeds the Airflow metadata in the bundle itself. `task_handler_bundle_name` is used only by +mixed-language Dags, to locate the task handlers for the `@task.stub` tasks of a Python Dag; Dags defined +natively in a language SDK do not use it. + +This config, the packed bundles and Node.js 22 or later must be present wherever tasks execute and on the +Dag processor, which checks the stub tasks of each Python Dag against the task handlers the bundles +register. The API server does not need them. TypeScript entrypoint: From c7055408b06e76e24f79c52effba6c13ed859a09 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Fri, 2 Oct 2026 11:27:27 +0000 Subject: [PATCH 08/94] Validate task_handler_bundle_name only for routed coordinators A coordinator no queue routes to is never used, so its kwargs are not checked. Add from_config tests for an accepted name, an unrouted coordinator and a spec without the kwarg. --- .../airflow/sdk/execution_time/coordinator.py | 17 +++--- .../execution_time/test_coordinator.py | 56 ++++++++++++++++++- 2 files changed, 63 insertions(+), 10 deletions(-) diff --git a/task-sdk/src/airflow/sdk/execution_time/coordinator.py b/task-sdk/src/airflow/sdk/execution_time/coordinator.py index 8aad5c6899dbc..778ea29aea899 100644 --- a/task-sdk/src/airflow/sdk/execution_time/coordinator.py +++ b/task-sdk/src/airflow/sdk/execution_time/coordinator.py @@ -257,23 +257,22 @@ def from_config(cls) -> Self: """ Load coordinator specs from configuration without initialization. - Every ``queue_to_coordinator`` key and every ``task_handler_bundle_name`` - kwarg is validated here, so a typo fails at config load rather than on the - first task routed to the coordinator. + Every ``queue_to_coordinator`` key and the ``task_handler_bundle_name`` of + every routed coordinator are validated here, so a typo fails at config load + rather than on the first task routed to the coordinator. """ coordinator_specs = { k: _CoordinatorSpec.model_validate(v) for k, v in conf.getjson("sdk", "coordinators", fallback={}).items() } queue_to_coordinator = conf.getjson("sdk", "queue_to_coordinator", fallback={}) - for key in queue_to_coordinator.values(): - if key not in coordinator_specs: + for key in set(queue_to_coordinator.values()): + if (spec := coordinator_specs.get(key)) is None: raise ValueError(f"[sdk] queue_to_coordinator references invalid coordinator key: {key!r}") - for key, spec in coordinator_specs.items(): bundle_name = spec.kwargs.get("task_handler_bundle_name") - if bundle_name is None: - continue - if not isinstance(bundle_name, str) or not DagBundlesManager.is_bundle_configured(bundle_name): + if bundle_name is not None and ( + not isinstance(bundle_name, str) or not DagBundlesManager.is_bundle_configured(bundle_name) + ): raise InvalidCoordinatorError( f"[sdk] coordinators {key!r} sets task_handler_bundle_name={bundle_name!r}, " f"which is not a bundle in [dag_processor] dag_bundle_config_list" diff --git a/task-sdk/tests/task_sdk/execution_time/test_coordinator.py b/task-sdk/tests/task_sdk/execution_time/test_coordinator.py index 3659bce784dbf..ccd6a22280e8d 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_coordinator.py +++ b/task-sdk/tests/task_sdk/execution_time/test_coordinator.py @@ -138,7 +138,7 @@ def test_from_config_rejects_invalid_queue_mapping(self, sdk_config): @pytest.mark.parametrize("bundle_name", ["ghost", ["ghost"]], ids=["unconfigured", "not-a-string"]) def test_from_config_rejects_invalid_task_handler_bundle_name(self, sdk_config, bundle_name): - """Validated for every spec, routed or not, without importing or constructing the coordinator.""" + """Validated without importing or constructing the coordinator.""" sdk_config( coordinators=json.dumps( { @@ -148,10 +148,64 @@ def test_from_config_rejects_invalid_task_handler_bundle_name(self, sdk_config, } } ), + queue_to_coordinator=json.dumps({"queue-boom": "boom"}), ) with pytest.raises(InvalidCoordinatorError, match=r"'boom' sets task_handler_bundle_name="): CoordinatorManager.from_config() + def test_from_config_accepts_configured_task_handler_bundle_name(self, sdk_config): + sdk_config( + coordinators=json.dumps( + { + "boom": { + "classpath": f"{_ExplodingCoordinator.__module__}._ExplodingCoordinator", + "kwargs": {"task_handler_bundle_name": "handlers"}, + } + } + ), + queue_to_coordinator=json.dumps({"queue-boom": "boom"}), + ) + bundles = [ + { + "name": "handlers", + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {}, + } + ] + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(bundles)}): + manager = CoordinatorManager.from_config() + assert manager._queue_to_coordinator == {"queue-boom": "boom"} + assert manager._created_coordinators == {} + + def test_from_config_ignores_task_handler_bundle_name_of_unrouted_coordinator(self, sdk_config): + sdk_config( + coordinators=json.dumps( + { + "boom": { + "classpath": f"{_ExplodingCoordinator.__module__}._ExplodingCoordinator", + "kwargs": {"task_handler_bundle_name": "ghost"}, + } + } + ), + ) + manager = CoordinatorManager.from_config() + assert set(manager._coordinator_specs) == {"boom"} + + @mock.patch( + "airflow.sdk.execution_time.coordinator.DagBundlesManager.is_bundle_configured", autospec=True + ) + def test_from_config_skips_bundle_lookup_without_task_handler_bundle_name( + self, is_bundle_configured, sdk_config + ): + sdk_config( + coordinators=json.dumps( + {"alpha": {"classpath": f"{_CoordinatorA.__module__}._CoordinatorA", "kwargs": {}}} + ), + queue_to_coordinator=json.dumps({"queue-a": "alpha"}), + ) + CoordinatorManager.from_config() + is_bundle_configured.assert_not_called() + @pytest.mark.parametrize( ("coordinator_spec", "expected_match"), [ From 9d00ead095fa6ab797d334eacbb84d02c2bd7b0e Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Fri, 2 Oct 2026 11:27:39 +0000 Subject: [PATCH 09/94] Add newsfragment for the removed coordinator artifact roots --- airflow-core/newsfragments/73970.significant.rst | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 airflow-core/newsfragments/73970.significant.rst diff --git a/airflow-core/newsfragments/73970.significant.rst b/airflow-core/newsfragments/73970.significant.rst new file mode 100644 index 0000000000000..886d1d6ca0a8f --- /dev/null +++ b/airflow-core/newsfragments/73970.significant.rst @@ -0,0 +1,5 @@ +Language SDK coordinators load task handlers from a Dag bundle + +The experimental ``jars_root`` (Java) and ``executables_root`` (Go) coordinator kwargs are removed. +Put the artifacts in a Dag bundle registered in ``[dag_processor] dag_bundle_config_list`` and set ``task_handler_bundle_name`` in the coordinator's ``kwargs``, or leave it unset to use the task's own Dag bundle. +A config that still sets a removed kwarg fails with ``Cannot instantiate coordinator ''``. From 68f03d615b81365140d804b742a9ffaaa6b6c825 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Fri, 2 Oct 2026 11:27:44 +0000 Subject: [PATCH 10/94] Mark ADR-0013 as accepted --- .../adr/lang-sdk/0013-persisted-task-handler-bindings.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md b/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md index cbcace7df5231..016f804129ab9 100644 --- a/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md +++ b/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md @@ -21,7 +21,7 @@ ## Status -Proposed +Accepted ## Context From c174aa41637140d3b9d233af38a9967ca6810f75 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 11:42:36 +0000 Subject: [PATCH 11/94] Add lang-SDK task handler binding tables ADR-0013 records at parse time which Language SDK artifact runs each stub task, so task execution can stop scanning artifact bundles. This adds the two tables and their migration. Later layers write and read them. The unique key and the owning-file index use an md5 hash of each path instead of the path. A 2000-character path is over MySQL's 3072-byte key limit and Postgres's 2704-byte btree entry limit. This follows DagPriorityParsingRequest. --- airflow-core/docs/migrations-ref.rst | 5 +- ..._3_4_0_add_lang_sdk_task_handler_tables.py | 91 +++++++++++ airflow-core/src/airflow/models/__init__.py | 1 + .../airflow/models/lang_sdk_task_handler.py | 116 ++++++++++++++ airflow-core/src/airflow/utils/db.py | 2 +- .../unit/models/test_lang_sdk_task_handler.py | 147 ++++++++++++++++++ .../tests/unit/utils/test_db_cleanup.py | 2 + 7 files changed, 362 insertions(+), 2 deletions(-) create mode 100644 airflow-core/src/airflow/migrations/versions/0142_3_4_0_add_lang_sdk_task_handler_tables.py create mode 100644 airflow-core/src/airflow/models/lang_sdk_task_handler.py create mode 100644 airflow-core/tests/unit/models/test_lang_sdk_task_handler.py diff --git a/airflow-core/docs/migrations-ref.rst b/airflow-core/docs/migrations-ref.rst index afd3ed01e8c32..aa57aaf1472d2 100644 --- a/airflow-core/docs/migrations-ref.rst +++ b/airflow-core/docs/migrations-ref.rst @@ -39,7 +39,10 @@ Here's the list of all the Database Migrations that are executed via when you ru +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | Revision ID | Revises ID | Airflow Version | Description | +=========================+==================+===================+==============================================================+ -| ``90e4d18ccadf`` (head) | ``e5a91c7f42b3`` | ``3.4.0`` | Add timetable_asset_gated to DagModel. | +| ``ffff41014d89`` (head) | ``90e4d18ccadf`` | ``3.4.0`` | Add lang_sdk_task_handler_artifact and lang_sdk_task_handler | +| | | | tables. | ++-------------------------+------------------+-------------------+--------------------------------------------------------------+ +| ``90e4d18ccadf`` | ``e5a91c7f42b3`` | ``3.4.0`` | Add timetable_asset_gated to DagModel. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | ``e5a91c7f42b3`` | ``ca8499dc1004`` | ``3.4.0`` | Add language column to dag_code. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ diff --git a/airflow-core/src/airflow/migrations/versions/0142_3_4_0_add_lang_sdk_task_handler_tables.py b/airflow-core/src/airflow/migrations/versions/0142_3_4_0_add_lang_sdk_task_handler_tables.py new file mode 100644 index 0000000000000..380e5f1dc2a71 --- /dev/null +++ b/airflow-core/src/airflow/migrations/versions/0142_3_4_0_add_lang_sdk_task_handler_tables.py @@ -0,0 +1,91 @@ +# +# 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. + +""" +Add lang_sdk_task_handler_artifact and lang_sdk_task_handler tables. + +Revision ID: ffff41014d89 +Revises: 90e4d18ccadf +Create Date: 2026-09-30 12:00:00.000000 + +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +from airflow.migrations.db_types import StringID +from airflow.utils.sqlalchemy import UtcDateTime + +# revision identifiers, used by Alembic. +revision = "ffff41014d89" +down_revision = "90e4d18ccadf" +branch_labels = None +depends_on = None +airflow_version = "3.4.0" + + +def upgrade(): + """Add lang_sdk_task_handler_artifact and lang_sdk_task_handler tables.""" + op.create_table( + "lang_sdk_task_handler_artifact", + sa.Column("id", sa.Uuid(), nullable=False), + sa.Column("bundle_name", StringID(), nullable=False), + sa.Column("relative_fileloc", sa.String(length=2000), nullable=False), + sa.Column("relative_fileloc_hash", sa.String(length=32), nullable=False), + sa.Column("size_bytes", sa.BigInteger(), nullable=False), + sa.Column("cache_digest", sa.String(length=128), nullable=False), + sa.Column("last_probed_at", UtcDateTime(), nullable=False), + sa.PrimaryKeyConstraint("id", name=op.f("lang_sdk_task_handler_artifact_pkey")), + sa.UniqueConstraint( + "bundle_name", + "relative_fileloc_hash", + name=op.f("lang_sdk_task_handler_artifact_bundle_fileloc_uq"), + ), + ) + op.create_table( + "lang_sdk_task_handler", + sa.Column("dag_id", StringID(), nullable=False), + sa.Column("task_id", StringID(), nullable=False), + sa.Column("artifact_id", sa.Uuid(), nullable=False), + sa.Column("dag_bundle_name", StringID(), nullable=False), + sa.Column("dag_relative_fileloc", sa.String(length=2000), nullable=False), + sa.Column("dag_relative_fileloc_hash", sa.String(length=32), nullable=False), + sa.Column("handler_params", sa.JSON(), nullable=False), + sa.PrimaryKeyConstraint("dag_id", "task_id", name=op.f("lang_sdk_task_handler_pkey")), + sa.ForeignKeyConstraint( + columns=("dag_id",), + refcolumns=["dag.dag_id"], + name="lang_sdk_task_handler_dag_id_fkey", + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + columns=("artifact_id",), + refcolumns=["lang_sdk_task_handler_artifact.id"], + name="lang_sdk_task_handler_artifact_id_fkey", + ), + sa.Index("idx_lang_sdk_task_handler_dag_file", "dag_bundle_name", "dag_relative_fileloc_hash"), + sa.Index("idx_lang_sdk_task_handler_artifact_id", "artifact_id"), + ) + + +def downgrade(): + """Drop lang_sdk_task_handler and lang_sdk_task_handler_artifact tables.""" + op.drop_table("lang_sdk_task_handler") + op.drop_table("lang_sdk_task_handler_artifact") diff --git a/airflow-core/src/airflow/models/__init__.py b/airflow-core/src/airflow/models/__init__.py index f5bffb4e57173..f806537c684df 100644 --- a/airflow-core/src/airflow/models/__init__.py +++ b/airflow-core/src/airflow/models/__init__.py @@ -71,6 +71,7 @@ def import_all_models(): import airflow.models.dagwarning import airflow.models.deadline_alert import airflow.models.errors + import airflow.models.lang_sdk_task_handler import airflow.models.revoked_token import airflow.models.serialized_dag import airflow.models.task_state_store diff --git a/airflow-core/src/airflow/models/lang_sdk_task_handler.py b/airflow-core/src/airflow/models/lang_sdk_task_handler.py new file mode 100644 index 0000000000000..5fda05ade2376 --- /dev/null +++ b/airflow-core/src/airflow/models/lang_sdk_task_handler.py @@ -0,0 +1,116 @@ +# +# 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 collections.abc import Callable +from datetime import datetime +from typing import TYPE_CHECKING, Any +from uuid import UUID + +import sqlalchemy as sa +import uuid6 +from sqlalchemy import BigInteger, ForeignKeyConstraint, Index, String, UniqueConstraint, Uuid +from sqlalchemy.orm import Mapped, mapped_column, validates + +from airflow._shared.timezones import timezone +from airflow.models.base import Base, StringID +from airflow.utils.hashlib_wrapper import md5 +from airflow.utils.sqlalchemy import UtcDateTime + +if TYPE_CHECKING: + from sqlalchemy.engine.default import DefaultExecutionContext + + +def compute_fileloc_hash(relative_fileloc: str) -> str: + """ + Return the md5 hex digest that the task handler tables index in place of a relative file location. + + A 2000-character path exceeds MySQL's 3072-byte key limit and Postgres's 2704-byte btree entry limit. + """ + return md5(relative_fileloc.encode()).hexdigest() + + +def _build_fileloc_hash_default(path_key: str) -> Callable[[DefaultExecutionContext], str]: + def compute(context: DefaultExecutionContext) -> str: + return compute_fileloc_hash(context.get_current_parameters()[path_key]) + + return compute + + +class LangSDKTaskHandlerArtifact(Base): + """A Language SDK artifact in a task handler Dag bundle, cached with its fingerprint.""" + + __tablename__ = "lang_sdk_task_handler_artifact" + + id: Mapped[UUID] = mapped_column(Uuid(), primary_key=True, default=uuid6.uuid7) + bundle_name: Mapped[str] = mapped_column(StringID(), nullable=False) + relative_fileloc: Mapped[str] = mapped_column(String(2000), nullable=False) + relative_fileloc_hash: Mapped[str] = mapped_column( + String(32), nullable=False, default=_build_fileloc_hash_default("relative_fileloc") + ) + size_bytes: Mapped[int] = mapped_column(BigInteger, nullable=False) + cache_digest: Mapped[str] = mapped_column(String(128), nullable=False) + last_probed_at: Mapped[datetime] = mapped_column(UtcDateTime, nullable=False, default=timezone.utcnow) + + __table_args__ = ( + UniqueConstraint( + bundle_name, + relative_fileloc_hash, + name="lang_sdk_task_handler_artifact_bundle_fileloc_uq", + ), + ) + + +class LangSDKTaskHandler(Base): + """The artifact that runs one stub task, and the task handler parameters it declares.""" + + __tablename__ = "lang_sdk_task_handler" + + dag_id: Mapped[str] = mapped_column(StringID(), primary_key=True) + task_id: Mapped[str] = mapped_column(StringID(), primary_key=True) + artifact_id: Mapped[UUID] = mapped_column(Uuid(), nullable=False) + dag_bundle_name: Mapped[str] = mapped_column(StringID(), nullable=False) + dag_relative_fileloc: Mapped[str] = mapped_column(String(2000), nullable=False) + dag_relative_fileloc_hash: Mapped[str] = mapped_column( + String(32), nullable=False, default=_build_fileloc_hash_default("dag_relative_fileloc") + ) + handler_params: Mapped[list[dict[str, Any]]] = mapped_column(sa.JSON(), nullable=False) + + __table_args__ = ( + ForeignKeyConstraint( + (dag_id,), + ["dag.dag_id"], + name="lang_sdk_task_handler_dag_id_fkey", + ondelete="CASCADE", + ), + # No ON DELETE: deleting an artifact that a handler still references must fail. + ForeignKeyConstraint( + (artifact_id,), + ["lang_sdk_task_handler_artifact.id"], + name="lang_sdk_task_handler_artifact_id_fkey", + ), + Index("idx_lang_sdk_task_handler_dag_file", dag_bundle_name, dag_relative_fileloc_hash), + Index("idx_lang_sdk_task_handler_artifact_id", artifact_id), + ) + + @validates("dag_relative_fileloc") + def _sync_dag_relative_fileloc_hash(self, key: str, dag_relative_fileloc: str) -> str: + # Not an onupdate default: that fires on every UPDATE, even one that leaves the path alone. + # A Core UPDATE or upsert that changes the path must set the hash itself. + self.dag_relative_fileloc_hash = compute_fileloc_hash(dag_relative_fileloc) + return dag_relative_fileloc diff --git a/airflow-core/src/airflow/utils/db.py b/airflow-core/src/airflow/utils/db.py index f160cb0597d84..56adf7b831731 100644 --- a/airflow-core/src/airflow/utils/db.py +++ b/airflow-core/src/airflow/utils/db.py @@ -117,7 +117,7 @@ class MappedClassProtocol(Protocol): "3.1.8": "509b94a1042d", "3.2.0": "1d6611b6ab7c", "3.3.0": "d2f4e1b3c5a7", - "3.4.0": "90e4d18ccadf", + "3.4.0": "ffff41014d89", } # Prefix used to identify tables holding data moved during migration. diff --git a/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py b/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py new file mode 100644 index 0000000000000..52d5d1f033843 --- /dev/null +++ b/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py @@ -0,0 +1,147 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import pytest +import uuid6 +from sqlalchemy import delete, insert, select +from sqlalchemy.exc import IntegrityError + +from airflow.models.dag import DagModel +from airflow.models.lang_sdk_task_handler import ( + LangSDKTaskHandler, + LangSDKTaskHandlerArtifact, + compute_fileloc_hash, +) + +pytestmark = pytest.mark.db_test + +ARTIFACT_BUNDLE = "java-task-handlers" +# 6000 bytes in UTF-8, over both the MySQL key limit and the Postgres btree entry limit. +LONG_NON_ASCII_PATH = "任" * 2000 + + +def _make_artifact( + *, bundle_name: str = ARTIFACT_BUNDLE, relative_fileloc: str = "etl.jar" +) -> LangSDKTaskHandlerArtifact: + return LangSDKTaskHandlerArtifact( + bundle_name=bundle_name, relative_fileloc=relative_fileloc, size_bytes=1024, cache_digest="0" * 64 + ) + + +def _make_handler(*, dag_id: str, artifact: LangSDKTaskHandlerArtifact) -> LangSDKTaskHandler: + return LangSDKTaskHandler( + dag_id=dag_id, + task_id="extract", + artifact_id=artifact.id, + dag_bundle_name="testing", + dag_relative_fileloc=f"{dag_id}.py", + handler_params=[{"name": "path", "value_schema": None, "required": True}], + ) + + +def _add_dags_and_artifact(session, *dag_ids: str) -> LangSDKTaskHandlerArtifact: + session.add_all(DagModel(dag_id=dag_id, bundle_name="testing") for dag_id in dag_ids) + artifact = _make_artifact() + session.add(artifact) + session.flush() + return artifact + + +def test_deleting_dag_deletes_its_handlers(testing_dag_bundle, session): + artifact = _add_dags_and_artifact(session, "dag_a", "dag_b") + session.add_all(_make_handler(dag_id=dag_id, artifact=artifact) for dag_id in ("dag_a", "dag_b")) + session.flush() + + session.execute(delete(DagModel).where(DagModel.dag_id == "dag_a")) + + assert session.scalars(select(LangSDKTaskHandler.dag_id)).all() == ["dag_b"] + assert session.scalars(select(LangSDKTaskHandlerArtifact.id)).all() == [artifact.id] + + +@pytest.mark.parametrize( + "relative_fileloc", + [ + pytest.param("etl.jar", id="short"), + pytest.param(LONG_NON_ASCII_PATH, id="long-non-ascii"), + ], +) +def test_artifact_path_is_unique_per_bundle(session, relative_fileloc): + session.add_all( + [ + _make_artifact(relative_fileloc=relative_fileloc), + _make_artifact(bundle_name="go-task-handlers", relative_fileloc=relative_fileloc), + ] + ) + session.flush() + + session.add(_make_artifact(relative_fileloc=relative_fileloc)) + with pytest.raises(IntegrityError): + session.flush() + + +def test_core_insert_fills_fileloc_hashes(testing_dag_bundle, session): + session.add(DagModel(dag_id="dag_a", bundle_name="testing")) + session.flush() + artifact_id = uuid6.uuid7() + + session.execute( + insert(LangSDKTaskHandlerArtifact).values( + id=artifact_id, + bundle_name=ARTIFACT_BUNDLE, + relative_fileloc="etl.jar", + size_bytes=1024, + cache_digest="0" * 64, + ) + ) + session.execute( + insert(LangSDKTaskHandler).values( + dag_id="dag_a", + task_id="extract", + artifact_id=artifact_id, + dag_bundle_name="testing", + dag_relative_fileloc="dags/etl.py", + handler_params=[], + ) + ) + + assert session.scalar(select(LangSDKTaskHandlerArtifact.relative_fileloc_hash)) == compute_fileloc_hash( + "etl.jar" + ) + assert session.scalar(select(LangSDKTaskHandler.dag_relative_fileloc_hash)) == compute_fileloc_hash( + "dags/etl.py" + ) + + +def test_dag_relative_fileloc_hash_follows_orm_updates(testing_dag_bundle, session): + artifact = _add_dags_and_artifact(session, "dag_a") + handler = _make_handler(dag_id="dag_a", artifact=artifact) + session.add(handler) + session.flush() + + handler.handler_params = [] + session.flush() + assert session.scalar(select(LangSDKTaskHandler.dag_relative_fileloc_hash)) == compute_fileloc_hash( + "dag_a.py" + ) + + handler.dag_relative_fileloc = "moved/dag_a.py" + session.flush() + assert session.scalar(select(LangSDKTaskHandler.dag_relative_fileloc_hash)) == compute_fileloc_hash( + "moved/dag_a.py" + ) diff --git a/airflow-core/tests/unit/utils/test_db_cleanup.py b/airflow-core/tests/unit/utils/test_db_cleanup.py index 375461b6848ac..41ec3ff7de9a7 100644 --- a/airflow-core/tests/unit/utils/test_db_cleanup.py +++ b/airflow-core/tests/unit/utils/test_db_cleanup.py @@ -1087,6 +1087,7 @@ def test_no_models_missing(self): # leave alone - per-asset key/value state, upserted in place (PK is asset_id+key), # so it is bounded and current, not accumulating history; removed with its asset "asset_state_store", + "lang_sdk_task_handler_artifact", # parse-time cache of task handler artifacts, not run data # Purged indirectly: each of these hangs off a cleaned table by an # ON DELETE CASCADE foreign key, so the rows go when the parent does. # cascade from dag_run once the partition run has fired; while it is still @@ -1098,6 +1099,7 @@ def test_no_models_missing(self): "hitl_detail", # cascade from task_instance "hitl_detail_history", # cascade from task_instance_history "job_team", # cascade from job + "lang_sdk_task_handler", # cascade from dag "task_inlet_asset_reference", # cascade from dag } From 6762729ecfcad603a699edf55fa3b8a7ce7f2fee Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 11:45:47 +0000 Subject: [PATCH 12/94] Match ADR-0013 task handler table SQL to the migration The tables as built need an md5 hash column beside each long path, a wider cache_digest, and the repo's constraint names. Later layers read this SQL as the spec, so it now matches the migration. --- .../0013-persisted-task-handler-bindings.md | 41 ++++++++++--------- 1 file changed, 22 insertions(+), 19 deletions(-) diff --git a/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md b/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md index 016f804129ab9..d4d242a2152e9 100644 --- a/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md +++ b/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md @@ -157,31 +157,34 @@ backs many handlers and its fingerprint must have exactly one value. ```sql CREATE TABLE lang_sdk_task_handler_artifact ( - id UUID NOT NULL, - bundle_name VARCHAR(250) NOT NULL, -- the task_handler_bundle_name it was found in - relative_fileloc VARCHAR(2000) NOT NULL, -- path within that bundle - size_bytes BIGINT NOT NULL, -- cheap fingerprint tier - cache_digest VARCHAR(64) NOT NULL, -- content fingerprint tier; see "The fast path" - last_probed_at TIMESTAMP NOT NULL, - PRIMARY KEY (id), - CONSTRAINT lstha_bundle_fileloc_uq UNIQUE (bundle_name, relative_fileloc) + id UUID NOT NULL, + bundle_name VARCHAR(250) NOT NULL, -- the task_handler_bundle_name it was found in + relative_fileloc VARCHAR(2000) NOT NULL, -- path within that bundle + relative_fileloc_hash VARCHAR(32) NOT NULL, -- md5 of relative_fileloc; the path is too long to index + size_bytes BIGINT NOT NULL, -- cheap fingerprint tier + cache_digest VARCHAR(128) NOT NULL, -- content fingerprint tier; see "The fast path" + last_probed_at TIMESTAMP NOT NULL, + CONSTRAINT lang_sdk_task_handler_artifact_pkey PRIMARY KEY (id), + CONSTRAINT lang_sdk_task_handler_artifact_bundle_fileloc_uq UNIQUE (bundle_name, relative_fileloc_hash) ); CREATE TABLE lang_sdk_task_handler ( - dag_id VARCHAR(250) NOT NULL, - task_id VARCHAR(250) NOT NULL, - artifact_id UUID NOT NULL, - dag_bundle_name VARCHAR(250) NOT NULL, -- the *Python* file that owns this row - dag_relative_fileloc VARCHAR(2000) NOT NULL, -- ditto - handler_params JSON NOT NULL, -- list[TaskHandlerParam], ordered - PRIMARY KEY (dag_id, task_id), - CONSTRAINT lsth_dag_fkey FOREIGN KEY (dag_id) + dag_id VARCHAR(250) NOT NULL, + task_id VARCHAR(250) NOT NULL, + artifact_id UUID NOT NULL, + dag_bundle_name VARCHAR(250) NOT NULL, -- the *Python* file that owns this row + dag_relative_fileloc VARCHAR(2000) NOT NULL, -- ditto + dag_relative_fileloc_hash VARCHAR(32) NOT NULL, -- md5 of dag_relative_fileloc + handler_params JSON NOT NULL, -- list[TaskHandlerParam], ordered + CONSTRAINT lang_sdk_task_handler_pkey PRIMARY KEY (dag_id, task_id), + CONSTRAINT lang_sdk_task_handler_dag_id_fkey FOREIGN KEY (dag_id) REFERENCES dag (dag_id) ON DELETE CASCADE, - CONSTRAINT lsth_artifact_fkey FOREIGN KEY (artifact_id) + CONSTRAINT lang_sdk_task_handler_artifact_id_fkey FOREIGN KEY (artifact_id) REFERENCES lang_sdk_task_handler_artifact (id) ); -CREATE INDEX idx_lsth_dag_file ON lang_sdk_task_handler (dag_bundle_name, dag_relative_fileloc); -CREATE INDEX idx_lsth_artifact_id ON lang_sdk_task_handler (artifact_id); +CREATE INDEX idx_lang_sdk_task_handler_dag_file + ON lang_sdk_task_handler (dag_bundle_name, dag_relative_fileloc_hash); +CREATE INDEX idx_lang_sdk_task_handler_artifact_id ON lang_sdk_task_handler (artifact_id); ``` `PRIMARY KEY (dag_id, task_id)` is the conflict guard: two artifacts claiming the same task cannot From dc15082f1350d02345d14516edb7083a00a09978 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 14:16:32 +0000 Subject: [PATCH 13/94] Keep the artifact path hash in sync on ORM updates An ORM update of an artifact's relative_fileloc kept the old hash, so the unique key no longer matched the path. The handler table already resyncs its hash this way. --- .../airflow/models/lang_sdk_task_handler.py | 10 +++- .../unit/models/test_lang_sdk_task_handler.py | 53 +++++++++++++++---- 2 files changed, 50 insertions(+), 13 deletions(-) diff --git a/airflow-core/src/airflow/models/lang_sdk_task_handler.py b/airflow-core/src/airflow/models/lang_sdk_task_handler.py index 5fda05ade2376..d53301fe05b4a 100644 --- a/airflow-core/src/airflow/models/lang_sdk_task_handler.py +++ b/airflow-core/src/airflow/models/lang_sdk_task_handler.py @@ -45,6 +45,9 @@ def compute_fileloc_hash(relative_fileloc: str) -> str: return md5(relative_fileloc.encode()).hexdigest() +# The column default fills a hash on INSERT, ORM or Core. ORM updates resync it in each model's +# @validates hook instead of an onupdate default, which would fire on every UPDATE, even one that +# leaves the path alone. A Core UPDATE or upsert that changes a path must set its hash itself. def _build_fileloc_hash_default(path_key: str) -> Callable[[DefaultExecutionContext], str]: def compute(context: DefaultExecutionContext) -> str: return compute_fileloc_hash(context.get_current_parameters()[path_key]) @@ -75,6 +78,11 @@ class LangSDKTaskHandlerArtifact(Base): ), ) + @validates("relative_fileloc") + def _sync_relative_fileloc_hash(self, key: str, relative_fileloc: str) -> str: + self.relative_fileloc_hash = compute_fileloc_hash(relative_fileloc) + return relative_fileloc + class LangSDKTaskHandler(Base): """The artifact that runs one stub task, and the task handler parameters it declares.""" @@ -110,7 +118,5 @@ class LangSDKTaskHandler(Base): @validates("dag_relative_fileloc") def _sync_dag_relative_fileloc_hash(self, key: str, dag_relative_fileloc: str) -> str: - # Not an onupdate default: that fires on every UPDATE, even one that leaves the path alone. - # A Core UPDATE or upsert that changes the path must set the hash itself. self.dag_relative_fileloc_hash = compute_fileloc_hash(dag_relative_fileloc) return dag_relative_fileloc diff --git a/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py b/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py index 52d5d1f033843..b99b63bacf70f 100644 --- a/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py +++ b/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py @@ -128,20 +128,51 @@ def test_core_insert_fills_fileloc_hashes(testing_dag_bundle, session): ) -def test_dag_relative_fileloc_hash_follows_orm_updates(testing_dag_bundle, session): - artifact = _add_dags_and_artifact(session, "dag_a") - handler = _make_handler(dag_id="dag_a", artifact=artifact) +def _add_artifact(session) -> LangSDKTaskHandlerArtifact: + artifact = _make_artifact() + session.add(artifact) + session.flush() + return artifact + + +def _add_handler(session) -> LangSDKTaskHandler: + handler = _make_handler(dag_id="dag_a", artifact=_add_dags_and_artifact(session, "dag_a")) session.add(handler) session.flush() + return handler + + +@pytest.mark.parametrize( + ("add_row", "path_attr", "hash_column", "other_attr", "other_value"), + [ + pytest.param( + _add_artifact, + "relative_fileloc", + LangSDKTaskHandlerArtifact.relative_fileloc_hash, + "size_bytes", + 2048, + id="artifact", + ), + pytest.param( + _add_handler, + "dag_relative_fileloc", + LangSDKTaskHandler.dag_relative_fileloc_hash, + "handler_params", + [], + id="handler", + ), + ], +) +def test_fileloc_hash_follows_orm_updates( + testing_dag_bundle, session, add_row, path_attr, hash_column, other_attr, other_value +): + row = add_row(session) + original_path = getattr(row, path_attr) - handler.handler_params = [] + setattr(row, other_attr, other_value) session.flush() - assert session.scalar(select(LangSDKTaskHandler.dag_relative_fileloc_hash)) == compute_fileloc_hash( - "dag_a.py" - ) + assert session.scalar(select(hash_column)) == compute_fileloc_hash(original_path) - handler.dag_relative_fileloc = "moved/dag_a.py" + setattr(row, path_attr, "moved/file") session.flush() - assert session.scalar(select(LangSDKTaskHandler.dag_relative_fileloc_hash)) == compute_fileloc_hash( - "moved/dag_a.py" - ) + assert session.scalar(select(hash_column)) == compute_fileloc_hash("moved/file") From ab8f0257223203028a12e9bbeba836694c8fc446 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 14:18:11 +0000 Subject: [PATCH 14/94] Test that a referenced task handler artifact cannot be deleted The handler-to-artifact foreign key has no ON DELETE on purpose, so a cleanup that races with a new binding fails instead of dropping it. Only the cascade from dag was tested. --- .../tests/unit/models/test_lang_sdk_task_handler.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py b/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py index b99b63bacf70f..261a157443167 100644 --- a/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py +++ b/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py @@ -74,6 +74,17 @@ def test_deleting_dag_deletes_its_handlers(testing_dag_bundle, session): assert session.scalars(select(LangSDKTaskHandlerArtifact.id)).all() == [artifact.id] +def test_deleting_referenced_artifact_fails(testing_dag_bundle, session): + artifact = _add_dags_and_artifact(session, "dag_a") + session.add(_make_handler(dag_id="dag_a", artifact=artifact)) + session.flush() + + with pytest.raises(IntegrityError): + session.execute( + delete(LangSDKTaskHandlerArtifact).where(LangSDKTaskHandlerArtifact.id == artifact.id) + ) + + @pytest.mark.parametrize( "relative_fileloc", [ From 5da25f7206e6969fcf3ca1edd394740a4e825d08 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 14:18:53 +0000 Subject: [PATCH 15/94] Describe LangSDKTaskHandler as a stub task binding The old docstring read as if the row were the artifact. A row is the binding of one stub task to an artifact. --- airflow-core/src/airflow/models/lang_sdk_task_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow-core/src/airflow/models/lang_sdk_task_handler.py b/airflow-core/src/airflow/models/lang_sdk_task_handler.py index d53301fe05b4a..986eacaad9cea 100644 --- a/airflow-core/src/airflow/models/lang_sdk_task_handler.py +++ b/airflow-core/src/airflow/models/lang_sdk_task_handler.py @@ -85,7 +85,7 @@ def _sync_relative_fileloc_hash(self, key: str, relative_fileloc: str) -> str: class LangSDKTaskHandler(Base): - """The artifact that runs one stub task, and the task handler parameters it declares.""" + """The binding of one stub task, by ``(dag_id, task_id)``, to its artifact and declared handler parameters.""" __tablename__ = "lang_sdk_task_handler" From 396cac2600dc7c953ff807645494434bbe5dcdc2 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 14:27:53 +0000 Subject: [PATCH 16/94] Regenerate the task handler migration revision id The id and create date of migration 0142 were not made by alembic. They are now alembic's own rev_id() and a real timestamp, and the hooks updated the revision map and migration reference to match. --- airflow-core/docs/migrations-ref.rst | 2 +- .../versions/0142_3_4_0_add_lang_sdk_task_handler_tables.py | 6 +++--- airflow-core/src/airflow/utils/db.py | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/airflow-core/docs/migrations-ref.rst b/airflow-core/docs/migrations-ref.rst index aa57aaf1472d2..70847aee7d4d3 100644 --- a/airflow-core/docs/migrations-ref.rst +++ b/airflow-core/docs/migrations-ref.rst @@ -39,7 +39,7 @@ Here's the list of all the Database Migrations that are executed via when you ru +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | Revision ID | Revises ID | Airflow Version | Description | +=========================+==================+===================+==============================================================+ -| ``ffff41014d89`` (head) | ``90e4d18ccadf`` | ``3.4.0`` | Add lang_sdk_task_handler_artifact and lang_sdk_task_handler | +| ``f7ed13533d23`` (head) | ``90e4d18ccadf`` | ``3.4.0`` | Add lang_sdk_task_handler_artifact and lang_sdk_task_handler | | | | | tables. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | ``90e4d18ccadf`` | ``e5a91c7f42b3`` | ``3.4.0`` | Add timetable_asset_gated to DagModel. | diff --git a/airflow-core/src/airflow/migrations/versions/0142_3_4_0_add_lang_sdk_task_handler_tables.py b/airflow-core/src/airflow/migrations/versions/0142_3_4_0_add_lang_sdk_task_handler_tables.py index 380e5f1dc2a71..ca984a7be2c9c 100644 --- a/airflow-core/src/airflow/migrations/versions/0142_3_4_0_add_lang_sdk_task_handler_tables.py +++ b/airflow-core/src/airflow/migrations/versions/0142_3_4_0_add_lang_sdk_task_handler_tables.py @@ -19,9 +19,9 @@ """ Add lang_sdk_task_handler_artifact and lang_sdk_task_handler tables. -Revision ID: ffff41014d89 +Revision ID: f7ed13533d23 Revises: 90e4d18ccadf -Create Date: 2026-09-30 12:00:00.000000 +Create Date: 2026-09-30 14:19:35.109799 """ @@ -34,7 +34,7 @@ from airflow.utils.sqlalchemy import UtcDateTime # revision identifiers, used by Alembic. -revision = "ffff41014d89" +revision = "f7ed13533d23" down_revision = "90e4d18ccadf" branch_labels = None depends_on = None diff --git a/airflow-core/src/airflow/utils/db.py b/airflow-core/src/airflow/utils/db.py index 56adf7b831731..689d57ded759f 100644 --- a/airflow-core/src/airflow/utils/db.py +++ b/airflow-core/src/airflow/utils/db.py @@ -117,7 +117,7 @@ class MappedClassProtocol(Protocol): "3.1.8": "509b94a1042d", "3.2.0": "1d6611b6ab7c", "3.3.0": "d2f4e1b3c5a7", - "3.4.0": "ffff41014d89", + "3.4.0": "f7ed13533d23", } # Prefix used to identify tables holding data moved during migration. From 61b2d76b38e31edce15a9109ef91509e225da649 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 15:15:51 +0000 Subject: [PATCH 17/94] Store the task handler binding mode beside its params A task handler declaration now says whether it binds stub arguments positionally, by name, or by name or as a whole. The fast path checks a changed Python file against the stored declaration without running the runtime, so it needs the mode as well as the params. --- ..._3_4_0_add_lang_sdk_task_handler_tables.py | 1 + .../airflow/models/lang_sdk_task_handler.py | 1 + .../unit/models/test_lang_sdk_task_handler.py | 21 ++++++++++++++++++- 3 files changed, 22 insertions(+), 1 deletion(-) diff --git a/airflow-core/src/airflow/migrations/versions/0142_3_4_0_add_lang_sdk_task_handler_tables.py b/airflow-core/src/airflow/migrations/versions/0142_3_4_0_add_lang_sdk_task_handler_tables.py index ca984a7be2c9c..a6112e45a5b24 100644 --- a/airflow-core/src/airflow/migrations/versions/0142_3_4_0_add_lang_sdk_task_handler_tables.py +++ b/airflow-core/src/airflow/migrations/versions/0142_3_4_0_add_lang_sdk_task_handler_tables.py @@ -67,6 +67,7 @@ def upgrade(): sa.Column("dag_bundle_name", StringID(), nullable=False), sa.Column("dag_relative_fileloc", sa.String(length=2000), nullable=False), sa.Column("dag_relative_fileloc_hash", sa.String(length=32), nullable=False), + sa.Column("handler_binding", sa.String(length=20), nullable=False), sa.Column("handler_params", sa.JSON(), nullable=False), sa.PrimaryKeyConstraint("dag_id", "task_id", name=op.f("lang_sdk_task_handler_pkey")), sa.ForeignKeyConstraint( diff --git a/airflow-core/src/airflow/models/lang_sdk_task_handler.py b/airflow-core/src/airflow/models/lang_sdk_task_handler.py index 986eacaad9cea..0489e21c1b2f2 100644 --- a/airflow-core/src/airflow/models/lang_sdk_task_handler.py +++ b/airflow-core/src/airflow/models/lang_sdk_task_handler.py @@ -97,6 +97,7 @@ class LangSDKTaskHandler(Base): dag_relative_fileloc_hash: Mapped[str] = mapped_column( String(32), nullable=False, default=_build_fileloc_hash_default("dag_relative_fileloc") ) + handler_binding: Mapped[str] = mapped_column(String(20), nullable=False) handler_params: Mapped[list[dict[str, Any]]] = mapped_column(sa.JSON(), nullable=False) __table_args__ = ( diff --git a/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py b/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py index 261a157443167..4f03fa8008683 100644 --- a/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py +++ b/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py @@ -34,6 +34,10 @@ ARTIFACT_BUNDLE = "java-task-handlers" # 6000 bytes in UTF-8, over both the MySQL key limit and the Postgres btree entry limit. LONG_NON_ASCII_PATH = "任" * 2000 +HANDLER_PARAMS = [ + {"name": "path", "value_schema": {"type": "string"}, "required": True, "exact_name": True}, + {"name": None, "value_schema": None, "required": False, "exact_name": False}, +] def _make_artifact( @@ -51,7 +55,8 @@ def _make_handler(*, dag_id: str, artifact: LangSDKTaskHandlerArtifact) -> LangS artifact_id=artifact.id, dag_bundle_name="testing", dag_relative_fileloc=f"{dag_id}.py", - handler_params=[{"name": "path", "value_schema": None, "required": True}], + handler_binding="positional", + handler_params=HANDLER_PARAMS, ) @@ -74,6 +79,19 @@ def test_deleting_dag_deletes_its_handlers(testing_dag_bundle, session): assert session.scalars(select(LangSDKTaskHandlerArtifact.id)).all() == [artifact.id] +def test_handler_stores_its_declaration(testing_dag_bundle, session): + artifact = _add_dags_and_artifact(session, "dag_a") + session.add(_make_handler(dag_id="dag_a", artifact=artifact)) + session.flush() + session.expire_all() + + stored = session.execute( + select(LangSDKTaskHandler.handler_binding, LangSDKTaskHandler.handler_params) + ).one() + + assert tuple(stored) == ("positional", HANDLER_PARAMS) + + def test_deleting_referenced_artifact_fails(testing_dag_bundle, session): artifact = _add_dags_and_artifact(session, "dag_a") session.add(_make_handler(dag_id="dag_a", artifact=artifact)) @@ -127,6 +145,7 @@ def test_core_insert_fills_fileloc_hashes(testing_dag_bundle, session): artifact_id=artifact_id, dag_bundle_name="testing", dag_relative_fileloc="dags/etl.py", + handler_binding="named_or_whole", handler_params=[], ) ) From fce1160d9b86b435717916d350c01ae8fc910231 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 15:16:54 +0000 Subject: [PATCH 18/94] Add handler_binding to ADR-0013's task handler table SQL The migration now stores the declaration's binding mode, and later layers read this SQL as the spec. --- .../adr/lang-sdk/0013-persisted-task-handler-bindings.md | 1 + 1 file changed, 1 insertion(+) diff --git a/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md b/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md index d4d242a2152e9..780a5bc4d7377 100644 --- a/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md +++ b/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md @@ -175,6 +175,7 @@ CREATE TABLE lang_sdk_task_handler ( dag_bundle_name VARCHAR(250) NOT NULL, -- the *Python* file that owns this row dag_relative_fileloc VARCHAR(2000) NOT NULL, -- ditto dag_relative_fileloc_hash VARCHAR(32) NOT NULL, -- md5 of dag_relative_fileloc + handler_binding VARCHAR(20) NOT NULL, -- positional | named | named_or_whole handler_params JSON NOT NULL, -- list[TaskHandlerParam], ordered CONSTRAINT lang_sdk_task_handler_pkey PRIMARY KEY (dag_id, task_id), CONSTRAINT lang_sdk_task_handler_dag_id_fkey FOREIGN KEY (dag_id) From d18d0afe556b80105510f04d8c3102a8a664b045 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 1 Oct 2026 06:01:15 +0000 Subject: [PATCH 19/94] Make session keyword-only in the task handler model test helpers airflow-core takes session keyword-only so that it can never be passed in the wrong position, and review flagged the new test helpers for taking it positionally. --- .../unit/models/test_lang_sdk_task_handler.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py b/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py index 4f03fa8008683..b696487f0c108 100644 --- a/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py +++ b/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py @@ -60,7 +60,7 @@ def _make_handler(*, dag_id: str, artifact: LangSDKTaskHandlerArtifact) -> LangS ) -def _add_dags_and_artifact(session, *dag_ids: str) -> LangSDKTaskHandlerArtifact: +def _add_dags_and_artifact(*dag_ids: str, session) -> LangSDKTaskHandlerArtifact: session.add_all(DagModel(dag_id=dag_id, bundle_name="testing") for dag_id in dag_ids) artifact = _make_artifact() session.add(artifact) @@ -69,7 +69,7 @@ def _add_dags_and_artifact(session, *dag_ids: str) -> LangSDKTaskHandlerArtifact def test_deleting_dag_deletes_its_handlers(testing_dag_bundle, session): - artifact = _add_dags_and_artifact(session, "dag_a", "dag_b") + artifact = _add_dags_and_artifact("dag_a", "dag_b", session=session) session.add_all(_make_handler(dag_id=dag_id, artifact=artifact) for dag_id in ("dag_a", "dag_b")) session.flush() @@ -80,7 +80,7 @@ def test_deleting_dag_deletes_its_handlers(testing_dag_bundle, session): def test_handler_stores_its_declaration(testing_dag_bundle, session): - artifact = _add_dags_and_artifact(session, "dag_a") + artifact = _add_dags_and_artifact("dag_a", session=session) session.add(_make_handler(dag_id="dag_a", artifact=artifact)) session.flush() session.expire_all() @@ -93,7 +93,7 @@ def test_handler_stores_its_declaration(testing_dag_bundle, session): def test_deleting_referenced_artifact_fails(testing_dag_bundle, session): - artifact = _add_dags_and_artifact(session, "dag_a") + artifact = _add_dags_and_artifact("dag_a", session=session) session.add(_make_handler(dag_id="dag_a", artifact=artifact)) session.flush() @@ -158,15 +158,15 @@ def test_core_insert_fills_fileloc_hashes(testing_dag_bundle, session): ) -def _add_artifact(session) -> LangSDKTaskHandlerArtifact: +def _add_artifact(*, session) -> LangSDKTaskHandlerArtifact: artifact = _make_artifact() session.add(artifact) session.flush() return artifact -def _add_handler(session) -> LangSDKTaskHandler: - handler = _make_handler(dag_id="dag_a", artifact=_add_dags_and_artifact(session, "dag_a")) +def _add_handler(*, session) -> LangSDKTaskHandler: + handler = _make_handler(dag_id="dag_a", artifact=_add_dags_and_artifact("dag_a", session=session)) session.add(handler) session.flush() return handler @@ -196,7 +196,7 @@ def _add_handler(session) -> LangSDKTaskHandler: def test_fileloc_hash_follows_orm_updates( testing_dag_bundle, session, add_row, path_attr, hash_column, other_attr, other_value ): - row = add_row(session) + row = add_row(session=session) original_path = getattr(row, path_attr) setattr(row, other_attr, other_value) From 7b572ab02d506f7a5c4eb2b8cb67a0f212fe094e Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 1 Oct 2026 07:04:39 +0000 Subject: [PATCH 20/94] Cache an artifact's task handlers on its row, not on each binding The fast path trusts an artifact's fingerprint as the cache key, but a binding stored only what its own file's probe asked about, so one file probing a rebuilt or colliding artifact left another file's bindings stale. Once the runtime is asked for every handler, its answer depends on the artifact alone, so it belongs beside the fingerprint and is written by the same probe. The binding's own copy then has no reader. --- ..._3_4_0_add_lang_sdk_task_handler_tables.py | 3 +- .../airflow/models/lang_sdk_task_handler.py | 12 ++-- .../unit/models/test_lang_sdk_task_handler.py | 56 ++++++++++++------- 3 files changed, 45 insertions(+), 26 deletions(-) diff --git a/airflow-core/src/airflow/migrations/versions/0142_3_4_0_add_lang_sdk_task_handler_tables.py b/airflow-core/src/airflow/migrations/versions/0142_3_4_0_add_lang_sdk_task_handler_tables.py index a6112e45a5b24..c09bc759f49d9 100644 --- a/airflow-core/src/airflow/migrations/versions/0142_3_4_0_add_lang_sdk_task_handler_tables.py +++ b/airflow-core/src/airflow/migrations/versions/0142_3_4_0_add_lang_sdk_task_handler_tables.py @@ -51,6 +51,7 @@ def upgrade(): sa.Column("relative_fileloc_hash", sa.String(length=32), nullable=False), sa.Column("size_bytes", sa.BigInteger(), nullable=False), sa.Column("cache_digest", sa.String(length=128), nullable=False), + sa.Column("task_handlers", sa.JSON(), nullable=False), sa.Column("last_probed_at", UtcDateTime(), nullable=False), sa.PrimaryKeyConstraint("id", name=op.f("lang_sdk_task_handler_artifact_pkey")), sa.UniqueConstraint( @@ -67,8 +68,6 @@ def upgrade(): sa.Column("dag_bundle_name", StringID(), nullable=False), sa.Column("dag_relative_fileloc", sa.String(length=2000), nullable=False), sa.Column("dag_relative_fileloc_hash", sa.String(length=32), nullable=False), - sa.Column("handler_binding", sa.String(length=20), nullable=False), - sa.Column("handler_params", sa.JSON(), nullable=False), sa.PrimaryKeyConstraint("dag_id", "task_id", name=op.f("lang_sdk_task_handler_pkey")), sa.ForeignKeyConstraint( columns=("dag_id",), diff --git a/airflow-core/src/airflow/models/lang_sdk_task_handler.py b/airflow-core/src/airflow/models/lang_sdk_task_handler.py index 0489e21c1b2f2..f60aaa602ed05 100644 --- a/airflow-core/src/airflow/models/lang_sdk_task_handler.py +++ b/airflow-core/src/airflow/models/lang_sdk_task_handler.py @@ -56,7 +56,12 @@ def compute(context: DefaultExecutionContext) -> str: class LangSDKTaskHandlerArtifact(Base): - """A Language SDK artifact in a task handler Dag bundle, cached with its fingerprint.""" + """ + A Language SDK artifact in a task handler Dag bundle, cached with its fingerprint and probe answer. + + ``task_handlers`` is the runtime's whole answer, ``{dag_id: [TaskHandlerDeclaration as JSON, ...]}``. + It is written only from a probe and together with the fingerprint, so the two always match. + """ __tablename__ = "lang_sdk_task_handler_artifact" @@ -68,6 +73,7 @@ class LangSDKTaskHandlerArtifact(Base): ) size_bytes: Mapped[int] = mapped_column(BigInteger, nullable=False) cache_digest: Mapped[str] = mapped_column(String(128), nullable=False) + task_handlers: Mapped[dict[str, list[dict[str, Any]]]] = mapped_column(sa.JSON(), nullable=False) last_probed_at: Mapped[datetime] = mapped_column(UtcDateTime, nullable=False, default=timezone.utcnow) __table_args__ = ( @@ -85,7 +91,7 @@ def _sync_relative_fileloc_hash(self, key: str, relative_fileloc: str) -> str: class LangSDKTaskHandler(Base): - """The binding of one stub task, by ``(dag_id, task_id)``, to its artifact and declared handler parameters.""" + """The binding of one stub task, by ``(dag_id, task_id)``, to the artifact that registers its handler.""" __tablename__ = "lang_sdk_task_handler" @@ -97,8 +103,6 @@ class LangSDKTaskHandler(Base): dag_relative_fileloc_hash: Mapped[str] = mapped_column( String(32), nullable=False, default=_build_fileloc_hash_default("dag_relative_fileloc") ) - handler_binding: Mapped[str] = mapped_column(String(20), nullable=False) - handler_params: Mapped[list[dict[str, Any]]] = mapped_column(sa.JSON(), nullable=False) __table_args__ = ( ForeignKeyConstraint( diff --git a/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py b/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py index b696487f0c108..f1b549ebc2818 100644 --- a/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py +++ b/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py @@ -17,6 +17,8 @@ # under the License. from __future__ import annotations +from typing import Any + import pytest import uuid6 from sqlalchemy import delete, insert, select @@ -34,17 +36,32 @@ ARTIFACT_BUNDLE = "java-task-handlers" # 6000 bytes in UTF-8, over both the MySQL key limit and the Postgres btree entry limit. LONG_NON_ASCII_PATH = "任" * 2000 -HANDLER_PARAMS = [ - {"name": "path", "value_schema": {"type": "string"}, "required": True, "exact_name": True}, - {"name": None, "value_schema": None, "required": False, "exact_name": False}, -] +TASK_HANDLERS = { + "etl": [ + { + "task_id": "extract", + "binding": "positional", + "params": [ + {"name": "path", "value_schema": {"type": "string"}, "required": True, "exact_name": True}, + {"name": None, "value_schema": None, "required": False, "exact_name": False}, + ], + } + ], +} def _make_artifact( - *, bundle_name: str = ARTIFACT_BUNDLE, relative_fileloc: str = "etl.jar" + *, + bundle_name: str = ARTIFACT_BUNDLE, + relative_fileloc: str = "etl.jar", + task_handlers: dict[str, list[dict[str, Any]]] | None = None, ) -> LangSDKTaskHandlerArtifact: return LangSDKTaskHandlerArtifact( - bundle_name=bundle_name, relative_fileloc=relative_fileloc, size_bytes=1024, cache_digest="0" * 64 + bundle_name=bundle_name, + relative_fileloc=relative_fileloc, + size_bytes=1024, + cache_digest="0" * 64, + task_handlers=TASK_HANDLERS if task_handlers is None else task_handlers, ) @@ -55,8 +72,6 @@ def _make_handler(*, dag_id: str, artifact: LangSDKTaskHandlerArtifact) -> LangS artifact_id=artifact.id, dag_bundle_name="testing", dag_relative_fileloc=f"{dag_id}.py", - handler_binding="positional", - handler_params=HANDLER_PARAMS, ) @@ -79,17 +94,19 @@ def test_deleting_dag_deletes_its_handlers(testing_dag_bundle, session): assert session.scalars(select(LangSDKTaskHandlerArtifact.id)).all() == [artifact.id] -def test_handler_stores_its_declaration(testing_dag_bundle, session): - artifact = _add_dags_and_artifact("dag_a", session=session) - session.add(_make_handler(dag_id="dag_a", artifact=artifact)) +@pytest.mark.parametrize( + "task_handlers", + [ + pytest.param(TASK_HANDLERS, id="handlers"), + pytest.param({}, id="none"), + ], +) +def test_artifact_stores_its_task_handlers(session, task_handlers): + session.add(_make_artifact(task_handlers=task_handlers)) session.flush() session.expire_all() - stored = session.execute( - select(LangSDKTaskHandler.handler_binding, LangSDKTaskHandler.handler_params) - ).one() - - assert tuple(stored) == ("positional", HANDLER_PARAMS) + assert session.scalar(select(LangSDKTaskHandlerArtifact.task_handlers)) == task_handlers def test_deleting_referenced_artifact_fails(testing_dag_bundle, session): @@ -136,6 +153,7 @@ def test_core_insert_fills_fileloc_hashes(testing_dag_bundle, session): relative_fileloc="etl.jar", size_bytes=1024, cache_digest="0" * 64, + task_handlers={}, ) ) session.execute( @@ -145,8 +163,6 @@ def test_core_insert_fills_fileloc_hashes(testing_dag_bundle, session): artifact_id=artifact_id, dag_bundle_name="testing", dag_relative_fileloc="dags/etl.py", - handler_binding="named_or_whole", - handler_params=[], ) ) @@ -187,8 +203,8 @@ def _add_handler(*, session) -> LangSDKTaskHandler: _add_handler, "dag_relative_fileloc", LangSDKTaskHandler.dag_relative_fileloc_hash, - "handler_params", - [], + "dag_bundle_name", + "other-bundle", id="handler", ), ], From 03613039ea9cac9dfe8ef08b0647bd67e7dabc44 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 1 Oct 2026 07:05:11 +0000 Subject: [PATCH 21/94] Move the handler declarations to the artifact in ADR-0013's SQL The migration now caches the probe answer on the artifact row, and later layers read this SQL as the spec. --- .../lang-sdk/0013-persisted-task-handler-bindings.md | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md b/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md index 780a5bc4d7377..74a93c0e4863c 100644 --- a/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md +++ b/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md @@ -163,6 +163,7 @@ CREATE TABLE lang_sdk_task_handler_artifact ( relative_fileloc_hash VARCHAR(32) NOT NULL, -- md5 of relative_fileloc; the path is too long to index size_bytes BIGINT NOT NULL, -- cheap fingerprint tier cache_digest VARCHAR(128) NOT NULL, -- content fingerprint tier; see "The fast path" + task_handlers JSON NOT NULL, -- {dag_id: [TaskHandlerDeclaration, ...]}, the probe answer last_probed_at TIMESTAMP NOT NULL, CONSTRAINT lang_sdk_task_handler_artifact_pkey PRIMARY KEY (id), CONSTRAINT lang_sdk_task_handler_artifact_bundle_fileloc_uq UNIQUE (bundle_name, relative_fileloc_hash) @@ -175,8 +176,6 @@ CREATE TABLE lang_sdk_task_handler ( dag_bundle_name VARCHAR(250) NOT NULL, -- the *Python* file that owns this row dag_relative_fileloc VARCHAR(2000) NOT NULL, -- ditto dag_relative_fileloc_hash VARCHAR(32) NOT NULL, -- md5 of dag_relative_fileloc - handler_binding VARCHAR(20) NOT NULL, -- positional | named | named_or_whole - handler_params JSON NOT NULL, -- list[TaskHandlerParam], ordered CONSTRAINT lang_sdk_task_handler_pkey PRIMARY KEY (dag_id, task_id), CONSTRAINT lang_sdk_task_handler_dag_id_fkey FOREIGN KEY (dag_id) REFERENCES dag (dag_id) ON DELETE CASCADE, @@ -198,11 +197,10 @@ joining through `DagModel`: `dag.relative_fileloc` is not indexed, and the codeb that querying it means "a sequential scan of dag" ([`reassign_dags_with_unconfigured_bundles`](https://github.com/apache/airflow/blob/79991cd4db0c9346a28b23c453377f6df0c6b4ed/airflow-core/src/airflow/dag_processing/bundles/manager.py#L478)). -`handler_params` stores what the runtime declared, so a changed Python file can be re-validated -against a cached declaration with no subprocess. It is deliberately **not** called `arg_bindings`: -that name already denotes the Python side of the comparison (`XComArgBinding` / `LiteralArgBinding`, -carrying wiring and values, [`build_arg_bindings`](https://github.com/apache/airflow/blob/79991cd4db0c9346a28b23c453377f6df0c6b4ed/airflow-core/src/airflow/serialization/stub_arg_bindings.py#L221-L286)), -and reusing it would make the validation read as comparing a thing to itself. +`task_handlers` caches the artifact's whole probe answer, every handler it registers keyed by +`dag_id`, so a changed Python file can be re-validated against it with no subprocess. It is written +only from a probe, together with `size_bytes` and `cache_digest`, so the answer always belongs to the +fingerprint beside it. A binding row is then only a mapping from a stub task to its artifact. `cache_digest` is **opaque and coordinator-defined**, not "SHA-256 of the file". From 2c96b482aa64df25b03864d2b4657abe1411f9df Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 1 Oct 2026 07:11:58 +0000 Subject: [PATCH 22/94] Let a task handler artifact store no cache digest Some artifacts carry no digest: a Go bundle packed before digests.cache existed, or a third-party coordinator that computes none. Such a candidate is still probed and bound, so it needs an artifact row. A NULL digest never matches, so the candidate is probed on every parse. --- .../0142_3_4_0_add_lang_sdk_task_handler_tables.py | 2 +- .../src/airflow/models/lang_sdk_task_handler.py | 3 ++- .../tests/unit/models/test_lang_sdk_task_handler.py | 11 ++++++++++- 3 files changed, 13 insertions(+), 3 deletions(-) diff --git a/airflow-core/src/airflow/migrations/versions/0142_3_4_0_add_lang_sdk_task_handler_tables.py b/airflow-core/src/airflow/migrations/versions/0142_3_4_0_add_lang_sdk_task_handler_tables.py index c09bc759f49d9..c9679ede25bec 100644 --- a/airflow-core/src/airflow/migrations/versions/0142_3_4_0_add_lang_sdk_task_handler_tables.py +++ b/airflow-core/src/airflow/migrations/versions/0142_3_4_0_add_lang_sdk_task_handler_tables.py @@ -50,7 +50,7 @@ def upgrade(): sa.Column("relative_fileloc", sa.String(length=2000), nullable=False), sa.Column("relative_fileloc_hash", sa.String(length=32), nullable=False), sa.Column("size_bytes", sa.BigInteger(), nullable=False), - sa.Column("cache_digest", sa.String(length=128), nullable=False), + sa.Column("cache_digest", sa.String(length=128), nullable=True), sa.Column("task_handlers", sa.JSON(), nullable=False), sa.Column("last_probed_at", UtcDateTime(), nullable=False), sa.PrimaryKeyConstraint("id", name=op.f("lang_sdk_task_handler_artifact_pkey")), diff --git a/airflow-core/src/airflow/models/lang_sdk_task_handler.py b/airflow-core/src/airflow/models/lang_sdk_task_handler.py index f60aaa602ed05..86cb020f3f1d7 100644 --- a/airflow-core/src/airflow/models/lang_sdk_task_handler.py +++ b/airflow-core/src/airflow/models/lang_sdk_task_handler.py @@ -72,7 +72,8 @@ class LangSDKTaskHandlerArtifact(Base): String(32), nullable=False, default=_build_fileloc_hash_default("relative_fileloc") ) size_bytes: Mapped[int] = mapped_column(BigInteger, nullable=False) - cache_digest: Mapped[str] = mapped_column(String(128), nullable=False) + # NULL means the artifact stores no fingerprint and is probed again on every parse. + cache_digest: Mapped[str | None] = mapped_column(String(128), nullable=True) task_handlers: Mapped[dict[str, list[dict[str, Any]]]] = mapped_column(sa.JSON(), nullable=False) last_probed_at: Mapped[datetime] = mapped_column(UtcDateTime, nullable=False, default=timezone.utcnow) diff --git a/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py b/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py index f1b549ebc2818..510e2117dd299 100644 --- a/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py +++ b/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py @@ -54,13 +54,14 @@ def _make_artifact( *, bundle_name: str = ARTIFACT_BUNDLE, relative_fileloc: str = "etl.jar", + cache_digest: str | None = "0" * 64, task_handlers: dict[str, list[dict[str, Any]]] | None = None, ) -> LangSDKTaskHandlerArtifact: return LangSDKTaskHandlerArtifact( bundle_name=bundle_name, relative_fileloc=relative_fileloc, size_bytes=1024, - cache_digest="0" * 64, + cache_digest=cache_digest, task_handlers=TASK_HANDLERS if task_handlers is None else task_handlers, ) @@ -109,6 +110,14 @@ def test_artifact_stores_its_task_handlers(session, task_handlers): assert session.scalar(select(LangSDKTaskHandlerArtifact.task_handlers)) == task_handlers +def test_artifact_can_store_no_cache_digest(session): + session.add(_make_artifact(cache_digest=None)) + session.flush() + session.expire_all() + + assert session.scalar(select(LangSDKTaskHandlerArtifact.cache_digest)) is None + + def test_deleting_referenced_artifact_fails(testing_dag_bundle, session): artifact = _add_dags_and_artifact("dag_a", session=session) session.add(_make_handler(dag_id="dag_a", artifact=artifact)) From da3372e04cd3d092b2eacaca9d88c2849f7a4845 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 1 Oct 2026 07:12:27 +0000 Subject: [PATCH 23/94] Allow a NULL cache digest in ADR-0013's artifact table SQL The migration now lets an artifact store no digest, and later layers read this SQL as the spec. --- .../adr/lang-sdk/0013-persisted-task-handler-bindings.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md b/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md index 74a93c0e4863c..d62927452aba0 100644 --- a/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md +++ b/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md @@ -162,7 +162,7 @@ CREATE TABLE lang_sdk_task_handler_artifact ( relative_fileloc VARCHAR(2000) NOT NULL, -- path within that bundle relative_fileloc_hash VARCHAR(32) NOT NULL, -- md5 of relative_fileloc; the path is too long to index size_bytes BIGINT NOT NULL, -- cheap fingerprint tier - cache_digest VARCHAR(128) NOT NULL, -- content fingerprint tier; see "The fast path" + cache_digest VARCHAR(128) NULL, -- content fingerprint tier, NULL if none is stored; see "The fast path" task_handlers JSON NOT NULL, -- {dag_id: [TaskHandlerDeclaration, ...]}, the probe answer last_probed_at TIMESTAMP NOT NULL, CONSTRAINT lang_sdk_task_handler_artifact_pkey PRIMARY KEY (id), From db4f8470d08d564ccf61b737220b610d6768c3ae Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Tue, 29 Sep 2026 16:25:27 +0000 Subject: [PATCH 24/94] Move the Dag file processor's shared plumbing into a base class The native Dag parse and the task-handler probe start parse children other than the Python one, and they need the same parse log forwarding, request handling, readiness and cleanup. No behavior change. --- .../src/airflow/dag_processing/processor.py | 148 ++++++++++-------- 1 file changed, 82 insertions(+), 66 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/processor.py b/airflow-core/src/airflow/dag_processing/processor.py index 80e6d7161dec2..dbe6e5c2becb4 100644 --- a/airflow-core/src/airflow/dag_processing/processor.py +++ b/airflow-core/src/airflow/dag_processing/processor.py @@ -23,7 +23,7 @@ import traceback from collections.abc import Callable, Sequence from pathlib import Path -from typing import TYPE_CHECKING, Annotated, BinaryIO, ClassVar, Literal +from typing import TYPE_CHECKING, Annotated, Any, BinaryIO, ClassVar, Literal import attrs from pydantic import BaseModel, Field, TypeAdapter @@ -92,7 +92,7 @@ from airflow.sdk.definitions.context import Context from airflow.sdk.definitions.dag import DAG from airflow.sdk.definitions.mappedoperator import MappedOperator - from airflow.sdk.execution_time.supervisor import RequestHandler, RequestResult + from airflow.sdk.execution_time.supervisor import RequestHandler, RequestResult, ResponseSent from airflow.typing_compat import Self @@ -573,18 +573,17 @@ def in_process_api_server() -> InProcessExecutionAPI: @attrs.define(kw_only=True) -class DagFileProcessorProcess(WatchedSubprocess, LoggingMixin): +class BaseDagFileProcessorProcess(WatchedSubprocess, LoggingMixin): """ - Parses dags with Task SDK API. - - This class provides a wrapper and management around a subprocess to parse a specific DAG file. + Parse one Dag file in a child process for the Dag processor manager. - Since DAGs are written with the Task SDK, we need to parse them in a task SDK process such that - we can use the Task SDK definitions when serializing. This prevents potential conflicts with classes - in core Airflow. + The child's output goes to the file's parse log, and its requests are answered with + :attr:`client`. The parse is done once the child has exited and all its sockets are closed; + :attr:`parsing_result` then holds what it sent. Subclasses start the child and send it the + parse request. """ - logger_filehandle: BinaryIO + logger_filehandle: BinaryIO | None = None parsing_result: DagFileParsingResult | None = None decoder: ClassVar[TypeAdapter[ToManager]] = TypeAdapter[ToManager](ToManager) had_callbacks: bool = False # Track if this process was started with callbacks to prevent stale DAG detection false positives @@ -595,60 +594,6 @@ class DagFileProcessorProcess(WatchedSubprocess, LoggingMixin): bundle_name: str dag_file_rel_path: str - @classmethod - def start( # type: ignore[override] - cls, - *, - path: str | os.PathLike[str], - bundle_path: Path, - bundle_name: str, - dag_file_rel_path: str, - callbacks: list[CallbackRequest], - target: Callable[[], None] = _parse_file_entrypoint, - client: Client, - **kwargs, - ) -> Self: - logger = kwargs["logger"] - - # Parsing DAG files runs user code that can trigger macOS-unsafe ObjC - # initialization (secret backends, connection/variable lookups, HTTP - # clients). Fork+exec a clean interpreter there. Tests override `target` - # with a stub to exercise the base infrastructure; keep bare fork for those. - use_exec = target is _parse_file_entrypoint and supervisor._should_use_exec() - - # Pre-importing only helps the bare-fork child (it inherits the imports via - # copy-on-write). An exec'd child re-imports from scratch, so skip it there - # to avoid leaking user modules into the long-lived processor manager. - if not use_exec: - _pre_import_airflow_modules(os.fspath(path), logger) - - proc: Self = super().start( - target=target, - client=client, - bundle_name=bundle_name, - dag_file_rel_path=dag_file_rel_path, - use_exec=use_exec, - **kwargs, - ) - proc.had_callbacks = bool(callbacks) # Track if this process had callbacks - proc._on_child_started(callbacks, path, bundle_path, bundle_name) - return proc - - def _on_child_started( - self, - callbacks: list[CallbackRequest], - path: str | os.PathLike[str], - bundle_path: Path, - bundle_name: str, - ) -> None: - msg = DagFileParseRequest( - file=os.fspath(path), - bundle_path=bundle_path, - bundle_name=bundle_name, - callback_requests=callbacks, - ) - self.send_msg(msg, request_id=0) - def _get_target_loggers(self) -> tuple[FilteringBoundLogger, ...]: base = super()._get_target_loggers() if not self.subprocess_logs_to_stdout: @@ -674,11 +619,11 @@ def _create_log_forwarder( def _handle_parsing_result( self, msg: DagFileParsingResult, log: FilteringBoundLogger, req_id: int - ) -> RequestResult: + ) -> RequestResult | ResponseSent: self.parsing_result = msg return None, {} - _request_handlers: ClassVar[dict[type[BaseModel], RequestHandler[DagFileProcessorProcess]]] = { + _request_handlers: ClassVar[dict[type[BaseModel], RequestHandler[Any]]] = { **WatchedSubprocess._get_shared_request_handlers( DeleteVariable, GetConnection, @@ -720,6 +665,8 @@ def wait(self) -> int: def close(self): self.cleanup_sockets_after_kill() + if self.logger_filehandle is None: + return try: self.logger_filehandle.close() except OSError: @@ -728,3 +675,72 @@ def close(self): self.dag_file_rel_path, exc_info=True, ) + + +@attrs.define(kw_only=True) +class DagFileProcessorProcess(BaseDagFileProcessorProcess): + """ + Parses dags with Task SDK API. + + This class provides a wrapper and management around a subprocess to parse a specific DAG file. + + Since DAGs are written with the Task SDK, we need to parse them in a task SDK process such that + we can use the Task SDK definitions when serializing. This prevents potential conflicts with classes + in core Airflow. + """ + + logger_filehandle: BinaryIO + + @classmethod + def start( # type: ignore[override] + cls, + *, + path: str | os.PathLike[str], + bundle_path: Path, + bundle_name: str, + dag_file_rel_path: str, + callbacks: list[CallbackRequest], + target: Callable[[], None] = _parse_file_entrypoint, + client: Client, + **kwargs, + ) -> Self: + logger = kwargs["logger"] + + # Parsing DAG files runs user code that can trigger macOS-unsafe ObjC + # initialization (secret backends, connection/variable lookups, HTTP + # clients). Fork+exec a clean interpreter there. Tests override `target` + # with a stub to exercise the base infrastructure; keep bare fork for those. + use_exec = target is _parse_file_entrypoint and supervisor._should_use_exec() + + # Pre-importing only helps the bare-fork child (it inherits the imports via + # copy-on-write). An exec'd child re-imports from scratch, so skip it there + # to avoid leaking user modules into the long-lived processor manager. + if not use_exec: + _pre_import_airflow_modules(os.fspath(path), logger) + + proc: Self = super().start( + target=target, + client=client, + bundle_name=bundle_name, + dag_file_rel_path=dag_file_rel_path, + use_exec=use_exec, + **kwargs, + ) + proc.had_callbacks = bool(callbacks) # Track if this process had callbacks + proc._on_child_started(callbacks, path, bundle_path, bundle_name) + return proc + + def _on_child_started( + self, + callbacks: list[CallbackRequest], + path: str | os.PathLike[str], + bundle_path: Path, + bundle_name: str, + ) -> None: + msg = DagFileParseRequest( + file=os.fspath(path), + bundle_path=bundle_path, + bundle_name=bundle_name, + callback_requests=callbacks, + ) + self.send_msg(msg, request_id=0) From 7e26f0d11be3e98e894b51c3b2a0d2b864435210 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 1 Oct 2026 06:04:06 +0000 Subject: [PATCH 25/94] Test closing a Dag file processor that has no parse log file The base class makes logger_filehandle optional for parse children other than the Python one, but DagFileProcessorProcess still requires it, so nothing in this layer reached the close path without one. --- .../unit/dag_processing/test_processor.py | 22 ++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/airflow-core/tests/unit/dag_processing/test_processor.py b/airflow-core/tests/unit/dag_processing/test_processor.py index 0cb61aa02866a..9006876b477e9 100644 --- a/airflow-core/tests/unit/dag_processing/test_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_processor.py @@ -27,7 +27,7 @@ import uuid import zipfile from collections.abc import Callable, Iterable -from socket import socketpair +from socket import socket, socketpair from typing import TYPE_CHECKING, Any, BinaryIO from unittest.mock import MagicMock, PropertyMock, patch @@ -54,6 +54,7 @@ from airflow.dag_processing.dagbag import DagBag from airflow.dag_processing.manager import DagFileProcessorManager, process_parse_results from airflow.dag_processing.processor import ( + BaseDagFileProcessorProcess, DagFileParseRequest, DagFileParsingResult, DagFileProcessorProcess, @@ -2546,6 +2547,25 @@ def test_handle_request_get_variable_masks_value_with_key(self, proc): } +class TestBaseDagFileProcessorProcess: + @patch.object(BaseDagFileProcessorProcess, "cleanup_sockets_after_kill", autospec=True) + def test_close_without_log_file_handle_cleans_up_sockets(self, cleanup_sockets_after_kill): + proc = BaseDagFileProcessorProcess( + process_log=structlog.get_logger(), + id=uuid.uuid4(), + pid=1234, + process=MagicMock(spec=supervisor.ProcessTracker), + stdin=MagicMock(spec=socket), + client=MagicMock(spec=Client), + bundle_name="mybundle", + dag_file_rel_path="dags/my_dag.py", + ) + + proc.close() + + cleanup_sockets_after_kill.assert_called_once_with(proc) + + class TestMultiTeamCallbackMetrics: """Tests for team_name tag on dag.callback_exceptions in multi-team mode.""" From 4f4b3bdc55987258e5515875ab2ab7069004953b Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 10:45:40 +0000 Subject: [PATCH 26/94] Add task-handler parse messages and their parse-channel union ADR-0012 gives the Dag-file parse a second question for a Lang-SDK runtime: which task handlers an artifact registers for Dags that Python already owns. This adds its request and result, and ToSDKTaskHandlerProcessor, which shares every response with ToDagProcessor. ToManager gains the result. Nothing sends or receives these messages yet, so DagFileProcessorProcess still rejects the result as an unhandled request. The supervisor schema registry introspects the new union, so the snapshot gains four definitions. They are new bodies, so no VersionChange is needed. --- .../src/airflow/dag_processing/processor.py | 74 ++++++++- .../unit/dag_processing/test_processor.py | 49 +++++- .../generate_supervisor_schemas_snapshot.py | 3 +- .../sdk/execution_time/schema/__init__.py | 17 +- .../sdk/execution_time/schema/schema.json | 150 ++++++++++++++++++ .../execution_time/schema/test_migrator.py | 6 + .../test_supervisor_schemas_name_type_sync.py | 5 +- 7 files changed, 286 insertions(+), 18 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/processor.py b/airflow-core/src/airflow/dag_processing/processor.py index dbe6e5c2becb4..242c051e475a0 100644 --- a/airflow-core/src/airflow/dag_processing/processor.py +++ b/airflow-core/src/airflow/dag_processing/processor.py @@ -29,6 +29,7 @@ from pydantic import BaseModel, Field, TypeAdapter from airflow._shared.observability.metrics import stats +from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import ArgValueSchema # noqa: TC001 from airflow.callbacks.callback_requests import ( CallbackRequest, DagCallbackRequest, @@ -135,8 +136,66 @@ class DagFileParsingResult(BaseModel): type: Literal["DagFileParsingResult"] = "DagFileParsingResult" +class TaskHandlerParam(BaseModel): + """One parameter of a task handler.""" + + name: str + + value_schema: ArgValueSchema | None = None + """JSON Schema of the values the parameter accepts; ``None`` when the handler does not constrain it.""" + + required: bool + """Whether the handler declares no default for the parameter.""" + + +class TaskHandlerDeclaration(BaseModel): + """A task handler that a Lang-SDK artifact registers for one task.""" + + task_id: str + + params: list[TaskHandlerParam] + """In declaration order, because stub-task arguments bind to them by position.""" + + +class TaskHandlerParseRequest(BaseModel): + """ + Request for Task Handler Parsing. + + Asks a Lang-SDK runtime which task handlers an artifact registers for the given Dags. + """ + + file: str + """The artifact to ask.""" + + dag_ids: list[str] + + bundle_path: Path + + bundle_name: str + + type: Literal["TaskHandlerParseRequest"] = "TaskHandlerParseRequest" + + +class TaskHandlerParsingResult(BaseModel): + """ + Result of Task Handler Parsing. + + The task handlers a Lang-SDK artifact registers, keyed by Dag id. + """ + + fileloc: str + + task_handlers: dict[str, list[TaskHandlerDeclaration]] + """A requested Dag id the artifact registers no task handler for is omitted, not mapped to ``[]``.""" + + import_errors: dict[str, str] | None = None + warnings: list | None = None + type: Literal["TaskHandlerParsingResult"] = "TaskHandlerParsingResult" + + ToManager = Annotated[ DagFileParsingResult + | TaskHandlerParsingResult | GetConnection | GetVariable | GetVariableKeys @@ -155,9 +214,9 @@ class DagFileParsingResult(BaseModel): Field(discriminator="type"), ] -ToDagProcessor = Annotated[ - DagFileParseRequest - | ConnectionResult +# Answers to the child's requests, whichever parse it was started for. +_ParseSideResponses = ( + ConnectionResult | VariableResult | VariableKeysResult | TaskStatesResult @@ -169,8 +228,13 @@ class DagFileParsingResult(BaseModel): | XComCountResponse | XComResult | XComSequenceIndexResult - | XComSequenceSliceResult, - Field(discriminator="type"), + | XComSequenceSliceResult +) + +ToDagProcessor = Annotated[DagFileParseRequest | _ParseSideResponses, Field(discriminator="type")] + +ToSDKTaskHandlerProcessor = Annotated[ + TaskHandlerParseRequest | _ParseSideResponses, Field(discriminator="type") ] diff --git a/airflow-core/tests/unit/dag_processing/test_processor.py b/airflow-core/tests/unit/dag_processing/test_processor.py index 9006876b477e9..c66b6a2ed686a 100644 --- a/airflow-core/tests/unit/dag_processing/test_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_processor.py @@ -58,8 +58,11 @@ DagFileParseRequest, DagFileParsingResult, DagFileProcessorProcess, + TaskHandlerParseRequest, + TaskHandlerParsingResult, ToDagProcessor, ToManager, + ToSDKTaskHandlerProcessor, _execute_callbacks, _execute_dag_callbacks, _execute_email_callbacks, @@ -2355,15 +2358,26 @@ def get_type_names(union_type): + "\n\nEither handle these types in ToDagProcessor or update in_task_runner_but_not_in_dag_processing_process list." ) + def test_parse_request_unions_differ_only_in_request(self): + dag_processor_types = set(typing.get_args(typing.get_args(ToDagProcessor)[0])) + task_handler_types = set(typing.get_args(typing.get_args(ToSDKTaskHandlerProcessor)[0])) + + assert dag_processor_types - task_handler_types == {DagFileParseRequest} + assert task_handler_types - dag_processor_types == {TaskHandlerParseRequest} + class TestDagFileProcessorProcess: def test_registered_message_types(self): - expected = set(typing.get_args(typing.get_args(ToManager)[0])) + expected = set(typing.get_args(typing.get_args(ToManager)[0])) - {TaskHandlerParsingResult} assert set(DagFileProcessorProcess._request_handlers) == expected @pytest.mark.parametrize( "message_type", - sorted(set(typing.get_args(typing.get_args(ToManager)[0])) - {DagFileParsingResult}, key=str), + sorted( + set(typing.get_args(typing.get_args(ToManager)[0])) + - {DagFileParsingResult, TaskHandlerParsingResult}, + key=str, + ), ids=lambda message_type: message_type.__name__, ) def test_reuses_shared_request_handlers(self, message_type): @@ -2400,6 +2414,37 @@ def test_dispatch_parsing_result(self, send_msg, proc): assert proc.parsing_result is result send_msg.assert_called_once_with(proc, None, request_id=42, error=None) + @patch.object(DagFileProcessorProcess, "send_msg", autospec=True) + def test_rejects_task_handler_parsing_result(self, send_msg, proc): + msg = proc.decoder.validate_python( + { + "type": "TaskHandlerParsingResult", + "fileloc": "handlers.jar", + "task_handlers": { + "etl": [ + { + "task_id": "extract", + "params": [ + {"name": "day", "value_schema": {"type": "string"}, "required": True}, + {"name": "limit", "required": False}, + ], + } + ] + }, + } + ) + assert isinstance(msg, TaskHandlerParsingResult) + + proc._handle_request(msg, structlog.get_logger(), req_id=42) + + send_msg.assert_called_once_with( + proc, + None, + request_id=42, + error=comms.ErrorResponse(detail={"status_code": 400, "message": "Unhandled request"}), + ) + assert proc.parsing_result is None + @patch.object(DagFileProcessorProcess, "send_msg", autospec=True) def test_previous_successful_run_uses_process_id(self, send_msg, proc): proc.client.task_instances.get_previous_successful_dagrun.return_value = ( diff --git a/scripts/ci/prek/generate_supervisor_schemas_snapshot.py b/scripts/ci/prek/generate_supervisor_schemas_snapshot.py index a626f32718e06..106f12a836663 100755 --- a/scripts/ci/prek/generate_supervisor_schemas_snapshot.py +++ b/scripts/ci/prek/generate_supervisor_schemas_snapshot.py @@ -27,7 +27,8 @@ The snapshot is the head-version JSON Schema for every Pydantic class on the supervisor schema wire (the union members of ``ToTask``, -``ToSupervisor``, ``ToManager``, ``ToDagProcessor``). +``ToSupervisor``, ``ToManager``, ``ToDagProcessor``, +``ToSDKTaskHandlerProcessor``). The actual dump is delegated to ``dump_supervisor_schemas.py`` (the sibling stdout-only script). If the committed snapshot differs from diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/__init__.py b/task-sdk/src/airflow/sdk/execution_time/schema/__init__.py index 585bc84a33e2d..d2861f1665bc9 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/__init__.py +++ b/task-sdk/src/airflow/sdk/execution_time/schema/__init__.py @@ -32,11 +32,12 @@ :func:`registered_models_by_name` resolves a wire-shape ``type`` discriminator to the head Pydantic class. It is computed dynamically -from the four discriminated unions ``ToTask``, ``ToSupervisor`` -(task-execution channel) and ``ToManager``, ``ToDagProcessor`` -(dag-processing channel) so the registry is always in sync with the -actual unions ``CommsDecoder`` decodes against -- no hand-maintained -list to drift. Triggerer unions are intentionally excluded (the +from the five discriminated unions ``ToTask``, ``ToSupervisor`` +(task-execution channel) and ``ToManager``, ``ToDagProcessor``, +``ToSDKTaskHandlerProcessor`` (dag-processing channel) so the registry +is always in sync with the actual unions ``CommsDecoder`` decodes +against -- no hand-maintained list to drift. Triggerer unions are +intentionally excluded (the Triggerer channel is not handled by lang-SDK coordinators today). """ @@ -75,7 +76,7 @@ def registered_models_by_name() -> dict[str, type[BaseModel]]: Map every supervisor schema body's class name to the head Pydantic class. Single source of truth for the registry. Built once by walking the - four discriminated unions the supervisor decodes against; cached + five discriminated unions the supervisor decodes against; cached per-process because the registry only changes when a union member is added in ``comms.py`` or ``processor.py`` (which needs a restart anyway). :func:`resolve_body_class` looks up the wire-shape @@ -91,11 +92,11 @@ def registered_models_by_name() -> dict[str, type[BaseModel]]: single head class, so a name clash is a programmer error that must surface immediately rather than silently picking a winner. """ - from airflow.dag_processing.processor import ToDagProcessor, ToManager + from airflow.dag_processing.processor import ToDagProcessor, ToManager, ToSDKTaskHandlerProcessor from airflow.sdk.execution_time.comms import ToSupervisor, ToTask by_name: dict[str, type[BaseModel]] = {} - for source in (ToTask, ToSupervisor, ToManager, ToDagProcessor): + for source in (ToTask, ToSupervisor, ToManager, ToDagProcessor, ToSDKTaskHandlerProcessor): for model in _iter_model_types(source): existing = by_name.get(model.__name__) if existing is None: diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index d6ee95927efee..45f53e5e4675d 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -4156,6 +4156,156 @@ "title": "TaskCallbackRequest", "type": "object" }, + "TaskHandlerDeclaration": { + "description": "A task handler that a Lang-SDK artifact registers for one task.", + "properties": { + "task_id": { + "title": "Task Id", + "type": "string" + }, + "params": { + "items": { + "$ref": "#/$defs/TaskHandlerParam" + }, + "title": "Params", + "type": "array" + } + }, + "required": [ + "task_id", + "params" + ], + "title": "TaskHandlerDeclaration", + "type": "object" + }, + "TaskHandlerParam": { + "description": "One parameter of a task handler.", + "properties": { + "name": { + "title": "Name", + "type": "string" + }, + "value_schema": { + "anyOf": [ + { + "$ref": "#/$defs/ArgValueSchema" + }, + { + "type": "null" + } + ], + "default": null + }, + "required": { + "title": "Required", + "type": "boolean" + } + }, + "required": [ + "name", + "required" + ], + "title": "TaskHandlerParam", + "type": "object" + }, + "TaskHandlerParseRequest": { + "description": "Request for Task Handler Parsing.\n\nAsks a Lang-SDK runtime which task handlers an artifact registers for the given Dags.", + "properties": { + "file": { + "title": "File", + "type": "string" + }, + "dag_ids": { + "items": { + "type": "string" + }, + "title": "Dag Ids", + "type": "array" + }, + "bundle_path": { + "format": "path", + "title": "Bundle Path", + "type": "string" + }, + "bundle_name": { + "title": "Bundle Name", + "type": "string" + }, + "type": { + "const": "TaskHandlerParseRequest", + "default": "TaskHandlerParseRequest", + "title": "Type", + "type": "string" + } + }, + "required": [ + "file", + "dag_ids", + "bundle_path", + "bundle_name" + ], + "title": "TaskHandlerParseRequest", + "type": "object" + }, + "TaskHandlerParsingResult": { + "description": "Result of Task Handler Parsing.\n\nThe task handlers a Lang-SDK artifact registers, keyed by Dag id.", + "properties": { + "fileloc": { + "title": "Fileloc", + "type": "string" + }, + "task_handlers": { + "additionalProperties": { + "items": { + "$ref": "#/$defs/TaskHandlerDeclaration" + }, + "type": "array" + }, + "title": "Task Handlers", + "type": "object" + }, + "import_errors": { + "anyOf": [ + { + "additionalProperties": { + "type": "string" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Import Errors" + }, + "warnings": { + "anyOf": [ + { + "items": {}, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Warnings" + }, + "type": { + "const": "TaskHandlerParsingResult", + "default": "TaskHandlerParsingResult", + "title": "Type", + "type": "string" + } + }, + "required": [ + "fileloc", + "task_handlers" + ], + "title": "TaskHandlerParsingResult", + "type": "object" + }, "TaskInstanceState": { "description": "All possible states that a Task Instance can be in.\n\nNote that None is also allowed, so always use this in a type hint with Optional.", "enum": [ diff --git a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py index 3befcf14b421b..204ffe7335593 100644 --- a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py +++ b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py @@ -48,6 +48,7 @@ _SupervisorResponse, ) +from airflow.dag_processing.processor import TaskHandlerParseRequest from airflow.sdk import TaskInstanceState from airflow.sdk.execution_time.comms import TaskState from airflow.sdk.execution_time.schema import ( @@ -337,6 +338,11 @@ def test_extra_fields_in_body_do_not_affect_resolution(self, mock_registry): assert resolve_body_class(body) is _LangSdkRequest +class TestRealRegistry: + def test_resolves_task_handler_parse_request(self): + assert resolve_body_class({"type": "TaskHandlerParseRequest"}) is TaskHandlerParseRequest + + class TestLazyCadwynImport: """``cadwyn`` (which imports FastAPI/Starlette/Jinja2) must stay off the worker import path. diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor_schemas_name_type_sync.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor_schemas_name_type_sync.py index e1c3ef1311860..3ad92b430542e 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor_schemas_name_type_sync.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor_schemas_name_type_sync.py @@ -33,7 +33,7 @@ class by ``__name__``. If the two strings drift, a frame decodes against import pytest from pydantic import BaseModel -from airflow.dag_processing.processor import ToDagProcessor, ToManager +from airflow.dag_processing.processor import ToDagProcessor, ToManager, ToSDKTaskHandlerProcessor from airflow.jobs.triggerer_job_runner import ToTriggerRunner, ToTriggerSupervisor from airflow.sdk.execution_time.comms import ToSupervisor, ToTask @@ -45,7 +45,7 @@ def _members_of_union(union: object) -> tuple[type[BaseModel], ...]: return tuple(m for m in get_args(union) if isinstance(m, type) and issubclass(m, BaseModel)) -# All six supervisor discriminated unions. Triggerer's two unions are +# All seven supervisor discriminated unions. Triggerer's two unions are # not part of the lang-SDK-facing registry, but the same name/type # invariant is required for ``CommsDecoder`` to round-trip them. SUPERVISOR_UNIONS = [ @@ -53,6 +53,7 @@ def _members_of_union(union: object) -> tuple[type[BaseModel], ...]: pytest.param(ToSupervisor, id="ToSupervisor"), pytest.param(ToManager, id="ToManager"), pytest.param(ToDagProcessor, id="ToDagProcessor"), + pytest.param(ToSDKTaskHandlerProcessor, id="ToSDKTaskHandlerProcessor"), pytest.param(ToTriggerRunner, id="ToTriggerRunner"), pytest.param(ToTriggerSupervisor, id="ToTriggerSupervisor"), ] From 7e0327ce1e52487215744251f6f2b5f561025b18 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 10:46:26 +0000 Subject: [PATCH 27/94] Regenerate TS SDK supervisor types for task-handler messages check-ts-sdk-supervisor-schema regenerates supervisor.ts from the supervisor schema snapshot and fails on any difference. CI runs it with --all-files whenever a ts-sdk file changes, so leaving it stale after the snapshot gained the task-handler messages would fail the next unrelated ts-sdk change. The Type renumbering is the generator's own output. --- ts-sdk/src/generated/supervisor.ts | 125 ++++++++++++++++++++++------- 1 file changed, 97 insertions(+), 28 deletions(-) diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index 76e9e727815cf..1cd4be4f46046 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -595,15 +595,30 @@ export type Breadcrumbs = { [k: string]: unknown; }[]; export type Type76 = "TaskBreadcrumbsResult"; +export type TaskId10 = string; +export type Name16 = string; +export type Required = boolean; +export type Params2 = TaskHandlerParam[]; +export type File1 = string; +export type DagIds = string[]; +export type BundlePath1 = string; +export type BundleName5 = string; +export type Type77 = "TaskHandlerParseRequest"; +export type Fileloc1 = string; +export type ImportErrors1 = { + [k: string]: string; +} | null; +export type Warnings1 = unknown[] | null; +export type Type78 = "TaskHandlerParsingResult"; export type StartDate7 = string | null; -export type Type77 = "TaskRescheduleStartDate"; +export type Type79 = "TaskRescheduleStartDate"; export type State8 = "failed" | "skipped" | "removed"; export type EndDate8 = string | null; -export type Type78 = "TaskState"; +export type Type80 = "TaskState"; export type RenderedMapIndex5 = string | null; export type RetryReason1 = string | null; -export type Type79 = "TaskStateStoreResult"; -export type Type80 = "TaskStatesResult"; +export type Type81 = "TaskStateStoreResult"; +export type Type82 = "TaskStatesResult"; export type LogicalDate6 = string | null; export type RunAfter2 = string | null; export type Conf2 = { @@ -614,10 +629,10 @@ export type PartitionKey7 = string | null; export type Note2 = string | null; export type DagId22 = string; export type DagRunId = string; -export type Type81 = "TriggerDagRun"; +export type Type83 = "TriggerDagRun"; export type TiId9 = string; export type Note3 = string | null; -export type Type82 = "UpdateDagRunNote"; +export type Type84 = "UpdateDagRunNote"; export type TiId10 = string; /** * @minItems 1 @@ -626,22 +641,22 @@ export type ChosenOptions = [string, ...string[]]; export type ParamsInput = { [k: string]: unknown; } | null; -export type Type83 = "UpdateHITLDetail"; +export type Type85 = "UpdateHITLDetail"; export type TiId11 = string; -export type Type84 = "ValidateInletsAndOutlets"; +export type Type86 = "ValidateInletsAndOutlets"; export type Keys = string[]; export type TotalEntries = number; -export type Type85 = "VariableKeysResult"; +export type Type87 = "VariableKeysResult"; export type Key19 = string; export type Value2 = string | null; -export type Type86 = "VariableResult"; +export type Type88 = "VariableResult"; export type Len = number; -export type Type87 = "XComCountResponse"; +export type Type89 = "XComCountResponse"; export type Key20 = string; -export type Type88 = "XComResult"; -export type Type89 = "XComSequenceIndexResult"; +export type Type90 = "XComResult"; +export type Type91 = "XComSequenceIndexResult"; export type Root = JsonValue[]; -export type Type90 = "XComSequenceSliceResult"; +export type Type92 = "XComSequenceSliceResult"; export interface SupervisorWireSchema {} /** @@ -1824,6 +1839,60 @@ export interface TaskBreadcrumbsResult { breadcrumbs: Breadcrumbs; type?: Type76; } +/** + * A task handler that a Lang-SDK artifact registers for one task. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TaskHandlerDeclaration". + */ +export interface TaskHandlerDeclaration { + task_id: TaskId10; + params: Params2; +} +/** + * One parameter of a task handler. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TaskHandlerParam". + */ +export interface TaskHandlerParam { + name: Name16; + value_schema?: ArgValueSchema | null; + required: Required; +} +/** + * Request for Task Handler Parsing. + * + * Asks a Lang-SDK runtime which task handlers an artifact registers for the given Dags. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TaskHandlerParseRequest". + */ +export interface TaskHandlerParseRequest { + file: File1; + dag_ids: DagIds; + bundle_path: BundlePath1; + bundle_name: BundleName5; + type?: Type77; +} +/** + * Result of Task Handler Parsing. + * + * The task handlers a Lang-SDK artifact registers, keyed by Dag id. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TaskHandlerParsingResult". + */ +export interface TaskHandlerParsingResult { + fileloc: Fileloc1; + task_handlers: TaskHandlers; + import_errors?: ImportErrors1; + warnings?: Warnings1; + type?: Type78; +} +export interface TaskHandlers { + [k: string]: TaskHandlerDeclaration[]; +} /** * Response containing the first reschedule date for a task instance. * @@ -1832,7 +1901,7 @@ export interface TaskBreadcrumbsResult { */ export interface TaskRescheduleStartDate { start_date: StartDate7; - type?: Type77; + type?: Type79; } /** * Update a task's state. @@ -1847,7 +1916,7 @@ export interface TaskRescheduleStartDate { export interface TaskState { state: State8; end_date?: EndDate8; - type?: Type78; + type?: Type80; rendered_map_index?: RenderedMapIndex5; retry_reason?: RetryReason1; } @@ -1859,7 +1928,7 @@ export interface TaskState { */ export interface TaskStateStoreResult { value: JsonValue | null; - type?: Type79; + type?: Type81; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1867,7 +1936,7 @@ export interface TaskStateStoreResult { */ export interface TaskStatesResult { task_states: TaskStates; - type?: Type80; + type?: Type82; } export interface TaskStates { [k: string]: unknown; @@ -1885,7 +1954,7 @@ export interface TriggerDagRun { note?: Note2; dag_id: DagId22; run_id: DagRunId; - type?: Type81; + type?: Type83; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1894,7 +1963,7 @@ export interface TriggerDagRun { export interface UpdateDagRunNote { ti_id: TiId9; note: Note3; - type?: Type82; + type?: Type84; } /** * Update the response content part of an existing Human-in-the-loop response. @@ -1906,7 +1975,7 @@ export interface UpdateHITLDetail { ti_id: TiId10; chosen_options: ChosenOptions; params_input?: ParamsInput; - type?: Type83; + type?: Type85; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1914,7 +1983,7 @@ export interface UpdateHITLDetail { */ export interface ValidateInletsAndOutlets { ti_id: TiId11; - type?: Type84; + type?: Type86; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1923,7 +1992,7 @@ export interface ValidateInletsAndOutlets { export interface VariableKeysResult { keys: Keys; total_entries: TotalEntries; - type?: Type85; + type?: Type87; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1932,7 +2001,7 @@ export interface VariableKeysResult { export interface VariableResult { key: Key19; value: Value2; - type?: Type86; + type?: Type88; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1940,7 +2009,7 @@ export interface VariableResult { */ export interface XComCountResponse { len: Len; - type?: Type87; + type?: Type89; } /** * Response to ReadXCom request. @@ -1951,7 +2020,7 @@ export interface XComCountResponse { export interface XComResult { key: Key20; value: JsonValue | null; - type?: Type88; + type?: Type90; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1959,7 +2028,7 @@ export interface XComResult { */ export interface XComSequenceIndexResult { root: JsonValue; - type?: Type89; + type?: Type91; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1967,7 +2036,7 @@ export interface XComSequenceIndexResult { */ export interface XComSequenceSliceResult { root: Root; - type?: Type90; + type?: Type92; } /** Cadwyn schema version this SDK was generated against. From 2826ec45c7fdca28c0045f0ee0656dce59d45b00 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 10:47:03 +0000 Subject: [PATCH 28/94] Document the version rules for a new supervisor schema body The versions module now says a brand-new body needs no VersionChange, but the schema guide did not say why, and check-supervisor-schemas-versions still fails a local commit that changes the snapshot without touching versions/. The next author of a new body could not tell that failure from a real one, so the guide gives the rule, the local hook behaviour, and the task-handler union. --- .../sdk/execution_time/schema/AGENTS.md | 25 ++++++++++++------- 1 file changed, 16 insertions(+), 9 deletions(-) diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/AGENTS.md b/task-sdk/src/airflow/sdk/execution_time/schema/AGENTS.md index 84c821bad2bbc..362e628eda531 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/AGENTS.md +++ b/task-sdk/src/airflow/sdk/execution_time/schema/AGENTS.md @@ -15,10 +15,10 @@ models stay in their semantic homes: `ToTask` and `ToSupervisor` discriminated unions in `airflow.sdk.execution_time.comms`. - Dag-processing channel (manager ↔ parser-supervisor): the - `ToManager` and `ToDagProcessor` discriminated unions in - `airflow.dag_processing.processor`. + `ToManager`, `ToDagProcessor` and `ToSDKTaskHandlerProcessor` + discriminated unions in `airflow.dag_processing.processor`. -`registered_models_by_name()` introspects those four unions on first +`registered_models_by_name()` introspects those five unions on first call, so the snapshot the prek hook commits to `schema.json` always matches the exact set of classes `CommsDecoder` actually decodes against — there is no hand-maintained list to keep in sync. The @@ -58,12 +58,19 @@ version bump, and vice versa. Append the class to the relevant discriminated union in its semantic home — `ToTask` / `ToSupervisor` in `comms.py`, or `ToManager` / -`ToDagProcessor` in `processor.py`. That is the *only* registration -step; `registered_models_by_name()` picks it up automatically the next -time the snapshot hook runs. - -No `VersionChange` entry is required on the first introduction — the -head shape *is* the schema for the new body. +`ToDagProcessor` / `ToSDKTaskHandlerProcessor` in `processor.py`. That +is the *only* registration step; `registered_models_by_name()` picks it +up automatically the next time the snapshot hook runs. + +No `VersionChange` entry is required on the first introduction: a +lang-SDK pinned to an older version never sends the new body, so there +is nothing to strip on the way down. Only changes to bodies that +already existed at an earlier dated version need one. + +`check-supervisor-schemas-versions` still fails locally on a commit +that changes the snapshot without touching anything under `versions/`. +CI's `--all-files` run passes, because it hands the hook the +`versions/` files too. ### Adding a field to a registered body From a2b0c2936d4978218bc06d75ca060f30f40c8cc3 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 15:01:40 +0000 Subject: [PATCH 29/94] Add a binding mode to task-handler declarations The parse-time check has to bind stub-task arguments the way the runtime does, and names compared in order fit none of the SDKs. Go flat params bind by position and have no names. A Go struct binds by name in any order, ignoring case and underscores except for arg-tagged fields, and an untagged struct can take one unmatched argument as the whole value. Java's TaskArgs has names but binds by position; Java's TaskInput and TypeScript bind by name. TaskHandlerDeclaration gains a required binding (positional, named or named_or_whole), TaskHandlerParam.name becomes nullable for nameless positional params, and exact_name marks a name that matches only as spelled. The bodies are still unreleased and first introduced in 2026-10-30, so AddTaskHandlerParseMessages still covers them, and check-supervisor-schemas-versions was skipped: it fails on any snapshot change that touches nothing under versions/. --- .../src/airflow/dag_processing/processor.py | 18 ++++++++- .../unit/dag_processing/test_processor.py | 38 ++++++++++++++++++- .../sdk/execution_time/schema/schema.json | 26 ++++++++++++- 3 files changed, 77 insertions(+), 5 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/processor.py b/airflow-core/src/airflow/dag_processing/processor.py index 242c051e475a0..6afffe89427a2 100644 --- a/airflow-core/src/airflow/dag_processing/processor.py +++ b/airflow-core/src/airflow/dag_processing/processor.py @@ -139,7 +139,8 @@ class DagFileParsingResult(BaseModel): class TaskHandlerParam(BaseModel): """One parameter of a task handler.""" - name: str + name: str | None + """``None`` when the runtime has no name for this positional parameter.""" value_schema: ArgValueSchema | None = None """JSON Schema of the values the parameter accepts; ``None`` when the handler does not constrain it.""" @@ -147,14 +148,27 @@ class TaskHandlerParam(BaseModel): required: bool """Whether the handler declares no default for the parameter.""" + exact_name: bool = False + """Whether ``name`` matches only as spelled, not case-insensitively with underscores ignored.""" + class TaskHandlerDeclaration(BaseModel): """A task handler that a Lang-SDK artifact registers for one task.""" task_id: str + binding: Literal["positional", "named", "named_or_whole"] + """ + How stub-task arguments bind to ``params``. + + - ``positional``: by position; names are informative only. + - ``named``: by name in any order, case-insensitively with underscores ignored unless ``exact_name`` is set. + - ``named_or_whole``: as ``named``, except that when there is exactly one argument and it matches no + parameter, it is decoded as the whole value. + """ + params: list[TaskHandlerParam] - """In declaration order, because stub-task arguments bind to them by position.""" + """In declaration order; the order is significant only for ``positional`` binding.""" class TaskHandlerParseRequest(BaseModel): diff --git a/airflow-core/tests/unit/dag_processing/test_processor.py b/airflow-core/tests/unit/dag_processing/test_processor.py index c66b6a2ed686a..4cafb39b695fe 100644 --- a/airflow-core/tests/unit/dag_processing/test_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_processor.py @@ -33,7 +33,7 @@ import pytest import structlog -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError from sqlalchemy import select from structlog.typing import FilteringBoundLogger @@ -58,6 +58,7 @@ DagFileParseRequest, DagFileParsingResult, DagFileProcessorProcess, + TaskHandlerDeclaration, TaskHandlerParseRequest, TaskHandlerParsingResult, ToDagProcessor, @@ -2366,6 +2367,40 @@ def test_parse_request_unions_differ_only_in_request(self): assert task_handler_types - dag_processor_types == {TaskHandlerParseRequest} +class TestTaskHandlerDeclaration: + @pytest.mark.parametrize( + ("binding", "param", "expected_name", "expected_exact_name"), + [ + pytest.param("positional", {"name": None, "required": True}, None, False, id="positional"), + pytest.param( + "named", {"name": "day", "required": False, "exact_name": True}, "day", True, id="named" + ), + pytest.param( + "named_or_whole", {"name": "Day", "required": False}, "Day", False, id="named_or_whole" + ), + ], + ) + def test_decodes_binding(self, binding, param, expected_name, expected_exact_name): + declaration = TaskHandlerDeclaration.model_validate( + {"task_id": "extract", "binding": binding, "params": [param]} + ) + + assert declaration.binding == binding + assert declaration.params[0].name == expected_name + assert declaration.params[0].exact_name is expected_exact_name + + @pytest.mark.parametrize( + "declaration", + [ + pytest.param({"task_id": "extract", "params": []}, id="missing"), + pytest.param({"task_id": "extract", "binding": "keyword", "params": []}, id="unknown"), + ], + ) + def test_rejects_invalid_binding(self, declaration): + with pytest.raises(ValidationError, match="binding"): + TaskHandlerDeclaration.model_validate(declaration) + + class TestDagFileProcessorProcess: def test_registered_message_types(self): expected = set(typing.get_args(typing.get_args(ToManager)[0])) - {TaskHandlerParsingResult} @@ -2424,6 +2459,7 @@ def test_rejects_task_handler_parsing_result(self, send_msg, proc): "etl": [ { "task_id": "extract", + "binding": "positional", "params": [ {"name": "day", "value_schema": {"type": "string"}, "required": True}, {"name": "limit", "required": False}, diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index 45f53e5e4675d..055c154b09c56 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -4163,6 +4163,15 @@ "title": "Task Id", "type": "string" }, + "binding": { + "enum": [ + "positional", + "named", + "named_or_whole" + ], + "title": "Binding", + "type": "string" + }, "params": { "items": { "$ref": "#/$defs/TaskHandlerParam" @@ -4173,6 +4182,7 @@ }, "required": [ "task_id", + "binding", "params" ], "title": "TaskHandlerDeclaration", @@ -4182,8 +4192,15 @@ "description": "One parameter of a task handler.", "properties": { "name": { - "title": "Name", - "type": "string" + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Name" }, "value_schema": { "anyOf": [ @@ -4199,6 +4216,11 @@ "required": { "title": "Required", "type": "boolean" + }, + "exact_name": { + "default": false, + "title": "Exact Name", + "type": "boolean" } }, "required": [ From f3601d6e8a0395d5ace49113b91f0dfcef7f2b2e Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 15:02:36 +0000 Subject: [PATCH 30/94] Regenerate TS SDK supervisor types for the binding mode check-ts-sdk-supervisor-schema regenerates supervisor.ts from the supervisor schema snapshot and fails on any difference, so the snapshot's new binding, nullable name and exact_name fields have to land in the generated types too. --- ts-sdk/src/generated/supervisor.ts | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index 1cd4be4f46046..7888f13942a60 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -596,8 +596,10 @@ export type Breadcrumbs = { }[]; export type Type76 = "TaskBreadcrumbsResult"; export type TaskId10 = string; -export type Name16 = string; +export type Binding = "positional" | "named" | "named_or_whole"; +export type Name16 = string | null; export type Required = boolean; +export type ExactName = boolean; export type Params2 = TaskHandlerParam[]; export type File1 = string; export type DagIds = string[]; @@ -1847,6 +1849,7 @@ export interface TaskBreadcrumbsResult { */ export interface TaskHandlerDeclaration { task_id: TaskId10; + binding: Binding; params: Params2; } /** @@ -1859,6 +1862,7 @@ export interface TaskHandlerParam { name: Name16; value_schema?: ArgValueSchema | null; required: Required; + exact_name?: ExactName; } /** * Request for Task Handler Parsing. From 7160a104c3da1791bd44dfe5314db4d66096b499 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 15:03:33 +0000 Subject: [PATCH 31/94] Describe the task-handler binding mode in ADR-0012 ADR-0012 said handler params are compared by name in order because both sides bind positionally, but only the stub side does. The message shapes and Appendix B now show the binding mode, nullable names and exact_name, and what each mode means for matching. --- .../lang-sdk/0012-lang-sdk-parse-protocol.md | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md b/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md index e144b87445b58..7d548de464f32 100644 --- a/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md +++ b/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md @@ -96,12 +96,14 @@ class TaskHandlerParsingResult: class TaskHandlerDeclaration: task_id: str - params: list[TaskHandlerParam] # ordered — arg_bindings are positional + binding: Literal["positional", "named", "named_or_whole"] # how stub-task arguments bind to params + params: list[TaskHandlerParam] # ordered; the order matters only for "positional" class TaskHandlerParam: - name: str + name: str | None # None: the runtime has no name for this positional parameter value_schema: ArgValueSchema | None = None required: bool # the handler declares no default + exact_name: bool = False # match as spelled, not case-insensitively with underscores ignored ``` One request carries every `dag_id` that resolved to the same artifact under the same coordinator, so a file whose stubs all target one runtime costs one process. A `dag_id` the @@ -109,7 +111,7 @@ artifact registers nothing for is **omitted** from `task_handlers` rather than r coordinators that has to cover the stubs ([ADR-0011](0011-mixed-language-dag-processing.md)). `value_schema` reuses the `ArgValueSchema` definition `arg_bindings` already carries ([ADR-0007](0007-taskflow-across-language-boundary.md)), so both sides of a comparison are the -same type. Two properties matter to validation: the field is nullable on both sides, and `params` is ordered. Appendix B says what that forces. +same type. Two properties matter to validation: the field is nullable on both sides, and each declaration names its binding mode. Appendix B says what that forces. `task_handlers` is the counterpart to `DagFileParsingResult.serialized_dags`, but fully typed. `serialized_dags` is `list[LazyDeserializedDAG]`, which is an opaque object in the schema snapshot. A handler declaration carries no Dag, so it code-generates and schema-validates in every SDK, and nothing on this path needs a DagSerialization implementation. @@ -208,13 +210,18 @@ bodies that never differ. Reusing `ToManager` leaves one reply union with one ne A boolean on `DagFileParseRequest` was the other alternative. It cannot work: a `DagRef` and a `TaskHandlerRef` are different payloads, not two subsets of one, so the flag would select between shapes the result type cannot both hold. -### Appendix B — What the nullable, ordered parameter list forces +### Appendix B — What the nullable schema and the binding mode force `value_schema` is nullable on both sides. An unannotated `@task.stub` parameter produces `value_schema: null` today, so validation compares schemas only where neither side is null, and falls back to name-and-arity otherwise. A strict comparison would turn every untyped stub argument into a parse error. -`params` is ordered because `LiteralArgBinding` and `XComArgBinding` are each documented as "one positional stub-task argument". Position is part of the contract, not incidental, -and both sides bind positionally. +The stub side always has names and positions: `LiteralArgBinding` and `XComArgBinding` are each documented as "one positional stub-task argument". The handler side binds the way +its runtime does, so each declaration names its `binding` and the check follows it: + +- `positional`: by position. Names are informative only, and absent where the runtime has none (Go flat params). Java's `TaskArgs` binds this way although it has names. +- `named`: by name in any order, case-insensitively with underscores ignored unless `exact_name` is set (Go `arg:` tags, explicit Java names). A Go struct with `arg:` tags, Java's + `TaskInput` and TypeScript bind this way. +- `named_or_whole`: as `named`, except that when there is exactly one argument and it matches no parameter, it is decoded as the whole value. An untagged Go struct binds this way. A declaration carries no class, method, or source location. [ADR-0006](0006-no-lang-sdk-source-display.md) rules out Lang-SDK source display, and putting it on the wire would invite a consumer to render it. It carries no `dag_id` either — the `task_handlers` key supplies it, so a declaration cannot disagree with the bucket it arrived in. From de96d42d5368f885f84df68b2c220f0111a5855b Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 15:06:11 +0000 Subject: [PATCH 32/94] Describe the binding-mode parse check in ADR-0011 and ADR-0013 Both ADRs still said the parse-time check compares stub and handler names in order because both sides bind positionally. ADR-0012 now gives each declaration a binding, so their comparison diagrams, the ADR-0011 mismatch consequence and Appendix B, and the ADR-0013 declaration sketch and validation step say the check binds by position or by folded or exact name, as that binding says. --- .../adr/lang-sdk/0011-mixed-language-dag-processing.md | 10 ++++++---- .../lang-sdk/0013-persisted-task-handler-bindings.md | 9 ++++++--- 2 files changed, 12 insertions(+), 7 deletions(-) diff --git a/airflow-core/adr/lang-sdk/0011-mixed-language-dag-processing.md b/airflow-core/adr/lang-sdk/0011-mixed-language-dag-processing.md index c2acaadbedf2a..b44d2a2074820 100644 --- a/airflow-core/adr/lang-sdk/0011-mixed-language-dag-processing.md +++ b/airflow-core/adr/lang-sdk/0011-mixed-language-dag-processing.md @@ -162,7 +162,8 @@ DagFileProcessorProcess(etl.py) ← manager spawn │ │ Python Dag "etl" (stub tasks) TaskHandlerDeclaration │ │ │ ────────────────────────────── ────────────────────────────── │ │ │ task_id ↔ task_id (sets must match) │ - │ │ arg_bindings[*].name ↔ params[*].name (in order) │ + │ │ arg_bindings[*] ↔ params[*] (per binding) │ + │ │ by position, or by folded or exact name │ │ │ arg_bindings[*].value_schema ↔ params[*].value_schema │ │ │ compared only where neither side is null │ │ │ │ @@ -197,8 +198,8 @@ There is no fourth row. A `TaskHandlerRef` has no Dag, so no `DagImporter` — a - No importer knows about coordinators. `PythonDagImporter` is unchanged by this ADR; the stub-to-handler comparison sits in `_parse_file`, above every importer. - A mixed-language `dag_id` never appears in Dag processing results. No `Dag` registration exists for a `dag_id` a Python file already owns, so everything downstream sees exactly one record per `dag_id`, with no flag to interpret. -- Stub/implementation mismatches — missing handler, extra handler, parameter name or order, incompatible schema — surface as import errors against the Python file at parse time, - alongside the errors the parse already reports. An unannotated stub argument is checked by name and position only. +- Stub/implementation mismatches — missing handler, extra handler, an argument that does not bind under the declaration's `binding`, incompatible schema — surface as import errors + against the Python file at parse time, alongside the errors the parse already reports. An unannotated stub argument is checked only for how it binds, by position or by name. - The Python Dag and Lang-SDK artifact can live in different DagBundles. - A single Dag can have stubs targeting different queues, some Java, some Go. Each resolves to its own coordinator instance, and validation unions their declarations per `dag_id` before comparing task ids. @@ -237,7 +238,8 @@ with one `bundle.serve()`. One artifact can hold both, so neither the file nor t Handler declarations from every process spawned in Step 4 are unioned per `dag_id` before comparison, since one Dag's stubs can target several queues. - `task_id` sets must match exactly. A missing or extra handler is an error. -- `arg_bindings[*].name` against `params[*].name`, in order — both sides bind positionally. +- `arg_bindings` against `params` as each declaration's `binding` says: by position for `positional`, by name for `named` and `named_or_whole`, compared case-insensitively with + underscores ignored unless `exact_name` is set ([ADR-0012](0012-lang-sdk-parse-protocol.md) Appendix B). - `arg_bindings[*].value_schema` against `params[*].value_schema`, compared only where neither side is null. An unannotated `@task.stub` parameter produces `null` today, so a strict comparison would make every untyped stub argument a parse error. diff --git a/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md b/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md index d62927452aba0..b3813b2858c89 100644 --- a/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md +++ b/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md @@ -259,13 +259,15 @@ class SDKTaskHandlerParsingResult(BaseModel): # runtime -> parent, on ToManager class TaskHandlerDeclaration(BaseModel): task_id: str - params: list[TaskHandlerParam] # ordered; arg bindings are positional + binding: Literal["positional", "named", "named_or_whole"] # how stub-task arguments bind to params + params: list[TaskHandlerParam] # ordered; the order matters only for "positional" class TaskHandlerParam(BaseModel): - name: str + name: str | None # None: the runtime has no name for this positional parameter value_schema: JSONSchema | None = None required: bool # the handler declares no default + exact_name: bool = False # match as spelled, not case-insensitively with underscores ignored ``` A `dag_id` the artifact registers nothing for is **omitted** from `task_handlers` rather than returned @@ -399,7 +401,8 @@ DagFileProcessorProcess(etl.py) [no DB — client con │ ├─7─ VALIDATE per dag_id, unioned across coordinators │ task_id sets must match exactly - │ arg_bindings[*].name ↔ handler_params[*].name, in order + │ arg_bindings[*] ↔ handler_params[*], per the declaration's binding: + │ by position, or by folded or exact name │ arg_bindings[*].schema ↔ handler_params[*].value_schema, │ compared only where neither is null │ two candidates claiming one (dag_id, task_id) → import error From 06386fb2291dcff82287c5f811f44debf1ae8999 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 15:45:06 +0000 Subject: [PATCH 33/94] Add a named_open binding for non-exhaustive task-handler params The TypeScript SDK erases types, so a TS handler knows only its explicit withArgNames renames, never its full parameter list. Under named binding the parse-time check would flag every stub argument the declaration does not list. named_open binds by name like named, but only the listed names are checked, and an argument none of them names is not a mismatch. The bodies are still unreleased and first introduced in 2026-10-30, so AddTaskHandlerParseMessages still covers them, and check-supervisor-schemas-versions was skipped: it fails on any snapshot change that touches nothing under versions/. --- airflow-core/src/airflow/dag_processing/processor.py | 4 +++- airflow-core/tests/unit/dag_processing/test_processor.py | 1 + task-sdk/src/airflow/sdk/execution_time/schema/schema.json | 3 ++- 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/processor.py b/airflow-core/src/airflow/dag_processing/processor.py index 6afffe89427a2..75517ae378bc4 100644 --- a/airflow-core/src/airflow/dag_processing/processor.py +++ b/airflow-core/src/airflow/dag_processing/processor.py @@ -157,7 +157,7 @@ class TaskHandlerDeclaration(BaseModel): task_id: str - binding: Literal["positional", "named", "named_or_whole"] + binding: Literal["positional", "named", "named_or_whole", "named_open"] """ How stub-task arguments bind to ``params``. @@ -165,6 +165,8 @@ class TaskHandlerDeclaration(BaseModel): - ``named``: by name in any order, case-insensitively with underscores ignored unless ``exact_name`` is set. - ``named_or_whole``: as ``named``, except that when there is exactly one argument and it matches no parameter, it is decoded as the whole value. + - ``named_open``: as ``named``, but ``params`` need not list every parameter the handler takes. Only the + listed names are checked, and an argument that none of them names is not a mismatch. """ params: list[TaskHandlerParam] diff --git a/airflow-core/tests/unit/dag_processing/test_processor.py b/airflow-core/tests/unit/dag_processing/test_processor.py index 4cafb39b695fe..832e4c0c2e5f4 100644 --- a/airflow-core/tests/unit/dag_processing/test_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_processor.py @@ -2378,6 +2378,7 @@ class TestTaskHandlerDeclaration: pytest.param( "named_or_whole", {"name": "Day", "required": False}, "Day", False, id="named_or_whole" ), + pytest.param("named_open", {"name": "day", "required": False}, "day", False, id="named_open"), ], ) def test_decodes_binding(self, binding, param, expected_name, expected_exact_name): diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index 055c154b09c56..30f196e693cc4 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -4167,7 +4167,8 @@ "enum": [ "positional", "named", - "named_or_whole" + "named_or_whole", + "named_open" ], "title": "Binding", "type": "string" From 12b142d0ab9d8b1b2d7a6422ea7272efa6aec1d5 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 15:45:32 +0000 Subject: [PATCH 34/94] Regenerate TS SDK supervisor types for the named_open binding check-ts-sdk-supervisor-schema regenerates supervisor.ts from the supervisor schema snapshot and fails on any difference, so the snapshot's new named_open binding value has to land in the generated Binding type too. --- ts-sdk/src/generated/supervisor.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index 7888f13942a60..c34331ccbb056 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -596,7 +596,7 @@ export type Breadcrumbs = { }[]; export type Type76 = "TaskBreadcrumbsResult"; export type TaskId10 = string; -export type Binding = "positional" | "named" | "named_or_whole"; +export type Binding = "positional" | "named" | "named_or_whole" | "named_open"; export type Name16 = string | null; export type Required = boolean; export type ExactName = boolean; From 270011518de5d2eb2be3a6da193cad402613f8f6 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 15:47:15 +0000 Subject: [PATCH 35/94] Add the named_open binding to the Lang-SDK ADRs The declarations gained a fourth binding for handlers whose params list is not exhaustive. ADR-0012's message shapes and Appendix B list it and say TypeScript binds this way, and the binding lists in ADR-0011's Appendix B and ADR-0013's declaration sketch include it. --- .../adr/lang-sdk/0011-mixed-language-dag-processing.md | 4 ++-- airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md | 8 +++++--- .../adr/lang-sdk/0013-persisted-task-handler-bindings.md | 2 +- 3 files changed, 8 insertions(+), 6 deletions(-) diff --git a/airflow-core/adr/lang-sdk/0011-mixed-language-dag-processing.md b/airflow-core/adr/lang-sdk/0011-mixed-language-dag-processing.md index b44d2a2074820..c1cc9795473f0 100644 --- a/airflow-core/adr/lang-sdk/0011-mixed-language-dag-processing.md +++ b/airflow-core/adr/lang-sdk/0011-mixed-language-dag-processing.md @@ -238,8 +238,8 @@ with one `bundle.serve()`. One artifact can hold both, so neither the file nor t Handler declarations from every process spawned in Step 4 are unioned per `dag_id` before comparison, since one Dag's stubs can target several queues. - `task_id` sets must match exactly. A missing or extra handler is an error. -- `arg_bindings` against `params` as each declaration's `binding` says: by position for `positional`, by name for `named` and `named_or_whole`, compared case-insensitively with - underscores ignored unless `exact_name` is set ([ADR-0012](0012-lang-sdk-parse-protocol.md) Appendix B). +- `arg_bindings` against `params` as each declaration's `binding` says: by position for `positional`, by name for `named`, `named_or_whole` and `named_open`, compared + case-insensitively with underscores ignored unless `exact_name` is set. Under `named_open` only the listed names are checked ([ADR-0012](0012-lang-sdk-parse-protocol.md) Appendix B). - `arg_bindings[*].value_schema` against `params[*].value_schema`, compared only where neither side is null. An unannotated `@task.stub` parameter produces `null` today, so a strict comparison would make every untyped stub argument a parse error. diff --git a/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md b/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md index 7d548de464f32..e655703dce9f2 100644 --- a/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md +++ b/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md @@ -96,7 +96,7 @@ class TaskHandlerParsingResult: class TaskHandlerDeclaration: task_id: str - binding: Literal["positional", "named", "named_or_whole"] # how stub-task arguments bind to params + binding: Literal["positional", "named", "named_or_whole", "named_open"] # how stub-task arguments bind to params params: list[TaskHandlerParam] # ordered; the order matters only for "positional" class TaskHandlerParam: @@ -219,9 +219,11 @@ The stub side always has names and positions: `LiteralArgBinding` and `XComArgBi its runtime does, so each declaration names its `binding` and the check follows it: - `positional`: by position. Names are informative only, and absent where the runtime has none (Go flat params). Java's `TaskArgs` binds this way although it has names. -- `named`: by name in any order, case-insensitively with underscores ignored unless `exact_name` is set (Go `arg:` tags, explicit Java names). A Go struct with `arg:` tags, Java's - `TaskInput` and TypeScript bind this way. +- `named`: by name in any order, case-insensitively with underscores ignored unless `exact_name` is set (Go `arg:` tags, explicit Java names). A Go struct with `arg:` tags and + Java's `TaskInput` bind this way. - `named_or_whole`: as `named`, except that when there is exactly one argument and it matches no parameter, it is decoded as the whole value. An untagged Go struct binds this way. +- `named_open`: as `named`, but `params` need not list every parameter the handler takes. Only the listed names are checked, and an argument that none of them names is not a + mismatch. TypeScript binds this way: its types are erased, so a handler knows only its explicit `withArgNames` renames. A declaration carries no class, method, or source location. [ADR-0006](0006-no-lang-sdk-source-display.md) rules out Lang-SDK source display, and putting it on the wire would invite a consumer to render it. It carries no `dag_id` either — the `task_handlers` key supplies it, so a declaration cannot disagree with the bucket it arrived in. diff --git a/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md b/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md index b3813b2858c89..e778600776f77 100644 --- a/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md +++ b/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md @@ -259,7 +259,7 @@ class SDKTaskHandlerParsingResult(BaseModel): # runtime -> parent, on ToManager class TaskHandlerDeclaration(BaseModel): task_id: str - binding: Literal["positional", "named", "named_or_whole"] # how stub-task arguments bind to params + binding: Literal["positional", "named", "named_or_whole", "named_open"] # how arguments bind to params params: list[TaskHandlerParam] # ordered; the order matters only for "positional" From d33f9041f2241b4260cfabae691ce13ef2225332 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 1 Oct 2026 07:18:25 +0000 Subject: [PATCH 36/94] Ask a task handler runtime for every handler it registers A probe answered only for the Dag ids it was sent, so its answer depended on which file asked and could not be cached under the artifact's fingerprint alone. Without dag_ids the answer depends only on the artifact. TaskHandlerParseRequest is new in the unreleased schema version, so dropping the field needs no VersionChange. --- airflow-core/src/airflow/dag_processing/processor.py | 10 +++++----- .../airflow/sdk/execution_time/schema/schema.json | 12 ++---------- 2 files changed, 7 insertions(+), 15 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/processor.py b/airflow-core/src/airflow/dag_processing/processor.py index 75517ae378bc4..dfd680872b3a8 100644 --- a/airflow-core/src/airflow/dag_processing/processor.py +++ b/airflow-core/src/airflow/dag_processing/processor.py @@ -177,14 +177,12 @@ class TaskHandlerParseRequest(BaseModel): """ Request for Task Handler Parsing. - Asks a Lang-SDK runtime which task handlers an artifact registers for the given Dags. + Asks a Lang-SDK runtime for every task handler an artifact registers. """ file: str """The artifact to ask.""" - dag_ids: list[str] - bundle_path: Path bundle_name: str @@ -196,13 +194,15 @@ class TaskHandlerParsingResult(BaseModel): """ Result of Task Handler Parsing. - The task handlers a Lang-SDK artifact registers, keyed by Dag id. + Every task handler a Lang-SDK artifact registers, keyed by Dag id. + + The answer depends only on the artifact, never on the request. """ fileloc: str task_handlers: dict[str, list[TaskHandlerDeclaration]] - """A requested Dag id the artifact registers no task handler for is omitted, not mapped to ``[]``.""" + """Every Dag id the artifact registers a task handler for; ``{}`` when it registers none.""" import_errors: dict[str, str] | None = None warnings: list | None = None diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index 30f196e693cc4..6ca8e046cbdd3 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -4232,19 +4232,12 @@ "type": "object" }, "TaskHandlerParseRequest": { - "description": "Request for Task Handler Parsing.\n\nAsks a Lang-SDK runtime which task handlers an artifact registers for the given Dags.", + "description": "Request for Task Handler Parsing.\n\nAsks a Lang-SDK runtime for every task handler an artifact registers.", "properties": { "file": { "title": "File", "type": "string" }, - "dag_ids": { - "items": { - "type": "string" - }, - "title": "Dag Ids", - "type": "array" - }, "bundle_path": { "format": "path", "title": "Bundle Path", @@ -4263,7 +4256,6 @@ }, "required": [ "file", - "dag_ids", "bundle_path", "bundle_name" ], @@ -4271,7 +4263,7 @@ "type": "object" }, "TaskHandlerParsingResult": { - "description": "Result of Task Handler Parsing.\n\nThe task handlers a Lang-SDK artifact registers, keyed by Dag id.", + "description": "Result of Task Handler Parsing.\n\nEvery task handler a Lang-SDK artifact registers, keyed by Dag id.\n\nThe answer depends only on the artifact, never on the request.", "properties": { "fileloc": { "title": "Fileloc", From d86cead24b2002381a9b713d241fc8e1771e9e80 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 1 Oct 2026 07:18:43 +0000 Subject: [PATCH 37/94] Regenerate TS SDK supervisor types for the probe without dag_ids check-ts-sdk-supervisor-schema regenerates supervisor.ts from the supervisor schema snapshot and fails on any difference. CI runs it with --all-files whenever a ts-sdk file changes, so a stale copy would fail the next unrelated ts-sdk change. --- ts-sdk/src/generated/supervisor.ts | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index c34331ccbb056..bec0a5c37af3d 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -602,7 +602,6 @@ export type Required = boolean; export type ExactName = boolean; export type Params2 = TaskHandlerParam[]; export type File1 = string; -export type DagIds = string[]; export type BundlePath1 = string; export type BundleName5 = string; export type Type77 = "TaskHandlerParseRequest"; @@ -1867,14 +1866,13 @@ export interface TaskHandlerParam { /** * Request for Task Handler Parsing. * - * Asks a Lang-SDK runtime which task handlers an artifact registers for the given Dags. + * Asks a Lang-SDK runtime for every task handler an artifact registers. * * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema * via the `definition` "TaskHandlerParseRequest". */ export interface TaskHandlerParseRequest { file: File1; - dag_ids: DagIds; bundle_path: BundlePath1; bundle_name: BundleName5; type?: Type77; @@ -1882,7 +1880,9 @@ export interface TaskHandlerParseRequest { /** * Result of Task Handler Parsing. * - * The task handlers a Lang-SDK artifact registers, keyed by Dag id. + * Every task handler a Lang-SDK artifact registers, keyed by Dag id. + * + * The answer depends only on the artifact, never on the request. * * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema * via the `definition` "TaskHandlerParsingResult". From c9eaf687a1e9bd9899ef70707b37cf93319ddb7f Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 1 Oct 2026 07:18:51 +0000 Subject: [PATCH 38/94] Describe the all-handlers probe in ADR-0011 and ADR-0012 Both ADRs still had the request name the file's Dags and the runtime omit the rest. SDK authors and later layers read them as the spec for what a runtime answers, so they must say the answer depends only on the artifact. --- .../0011-mixed-language-dag-processing.md | 20 +++++++++---------- .../lang-sdk/0012-lang-sdk-parse-protocol.md | 12 +++++------ 2 files changed, 16 insertions(+), 16 deletions(-) diff --git a/airflow-core/adr/lang-sdk/0011-mixed-language-dag-processing.md b/airflow-core/adr/lang-sdk/0011-mixed-language-dag-processing.md index c1cc9795473f0..463e5d6960a2d 100644 --- a/airflow-core/adr/lang-sdk/0011-mixed-language-dag-processing.md +++ b/airflow-core/adr/lang-sdk/0011-mixed-language-dag-processing.md @@ -90,7 +90,7 @@ DagFileProcessorProcess(etl.py) ← manager spawn ├── _serialize_dags(bag) → is_stub tasks carry arg_bindings (ADR-0007) │ │ ┌─────────────────────────────────────────────────────────────────────┐ - │ │ Step 1: Collect the dag_ids to ask about — every Dag in this │ + │ │ Step 1: Collect the Dags to validate — every Dag in this │ │ │ file with at least one is_stub task → ["etl"] │ │ └─────────────────────────────────────────────────────────────────────┘ │ @@ -123,7 +123,7 @@ DagFileProcessorProcess(etl.py) ← manager spawn │ │ └── BundleScanner.scanBundles(roots) │ │ │ → "etl" → ResolvedBundle(analytics.jar, mainClass, ...) │ │ │ │ - │ │ Group the dag_ids by (coordinator, artifact) — one group, one │ + │ │ Group the stub tasks by (coordinator, artifact) — one group, one │ │ │ process, one request │ │ └─────────────────────────────────────────────────────────────────────┘ │ @@ -138,12 +138,11 @@ DagFileProcessorProcess(etl.py) ← manager spawn │ │ ├── in the child: _build_parse_task_handler_command() │ │ │ │ coordinator.parse_task_handler() — spawn JVM │ │ │ │ │ - │ │ │ ──TaskHandlerParseRequest(file=analytics.jar, │ - │ │ │ dag_ids=["etl"])─────▶ JVM │ + │ │ │ ──TaskHandlerParseRequest(file=analytics.jar)─────▶ JVM │ │ │ │ (ToSDKTaskHandlerProcessor) │ │ │ │ │ - │ │ │ JVM answers from its own TaskHandler registrations │ - │ │ │ whose dagId is one of the requested ids │ + │ │ │ JVM answers with every TaskHandler registration │ + │ │ │ in the artifact, or {} when there is none │ │ │ │ │ │ │ │ ◀─TaskHandlerParsingResult(task_handlers={ │ │ │ │ "etl": [extract, transform, load]})─────── JVM │ @@ -176,8 +175,8 @@ DagFileProcessorProcess(etl.py) ← manager spawn ``` The parse owns validation, not an importer. `PythonDagImporter` returns `airflow.sdk.DAG` objects and knows nothing about coordinators or queues, so `@task.stub` keeps working for -any importer that can produce a Dag carrying stub tasks. `_parse_file` is also the only place where the whole file's Dags are visible at once, which is what lets one request cover -every `dag_id` that resolved to the same artifact ([ADR-0012](0012-lang-sdk-parse-protocol.md)). +any importer that can produce a Dag carrying stub tasks. `_parse_file` is also the only place where the whole file's Dags are visible at once, which is what lets one request per +artifact serve every Dag in the file whose stubs resolved to it ([ADR-0012](0012-lang-sdk-parse-protocol.md)). Resolution goes through the coordinator registry, not the filesystem, so the Python Dag and the Lang-SDK artifact **do not need to be in the same DagBundle**. Nothing here needs an `airflow.sdk.DAG` round-trip either — validation compares against the Dag the Python parser already built. Appendix B states exactly what is compared. @@ -187,7 +186,7 @@ Resolution goes through the coordinator registry, not the filesystem, so the Pyt | Caller | Coordinator call | What comes back | Action | |--------------------------------------------|------------------------------------------------------|------------------------------------------|--------------------------------------------------| | `_parse_file` → `PythonDagImporter` | — (the Python file is parsed in process) | its own parsed Dags | PERSIST | -| `_parse_file`, per (coordinator, artifact) | `parse_task_handler`, scoped to that group's dag_ids | `TaskHandlerParsingResult` | VALIDATE only — not a Dag, so nothing to persist | +| `_parse_file`, per (coordinator, artifact) | `parse_task_handler`, for every handler it registers | `TaskHandlerParsingResult` | VALIDATE only — not a Dag, so nothing to persist | | `_parse_file` → `JavaDagImporter` | `parse_dag` | `DagFileParsingResult`, native Dags only | PERSIST | There is no fourth row. A `TaskHandlerRef` has no Dag, so no `DagImporter` — and nothing reading a `DagImporter`'s results — ever sees one. @@ -235,7 +234,8 @@ with one `bundle.serve()`. One artifact can hold both, so neither the file nor t ### Appendix B — What is compared -Handler declarations from every process spawned in Step 4 are unioned per `dag_id` before comparison, since one Dag's stubs can target several queues. +Handler declarations from every process spawned in Step 4 are unioned per `dag_id` before comparison, since one Dag's stubs can target several queues. Each reply covers +every Dag its artifact registers handlers for; only the parsed file's Dags are compared. - `task_id` sets must match exactly. A missing or extra handler is an error. - `arg_bindings` against `params` as each declaration's `binding` says: by position for `positional`, by name for `named`, `named_or_whole` and `named_open`, compared diff --git a/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md b/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md index e655703dce9f2..afda6794a794d 100644 --- a/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md +++ b/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md @@ -26,8 +26,8 @@ Proposed ## Context The Dag processor asks a Lang-SDK runtime two different questions. "Which Dags does this artifact define?" is answered over the messages [ADR-0004](0004-dag-parsing.md) already -defines. "Which task handlers does this artifact register for a `dag_id` Python already owns?" has no answer in those messages, because a `TaskHandler` registration carries no Dag -([ADR-0011](0011-mixed-language-dag-processing.md)). +defines. "Which task handlers does this artifact register, each for a `dag_id` Python already owns?" has no answer in those messages, because a `TaskHandler` registration carries +no Dag ([ADR-0011](0011-mixed-language-dag-processing.md)). This ADR defines the request that carries the second question, the subprocess classes that carry both, and the two parse-side entry points on the coordinator. @@ -82,7 +82,6 @@ relay is only type-safe because both hops speak the same pair, which is the reas ``` class TaskHandlerParseRequest: file: str # the artifact resolved for this coordinator - dag_ids: list[str] # every Dag in the parsed file with stub tasks that resolved here bundle_path: Path bundle_name: str type: Literal["TaskHandlerParseRequest"] @@ -106,9 +105,10 @@ class TaskHandlerParam: exact_name: bool = False # match as spelled, not case-insensitively with underscores ignored ``` -One request carries every `dag_id` that resolved to the same artifact under the same coordinator, so a file whose stubs all target one runtime costs one process. A `dag_id` the -artifact registers nothing for is **omitted** from `task_handlers` rather than returned empty: the key set is not required to match `dag_ids`, because it is the union across -coordinators that has to cover the stubs ([ADR-0011](0011-mixed-language-dag-processing.md)). +The request names no Dags. The runtime answers with every task handler the artifact registers, keyed by `dag_id`, and with `{}` when it registers none. The answer must depend only +on the artifact, never on the request, so the Dag processor can cache it per artifact ([ADR-0013](0013-persisted-task-handler-bindings.md)). One request per (coordinator, artifact) +pair serves every Dag in the parsed file whose stubs resolved there, so a file whose stubs all target one runtime costs one process. The key set is not required to match the file's +Dags: it can hold other files' Dags, and it is the union across coordinators that has to cover the stubs ([ADR-0011](0011-mixed-language-dag-processing.md)). `value_schema` reuses the `ArgValueSchema` definition `arg_bindings` already carries ([ADR-0007](0007-taskflow-across-language-boundary.md)), so both sides of a comparison are the same type. Two properties matter to validation: the field is nullable on both sides, and each declaration names its binding mode. Appendix B says what that forces. From 1bd56509d092cb936d8a10803394fa56bbd99326 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 1 Oct 2026 17:28:15 +0000 Subject: [PATCH 39/94] Bind task handlers only by position or by name Four binding modes asked the parse-time check for more than the runtimes can state. TypeScript erases types, so named_open existed only to carry rename lists, and named_or_whole differed from named only for one unmatched argument, which the runtime already accepts as the whole value. Once a named mismatch is a warning, as it is at run time, positional and named cover every SDK. A runtime that cannot list a handler's params sends null instead, so only the handler's presence is checked. required goes too: a positional count is compared with and without the defaulted arguments, and an unfilled named param only warns, so nothing reads it. params gets its own title because go-jsonschema names the generated type of a nullable list by its title, and Params is already the HITL params map. The bodies are unreleased and new in 2026-10-30, so they need no VersionChange. check-supervisor-schemas-versions was skipped because it fails on any snapshot change that touches nothing under versions/. --- .../src/airflow/dag_processing/processor.py | 29 ++++++++++------ .../unit/dag_processing/test_processor.py | 34 ++++++++++++------- .../sdk/execution_time/schema/schema.json | 28 +++++++-------- 3 files changed, 54 insertions(+), 37 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/processor.py b/airflow-core/src/airflow/dag_processing/processor.py index dfd680872b3a8..b3b177f451447 100644 --- a/airflow-core/src/airflow/dag_processing/processor.py +++ b/airflow-core/src/airflow/dag_processing/processor.py @@ -136,6 +136,9 @@ class DagFileParsingResult(BaseModel): type: Literal["DagFileParsingResult"] = "DagFileParsingResult" +TaskHandlerBindingMode = Literal["positional", "named"] + + class TaskHandlerParam(BaseModel): """One parameter of a task handler.""" @@ -145,9 +148,6 @@ class TaskHandlerParam(BaseModel): value_schema: ArgValueSchema | None = None """JSON Schema of the values the parameter accepts; ``None`` when the handler does not constrain it.""" - required: bool - """Whether the handler declares no default for the parameter.""" - exact_name: bool = False """Whether ``name`` matches only as spelled, not case-insensitively with underscores ignored.""" @@ -157,20 +157,27 @@ class TaskHandlerDeclaration(BaseModel): task_id: str - binding: Literal["positional", "named", "named_or_whole", "named_open"] + binding: TaskHandlerBindingMode """ How stub-task arguments bind to ``params``. - - ``positional``: by position; names are informative only. + - ``positional``: by position; names are informative only. An argument count that matches ``params`` + neither with every argument nor after dropping the defaulted ones, or a value type the param does not + accept, makes the Dag fail to import. - ``named``: by name in any order, case-insensitively with underscores ignored unless ``exact_name`` is set. - - ``named_or_whole``: as ``named``, except that when there is exactly one argument and it matches no - parameter, it is decoded as the whole value. - - ``named_open``: as ``named``, but ``params`` need not list every parameter the handler takes. Only the - listed names are checked, and an argument that none of them names is not a mismatch. + An argument no param takes, or a param no argument fills, is logged as a warning, and the task still + runs. When no param matches and exactly one argument was passed, it may be the whole value and is not + warned about, unless ``params`` is empty, a param sets ``exact_name``, or the argument cannot be an + object. A value type a param does not accept makes the Dag fail to import. """ - params: list[TaskHandlerParam] - """In declaration order; the order is significant only for ``positional`` binding.""" + # The title keeps Go's generated type for this list apart from the HITL ``Params`` map. + params: Annotated[list[TaskHandlerParam] | None, Field(title="Task Handler Params")] + """ + In declaration order; the order is significant only for ``positional`` binding. + + ``None`` when the runtime cannot list the handler's parameters, so only the handler's presence is checked. + """ class TaskHandlerParseRequest(BaseModel): diff --git a/airflow-core/tests/unit/dag_processing/test_processor.py b/airflow-core/tests/unit/dag_processing/test_processor.py index 832e4c0c2e5f4..8887440ecb527 100644 --- a/airflow-core/tests/unit/dag_processing/test_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_processor.py @@ -59,6 +59,7 @@ DagFileParsingResult, DagFileProcessorProcess, TaskHandlerDeclaration, + TaskHandlerParam, TaskHandlerParseRequest, TaskHandlerParsingResult, ToDagProcessor, @@ -2369,32 +2370,41 @@ def test_parse_request_unions_differ_only_in_request(self): class TestTaskHandlerDeclaration: @pytest.mark.parametrize( - ("binding", "param", "expected_name", "expected_exact_name"), + ("binding", "param", "expected"), [ - pytest.param("positional", {"name": None, "required": True}, None, False, id="positional"), + pytest.param("positional", {"name": None}, TaskHandlerParam(name=None), id="positional"), pytest.param( - "named", {"name": "day", "required": False, "exact_name": True}, "day", True, id="named" + "named", + {"name": "day", "exact_name": True}, + TaskHandlerParam(name="day", exact_name=True), + id="named", ), - pytest.param( - "named_or_whole", {"name": "Day", "required": False}, "Day", False, id="named_or_whole" - ), - pytest.param("named_open", {"name": "day", "required": False}, "day", False, id="named_open"), ], ) - def test_decodes_binding(self, binding, param, expected_name, expected_exact_name): + def test_decodes_binding(self, binding, param, expected): declaration = TaskHandlerDeclaration.model_validate( {"task_id": "extract", "binding": binding, "params": [param]} ) assert declaration.binding == binding - assert declaration.params[0].name == expected_name - assert declaration.params[0].exact_name is expected_exact_name + assert declaration.params == [expected] + + def test_decodes_unlisted_params(self): + declaration = TaskHandlerDeclaration.model_validate( + {"task_id": "extract", "binding": "named", "params": None} + ) + + assert declaration.params is None @pytest.mark.parametrize( "declaration", [ pytest.param({"task_id": "extract", "params": []}, id="missing"), pytest.param({"task_id": "extract", "binding": "keyword", "params": []}, id="unknown"), + pytest.param( + {"task_id": "extract", "binding": "named_or_whole", "params": []}, id="named_or_whole" + ), + pytest.param({"task_id": "extract", "binding": "named_open", "params": []}, id="named_open"), ], ) def test_rejects_invalid_binding(self, declaration): @@ -2462,8 +2472,8 @@ def test_rejects_task_handler_parsing_result(self, send_msg, proc): "task_id": "extract", "binding": "positional", "params": [ - {"name": "day", "value_schema": {"type": "string"}, "required": True}, - {"name": "limit", "required": False}, + {"name": "day", "value_schema": {"type": "string"}}, + {"name": "limit"}, ], } ] diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index 6ca8e046cbdd3..98b9be109d26f 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -4166,19 +4166,24 @@ "binding": { "enum": [ "positional", - "named", - "named_or_whole", - "named_open" + "named" ], "title": "Binding", "type": "string" }, "params": { - "items": { - "$ref": "#/$defs/TaskHandlerParam" - }, - "title": "Params", - "type": "array" + "anyOf": [ + { + "items": { + "$ref": "#/$defs/TaskHandlerParam" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Task Handler Params" } }, "required": [ @@ -4214,10 +4219,6 @@ ], "default": null }, - "required": { - "title": "Required", - "type": "boolean" - }, "exact_name": { "default": false, "title": "Exact Name", @@ -4225,8 +4226,7 @@ } }, "required": [ - "name", - "required" + "name" ], "title": "TaskHandlerParam", "type": "object" From a4598c4102a607cf9bd18e5de3341166f00d52d5 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 1 Oct 2026 17:28:40 +0000 Subject: [PATCH 40/94] Regenerate TS SDK supervisor types for the two binding modes check-ts-sdk-supervisor-schema regenerates supervisor.ts from the supervisor schema snapshot and fails on any difference, so the generated Binding and TaskHandlerParam types must follow the snapshot's two modes, nullable params and dropped required. --- ts-sdk/src/generated/supervisor.ts | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index bec0a5c37af3d..4567d65a39835 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -596,11 +596,10 @@ export type Breadcrumbs = { }[]; export type Type76 = "TaskBreadcrumbsResult"; export type TaskId10 = string; -export type Binding = "positional" | "named" | "named_or_whole" | "named_open"; +export type Binding = "positional" | "named"; +export type TaskHandlerParams = TaskHandlerParam[] | null; export type Name16 = string | null; -export type Required = boolean; export type ExactName = boolean; -export type Params2 = TaskHandlerParam[]; export type File1 = string; export type BundlePath1 = string; export type BundleName5 = string; @@ -1849,7 +1848,7 @@ export interface TaskBreadcrumbsResult { export interface TaskHandlerDeclaration { task_id: TaskId10; binding: Binding; - params: Params2; + params: TaskHandlerParams; } /** * One parameter of a task handler. @@ -1860,7 +1859,6 @@ export interface TaskHandlerDeclaration { export interface TaskHandlerParam { name: Name16; value_schema?: ArgValueSchema | null; - required: Required; exact_name?: ExactName; } /** From 1d793f9974d1df14ce93accfedc33c894a19b3ff Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Fri, 2 Oct 2026 02:14:10 +0000 Subject: [PATCH 41/94] Regenerate Go SDK models for task-handler messages check-go-sdk-generated-drift regenerates the Go SDK's coordinator-protocol models from the supervisor schema snapshot and fails on any difference, so the models must follow the snapshot's task-handler messages and binding modes. go-jsonschema moves some type blocks, so the diff is larger than the new types. --- .../execution/genmodels/discriminators.gen.go | 8 + go-sdk/pkg/execution/genmodels/models.gen.go | 348 +++++++++++------- 2 files changed, 218 insertions(+), 138 deletions(-) diff --git a/go-sdk/pkg/execution/genmodels/discriminators.gen.go b/go-sdk/pkg/execution/genmodels/discriminators.gen.go index 09da3118f55da..0fe7a014ca505 100644 --- a/go-sdk/pkg/execution/genmodels/discriminators.gen.go +++ b/go-sdk/pkg/execution/genmodels/discriminators.gen.go @@ -100,6 +100,8 @@ const ( TypeTICount = "TICount" TypeTaskBreadcrumbsResult = "TaskBreadcrumbsResult" TypeTaskCallbackRequest = "TaskCallbackRequest" + TypeTaskHandlerParseRequest = "TaskHandlerParseRequest" + TypeTaskHandlerParsingResult = "TaskHandlerParsingResult" TypeTaskRescheduleStartDate = "TaskRescheduleStartDate" TypeTaskState = "TaskState" TypeTaskStateStoreResult = "TaskStateStoreResult" @@ -350,6 +352,12 @@ func EnsureType(m any) any { case TaskCallbackRequest: b.Type = TypeTaskCallbackRequest return b + case TaskHandlerParseRequest: + b.Type = TypeTaskHandlerParseRequest + return b + case TaskHandlerParsingResult: + b.Type = TypeTaskHandlerParsingResult + return b case TaskRescheduleStartDate: b.Type = TypeTaskRescheduleStartDate return b diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go b/go-sdk/pkg/execution/genmodels/models.gen.go index dee470082f124..0b0b4eed57d51 100644 --- a/go-sdk/pkg/execution/genmodels/models.gen.go +++ b/go-sdk/pkg/execution/genmodels/models.gen.go @@ -888,6 +888,86 @@ type GetAssetEventByAssetAlias struct { Type string `msgpack:"type,omitempty"` } +// Variable schema for responses with fields that are needed for Runtime. +type VariableResponse struct { + // Key corresponds to the JSON schema field "key". + Key string `msgpack:"key"` + + // Value corresponds to the JSON schema field "value". + Value interface{} `msgpack:"value"` +} + +type TaskArgBinding interface{} + +type TriggerKwargs map[string]interface{} + +type Warnings []interface{} + +// Lazily build information from the serialized DAG structure. +// +// An object that will present "enough" of the DAG like interface to update DAG db +// models etc, without having +// to deserialize the full DAG and Task hierarchy. +type LazyDeserializedDAG struct { + // Data corresponds to the JSON schema field "data". + Data Data `msgpack:"data"` + + // LastLoaded corresponds to the JSON schema field "last_loaded". + LastLoaded interface{} `msgpack:"last_loaded,omitempty"` +} + +type ImportErrors map[string]string + +// Schema for TaskInstance model with minimal required fields needed for Runtime. +type TaskInstance struct { + // ContextCarrier corresponds to the JSON schema field "context_carrier". + ContextCarrier *ContextCarrier `msgpack:"context_carrier,omitempty"` + + // DagID corresponds to the JSON schema field "dag_id". + DagID string `msgpack:"dag_id"` + + // DagVersionID corresponds to the JSON schema field "dag_version_id". + DagVersionID string `msgpack:"dag_version_id"` + + // Hostname corresponds to the JSON schema field "hostname". + Hostname interface{} `msgpack:"hostname,omitempty"` + + // ID corresponds to the JSON schema field "id". + ID string `msgpack:"id"` + + // MapIndex corresponds to the JSON schema field "map_index". + MapIndex *int `msgpack:"map_index,omitempty"` + + // Queue corresponds to the JSON schema field "queue". + Queue string `msgpack:"queue,omitempty"` + + // RunID corresponds to the JSON schema field "run_id". + RunID string `msgpack:"run_id"` + + // TaskID corresponds to the JSON schema field "task_id". + TaskID string `msgpack:"task_id"` + + // TryNumber corresponds to the JSON schema field "try_number". + TryNumber int `msgpack:"try_number"` +} + +type Params map[string]interface{} + +// Schema for a Human-in-the-loop users. +type HITLUser struct { + // ID corresponds to the JSON schema field "id". + ID string `msgpack:"id"` + + // Name corresponds to the JSON schema field "name". + Name string `msgpack:"name"` +} + +type VersionData map[string]interface{} + +type NextKwargs map[string]interface{} + +type JsonValue interface{} + type GetAssetStateStoreByName struct { // Key corresponds to the JSON schema field "key". Key string `msgpack:"key"` @@ -926,6 +1006,12 @@ type GetConnection struct { Type string `msgpack:"type,omitempty"` } +type LogicalDates []time.Time + +type RunIds []string + +type States []string + type GetDRCount struct { // DagID corresponds to the JSON schema field "dag_id". DagID string `msgpack:"dag_id"` @@ -1004,6 +1090,22 @@ type GetPreviousDagRun struct { Type string `msgpack:"type,omitempty"` } +type TaskInstanceState string + +const TaskInstanceStateRemoved TaskInstanceState = "removed" +const TaskInstanceStateScheduled TaskInstanceState = "scheduled" +const TaskInstanceStateQueued TaskInstanceState = "queued" +const TaskInstanceStateRunning TaskInstanceState = "running" +const TaskInstanceStateSuccess TaskInstanceState = "success" +const TaskInstanceStateRestarting TaskInstanceState = "restarting" +const TaskInstanceStateFailed TaskInstanceState = "failed" +const TaskInstanceStateUpForRetry TaskInstanceState = "up_for_retry" +const TaskInstanceStateUpForReschedule TaskInstanceState = "up_for_reschedule" +const TaskInstanceStateUpstreamFailed TaskInstanceState = "upstream_failed" +const TaskInstanceStateSkipped TaskInstanceState = "skipped" +const TaskInstanceStateDeferred TaskInstanceState = "deferred" +const TaskInstanceStateAwaitingInput TaskInstanceState = "awaiting_input" + // Request to get previous task instance. type GetPreviousTI struct { // DagID corresponds to the JSON schema field "dag_id". @@ -1025,6 +1127,8 @@ type GetPreviousTI struct { Type string `msgpack:"type,omitempty"` } +type TaskIds []string + type GetTICount struct { // DagID corresponds to the JSON schema field "dag_id". DagID string `msgpack:"dag_id"` @@ -1249,17 +1353,6 @@ type HITLDetailRequestResult struct { Type string `msgpack:"type,omitempty"` } -// Schema for a Human-in-the-loop users. -type HITLUser struct { - // ID corresponds to the JSON schema field "id". - ID string `msgpack:"id"` - - // Name corresponds to the JSON schema field "name". - Name string `msgpack:"name"` -} - -type ImportErrors map[string]string - type InactiveAssets []AssetProfile // Response of InactiveAssets requests. @@ -1271,19 +1364,44 @@ type InactiveAssetsResult struct { Type string `msgpack:"type,omitempty"` } -type JsonValue interface{} +// Response schema for TaskInstance run context. +type TIRunContext struct { + // ArgBindings corresponds to the JSON schema field "arg_bindings". + ArgBindings *ArgBindings `msgpack:"arg_bindings,omitempty"` -// Lazily build information from the serialized DAG structure. -// -// An object that will present "enough" of the DAG like interface to update DAG db -// models etc, without having -// to deserialize the full DAG and Task hierarchy. -type LazyDeserializedDAG struct { - // Data corresponds to the JSON schema field "data". - Data Data `msgpack:"data"` + // Connections corresponds to the JSON schema field "connections". + Connections []ConnectionResponse `msgpack:"connections,omitempty"` - // LastLoaded corresponds to the JSON schema field "last_loaded". - LastLoaded interface{} `msgpack:"last_loaded,omitempty"` + // DagRun corresponds to the JSON schema field "dag_run". + DagRun DagRun `msgpack:"dag_run"` + + // MaxTries corresponds to the JSON schema field "max_tries". + MaxTries int `msgpack:"max_tries"` + + // MultiTeam corresponds to the JSON schema field "multi_team". + MultiTeam bool `msgpack:"multi_team,omitempty"` + + // NextKwargs corresponds to the JSON schema field "next_kwargs". + NextKwargs *NextKwargs `msgpack:"next_kwargs,omitempty"` + + // NextMethod corresponds to the JSON schema field "next_method". + NextMethod interface{} `msgpack:"next_method,omitempty"` + + // ShouldRetry corresponds to the JSON schema field "should_retry". + ShouldRetry bool `msgpack:"should_retry,omitempty"` + + // StartDate corresponds to the JSON schema field "start_date". + StartDate interface{} `msgpack:"start_date,omitempty"` + + // TaskRescheduleCount corresponds to the JSON schema field + // "task_reschedule_count". + TaskRescheduleCount int `msgpack:"task_reschedule_count,omitempty"` + + // Variables corresponds to the JSON schema field "variables". + Variables []VariableResponse `msgpack:"variables,omitempty"` + + // XcomKeysToClear corresponds to the JSON schema field "xcom_keys_to_clear". + XcomKeysToClear []string `msgpack:"xcom_keys_to_clear,omitempty"` } // One positional stub-task argument carrying an inline literal from the Dag file. @@ -1304,8 +1422,6 @@ type LiteralArgBinding struct { ValueSchema *ArgValueSchema `msgpack:"value_schema,omitempty"` } -type LogicalDates []time.Time - // Add a new value to be redacted in task logs. type MaskSecret struct { // Name corresponds to the JSON schema field "name". @@ -1318,8 +1434,6 @@ type MaskSecret struct { Value JsonValue `msgpack:"value"` } -type NextKwargs map[string]interface{} - type OKResponse struct { // Ok corresponds to the JSON schema field "ok". Ok bool `msgpack:"ok"` @@ -1328,12 +1442,6 @@ type OKResponse struct { Type string `msgpack:"type,omitempty"` } -type OutletEvents []map[string]interface{} - -type Params map[string]interface{} - -type ParamsInput map[string]interface{} - type PrevSuccessfulDagRunResult struct { // DataIntervalEnd corresponds to the JSON schema field "data_interval_end". DataIntervalEnd interface{} `msgpack:"data_interval_end,omitempty"` @@ -1416,8 +1524,6 @@ type PutVariable struct { Value interface{} `msgpack:"value"` } -type RenderedFields map[string]JsonValue - // Update a task instance state to reschedule/up_for_reschedule. type RescheduleTask struct { // EndDate corresponds to the JSON schema field "end_date". @@ -1459,8 +1565,6 @@ type RetryTask struct { Type string `msgpack:"type,omitempty"` } -type RunIds []string - type SentFDs struct { // Fds corresponds to the JSON schema field "fds". Fds []int `msgpack:"fds"` @@ -1497,6 +1601,8 @@ type SetAssetStateStoreByURI struct { Value JsonValue `msgpack:"value"` } +type RenderedFields map[string]JsonValue + // Payload for setting RTIF for a task instance. type SetRenderedFields struct { // RenderedFields corresponds to the JSON schema field "rendered_fields". @@ -1594,7 +1700,9 @@ type StartupDetails struct { Type string `msgpack:"type,omitempty"` } -type States []string +type OutletEvents []map[string]interface{} + +type TaskOutlets []AssetProfile // Update a task's state to success. Includes task_outlets and outlet_events for // registering asset events. @@ -1627,47 +1735,7 @@ type TICount struct { Type string `msgpack:"type,omitempty"` } -// Response schema for TaskInstance run context. -type TIRunContext struct { - // ArgBindings corresponds to the JSON schema field "arg_bindings". - ArgBindings *ArgBindings `msgpack:"arg_bindings,omitempty"` - - // Connections corresponds to the JSON schema field "connections". - Connections []ConnectionResponse `msgpack:"connections,omitempty"` - - // DagRun corresponds to the JSON schema field "dag_run". - DagRun DagRun `msgpack:"dag_run"` - - // MaxTries corresponds to the JSON schema field "max_tries". - MaxTries int `msgpack:"max_tries"` - - // MultiTeam corresponds to the JSON schema field "multi_team". - MultiTeam bool `msgpack:"multi_team,omitempty"` - - // NextKwargs corresponds to the JSON schema field "next_kwargs". - NextKwargs *NextKwargs `msgpack:"next_kwargs,omitempty"` - - // NextMethod corresponds to the JSON schema field "next_method". - NextMethod interface{} `msgpack:"next_method,omitempty"` - - // ShouldRetry corresponds to the JSON schema field "should_retry". - ShouldRetry bool `msgpack:"should_retry,omitempty"` - - // StartDate corresponds to the JSON schema field "start_date". - StartDate interface{} `msgpack:"start_date,omitempty"` - - // TaskRescheduleCount corresponds to the JSON schema field - // "task_reschedule_count". - TaskRescheduleCount int `msgpack:"task_reschedule_count,omitempty"` - - // Variables corresponds to the JSON schema field "variables". - Variables []VariableResponse `msgpack:"variables,omitempty"` - - // XcomKeysToClear corresponds to the JSON schema field "xcom_keys_to_clear". - XcomKeysToClear []string `msgpack:"xcom_keys_to_clear,omitempty"` -} - -type TaskArgBinding interface{} +type TaskBreadcrumbsResultBreadcrumbsElem map[string]interface{} type TaskBreadcrumbsResult struct { // Breadcrumbs corresponds to the JSON schema field "breadcrumbs". @@ -1677,8 +1745,6 @@ type TaskBreadcrumbsResult struct { Type string `msgpack:"type,omitempty"` } -type TaskBreadcrumbsResultBreadcrumbsElem map[string]interface{} - // Task callback status information. // // A Class with information about the success/failure TI callback to be executed. @@ -1714,58 +1780,77 @@ type TaskCallbackRequest struct { VersionData *VersionData `msgpack:"version_data,omitempty"` } -type TaskIds []string - -// Schema for TaskInstance model with minimal required fields needed for Runtime. -type TaskInstance struct { - // ContextCarrier corresponds to the JSON schema field "context_carrier". - ContextCarrier *ContextCarrier `msgpack:"context_carrier,omitempty"` +type TaskHandlerDeclarationBinding string - // DagID corresponds to the JSON schema field "dag_id". - DagID string `msgpack:"dag_id"` +const TaskHandlerDeclarationBindingPositional TaskHandlerDeclarationBinding = "positional" +const TaskHandlerDeclarationBindingNamed TaskHandlerDeclarationBinding = "named" - // DagVersionID corresponds to the JSON schema field "dag_version_id". - DagVersionID string `msgpack:"dag_version_id"` +// One parameter of a task handler. +type TaskHandlerParam struct { + // ExactName corresponds to the JSON schema field "exact_name". + ExactName bool `msgpack:"exact_name,omitempty"` - // Hostname corresponds to the JSON schema field "hostname". - Hostname interface{} `msgpack:"hostname,omitempty"` + // Name corresponds to the JSON schema field "name". + Name interface{} `msgpack:"name"` - // ID corresponds to the JSON schema field "id". - ID string `msgpack:"id"` + // ValueSchema corresponds to the JSON schema field "value_schema". + ValueSchema *ArgValueSchema `msgpack:"value_schema,omitempty"` +} - // MapIndex corresponds to the JSON schema field "map_index". - MapIndex *int `msgpack:"map_index,omitempty"` +type TaskHandlerParams []TaskHandlerParam - // Queue corresponds to the JSON schema field "queue". - Queue string `msgpack:"queue,omitempty"` +// A task handler that a Lang-SDK artifact registers for one task. +type TaskHandlerDeclaration struct { + // Binding corresponds to the JSON schema field "binding". + Binding TaskHandlerDeclarationBinding `msgpack:"binding"` - // RunID corresponds to the JSON schema field "run_id". - RunID string `msgpack:"run_id"` + // Params corresponds to the JSON schema field "params". + Params *TaskHandlerParams `msgpack:"params"` // TaskID corresponds to the JSON schema field "task_id". TaskID string `msgpack:"task_id"` +} - // TryNumber corresponds to the JSON schema field "try_number". - TryNumber int `msgpack:"try_number"` +// Request for Task Handler Parsing. +// +// Asks a Lang-SDK runtime for every task handler an artifact registers. +type TaskHandlerParseRequest struct { + // BundleName corresponds to the JSON schema field "bundle_name". + BundleName string `msgpack:"bundle_name"` + + // BundlePath corresponds to the JSON schema field "bundle_path". + BundlePath string `msgpack:"bundle_path"` + + // File corresponds to the JSON schema field "file". + File string `msgpack:"file"` + + // Type corresponds to the JSON schema field "type". + Type string `msgpack:"type,omitempty"` } -type TaskInstanceState string +type TaskHandlers map[string][]TaskHandlerDeclaration -const TaskInstanceStateAwaitingInput TaskInstanceState = "awaiting_input" -const TaskInstanceStateDeferred TaskInstanceState = "deferred" -const TaskInstanceStateFailed TaskInstanceState = "failed" -const TaskInstanceStateQueued TaskInstanceState = "queued" -const TaskInstanceStateRemoved TaskInstanceState = "removed" -const TaskInstanceStateRestarting TaskInstanceState = "restarting" -const TaskInstanceStateRunning TaskInstanceState = "running" -const TaskInstanceStateScheduled TaskInstanceState = "scheduled" -const TaskInstanceStateSkipped TaskInstanceState = "skipped" -const TaskInstanceStateSuccess TaskInstanceState = "success" -const TaskInstanceStateUpForReschedule TaskInstanceState = "up_for_reschedule" -const TaskInstanceStateUpForRetry TaskInstanceState = "up_for_retry" -const TaskInstanceStateUpstreamFailed TaskInstanceState = "upstream_failed" +// Result of Task Handler Parsing. +// +// Every task handler a Lang-SDK artifact registers, keyed by Dag id. +// +// The answer depends only on the artifact, never on the request. +type TaskHandlerParsingResult struct { + // Fileloc corresponds to the JSON schema field "fileloc". + Fileloc string `msgpack:"fileloc"` -type TaskOutlets []AssetProfile + // ImportErrors corresponds to the JSON schema field "import_errors". + ImportErrors *ImportErrors `msgpack:"import_errors,omitempty"` + + // TaskHandlers corresponds to the JSON schema field "task_handlers". + TaskHandlers TaskHandlers `msgpack:"task_handlers"` + + // Type corresponds to the JSON schema field "type". + Type string `msgpack:"type,omitempty"` + + // Warnings corresponds to the JSON schema field "warnings". + Warnings *Warnings `msgpack:"warnings,omitempty"` +} // Response containing the first reschedule date for a task instance. type TaskRescheduleStartDate struct { @@ -1776,6 +1861,12 @@ type TaskRescheduleStartDate struct { Type string `msgpack:"type,omitempty"` } +type TaskStateState string + +const TaskStateStateFailed TaskStateState = "failed" +const TaskStateStateSkipped TaskStateState = "skipped" +const TaskStateStateRemoved TaskStateState = "removed" + // Update a task's state. // // If a process exits without sending one of these the state will be derived from @@ -1799,12 +1890,6 @@ type TaskState struct { Type string `msgpack:"type,omitempty"` } -type TaskStateState string - -const TaskStateStateFailed TaskStateState = "failed" -const TaskStateStateRemoved TaskStateState = "removed" -const TaskStateStateSkipped TaskStateState = "skipped" - // Response to GetTaskStateStore; wraps the generated API response for supervisor // to worker comms. type TaskStateStoreResult struct { @@ -1854,8 +1939,6 @@ type TriggerDagRun struct { Type string `msgpack:"type,omitempty"` } -type TriggerKwargs map[string]interface{} - type UpdateDagRunNote struct { // Note corresponds to the JSON schema field "note". Note interface{} `msgpack:"note"` @@ -1867,6 +1950,8 @@ type UpdateDagRunNote struct { Type string `msgpack:"type,omitempty"` } +type ParamsInput map[string]interface{} + // Update the response content part of an existing Human-in-the-loop response. type UpdateHITLDetail struct { // ChosenOptions corresponds to the JSON schema field "chosen_options". @@ -1882,19 +1967,6 @@ type UpdateHITLDetail struct { Type string `msgpack:"type,omitempty"` } -// Variable schema for responses with fields that are needed for Runtime. -type VariableResponse struct { - // Key corresponds to the JSON schema field "key". - Key string `msgpack:"key"` - - // Value corresponds to the JSON schema field "value". - Value interface{} `msgpack:"value"` -} - -type Warnings []interface{} - -type VersionData map[string]interface{} - type ValidateInletsAndOutlets struct { // TIID corresponds to the JSON schema field "ti_id". TIID string `msgpack:"ti_id"` From 07b5199cce6ef24dcace21c88f786e0ade0288d1 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 1 Oct 2026 17:30:46 +0000 Subject: [PATCH 42/94] Describe the two task handler binding modes in the ADRs The declarations now bind only by position or by name, and a named mismatch warns instead of failing the import, so the ADRs must say which mismatches are import errors. ADR-0012's message sketch and Appendix B, ADR-0011's comparison rules and ADR-0013's declaration sketch drop named_or_whole, named_open and required, and say what a null params means. --- .../0011-mixed-language-dag-processing.md | 14 ++++++++----- .../lang-sdk/0012-lang-sdk-parse-protocol.md | 20 ++++++++++--------- .../0013-persisted-task-handler-bindings.md | 5 ++--- 3 files changed, 22 insertions(+), 17 deletions(-) diff --git a/airflow-core/adr/lang-sdk/0011-mixed-language-dag-processing.md b/airflow-core/adr/lang-sdk/0011-mixed-language-dag-processing.md index 463e5d6960a2d..703a886e5e997 100644 --- a/airflow-core/adr/lang-sdk/0011-mixed-language-dag-processing.md +++ b/airflow-core/adr/lang-sdk/0011-mixed-language-dag-processing.md @@ -163,6 +163,7 @@ DagFileProcessorProcess(etl.py) ← manager spawn │ │ task_id ↔ task_id (sets must match) │ │ │ arg_bindings[*] ↔ params[*] (per binding) │ │ │ by position, or by folded or exact name │ + │ │ unmatched by name → a warning, not an error │ │ │ arg_bindings[*].value_schema ↔ params[*].value_schema │ │ │ compared only where neither side is null │ │ │ │ @@ -197,8 +198,9 @@ There is no fourth row. A `TaskHandlerRef` has no Dag, so no `DagImporter` — a - No importer knows about coordinators. `PythonDagImporter` is unchanged by this ADR; the stub-to-handler comparison sits in `_parse_file`, above every importer. - A mixed-language `dag_id` never appears in Dag processing results. No `Dag` registration exists for a `dag_id` a Python file already owns, so everything downstream sees exactly one record per `dag_id`, with no flag to interpret. -- Stub/implementation mismatches — missing handler, extra handler, an argument that does not bind under the declaration's `binding`, incompatible schema — surface as import errors - against the Python file at parse time, alongside the errors the parse already reports. An unannotated stub argument is checked only for how it binds, by position or by name. +- Stub/implementation mismatches — missing handler, extra handler, a `positional` argument count that does not match, incompatible schema — surface as import errors against + the Python file at parse time, alongside the errors the parse already reports. A `named` argument or parameter that matches nothing is only logged as a warning, because the + runtime runs the task anyway. An unannotated stub argument is checked only for how it binds, by position or by name. - The Python Dag and Lang-SDK artifact can live in different DagBundles. - A single Dag can have stubs targeting different queues, some Java, some Go. Each resolves to its own coordinator instance, and validation unions their declarations per `dag_id` before comparing task ids. @@ -238,10 +240,12 @@ Handler declarations from every process spawned in Step 4 are unioned per `dag_i every Dag its artifact registers handlers for; only the parsed file's Dags are compared. - `task_id` sets must match exactly. A missing or extra handler is an error. -- `arg_bindings` against `params` as each declaration's `binding` says: by position for `positional`, by name for `named`, `named_or_whole` and `named_open`, compared - case-insensitively with underscores ignored unless `exact_name` is set. Under `named_open` only the listed names are checked ([ADR-0012](0012-lang-sdk-parse-protocol.md) Appendix B). +- `arg_bindings` against `params` as each declaration's `binding` says: by position for `positional`, by name for `named`, compared case-insensitively with underscores ignored + unless `exact_name` is set. A `positional` count mismatch is an error. Under `named`, an argument or parameter that matches nothing is only logged as a warning, and a lone + unmatched argument, which may be the whole value, is not ([ADR-0012](0012-lang-sdk-parse-protocol.md) Appendix B). When `params` is `None`, only the handler's presence is + checked. - `arg_bindings[*].value_schema` against `params[*].value_schema`, compared only where neither side is null. An unannotated `@task.stub` parameter produces `null` today, so a strict comparison would make every untyped stub argument a parse error. -Any mismatch is reported against the Python file, which is the definition the author can act on, and travels back on `DagFileParsingResult.import_errors` with everything else the +Every error is reported against the Python file, which is the definition the author can act on, and travels back on `DagFileParsingResult.import_errors` with everything else the parse found. diff --git a/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md b/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md index afda6794a794d..39e8da7ce9617 100644 --- a/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md +++ b/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md @@ -95,13 +95,12 @@ class TaskHandlerParsingResult: class TaskHandlerDeclaration: task_id: str - binding: Literal["positional", "named", "named_or_whole", "named_open"] # how stub-task arguments bind to params - params: list[TaskHandlerParam] # ordered; the order matters only for "positional" + binding: Literal["positional", "named"] # how stub-task arguments bind to params + params: list[TaskHandlerParam] | None # ordered; the order matters only for "positional". None: the runtime cannot list them class TaskHandlerParam: name: str | None # None: the runtime has no name for this positional parameter value_schema: ArgValueSchema | None = None - required: bool # the handler declares no default exact_name: bool = False # match as spelled, not case-insensitively with underscores ignored ``` @@ -218,12 +217,15 @@ and falls back to name-and-arity otherwise. A strict comparison would turn every The stub side always has names and positions: `LiteralArgBinding` and `XComArgBinding` are each documented as "one positional stub-task argument". The handler side binds the way its runtime does, so each declaration names its `binding` and the check follows it: -- `positional`: by position. Names are informative only, and absent where the runtime has none (Go flat params). Java's `TaskArgs` binds this way although it has names. -- `named`: by name in any order, case-insensitively with underscores ignored unless `exact_name` is set (Go `arg:` tags, explicit Java names). A Go struct with `arg:` tags and - Java's `TaskInput` bind this way. -- `named_or_whole`: as `named`, except that when there is exactly one argument and it matches no parameter, it is decoded as the whole value. An untagged Go struct binds this way. -- `named_open`: as `named`, but `params` need not list every parameter the handler takes. Only the listed names are checked, and an argument that none of them names is not a - mismatch. TypeScript binds this way: its types are erased, so a handler knows only its explicit `withArgNames` renames. +- `positional`: by position. Names are informative only, and absent where the runtime has none (Go flat params). Java's `TaskArgs` binds this way although it has names. An + argument count that matches `params` neither with every argument nor after dropping the defaulted ones, or a value type a parameter does not accept, is an import error. +- `named`: by name in any order, case-insensitively with underscores ignored unless `exact_name` is set (Go `arg:` tags, explicit Java names). A Go struct, tagged or untagged, + and Java's `TaskInput` bind this way. An argument no parameter takes, or a parameter no argument fills, is logged as a warning, and the task still runs: the runtimes allow + both, and an unfilled field keeps its default. When no parameter matches and exactly one argument was passed, it may be the whole value and is not warned about, unless no + parameter is declared, a parameter sets `exact_name` (a Go `arg:` tag), or the argument cannot be an object. A value type a parameter does not accept is an import error. + +`params` is `None` when the runtime cannot list a handler's parameters, and then only the handler's presence is checked. TypeScript declares `named` with `params: None`: its +types are erased, so a handler cannot list what it takes. A declaration carries no class, method, or source location. [ADR-0006](0006-no-lang-sdk-source-display.md) rules out Lang-SDK source display, and putting it on the wire would invite a consumer to render it. It carries no `dag_id` either — the `task_handlers` key supplies it, so a declaration cannot disagree with the bucket it arrived in. diff --git a/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md b/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md index e778600776f77..47468bb06a926 100644 --- a/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md +++ b/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md @@ -259,14 +259,13 @@ class SDKTaskHandlerParsingResult(BaseModel): # runtime -> parent, on ToManager class TaskHandlerDeclaration(BaseModel): task_id: str - binding: Literal["positional", "named", "named_or_whole", "named_open"] # how arguments bind to params - params: list[TaskHandlerParam] # ordered; the order matters only for "positional" + binding: Literal["positional", "named"] # how arguments bind to params + params: list[TaskHandlerParam] | None # ordered; None: the runtime cannot list them class TaskHandlerParam(BaseModel): name: str | None # None: the runtime has no name for this positional parameter value_schema: JSONSchema | None = None - required: bool # the handler declares no default exact_name: bool = False # match as spelled, not case-insensitively with underscores ignored ``` From 32db985ad1d21333baad3b0e23f25e82df6d0524 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 12:29:49 +0000 Subject: [PATCH 43/94] Keep a nested supervised child's standard streams open A supervised child reopens stdin, stdout and stderr on its sockets and replaces sys.stdin and friends. When that child supervises a child of its own, as a Dag-parsing child starting a task-handler probe will, the grandchild replaces handles that already wrap fds 0-2, and dropping them closed the fds it had just dup2'd, so the grandchild lost its streams. The handles no longer own their fds; the interpreter's original ones are kept alive by sys.__stdin__ and friends and were never closed here. --- .../airflow/sdk/execution_time/supervisor.py | 5 +++-- .../execution_time/test_supervisor.py | 20 +++++++++++++++++++ 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index f2373b3bb835e..ecb3799bfb34a 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -390,8 +390,9 @@ def _reopen_std_io_handles(child_stdin, child_stdout, child_stderr): del sock # We open the socket/fd as binary, and then pass it to a TextIOWrapper so that it looks more like a - # normal sys.stdout etc. - binary = os.fdopen(fd, mode + "b") + # normal sys.stdout etc. The handle does not own the fd: a child this process supervises replaces + # it after dup2 onto the same fd, and closing it there would close the child's new stream. + binary = os.fdopen(fd, mode + "b", closefd=False) handle = io.TextIOWrapper(binary, line_buffering=True) setattr(sys, handle_name, handle) diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py index 12daa4173d58f..7b931734bb566 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -172,6 +172,7 @@ WatchedSubprocess, _make_process_nondumpable, _remote_logging_conn, + _reopen_std_io_handles, forward_to_log, in_process_api_server, make_buffered_socket_reader, @@ -5060,6 +5061,25 @@ def on_close(sock): log_write.close() +def test_a_nested_child_keeps_its_standard_streams(): + """A child of a supervised child reopens handles its parent already reopened, and keeps them open.""" + pid = os.fork() + if pid == 0: + code = 1 + try: + for _ in ("supervised child", "its child"): + pairs = [socket.socketpair() for _ in range(3)] + _reopen_std_io_handles(*(child_end for child_end, _ in pairs)) + for fd in (0, 1, 2): + os.fstat(fd) + code = 0 + finally: + os._exit(code) + _, status = os.waitpid(pid, 0) + + assert os.waitstatus_to_exitcode(status) == 0 + + def test_reinit_supervisor_comms(monkeypatch, client_with_ti_start, caplog): def subprocess_main(): # This is run in the subprocess! From 1ac623cb01a68af0d8dd08de53e53f6c4ff4d888 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Fri, 2 Oct 2026 03:26:39 +0000 Subject: [PATCH 44/94] Let a subprocess coordinator exec a task-handler parse runtime BaseCoordinator.parse_task_handler replaces the calling process with the runtime that reports an artifact's task handlers. The default raises NotImplementedError, so a coordinator opts in; SubprocessCoordinator does so through _build_parse_task_handler_command, under the artifact bundle's scan roots. SubprocessCoordinator.parse_task_handler takes the steps of the native parse_dag (#74042) in the same order, with the same helpers: the pid based connection check, close-on-exec above stderr and the runtime env. On Linux the runtime also gets PR_SET_PDEATHSIG, so it dies with the process that reads its sockets, which the manager kills by pid when a Dag-parsing child times out. CoordinatorManager.get_coordinator finds a coordinator by its config key, which is how a parse child that has only the key rebuilds it. for_queue now builds its coordinator through it too. --- .../airflow/sdk/coordinators/_subprocess.py | 155 +++++++++++++++--- .../airflow/sdk/execution_time/coordinator.py | 59 +++++-- .../task_sdk/coordinators/test_subprocess.py | 135 ++++++++++++++- .../execution_time/test_coordinator.py | 58 +++++++ 4 files changed, 368 insertions(+), 39 deletions(-) diff --git a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py index 7ddddc9eeca4e..9a9b9f5a50e58 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py +++ b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py @@ -28,6 +28,7 @@ from __future__ import annotations import contextlib +import ctypes import ipaddress import itertools import os @@ -35,8 +36,9 @@ import signal import socket import subprocess +import sys import time -from typing import TYPE_CHECKING, TypeVar, cast +from typing import TYPE_CHECKING, NoReturn, TypeVar, cast import attrs import psutil @@ -47,11 +49,12 @@ 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.schema import get_schema_version_migrator from airflow.sdk.execution_time.supervisor import ActivitySubprocess, NeverRaised, ProcessTracker if TYPE_CHECKING: import pathlib - from collections.abc import Sequence + from collections.abc import Callable, Sequence from structlog.typing import FilteringBoundLogger from typing_extensions import Self @@ -101,21 +104,19 @@ def _socket_address(value: tuple | str) -> tuple[str, int] | None: return host, int(port) -def _connection_owned_by_process_tree( - peer: tuple[str, int], local: tuple[str, int], proc: subprocess.Popen -) -> bool: +def _connection_owned_by_process_tree(peer: tuple[str, int], local: tuple[str, int], pid: int) -> bool: """ Return whether ``peer`` <-> ``local`` is an established connection in the child's process tree. The launched child may itself spawn the process that connects back to the supervisor — a JVM launcher, a shell wrapper, or any runtime that forks a worker — so the connecting peer can legitimately belong to a *descendant* of - ``proc.pid`` rather than ``proc.pid`` itself. Every process in the subtree - rooted at ``proc.pid`` is part of the task and is trusted; a process outside - that subtree (e.g. an unrelated local process racing for the port) is not. + ``pid`` rather than ``pid`` itself. Every process in the subtree rooted at + ``pid`` is part of the task and is trusted; a process outside that subtree + (e.g. an unrelated local process racing for the port) is not. """ try: - root = psutil.Process(proc.pid) + root = psutil.Process(pid) processes = [root, *root.children(recursive=True)] except (psutil.AccessDenied, psutil.NoSuchProcess, psutil.ZombieProcess, OSError): return False @@ -153,13 +154,68 @@ def _is_connection_from_process( return False deadline = time.monotonic() + verify_timeout while True: - if _connection_owned_by_process_tree(peer, local, proc): + if _connection_owned_by_process_tree(peer, local, proc.pid): return True if time.monotonic() >= deadline: return False time.sleep(poll_interval) +def _is_connection_from_pid(conn: socket.socket, pid: int) -> bool: + """ + Return whether the accepted TCP connection comes from ``pid``'s process tree, checking once. + + A connection that is not visible yet is reported as not owned, so the caller retries later + instead of waiting here. + """ + peer = _socket_address(conn.getpeername()) + local = _socket_address(conn.getsockname()) + return peer is not None and local is not None and _connection_owned_by_process_tree(peer, local, pid) + + +def _set_close_on_exec_above_stderr() -> None: + """Mark every file descriptor above 2 close-on-exec, as ``subprocess.Popen(close_fds=True)`` does.""" + fd_dir = "/proc/self/fd" if os.path.isdir("/proc/self/fd") else "/dev/fd" + for name in os.listdir(fd_dir): + if (fd := int(name)) > 2: + # The descriptor listing the directory is already closed. + with contextlib.suppress(OSError): + os.set_inheritable(fd, False) + + +def _build_runtime_env() -> dict[str, str]: + """Return the environment for a language SDK runtime, with the Airflow config it reads resolved.""" + # A language SDK runtime cannot read Airflow's config, so the options it needs are resolved here. + # An option without a config default gets the fallback Python uses. StartupDetails arrives too + # late: logs may already be produced. + return { + **os.environ, + "AIRFLOW__LOGGING__LOGGING_LEVEL": conf.get("logging", "logging_level", fallback="INFO"), + "AIRFLOW__LOGGING__NAMESPACE_LEVELS": conf.get("logging", "namespace_levels", fallback=""), + } + + +# From +_PR_SET_PDEATHSIG = 1 + + +def _set_parent_death_signal() -> None: + """ + On Linux, have the kernel SIGKILL this process when the thread that forked it exits. + + The setting survives the exec, so the runtime dies with the process that reads its sockets, + even when that process is killed without the chance to kill the runtime first. + """ + if sys.platform != "linux": + return + try: + libc = ctypes.CDLL(None, use_errno=True) + if libc.prctl(_PR_SET_PDEATHSIG, signal.SIGKILL, 0, 0, 0) != 0: + log.warning("Failed to set PR_SET_PDEATHSIG", errno=ctypes.get_errno()) + except Exception: + log.warning("Unable to set PR_SET_PDEATHSIG", exc_info=True) + + def _accept_connections( servers: dict[str, socket.socket], drains: dict[str, socket.socket], @@ -310,15 +366,6 @@ def start( # type: ignore[override] stdout_r, stdout_w = tracker.track(*socket.socketpair()) stderr_r, stderr_w = tracker.track(*socket.socketpair()) - # A language SDK runtime cannot read Airflow's config, so propagate the - # resolved log levels via the environment at launch. StartupDetails - # arrives too late, the logs might already be produced by then. - env = { - **os.environ, - "AIRFLOW__LOGGING__LOGGING_LEVEL": conf.get("logging", "logging_level", fallback="INFO"), - "AIRFLOW__LOGGING__NAMESPACE_LEVELS": conf.get("logging", "namespace_levels", fallback=""), - } - proc = subprocess.Popen( [ *command, @@ -327,7 +374,7 @@ def start( # type: ignore[override] ], stdout=stdout_w.fileno(), stderr=stderr_w.fileno(), - env=env, + env=_build_runtime_env(), ) tracker.track(proc) for soc in tracker.untrack(stdout_w, stderr_w): @@ -411,6 +458,9 @@ class SubprocessCoordinator(BaseCoordinator): connections, draining startup output, and tearing everything down on failure — is handled here. + A subclass that also reports its artifacts' task handlers implements + :meth:`_build_parse_task_handler_command`. + :param task_startup_timeout: Maximum time the coordinator waits for the subprocess to connect to both servers, in seconds. The default is 10 seconds. @@ -448,9 +498,12 @@ def _resolve_artifact_bundle( return bundle def _get_scan_roots(self) -> tuple[pathlib.Path, ...]: - """Return the artifact roots resolved for the active task.""" + """Return the artifact roots resolved for the active task or parse.""" if self._active_scan_roots is None: - raise RuntimeError("_get_scan_roots requires an active task; call it during execute_task.") + raise RuntimeError( + "_get_scan_roots requires an active task or parse; " + "call it during execute_task or parse_task_handler." + ) return self._active_scan_roots def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], str | None]: @@ -467,11 +520,65 @@ def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], """ raise NotImplementedError + def _build_parse_task_handler_command(self, *, path: pathlib.Path) -> tuple[list[str], str | None]: + """ + Build the command that reports the task handlers of the artifact at *path*, with its schema version. + + Subclasses can retrieve the directories to scan for artifacts with + :meth:`_get_scan_roots`; for a parse they are the root of the Dag bundle holding + *path*. The contract is that of :meth:`_build_execute_task_command`: *command* + MUST NOT include the ``--comm`` / ``--logs`` flags. + """ + raise NotImplementedError(f"{type(self).__name__} does not parse task handlers") + + def parse_task_handler( + self, + *, + path: pathlib.Path, + bundle_path: pathlib.Path, + comm_address: tuple[str, int], + logs_address: tuple[str, int], + report_schema_version: Callable[[str | None], None], + ) -> NoReturn: + """ + Replace the current process with the runtime that reports the task handlers of the artifact at *path*. + + Call this in a child process whose standard streams are already set up; the runtime + inherits them and no other file descriptor. The command and its supervisor wire-schema + version are resolved against *bundle_path*, and the version is passed to + *report_schema_version* just before the exec. The runtime connects back to + *comm_address* and *logs_address*. On Linux the runtime dies with the thread that forked + the calling process. + + :raises Exception: when the command cannot be resolved or started; the process is then + unchanged. + """ + with self._set_scan_roots([bundle_path]): + command, schema_version = self._build_parse_task_handler_command(path=path) + if schema_version is not None: + get_schema_version_migrator().resolve_version(schema_version) + argv = [ + *command, + f"--comm={comm_address[0]}:{comm_address[1]}", + f"--logs={logs_address[0]}:{logs_address[1]}", + ] + report_schema_version(schema_version) + # Python ignores these at startup and exec keeps ignored signals; subprocess.Popen resets + # them the same way for the task runtime. + for name in ("SIGPIPE", "SIGXFZ", "SIGXFSZ"): + if (sig := getattr(signal, name, None)) is not None: + signal.signal(sig, signal.SIG_DFL) + _set_parent_death_signal() + _set_close_on_exec_above_stderr() + os.execvpe(argv[0], argv, _build_runtime_env()) + @contextlib.contextmanager def _set_scan_roots(self, roots: Sequence[pathlib.Path]): - """Expose *roots* to the command builder for the duration of the task.""" + """Expose *roots* to the command builder for the duration of the task or parse.""" if self._active_scan_roots is not None: - raise RuntimeError("SubprocessCoordinator.execute_task is not re-entrant.") + raise RuntimeError( + "SubprocessCoordinator.execute_task and parse_task_handler are not re-entrant." + ) self._active_scan_roots = tuple(roots) try: yield diff --git a/task-sdk/src/airflow/sdk/execution_time/coordinator.py b/task-sdk/src/airflow/sdk/execution_time/coordinator.py index 778ea29aea899..904faf5068de3 100644 --- a/task-sdk/src/airflow/sdk/execution_time/coordinator.py +++ b/task-sdk/src/airflow/sdk/execution_time/coordinator.py @@ -42,7 +42,7 @@ import functools import os import signal -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, NoReturn import attrs import pydantic @@ -53,7 +53,8 @@ from airflow.sdk.configuration import conf if TYPE_CHECKING: - from collections.abc import Generator, Mapping + import pathlib + from collections.abc import Callable, Generator, Mapping from os import PathLike from structlog.typing import FilteringBoundLogger @@ -108,6 +109,30 @@ def execute_task( """ raise NotImplementedError + def parse_task_handler( + self, + *, + path: pathlib.Path, + bundle_path: pathlib.Path, + comm_address: tuple[str, int], + logs_address: tuple[str, int], + report_schema_version: Callable[[str | None], None], + ) -> NoReturn: + """ + Replace the current process with the runtime that reports the task handlers of the artifact at *path*. + + Call this in a child process whose standard streams are already set up; the runtime inherits + them and no other file descriptor. *bundle_path* is the root of the Dag bundle holding *path*. + The runtime connects back to *comm_address* and *logs_address* and answers a + ``TaskHandlerParseRequest``. Its supervisor wire-schema version is passed to + *report_schema_version* just before the exec. + + A coordinator opts in by overriding this; the default raises :class:`NotImplementedError`. + + :raises Exception: when the runtime cannot be resolved or started. + """ + raise NotImplementedError(f"{type(self).__name__} does not parse task handlers") + class _CoordinatorSpec(pydantic.BaseModel): classpath: str @@ -279,13 +304,6 @@ def from_config(cls) -> Self: ) return cls(coordinator_specs=coordinator_specs, queue_to_coordinator=queue_to_coordinator) - def _find_queue(self, key: str) -> BaseCoordinator: - with contextlib.suppress(KeyError): - return self._created_coordinators[key] - spec = self._coordinator_specs[key] - coordinator = self._created_coordinators[key] = import_string(spec.classpath)(**spec.kwargs) - return coordinator - def for_queue(self, queue: str) -> BaseCoordinator: """ Find the coordinator for *queue*. @@ -297,15 +315,32 @@ def for_queue(self, queue: str) -> BaseCoordinator: except KeyError: log.debug("Queue not configured to a coordinator; defaulting to Python", queue=queue) return _build_python_coordinator() + if key not in self._coordinator_specs: + raise InvalidCoordinatorError(f"Queue {queue!r} configured to nonexistent coordinator") + coordinator = self.get_coordinator(key) + log.debug("Coordinator found for queue", coordinator=coordinator, queue=queue) + return coordinator + + def get_coordinator(self, key: str) -> BaseCoordinator: + """ + Return the coordinator configured under *key* in ``[sdk] coordinators``, building it on first use. + + :raises InvalidCoordinatorError: when *key* is not configured, or its coordinator cannot be + imported or built. + """ + with contextlib.suppress(KeyError): + return self._created_coordinators[key] try: - coordinator = self._find_queue(key) + spec = self._coordinator_specs[key] except KeyError: - raise InvalidCoordinatorError(f"Queue {queue!r} configured to nonexistent coordinator") + raise InvalidCoordinatorError(f"No coordinator {key!r} in [sdk] coordinators") + try: + coordinator = import_string(spec.classpath)(**spec.kwargs) except ImportError: raise InvalidCoordinatorError(f"Cannot import coordinator {key!r}") except TypeError: raise InvalidCoordinatorError(f"Cannot instantiate coordinator {key!r}") - log.debug("Coordinator found for queue", coordinator=coordinator, queue=queue) + self._created_coordinators[key] = coordinator return coordinator def extra_for_queue(self, queue: str) -> dict[str, Any] | None: diff --git a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py index f836c6797350c..81461785d9fa4 100644 --- a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py +++ b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py @@ -20,6 +20,7 @@ import contextlib import os import pathlib +import signal import socket import subprocess import sys @@ -39,6 +40,7 @@ SubprocessCoordinator, _accept_connections, _connection_owned_by_process_tree, + _is_connection_from_pid, _is_connection_from_process, _PopenActivitySubprocess, _ResourceTracker, @@ -494,13 +496,22 @@ def test_rejects_when_ownership_never_confirmed(self): ) def test_owned_by_tree_returns_false_when_process_gone(self): - mock_proc = MagicMock(spec=subprocess.Popen) - mock_proc.pid = 999999 with patch( "airflow.sdk.coordinators._subprocess.psutil.Process", side_effect=psutil.NoSuchProcess(999999), ): - assert _connection_owned_by_process_tree(("127.0.0.1", 1), ("127.0.0.1", 2), mock_proc) is False + assert _connection_owned_by_process_tree(("127.0.0.1", 1), ("127.0.0.1", 2), 999999) is False + + @pytest.mark.parametrize("owned", [True, False]) + @patch("airflow.sdk.coordinators._subprocess._connection_owned_by_process_tree", autospec=True) + def test_is_connection_from_pid_checks_once(self, mock_owned, owned): + conn = MagicMock(spec=socket.socket) + conn.getpeername.return_value = ("127.0.0.1", 5000) + conn.getsockname.return_value = ("127.0.0.1", 6000) + mock_owned.return_value = owned + + assert _is_connection_from_pid(conn, 999) is owned + mock_owned.assert_called_once_with(("127.0.0.1", 5000), ("127.0.0.1", 6000), 999) class TestResourceTracker: @@ -1020,3 +1031,121 @@ def test_execute_task_is_not_reentrant(self, mock_client, tmp_path): client=mock_client, subprocess_logs_to_stdout=False, ) + + +@attrs.define(kw_only=True) +class _TaskHandlerParsingCoordinator(_StubSubprocessCoordinator): + def _build_parse_task_handler_command(self, *, path): + self.recorded_roots.append(list(self._get_scan_roots())) + return [*self.command, os.fspath(path)], self.schema_version + + +def _parse_task_handler(coordinator: SubprocessCoordinator, tmp_path: pathlib.Path, reported: list) -> None: + coordinator.parse_task_handler( + path=tmp_path / "handlers.artifact", + bundle_path=tmp_path, + comm_address=("127.0.0.1", 1001), + logs_address=("127.0.0.1", 1002), + report_schema_version=reported.append, + ) + + +def _is_running(process: psutil.Process) -> bool: + try: + return process.status() != psutil.STATUS_ZOMBIE + except psutil.NoSuchProcess: + return False + + +class TestParseTaskHandler: + def test_build_parse_task_handler_command_default_raises(self, tmp_path): + coordinator = _StubSubprocessCoordinator(command=["x"]) + + with pytest.raises( + NotImplementedError, match="_StubSubprocessCoordinator does not parse task handlers" + ): + coordinator._build_parse_task_handler_command(path=tmp_path / "handlers.artifact") + + @patch("airflow.sdk.coordinators._subprocess._set_close_on_exec_above_stderr", autospec=True) + @patch("airflow.sdk.coordinators._subprocess._set_parent_death_signal", autospec=True) + @patch("airflow.sdk.coordinators._subprocess.signal.signal", autospec=True) + @patch( + "airflow.sdk.coordinators._subprocess.os.execvpe", autospec=True, side_effect=OSError("exec failed") + ) + def test_resolves_the_command_under_the_bundle_root_then_execs_it( + self, mock_execvpe, mock_signal, mock_death_signal, mock_close_on_exec, tmp_path + ): + coordinator = _TaskHandlerParsingCoordinator(command=["runtime"], schema_version="2026-06-16") + reported: list[str | None] = [] + + with pytest.raises(OSError, match="exec failed"): + _parse_task_handler(coordinator, tmp_path, reported) + + assert coordinator.recorded_roots == [[tmp_path]] + assert coordinator._active_scan_roots is None + assert reported == ["2026-06-16"] + argv = [ + "runtime", + os.fspath(tmp_path / "handlers.artifact"), + "--comm=127.0.0.1:1001", + "--logs=127.0.0.1:1002", + ] + mock_execvpe.assert_called_once_with("runtime", argv, ANY) + assert "AIRFLOW__LOGGING__LOGGING_LEVEL" in mock_execvpe.call_args.args[2] + restored = {c.args[0] for c in mock_signal.call_args_list if c.args[1] == signal.SIG_DFL} + assert {signal.SIGPIPE, signal.SIGXFSZ} <= restored + mock_death_signal.assert_called_once_with() + mock_close_on_exec.assert_called_once_with() + + @patch("airflow.sdk.coordinators._subprocess.os.execvpe", autospec=True) + def test_rejects_an_unknown_schema_version_before_reporting(self, mock_execvpe, tmp_path): + coordinator = _TaskHandlerParsingCoordinator(command=["runtime"], schema_version="1999-01-01") + reported: list[str | None] = [] + + with pytest.raises(ValueError, match="'1999-01-01' not found"): + _parse_task_handler(coordinator, tmp_path, reported) + + assert reported == [] + mock_execvpe.assert_not_called() + + @pytest.mark.skipif(sys.platform != "linux", reason="PR_SET_PDEATHSIG is Linux-only") + def test_the_runtime_dies_with_the_process_that_started_it(self, tmp_path): + coordinator = _TaskHandlerParsingCoordinator(command=["/bin/sh", "-c", "exec sleep 60"]) + read_fd, write_fd = os.pipe() + parent_pid = os.fork() + if parent_pid == 0: + # The process reading the runtime's sockets: it starts the runtime, reports its pid and waits. + try: + os.close(read_fd) + runtime_pid = os.fork() + if runtime_pid == 0: + try: + _parse_task_handler(coordinator, tmp_path, []) + finally: + os._exit(1) + os.write(write_fd, str(runtime_pid).encode()) + os.close(write_fd) + os.waitpid(runtime_pid, 0) + finally: + os._exit(0) + + os.close(write_fd) + with os.fdopen(read_fd) as reader: + runtime = psutil.Process(int(reader.read())) + try: + # Polls: the runtime is not a child to wait for, and time_machine cannot reach another process. + deadline = time.monotonic() + 30 + while runtime.name() != "sleep": + assert time.monotonic() < deadline, "the runtime did not start" + time.sleep(0.05) + + os.kill(parent_pid, signal.SIGKILL) + os.waitpid(parent_pid, 0) + + deadline = time.monotonic() + 10 + while _is_running(runtime): + assert time.monotonic() < deadline, "the runtime outlived the process that started it" + time.sleep(0.05) + finally: + with contextlib.suppress(psutil.NoSuchProcess): + runtime.kill() diff --git a/task-sdk/tests/task_sdk/execution_time/test_coordinator.py b/task-sdk/tests/task_sdk/execution_time/test_coordinator.py index ccd6a22280e8d..cc8a81410a04f 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_coordinator.py +++ b/task-sdk/tests/task_sdk/execution_time/test_coordinator.py @@ -244,6 +244,53 @@ def test_for_queue_raises_when_spec_missing(self): ): manager.for_queue("queue-bad") + def test_get_coordinator_builds_the_keyed_coordinator_once(self, sdk_config): + sdk_config( + coordinators=json.dumps( + { + "alpha": { + "classpath": f"{_CoordinatorA.__module__}._CoordinatorA", + "kwargs": {"label": "x"}, + } + } + ), + ) + manager = CoordinatorManager.from_config() + + coordinator = manager.get_coordinator("alpha") + + assert isinstance(coordinator, _CoordinatorA) + assert coordinator.label == "x" + assert manager.get_coordinator("alpha") is coordinator + + @pytest.mark.parametrize( + ("coordinators", "expected_match"), + [ + pytest.param({}, r"No coordinator 'bad-coord' in \[sdk\] coordinators", id="not-configured"), + pytest.param( + {"bad-coord": {"classpath": "nonexistent.module.MissingCoordinator"}}, + r"Cannot import coordinator 'bad-coord'", + id="classpath-cannot-be-imported", + ), + pytest.param( + { + "bad-coord": { + "classpath": f"{_CoordinatorA.__module__}._CoordinatorA", + "kwargs": {"unknown_param": "x"}, + } + }, + r"Cannot instantiate coordinator 'bad-coord'", + id="kwargs-dont-match-constructor", + ), + ], + ) + def test_get_coordinator_raises_invalid_coordinator_error(self, sdk_config, coordinators, expected_match): + sdk_config(coordinators=json.dumps(coordinators)) + manager = CoordinatorManager.from_config() + + with pytest.raises(InvalidCoordinatorError, match=expected_match): + manager.get_coordinator("bad-coord") + def test_get_coordinator_manager_is_cached(self, monkeypatch): monkeypatch.delenv("AIRFLOW__SDK__COORDINATORS", raising=False) @@ -450,6 +497,17 @@ def test_sigterm_inside_context_does_not_kill(self): assert signal.getsignal(signal.SIGTERM) is signal.SIG_DFL +def test_a_coordinator_parses_no_task_handlers_by_default(tmp_path): + with pytest.raises(NotImplementedError, match="_CoordinatorB does not parse task handlers"): + _CoordinatorB().parse_task_handler( + path=tmp_path / "handlers.artifact", + bundle_path=tmp_path, + comm_address=("127.0.0.1", 1001), + logs_address=("127.0.0.1", 1002), + report_schema_version=lambda schema_version: None, + ) + + class TestPythonCoordinatorWarmShutdown: """The Python coordinator must wrap start() and wait() in the warm-shutdown handlers.""" From 4cb786b1d35e8f5fa71337018c77274de489c3b4 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Fri, 2 Oct 2026 03:34:50 +0000 Subject: [PATCH 45/94] Probe a Lang-SDK artifact for its task handlers The Dag-file parse will ask a coordinator's runtime which task handlers an artifact registers (ADR-0011, ADR-0012). LangSDKTaskHandlerProcessorProcess does that: it sends TaskHandlerParseRequest and collects one TaskHandlerParsingResult, and a failed start, a missing result, an invalid frame or message, or a timeout becomes an import error on it. Nothing calls it yet. It is written like LangSDKDagFileProcessorProcess of the native Dag parse (#74035), with the same method names and order, so the two can later share one base. Its parse child finds the coordinator by its config key and reports a failed start with its own message. The runtime's channel accepts ToManager messages only, and a DagFileParsingResult on it is an unhandled request. The probe runs in a Dag-parsing child, which has no API client, so the runtime's requests are relayed up that child's supervisor channel and the replies, errors included, sent back unchanged. With neither a client nor a parent, they get an error. The ADRs name the probe and its entrypoint as the code does. --- .../0011-mixed-language-dag-processing.md | 4 +- .../lang-sdk/0012-lang-sdk-parse-protocol.md | 4 +- .../0013-persisted-task-handler-bindings.md | 4 +- .../dag_processing/task_handler_processor.py | 550 +++++++++++++ .../fake_task_handler_runtime.py | 120 +++ .../test_task_handler_processor.py | 777 ++++++++++++++++++ generated/known_sdk_imports_in_core.txt | 1 + 7 files changed, 1454 insertions(+), 6 deletions(-) create mode 100644 airflow-core/src/airflow/dag_processing/task_handler_processor.py create mode 100644 airflow-core/tests/unit/dag_processing/fake_task_handler_runtime.py create mode 100644 airflow-core/tests/unit/dag_processing/test_task_handler_processor.py diff --git a/airflow-core/adr/lang-sdk/0011-mixed-language-dag-processing.md b/airflow-core/adr/lang-sdk/0011-mixed-language-dag-processing.md index 703a886e5e997..5dc9518068130 100644 --- a/airflow-core/adr/lang-sdk/0011-mixed-language-dag-processing.md +++ b/airflow-core/adr/lang-sdk/0011-mixed-language-dag-processing.md @@ -130,8 +130,8 @@ DagFileProcessorProcess(etl.py) ← manager spawn │ ┌─────────────────────────────────────────────────────────────────────┐ │ │ Step 4: Query each group — one request, one response │ │ │ │ - │ │ SDKTaskHandlerProcessorProcess.start( │ - │ │ target=_parse_task_handler_entrypoint, │ + │ │ LangSDKTaskHandlerProcessorProcess.start( │ + │ │ target=_start_task_handler_runtime_entrypoint, │ │ │ coordinator=JavaCoordinator("jdk-11"), │ │ │ path=analytics.jar) │ │ │ │ │ diff --git a/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md b/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md index 39e8da7ce9617..34c18508e4567 100644 --- a/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md +++ b/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md @@ -130,8 +130,8 @@ WatchedSubprocess │ └── coordinator.parse_dag() — spawn runtime, forward fd 0 ⇄ comm socket │ same request and result types as its base class │ - └── SDKTaskHandlerProcessorProcess (new — ADR-0011) - target = _parse_task_handler_entrypoint + └── LangSDKTaskHandlerProcessorProcess (new — ADR-0011) + target = _start_task_handler_runtime_entrypoint └── coordinator.parse_task_handler() — same forwarding TaskHandlerParseRequest → TaskHandlerParsingResult ``` diff --git a/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md b/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md index 47468bb06a926..6a62c732497c3 100644 --- a/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md +++ b/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md @@ -390,8 +390,8 @@ DagFileProcessorProcess(etl.py) [no DB — client con │ anything differs ──▶ probe │ ├─6─ PROBE, per differing candidate — one subprocess - │ SDKTaskHandlerProcessorProcess.start( - │ target=_parse_task_handler_entrypoint, + │ LangSDKTaskHandlerProcessorProcess.start( + │ target=_start_task_handler_runtime_entrypoint, │ coordinator=JavaCoordinator("jdk-17"), │ path=) │ ──SDKTaskHandlerParseRequest(file=…, dag_ids=["etl"])──▶ runtime diff --git a/airflow-core/src/airflow/dag_processing/task_handler_processor.py b/airflow-core/src/airflow/dag_processing/task_handler_processor.py new file mode 100644 index 0000000000000..eca802ee00066 --- /dev/null +++ b/airflow-core/src/airflow/dag_processing/task_handler_processor.py @@ -0,0 +1,550 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Ask the Lang-SDK runtime of a coordinator which task handlers an artifact registers.""" + +from __future__ import annotations + +import functools +import os +import selectors +import signal +import time +from pathlib import Path +from socket import socket +from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Literal, cast, get_args + +import attrs +import msgspec +from pydantic import BaseModel, Field, TypeAdapter +from uuid6 import uuid7 + +from airflow import settings +from airflow.configuration import conf +from airflow.dag_processing.processor import ( + BaseDagFileProcessorProcess, + DagFileParsingResult, + TaskHandlerParseRequest, + TaskHandlerParsingResult, + ToManager, +) +from airflow.sdk.coordinators._subprocess import _is_connection_from_pid, _start_server +from airflow.sdk.exceptions import AirflowRuntimeError +from airflow.sdk.execution_time import supervisor, task_runner +from airflow.sdk.execution_time.comms import CommsDecoder, ErrorResponse, MaskSecret, _RequestFrame +from airflow.sdk.execution_time.coordinator import get_coordinator_manager +from airflow.sdk.execution_time.supervisor import ( + ResponseSent, + length_prefixed_frame_reader, + make_buffered_socket_reader, + process_log_messages_from_subprocess, + register_request_method, +) + +if TYPE_CHECKING: + from collections.abc import Generator + + from structlog.typing import FilteringBoundLogger + + from airflow.sdk.api.client import Client + from airflow.sdk.execution_time.supervisor import RequestHandler, RequestResult + from airflow.typing_compat import Self + +# How long a runtime may keep running after its parse result, as Node does while a handle stays open. +_EXIT_GRACE_PERIOD = 5.0 + + +# StartTaskHandlerRuntime, LangSDKRuntimeSchemaVersion and LangSDKRuntimeStartFailed pass only between +# the probe and its forked child before the exec, so they are not part of the supervisor schema the +# runtimes speak. + + +class StartTaskHandlerRuntime(BaseModel): + """Ask the parse child to exec the runtime of the coordinator configured under *coordinator*.""" + + file: str + bundle_path: Path + coordinator: str + """The coordinator's key in ``[sdk] coordinators``.""" + comm_address: tuple[str, int] + logs_address: tuple[str, int] + type: Literal["StartTaskHandlerRuntime"] = "StartTaskHandlerRuntime" + + +class LangSDKRuntimeSchemaVersion(BaseModel): + """The schema version and the import timeout of the runtime the parse child is about to exec.""" + + schema_version: str | None + import_timeout: float | None = None + """Seconds from the start of the parse; ``None`` means no timeout.""" + type: Literal["LangSDKRuntimeSchemaVersion"] = "LangSDKRuntimeSchemaVersion" + + +class LangSDKRuntimeStartFailed(BaseModel): + """Why the parse child could not start the runtime.""" + + error: str + type: Literal["LangSDKRuntimeStartFailed"] = "LangSDKRuntimeStartFailed" + + +def _get_import_timeout(path: str) -> float | None: + """Return the ``get_dagbag_import_timeout`` policy's timeout for *path*; ``None`` means none.""" + timeout = settings.get_dagbag_import_timeout(path) + if not isinstance(timeout, (int, float)): + raise TypeError(f"Value ({timeout}) from get_dagbag_import_timeout must be int or float") + return timeout if timeout > 0 else None + + +def _start_task_handler_runtime_entrypoint() -> None: + """Exec the runtime that probes the artifact named by the start request, or report why it cannot start.""" + os.environ["_AIRFLOW_PROCESS_CONTEXT"] = "client" + # fd 0 becomes the runtime's stdin, so the request channel moves to a close-on-exec copy. + comms = CommsDecoder[StartTaskHandlerRuntime, LangSDKRuntimeSchemaVersion | LangSDKRuntimeStartFailed]( + socket=socket(fileno=os.dup(0)), + body_decoder=TypeAdapter(StartTaskHandlerRuntime), + ) + devnull = os.open(os.devnull, os.O_RDONLY) + os.dup2(devnull, 0) + os.close(devnull) + + msg = comms._get_response() + if not isinstance(msg, StartTaskHandlerRuntime): + raise RuntimeError(f"Required first message to be a StartTaskHandlerRuntime, it was {msg}") + + def report_schema_version(schema_version: str | None) -> None: + comms.send(LangSDKRuntimeSchemaVersion(schema_version=schema_version, import_timeout=import_timeout)) + + try: + # The policy is user code: it runs in this child, where a failure is only this file's import error. + import_timeout = _get_import_timeout(msg.file) + get_coordinator_manager().get_coordinator(msg.coordinator).parse_task_handler( + path=Path(msg.file), + bundle_path=msg.bundle_path, + comm_address=msg.comm_address, + logs_address=msg.logs_address, + report_schema_version=report_schema_version, + ) + except Exception as e: + comms.send(LangSDKRuntimeStartFailed(error=f"{type(e).__name__}: {e}")) + + +_Channel = Literal["comm", "logs"] + + +# The requests relayed to the Dag processor without a client: those it answers, but not the Dag parse's +# result. MaskSecret is not relayed: its handler masks the secret here, and mask_secret sends it on to a +# parent. +_PARENT_REQUESTS = frozenset(BaseDagFileProcessorProcess._request_handlers) - { + DagFileParsingResult, + MaskSecret, +} + + +@attrs.define(kw_only=True) +class LangSDKTaskHandlerProcessorProcess(BaseDagFileProcessorProcess): + """ + Ask a coordinator's Lang-SDK runtime for every task handler an artifact registers. + + The forked parse child finds the coordinator, reports the runtime's schema version and execs the + runtime. The runtime connects back to two listeners this process owns and answers the + ``TaskHandlerParseRequest`` itself, so the request is sent once it has connected. A failed start, + a missing result, an invalid frame or message, or a timeout is an import error on the result, + keyed by the artifact's path in its Dag bundle. + + The runtime's requests are answered with :attr:`client`. Without one, as in a Dag-parsing child, + they are relayed up the supervisor channel of the process this runs in, or get an error when there + is none. On Linux the runtime is killed when the thread that started this process exits, so start + it from a thread that outlives the probe. + """ + + client: Client | None = None # type: ignore[assignment] + """Answers the runtime's requests; without one, they are relayed to the parent process if there is one.""" + + parsing_result: TaskHandlerParsingResult | None = None # type: ignore[assignment] + + decoder = TypeAdapter( + Annotated[ + LangSDKRuntimeSchemaVersion | LangSDKRuntimeStartFailed | get_args(ToManager)[0], + Field(discriminator="type"), + ] + ) + + coordinator: str + """The coordinator's key in ``[sdk] coordinators``.""" + + _listeners: dict[_Channel, socket] + _parse_request: TaskHandlerParseRequest + _runtime_schema_version: str | None = attrs.field(default=None, init=False) + _import_timeout: float | None = attrs.field(default=None, init=False) + _schema_version_reported: bool = attrs.field(default=False, init=False) + _parsing_result_monotonic: float | None = attrs.field(default=None, init=False) + _unverified_connections: list[tuple[socket, _Channel]] = attrs.field(factory=list, init=False) + + @classmethod + def start( # type: ignore[override] + cls, + *, + coordinator: str, + path: str | os.PathLike[str], + bundle_path: Path, + bundle_name: str, + artifact_rel_path: str, + **kwargs, + ) -> Self: + """ + Start probing the artifact at *path* for every task handler it registers. + + *bundle_path* and *bundle_name* are those of the Dag bundle holding the artifact, and + *artifact_rel_path* is the artifact's path in it. + """ + listeners: dict[_Channel, socket] = {"comm": _start_server(), "logs": _start_server()} + try: + for listener in listeners.values(): + listener.setblocking(False) + parse_request = TaskHandlerParseRequest( + file=os.fspath(path), bundle_path=bundle_path, bundle_name=bundle_name + ) + proc = super().start( + target=_start_task_handler_runtime_entrypoint, + use_exec=supervisor._should_use_exec(), + new_process_group=True, + coordinator=coordinator, + bundle_name=bundle_name, + dag_file_rel_path=artifact_rel_path, + listeners=listeners, + parse_request=parse_request, + **kwargs, + ) + except BaseException: + for listener in listeners.values(): + listener.close() + raise + for channel, listener in listeners.items(): + proc._open_sockets[listener] = f"{channel}-listener" + proc.selector.register( + listener, + selectors.EVENT_READ, + (functools.partial(proc._accept_connection, channel=channel), proc._on_socket_closed), + ) + proc.send_msg( + StartTaskHandlerRuntime( + file=parse_request.file, + bundle_path=bundle_path, + coordinator=coordinator, + comm_address=listeners["comm"].getsockname()[:2], + logs_address=listeners["logs"].getsockname()[:2], + ), + request_id=0, + ) + return proc + + @classmethod + def run( + cls, + *, + coordinator: str, + path: str | os.PathLike[str], + bundle_path: Path, + bundle_name: str, + artifact_rel_path: str, + logger: FilteringBoundLogger, + ) -> TaskHandlerParsingResult: + """ + Probe the artifact at *path* as :meth:`start` does, and wait for the result. + + The artifact's ``get_dagbag_import_timeout`` bounds the probe, and + ``[dag_processor] dag_file_processor_timeout`` until the parse child resolves it. + """ + return cls._run_to_completion( + coordinator=coordinator, + path=path, + bundle_path=bundle_path, + bundle_name=bundle_name, + artifact_rel_path=artifact_rel_path, + logger=logger, + ) + + @classmethod + def _run_to_completion( + cls, *, logger: FilteringBoundLogger, **start_kwargs: Any + ) -> TaskHandlerParsingResult: + """ + Probe outside a caller's selector loop and wait for the result. + + The import timeout the parse child reports bounds the probe, and + ``[dag_processor] dag_file_processor_timeout`` until it is reported. + """ + processor_timeout = conf.getfloat("dag_processor", "dag_file_processor_timeout") + with selectors.DefaultSelector() as selector: + proc = cls.start(id=uuid7(), selector=selector, logger=logger, **start_kwargs) + try: + while not proc.is_ready: + timeout = proc._import_timeout if proc._schema_version_reported else processor_timeout + if timeout is not None and time.monotonic() - proc.start_time > timeout: + # Unlike is_ready, this does not wait for an exited runtime's leftover processes, + # which can hold its sockets open. close() closes them. + proc._time_out(timeout) + break + proc._service_subprocess(max_wait_time=0.1) + except BaseException: + proc._kill_runtime() + raise + finally: + proc.close() + return cast("TaskHandlerParsingResult", proc.parsing_result) + + def _accept_connection(self, listener: socket, *, channel: _Channel) -> bool: + try: + conn, _ = listener.accept() + except (BlockingIOError, InterruptedError): + return True + conn.setblocking(True) + self._unverified_connections.append((conn, channel)) + self._verify_connections() + return True + + def _verify_connections(self) -> None: + """ + Use each accepted connection once it is confirmed to come from the runtime. + + A connection that is not visible yet stays pending and is checked again on the next + ``is_ready`` poll, so the caller's loop never waits here. + """ + pending = [] + for conn, channel in self._unverified_connections: + if channel not in self._listeners: + # The runtime already connected this channel. + conn.close() + continue + try: + owned = _is_connection_from_pid(conn, self.pid) + except OSError: + conn.close() + continue + if not owned: + pending.append((conn, channel)) + continue + self._close_listener(channel) + if channel == "comm": + self._register_comm(conn) + else: + self._register_logs(conn) + self._unverified_connections = pending + + def _close_listener(self, channel: _Channel) -> None: + if (listener := self._listeners.pop(channel, None)) is not None: + self._on_socket_closed(listener) + listener.close() + + def _close_listeners(self) -> None: + """Close the listeners of a runtime that did not connect, and connections never verified.""" + for channel in list(self._listeners): + self._close_listener(channel) + for conn, _ in self._unverified_connections: + conn.close() + self._unverified_connections = [] + + def _register_comm(self, conn: socket) -> None: + self.stdin = conn + self._open_sockets[conn] = "requests" + read_frame, on_close = length_prefixed_frame_reader( + self._handle_valid_requests(), on_close=self._on_socket_closed + ) + + def read_valid_frame(sock: socket) -> bool: + # A frame that does not decode would otherwise escape the caller's selector loop. + try: + return read_frame(sock) + except msgspec.DecodeError as e: + self._fail_on_invalid_message(f"The Lang-SDK runtime sent an invalid frame: {e}") + return False + + self.selector.register(conn, selectors.EVENT_READ, (read_valid_frame, on_close)) + # The parse child reports the version and waits for the reply before it execs the runtime, + # so the version is known here. It is set only now, so the child's messages are not migrated. + self._subprocess_schema_version = self._runtime_schema_version + self.send_msg(self._parse_request, request_id=0) + + def _handle_valid_requests(self) -> Generator[None, _RequestFrame, None]: + """ + Pass each request on to ``handle_requests``, or kill the runtime at one that does not validate. + + ``handle_requests`` would only log such a request, and the runtime would wait for a reply. The + runtime speaks ``ToManager`` only; the start messages come from the parse child. + """ + requests = self.handle_requests(self.process_log) + next(requests) + while True: + frame = yield + try: + BaseDagFileProcessorProcess.decoder.validate_python(self._deserialize_request(frame.body)) + except ValueError as e: + self._fail_on_invalid_message( + f"The Lang-SDK runtime sent a message that does not validate: {e}" + ) + return + requests.send(frame) + + def _fail_on_invalid_message(self, message: str) -> None: + """Kill the runtime; *message* is the import error unless a parse result was already received.""" + if self.parsing_result is None: + self._set_import_error(message) + else: + self.process_log.warning( + "Ignoring an invalid message from the Lang-SDK runtime after its parse result", error=message + ) + self._kill_runtime() + + def _register_logs(self, conn: socket) -> None: + self._open_sockets[conn] = "logs" + self.selector.register( + conn, + selectors.EVENT_READ, + make_buffered_socket_reader( + process_log_messages_from_subprocess(self._get_target_loggers()), + on_close=self._on_socket_closed, + ), + ) + + def _set_import_error(self, message: str) -> None: + self.parsing_result = TaskHandlerParsingResult( + fileloc=self._parse_request.file, + task_handlers={}, + import_errors={self.dag_file_rel_path: message}, + ) + + def _handle_runtime_schema_version( + self, msg: LangSDKRuntimeSchemaVersion, log: FilteringBoundLogger, req_id: int + ) -> RequestResult | ResponseSent: + if self._schema_version_reported: + self._reject_request(msg, log, req_id) + return ResponseSent.ALREADY_SENT + self._runtime_schema_version = msg.schema_version + self._import_timeout = msg.import_timeout + self._schema_version_reported = True + return None, {} + + def _handle_start_failed( + self, msg: LangSDKRuntimeStartFailed, log: FilteringBoundLogger, req_id: int + ) -> RequestResult | ResponseSent: + self._set_import_error(f"Cannot start the Lang-SDK runtime: {msg.error}") + return None, {} + + def _handle_parsing_result( # type: ignore[override] + self, msg: TaskHandlerParsingResult, log: FilteringBoundLogger, req_id: int + ) -> RequestResult | ResponseSent: + if self.parsing_result is not None: + log.warning("Ignoring another parse result from the Lang-SDK runtime", fileloc=msg.fileloc) + self.send_msg( + None, + request_id=req_id, + error=ErrorResponse(detail={"message": "A parse result was already received"}), + ) + return ResponseSent.ALREADY_SENT + self.parsing_result = msg + self._parsing_result_monotonic = time.monotonic() + return None, {} + + _request_handlers: ClassVar[dict[type[BaseModel], RequestHandler[Any]]] = { + **{ + message_type: handler + for message_type, handler in BaseDagFileProcessorProcess._request_handlers.items() + if message_type is not DagFileParsingResult + }, + **dict( + [ + register_request_method(LangSDKRuntimeSchemaVersion, _handle_runtime_schema_version), + register_request_method(LangSDKRuntimeStartFailed, _handle_start_failed), + register_request_method(TaskHandlerParsingResult, _handle_parsing_result), + ] + ), + } + + def _handle_request(self, msg, log: FilteringBoundLogger, req_id: int) -> None: + if self.client is None and type(msg) in _PARENT_REQUESTS: + self._relay_request(msg, req_id) + return + super()._handle_request(msg, log, req_id) + + def _relay_request(self, msg: BaseModel, req_id: int) -> None: + """Answer the runtime's request through the supervisor channel of this process, if it has one.""" + comms = getattr(task_runner, "SUPERVISOR_COMMS", None) + if comms is None: + self.send_msg( + None, + request_id=req_id, + error=ErrorResponse( + detail={"message": f"{type(msg).__name__} is answered only in the Dag processor"} + ), + ) + return + try: + response = comms.send(msg) + except AirflowRuntimeError as e: + self.send_msg(None, request_id=req_id, error=e.error) + return + # Only the fields the parent sent, under their wire names, so the runtime gets the same body. + self.send_msg(response, request_id=req_id, exclude_unset=True, by_alias=True) + + @property + def is_ready(self) -> bool: + self._verify_connections() + if ( + self._parsing_result_monotonic is not None + and self._exit_code is None + and time.monotonic() - self._parsing_result_monotonic > _EXIT_GRACE_PERIOD + ): + self.process_log.warning("The Lang-SDK runtime did not exit after its parse result; killing it") + self._kill_runtime() + if ( + self._import_timeout is not None + and self.parsing_result is None + and self._exit_code is None + and time.monotonic() - self.start_time > self._import_timeout + ): + self._time_out(self._import_timeout) + if self._check_subprocess_exit() is None: + return False + self._close_listeners() + if not super().is_ready: + return False + if self.parsing_result is None: + self._set_import_error( + f"The Lang-SDK runtime exited with code {self._exit_code} without a parse result" + ) + return True + + def _time_out(self, timeout: float) -> None: + if self.parsing_result is None: + self._set_import_error( + f"The Lang-SDK runtime did not parse {self._parse_request.file} within {timeout}s" + ) + self._kill_runtime() + + def _kill_runtime(self) -> None: + """Kill the runtime and wait for it, without servicing its sockets, whose handler may have failed.""" + if self._exit_code is not None: + return + try: + self._signal_subprocess(signal.SIGKILL) + self._exit_code = self._process.wait(timeout=None) + except (self._process.ProcessNotFound, ProcessLookupError): + self._exit_code = -1 + + def close(self) -> None: + # A listener has nothing to drain, and cleanup would call its accept handler forever. + self._close_listeners() + super().close() diff --git a/airflow-core/tests/unit/dag_processing/fake_task_handler_runtime.py b/airflow-core/tests/unit/dag_processing/fake_task_handler_runtime.py new file mode 100644 index 0000000000000..151d30d4a8dfb --- /dev/null +++ b/airflow-core/tests/unit/dag_processing/fake_task_handler_runtime.py @@ -0,0 +1,120 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""A coordinator whose artifacts name the command that parses them, and a runtime a test can play.""" + +from __future__ import annotations + +import contextlib +import json +import socket +from pathlib import Path +from typing import TYPE_CHECKING, Any +from unittest import mock + +import attrs +from pydantic import TypeAdapter + +from airflow.dag_processing.processor import ( + TaskHandlerParseRequest, + TaskHandlerParsingResult, + ToManager, + ToSDKTaskHandlerProcessor, +) +from airflow.sdk.coordinators._subprocess import SubprocessCoordinator +from airflow.sdk.execution_time import supervisor +from airflow.sdk.execution_time.comms import CommsDecoder +from airflow.sdk.execution_time.coordinator import reset_coordinator_manager + +from tests_common.test_utils.config import conf_vars + +if TYPE_CHECKING: + from collections.abc import Callable, Iterator, Sequence + + +@attrs.define(kw_only=True) +class FakeCoordinator(SubprocessCoordinator): + """ + The artifact's JSON names the command that parses it. + + ``argv`` is the command, ``schema_version`` its schema version, and ``command_error`` an error to + raise instead. + """ + + def _build_parse_task_handler_command(self, *, path: Path) -> tuple[list[str], str | None]: + spec = json.loads(path.read_text()) + if error := spec.get("command_error"): + raise FileNotFoundError(error) + return spec.get("argv", ["/bin/false"]), spec.get("schema_version") + + +@contextlib.contextmanager +def fake_coordinator(**kwargs: Any) -> Iterator[None]: + """ + Configure a ``FakeCoordinator`` as ``fake``, with fresh coordinators inside and after the block. + + The parse child is a bare fork even on macOS, so it sees the test's patches and config. + """ + spec = {"fake": {"classpath": f"{__name__}.FakeCoordinator", "kwargs": kwargs}} + reset_coordinator_manager() + try: + with ( + conf_vars({("sdk", "coordinators"): json.dumps(spec)}), + mock.patch.object(supervisor, "_should_use_exec", return_value=False), + ): + yield + finally: + reset_coordinator_manager() + + +def write_artifact(path: Path, **spec: Any) -> Path: + path.write_text(json.dumps(spec)) + return path + + +def play_runtime( + reply: Callable[[TaskHandlerParseRequest, CommsDecoder], TaskHandlerParsingResult | None], + *, + schema_version: str | None = None, + log_lines: Sequence[dict[str, Any]] = (), +) -> Callable[..., None]: + """ + Return a ``parse_task_handler`` that plays the runtime in the parse child instead of exec'ing one. + + It connects back as a runtime does, writes *log_lines* to its logs channel and answers the parse + request with what *reply* returns; ``None`` sends no result. + """ + + def parse_task_handler( + coordinator, *, comm_address, logs_address, report_schema_version, **kwargs + ) -> None: + report_schema_version(schema_version) + with ( + socket.create_connection(comm_address) as comm, + socket.create_connection(logs_address) as logs, + ): + for line in log_lines: + logs.sendall(json.dumps(line).encode() + b"\n") + comms = CommsDecoder[ToSDKTaskHandlerProcessor, ToManager]( + socket=comm, body_decoder=TypeAdapter(ToSDKTaskHandlerProcessor) + ) + request = comms._get_response() + assert isinstance(request, TaskHandlerParseRequest) + if (result := reply(request, comms)) is not None: + comms.send(result) + + return parse_task_handler diff --git a/airflow-core/tests/unit/dag_processing/test_task_handler_processor.py b/airflow-core/tests/unit/dag_processing/test_task_handler_processor.py new file mode 100644 index 0000000000000..4a79c115588aa --- /dev/null +++ b/airflow-core/tests/unit/dag_processing/test_task_handler_processor.py @@ -0,0 +1,777 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import io +import os +import selectors +import signal +import socket +import sys +import threading +import time +import uuid +from pathlib import Path +from unittest.mock import ANY, MagicMock, patch + +import psutil +import pytest +import structlog +from pydantic import TypeAdapter +from structlog.typing import FilteringBoundLogger + +from airflow.configuration import conf +from airflow.dag_processing.processor import ( + BaseDagFileProcessorProcess, + DagFileParseRequest, + DagFileParsingResult, + DagFileProcessorProcess, + TaskHandlerDeclaration, + TaskHandlerParam, + TaskHandlerParseRequest, + TaskHandlerParsingResult, + ToDagProcessor, + ToManager, +) +from airflow.dag_processing.task_handler_processor import ( + LangSDKRuntimeSchemaVersion, + LangSDKTaskHandlerProcessorProcess, + _get_import_timeout, +) +from airflow.sdk.api.client import Client, VariableOperations +from airflow.sdk.api.datamodels._generated import VariableResponse +from airflow.sdk.coordinators._subprocess import SubprocessCoordinator +from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType +from airflow.sdk.execution_time import supervisor, task_runner +from airflow.sdk.execution_time.comms import ( + CommsDecoder, + ErrorResponse, + GetVariable, + MaskSecret, + VariableResult, + _RequestFrame, +) + +from tests_common.test_utils.config import conf_vars +from unit.dag_processing.fake_task_handler_runtime import ( + FakeCoordinator, + fake_coordinator, + play_runtime, + write_artifact, +) + +# The oldest supervisor schema version, so the parse request is downgraded. +OLDEST_SCHEMA_VERSION = "2026-06-16" + + +def _reply_with(*task_ids: str, **result): + """Reply with a handler for each of *task_ids* under the Dag ``etl``.""" + + def reply(request: TaskHandlerParseRequest, comms) -> TaskHandlerParsingResult: + declarations = [ + TaskHandlerDeclaration(task_id=task_id, binding="positional", params=[]) for task_id in task_ids + ] + return TaskHandlerParsingResult(fileloc=request.file, task_handlers={"etl": declarations}, **result) + + return reply + + +def _get_task_ids(result: TaskHandlerParsingResult) -> list[str]: + return [declaration.task_id for declaration in result.task_handlers["etl"]] + + +def _get_open_fds() -> set[int]: + # Without /proc, as on macOS, this is empty, so the fd leak checks pass trivially. + return {int(fd) for fd in os.listdir("/proc/self/fd")} if os.path.isdir("/proc/self/fd") else set() + + +@pytest.fixture(autouse=True) +def _coordinator(): + with fake_coordinator(): + yield + + +@pytest.fixture +def supervisor_comms(monkeypatch): + """Give this process a supervisor channel, as a Dag-parsing child has.""" + comms = MagicMock(spec=CommsDecoder) + monkeypatch.setattr(task_runner, "SUPERVISOR_COMMS", comms, raising=False) + return comms + + +def _start(tmp_path, selector, *, client: Client | None = None, **spec) -> LangSDKTaskHandlerProcessorProcess: + return LangSDKTaskHandlerProcessorProcess.start( + id=uuid.uuid4(), + coordinator="fake", + path=write_artifact(tmp_path / "etl.artifact", **spec), + bundle_path=tmp_path, + bundle_name="task-handlers", + artifact_rel_path="etl.artifact", + selector=selector, + logger=structlog.get_logger(), + client=client, + ) + + +@pytest.fixture +def parse(tmp_path): + """Probe ``etl.artifact`` under a caller's selector loop, and check that nothing is left open.""" + + def _parse(**kwargs) -> LangSDKTaskHandlerProcessorProcess: + fds_before = _get_open_fds() + with selectors.DefaultSelector() as selector: + proc = _start(tmp_path, selector, **kwargs) + deadline = time.monotonic() + 30 + while not proc.is_ready: + assert time.monotonic() < deadline, "the Lang-SDK parse did not finish" + proc._service_subprocess(max_wait_time=0.1) + assert selector.get_map() == {} + proc.close() + assert _get_open_fds() <= fds_before + return proc + + return _parse + + +def _block_until_killed(comms) -> None: + """Hang as a stuck runtime does, until it is killed or the parent closes the comm socket.""" + comms.socket.recv(1) + + +def _send_an_invalid_frame(request, comms) -> None: + comms.socket.sendall(bytes.fromhex("00000003c1c1c1")) + _block_until_killed(comms) + + +def _send_a_frame(comms, body: dict) -> None: + comms.socket.sendall(_RequestFrame(id=1, body=body).as_bytes()) + _block_until_killed(comms) + + +def _send_an_invalid_result(request, comms) -> None: + _send_a_frame( + comms, {"type": "TaskHandlerParsingResult", "fileloc": request.file, "task_handlers": "none"} + ) + + +def _send_a_start_message(request, comms) -> None: + _send_a_frame(comms, LangSDKRuntimeSchemaVersion(schema_version=None).model_dump()) + + +def _run(tmp_path, *, coordinator: str = "fake", **spec) -> TaskHandlerParsingResult: + return LangSDKTaskHandlerProcessorProcess.run( + coordinator=coordinator, + path=write_artifact(tmp_path / "etl.artifact", **spec), + bundle_path=tmp_path, + bundle_name="task-handlers", + artifact_rel_path="etl.artifact", + logger=structlog.get_logger(), + ) + + +class TestLangSDKTaskHandlerProcessorProcess: + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_collects_the_result_the_runtime_returns( + self, mock_parse_task_handler, parse, tmp_path, cap_structlog + ): + mock_parse_task_handler.side_effect = play_runtime( + _reply_with("extract"), + schema_version=OLDEST_SCHEMA_VERSION, + log_lines=[{"event": "Registering handlers", "level": "info"}], + ) + + proc = parse() + + assert proc.parsing_result.fileloc == os.fspath(tmp_path / "etl.artifact") + assert proc.parsing_result.import_errors is None + assert _get_task_ids(proc.parsing_result) == ["extract"] + assert proc._subprocess_schema_version == OLDEST_SCHEMA_VERSION + assert "Registering handlers" in cap_structlog + + @patch("airflow.dag_processing.task_handler_processor._is_connection_from_pid", autospec=True) + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_connection_is_used_once_it_is_verified(self, mock_parse_task_handler, mock_owned, tmp_path): + mock_parse_task_handler.side_effect = play_runtime(_reply_with("extract")) + mock_owned.return_value = False + with selectors.DefaultSelector() as selector: + proc = _start(tmp_path, selector) + child_stdin = proc.stdin + deadline = time.monotonic() + 30 + while len(proc._unverified_connections) < 2: + assert proc.stdin is child_stdin, "an unverified connection was used" + assert time.monotonic() < deadline, "the runtime did not connect" + proc._service_subprocess(max_wait_time=0.1) + assert proc.stdin is child_stdin + + mock_owned.return_value = True + while not proc.is_ready: + assert time.monotonic() < deadline, "the Lang-SDK parse did not finish" + proc._service_subprocess(max_wait_time=0.1) + proc.close() + + assert _get_task_ids(proc.parsing_result) == ["extract"] + + @patch.object(BaseDagFileProcessorProcess, "start", autospec=True, side_effect=OSError("fork failed")) + def test_a_start_that_fails_before_the_fork_closes_its_listeners(self, mock_start, tmp_path): + fds_before = _get_open_fds() + + with selectors.DefaultSelector() as selector: + with pytest.raises(OSError, match="fork failed"): + _start(tmp_path, selector) + + listeners = mock_start.call_args.kwargs["listeners"] + assert [listener.fileno() for listener in listeners.values()] == [-1, -1] + assert _get_open_fds() <= fds_before + + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_requests_are_answered_by_the_client(self, mock_parse_task_handler, parse): + def reply(request, comms): + variable = comms.send(GetVariable(key="probe_var")) + return _reply_with(variable.value)(request, comms) + + mock_parse_task_handler.side_effect = play_runtime(reply) + client = MagicMock(spec=Client) + client.variables = MagicMock(spec=VariableOperations) + client.variables.get.return_value = VariableResponse(key="probe_var", value="from-the-client") + + proc = parse(client=client) + + assert _get_task_ids(proc.parsing_result) == ["from-the-client"] + + @pytest.mark.parametrize( + ("spec", "reply", "error"), + [ + pytest.param( + {"command_error": "no runtime"}, + None, + "Cannot start the Lang-SDK runtime: FileNotFoundError: no runtime", + id="command-not-resolved", + ), + pytest.param( + {"argv": ["/no/such/runtime"]}, + None, + "Cannot start the Lang-SDK runtime: FileNotFoundError: " + "[Errno 2] No such file or directory: '/no/such/runtime'", + id="exec-failed", + ), + pytest.param( + {"argv": ["/bin/sh", "-c", "exit 3"]}, + None, + "The Lang-SDK runtime exited with code 3 without a parse result", + id="exits-before-connecting", + ), + pytest.param( + {}, + lambda request, comms: None, + "The Lang-SDK runtime exited with code 0 without a parse result", + id="exits-without-a-result", + ), + pytest.param( + {}, + _send_an_invalid_frame, + "The Lang-SDK runtime sent an invalid frame: MessagePack data is malformed: " + "invalid opcode '\\xc1' (byte 0)", + id="invalid-frame", + ), + ], + ) + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_failed_parse_is_an_import_error( + self, mock_parse_task_handler, parse, tmp_path, spec, reply, error + ): + mock_parse_task_handler.side_effect = ( + play_runtime(reply) if reply else SubprocessCoordinator.parse_task_handler + ) + + proc = parse(**spec) + + assert proc.parsing_result == TaskHandlerParsingResult( + fileloc=os.fspath(tmp_path / "etl.artifact"), + task_handlers={}, + import_errors={"etl.artifact": error}, + ) + + def test_a_coordinator_that_is_not_configured_is_an_import_error(self, tmp_path): + result = _run(tmp_path, coordinator="missing") + + assert result == TaskHandlerParsingResult( + fileloc=os.fspath(tmp_path / "etl.artifact"), + task_handlers={}, + import_errors={ + "etl.artifact": "Cannot start the Lang-SDK runtime: " + "InvalidCoordinatorError: No coordinator 'missing' in [sdk] coordinators" + }, + ) + + @patch.object( + FakeCoordinator, + "_build_parse_task_handler_command", + SubprocessCoordinator._build_parse_task_handler_command, + ) + def test_a_coordinator_that_does_not_parse_task_handlers_is_an_import_error(self, parse): + proc = parse() + + assert proc.parsing_result.import_errors == { + "etl.artifact": "Cannot start the Lang-SDK runtime: " + "NotImplementedError: FakeCoordinator does not parse task handlers" + } + + @pytest.mark.parametrize( + ("reply", "error"), + [ + pytest.param( + _send_an_invalid_result, + "TaskHandlerParsingResult.task_handlers\n Input should be a valid dictionary", + id="result", + ), + pytest.param( + _send_a_start_message, "does not match any of the expected tags", id="start-message" + ), + ], + ) + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_message_that_does_not_validate_is_an_import_error( + self, mock_parse_task_handler, parse, reply, error + ): + mock_parse_task_handler.side_effect = play_runtime(reply) + + proc = parse() + + [message] = proc.parsing_result.import_errors.values() + assert message.startswith("The Lang-SDK runtime sent a message that does not validate: ") + assert error in message + assert proc._exit_code == -signal.SIGKILL + + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_dag_file_parsing_result_is_not_a_task_handler_result(self, mock_parse_task_handler, tmp_path): + def reply(request, comms): + with pytest.raises(AirflowRuntimeError, match="Unhandled request"): + comms.send(DagFileParsingResult(fileloc=request.file, serialized_dags=[])) + + mock_parse_task_handler.side_effect = play_runtime(reply) + + result = _run(tmp_path) + + assert result.import_errors == { + "etl.artifact": "The Lang-SDK runtime exited with code 0 without a parse result" + } + + @patch.object( + FakeCoordinator, "parse_task_handler", autospec=True, side_effect=play_runtime(_send_an_invalid_frame) + ) + def test_killing_the_runtime_is_not_reported_as_out_of_memory( + self, mock_parse_task_handler, parse, cap_structlog + ): + proc = parse() + + assert proc._exit_code == -signal.SIGKILL + assert not any("Likely out of memory" in str(entry.get("event")) for entry in cap_structlog.entries) + + @patch("airflow.dag_processing.task_handler_processor._EXIT_GRACE_PERIOD", 0.5) + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_runtime_that_runs_on_after_its_result_is_killed( + self, mock_parse_task_handler, parse, cap_structlog + ): + def reply(request, comms): + comms.send(_reply_with("extract")(request, comms)) + _block_until_killed(comms) + + mock_parse_task_handler.side_effect = play_runtime(reply) + + proc = parse() + + assert _get_task_ids(proc.parsing_result) == ["extract"] + assert proc._exit_code == -signal.SIGKILL + assert "The Lang-SDK runtime did not exit after its parse result; killing it" in cap_structlog + + @pytest.mark.parametrize( + ("policy", "error"), + [ + pytest.param(RuntimeError("policy bug"), "RuntimeError: policy bug", id="raises"), + pytest.param( + lambda path: "30", + "TypeError: Value (30) from get_dagbag_import_timeout must be int or float", + id="not-a-number", + ), + ], + ) + @patch("airflow.settings.get_dagbag_import_timeout", autospec=True) + def test_a_failing_import_timeout_policy_is_an_import_error(self, mock_timeout, parse, policy, error): + mock_timeout.side_effect = policy + + proc = parse() + + assert proc.parsing_result.import_errors == { + "etl.artifact": f"Cannot start the Lang-SDK runtime: {error}" + } + + @pytest.mark.skipif(not Path("/proc/self/fd").is_dir(), reason="reads /proc") + @pytest.mark.parametrize("use_exec", [False, True], ids=["fork", "spawn"]) + def test_the_runtime_inherits_only_its_standard_streams(self, monkeypatch, tmp_path, use_exec): + if use_exec: + # The spawned interpreter finds the coordinator again from its environment. + monkeypatch.setattr(supervisor, "_should_use_exec", lambda: True) + monkeypatch.setenv("PYTHONPATH", os.pathsep.join(sys.path)) + monkeypatch.setenv("AIRFLOW__SDK__COORDINATORS", conf.get("sdk", "coordinators")) + with selectors.DefaultSelector() as selector: + proc = _start(tmp_path, selector, argv=["/bin/sh", "-c", "exec sleep 30"]) + deadline = time.monotonic() + 30 + while psutil.Process(proc.pid).name() != "sleep": + assert time.monotonic() < deadline, "the runtime did not start" + proc._service_subprocess(max_wait_time=0.1) + fd_dir = Path(f"/proc/{proc.pid}/fd") + fds = {fd.name: os.readlink(fd) for fd in fd_dir.iterdir()} + proc.kill(signal.SIGKILL) + proc.close() + + assert sorted(fds) == ["0", "1", "2"] + assert fds["0"] == "/dev/null" + + +def _probe_from_a_dag_parsing_child() -> None: + """Stand in for ``_parse_file_entrypoint``: probe the file it is asked to parse, and return the result.""" + comms_decoder = CommsDecoder[ToDagProcessor, ToManager](body_decoder=TypeAdapter(ToDagProcessor)) + request = comms_decoder._get_response() + assert isinstance(request, DagFileParseRequest) + task_runner.SUPERVISOR_COMMS = comms_decoder # type: ignore[assignment] + + result = LangSDKTaskHandlerProcessorProcess.run( + coordinator="fake", + path=request.file, + bundle_path=request.bundle_path, + bundle_name=request.bundle_name, + artifact_rel_path="etl.artifact", + logger=structlog.get_logger(logger_name="task"), + ) + comms_decoder.send( + DagFileParsingResult( + fileloc=request.file, serialized_dags=[], warnings=[result.model_dump(mode="json")] + ) + ) + + +class TestRequestsWithoutAClient: + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_request_is_relayed_to_the_parent(self, mock_parse_task_handler, parse, supervisor_comms): + def reply(request, comms): + variable = comms.send(GetVariable(key="probe_var")) + return _reply_with(variable.value)(request, comms) + + mock_parse_task_handler.side_effect = play_runtime(reply, schema_version=OLDEST_SCHEMA_VERSION) + # As decoded from the parent's frame, which always carries its type. + supervisor_comms.send.return_value = VariableResult.model_validate( + {"key": "probe_var", "value": "from-the-parent", "type": "VariableResult"} + ) + + proc = parse() + + assert _get_task_ids(proc.parsing_result) == ["from-the-parent"] + supervisor_comms.send.assert_called_once_with(GetVariable(key="probe_var")) + + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_the_parent_s_error_reaches_the_runtime(self, mock_parse_task_handler, parse, supervisor_comms): + def reply(request, comms): + with pytest.raises(AirflowRuntimeError) as ctx: + comms.send(GetVariable(key="probe_var")) + error = ctx.value.error + return _reply_with(f"{error.error.value}:{error.detail['key']}")(request, comms) + + mock_parse_task_handler.side_effect = play_runtime(reply) + supervisor_comms.send.side_effect = AirflowRuntimeError( + ErrorResponse(error=ErrorType.VARIABLE_NOT_FOUND, detail={"key": "probe_var"}) + ) + + proc = parse() + + assert _get_task_ids(proc.parsing_result) == [f"{ErrorType.VARIABLE_NOT_FOUND.value}:probe_var"] + + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_request_gets_an_error_without_a_parent(self, mock_parse_task_handler, monkeypatch, tmp_path): + monkeypatch.delattr(task_runner, "SUPERVISOR_COMMS", raising=False) + + def reply(request, comms): + with pytest.raises(AirflowRuntimeError) as ctx: + comms.send(GetVariable(key="probe_var")) + return _reply_with(ctx.value.error.detail["message"])(request, comms) + + mock_parse_task_handler.side_effect = play_runtime(reply) + + result = _run(tmp_path) + + assert _get_task_ids(result) == ["GetVariable is answered only in the Dag processor"] + + @patch("airflow.sdk.log._secrets_masker", autospec=True) + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_secret_is_masked_here_and_by_the_parent( + self, mock_parse_task_handler, mock_secrets_masker, parse, supervisor_comms + ): + def reply(request, comms): + comms.send(MaskSecret(value="probe-secret", name="probe_conn")) + return _reply_with("extract")(request, comms) + + mock_parse_task_handler.side_effect = play_runtime(reply) + + proc = parse() + + assert _get_task_ids(proc.parsing_result) == ["extract"] + mock_secrets_masker.return_value.add_mask.assert_called_once_with("probe-secret", "probe_conn") + supervisor_comms.send.assert_called_once_with(MaskSecret(value="probe-secret", name="probe_conn")) + + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_runtime_request_is_answered_by_the_dag_processor(self, mock_parse_task_handler, tmp_path): + def reply(request, comms): + variable = comms.send(GetVariable(key="probe_var")) + return _reply_with(variable.value)(request, comms) + + mock_parse_task_handler.side_effect = play_runtime(reply) + client = MagicMock(spec=Client) + client.variables = MagicMock(spec=VariableOperations) + client.variables.get.return_value = VariableResponse(key="probe_var", value="from-the-dag-processor") + artifact = write_artifact(tmp_path / "etl.artifact") + + proc = DagFileProcessorProcess.start( + id=1, + path=artifact, + bundle_path=tmp_path, + bundle_name="task-handlers", + dag_file_rel_path="etl.artifact", + callbacks=[], + target=_probe_from_a_dag_parsing_child, + logger=structlog.get_logger(), + logger_filehandle=io.BytesIO(), + client=client, + ) + deadline = time.monotonic() + 30 + while not proc.is_ready: + assert time.monotonic() < deadline, "the Dag-parsing child did not finish" + proc._service_subprocess(max_wait_time=0.1) + proc.close() + + client.variables.get.assert_called_once_with("probe_var") + [probe_result] = proc.parsing_result.warnings + assert TaskHandlerParsingResult.model_validate(probe_result) == TaskHandlerParsingResult( + fileloc=os.fspath(artifact), + task_handlers={ + "etl": [ + TaskHandlerDeclaration(task_id="from-the-dag-processor", binding="positional", params=[]) + ] + }, + ) + + +class TestRun: + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_returns_every_handler_the_runtime_registers(self, mock_parse_task_handler, tmp_path): + def reply(request, comms): + # Echo what the request carries, so the test can check it. + param = TaskHandlerParam(name=request.bundle_name, value_schema={"type": "string"}) + return TaskHandlerParsingResult( + fileloc=request.file, + task_handlers={ + dag_id: [ + TaskHandlerDeclaration( + task_id=os.fspath(request.bundle_path), binding="positional", params=[param] + ) + ] + for dag_id in ("etl", "report") + }, + ) + + mock_parse_task_handler.side_effect = play_runtime(reply) + + result = _run(tmp_path) + + param = TaskHandlerParam(name="task-handlers", value_schema={"type": "string"}) + declaration = TaskHandlerDeclaration( + task_id=os.fspath(tmp_path), binding="positional", params=[param] + ) + assert result == TaskHandlerParsingResult( + fileloc=os.fspath(tmp_path / "etl.artifact"), + task_handlers={"etl": [declaration], "report": [declaration]}, + ) + + @pytest.mark.execution_timeout(30) + @pytest.mark.parametrize("connected", [False, True], ids=["before-connecting", "after-connecting"]) + @patch("airflow.settings.get_dagbag_import_timeout", autospec=True, return_value=1) + @patch.object( + LangSDKTaskHandlerProcessorProcess, + "close", + autospec=True, + side_effect=LangSDKTaskHandlerProcessorProcess.close, + ) + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_parse_past_the_import_timeout_is_killed( + self, mock_parse_task_handler, mock_close, mock_timeout, tmp_path, connected + ): + mock_parse_task_handler.side_effect = ( + play_runtime(lambda request, comms: _block_until_killed(comms)) + if connected + else SubprocessCoordinator.parse_task_handler + ) + # A runtime that never connects leaves both listeners open when it is killed. + fds_before = _get_open_fds() + + result = _run(tmp_path, argv=["/bin/sh", "-c", "exec sleep 60"]) + + assert result.import_errors == { + "etl.artifact": f"The Lang-SDK runtime did not parse {tmp_path / 'etl.artifact'} within 1.0s" + } + [proc] = [c.args[0] for c in mock_close.call_args_list] + assert proc._exit_code == -9 + assert not proc._open_sockets + assert _get_open_fds() <= fds_before + + @patch("airflow.settings.get_dagbag_import_timeout", autospec=True, return_value=1) + @patch.object( + LangSDKTaskHandlerProcessorProcess, + "close", + autospec=True, + side_effect=LangSDKTaskHandlerProcessorProcess.close, + ) + def test_the_import_timeout_holds_after_the_runtime_exits(self, mock_close, mock_timeout, tmp_path): + # The runtime exits, and the process it leaves behind keeps its output open. + result = _run(tmp_path, argv=["/bin/sh", "-c", "sleep 30 & exit 0"]) + [proc] = [c.args[0] for c in mock_close.call_args_list] + os.killpg(proc.pid, signal.SIGKILL) + + assert result.import_errors == { + "etl.artifact": f"The Lang-SDK runtime did not parse {tmp_path / 'etl.artifact'} within 1.0s" + } + assert proc._exit_code == 0 + assert not proc._open_sockets + + @pytest.mark.execution_timeout(30) + @conf_vars({("dag_processor", "dag_file_processor_timeout"): "1"}) + @patch.object( + FakeCoordinator, + "_build_parse_task_handler_command", + autospec=True, + side_effect=lambda self, *, path: threading.Event().wait(), + ) + def test_the_dag_file_processor_timeout_applies_until_the_import_timeout_is_reported( + self, mock_build_command, tmp_path + ): + result = _run(tmp_path) + + assert result.import_errors == { + "etl.artifact": f"The Lang-SDK runtime did not parse {tmp_path / 'etl.artifact'} within 1.0s" + } + + +@pytest.mark.parametrize(("configured", "expected"), [(30, 30), (0.5, 0.5), (0, None), (-1, None)]) +@patch("airflow.settings.get_dagbag_import_timeout", autospec=True) +def test_only_a_positive_import_timeout_applies(mock_timeout, configured, expected): + mock_timeout.return_value = configured + + assert _get_import_timeout("/b/etl.artifact") == expected + mock_timeout.assert_called_once_with("/b/etl.artifact") + + +def _make_process(**kwargs) -> LangSDKTaskHandlerProcessorProcess: + return LangSDKTaskHandlerProcessorProcess( + id=uuid.uuid4(), + pid=1, + stdin=MagicMock(spec=socket.socket), + process=MagicMock(spec=supervisor.ProcessTracker), + process_log=MagicMock(spec=FilteringBoundLogger), + selector=MagicMock(spec=selectors.BaseSelector), + bundle_name="task-handlers", + dag_file_rel_path="etl.artifact", + coordinator="fake", + listeners={}, + parse_request=TaskHandlerParseRequest( + file="/b/etl.artifact", bundle_path=Path("/b"), bundle_name="task-handlers" + ), + **kwargs, + ) + + +def _build_result(*task_ids: str) -> TaskHandlerParsingResult: + return TaskHandlerParsingResult( + fileloc="/b/etl.artifact", + task_handlers={ + "etl": [ + TaskHandlerDeclaration(task_id=task_id, binding="positional", params=[]) + for task_id in task_ids + ] + }, + ) + + +@patch.object(LangSDKTaskHandlerProcessorProcess, "send_msg", autospec=True) +def test_the_first_parse_result_wins(mock_send_msg): + proc = _make_process() + + proc._handle_request(_build_result("first"), structlog.get_logger(), 1) + proc._handle_request(_build_result("second"), structlog.get_logger(), 2) + + assert _get_task_ids(proc.parsing_result) == ["first"] + assert mock_send_msg.call_args.kwargs["error"].detail == { + "message": "A parse result was already received" + } + + +@pytest.mark.parametrize( + "invalid_frame", + [ + pytest.param(bytes.fromhex("00000003c1c1c1"), id="does-not-decode"), + pytest.param( + _RequestFrame( + id=2, + body={ + "type": "TaskHandlerParsingResult", + "fileloc": "/b/etl.artifact", + "task_handlers": "none", + }, + ).as_bytes(), + id="does-not-validate", + ), + ], +) +@patch.object(LangSDKTaskHandlerProcessorProcess, "_kill_runtime", autospec=True) +@patch.object(LangSDKTaskHandlerProcessorProcess, "send_msg", autospec=True) +def test_an_invalid_message_after_the_parse_result_keeps_it(mock_send_msg, mock_kill_runtime, invalid_frame): + proc = _make_process() + runtime, conn = socket.socketpair() + with runtime, conn: + proc._register_comm(conn) + read_frame, _ = proc.selector.register.call_args.args[2] + runtime.sendall(_RequestFrame(id=1, body=_build_result("extract").model_dump(mode="json")).as_bytes()) + assert read_frame(conn) + runtime.sendall(invalid_frame) + assert not read_frame(conn) + + assert _get_task_ids(proc.parsing_result) == ["extract"] + assert proc.parsing_result.import_errors is None + proc.process_log.warning.assert_called_once_with( + "Ignoring an invalid message from the Lang-SDK runtime after its parse result", error=ANY + ) + mock_kill_runtime.assert_called_once_with(proc) + + +@patch.object(LangSDKTaskHandlerProcessorProcess, "send_msg", autospec=True) +def test_the_schema_version_is_reported_once(mock_send_msg): + proc = _make_process() + + proc._handle_request( + LangSDKRuntimeSchemaVersion(schema_version=OLDEST_SCHEMA_VERSION), structlog.get_logger(), 1 + ) + proc._handle_request(LangSDKRuntimeSchemaVersion(schema_version=None), structlog.get_logger(), 2) + + assert proc._runtime_schema_version == OLDEST_SCHEMA_VERSION + assert mock_send_msg.call_args.kwargs["error"].detail["message"] == "Unhandled request" diff --git a/generated/known_sdk_imports_in_core.txt b/generated/known_sdk_imports_in_core.txt index b906ba258554b..dd4fa33cb2381 100644 --- a/generated/known_sdk_imports_in_core.txt +++ b/generated/known_sdk_imports_in_core.txt @@ -6,6 +6,7 @@ airflow-core/src/airflow/configuration.py::1 airflow-core/src/airflow/dag_processing/dagbag.py::3 airflow-core/src/airflow/dag_processing/manager.py::4 airflow-core/src/airflow/dag_processing/processor.py::14 +airflow-core/src/airflow/dag_processing/task_handler_processor.py::8 airflow-core/src/airflow/exceptions.py::1 airflow-core/src/airflow/executors/base_executor.py::3 airflow-core/src/airflow/jobs/triggerer_job_runner.py::18 From bf185776a4258c420b4dcf4aef248aa6e60bed61 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Fri, 2 Oct 2026 03:36:21 +0000 Subject: [PATCH 46/94] Read a Lang-SDK runtime's frames without blocking the parse loop The frame reader reads a frame's payload in the same callback as its length header. On the blocking comm socket, a runtime that stopped after the header blocked the loop in that read, so the import timeout never fired and the runtime was never killed. Reads now return at once and the reader completes the frame on a later callback; replies still go out on the blocking socket. --- .../dag_processing/task_handler_processor.py | 27 ++++++++++++++++--- .../test_task_handler_processor.py | 18 +++++++++++++ 2 files changed, 42 insertions(+), 3 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/task_handler_processor.py b/airflow-core/src/airflow/dag_processing/task_handler_processor.py index eca802ee00066..54bf7c4f04d8f 100644 --- a/airflow-core/src/airflow/dag_processing/task_handler_processor.py +++ b/airflow-core/src/airflow/dag_processing/task_handler_processor.py @@ -24,7 +24,7 @@ import signal import time from pathlib import Path -from socket import socket +from socket import MSG_DONTWAIT, socket from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Literal, cast, get_args import attrs @@ -144,6 +144,24 @@ def report_schema_version(schema_version: str | None) -> None: _Channel = Literal["comm", "logs"] +class _ReadsWithoutWaiting: + """ + The runtime's comm socket, with reads that return at once instead of waiting for data. + + A runtime that stops in the middle of a frame then cannot block the caller's loop, which keeps + checking the import timeout. Replies are sent on the socket itself, which stays blocking. + """ + + def __init__(self, sock: socket) -> None: + self._sock = sock + + def recv(self, bufsize: int) -> bytes: + return self._sock.recv(bufsize, MSG_DONTWAIT) + + def recv_into(self, buffer: memoryview) -> int: + return self._sock.recv_into(buffer, 0, MSG_DONTWAIT) + + # The requests relayed to the Dag processor without a client: those it answers, but not the Dag parse's # result. MaskSecret is not relayed: its handler masks the secret here, and mask_secret sends it on to a # parent. @@ -365,10 +383,13 @@ def _register_comm(self, conn: socket) -> None: ) def read_valid_frame(sock: socket) -> bool: - # A frame that does not decode would otherwise escape the caller's selector loop. try: - return read_frame(sock) + return read_frame(cast("socket", _ReadsWithoutWaiting(sock))) + except BlockingIOError: + # The rest of the frame has not arrived; the reader keeps what it has read so far. + return True except msgspec.DecodeError as e: + # A frame that does not decode would otherwise escape the caller's selector loop. self._fail_on_invalid_message(f"The Lang-SDK runtime sent an invalid frame: {e}") return False diff --git a/airflow-core/tests/unit/dag_processing/test_task_handler_processor.py b/airflow-core/tests/unit/dag_processing/test_task_handler_processor.py index 4a79c115588aa..288c2876494c6 100644 --- a/airflow-core/tests/unit/dag_processing/test_task_handler_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_task_handler_processor.py @@ -636,6 +636,24 @@ def test_a_parse_past_the_import_timeout_is_killed( assert not proc._open_sockets assert _get_open_fds() <= fds_before + @pytest.mark.execution_timeout(30) + @patch("airflow.settings.get_dagbag_import_timeout", autospec=True, return_value=1) + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_runtime_that_stops_mid_frame_is_timed_out( + self, mock_parse_task_handler, mock_timeout, tmp_path + ): + def reply(request, comms): + comms.socket.sendall((100).to_bytes(4, byteorder="big")) + _block_until_killed(comms) + + mock_parse_task_handler.side_effect = play_runtime(reply) + + result = _run(tmp_path) + + assert result.import_errors == { + "etl.artifact": f"The Lang-SDK runtime did not parse {tmp_path / 'etl.artifact'} within 1.0s" + } + @patch("airflow.settings.get_dagbag_import_timeout", autospec=True, return_value=1) @patch.object( LangSDKTaskHandlerProcessorProcess, From 896cc6faf9b8bf45bc692fec741512a7cc74d88c Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Fri, 2 Oct 2026 03:37:09 +0000 Subject: [PATCH 47/94] Kill what a Lang-SDK runtime leaves behind when it exits A runtime that exited while a child it started still held its output left that child running: the kill is skipped once the runtime has exited, and its process group cannot be looked up after it is reaped. The group is now killed when the runtime is seen to exit, before it is reaped, while its zombie still reserves the group id. A leftover outside the group, which escaped with setsid, still holds the sockets. is_ready now applies the import timeout to that case too, so a caller driving its own selector loop is not stuck waiting for EOF, and the standalone loop keeps only the timeout that applies before the parse child reports one. --- .../dag_processing/task_handler_processor.py | 41 ++++++++-- .../test_task_handler_processor.py | 80 ++++++++++++++----- 2 files changed, 96 insertions(+), 25 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/task_handler_processor.py b/airflow-core/src/airflow/dag_processing/task_handler_processor.py index 54bf7c4f04d8f..c99557b20b626 100644 --- a/airflow-core/src/airflow/dag_processing/task_handler_processor.py +++ b/airflow-core/src/airflow/dag_processing/task_handler_processor.py @@ -18,6 +18,7 @@ from __future__ import annotations +import contextlib import functools import os import selectors @@ -29,6 +30,7 @@ import attrs import msgspec +import psutil from pydantic import BaseModel, Field, TypeAdapter from uuid6 import uuid7 @@ -180,7 +182,8 @@ class LangSDKTaskHandlerProcessorProcess(BaseDagFileProcessorProcess): runtime. The runtime connects back to two listeners this process owns and answers the ``TaskHandlerParseRequest`` itself, so the request is sent once it has connected. A failed start, a missing result, an invalid frame or message, or a timeout is an import error on the result, - keyed by the artifact's path in its Dag bundle. + keyed by the artifact's path in its Dag bundle. Processes the runtime leaves in its process group + are killed when it exits. The runtime's requests are answered with :attr:`client`. Without one, as in a Dag-parsing child, they are relayed up the supervisor channel of the process this runs in, or get an error when there @@ -310,11 +313,12 @@ def _run_to_completion( proc = cls.start(id=uuid7(), selector=selector, logger=logger, **start_kwargs) try: while not proc.is_ready: - timeout = proc._import_timeout if proc._schema_version_reported else processor_timeout - if timeout is not None and time.monotonic() - proc.start_time > timeout: - # Unlike is_ready, this does not wait for an exited runtime's leftover processes, - # which can hold its sockets open. close() closes them. - proc._time_out(timeout) + # is_ready applies the import timeout once the parse child has reported it. + if ( + not proc._schema_version_reported + and time.monotonic() - proc.start_time > processor_timeout + ): + proc._time_out(processor_timeout) break proc._service_subprocess(max_wait_time=0.1) except BaseException: @@ -540,6 +544,14 @@ def is_ready(self) -> bool: if self._check_subprocess_exit() is None: return False self._close_listeners() + if ( + self._open_sockets + and self._import_timeout is not None + and time.monotonic() - self.start_time > self._import_timeout + ): + # A process the runtime left outside its process group holds these open. + self._time_out(self._import_timeout) + self.cleanup_sockets_after_kill() if not super().is_ready: return False if self.parsing_result is None: @@ -555,6 +567,23 @@ def _time_out(self, timeout: float) -> None: ) self._kill_runtime() + def _check_subprocess_exit( + self, raise_on_timeout: bool = False, expect_signal: None | int = None + ) -> int | None: + if self._exit_code is None and self._is_runtime_exited(): + # Until the exited runtime is reaped below, its pid, and so its process group id, cannot be + # reused, so this reaches only the processes it left behind. + with contextlib.suppress(ProcessLookupError, PermissionError): + os.killpg(self.pid, signal.SIGKILL) + return super()._check_subprocess_exit(raise_on_timeout=raise_on_timeout, expect_signal=expect_signal) + + def _is_runtime_exited(self) -> bool: + """Return whether the runtime has exited and is not reaped yet.""" + try: + return psutil.Process(self.pid).status() == psutil.STATUS_ZOMBIE + except psutil.Error: + return False + def _kill_runtime(self) -> None: """Kill the runtime and wait for it, without servicing its sockets, whose handler may have failed.""" if self._exit_code is not None: diff --git a/airflow-core/tests/unit/dag_processing/test_task_handler_processor.py b/airflow-core/tests/unit/dag_processing/test_task_handler_processor.py index 288c2876494c6..6a99f43621436 100644 --- a/airflow-core/tests/unit/dag_processing/test_task_handler_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_task_handler_processor.py @@ -17,6 +17,7 @@ # under the License. from __future__ import annotations +import contextlib import io import os import selectors @@ -100,6 +101,29 @@ def _get_open_fds() -> set[int]: return {int(fd) for fd in os.listdir("/proc/self/fd")} if os.path.isdir("/proc/self/fd") else set() +# A runtime that exits once a child has left its process group; the child keeps the runtime's +# stdout and stderr open and writes its pid to argv[1]. +_LEAVE_A_CHILD_OUTSIDE_THE_GROUP = """ +import os, sys, time +read_fd, write_fd = os.pipe() +if os.fork(): + os.read(read_fd, 1) + os._exit(0) +os.setsid() +with open(sys.argv[1], "w") as f: + f.write(str(os.getpid())) +os.write(write_fd, b"x") +time.sleep(30) +""" + + +def _is_running(pid: int) -> bool: + try: + return psutil.Process(pid).status() != psutil.STATUS_ZOMBIE + except psutil.NoSuchProcess: + return False + + @pytest.fixture(autouse=True) def _coordinator(): with fake_coordinator(): @@ -399,6 +423,43 @@ def reply(request, comms): assert proc._exit_code == -signal.SIGKILL assert "The Lang-SDK runtime did not exit after its parse result; killing it" in cap_structlog + @patch("airflow.settings.get_dagbag_import_timeout", autospec=True, return_value=5) + def test_what_the_runtime_leaves_in_its_process_group_is_killed_when_it_exits( + self, mock_timeout, parse, tmp_path + ): + pid_file = tmp_path / "leftover.pid" + + proc = parse(argv=["/bin/sh", "-c", f"sleep 30 & echo $! > {pid_file}; exit 0"]) + + leftover = int(pid_file.read_text()) + try: + assert proc.parsing_result.import_errors == { + "etl.artifact": "The Lang-SDK runtime exited with code 0 without a parse result" + } + # A poll: the leftover is not a child to wait for, and time_machine cannot reach another process. + deadline = time.monotonic() + 10 + while _is_running(leftover): + assert time.monotonic() < deadline, "the runtime's leftover process is still running" + time.sleep(0.05) + finally: + with contextlib.suppress(ProcessLookupError): + os.kill(leftover, signal.SIGKILL) + + @patch("airflow.settings.get_dagbag_import_timeout", autospec=True, return_value=1) + def test_the_import_timeout_holds_after_the_runtime_exits(self, mock_timeout, parse, tmp_path): + pid_file = tmp_path / "leftover.pid" + + try: + proc = parse(argv=[sys.executable, "-c", _LEAVE_A_CHILD_OUTSIDE_THE_GROUP, os.fspath(pid_file)]) + finally: + if pid_file.exists(): + os.kill(int(pid_file.read_text()), signal.SIGKILL) + + assert proc.parsing_result.import_errors == { + "etl.artifact": f"The Lang-SDK runtime did not parse {tmp_path / 'etl.artifact'} within 1.0s" + } + assert proc._exit_code == 0 + @pytest.mark.parametrize( ("policy", "error"), [ @@ -654,25 +715,6 @@ def reply(request, comms): "etl.artifact": f"The Lang-SDK runtime did not parse {tmp_path / 'etl.artifact'} within 1.0s" } - @patch("airflow.settings.get_dagbag_import_timeout", autospec=True, return_value=1) - @patch.object( - LangSDKTaskHandlerProcessorProcess, - "close", - autospec=True, - side_effect=LangSDKTaskHandlerProcessorProcess.close, - ) - def test_the_import_timeout_holds_after_the_runtime_exits(self, mock_close, mock_timeout, tmp_path): - # The runtime exits, and the process it leaves behind keeps its output open. - result = _run(tmp_path, argv=["/bin/sh", "-c", "sleep 30 & exit 0"]) - [proc] = [c.args[0] for c in mock_close.call_args_list] - os.killpg(proc.pid, signal.SIGKILL) - - assert result.import_errors == { - "etl.artifact": f"The Lang-SDK runtime did not parse {tmp_path / 'etl.artifact'} within 1.0s" - } - assert proc._exit_code == 0 - assert not proc._open_sockets - @pytest.mark.execution_timeout(30) @conf_vars({("dag_processor", "dag_file_processor_timeout"): "1"}) @patch.object( From 1b9321728acce664b86687039e2827e87272f8f6 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Fri, 2 Oct 2026 03:37:43 +0000 Subject: [PATCH 48/94] Do not exec a Lang-SDK runtime whose parent has already exited The parent death signal is only sent for a parent that exits after it is set. A parent that exited between the schema-version report and the prctl call left the runtime with no signal to come and nobody to read its sockets. The child now checks that its parent is unchanged once the signal is set. --- .../src/airflow/sdk/coordinators/_subprocess.py | 11 ++++++++--- .../task_sdk/coordinators/test_subprocess.py | 16 ++++++++++++++++ 2 files changed, 24 insertions(+), 3 deletions(-) diff --git a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py index 9a9b9f5a50e58..6e722a33a86d7 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py +++ b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py @@ -548,10 +548,10 @@ def parse_task_handler( version are resolved against *bundle_path*, and the version is passed to *report_schema_version* just before the exec. The runtime connects back to *comm_address* and *logs_address*. On Linux the runtime dies with the thread that forked - the calling process. + the calling process, and it is not started once that parent has exited. - :raises Exception: when the command cannot be resolved or started; the process is then - unchanged. + :raises Exception: when the command cannot be resolved or started, or the parent has + exited; the process is then unchanged. """ with self._set_scan_roots([bundle_path]): command, schema_version = self._build_parse_task_handler_command(path=path) @@ -562,6 +562,9 @@ def parse_task_handler( f"--comm={comm_address[0]}:{comm_address[1]}", f"--logs={logs_address[0]}:{logs_address[1]}", ] + # Taken before the report, which fails if the parent has already exited. A different parent once the + # death signal is set means the parent exited in between, and the kernel does not send it for that. + parent_pid = os.getppid() report_schema_version(schema_version) # Python ignores these at startup and exec keeps ignored signals; subprocess.Popen resets # them the same way for the task runtime. @@ -569,6 +572,8 @@ def parse_task_handler( if (sig := getattr(signal, name, None)) is not None: signal.signal(sig, signal.SIG_DFL) _set_parent_death_signal() + if os.getppid() != parent_pid: + raise RuntimeError("The process that started the runtime has exited") _set_close_on_exec_above_stderr() os.execvpe(argv[0], argv, _build_runtime_env()) diff --git a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py index 81461785d9fa4..8db8cbb186fd7 100644 --- a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py +++ b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py @@ -1097,6 +1097,22 @@ def test_resolves_the_command_under_the_bundle_root_then_execs_it( mock_death_signal.assert_called_once_with() mock_close_on_exec.assert_called_once_with() + @patch("airflow.sdk.coordinators._subprocess._set_close_on_exec_above_stderr", autospec=True) + @patch("airflow.sdk.coordinators._subprocess._set_parent_death_signal", autospec=True) + @patch("airflow.sdk.coordinators._subprocess.signal.signal", autospec=True) + @patch("airflow.sdk.coordinators._subprocess.os.getppid", autospec=True, side_effect=[4242, 1]) + @patch("airflow.sdk.coordinators._subprocess.os.execvpe", autospec=True) + def test_does_not_exec_once_the_parent_has_exited( + self, mock_execvpe, mock_getppid, mock_signal, mock_death_signal, mock_close_on_exec, tmp_path + ): + coordinator = _TaskHandlerParsingCoordinator(command=["runtime"]) + + with pytest.raises(RuntimeError, match="The process that started the runtime has exited"): + _parse_task_handler(coordinator, tmp_path, []) + + mock_death_signal.assert_called_once_with() + mock_execvpe.assert_not_called() + @patch("airflow.sdk.coordinators._subprocess.os.execvpe", autospec=True) def test_rejects_an_unknown_schema_version_before_reporting(self, mock_execvpe, tmp_path): coordinator = _TaskHandlerParsingCoordinator(command=["runtime"], schema_version="1999-01-01") From b45c6ec8df69b14f3cd5bdaec6ed5e52b6530631 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Fri, 2 Oct 2026 03:38:31 +0000 Subject: [PATCH 49/94] Clean up a Lang-SDK runtime parse whose start fails after the fork Only the fork itself was guarded. A failure while registering the listeners or sending the start request left the child waiting for that request and both listeners open, and the standalone run does not guard start() either. --- .../dag_processing/task_handler_processor.py | 37 +++++++++++-------- .../test_task_handler_processor.py | 18 +++++++++ 2 files changed, 39 insertions(+), 16 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/task_handler_processor.py b/airflow-core/src/airflow/dag_processing/task_handler_processor.py index c99557b20b626..20a60d0a200cd 100644 --- a/airflow-core/src/airflow/dag_processing/task_handler_processor.py +++ b/airflow-core/src/airflow/dag_processing/task_handler_processor.py @@ -253,23 +253,28 @@ def start( # type: ignore[override] for listener in listeners.values(): listener.close() raise - for channel, listener in listeners.items(): - proc._open_sockets[listener] = f"{channel}-listener" - proc.selector.register( - listener, - selectors.EVENT_READ, - (functools.partial(proc._accept_connection, channel=channel), proc._on_socket_closed), + try: + for channel, listener in listeners.items(): + proc._open_sockets[listener] = f"{channel}-listener" + proc.selector.register( + listener, + selectors.EVENT_READ, + (functools.partial(proc._accept_connection, channel=channel), proc._on_socket_closed), + ) + proc.send_msg( + StartTaskHandlerRuntime( + file=parse_request.file, + bundle_path=bundle_path, + coordinator=coordinator, + comm_address=listeners["comm"].getsockname()[:2], + logs_address=listeners["logs"].getsockname()[:2], + ), + request_id=0, ) - proc.send_msg( - StartTaskHandlerRuntime( - file=parse_request.file, - bundle_path=bundle_path, - coordinator=coordinator, - comm_address=listeners["comm"].getsockname()[:2], - logs_address=listeners["logs"].getsockname()[:2], - ), - request_id=0, - ) + except BaseException: + proc._kill_runtime() + proc.close() + raise return proc @classmethod diff --git a/airflow-core/tests/unit/dag_processing/test_task_handler_processor.py b/airflow-core/tests/unit/dag_processing/test_task_handler_processor.py index 6a99f43621436..6b39cfd89bcc5 100644 --- a/airflow-core/tests/unit/dag_processing/test_task_handler_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_task_handler_processor.py @@ -250,6 +250,24 @@ def test_a_connection_is_used_once_it_is_verified(self, mock_parse_task_handler, assert _get_task_ids(proc.parsing_result) == ["extract"] + @patch.object( + LangSDKTaskHandlerProcessorProcess, + "send_msg", + autospec=True, + side_effect=OSError("cannot send the start request"), + ) + def test_a_start_that_fails_after_the_fork_leaves_nothing_behind(self, mock_send_msg, tmp_path): + fds_before = _get_open_fds() + children_before = {child.pid for child in psutil.Process().children()} + + with selectors.DefaultSelector() as selector: + with pytest.raises(OSError, match="cannot send the start request"): + _start(tmp_path, selector) + assert selector.get_map() == {} + + assert {child.pid for child in psutil.Process().children()} == children_before + assert _get_open_fds() <= fds_before + @patch.object(BaseDagFileProcessorProcess, "start", autospec=True, side_effect=OSError("fork failed")) def test_a_start_that_fails_before_the_fork_closes_its_listeners(self, mock_start, tmp_path): fds_before = _get_open_fds() From aeff5da3dc4c8287453217d73987399eb76cbc50 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 15:10:28 +0000 Subject: [PATCH 50/94] Go SDK: describe task handler params with JSON Schema The Dag processor will check each stub task against its Go handler at parse time, so the binding plan has to say how its parameters bind and what values they accept. Plan.Declare reports flat params as positional and nameless, a lone struct as named (or named_or_whole when untagged, since only then can it take one argument whole), with arg: tags matched exactly, and a JSON Schema per param in the vocabulary build_arg_bindings emits. The schema states only what decoding enforces: nullable slices, maps and pointers, integer ranges the decoder checks, and null where Go accepts more than one JSON shape. --- go-sdk/pkg/binding/binding.go | 3 +- go-sdk/pkg/binding/declaration.go | 211 +++++++++++++++++++++ go-sdk/pkg/binding/declaration_test.go | 247 +++++++++++++++++++++++++ 3 files changed, 460 insertions(+), 1 deletion(-) create mode 100644 go-sdk/pkg/binding/declaration.go create mode 100644 go-sdk/pkg/binding/declaration_test.go diff --git a/go-sdk/pkg/binding/binding.go b/go-sdk/pkg/binding/binding.go index 066f479f1bbee..40be16698c83f 100644 --- a/go-sdk/pkg/binding/binding.go +++ b/go-sdk/pkg/binding/binding.go @@ -23,7 +23,8 @@ // Captured defaults may go unclaimed, and a sole untagged struct can decode one // unclaimed argument as a whole value. // -// Analyze validates a function once. Resolve binds each execution. +// Analyze validates a function once. Resolve binds each execution. Declare describes the +// parameters for the Dag processor's check of a stub task against its handler. // AnalyzePositional builds a plan in which a sole struct binds positionally too, as one whole // argument. package binding diff --git a/go-sdk/pkg/binding/declaration.go b/go-sdk/pkg/binding/declaration.go new file mode 100644 index 0000000000000..04b86342d5de5 --- /dev/null +++ b/go-sdk/pkg/binding/declaration.go @@ -0,0 +1,211 @@ +// 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. + +package binding + +import ( + "encoding/json" + "math" + "reflect" + "time" + + "github.com/apache/airflow/go-sdk/pkg/execution/genmodels" +) + +// Declare describes the parameters that the task function's TaskFlow arguments fill, so the Dag +// processor can check a stub task against its Go handler before the task runs. +// +// Flat parameters bind by position and carry no name. A lone struct binds by field name, +// exactly for an `arg:`-tagged field and ignoring case and underscores otherwise. An untagged +// lone struct can also take one unmatched argument as a whole value. +func (p *Plan) Declare(taskID string) genmodels.TaskHandlerDeclaration { + decl := genmodels.TaskHandlerDeclaration{ + TaskID: taskID, + Binding: genmodels.TaskHandlerDeclarationBindingPositional, + Params: []genmodels.TaskHandlerParam{}, + } + for _, plan := range p.params { + switch plan.kind { + case paramData: + decl.Params = append(decl.Params, genmodels.TaskHandlerParam{ + Required: true, + ValueSchema: valueSchema(plan.typ), + }) + case paramLoneStruct: + decl.Binding = genmodels.TaskHandlerDeclarationBindingNamedOrWhole + if plan.tagged { + decl.Binding = genmodels.TaskHandlerDeclarationBindingNamed + } + // An unfilled field keeps its zero value, so none is required. + for _, sf := range plan.fields { + decl.Params = append(decl.Params, genmodels.TaskHandlerParam{ + Name: sf.argName, + ExactName: sf.tagged, + ValueSchema: valueSchema(sf.fieldType), + }) + } + } + } + return decl +} + +// valueSchema returns the JSON Schema of the values a parameter of type t accepts, or nil when +// no schema can state them. +// +// It states only what decoding into t enforces, so a value the schema rejects is one the task +// would fail on. The vocabulary is the one build_arg_bindings emits for a stub parameter's +// Python annotation, so the common pairs, such as int and int, come out equal. +func valueSchema(t reflect.Type) *genmodels.ArgValueSchema { + fragment := schemaFragment(t, map[reflect.Type]bool{}) + if fragment == nil { + return nil + } + schema := make(genmodels.ArgValueSchema, len(fragment)) + for k, v := range fragment { + schema[k] = v + } + return &schema +} + +var ( + timeType = reflect.TypeFor[time.Time]() + jsonNumberType = reflect.TypeFor[json.Number]() +) + +func schemaFragment(t reflect.Type, visiting map[reflect.Type]bool) map[string]any { + // A type that contains itself is left unconstrained where it recurs. + if visiting[t] { + return nil + } + visiting[t] = true + defer delete(visiting, t) + + switch t.Kind() { + case reflect.Pointer: + return nullable(schemaFragment(t.Elem(), visiting)) + case reflect.Slice, reflect.Map: + // Decoding null into a slice or map leaves it nil. + return nullable(nonNullFragment(t, visiting)) + default: + return nonNullFragment(t, visiting) + } +} + +func nonNullFragment(t reflect.Type, visiting map[reflect.Type]bool) map[string]any { + switch { + case t == timeType: + return map[string]any{"type": "string", "format": "date-time"} + // json.Number takes a number or a numeric string. + case t == jsonNumberType: + return nil + case implementsJSONUnmarshaler(t): + return nil + case implementsTextUnmarshaler(t): + return map[string]any{"type": "string"} + } + + switch t.Kind() { + case reflect.String: + return map[string]any{"type": "string"} + case reflect.Bool: + return map[string]any{"type": "boolean"} + case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64: + return signedIntFragment(t.Bits()) + case reflect.Uint, + reflect.Uint8, + reflect.Uint16, + reflect.Uint32, + reflect.Uint64, + reflect.Uintptr: + return unsignedIntFragment(t.Bits()) + case reflect.Float32: + return map[string]any{"type": "number", "format": "float"} + case reflect.Float64: + return map[string]any{"type": "number", "format": "double"} + case reflect.Slice: + // A byte slice takes a base64 string as well as an array. + if t.Elem().Kind() == reflect.Uint8 { + return nil + } + return arrayFragment(t.Elem(), visiting) + case reflect.Array: + return arrayFragment(t.Elem(), visiting) + case reflect.Map: + var values any = true + if fragment := schemaFragment(t.Elem(), visiting); fragment != nil { + values = fragment + } + return map[string]any{"type": "object", "additionalProperties": values} + case reflect.Struct: + return map[string]any{"type": "object"} + default: + return nil + } +} + +func signedIntFragment(bits int) map[string]any { + switch bits { + case 64: + return map[string]any{"type": "integer", "format": "int64"} + case 32: + return map[string]any{"type": "integer", "format": "int32"} + default: + return map[string]any{ + "type": "integer", + "minimum": -(int64(1) << (bits - 1)), + "maximum": int64(1)<<(bits-1) - 1, + } + } +} + +func unsignedIntFragment(bits int) map[string]any { + return map[string]any{ + "type": "integer", + "minimum": uint64(0), + "maximum": uint64(math.MaxUint64) >> (64 - bits), + } +} + +func arrayFragment(elem reflect.Type, visiting map[reflect.Type]bool) map[string]any { + items := schemaFragment(elem, visiting) + if items == nil { + items = map[string]any{} + } + return map[string]any{"type": "array", "items": items} +} + +func nullable(fragment map[string]any) map[string]any { + if fragment == nil { + return nil + } + if branches, isUnion := fragment["anyOf"].([]any); isUnion { + for _, branch := range branches { + if b, ok := branch.(map[string]any); ok && b["type"] == "null" { + return fragment + } + } + } + return map[string]any{"anyOf": []any{fragment, map[string]any{"type": "null"}}} +} + +func implementsJSONUnmarshaler(t reflect.Type) bool { + return t.Implements(jsonUnmarshalerType) || reflect.PointerTo(t).Implements(jsonUnmarshalerType) +} + +func implementsTextUnmarshaler(t reflect.Type) bool { + return t.Implements(textUnmarshalerType) || reflect.PointerTo(t).Implements(textUnmarshalerType) +} diff --git a/go-sdk/pkg/binding/declaration_test.go b/go-sdk/pkg/binding/declaration_test.go new file mode 100644 index 0000000000000..94f8e70f2608a --- /dev/null +++ b/go-sdk/pkg/binding/declaration_test.go @@ -0,0 +1,247 @@ +// 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. + +package binding + +import ( + "encoding/json" + "math" + "reflect" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/apache/airflow/go-sdk/internal/contexttest" + "github.com/apache/airflow/go-sdk/pkg/execution/genmodels" +) + +type declRegion string + +type declLevel int + +func (l *declLevel) UnmarshalText([]byte) error { return nil } + +type declCustom struct{ V int } + +func (c *declCustom) UnmarshalJSON([]byte) error { return nil } + +type declTree []declTree + +type declEmbedded struct{ Count int } + +type declTaggedInput struct { + Region string `arg:"region_code"` + Threshold float64 + declEmbedded +} + +type declUntaggedInput struct { + RegionCode string + Labels map[string]string +} + +var ( + stringFragment = map[string]any{"type": "string"} + int64Fragment = map[string]any{"type": "integer", "format": "int64"} + doubleFragment = map[string]any{"type": "number", "format": "double"} +) + +func nullOr(fragment map[string]any) map[string]any { + return map[string]any{"anyOf": []any{fragment, map[string]any{"type": "null"}}} +} + +func schemaOf(fragment map[string]any) *genmodels.ArgValueSchema { + schema := make(genmodels.ArgValueSchema, len(fragment)) + for k, v := range fragment { + schema[k] = v + } + return &schema +} + +func TestValueSchema(t *testing.T) { + for name, tc := range map[string]struct { + typ reflect.Type + want map[string]any + }{ + "string": {reflect.TypeFor[string](), stringFragment}, + "named string": {reflect.TypeFor[declRegion](), stringFragment}, + "bool": {reflect.TypeFor[bool](), map[string]any{"type": "boolean"}}, + "int": {reflect.TypeFor[int](), int64Fragment}, + "int64": {reflect.TypeFor[int64](), int64Fragment}, + "int32": { + reflect.TypeFor[int32](), + map[string]any{"type": "integer", "format": "int32"}, + }, + "int8": { + reflect.TypeFor[int8](), + map[string]any{"type": "integer", "minimum": int64(-128), "maximum": int64(127)}, + }, + "int16": { + reflect.TypeFor[int16](), + map[string]any{"type": "integer", "minimum": int64(-32768), "maximum": int64(32767)}, + }, + "uint8": { + reflect.TypeFor[uint8](), + map[string]any{"type": "integer", "minimum": uint64(0), "maximum": uint64(255)}, + }, + "uint32": { + reflect.TypeFor[uint32](), + map[string]any{"type": "integer", "minimum": uint64(0), "maximum": uint64(math.MaxUint32)}, + }, + "uint64": { + reflect.TypeFor[uint64](), + map[string]any{"type": "integer", "minimum": uint64(0), "maximum": uint64(math.MaxUint64)}, + }, + "float32": { + reflect.TypeFor[float32](), + map[string]any{"type": "number", "format": "float"}, + }, + "float64": {reflect.TypeFor[float64](), doubleFragment}, + "time.Time": { + reflect.TypeFor[time.Time](), + map[string]any{"type": "string", "format": "date-time"}, + }, + "pointer to time.Time": { + reflect.TypeFor[*time.Time](), + nullOr(map[string]any{"type": "string", "format": "date-time"}), + }, + "text unmarshaler on the pointer": {reflect.TypeFor[declLevel](), stringFragment}, + "json unmarshaler": {reflect.TypeFor[declCustom](), nil}, + "json.Number": {reflect.TypeFor[json.Number](), nil}, + "byte slice": {reflect.TypeFor[[]byte](), nil}, + "any": {reflect.TypeFor[any](), nil}, + "pointer to any": {reflect.TypeFor[*any](), nil}, + "complex": {reflect.TypeFor[complex128](), nil}, + "slice": { + reflect.TypeFor[[]string](), + nullOr(map[string]any{"type": "array", "items": stringFragment}), + }, + "slice of any": { + reflect.TypeFor[[]any](), + nullOr(map[string]any{"type": "array", "items": map[string]any{}}), + }, + "array": { + reflect.TypeFor[[3]int](), + map[string]any{"type": "array", "items": int64Fragment}, + }, + "map": { + reflect.TypeFor[map[string]int](), + nullOr(map[string]any{"type": "object", "additionalProperties": int64Fragment}), + }, + "map of any": { + reflect.TypeFor[map[string]any](), + nullOr(map[string]any{"type": "object", "additionalProperties": true}), + }, + "struct": {reflect.TypeFor[declUntaggedInput](), map[string]any{"type": "object"}}, + "pointer": {reflect.TypeFor[*int](), nullOr(int64Fragment)}, + "pointer to pointer": {reflect.TypeFor[**int](), nullOr(int64Fragment)}, + "pointer to slice": { + reflect.TypeFor[*[]int](), + nullOr(map[string]any{"type": "array", "items": int64Fragment}), + }, + "recursive": { + reflect.TypeFor[declTree](), + nullOr(map[string]any{"type": "array", "items": map[string]any{}}), + }, + } { + t.Run(name, func(t *testing.T) { + got := valueSchema(tc.typ) + if tc.want == nil { + assert.Nil(t, got) + return + } + assert.Equal(t, schemaOf(tc.want), got) + }) + } +} + +func declare(t *testing.T, fn any) genmodels.TaskHandlerDeclaration { + t.Helper() + plan, err := Analyze(reflect.TypeOf(fn), "testFn") + require.NoError(t, err) + return plan.Declare("task") +} + +func TestDeclareFlatParams(t *testing.T) { + decl := declare(t, func( + actx contexttest.Context, region string, count int, note *string, config declUntaggedInput, + ) error { + return nil + }) + + assert.Equal(t, genmodels.TaskHandlerDeclaration{ + TaskID: "task", + Binding: genmodels.TaskHandlerDeclarationBindingPositional, + Params: []genmodels.TaskHandlerParam{ + {Required: true, ValueSchema: schemaOf(stringFragment)}, + {Required: true, ValueSchema: schemaOf(int64Fragment)}, + {Required: true, ValueSchema: schemaOf(nullOr(stringFragment))}, + {Required: true, ValueSchema: schemaOf(map[string]any{"type": "object"})}, + }, + }, decl) +} + +func TestDeclareTaggedStruct(t *testing.T) { + decl := declare(t, func(actx contexttest.Context, in declTaggedInput) error { return nil }) + + assert.Equal(t, genmodels.TaskHandlerDeclaration{ + TaskID: "task", + Binding: genmodels.TaskHandlerDeclarationBindingNamed, + Params: []genmodels.TaskHandlerParam{ + {Name: "region_code", ExactName: true, ValueSchema: schemaOf(stringFragment)}, + {Name: "Threshold", ValueSchema: schemaOf(doubleFragment)}, + {Name: "Count", ValueSchema: schemaOf(int64Fragment)}, + }, + }, decl) +} + +func TestDeclareUntaggedStruct(t *testing.T) { + want := genmodels.TaskHandlerDeclaration{ + TaskID: "task", + Binding: genmodels.TaskHandlerDeclarationBindingNamedOrWhole, + Params: []genmodels.TaskHandlerParam{ + {Name: "RegionCode", ValueSchema: schemaOf(stringFragment)}, + { + Name: "Labels", + ValueSchema: schemaOf(nullOr(map[string]any{ + "type": "object", "additionalProperties": stringFragment, + })), + }, + }, + } + for name, fn := range map[string]any{ + "value": func(actx contexttest.Context, in declUntaggedInput) error { return nil }, + "pointer": func(actx contexttest.Context, in *declUntaggedInput) error { return nil }, + } { + t.Run(name, func(t *testing.T) { + assert.Equal(t, want, declare(t, fn)) + }) + } +} + +func TestDeclareContextOnly(t *testing.T) { + decl := declare(t, func(actx contexttest.Context) error { return nil }) + + // A nil slice would go out as null, which the Dag processor rejects. + assert.Equal(t, genmodels.TaskHandlerDeclaration{ + TaskID: "task", + Binding: genmodels.TaskHandlerDeclarationBindingPositional, + Params: []genmodels.TaskHandlerParam{}, + }, decl) +} From 4619601b5ded44f5278fe7a55dbbd616ce358b96 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 15:14:59 +0000 Subject: [PATCH 51/94] Go SDK: answer TaskHandlerParseRequest with registered handlers The Dag processor asks a Lang-SDK runtime which task handlers an artifact registers for the Dags a Python file owns, and today a Go bundle exits on that first frame as an unknown message. Serve now replies with one TaskHandlerParsingResult: each requested Dag's handler declarations in registration order, with a Dag that has none left out. It waits for the parent's acknowledgement before returning, so the result is recorded before the process exits and closes the socket. Empty params and task_handlers go out as empty values, not null, which the parent would reject. --- go-sdk/README.md | 3 + go-sdk/airflow/bundle.go | 9 +- go-sdk/airflow/serve.go | 5 +- go-sdk/internal/bundle/task.go | 15 ++ go-sdk/pkg/execution/integration_test.go | 129 +++++++++++++++++- go-sdk/pkg/execution/messages.go | 6 + go-sdk/pkg/execution/messages_test.go | 25 ++++ go-sdk/pkg/execution/server.go | 31 +++-- go-sdk/pkg/execution/task_handler_parse.go | 50 +++++++ .../pkg/execution/task_handler_parse_test.go | 76 +++++++++++ 10 files changed, 332 insertions(+), 17 deletions(-) create mode 100644 go-sdk/pkg/execution/task_handler_parse.go create mode 100644 go-sdk/pkg/execution/task_handler_parse_test.go diff --git a/go-sdk/README.md b/go-sdk/README.md index 400726c1acb47..9fdacd4a78617 100644 --- a/go-sdk/README.md +++ b/go-sdk/README.md @@ -389,6 +389,9 @@ Python supervisor / task runner protocol on the comm socket, with structured JSON-line logs on the logs socket. - The Python runtime is the worker. It proxies every `GetConnection` / `GetVariable` / `GetXCom` / `SetXCom` call through to the Execution API. The Go binary just runs the task function. +- The Dag processor launches the binary the same way to check a Python Dag's stub tasks. It sends a + `TaskHandlerParseRequest` that names Dag ids. For each of those Dags, the binary replies with the task + handlers it registers, in registration order, and how each handler's parameters bind. Then it exits. The Go side of the protocol is implemented in `pkg/execution/`. On the Python side it is the `ExecutableCoordinator` in `task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py`. diff --git a/go-sdk/airflow/bundle.go b/go-sdk/airflow/bundle.go index 8f82b07121138..58f04354ea358 100644 --- a/go-sdk/airflow/bundle.go +++ b/go-sdk/airflow/bundle.go @@ -125,18 +125,15 @@ func (b *BundleRef) Register(items ...Registerable) { } // taskHandlerMap holds the registered task handlers by dag_id and task_id. -// It also keeps registration order. The --airflow-metadata manifest lists the tasks of each Dag -// in that order. +// It also keeps registration order. The --airflow-metadata manifest, and the reply to the Dag +// processor's task handler parse, list the tasks of each Dag in that order. type taskHandlerMap struct { mu sync.RWMutex handlers map[string]map[string]bundle.Task order []bundle.TaskHandlerInfo } -var ( - _ bundle.Bundle = (*taskHandlerMap)(nil) - _ bundle.EnumerableBundle = (*taskHandlerMap)(nil) -) +var _ bundle.Registry = (*taskHandlerMap)(nil) func (m *taskHandlerMap) add(dagID, taskID string, task bundle.Task) { m.mu.Lock() diff --git a/go-sdk/airflow/serve.go b/go-sdk/airflow/serve.go index bcfb33ab934d1..910cd5a1398f2 100644 --- a/go-sdk/airflow/serve.go +++ b/go-sdk/airflow/serve.go @@ -57,8 +57,9 @@ const ( // The command-line flags of the executable decide what Serve does. // With --airflow-metadata it prints the bundle's manifest and returns, which is how // airflow-go-pack reads the Dag and task ids of the registered task handlers. -// With --comm and --logs, which the Airflow supervisor passes, it runs one task over the -// coordinator protocol. +// With --comm and --logs, which Airflow passes, it speaks the coordinator protocol: it either +// runs one task, or tells the Dag processor which task handlers it registers for the Dags the +// Dag processor asks about. // // main must exit with a non-zero status when Serve returns an error, because the exit status // is how the supervisor learns that the task failed: diff --git a/go-sdk/internal/bundle/task.go b/go-sdk/internal/bundle/task.go index d2353369be8e3..245f2d79308ee 100644 --- a/go-sdk/internal/bundle/task.go +++ b/go-sdk/internal/bundle/task.go @@ -26,6 +26,7 @@ import ( "runtime" "github.com/apache/airflow/go-sdk/pkg/binding" + "github.com/apache/airflow/go-sdk/pkg/execution/genmodels" "github.com/apache/airflow/go-sdk/pkg/sdkcontext" "github.com/apache/airflow/go-sdk/sdk" ) @@ -35,6 +36,9 @@ import ( // airflow.DagRef.Task wrap a plain Go function into a Task. type Task interface { Execute(ctx context.Context, logger *slog.Logger, args []binding.Arg) error + // Declare describes the parameters the task binds, for the Dag processor to check the + // Python stub task against. + Declare(taskID string) genmodels.TaskHandlerDeclaration } // Bundle looks up a registered task by dag_id and task_id. The coordinator @@ -57,6 +61,13 @@ type EnumerableBundle interface { ListTaskHandlers() []TaskHandlerInfo } +// Registry is what a bundle binary serves: a task run looks its handler up, and a task +// handler parse lists the handlers. +type Registry interface { + Bundle + EnumerableBundle +} + type taskFunction struct { fn reflect.Value fullName string @@ -107,6 +118,10 @@ func (f *taskFunction) Execute( return f.call(ctx, sdkClient, reflectArgs, logger) } +func (f *taskFunction) Declare(taskID string) genmodels.TaskHandlerDeclaration { + return f.plan.Declare(taskID) +} + func clientFrom(ctx context.Context) (sdk.Client, error) { client, ok := ctx.Value(sdkcontext.SdkClientContextKey).(sdk.Client) if !ok { diff --git a/go-sdk/pkg/execution/integration_test.go b/go-sdk/pkg/execution/integration_test.go index c53bed89695b2..5518dbabe3e48 100644 --- a/go-sdk/pkg/execution/integration_test.go +++ b/go-sdk/pkg/execution/integration_test.go @@ -101,7 +101,29 @@ func (b testBundle) LookupTask(dagID, taskID string) (bundle.Task, bool) { return task, ok } -func buildBundle(t *testing.T, register func(testBundle)) bundle.Bundle { +// ListTaskHandlers lists nothing: a map keeps no registration order, and only a task +// handler parse lists handlers. The parse tests use handlerRegistry instead. +func (b testBundle) ListTaskHandlers() []bundle.TaskHandlerInfo { return nil } + +// handlerRegistry lists its task handlers in the order they were added, as airflow.Bundle does. +type handlerRegistry struct { + testBundle + order []bundle.TaskHandlerInfo +} + +func newHandlerRegistry() *handlerRegistry { return &handlerRegistry{testBundle: testBundle{}} } + +func (r *handlerRegistry) add(dagID, taskID string, fn any) { + if _, ok := r.testBundle[dagID]; !ok { + r.AddDag(dagID) + } + r.testBundle[dagID].AddTaskWithName(taskID, fn) + r.order = append(r.order, bundle.TaskHandlerInfo{DagID: dagID, TaskID: taskID}) +} + +func (r *handlerRegistry) ListTaskHandlers() []bundle.TaskHandlerInfo { return r.order } + +func buildBundle(t *testing.T, register func(testBundle)) bundle.Registry { t.Helper() b := testBundle{} register(b) @@ -893,3 +915,108 @@ func TestServeFailureAfterConnectClosesComm(t *testing.T) { _, err = readFrame(commConn) require.Error(t, err) } + +type loadInput struct { + Table string `arg:"table_name"` +} + +// startTaskHandlerParse runs Serve for a registry with two handlers, sends it a +// TaskHandlerParseRequest, and returns the comm connection and the reply frame. +func startTaskHandlerParse(t *testing.T) (net.Conn, IncomingFrame, chan error) { + t.Helper() + commAddr, logsAddr, commCh, logsCh, cleanup := startSupervisor(t) + t.Cleanup(cleanup) + + registry := newHandlerRegistry() + registry.add("etl", "extract", func(actx contexttest.Context, region string) error { + return nil + }) + registry.add("etl", "load", func(actx contexttest.Context, in loadInput) error { return nil }) + registry.add("etl", "notify", simpleTask) + + done := make(chan error, 1) + go func() { done <- Serve(registry, commAddr, logsAddr) }() + + commConn := <-commCh + t.Cleanup(func() { commConn.Close() }) + logsConn := <-logsCh + t.Cleanup(func() { logsConn.Close() }) + require.NoError(t, commConn.SetDeadline(time.Now().Add(10*time.Second))) + + payload, err := encodeRequest(0, map[string]any{ + "type": "TaskHandlerParseRequest", + "file": "/bundles/go/etl", + "dag_ids": []string{"etl", "missing"}, + "bundle_path": "/bundles/go", + "bundle_name": "go", + }) + require.NoError(t, err) + require.NoError(t, writeFrame(commConn, payload)) + + frame, err := readFrame(commConn) + require.NoError(t, err) + require.True(t, isNilRaw(frame.Err)) + return commConn, frame, done +} + +func TestServeTaskHandlerParseRequestEndToEnd(t *testing.T) { + commConn, frame, done := startTaskHandlerParse(t) + + // Decoded as a map, so the wire's keys, nulls and empty lists are what is compared. + assert.Equal(t, map[string]any{ + "type": "TaskHandlerParsingResult", + "fileloc": "/bundles/go/etl", + "task_handlers": map[string]any{ + "etl": []any{ + map[string]any{ + "task_id": "extract", + "binding": "positional", + "params": []any{map[string]any{ + "name": nil, + "required": true, + "value_schema": map[string]any{"type": "string"}, + }}, + }, + map[string]any{ + "task_id": "load", + "binding": "named", + "params": []any{map[string]any{ + "name": "table_name", + "exact_name": true, + "required": false, + "value_schema": map[string]any{"type": "string"}, + }}, + }, + map[string]any{"task_id": "notify", "binding": "positional", "params": []any{}}, + }, + }, + }, rawToMap(t, frame.Body)) + + // Serve returns only once the parent acknowledges the result. + select { + case err := <-done: + t.Fatalf("Serve returned before the result was acknowledged: %v", err) + case <-time.After(100 * time.Millisecond): + } + require.NoError(t, writeFrame(commConn, encodeResponseFrame(t, frame.ID, nil, nil))) + + select { + case err := <-done: + require.NoError(t, err) + case <-time.After(2 * time.Second): + t.Fatal("Serve did not return after the result was acknowledged") + } +} + +func TestServeTaskHandlerParseFailsWithoutAcknowledgement(t *testing.T) { + commConn, _, done := startTaskHandlerParse(t) + + require.NoError(t, commConn.Close()) + + select { + case err := <-done: + require.ErrorContains(t, err, "sending task handler parse result") + case <-time.After(2 * time.Second): + t.Fatal("Serve did not return after the parent closed the socket") + } +} diff --git a/go-sdk/pkg/execution/messages.go b/go-sdk/pkg/execution/messages.go index 1dd596201258a..77df9add847a6 100644 --- a/go-sdk/pkg/execution/messages.go +++ b/go-sdk/pkg/execution/messages.go @@ -75,6 +75,12 @@ func decodeIncomingBody(raw msgpack.RawMessage) (any, error) { return nil, fmt.Errorf("decoding StartupDetails: %w", err) } return &msg, nil + case genmodels.TypeTaskHandlerParseRequest: + var msg genmodels.TaskHandlerParseRequest + if err := msgpack.Unmarshal(raw, &msg); err != nil { + return nil, fmt.Errorf("decoding TaskHandlerParseRequest: %w", err) + } + return &msg, nil case genmodels.TypeDagFileParseRequest: var msg genmodels.DagFileParseRequest if err := msgpack.Unmarshal(raw, &msg); err != nil { diff --git a/go-sdk/pkg/execution/messages_test.go b/go-sdk/pkg/execution/messages_test.go index fa81df4ea66d2..f86baaa0f5fa6 100644 --- a/go-sdk/pkg/execution/messages_test.go +++ b/go-sdk/pkg/execution/messages_test.go @@ -462,6 +462,31 @@ func TestDecodeIncomingBodyDispatch(t *testing.T) { assert.True(t, ok) }) + t.Run("TaskHandlerParseRequest", func(t *testing.T) { + raw := marshalBody(t, map[string]any{ + "type": "TaskHandlerParseRequest", + "file": "/bundles/go/etl", + "dag_ids": []string{"etl"}, + "bundle_path": "/bundles/go", + "bundle_name": "go", + }) + result, err := decodeIncomingBody(raw) + require.NoError(t, err) + assert.Equal(t, &genmodels.TaskHandlerParseRequest{ + Type: "TaskHandlerParseRequest", + File: "/bundles/go/etl", + DagIds: []string{"etl"}, + BundlePath: "/bundles/go", + BundleName: "go", + }, result) + }) + + t.Run("malformed TaskHandlerParseRequest", func(t *testing.T) { + raw := marshalBody(t, map[string]any{"type": "TaskHandlerParseRequest", "dag_ids": "etl"}) + _, err := decodeIncomingBody(raw) + assert.ErrorContains(t, err, "decoding TaskHandlerParseRequest") + }) + t.Run("ErrorResponse", func(t *testing.T) { raw := marshalBody(t, map[string]any{"type": "ErrorResponse", "error": "GENERIC_ERROR"}) result, err := decodeIncomingBody(raw) diff --git a/go-sdk/pkg/execution/server.go b/go-sdk/pkg/execution/server.go index 0f56e005315ca..7973c63aa3949 100644 --- a/go-sdk/pkg/execution/server.go +++ b/go-sdk/pkg/execution/server.go @@ -20,8 +20,10 @@ // the Airflow supervisor (Python ExecutableCoordinator), the Serve method of // airflow.BundleRef dispatches here. // -// The first inbound frame on the comm socket is a StartupDetails message -// that drives multi-round task execution. +// The first inbound frame on the comm socket selects what the runtime does. A +// StartupDetails message drives multi-round task execution. A +// TaskHandlerParseRequest from the Dag processor gets one TaskHandlerParsingResult +// that declares the task handlers registered for the requested Dags. // // See go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md. package execution @@ -58,11 +60,12 @@ const terminalSendTimeout = 30 * time.Second // comm and logs sockets, installs an slog handler that writes JSON-line // records to the logs connection, and dispatches on the first frame. // -// Serve returns nil on a clean shutdown: the task ran and its terminal -// TaskState/SucceedTask frame was delivered, and the caller should exit 0. A -// non-nil error indicates a protocol-level failure (connection loss, -// malformed frames, unknown first message type) that happens before or -// instead of delivering a terminal frame. +// Serve returns nil on a clean shutdown, and the caller should exit 0: either +// the task ran and its terminal TaskState/SucceedTask frame was delivered, or +// the task handler parse result was delivered and acknowledged. A non-nil error +// indicates a protocol-level failure (connection loss, malformed frames, +// unknown first message type) that happens before or instead of delivering a +// terminal frame. // // Failure-signaling contract: the caller (main) must turn a non-nil error // into a non-zero process exit. The supervisor derives the task's final state @@ -73,7 +76,7 @@ const terminalSendTimeout = 30 * time.Second // fails closed without needing to send a frame; the post-connect paths below // log the reason at Error first so it still reaches the supervisor's log // stream over the already-connected logs socket. -func Serve(b bundle.Bundle, commAddr, logsAddr string) error { +func Serve(b bundle.Registry, commAddr, logsAddr string) error { if commAddr == "" { return fmt.Errorf("missing --comm=host:port argument") } @@ -167,6 +170,18 @@ func Serve(b bundle.Bundle, commAddr, logsAddr string) error { } logger.Debug("Task execution complete") + case *genmodels.TaskHandlerParseRequest: + logger.Debug("Task handler parse mode", "dag_ids", msg.DagIds) + result := declareTaskHandlers(b, msg) + sendCtx, cancel := context.WithTimeout(ctx, terminalSendTimeout) + defer cancel() + // Waiting for the acknowledgement means the parent holds the result before this + // process exits and closes the socket under it. + if _, err := comm.Communicate(sendCtx, result); err != nil { + return fmt.Errorf("sending task handler parse result: %w", err) + } + logger.Debug("Task handler parse complete") + default: logger.Error("Unexpected initial message type", "type", fmt.Sprintf("%T", body)) return fmt.Errorf("unexpected initial message type: %T", body) diff --git a/go-sdk/pkg/execution/task_handler_parse.go b/go-sdk/pkg/execution/task_handler_parse.go new file mode 100644 index 0000000000000..17d0359b2e2d7 --- /dev/null +++ b/go-sdk/pkg/execution/task_handler_parse.go @@ -0,0 +1,50 @@ +// 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. + +package execution + +import ( + "github.com/apache/airflow/go-sdk/internal/bundle" + "github.com/apache/airflow/go-sdk/pkg/execution/genmodels" +) + +// declareTaskHandlers answers a TaskHandlerParseRequest with the declarations of the task +// handlers registered for the requested Dags, each Dag's in registration order. A requested +// Dag with no handler is left out rather than mapped to an empty list. +func declareTaskHandlers( + b bundle.Registry, + req *genmodels.TaskHandlerParseRequest, +) genmodels.TaskHandlerParsingResult { + requested := make(map[string]bool, len(req.DagIds)) + for _, dagID := range req.DagIds { + requested[dagID] = true + } + // A nil map would go out as null, which the Dag processor rejects. + handlers := genmodels.TaskHandlers{} + for _, info := range b.ListTaskHandlers() { + if !requested[info.DagID] { + continue + } + // Declare only what a task run could find. + task, ok := b.LookupTask(info.DagID, info.TaskID) + if !ok { + continue + } + handlers[info.DagID] = append(handlers[info.DagID], task.Declare(info.TaskID)) + } + return genmodels.TaskHandlerParsingResult{Fileloc: req.File, TaskHandlers: handlers} +} diff --git a/go-sdk/pkg/execution/task_handler_parse_test.go b/go-sdk/pkg/execution/task_handler_parse_test.go new file mode 100644 index 0000000000000..f077e27692c41 --- /dev/null +++ b/go-sdk/pkg/execution/task_handler_parse_test.go @@ -0,0 +1,76 @@ +// 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. + +package execution + +import ( + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/apache/airflow/go-sdk/internal/bundle" + "github.com/apache/airflow/go-sdk/internal/contexttest" + "github.com/apache/airflow/go-sdk/pkg/execution/genmodels" +) + +func TestDeclareTaskHandlers(t *testing.T) { + registry := newHandlerRegistry() + registry.add("etl", "transform", func(actx contexttest.Context, region string) error { + return nil + }) + registry.add("report", "render", simpleTask) + registry.add("etl", "extract", simpleTask) + // Listed, but a task run could not find it. + registry.order = append(registry.order, bundle.TaskHandlerInfo{DagID: "etl", TaskID: "ghost"}) + + result := declareTaskHandlers(registry, &genmodels.TaskHandlerParseRequest{ + File: "/bundles/go/etl", + DagIds: []string{"etl", "missing"}, + }) + + assert.Equal(t, genmodels.TaskHandlerParsingResult{ + Fileloc: "/bundles/go/etl", + TaskHandlers: genmodels.TaskHandlers{ + "etl": { + { + TaskID: "transform", + Binding: genmodels.TaskHandlerDeclarationBindingPositional, + Params: []genmodels.TaskHandlerParam{ + {Required: true, ValueSchema: &genmodels.ArgValueSchema{"type": "string"}}, + }, + }, + { + TaskID: "extract", + Binding: genmodels.TaskHandlerDeclarationBindingPositional, + Params: []genmodels.TaskHandlerParam{}, + }, + }, + }, + }, result) +} + +func TestDeclareTaskHandlersMatchingNothing(t *testing.T) { + registry := newHandlerRegistry() + registry.add("report", "render", simpleTask) + + result := declareTaskHandlers(registry, &genmodels.TaskHandlerParseRequest{ + DagIds: []string{"etl"}, + }) + + // A nil map would go out as null, which the Dag processor rejects. + assert.Equal(t, genmodels.TaskHandlers{}, result.TaskHandlers) +} From fc06e0786b70fa5f72723924ab57f6669e77aa35 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 15:41:32 +0000 Subject: [PATCH 52/94] Go SDK: clarify parse timeout and nested-null docs, test anyOf Review follow-up. terminalSendTimeout now also bounds the parse result's send and acknowledgement, and its comment only named the task path. valueSchema claimed to state everything decoding enforces, but a null element or map value decodes to the zero value and is not stated, as in the runtime's own argument check. The Serve wire test had no anyOf schema, so a nullable param now proves a union survives the msgpack round trip. --- go-sdk/pkg/binding/declaration.go | 6 ++++-- go-sdk/pkg/execution/integration_test.go | 22 ++++++++++++++++------ go-sdk/pkg/execution/server.go | 9 +++++---- 3 files changed, 25 insertions(+), 12 deletions(-) diff --git a/go-sdk/pkg/binding/declaration.go b/go-sdk/pkg/binding/declaration.go index 04b86342d5de5..9d75d41af077d 100644 --- a/go-sdk/pkg/binding/declaration.go +++ b/go-sdk/pkg/binding/declaration.go @@ -67,8 +67,10 @@ func (p *Plan) Declare(taskID string) genmodels.TaskHandlerDeclaration { // no schema can state them. // // It states only what decoding into t enforces, so a value the schema rejects is one the task -// would fail on. The vocabulary is the one build_arg_bindings emits for a stub parameter's -// Python annotation, so the common pairs, such as int and int, come out equal. +// would fail on. The exception is a null element or map value, which decoding turns into the +// zero value; like the runtime's own argument check, the schema states null only at the top level. +// The vocabulary is the one build_arg_bindings emits for a stub parameter's Python annotation, +// so the common pairs, such as int and int, come out equal. func valueSchema(t reflect.Type) *genmodels.ArgValueSchema { fragment := schemaFragment(t, map[reflect.Type]bool{}) if fragment == nil { diff --git a/go-sdk/pkg/execution/integration_test.go b/go-sdk/pkg/execution/integration_test.go index 5518dbabe3e48..94fd0a0d9c274 100644 --- a/go-sdk/pkg/execution/integration_test.go +++ b/go-sdk/pkg/execution/integration_test.go @@ -928,7 +928,7 @@ func startTaskHandlerParse(t *testing.T) (net.Conn, IncomingFrame, chan error) { t.Cleanup(cleanup) registry := newHandlerRegistry() - registry.add("etl", "extract", func(actx contexttest.Context, region string) error { + registry.add("etl", "extract", func(actx contexttest.Context, region string, day *int) error { return nil }) registry.add("etl", "load", func(actx contexttest.Context, in loadInput) error { return nil }) @@ -971,11 +971,21 @@ func TestServeTaskHandlerParseRequestEndToEnd(t *testing.T) { map[string]any{ "task_id": "extract", "binding": "positional", - "params": []any{map[string]any{ - "name": nil, - "required": true, - "value_schema": map[string]any{"type": "string"}, - }}, + "params": []any{ + map[string]any{ + "name": nil, + "required": true, + "value_schema": map[string]any{"type": "string"}, + }, + map[string]any{ + "name": nil, + "required": true, + "value_schema": map[string]any{"anyOf": []any{ + map[string]any{"type": "integer", "format": "int64"}, + map[string]any{"type": "null"}, + }}, + }, + }, }, map[string]any{ "task_id": "load", diff --git a/go-sdk/pkg/execution/server.go b/go-sdk/pkg/execution/server.go index 7973c63aa3949..789f998704d75 100644 --- a/go-sdk/pkg/execution/server.go +++ b/go-sdk/pkg/execution/server.go @@ -50,10 +50,11 @@ import ( const dialTimeout = 30 * time.Second // terminalSendTimeout bounds the write of the final TaskState/SucceedTask -// frame. The supervisor normally drains the comm socket promptly, but a -// half-open connection (the supervisor gone without a clean close) could -// otherwise wedge the runtime on a blocked write; the deadline turns that -// into a fast failure -- and thus a non-zero exit -- instead of a hang. +// frame, and the send and acknowledgement of a task handler parse result. The +// supervisor normally drains the comm socket promptly, but a half-open +// connection (the supervisor gone without a clean close) could otherwise wedge +// the runtime on a blocked write or wait; the deadline turns that into a fast +// failure -- and thus a non-zero exit -- instead of a hang. const terminalSendTimeout = 30 * time.Second // Serve runs the bundle binary in coordinator mode. It dials the supervisor's From f3d803492c5c94b13cd3367b83c336049126422d Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 15:54:33 +0000 Subject: [PATCH 53/94] Go SDK: record integrity and cache digests in bundle metadata The Dag processor will skip re-probing an artifact that has not changed since its last probe, and has to decide that from a stored value, because hashing every candidate on every parse is what the fast path exists to avoid. The trailer's binary_sha256 cannot serve: it misses a change to the source or the metadata. The cache digest covers all three regions without covering itself, so it is stable across packs of the same inputs and changes with any of them. integrity is the binary hash the packer already needs for that, recorded as ADR-0014 names it. --- go-sdk/cmd/airflow-go-pack/pack.go | 79 +++++++++-- .../airflow-go-pack/pack_integration_test.go | 11 ++ go-sdk/cmd/airflow-go-pack/pack_test.go | 127 +++++++++++++++++- .../airflowmetadata/airflowmetadata.go | 4 +- go-sdk/internal/bundlefooter/footer.go | 8 +- task-sdk/docs/airflow-metadata.schema.json | 17 +++ task-sdk/docs/executable-bundle-spec.rst | 17 +++ 7 files changed, 246 insertions(+), 17 deletions(-) diff --git a/go-sdk/cmd/airflow-go-pack/pack.go b/go-sdk/cmd/airflow-go-pack/pack.go index f4be6fa8c4578..39c8764d0bb0a 100644 --- a/go-sdk/cmd/airflow-go-pack/pack.go +++ b/go-sdk/cmd/airflow-go-pack/pack.go @@ -19,6 +19,8 @@ package main import ( "bytes" + "crypto/sha256" + "encoding/hex" "errors" "fmt" "go/ast" @@ -173,13 +175,21 @@ func runPack(stdout, stderr io.Writer, opts *packOptions) error { } warnOnSuspiciousIDs(stderr, meta) - manifest, err := renderManifest(meta, filepath.Base(sourcePath)) + sourceBytes, err := os.ReadFile(sourcePath) + if err != nil { + return fmt.Errorf("reading source file: %w", err) + } + baseManifest, err := renderManifest(meta, filepath.Base(sourcePath), nil) if err != nil { return fmt.Errorf("rendering manifest: %w", err) } - sourceBytes, err := os.ReadFile(sourcePath) + digests, err := computeDigests(execPath, sourceBytes, baseManifest) if err != nil { - return fmt.Errorf("reading source file: %w", err) + return err + } + manifest, err := renderManifest(meta, filepath.Base(sourcePath), &digests) + if err != nil { + return fmt.Errorf("rendering manifest: %w", err) } // Assemble the bundle through a temp file and atomically move it into @@ -462,12 +472,53 @@ func runIntrospect(execPath string, flag string) ([]byte, error) { return stdout.Bytes(), nil } +// cacheDigestDomain starts the cache digest's input, so the digest cannot equal +// a SHA-256 taken of the same bytes for another purpose, and a later definition +// can change the version. +const cacheDigestDomain = "airflow-go-pack cache digest v1\x00" + +// bundleDigests is the manifest's digests mapping, each a lower-case hex SHA-256. +type bundleDigests struct { + // Integrity covers the binary region; it equals the trailer's binary_sha256. + Integrity string + // Cache covers every region of the bundle except its own value. It is + // SHA-256(cacheDigestDomain || SHA-256(binary) || SHA-256(source) || + // SHA-256(manifest rendered without digests)), with each inner digest as + // its 32 raw bytes. Packing the same inputs again gives the same value, and + // a change to any of them gives a different one. + Cache string +} + +func computeDigests(execPath string, source, baseManifest []byte) (bundleDigests, error) { + binaryHash, err := bundlefooter.HashFile(execPath) + if err != nil { + return bundleDigests{}, fmt.Errorf("hashing executable %s: %w", execPath, err) + } + sourceHash := sha256.Sum256(source) + manifestHash := sha256.Sum256(baseManifest) + + h := sha256.New() + h.Write([]byte(cacheDigestDomain)) + h.Write(binaryHash[:]) + h.Write(sourceHash[:]) + h.Write(manifestHash[:]) + return bundleDigests{ + Integrity: hex.EncodeToString(binaryHash[:]), + Cache: hex.EncodeToString(h.Sum(nil)), + }, nil +} + // renderManifest serialises the airflow-metadata manifest as deterministic, // sorted-key YAML matching airflow-metadata.schema.json. It injects the schema's -// source field (the filename the manifest is built from), which the producer's -// Manifest omits because only the packer knows it; every other field is copied -// from the introspected manifest verbatim. -func renderManifest(meta airflowmetadata.Manifest, sourceName string) ([]byte, error) { +// source field (the filename the manifest is built from) and, when digests is +// not nil, the digests mapping, both of which the producer's Manifest omits +// because only the packer knows them; every other field is copied from the +// introspected manifest verbatim. +func renderManifest( + meta airflowmetadata.Manifest, + sourceName string, + digests *bundleDigests, +) ([]byte, error) { version := meta.AirflowBundleMetadataVersion if version == "" { version = airflowmetadata.FormatVersion @@ -514,9 +565,21 @@ func renderManifest(meta airflowmetadata.Manifest, sourceName string) ([]byte, e }, }, scalar("source"), quotedScalar(sourceName), - scalar("dags"), dagsNode, }, } + if digests != nil { + manifest.Content = append(manifest.Content, + scalar("digests"), + &yaml.Node{ + Kind: yaml.MappingNode, + Content: []*yaml.Node{ + scalar("integrity"), quotedScalar(digests.Integrity), + scalar("cache"), quotedScalar(digests.Cache), + }, + }, + ) + } + manifest.Content = append(manifest.Content, scalar("dags"), dagsNode) root.Content = []*yaml.Node{manifest} var buf bytes.Buffer diff --git a/go-sdk/cmd/airflow-go-pack/pack_integration_test.go b/go-sdk/cmd/airflow-go-pack/pack_integration_test.go index e1cc29fe68e67..1d75ab6c41936 100644 --- a/go-sdk/cmd/airflow-go-pack/pack_integration_test.go +++ b/go-sdk/cmd/airflow-go-pack/pack_integration_test.go @@ -19,6 +19,8 @@ package main import ( "bytes" + "crypto/sha256" + "encoding/hex" "encoding/json" "io" "os" @@ -138,6 +140,12 @@ func TestPack_CrossArchExecutableWithMetadataFile(t *testing.T) { sdkVersion := m[1] assert.Regexp(t, `^(\(devel\)|v[0-9].*)$`, sdkVersion, `sdk.version must be "(devel)" or a v-prefixed module version`) + // The cache digest covers the environment-dependent binary and sdk.version + // too, so it is folded in the same way. + cacheLine := regexp.MustCompile(`(?m)^ cache: "([0-9a-f]{64})"$`) + c := cacheLine.FindStringSubmatch(string(metadata)) + require.NotNil(t, c, "manifest must contain a digests.cache line:\n%s", metadata) + crossHash := sha256.Sum256(crossBytes) expectedManifest := `airflow_bundle_metadata_version: "1.0" sdk: @@ -145,6 +153,9 @@ sdk: version: "` + sdkVersion + `" supervisor_schema_version: "` + execution.SupervisorSchemaVersion + `" source: "main.go" +digests: + integrity: "` + hex.EncodeToString(crossHash[:]) + `" + cache: "` + c[1] + `" dags: concurrent_xcom_dag: tasks: diff --git a/go-sdk/cmd/airflow-go-pack/pack_test.go b/go-sdk/cmd/airflow-go-pack/pack_test.go index 999d5c0f01771..b001d6b6619fb 100644 --- a/go-sdk/cmd/airflow-go-pack/pack_test.go +++ b/go-sdk/cmd/airflow-go-pack/pack_test.go @@ -18,6 +18,9 @@ package main import ( + "bytes" + "crypto/sha256" + "encoding/hex" "errors" "io" "os" @@ -28,6 +31,7 @@ import ( "github.com/spf13/cobra" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "gopkg.in/yaml.v3" "github.com/apache/airflow/go-sdk/internal/airflowmetadata" "github.com/apache/airflow/go-sdk/internal/bundlefooter" @@ -47,9 +51,9 @@ func TestRenderManifest_DeterministicDagOrdering(t *testing.T) { }, } - got1, err := renderManifest(meta, "main.go") + got1, err := renderManifest(meta, "main.go", nil) require.NoError(t, err) - got2, err := renderManifest(meta, "main.go") + got2, err := renderManifest(meta, "main.go", nil) require.NoError(t, err) assert.Equal(t, got1, got2, "manifest should be byte-identical for identical input") @@ -87,7 +91,7 @@ func TestRenderManifest_QuotesValuesNotKeys(t *testing.T) { }, } - got, err := renderManifest(meta, "main.go") + got, err := renderManifest(meta, "main.go", nil) require.NoError(t, err) // Task values that look like scalars are quoted. @@ -98,6 +102,36 @@ func TestRenderManifest_QuotesValuesNotKeys(t *testing.T) { assert.NotContains(t, string(got), `"my_dag"`) } +func TestRenderManifest_Digests(t *testing.T) { + meta := airflowmetadata.Manifest{ + AirflowBundleMetadataVersion: "1.0", + SDK: airflowmetadata.SDK{ + Language: "go", + Version: "0.1.0", + SupervisorSchemaVersion: "2026-06-16", + }, + Dags: map[string]airflowmetadata.Dag{"my_dag": {Tasks: []string{"t1"}}}, + } + + got, err := renderManifest(meta, "main.go", &bundleDigests{Integrity: "aa11", Cache: "bb22"}) + require.NoError(t, err) + + assert.Equal(t, `airflow_bundle_metadata_version: "1.0" +sdk: + language: "go" + version: "0.1.0" + supervisor_schema_version: "2026-06-16" +source: "main.go" +digests: + integrity: "aa11" + cache: "bb22" +dags: + my_dag: + tasks: + - "t1" +`, string(got)) +} + func TestRenderManifest_EmptyDags(t *testing.T) { meta := airflowmetadata.Manifest{ AirflowBundleMetadataVersion: "1.0", @@ -108,7 +142,7 @@ func TestRenderManifest_EmptyDags(t *testing.T) { }, Dags: map[string]airflowmetadata.Dag{}, } - got, err := renderManifest(meta, "main.go") + got, err := renderManifest(meta, "main.go", nil) require.NoError(t, err) assert.Contains(t, string(got), "dags: {}") } @@ -417,6 +451,91 @@ func TestRunPack_AcceptsYAMLMetadataFile(t *testing.T) { ) } +// packDigests packs the given executable, source and JSON manifest in dir, and +// returns the digests the packed bundle's manifest records. +func packDigests(t *testing.T, dir string, exe, source, meta []byte) (integrity, cache string) { + t.Helper() + exePath := filepath.Join(dir, "exe") + require.NoError(t, os.WriteFile(exePath, exe, 0o755)) + sourcePath := filepath.Join(dir, "main.go") + require.NoError(t, os.WriteFile(sourcePath, source, 0o644)) + metaPath := filepath.Join(dir, "airflow-metadata.json") + require.NoError(t, os.WriteFile(metaPath, meta, 0o644)) + out := filepath.Join(dir, "bundle") + + require.NoError(t, runPack(io.Discard, io.Discard, &packOptions{ + executable: exePath, + source: sourcePath, + airflowMetadata: metaPath, + output: out, + })) + + _, gotMeta, err := bundlefooter.Read(out) + require.NoError(t, err) + var parsed struct { + Digests struct { + Integrity string `yaml:"integrity"` + Cache string `yaml:"cache"` + } `yaml:"digests"` + } + require.NoError(t, yaml.Unmarshal(gotMeta, &parsed)) + return parsed.Digests.Integrity, parsed.Digests.Cache +} + +func TestRunPack_RecordsDigests(t *testing.T) { + exe := []byte("binary-bytes") + source := []byte("package main\nfunc main() {}\n") + meta := []byte(`{"airflow_bundle_metadata_version":"1.0",` + + `"sdk":{"language":"go","version":"0.1.0","supervisor_schema_version":"2026-06-16"},` + + `"dags":{"my_dag":{"tasks":["t1"]}}}`) + + integrity, cache := packDigests(t, t.TempDir(), exe, source, meta) + + binaryHash := sha256.Sum256(exe) + assert.Equal( + t, + hex.EncodeToString(binaryHash[:]), + integrity, + "integrity is the trailer's binary_sha256", + ) + assert.Len(t, cache, sha256.Size*2) + assert.NotEqual(t, integrity, cache) + + t.Run("a repack of the same inputs", func(t *testing.T) { + gotIntegrity, gotCache := packDigests(t, t.TempDir(), exe, source, meta) + assert.Equal(t, integrity, gotIntegrity) + assert.Equal(t, cache, gotCache) + }) + + oneByte := func(b []byte, i int) []byte { + changed := bytes.Clone(b) + changed[i] ^= 1 + return changed + } + for name, tc := range map[string]struct { + exe, source, meta []byte + integrityChanges bool + }{ + "one source byte": {exe: exe, source: oneByte(source, len(source)-2), meta: meta}, + "one binary byte": {exe: oneByte(exe, 0), source: source, meta: meta, integrityChanges: true}, + "the manifest": { + exe: exe, + source: source, + meta: bytes.Replace(meta, []byte(`"0.1.0"`), []byte(`"0.1.1"`), 1), + }, + } { + t.Run(name, func(t *testing.T) { + gotIntegrity, gotCache := packDigests(t, t.TempDir(), tc.exe, tc.source, tc.meta) + assert.NotEqual(t, cache, gotCache) + if tc.integrityChanges { + assert.NotEqual(t, integrity, gotIntegrity) + } else { + assert.Equal(t, integrity, gotIntegrity) + } + }) + } +} + // With no explicit --output, packing "./bundle" from a package dir named // "bundle" derives a default output that collides with the pre-built binary. func TestRunPack_RejectsDefaultOutputAliasingExecutable(t *testing.T) { diff --git a/go-sdk/internal/airflowmetadata/airflowmetadata.go b/go-sdk/internal/airflowmetadata/airflowmetadata.go index d11412d722c44..2045daebb17f5 100644 --- a/go-sdk/internal/airflowmetadata/airflowmetadata.go +++ b/go-sdk/internal/airflowmetadata/airflowmetadata.go @@ -28,8 +28,8 @@ const FormatVersion = "1.0" // Manifest is the shape printed by a bundle binary's --airflow-metadata flag // (YAML by default, JSON under --format json). It mirrors -// airflow-metadata.schema.json minus the source field, which only the packer -// can resolve from the build inputs. The yaml tags let the packer's +// airflow-metadata.schema.json minus the source and digests fields, which only +// the packer can resolve from the build inputs. The yaml tags let the packer's // --airflow-metadata flag decode a captured manifest in either JSON or YAML // (the airflow-metadata.yaml in a bundle). type Manifest struct { diff --git a/go-sdk/internal/bundlefooter/footer.go b/go-sdk/internal/bundlefooter/footer.go index 6641ab3af08bc..d87d1b8b57389 100644 --- a/go-sdk/internal/bundlefooter/footer.go +++ b/go-sdk/internal/bundlefooter/footer.go @@ -98,7 +98,7 @@ func Append(execPath string, source, metadata []byte) error { // Compute binary_sha256 over the file as it stands now: at this point the // whole file is the binary region, so the digest matches what a reader // recomputes over [0, source_start) after the append. - binaryHash, err := hashFile(execPath) + binaryHash, err := HashFile(execPath) if err != nil { return fmt.Errorf("bundlefooter: hashing binary region of %s: %w", execPath, err) } @@ -247,8 +247,10 @@ func decodeTrailer(b [TrailerSize]byte) (Trailer, error) { return t, nil } -// hashFile computes SHA-256 over the entire contents of the file at path. -func hashFile(path string) ([sha256.Size]byte, error) { +// HashFile computes SHA-256 over the entire contents of the file at path. For an +// executable that has not had its trailer appended yet, that is the binary_sha256 +// Append records. +func HashFile(path string) ([sha256.Size]byte, error) { var sum [sha256.Size]byte f, err := os.Open(path) if err != nil { diff --git a/task-sdk/docs/airflow-metadata.schema.json b/task-sdk/docs/airflow-metadata.schema.json index d4a5cd240e712..b4619b0df3615 100644 --- a/task-sdk/docs/airflow-metadata.schema.json +++ b/task-sdk/docs/airflow-metadata.schema.json @@ -40,6 +40,23 @@ "description": "Original filename of the primary Dag source file (e.g. 'example.go'). Bundle formats with an embedded source region use this as its display name; other formats treat it as the logical authoring name.", "minLength": 1 }, + "digests": { + "type": "object", + "description": "Digests the packer records, so a consumer can compare a stored value instead of computing one. Each value is a lower-case hex SHA-256.", + "additionalProperties": true, + "properties": { + "integrity": { + "type": "string", + "description": "SHA-256 of the binary region: the same value as the trailer's binary_sha256. Task execution still recomputes the trailer's digest.", + "pattern": "^[0-9a-f]{64}$" + }, + "cache": { + "type": "string", + "description": "Stays the same across packs of the same inputs and changes when any region of the bundle changes. It is read, never recomputed, so it does not prove the bundle is intact. The producer defines how it is computed; consumers only compare it for equality.", + "pattern": "^[0-9a-f]{64}$" + } + } + }, "dags": { "type": "object", "description": "Mapping of dag_id to DAG entry. Every dag_id the bundle exposes must appear here.", diff --git a/task-sdk/docs/executable-bundle-spec.rst b/task-sdk/docs/executable-bundle-spec.rst index 27e48a0d04fde..b1f874ec7b4eb 100644 --- a/task-sdk/docs/executable-bundle-spec.rst +++ b/task-sdk/docs/executable-bundle-spec.rst @@ -221,6 +221,23 @@ Top-level keys: source-view panel and pick a syntax-highlighting mode from the extension. +``digests`` (mapping, optional) + Digests the packer records, so a consumer can compare a stored value + instead of computing one. Each value is a lower-case hex SHA-256. + + - ``integrity`` (string): SHA-256 of the binary region, the same value as + the trailer's ``binary_sha256``. Task execution still recomputes the + trailer's digest; this copy does not replace that check. + - ``cache`` (string): stays the same across packs of the same inputs and + changes when any region of the bundle changes. The Dag processor reads + it to tell whether a bundle changed since it last probed it. It is read, + never recomputed, so it does not prove the bundle is intact. The + producer defines how it is computed, and consumers only compare it for + equality. ``airflow-go-pack`` computes the SHA-256 of the bytes + ``airflow-go-pack cache digest v1`` and a NUL, followed by the raw + SHA-256 digests of the binary region, the source region, and the + manifest as rendered without the ``digests`` key, in that order. + ``dags`` (mapping, required) Mapping of ``dag_id`` to a *DAG entry*. Every ``dag_id`` the bundle exposes MUST appear here. The scanner uses these keys to match a DAG From c440bf74d680765b71e9760c16232e91e25a9a9e Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 15:57:44 +0000 Subject: [PATCH 54/94] Let the Executable coordinator probe a bundle's task handlers The Dag processor starts a coordinator's runtime to ask an artifact which task handlers it registers, and a coordinator opts in by building that command. The Go runtime now answers the request, so ExecutableCoordinator builds it: the artifact the Dag processor names, after the same trailer and digest check a task gets, with no dag_id lookup in the metadata. The real-bundle test is gated by AIRFLOW_LANG_SDK_REAL_PROBE_TESTS because it needs a Go toolchain and builds the example bundle. --- .../test_task_handler_processor_go.py | 140 ++++++++++++++++++ .../coordinators/executable/coordinator.py | 12 ++ .../executable/test_coordinator.py | 34 +++++ 3 files changed, 186 insertions(+) create mode 100644 airflow-core/tests/unit/dag_processing/test_task_handler_processor_go.py diff --git a/airflow-core/tests/unit/dag_processing/test_task_handler_processor_go.py b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_go.py new file mode 100644 index 0000000000000..d5692e478f545 --- /dev/null +++ b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_go.py @@ -0,0 +1,140 @@ +# 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. +"""Probe a packed Go SDK example bundle for its task handlers, end to end.""" + +from __future__ import annotations + +import json +import os +import shutil +import subprocess +from typing import TYPE_CHECKING +from unittest import mock + +import pytest +import structlog + +from airflow.dag_processing.processor import TaskHandlerDeclaration, TaskHandlerParam +from airflow.dag_processing.task_handler_processor import LangSDKTaskHandlerProcessorProcess +from airflow.sdk.execution_time import supervisor +from airflow.sdk.execution_time.coordinator import reset_coordinator_manager + +from tests_common.test_utils.config import conf_vars +from tests_common.test_utils.paths import AIRFLOW_ROOT_PATH + +if TYPE_CHECKING: + from pathlib import Path + +pytestmark = pytest.mark.skipif( + os.environ.get("AIRFLOW_LANG_SDK_REAL_PROBE_TESTS") != "1", + reason="set AIRFLOW_LANG_SDK_REAL_PROBE_TESTS=1 to build and probe a real Go bundle", +) + +STRING = {"type": "string"} +INT64 = {"type": "integer", "format": "int64"} +DOUBLE = {"type": "number", "format": "double"} + + +def _nullable(schema: dict) -> dict: + return {"anyOf": [schema, {"type": "null"}]} + + +@pytest.fixture(scope="module") +def go_bundle(tmp_path_factory) -> Path: + if shutil.which("go") is None: + pytest.skip("needs a Go toolchain on PATH") + bundle = tmp_path_factory.mktemp("go-task-handlers") / "example_dags" + completed = subprocess.run( + ["go", "tool", "airflow-go-pack", "--output", os.fspath(bundle), "./example/bundle"], + cwd=AIRFLOW_ROOT_PATH / "go-sdk", + env={**os.environ, "CGO_ENABLED": "0"}, + capture_output=True, + text=True, + check=False, + ) + assert completed.returncode == 0, completed.stderr + return bundle + + +@pytest.fixture(autouse=True) +def _go_coordinator(monkeypatch, tmp_path): + # Without cgo the Go SDK reads the current user from USER and HOME, and fails without them. + for name, value in (("USER", "airflow"), ("HOME", os.fspath(tmp_path))): + if not os.environ.get(name): + monkeypatch.setenv(name, value) + spec = {"go": {"classpath": "airflow.sdk.coordinators.executable.ExecutableCoordinator"}} + reset_coordinator_manager() + try: + # A bare fork, so the parse child sees this config. + with ( + conf_vars({("sdk", "coordinators"): json.dumps(spec)}), + mock.patch.object(supervisor, "_should_use_exec", return_value=False), + ): + yield + finally: + reset_coordinator_manager() + + +def test_probes_the_task_handlers_of_a_packed_go_bundle(go_bundle): + result = LangSDKTaskHandlerProcessorProcess.run( + coordinator="go", + path=go_bundle, + bundle_path=go_bundle.parent, + bundle_name="go-task-handlers", + artifact_rel_path=go_bundle.name, + dag_ids=["simple_dag", "taskflow_binding_dag", "not_registered"], + logger=structlog.get_logger(), + ) + + assert result.import_errors is None + assert set(result.task_handlers) == {"simple_dag", "taskflow_binding_dag"} + assert [d.task_id for d in result.task_handlers["simple_dag"]] == ["extract", "transform", "load"] + + declarations = {d.task_id: d for d in result.task_handlers["taskflow_binding_dag"]} + assert declarations["via_flat_args"] == TaskHandlerDeclaration( + task_id="via_flat_args", + binding="positional", + params=[ + TaskHandlerParam(name=None, required=True, value_schema=schema) + for schema in ( + STRING, + INT64, + DOUBLE, + {"type": "boolean"}, + _nullable({"type": "array", "items": STRING}), + {"type": "object"}, + _nullable({"type": "array", "items": INT64}), + _nullable(STRING), + ) + ], + ) + assert declarations["via_struct_arg_tag"] == TaskHandlerDeclaration( + task_id="via_struct_arg_tag", + binding="named", + params=[ + TaskHandlerParam(name="region_code", exact_name=True, required=False, value_schema=STRING), + TaskHandlerParam(name="threshold", exact_name=True, required=False, value_schema=DOUBLE), + ], + ) + assert declarations["via_flat_map"] == TaskHandlerDeclaration( + task_id="via_flat_map", + binding="named_or_whole", + params=[ + TaskHandlerParam(name="Region", required=False, value_schema=STRING), + TaskHandlerParam(name="Count", required=False, value_schema=INT64), + ], + ) diff --git a/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py index 237cc8dbbdf90..c2be8dd296e6c 100644 --- a/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py @@ -326,6 +326,9 @@ class ExecutableCoordinator(SubprocessCoordinator): """ Coordinator that launches a native executable subprocess for task execution. + The Dag processor also launches an executable bundle to ask which task handlers + it registers for the Dags of a Python file. + Configuration is taken from the ``[sdk] coordinators`` entry that constructs this instance:: @@ -349,3 +352,12 @@ def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], roots = self._get_scan_roots() bundle = _Bundle.find(roots, what.dag_id) return [str(bundle.path)], bundle.schema_version + + def _build_parse_task_handler_command(self, *, path: pathlib.Path) -> tuple[list[str], str | None]: + # The same trailer and digest check a task gets before its bundle runs. + if (metadata := _read_bundle_metadata(path)) is None: + raise ValueError( + f"{path} is not a valid executable bundle: the AFBNDL01 trailer is missing, " + "or its binary digest or metadata is invalid" + ) + return [os.fspath(path)], extract_supervisor_schema_version(metadata) 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 7071c9134bfb9..9c0eee769b42b 100644 --- a/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py @@ -409,6 +409,40 @@ def test_raises_when_dag_id_not_found(self, tmp_path): coordinator._build_execute_task_command(what=ti) +class TestBuildParseTaskHandlerCommand: + def test_returns_the_bundle_and_its_schema_version(self, tmp_path): + # The Dag ids in the metadata play no part: the Dag processor names the artifact. + bundle = _build_bundle(tmp_path / "etl", dag_ids=["other_dag"]) + + command, schema_version = ExecutableCoordinator()._build_parse_task_handler_command(path=bundle) + + assert command == [str(bundle)] + assert schema_version == "2026-06-16" + + @pytest.mark.parametrize( + "build", + [ + pytest.param(_make_executable, id="no-trailer"), + pytest.param( + lambda path: _build_bundle(path, binary_sha256=b"\x00" * 32), id="binary-digest-mismatch" + ), + ], + ) + def test_rejects_a_file_that_is_not_a_valid_bundle(self, tmp_path, build): + bundle = build(tmp_path / "etl") + + with pytest.raises(ValueError, match="is not a valid executable bundle"): + ExecutableCoordinator()._build_parse_task_handler_command(path=bundle) + + def test_rejects_a_bundle_without_a_schema_version(self, tmp_path): + metadata = _make_metadata(["etl"]) + del metadata["sdk"]["supervisor_schema_version"] + bundle = _build_bundle(tmp_path / "etl", metadata=metadata) + + with pytest.raises(ValueError, match="supervisor_schema_version"): + ExecutableCoordinator()._build_parse_task_handler_command(path=bundle) + + @pytest.fixture def bundles_dir(tmp_path): _build_bundle(tmp_path / "my_bundle", dag_ids=["tutorial_dag"]) From eb196107f46181213b1bd699a384bc4871f587d4 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 16:00:08 +0000 Subject: [PATCH 55/94] Read a Go bundle's stored cache digest without hashing it The Dag processor's fast path has to tell whether an executable bundle changed since its last probe without paying for a hash of the binary on every parse, so it compares the cache digest airflow-go-pack stored. Anything that is not a readable bundle with that digest reads as unknown, which the caller treats as changed. The real-bundle test also checks that a repack of the same inputs keeps the digest and a one-byte source change does not. --- .../test_task_handler_processor_go.py | 34 ++++++++++++---- .../coordinators/executable/coordinator.py | 29 ++++++++++++++ .../executable/test_coordinator.py | 39 +++++++++++++++++++ 3 files changed, 94 insertions(+), 8 deletions(-) diff --git a/airflow-core/tests/unit/dag_processing/test_task_handler_processor_go.py b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_go.py index d5692e478f545..b68c555b5c78d 100644 --- a/airflow-core/tests/unit/dag_processing/test_task_handler_processor_go.py +++ b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_go.py @@ -14,7 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""Probe a packed Go SDK example bundle for its task handlers, end to end.""" +"""Pack the Go SDK example bundle, probe it for its task handlers, and read its cache digest.""" from __future__ import annotations @@ -30,6 +30,7 @@ from airflow.dag_processing.processor import TaskHandlerDeclaration, TaskHandlerParam from airflow.dag_processing.task_handler_processor import LangSDKTaskHandlerProcessorProcess +from airflow.sdk.coordinators.executable.coordinator import read_cache_digest from airflow.sdk.execution_time import supervisor from airflow.sdk.execution_time.coordinator import reset_coordinator_manager @@ -53,14 +54,13 @@ def _nullable(schema: dict) -> dict: return {"anyOf": [schema, {"type": "null"}]} -@pytest.fixture(scope="module") -def go_bundle(tmp_path_factory) -> Path: - if shutil.which("go") is None: - pytest.skip("needs a Go toolchain on PATH") - bundle = tmp_path_factory.mktemp("go-task-handlers") / "example_dags" +GO_SDK_PATH = AIRFLOW_ROOT_PATH / "go-sdk" + + +def _pack(bundle: Path, *flags: str) -> Path: completed = subprocess.run( - ["go", "tool", "airflow-go-pack", "--output", os.fspath(bundle), "./example/bundle"], - cwd=AIRFLOW_ROOT_PATH / "go-sdk", + ["go", "tool", "airflow-go-pack", "--output", os.fspath(bundle), *flags, "./example/bundle"], + cwd=GO_SDK_PATH, env={**os.environ, "CGO_ENABLED": "0"}, capture_output=True, text=True, @@ -70,6 +70,13 @@ def go_bundle(tmp_path_factory) -> Path: return bundle +@pytest.fixture(scope="module") +def go_bundle(tmp_path_factory) -> Path: + if shutil.which("go") is None: + pytest.skip("needs a Go toolchain on PATH") + return _pack(tmp_path_factory.mktemp("go-task-handlers") / "example_dags") + + @pytest.fixture(autouse=True) def _go_coordinator(monkeypatch, tmp_path): # Without cgo the Go SDK reads the current user from USER and HOME, and fails without them. @@ -138,3 +145,14 @@ def test_probes_the_task_handlers_of_a_packed_go_bundle(go_bundle): TaskHandlerParam(name="Count", required=False, value_schema=INT64), ], ) + + +def test_a_repack_keeps_the_cache_digest_until_a_source_byte_changes(go_bundle, tmp_path): + digest = read_cache_digest(go_bundle) + assert digest is not None + + assert read_cache_digest(_pack(tmp_path / "unchanged")) == digest + + source = tmp_path / "main.go" + source.write_bytes((GO_SDK_PATH / "example" / "bundle" / "main.go").read_bytes() + b"\n") + assert read_cache_digest(_pack(tmp_path / "changed", "--source", os.fspath(source))) != digest diff --git a/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py index c2be8dd296e6c..e1be4ec3bae82 100644 --- a/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py @@ -245,6 +245,35 @@ def _read_bundle_metadata(path: pathlib.Path) -> dict[str, Any] | None: return None +def read_cache_digest(path: pathlib.Path) -> str | None: + """ + Return the cache digest stored in the metadata of the executable bundle at *path*. + + The digest changes whenever any region of the bundle changes, so the Dag processor can + compare it with the one it saw when it last probed the bundle. It is read, not computed: + the binary region is not hashed, so a match does not prove the bundle is intact. + + Returns ``None`` when the file is not a readable bundle or its metadata has no cache digest, + and the caller should treat the bundle as changed. + """ + try: + with open(path, "rb") as f: + footer = _Footer.read(f, path, os.fstat(f.fileno()).st_size) + if footer is None: + return None + f.seek(footer.metadata_start) + metadata_bytes = f.read(footer.metadata_len) + metadata = parse_metadata_mapping(metadata_bytes, source="bundle metadata") + except (OSError, ValueError) as exc: + log.debug("Cannot read the bundle's cache digest", path=str(path), error=str(exc)) + return None + digests = metadata.get("digests") + if not isinstance(digests, dict): + return None + cache = digests.get("cache") + return cache if isinstance(cache, str) and cache else None + + def _dag_ids(metadata: dict[str, Any]) -> set[str]: dags = metadata.get("dags") if not isinstance(dags, dict): 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 9c0eee769b42b..4737427727355 100644 --- a/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py @@ -39,6 +39,7 @@ _BinaryDigestCache, _Bundle, _digest_cache, + read_cache_digest, ) from airflow.sdk.execution_time.coordinator import BaseCoordinator from airflow.sdk.execution_time.supervisor import ActivitySubprocess @@ -409,6 +410,44 @@ def test_raises_when_dag_id_not_found(self, tmp_path): coordinator._build_execute_task_command(what=ti) +_CACHE_DIGEST = "c" * 64 + + +def _metadata_with_digests(**digests: str) -> dict: + return {**_make_metadata(["etl"]), "digests": digests} + + +class TestReadCacheDigest: + def test_returns_the_stored_digest(self, tmp_path): + metadata = _metadata_with_digests(integrity="a" * 64, cache=_CACHE_DIGEST) + bundle = _build_bundle(tmp_path / "etl", metadata=metadata) + + assert read_cache_digest(bundle) == _CACHE_DIGEST + + def test_does_not_hash_the_binary_region(self, tmp_path): + # A digest mismatch is caught when the bundle is launched, not here. + metadata = _metadata_with_digests(cache=_CACHE_DIGEST) + bundle = _build_bundle(tmp_path / "etl", metadata=metadata, binary_sha256=b"\x00" * 32) + + assert read_cache_digest(bundle) == _CACHE_DIGEST + + @pytest.mark.parametrize( + "build", + [ + pytest.param(lambda path: _build_bundle(path), id="no-digests"), + pytest.param( + lambda path: _build_bundle(path, metadata=_metadata_with_digests(integrity="a" * 64)), + id="no-cache-digest", + ), + pytest.param(lambda path: _build_bundle(path, metadata=b"\xff\xfe"), id="undecodable-metadata"), + pytest.param(_make_executable, id="not-a-bundle"), + pytest.param(lambda path: path, id="missing-file"), + ], + ) + def test_returns_none_without_a_readable_digest(self, tmp_path, build): + assert read_cache_digest(build(tmp_path / "etl")) is None + + class TestBuildParseTaskHandlerCommand: def test_returns_the_bundle_and_its_schema_version(self, tmp_path): # The Dag ids in the metadata play no part: the Dag processor names the artifact. From 6cf833f9f89f83f09cc66ec3fbc1def3ddb2aa00 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 16:02:00 +0000 Subject: [PATCH 56/94] Document the task handler parse for Language SDK authors The Go SDK is the first runtime to answer TaskHandlerParseRequest, and the Java and TypeScript runtimes follow the spec and the contributor guide. Neither said that the first frame can be a parse request, what the reply must contain, or which coordinator hook starts the runtime for it. The spec stays at 1.0: the lifecycle adds steps and renames no term. --- contributing-docs/30_new_language_sdk.rst | 54 ++++++++++++++++++++++- task-sdk/docs/lang-sdk-spec.rst | 18 ++++++++ 2 files changed, 70 insertions(+), 2 deletions(-) diff --git a/contributing-docs/30_new_language_sdk.rst b/contributing-docs/30_new_language_sdk.rst index 86c65efc716ea..63190f19527c2 100644 --- a/contributing-docs/30_new_language_sdk.rst +++ b/contributing-docs/30_new_language_sdk.rst @@ -126,6 +126,28 @@ or the task's own Dag bundle when it is unset, pinned for the whole task. Subclasses should scan those directories rather than locating artifacts themselves. +SubprocessCoordinator: implementing ``_build_parse_task_handler_command`` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +When a Python Dag file has stub tasks, the Dag processor asks the artifact that +implements them which task handlers it registers, and checks each stub task +against its handler. A coordinator opts in by building the command that starts +its runtime for one artifact: + +.. code-block:: python + + def _build_parse_task_handler_command(self, *, path: pathlib.Path) -> tuple[list[str], str | None]: ... + +*path* is the artifact the Dag processor picked, so the method does not search +for one. The returned pair follows the rules of ``_build_execute_task_command``: +no ``--comm`` or ``--logs`` flags, and the schema version the runtime +understands. The runtime then answers as described in +`Answering TaskHandlerParseRequest`_. + +The default raises ``NotImplementedError``, so a coordinator that does not +implement it cannot be probed. ``ExecutableCoordinator`` implements it for +executable bundles. + Supervisor Schema ~~~~~~~~~~~~~~~~~ @@ -195,8 +217,10 @@ as soon as possible. The supervisor verifies that the connecting peer belongs to the launched process tree, so the SDK MUST connect from the same process or one of its descendants. -Once both connections are accepted, the supervisor sends a ``StartupDetails`` -message on the comm socket to initiate execution. +Once both connections are accepted, the supervisor sends the first message on +the comm socket. ``StartupDetails`` starts a task. ``TaskHandlerParseRequest`` +comes from the Dag processor instead, and asks which task handlers the runtime +registers (see `Answering TaskHandlerParseRequest`_). Wire protocol ~~~~~~~~~~~~~ @@ -273,6 +297,32 @@ calls themselves) can already produce log records before the ``--logs`` socket in stage 3 exists to carry them. See `Logging`_ below for how to handle that gap. +Answering ``TaskHandlerParseRequest`` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The request names Dag ids. The SDK sends back one ``TaskHandlerParsingResult``, +waits for the supervisor's response to it, and exits. No user code runs; the +answer comes from the SDK's own task handler registrations. + +* ``task_handlers`` maps each requested Dag id to the declarations of its + handlers, in registration order. A requested Dag with no handler is left out + rather than mapped to an empty list. +* Each ``TaskHandlerDeclaration`` states in ``binding`` how stub-task arguments + bind to its ``params``: ``positional``, ``named`` (by name in any order, + ignoring case and underscores unless a param sets ``exact_name``), + ``named_or_whole`` (as ``named``, and one argument that matches no param is + decoded as the whole value), or ``named_open`` (as ``named``, but only the + listed names are checked). +* Each ``TaskHandlerParam`` has a ``name`` (``null`` when the SDK has no name + for a positional parameter), ``required`` (the handler has no default for it), + and a ``value_schema``: the JSON Schema of the values it accepts, in the + vocabulary ``@task.stub`` uses for Python annotations, or ``null`` when the SDK + cannot state one. +* An empty ``params`` list or ``task_handlers`` mapping is sent empty, never as + ``null``. + +``go-sdk/pkg/execution`` is a reference implementation. + Logging ~~~~~~~ diff --git a/task-sdk/docs/lang-sdk-spec.rst b/task-sdk/docs/lang-sdk-spec.rst index fed3c2d6f153f..b63043f829270 100644 --- a/task-sdk/docs/lang-sdk-spec.rst +++ b/task-sdk/docs/lang-sdk-spec.rst @@ -165,3 +165,21 @@ wrote. How ``fn`` reaches ``Context`` and ``Client`` is each SDK's own choice getters that read the scope opened at step 6 — and all this spec asks is that ``fn`` receives the same pair the SDK built at step 3. The terminal state at step 7 is one of ``SucceedTask``, ``RetryTask``, or ``TaskState``, and it is reported exactly once. + +Task handler parse lifecycle +---------------------------- + +The Dag processor starts the same process to check a Python Dag's ``@task.stub`` tasks against the +``TaskHandler`` registrations that implement them: + +.. code-block:: text + + 1 process start + 2 receive TaskHandlerParseRequest, which names dag_ids + 3 declare each registered TaskHandler whose dagId is one of them, in registration order + 4 send one TaskHandlerParsingResult, and wait for the supervisor's reply to it + 5 exit + +No ``fn`` runs: the answer comes from the registrations alone. A declaration states how the stub +task's arguments bind to the handler's parameters and which values each parameter accepts. A +requested ``dag_id`` with no ``TaskHandler`` is left out of the result rather than listed empty. From 9ecf47d459cb56f00487c9c14abfbb79f60a4ba7 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 16:33:08 +0000 Subject: [PATCH 57/94] Go SDK: take the bundle digests from the staged executable Review follow-up. The digests were computed from the input executable while the trailer's binary_sha256 is computed from the staged copy, so an executable replaced mid-pack would leave digests.integrity disagreeing with the trailer. Both now come from the copy. --- go-sdk/cmd/airflow-go-pack/inspect_test.go | 2 +- go-sdk/cmd/airflow-go-pack/pack.go | 43 +++++++++++++++------- go-sdk/cmd/airflow-go-pack/pack_test.go | 32 +++++++++++++++- 3 files changed, 61 insertions(+), 16 deletions(-) diff --git a/go-sdk/cmd/airflow-go-pack/inspect_test.go b/go-sdk/cmd/airflow-go-pack/inspect_test.go index 8208136e80176..5841f3d360e7a 100644 --- a/go-sdk/cmd/airflow-go-pack/inspect_test.go +++ b/go-sdk/cmd/airflow-go-pack/inspect_test.go @@ -42,7 +42,7 @@ func TestInspectCmd(t *testing.T) { " - \"t1\"\n", ) bundle := filepath.Join(dir, "bundle") - require.NoError(t, writeBundle(exe, bundle, source, manifest)) + require.NoError(t, writeBundle(exe, bundle, source, fixedManifest(manifest))) for _, tc := range []struct { name string diff --git a/go-sdk/cmd/airflow-go-pack/pack.go b/go-sdk/cmd/airflow-go-pack/pack.go index 39c8764d0bb0a..23d542832d5bb 100644 --- a/go-sdk/cmd/airflow-go-pack/pack.go +++ b/go-sdk/cmd/airflow-go-pack/pack.go @@ -179,24 +179,30 @@ func runPack(stdout, stderr io.Writer, opts *packOptions) error { if err != nil { return fmt.Errorf("reading source file: %w", err) } - baseManifest, err := renderManifest(meta, filepath.Base(sourcePath), nil) - if err != nil { - return fmt.Errorf("rendering manifest: %w", err) - } - digests, err := computeDigests(execPath, sourceBytes, baseManifest) - if err != nil { - return err - } - manifest, err := renderManifest(meta, filepath.Base(sourcePath), &digests) - if err != nil { - return fmt.Errorf("rendering manifest: %w", err) + // The digests are taken over the staged copy, the file the trailer's + // binary_sha256 is computed from, so the two agree even if the executable + // changes while packing. + renderWithDigests := func(stagedExec string) ([]byte, error) { + baseManifest, err := renderManifest(meta, filepath.Base(sourcePath), nil) + if err != nil { + return nil, fmt.Errorf("rendering manifest: %w", err) + } + digests, err := computeDigests(stagedExec, sourceBytes, baseManifest) + if err != nil { + return nil, err + } + manifest, err := renderManifest(meta, filepath.Base(sourcePath), &digests) + if err != nil { + return nil, fmt.Errorf("rendering manifest: %w", err) + } + return manifest, nil } // Assemble the bundle through a temp file and atomically move it into // place: we never mutate the build artefact or the user-supplied // --executable, and a failed pack never leaves a truncated or half-written // file at output. - if err := writeBundle(execPath, output, sourceBytes, manifest); err != nil { + if err := writeBundle(execPath, output, sourceBytes, renderWithDigests); err != nil { return err } @@ -676,11 +682,16 @@ func sameFile(a, b string) (bool, error) { // writeBundle assembles the bundle at output by copying the executable to a // temporary file in output's directory, appending the source+manifest footer -// to that copy, then atomically renaming it into place. Writing through a +// to that copy, then atomically renaming it into place. renderMetadata builds +// the manifest from the copy before the footer is appended. Writing through a // temp file keeps a failed pack from leaving a truncated or half-written // artefact at output, and guarantees the file being copied is never the same // open file as the destination. -func writeBundle(execPath, output string, source, metadata []byte) error { +func writeBundle( + execPath, output string, + source []byte, + renderMetadata func(stagedExec string) ([]byte, error), +) error { outDir := filepath.Dir(output) // The temp file and the atomic rename both live in output's directory, so it // must exist. Create it for the user (e.g. --output ./bin/bundle with no @@ -704,6 +715,10 @@ func writeBundle(execPath, output string, source, metadata []byte) error { if err := copyFile(execPath, tmpPath, 0o755); err != nil { return fmt.Errorf("writing %s: %w", output, err) } + metadata, err := renderMetadata(tmpPath) + if err != nil { + return err + } if err := bundlefooter.Append(tmpPath, source, metadata); err != nil { return err } diff --git a/go-sdk/cmd/airflow-go-pack/pack_test.go b/go-sdk/cmd/airflow-go-pack/pack_test.go index b001d6b6619fb..20c030dcb67ad 100644 --- a/go-sdk/cmd/airflow-go-pack/pack_test.go +++ b/go-sdk/cmd/airflow-go-pack/pack_test.go @@ -451,6 +451,11 @@ func TestRunPack_AcceptsYAMLMetadataFile(t *testing.T) { ) } +// fixedManifest is a writeBundle renderMetadata that ignores the staged executable. +func fixedManifest(manifest []byte) func(string) ([]byte, error) { + return func(string) ([]byte, error) { return manifest, nil } +} + // packDigests packs the given executable, source and JSON manifest in dir, and // returns the digests the packed bundle's manifest records. func packDigests(t *testing.T, dir string, exe, source, meta []byte) (integrity, cache string) { @@ -611,13 +616,38 @@ func TestTargetPlatform(t *testing.T) { // When the --output parent directory does not exist, the packer must create it // instead of failing with an opaque temp-file error. +// The manifest's digests must describe the copy the trailer is computed from, +// not the input executable, which can change while packing. +func TestWriteBundle_RendersMetadataFromTheStagedCopy(t *testing.T) { + dir := t.TempDir() + exe := filepath.Join(dir, "input-bin") + require.NoError(t, os.WriteFile(exe, []byte("binary-bytes"), 0o755)) + + var staged string + var stagedBytes []byte + require.NoError(t, writeBundle(exe, filepath.Join(dir, "bundle"), []byte("source"), + func(path string) ([]byte, error) { + staged = path + var err error + stagedBytes, err = os.ReadFile(path) + return []byte("manifest"), err + }, + )) + + assert.NotEqual(t, exe, staged) + assert.Equal(t, []byte("binary-bytes"), stagedBytes) +} + func TestWriteBundle_CreatesMissingOutputDir(t *testing.T) { dir := t.TempDir() exe := filepath.Join(dir, "input-bin") require.NoError(t, os.WriteFile(exe, []byte("binary-bytes"), 0o755)) output := filepath.Join(dir, "bin", "nested", "bundle") - require.NoError(t, writeBundle(exe, output, []byte("source"), []byte("manifest"))) + require.NoError( + t, + writeBundle(exe, output, []byte("source"), fixedManifest([]byte("manifest"))), + ) info, err := os.Stat(output) require.NoError(t, err, "bundle should be written into the created directory") From 262ca91a98a8710319023547de27bd78c7b6514b Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 16:33:57 +0000 Subject: [PATCH 58/94] Resolve the probed bundle path and tighten the digest docs Review follow-up. The parse command returned the path as given, and exec searches PATH for a bare relative name, so it is now absolute as the task command is. Its error names the unreadable-file case too. An empty stored cache digest is now tested as unknown. The contributor guide lost the exactly-one-argument condition of named_or_whole, and the digests docs said every value is a SHA-256 while allowing keys a consumer ignores. --- contributing-docs/30_new_language_sdk.rst | 6 +++--- task-sdk/docs/airflow-metadata.schema.json | 2 +- task-sdk/docs/executable-bundle-spec.rst | 3 ++- .../sdk/coordinators/executable/coordinator.py | 7 ++++--- .../coordinators/executable/test_coordinator.py | 14 +++++++++++++- 5 files changed, 23 insertions(+), 9 deletions(-) diff --git a/contributing-docs/30_new_language_sdk.rst b/contributing-docs/30_new_language_sdk.rst index 63190f19527c2..806bf94739253 100644 --- a/contributing-docs/30_new_language_sdk.rst +++ b/contributing-docs/30_new_language_sdk.rst @@ -310,9 +310,9 @@ answer comes from the SDK's own task handler registrations. * Each ``TaskHandlerDeclaration`` states in ``binding`` how stub-task arguments bind to its ``params``: ``positional``, ``named`` (by name in any order, ignoring case and underscores unless a param sets ``exact_name``), - ``named_or_whole`` (as ``named``, and one argument that matches no param is - decoded as the whole value), or ``named_open`` (as ``named``, but only the - listed names are checked). + ``named_or_whole`` (as ``named``, except that when there is exactly one + argument and it matches no param, it is decoded as the whole value), or + ``named_open`` (as ``named``, but only the listed names are checked). * Each ``TaskHandlerParam`` has a ``name`` (``null`` when the SDK has no name for a positional parameter), ``required`` (the handler has no default for it), and a ``value_schema``: the JSON Schema of the values it accepts, in the diff --git a/task-sdk/docs/airflow-metadata.schema.json b/task-sdk/docs/airflow-metadata.schema.json index b4619b0df3615..d2a8f35ba000a 100644 --- a/task-sdk/docs/airflow-metadata.schema.json +++ b/task-sdk/docs/airflow-metadata.schema.json @@ -42,7 +42,7 @@ }, "digests": { "type": "object", - "description": "Digests the packer records, so a consumer can compare a stored value instead of computing one. Each value is a lower-case hex SHA-256.", + "description": "Digests the packer records, so a consumer can compare a stored value instead of computing one. Consumers ignore keys they do not know.", "additionalProperties": true, "properties": { "integrity": { diff --git a/task-sdk/docs/executable-bundle-spec.rst b/task-sdk/docs/executable-bundle-spec.rst index b1f874ec7b4eb..83d91fbc5ca77 100644 --- a/task-sdk/docs/executable-bundle-spec.rst +++ b/task-sdk/docs/executable-bundle-spec.rst @@ -223,7 +223,8 @@ Top-level keys: ``digests`` (mapping, optional) Digests the packer records, so a consumer can compare a stored value - instead of computing one. Each value is a lower-case hex SHA-256. + instead of computing one. Both are lower-case hex SHA-256 digests, and + consumers ignore keys they do not know. - ``integrity`` (string): SHA-256 of the binary region, the same value as the trailer's ``binary_sha256``. Task execution still recomputes the diff --git a/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py index e1be4ec3bae82..9ca3b7c5d40db 100644 --- a/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py @@ -386,7 +386,8 @@ def _build_parse_task_handler_command(self, *, path: pathlib.Path) -> tuple[list # The same trailer and digest check a task gets before its bundle runs. if (metadata := _read_bundle_metadata(path)) is None: raise ValueError( - f"{path} is not a valid executable bundle: the AFBNDL01 trailer is missing, " - "or its binary digest or metadata is invalid" + f"{path} is not a valid executable bundle: it cannot be read, has no AFBNDL01 " + "trailer, or its binary digest or metadata is invalid" ) - return [os.fspath(path)], extract_supervisor_schema_version(metadata) + # Absolute, as for a task, so exec never searches PATH for it. + return [os.fspath(path.resolve())], extract_supervisor_schema_version(metadata) 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 4737427727355..8413dfada9422 100644 --- a/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py @@ -439,6 +439,10 @@ def test_does_not_hash_the_binary_region(self, tmp_path): lambda path: _build_bundle(path, metadata=_metadata_with_digests(integrity="a" * 64)), id="no-cache-digest", ), + pytest.param( + lambda path: _build_bundle(path, metadata=_metadata_with_digests(cache="")), + id="empty-cache-digest", + ), pytest.param(lambda path: _build_bundle(path, metadata=b"\xff\xfe"), id="undecodable-metadata"), pytest.param(_make_executable, id="not-a-bundle"), pytest.param(lambda path: path, id="missing-file"), @@ -455,9 +459,17 @@ def test_returns_the_bundle_and_its_schema_version(self, tmp_path): command, schema_version = ExecutableCoordinator()._build_parse_task_handler_command(path=bundle) - assert command == [str(bundle)] + assert command == [str(bundle.resolve())] assert schema_version == "2026-06-16" + def test_returns_an_absolute_path(self, tmp_path, monkeypatch): + _build_bundle(tmp_path / "etl") + monkeypatch.chdir(tmp_path) + + command, _ = ExecutableCoordinator()._build_parse_task_handler_command(path=Path("etl")) + + assert command == [str((tmp_path / "etl").resolve())] + @pytest.mark.parametrize( "build", [ From f731a54bab29ffc8526bc5aeb47a4093894be5a3 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 1 Oct 2026 05:59:16 +0000 Subject: [PATCH 59/94] Move Go SDK comments back to the code they describe The staged-copy test was inserted under the output directory test's doc comment, so that comment documented the wrong test and repeated a reason pack.go already gives. The tests also repeated why the parse reply must not send null, while the code that builds the empty slice said nothing. --- go-sdk/cmd/airflow-go-pack/pack_test.go | 6 ++---- go-sdk/pkg/binding/declaration.go | 3 ++- go-sdk/pkg/binding/declaration_test.go | 1 - go-sdk/pkg/execution/task_handler_parse_test.go | 1 - 4 files changed, 4 insertions(+), 7 deletions(-) diff --git a/go-sdk/cmd/airflow-go-pack/pack_test.go b/go-sdk/cmd/airflow-go-pack/pack_test.go index 20c030dcb67ad..ce5f90894dc4c 100644 --- a/go-sdk/cmd/airflow-go-pack/pack_test.go +++ b/go-sdk/cmd/airflow-go-pack/pack_test.go @@ -614,10 +614,6 @@ func TestTargetPlatform(t *testing.T) { }) } -// When the --output parent directory does not exist, the packer must create it -// instead of failing with an opaque temp-file error. -// The manifest's digests must describe the copy the trailer is computed from, -// not the input executable, which can change while packing. func TestWriteBundle_RendersMetadataFromTheStagedCopy(t *testing.T) { dir := t.TempDir() exe := filepath.Join(dir, "input-bin") @@ -638,6 +634,8 @@ func TestWriteBundle_RendersMetadataFromTheStagedCopy(t *testing.T) { assert.Equal(t, []byte("binary-bytes"), stagedBytes) } +// When the --output parent directory does not exist, the packer must create it +// instead of failing with an opaque temp-file error. func TestWriteBundle_CreatesMissingOutputDir(t *testing.T) { dir := t.TempDir() exe := filepath.Join(dir, "input-bin") diff --git a/go-sdk/pkg/binding/declaration.go b/go-sdk/pkg/binding/declaration.go index 9d75d41af077d..e402095308dbe 100644 --- a/go-sdk/pkg/binding/declaration.go +++ b/go-sdk/pkg/binding/declaration.go @@ -36,7 +36,8 @@ func (p *Plan) Declare(taskID string) genmodels.TaskHandlerDeclaration { decl := genmodels.TaskHandlerDeclaration{ TaskID: taskID, Binding: genmodels.TaskHandlerDeclarationBindingPositional, - Params: []genmodels.TaskHandlerParam{}, + // A nil slice would go out as null, which the Dag processor rejects. + Params: []genmodels.TaskHandlerParam{}, } for _, plan := range p.params { switch plan.kind { diff --git a/go-sdk/pkg/binding/declaration_test.go b/go-sdk/pkg/binding/declaration_test.go index 94f8e70f2608a..c83547a4a4325 100644 --- a/go-sdk/pkg/binding/declaration_test.go +++ b/go-sdk/pkg/binding/declaration_test.go @@ -238,7 +238,6 @@ func TestDeclareUntaggedStruct(t *testing.T) { func TestDeclareContextOnly(t *testing.T) { decl := declare(t, func(actx contexttest.Context) error { return nil }) - // A nil slice would go out as null, which the Dag processor rejects. assert.Equal(t, genmodels.TaskHandlerDeclaration{ TaskID: "task", Binding: genmodels.TaskHandlerDeclarationBindingPositional, diff --git a/go-sdk/pkg/execution/task_handler_parse_test.go b/go-sdk/pkg/execution/task_handler_parse_test.go index f077e27692c41..3cd2d966b60d2 100644 --- a/go-sdk/pkg/execution/task_handler_parse_test.go +++ b/go-sdk/pkg/execution/task_handler_parse_test.go @@ -71,6 +71,5 @@ func TestDeclareTaskHandlersMatchingNothing(t *testing.T) { DagIds: []string{"etl"}, }) - // A nil map would go out as null, which the Dag processor rejects. assert.Equal(t, genmodels.TaskHandlers{}, result.TaskHandlers) } From a9d24bac0540876e25b7b1b7e43ff95557f9e0fc Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 1 Oct 2026 05:59:40 +0000 Subject: [PATCH 60/94] Tidy names, mocks and cases in the Go probe and digest tests An unspecced patch of _should_use_exec accepts any call, noun-named helpers read as attributes, and the two stored-digest tests differed only in their bundle inputs. --- .../test_task_handler_processor_go.py | 10 +++--- .../executable/test_coordinator.py | 31 +++++++++++-------- 2 files changed, 23 insertions(+), 18 deletions(-) diff --git a/airflow-core/tests/unit/dag_processing/test_task_handler_processor_go.py b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_go.py index b68c555b5c78d..10ac65892639c 100644 --- a/airflow-core/tests/unit/dag_processing/test_task_handler_processor_go.py +++ b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_go.py @@ -50,7 +50,7 @@ DOUBLE = {"type": "number", "format": "double"} -def _nullable(schema: dict) -> dict: +def _make_nullable(schema: dict) -> dict: return {"anyOf": [schema, {"type": "null"}]} @@ -89,7 +89,7 @@ def _go_coordinator(monkeypatch, tmp_path): # A bare fork, so the parse child sees this config. with ( conf_vars({("sdk", "coordinators"): json.dumps(spec)}), - mock.patch.object(supervisor, "_should_use_exec", return_value=False), + mock.patch.object(supervisor, "_should_use_exec", autospec=True, return_value=False), ): yield finally: @@ -122,10 +122,10 @@ def test_probes_the_task_handlers_of_a_packed_go_bundle(go_bundle): INT64, DOUBLE, {"type": "boolean"}, - _nullable({"type": "array", "items": STRING}), + _make_nullable({"type": "array", "items": STRING}), {"type": "object"}, - _nullable({"type": "array", "items": INT64}), - _nullable(STRING), + _make_nullable({"type": "array", "items": INT64}), + _make_nullable(STRING), ) ], ) 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 8413dfada9422..ffe4e432dd646 100644 --- a/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py @@ -413,21 +413,26 @@ def test_raises_when_dag_id_not_found(self, tmp_path): _CACHE_DIGEST = "c" * 64 -def _metadata_with_digests(**digests: str) -> dict: +def _make_metadata_with_digests(**digests: str) -> dict: return {**_make_metadata(["etl"]), "digests": digests} class TestReadCacheDigest: - def test_returns_the_stored_digest(self, tmp_path): - metadata = _metadata_with_digests(integrity="a" * 64, cache=_CACHE_DIGEST) - bundle = _build_bundle(tmp_path / "etl", metadata=metadata) - - assert read_cache_digest(bundle) == _CACHE_DIGEST - - def test_does_not_hash_the_binary_region(self, tmp_path): - # A digest mismatch is caught when the bundle is launched, not here. - metadata = _metadata_with_digests(cache=_CACHE_DIGEST) - bundle = _build_bundle(tmp_path / "etl", metadata=metadata, binary_sha256=b"\x00" * 32) + @pytest.mark.parametrize( + "bundle_kwargs", + [ + pytest.param( + {"metadata": _make_metadata_with_digests(integrity="a" * 64, cache=_CACHE_DIGEST)}, + id="valid-bundle", + ), + pytest.param( + {"metadata": _make_metadata_with_digests(cache=_CACHE_DIGEST), "binary_sha256": b"\x00" * 32}, + id="binary-digest-mismatch-not-checked", + ), + ], + ) + def test_returns_the_stored_digest(self, tmp_path, bundle_kwargs): + bundle = _build_bundle(tmp_path / "etl", **bundle_kwargs) assert read_cache_digest(bundle) == _CACHE_DIGEST @@ -436,11 +441,11 @@ def test_does_not_hash_the_binary_region(self, tmp_path): [ pytest.param(lambda path: _build_bundle(path), id="no-digests"), pytest.param( - lambda path: _build_bundle(path, metadata=_metadata_with_digests(integrity="a" * 64)), + lambda path: _build_bundle(path, metadata=_make_metadata_with_digests(integrity="a" * 64)), id="no-cache-digest", ), pytest.param( - lambda path: _build_bundle(path, metadata=_metadata_with_digests(cache="")), + lambda path: _build_bundle(path, metadata=_make_metadata_with_digests(cache="")), id="empty-cache-digest", ), pytest.param(lambda path: _build_bundle(path, metadata=b"\xff\xfe"), id="undecodable-metadata"), From f479c8aef71842b657f240c43b56f1dfce49ee96 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 1 Oct 2026 07:33:06 +0000 Subject: [PATCH 61/94] Go SDK: answer a task handler parse with every handler The parse request no longer names Dags, and the Dag processor reuses an artifact's answer for every Dag file that uses the artifact, so the answer must not depend on what was asked. The genmodels regeneration also picks up UpdateDagRunNote (#69403), which main's schema gained without a Go regen. --- .../test_task_handler_processor_go.py | 8 +++++-- go-sdk/README.md | 4 ++-- go-sdk/pkg/execution/integration_test.go | 1 - go-sdk/pkg/execution/messages_test.go | 4 +--- go-sdk/pkg/execution/server.go | 4 ++-- go-sdk/pkg/execution/task_handler_parse.go | 12 ++-------- .../pkg/execution/task_handler_parse_test.go | 24 ++++++++++--------- 7 files changed, 26 insertions(+), 31 deletions(-) diff --git a/airflow-core/tests/unit/dag_processing/test_task_handler_processor_go.py b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_go.py index 10ac65892639c..3294eaf2b9aad 100644 --- a/airflow-core/tests/unit/dag_processing/test_task_handler_processor_go.py +++ b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_go.py @@ -103,12 +103,16 @@ def test_probes_the_task_handlers_of_a_packed_go_bundle(go_bundle): bundle_path=go_bundle.parent, bundle_name="go-task-handlers", artifact_rel_path=go_bundle.name, - dag_ids=["simple_dag", "taskflow_binding_dag", "not_registered"], logger=structlog.get_logger(), ) assert result.import_errors is None - assert set(result.task_handlers) == {"simple_dag", "taskflow_binding_dag"} + assert set(result.task_handlers) == { + "simple_dag", + "concurrent_xcom_dag", + "taskflow_binding_dag", + "variable_write_dag", + } assert [d.task_id for d in result.task_handlers["simple_dag"]] == ["extract", "transform", "load"] declarations = {d.task_id: d for d in result.task_handlers["taskflow_binding_dag"]} diff --git a/go-sdk/README.md b/go-sdk/README.md index 9fdacd4a78617..adb0feed4c11a 100644 --- a/go-sdk/README.md +++ b/go-sdk/README.md @@ -390,8 +390,8 @@ Python supervisor / task runner - The Python runtime is the worker. It proxies every `GetConnection` / `GetVariable` / `GetXCom` / `SetXCom` call through to the Execution API. The Go binary just runs the task function. - The Dag processor launches the binary the same way to check a Python Dag's stub tasks. It sends a - `TaskHandlerParseRequest` that names Dag ids. For each of those Dags, the binary replies with the task - handlers it registers, in registration order, and how each handler's parameters bind. Then it exits. + `TaskHandlerParseRequest`, and the binary replies with every task handler it registers, keyed by Dag + id in registration order, and how each handler's parameters bind. Then it exits. The Go side of the protocol is implemented in `pkg/execution/`. On the Python side it is the `ExecutableCoordinator` in `task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py`. diff --git a/go-sdk/pkg/execution/integration_test.go b/go-sdk/pkg/execution/integration_test.go index 94fd0a0d9c274..84fb4af0b6c9f 100644 --- a/go-sdk/pkg/execution/integration_test.go +++ b/go-sdk/pkg/execution/integration_test.go @@ -946,7 +946,6 @@ func startTaskHandlerParse(t *testing.T) (net.Conn, IncomingFrame, chan error) { payload, err := encodeRequest(0, map[string]any{ "type": "TaskHandlerParseRequest", "file": "/bundles/go/etl", - "dag_ids": []string{"etl", "missing"}, "bundle_path": "/bundles/go", "bundle_name": "go", }) diff --git a/go-sdk/pkg/execution/messages_test.go b/go-sdk/pkg/execution/messages_test.go index f86baaa0f5fa6..3cf8e1a742e85 100644 --- a/go-sdk/pkg/execution/messages_test.go +++ b/go-sdk/pkg/execution/messages_test.go @@ -466,7 +466,6 @@ func TestDecodeIncomingBodyDispatch(t *testing.T) { raw := marshalBody(t, map[string]any{ "type": "TaskHandlerParseRequest", "file": "/bundles/go/etl", - "dag_ids": []string{"etl"}, "bundle_path": "/bundles/go", "bundle_name": "go", }) @@ -475,14 +474,13 @@ func TestDecodeIncomingBodyDispatch(t *testing.T) { assert.Equal(t, &genmodels.TaskHandlerParseRequest{ Type: "TaskHandlerParseRequest", File: "/bundles/go/etl", - DagIds: []string{"etl"}, BundlePath: "/bundles/go", BundleName: "go", }, result) }) t.Run("malformed TaskHandlerParseRequest", func(t *testing.T) { - raw := marshalBody(t, map[string]any{"type": "TaskHandlerParseRequest", "dag_ids": "etl"}) + raw := marshalBody(t, map[string]any{"type": "TaskHandlerParseRequest", "file": 42}) _, err := decodeIncomingBody(raw) assert.ErrorContains(t, err, "decoding TaskHandlerParseRequest") }) diff --git a/go-sdk/pkg/execution/server.go b/go-sdk/pkg/execution/server.go index 789f998704d75..2596a53ceb823 100644 --- a/go-sdk/pkg/execution/server.go +++ b/go-sdk/pkg/execution/server.go @@ -23,7 +23,7 @@ // The first inbound frame on the comm socket selects what the runtime does. A // StartupDetails message drives multi-round task execution. A // TaskHandlerParseRequest from the Dag processor gets one TaskHandlerParsingResult -// that declares the task handlers registered for the requested Dags. +// that declares every registered task handler. // // See go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md. package execution @@ -172,7 +172,7 @@ func Serve(b bundle.Registry, commAddr, logsAddr string) error { logger.Debug("Task execution complete") case *genmodels.TaskHandlerParseRequest: - logger.Debug("Task handler parse mode", "dag_ids", msg.DagIds) + logger.Debug("Task handler parse mode", "file", msg.File) result := declareTaskHandlers(b, msg) sendCtx, cancel := context.WithTimeout(ctx, terminalSendTimeout) defer cancel() diff --git a/go-sdk/pkg/execution/task_handler_parse.go b/go-sdk/pkg/execution/task_handler_parse.go index 17d0359b2e2d7..1e2d5aca2f5f0 100644 --- a/go-sdk/pkg/execution/task_handler_parse.go +++ b/go-sdk/pkg/execution/task_handler_parse.go @@ -22,23 +22,15 @@ import ( "github.com/apache/airflow/go-sdk/pkg/execution/genmodels" ) -// declareTaskHandlers answers a TaskHandlerParseRequest with the declarations of the task -// handlers registered for the requested Dags, each Dag's in registration order. A requested -// Dag with no handler is left out rather than mapped to an empty list. +// declareTaskHandlers answers a TaskHandlerParseRequest with the declarations of every +// registered task handler, keyed by Dag id, each Dag's in registration order. func declareTaskHandlers( b bundle.Registry, req *genmodels.TaskHandlerParseRequest, ) genmodels.TaskHandlerParsingResult { - requested := make(map[string]bool, len(req.DagIds)) - for _, dagID := range req.DagIds { - requested[dagID] = true - } // A nil map would go out as null, which the Dag processor rejects. handlers := genmodels.TaskHandlers{} for _, info := range b.ListTaskHandlers() { - if !requested[info.DagID] { - continue - } // Declare only what a task run could find. task, ok := b.LookupTask(info.DagID, info.TaskID) if !ok { diff --git a/go-sdk/pkg/execution/task_handler_parse_test.go b/go-sdk/pkg/execution/task_handler_parse_test.go index 3cd2d966b60d2..6b6e7169f981e 100644 --- a/go-sdk/pkg/execution/task_handler_parse_test.go +++ b/go-sdk/pkg/execution/task_handler_parse_test.go @@ -37,10 +37,10 @@ func TestDeclareTaskHandlers(t *testing.T) { // Listed, but a task run could not find it. registry.order = append(registry.order, bundle.TaskHandlerInfo{DagID: "etl", TaskID: "ghost"}) - result := declareTaskHandlers(registry, &genmodels.TaskHandlerParseRequest{ - File: "/bundles/go/etl", - DagIds: []string{"etl", "missing"}, - }) + result := declareTaskHandlers( + registry, + &genmodels.TaskHandlerParseRequest{File: "/bundles/go/etl"}, + ) assert.Equal(t, genmodels.TaskHandlerParsingResult{ Fileloc: "/bundles/go/etl", @@ -59,17 +59,19 @@ func TestDeclareTaskHandlers(t *testing.T) { Params: []genmodels.TaskHandlerParam{}, }, }, + "report": { + { + TaskID: "render", + Binding: genmodels.TaskHandlerDeclarationBindingPositional, + Params: []genmodels.TaskHandlerParam{}, + }, + }, }, }, result) } -func TestDeclareTaskHandlersMatchingNothing(t *testing.T) { - registry := newHandlerRegistry() - registry.add("report", "render", simpleTask) - - result := declareTaskHandlers(registry, &genmodels.TaskHandlerParseRequest{ - DagIds: []string{"etl"}, - }) +func TestDeclareTaskHandlersWithoutHandlers(t *testing.T) { + result := declareTaskHandlers(newHandlerRegistry(), &genmodels.TaskHandlerParseRequest{}) assert.Equal(t, genmodels.TaskHandlers{}, result.TaskHandlers) } From 5155067e2c7f655e695c332cca78a3921c2a4fdc Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 1 Oct 2026 07:33:16 +0000 Subject: [PATCH 62/94] State that a task handler parse answer depends only on the artifact A runtime that still filtered or varied its answer would break other Dag files, because the Dag processor reuses one artifact's answer for every Dag file that uses it. --- contributing-docs/30_new_language_sdk.rst | 12 ++++++------ task-sdk/docs/lang-sdk-spec.rst | 11 ++++++----- 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/contributing-docs/30_new_language_sdk.rst b/contributing-docs/30_new_language_sdk.rst index 806bf94739253..37be47e9f84c5 100644 --- a/contributing-docs/30_new_language_sdk.rst +++ b/contributing-docs/30_new_language_sdk.rst @@ -300,13 +300,13 @@ gap. Answering ``TaskHandlerParseRequest`` ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -The request names Dag ids. The SDK sends back one ``TaskHandlerParsingResult``, -waits for the supervisor's response to it, and exits. No user code runs; the -answer comes from the SDK's own task handler registrations. +The SDK sends back one ``TaskHandlerParsingResult``, waits for the supervisor's +response to it, and exits. No user code runs; the answer comes from the SDK's +own task handler registrations. It MUST declare every registered handler and +depend only on the artifact, never on the request. -* ``task_handlers`` maps each requested Dag id to the declarations of its - handlers, in registration order. A requested Dag with no handler is left out - rather than mapped to an empty list. +* ``task_handlers`` maps each Dag id the artifact registers handlers for to + their declarations, in registration order. * Each ``TaskHandlerDeclaration`` states in ``binding`` how stub-task arguments bind to its ``params``: ``positional``, ``named`` (by name in any order, ignoring case and underscores unless a param sets ``exact_name``), diff --git a/task-sdk/docs/lang-sdk-spec.rst b/task-sdk/docs/lang-sdk-spec.rst index b63043f829270..93b2c1d4296e6 100644 --- a/task-sdk/docs/lang-sdk-spec.rst +++ b/task-sdk/docs/lang-sdk-spec.rst @@ -175,11 +175,12 @@ The Dag processor starts the same process to check a Python Dag's ``@task.stub`` .. code-block:: text 1 process start - 2 receive TaskHandlerParseRequest, which names dag_ids - 3 declare each registered TaskHandler whose dagId is one of them, in registration order + 2 receive TaskHandlerParseRequest + 3 declare every registered TaskHandler, keyed by dagId, in registration order 4 send one TaskHandlerParsingResult, and wait for the supervisor's reply to it 5 exit -No ``fn`` runs: the answer comes from the registrations alone. A declaration states how the stub -task's arguments bind to the handler's parameters and which values each parameter accepts. A -requested ``dag_id`` with no ``TaskHandler`` is left out of the result rather than listed empty. +No ``fn`` runs: the answer comes from the registrations alone. It must depend only on the artifact, +never on the request. A declaration states how the stub task's arguments bind to the handler's +parameters and which values each parameter accepts. An artifact that registers no ``TaskHandler`` +answers with an empty mapping. From 3d9eaf255fe5b345127a30f193e6881bc173f80d Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 1 Oct 2026 17:39:41 +0000 Subject: [PATCH 63/94] Go SDK: declare every lone struct handler as named The supervisor schema now has only positional and named bindings, nullable params, and no required. An untagged struct declares named like a tagged one: the Dag processor no longer warns about one unmatched argument, which covers the whole-value case named_or_whole stood for. Params goes out as an empty list rather than null, since null would say that the handler's params cannot be listed. The genmodels are regenerated with go generate from the new snapshot. --- .../test_task_handler_processor_go.py | 18 ++++++++++----- go-sdk/pkg/binding/declaration.go | 23 +++++++------------ go-sdk/pkg/binding/declaration_test.go | 18 +++++++-------- go-sdk/pkg/execution/integration_test.go | 5 +--- .../pkg/execution/task_handler_parse_test.go | 8 +++---- 5 files changed, 34 insertions(+), 38 deletions(-) diff --git a/airflow-core/tests/unit/dag_processing/test_task_handler_processor_go.py b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_go.py index 3294eaf2b9aad..6dd1694d17ac9 100644 --- a/airflow-core/tests/unit/dag_processing/test_task_handler_processor_go.py +++ b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_go.py @@ -120,7 +120,7 @@ def test_probes_the_task_handlers_of_a_packed_go_bundle(go_bundle): task_id="via_flat_args", binding="positional", params=[ - TaskHandlerParam(name=None, required=True, value_schema=schema) + TaskHandlerParam(name=None, value_schema=schema) for schema in ( STRING, INT64, @@ -137,16 +137,22 @@ def test_probes_the_task_handlers_of_a_packed_go_bundle(go_bundle): task_id="via_struct_arg_tag", binding="named", params=[ - TaskHandlerParam(name="region_code", exact_name=True, required=False, value_schema=STRING), - TaskHandlerParam(name="threshold", exact_name=True, required=False, value_schema=DOUBLE), + TaskHandlerParam(name="region_code", exact_name=True, value_schema=STRING), + TaskHandlerParam(name="threshold", exact_name=True, value_schema=DOUBLE), ], ) + # The Dag passes this handler an argument it does not take, which only warns under named binding. + assert declarations["via_struct_more_args"] == TaskHandlerDeclaration( + task_id="via_struct_more_args", + binding="named", + params=[TaskHandlerParam(name="region_code", exact_name=True, value_schema=STRING)], + ) assert declarations["via_flat_map"] == TaskHandlerDeclaration( task_id="via_flat_map", - binding="named_or_whole", + binding="named", params=[ - TaskHandlerParam(name="Region", required=False, value_schema=STRING), - TaskHandlerParam(name="Count", required=False, value_schema=INT64), + TaskHandlerParam(name="Region", value_schema=STRING), + TaskHandlerParam(name="Count", value_schema=INT64), ], ) diff --git a/go-sdk/pkg/binding/declaration.go b/go-sdk/pkg/binding/declaration.go index e402095308dbe..66e73cb814727 100644 --- a/go-sdk/pkg/binding/declaration.go +++ b/go-sdk/pkg/binding/declaration.go @@ -29,31 +29,23 @@ import ( // Declare describes the parameters that the task function's TaskFlow arguments fill, so the Dag // processor can check a stub task against its Go handler before the task runs. // -// Flat parameters bind by position and carry no name. A lone struct binds by field name, -// exactly for an `arg:`-tagged field and ignoring case and underscores otherwise. An untagged -// lone struct can also take one unmatched argument as a whole value. +// Flat parameters bind by position and carry no name. A lone struct, tagged or not, binds by +// field name, exactly for an `arg:`-tagged field and ignoring case and underscores otherwise. func (p *Plan) Declare(taskID string) genmodels.TaskHandlerDeclaration { decl := genmodels.TaskHandlerDeclaration{ TaskID: taskID, Binding: genmodels.TaskHandlerDeclarationBindingPositional, - // A nil slice would go out as null, which the Dag processor rejects. - Params: []genmodels.TaskHandlerParam{}, } + // Not nil: a null list would tell the Dag processor that the params cannot be listed. + params := genmodels.TaskHandlerParams{} for _, plan := range p.params { switch plan.kind { case paramData: - decl.Params = append(decl.Params, genmodels.TaskHandlerParam{ - Required: true, - ValueSchema: valueSchema(plan.typ), - }) + params = append(params, genmodels.TaskHandlerParam{ValueSchema: valueSchema(plan.typ)}) case paramLoneStruct: - decl.Binding = genmodels.TaskHandlerDeclarationBindingNamedOrWhole - if plan.tagged { - decl.Binding = genmodels.TaskHandlerDeclarationBindingNamed - } - // An unfilled field keeps its zero value, so none is required. + decl.Binding = genmodels.TaskHandlerDeclarationBindingNamed for _, sf := range plan.fields { - decl.Params = append(decl.Params, genmodels.TaskHandlerParam{ + params = append(params, genmodels.TaskHandlerParam{ Name: sf.argName, ExactName: sf.tagged, ValueSchema: valueSchema(sf.fieldType), @@ -61,6 +53,7 @@ func (p *Plan) Declare(taskID string) genmodels.TaskHandlerDeclaration { } } } + decl.Params = ¶ms return decl } diff --git a/go-sdk/pkg/binding/declaration_test.go b/go-sdk/pkg/binding/declaration_test.go index c83547a4a4325..3e4ef2f693383 100644 --- a/go-sdk/pkg/binding/declaration_test.go +++ b/go-sdk/pkg/binding/declaration_test.go @@ -188,11 +188,11 @@ func TestDeclareFlatParams(t *testing.T) { assert.Equal(t, genmodels.TaskHandlerDeclaration{ TaskID: "task", Binding: genmodels.TaskHandlerDeclarationBindingPositional, - Params: []genmodels.TaskHandlerParam{ - {Required: true, ValueSchema: schemaOf(stringFragment)}, - {Required: true, ValueSchema: schemaOf(int64Fragment)}, - {Required: true, ValueSchema: schemaOf(nullOr(stringFragment))}, - {Required: true, ValueSchema: schemaOf(map[string]any{"type": "object"})}, + Params: &genmodels.TaskHandlerParams{ + {ValueSchema: schemaOf(stringFragment)}, + {ValueSchema: schemaOf(int64Fragment)}, + {ValueSchema: schemaOf(nullOr(stringFragment))}, + {ValueSchema: schemaOf(map[string]any{"type": "object"})}, }, }, decl) } @@ -203,7 +203,7 @@ func TestDeclareTaggedStruct(t *testing.T) { assert.Equal(t, genmodels.TaskHandlerDeclaration{ TaskID: "task", Binding: genmodels.TaskHandlerDeclarationBindingNamed, - Params: []genmodels.TaskHandlerParam{ + Params: &genmodels.TaskHandlerParams{ {Name: "region_code", ExactName: true, ValueSchema: schemaOf(stringFragment)}, {Name: "Threshold", ValueSchema: schemaOf(doubleFragment)}, {Name: "Count", ValueSchema: schemaOf(int64Fragment)}, @@ -214,8 +214,8 @@ func TestDeclareTaggedStruct(t *testing.T) { func TestDeclareUntaggedStruct(t *testing.T) { want := genmodels.TaskHandlerDeclaration{ TaskID: "task", - Binding: genmodels.TaskHandlerDeclarationBindingNamedOrWhole, - Params: []genmodels.TaskHandlerParam{ + Binding: genmodels.TaskHandlerDeclarationBindingNamed, + Params: &genmodels.TaskHandlerParams{ {Name: "RegionCode", ValueSchema: schemaOf(stringFragment)}, { Name: "Labels", @@ -241,6 +241,6 @@ func TestDeclareContextOnly(t *testing.T) { assert.Equal(t, genmodels.TaskHandlerDeclaration{ TaskID: "task", Binding: genmodels.TaskHandlerDeclarationBindingPositional, - Params: []genmodels.TaskHandlerParam{}, + Params: &genmodels.TaskHandlerParams{}, }, decl) } diff --git a/go-sdk/pkg/execution/integration_test.go b/go-sdk/pkg/execution/integration_test.go index 84fb4af0b6c9f..dd37b0347eefd 100644 --- a/go-sdk/pkg/execution/integration_test.go +++ b/go-sdk/pkg/execution/integration_test.go @@ -973,12 +973,10 @@ func TestServeTaskHandlerParseRequestEndToEnd(t *testing.T) { "params": []any{ map[string]any{ "name": nil, - "required": true, "value_schema": map[string]any{"type": "string"}, }, map[string]any{ - "name": nil, - "required": true, + "name": nil, "value_schema": map[string]any{"anyOf": []any{ map[string]any{"type": "integer", "format": "int64"}, map[string]any{"type": "null"}, @@ -992,7 +990,6 @@ func TestServeTaskHandlerParseRequestEndToEnd(t *testing.T) { "params": []any{map[string]any{ "name": "table_name", "exact_name": true, - "required": false, "value_schema": map[string]any{"type": "string"}, }}, }, diff --git a/go-sdk/pkg/execution/task_handler_parse_test.go b/go-sdk/pkg/execution/task_handler_parse_test.go index 6b6e7169f981e..47237b5a8e719 100644 --- a/go-sdk/pkg/execution/task_handler_parse_test.go +++ b/go-sdk/pkg/execution/task_handler_parse_test.go @@ -49,21 +49,21 @@ func TestDeclareTaskHandlers(t *testing.T) { { TaskID: "transform", Binding: genmodels.TaskHandlerDeclarationBindingPositional, - Params: []genmodels.TaskHandlerParam{ - {Required: true, ValueSchema: &genmodels.ArgValueSchema{"type": "string"}}, + Params: &genmodels.TaskHandlerParams{ + {ValueSchema: &genmodels.ArgValueSchema{"type": "string"}}, }, }, { TaskID: "extract", Binding: genmodels.TaskHandlerDeclarationBindingPositional, - Params: []genmodels.TaskHandlerParam{}, + Params: &genmodels.TaskHandlerParams{}, }, }, "report": { { TaskID: "render", Binding: genmodels.TaskHandlerDeclarationBindingPositional, - Params: []genmodels.TaskHandlerParam{}, + Params: &genmodels.TaskHandlerParams{}, }, }, }, From ccecc8115d5289120bd80f1616dda38b2a94dcf2 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 1 Oct 2026 17:40:10 +0000 Subject: [PATCH 64/94] Describe the two task handler bindings for Language SDK authors An SDK author picks a binding by what the Dag processor does with a mismatch, so the guide now says which mismatches fail the import under positional and which only warn under named, and that null params mean the SDK cannot list them. --- contributing-docs/30_new_language_sdk.rst | 30 ++++++++++++++--------- task-sdk/docs/lang-sdk-spec.rst | 7 +++--- 2 files changed, 23 insertions(+), 14 deletions(-) diff --git a/contributing-docs/30_new_language_sdk.rst b/contributing-docs/30_new_language_sdk.rst index 37be47e9f84c5..46ded9b149369 100644 --- a/contributing-docs/30_new_language_sdk.rst +++ b/contributing-docs/30_new_language_sdk.rst @@ -308,18 +308,26 @@ depend only on the artifact, never on the request. * ``task_handlers`` maps each Dag id the artifact registers handlers for to their declarations, in registration order. * Each ``TaskHandlerDeclaration`` states in ``binding`` how stub-task arguments - bind to its ``params``: ``positional``, ``named`` (by name in any order, - ignoring case and underscores unless a param sets ``exact_name``), - ``named_or_whole`` (as ``named``, except that when there is exactly one - argument and it matches no param, it is decoded as the whole value), or - ``named_open`` (as ``named``, but only the listed names are checked). + bind to its ``params``: + + * ``positional``: by position. An argument count or value type the handler + cannot take makes the Dag fail to import. + * ``named``: by name in any order, ignoring case and underscores unless a + param sets ``exact_name``. An argument or param that matches nothing is + logged as a warning and the task still runs, so the runtime must accept + both. When no param matches and exactly one argument was passed, it may be + the whole value and is not warned about, unless ``params`` is empty, a + param sets ``exact_name``, or the argument cannot be an object. A value + type a param does not accept makes the Dag fail to import. + +* ``params`` lists the handler's parameters in order, ``[]`` when it has none, + or is ``null`` when the SDK cannot list them. Then only the handler's + presence is checked. * Each ``TaskHandlerParam`` has a ``name`` (``null`` when the SDK has no name - for a positional parameter), ``required`` (the handler has no default for it), - and a ``value_schema``: the JSON Schema of the values it accepts, in the - vocabulary ``@task.stub`` uses for Python annotations, or ``null`` when the SDK - cannot state one. -* An empty ``params`` list or ``task_handlers`` mapping is sent empty, never as - ``null``. + for a positional parameter) and a ``value_schema``: the JSON Schema of the + values it accepts, in the vocabulary ``@task.stub`` uses for Python + annotations, or ``null`` when the SDK cannot state one. +* An empty ``task_handlers`` mapping is sent empty, never as ``null``. ``go-sdk/pkg/execution`` is a reference implementation. diff --git a/task-sdk/docs/lang-sdk-spec.rst b/task-sdk/docs/lang-sdk-spec.rst index 93b2c1d4296e6..fb99e68919706 100644 --- a/task-sdk/docs/lang-sdk-spec.rst +++ b/task-sdk/docs/lang-sdk-spec.rst @@ -181,6 +181,7 @@ The Dag processor starts the same process to check a Python Dag's ``@task.stub`` 5 exit No ``fn`` runs: the answer comes from the registrations alone. It must depend only on the artifact, -never on the request. A declaration states how the stub task's arguments bind to the handler's -parameters and which values each parameter accepts. An artifact that registers no ``TaskHandler`` -answers with an empty mapping. +never on the request. A declaration states whether the stub task's arguments bind to the handler's +parameters by position or by name, and which values each parameter accepts. An SDK that cannot list +a handler's parameters sends ``null`` for them, and then only the handler's presence is checked. An +artifact that registers no ``TaskHandler`` answers with an empty mapping. From d210b17c2935ba4ef856f93efc5b419a2ffa2816 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 16:00:44 +0000 Subject: [PATCH 65/94] TS SDK: answer TaskHandlerParseRequest from bundle.serve The Dag processor asks a Lang-SDK runtime which task handlers an artifact registers, and a TypeScript bundle exits on that first frame today. Types are erased, so a handler declares only its withArgNames renames, each an exact, optional name with no schema. Every other argument binds by folding, so the binding is named_open and the list is not exhaustive. The result is sent as a request and the runtime waits for the parent's acknowledgement before it exits, as the Go runtime does. --- ts-sdk/README.md | 4 +- ts-sdk/src/coordinator/comm-channel.ts | 7 +- ts-sdk/src/coordinator/protocol.ts | 4 + ts-sdk/src/coordinator/runtime.ts | 33 ++++- ts-sdk/src/coordinator/task-handler-parse.ts | 89 +++++++++++++ ts-sdk/src/sdk/bundle.ts | 9 ++ ts-sdk/tests/coordinator/integration.test.ts | 66 ++++++++-- ts-sdk/tests/coordinator/protocol.test.ts | 11 ++ .../coordinator/task-handler-parse.test.ts | 117 ++++++++++++++++++ 9 files changed, 327 insertions(+), 13 deletions(-) create mode 100644 ts-sdk/src/coordinator/task-handler-parse.ts create mode 100644 ts-sdk/tests/coordinator/task-handler-parse.test.ts diff --git a/ts-sdk/README.md b/ts-sdk/README.md index b105f768c6a41..2a7db7cef0391 100644 --- a/ts-sdk/README.md +++ b/ts-sdk/README.md @@ -287,7 +287,9 @@ Its values are Python names, which `tsc` cannot see and does not check. Airflow launches the bundled entrypoint with `--comm=host:port` and `--logs=host:port`. `bundle.serve()` connects to those sockets, receives the task startup message, finds the registered handler for the Dag/task pair, and -reports the terminal task state back to Airflow. +reports the terminal task state back to Airflow. Airflow can also start the bundle +to ask which task handlers it registers; `bundle.serve()` then answers with each +handler's task ID and `withArgNames` renames, without running any handler. See [`example/`](https://github.com/apache/airflow/tree/main/ts-sdk/example) for a coordinator-runtime example that packs a bundle with `airflow-ts-pack` and diff --git a/ts-sdk/src/coordinator/comm-channel.ts b/ts-sdk/src/coordinator/comm-channel.ts index 2c8eeb0a908b3..cc0f629fc4ea3 100644 --- a/ts-sdk/src/coordinator/comm-channel.ts +++ b/ts-sdk/src/coordinator/comm-channel.ts @@ -22,8 +22,8 @@ // // The channel is the sole reader on the socket. The task sends // requests and awaits id-correlated replies. The supervisor's only -// unprompted frame is the greeting (StartupDetails / -// DagFileParseRequest), pre-caught into the `greeting` promise that +// unprompted frame is the greeting (StartupDetails, DagFileParseRequest +// or TaskHandlerParseRequest), pre-caught into the `greeting` promise that // `connect()` awaits — the protocol sends nothing else // supervisor-initiated (comms.py: "No messages are sent to task // process except in response to a request"). @@ -35,7 +35,8 @@ import { Deferred } from "./deferred.js"; import type { LogChannel } from "./log-channel.js"; /** What `CommChannel.connect` resolves to: the live channel plus the - * supervisor's first frame (StartupDetails / DagFileParseRequest), + * supervisor's first frame (StartupDetails, DagFileParseRequest or + * TaskHandlerParseRequest), * already in hand so the caller never has to manage a "frame arrived * with no consumer" window. */ export interface CommConnection { diff --git a/ts-sdk/src/coordinator/protocol.ts b/ts-sdk/src/coordinator/protocol.ts index 323765036ea19..312b8ba34997e 100644 --- a/ts-sdk/src/coordinator/protocol.ts +++ b/ts-sdk/src/coordinator/protocol.ts @@ -32,6 +32,7 @@ import type { SucceedTask, TaskState, } from "../generated/supervisor.js"; +import type { TaskHandlerParseRequest, TaskHandlerParsingResult } from "../generated/supervisor.js"; export { SUPERVISOR_API_VERSION } from "../generated/supervisor.js"; @@ -75,6 +76,7 @@ type WithRequiredType = Omit & { export type MsgFromSupervisor = | WithRequiredType | WithRequiredType + | WithRequiredType | WithRequiredType; // -------- Frames to supervisor -------- @@ -86,6 +88,7 @@ export type RuntimeSucceedTask = WithRequiredType & { outlet_events: NonNullable; }; export type RuntimeDagFileParsingResult = WithRequiredType; +export type RuntimeTaskHandlerParsingResult = WithRequiredType; // -------- Decoder: raw map → typed message -------- @@ -94,6 +97,7 @@ export function asMsgFromSupervisor(raw: unknown): MsgFromSupervisor { switch (body.type) { case "StartupDetails": case "DagFileParseRequest": + case "TaskHandlerParseRequest": case "ErrorResponse": return body as unknown as MsgFromSupervisor; default: diff --git a/ts-sdk/src/coordinator/runtime.ts b/ts-sdk/src/coordinator/runtime.ts index bf9764f3e6618..f80d9bb816e41 100644 --- a/ts-sdk/src/coordinator/runtime.ts +++ b/ts-sdk/src/coordinator/runtime.ts @@ -33,6 +33,7 @@ // 1. Parse --comm / --logs from argv // 2. Connect both TCP sockets // 3. Read the first frame from comm: +// - TaskHandlerParseRequest → send TaskHandlerParsingResult, await its ack, exit // - DagFileParseRequest → respond with DagFileParsingResult, exit // - StartupDetails → run task, respond Succeed or Fail, exit // @@ -40,6 +41,7 @@ import { resolveArgs, type BoundArgs } from "./arg-binding.js"; import { createCoordinatorClient } from "./client.js"; import { CommChannel } from "./comm-channel.js"; import { LogChannel } from "./log-channel.js"; +import { declareTaskHandlers } from "./task-handler-parse.js"; import { AIRFLOW_METADATA_FLAG, AIRFLOW_METADATA_SENTINEL, @@ -165,7 +167,17 @@ export async function startCoordinator( const body = asMsgFromSupervisor(firstFrame.body); - if (body.type === "DagFileParseRequest") { + if (body.type === "TaskHandlerParseRequest") { + runtimeLogs.info("Received task handler parse request", { + file: body.file, + dag_ids: body.dag_ids, + }); + const result = declareTaskHandlers(bundle, body); + await sendParseResult(result, comm, runtimeLogs); + runtimeLogs.info("Declared task handlers", { + dag_ids: Object.keys(result.task_handlers), + }); + } else if (body.type === "DagFileParseRequest") { runtimeLogs.info("Received Dag parse request", { file: body.file, bundle_path: body.bundle_path, @@ -194,7 +206,7 @@ export async function startCoordinator( runtimeLogs.info("Task succeeded", { task_id: body.ti.task_id }); } } else { - const errMsg = `First frame must be DagFileParseRequest or StartupDetails, got ${body.type}`; + const errMsg = `First frame must be TaskHandlerParseRequest, DagFileParseRequest or StartupDetails, got ${body.type}`; runtimeLogs.error("Unexpected first frame", { type: body.type }); await sendSupervisorResponse(firstFrame.id, null, comm, runtimeLogs, { error: "protocol_error", @@ -382,6 +394,23 @@ async function sendSupervisorResponse( } } +/** + * Send a parse result as a request and wait for Airflow's acknowledgement, so + * the process never exits before the result is recorded. An acknowledgement + * carrying an error, including the socket closing first, throws. + */ +async function sendParseResult(body: unknown, comm: CommChannel, logs: LogChannel): Promise { + const ack = await comm.request(body, { timeoutMs: COORDINATOR_RESPONSE_TIMEOUT_MS }); + if (ack.error == null) return; + logs.error("Parse result was not acknowledged", { + response_type: responseType(body), + error: ack.error, + }); + throw new Error( + `Airflow did not acknowledge the ${responseType(body)}: ${JSON.stringify(ack.error)}`, + ); +} + function buildFailureResponse( details: StartupDetails, message: string, diff --git a/ts-sdk/src/coordinator/task-handler-parse.ts b/ts-sdk/src/coordinator/task-handler-parse.ts new file mode 100644 index 0000000000000..dd6c0977e3ede --- /dev/null +++ b/ts-sdk/src/coordinator/task-handler-parse.ts @@ -0,0 +1,89 @@ +/*! + * 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. + */ + +// The answer to a TaskHandlerParseRequest: the task handlers this bundle +// registers for the Dags a Python file declares, as the Dag processor checks +// them against that file's `@task.stub` calls. + +import type { TaskHandlerDeclaration, TaskHandlerParseRequest } from "../generated/supervisor.js"; +import type { RuntimeTaskHandlerParsingResult } from "./protocol.js"; +import { getArgNames } from "../sdk/arg-names.js"; +import { getBundleTaskHandlers, type Bundle } from "../sdk/bundle.js"; +import type { TaskFunction } from "../sdk/task.js"; + +/** + * Declare the task handlers registered for the requested Dags, each Dag's in + * registration order. A requested Dag with no handler is left out rather than + * mapped to an empty list, and a Dag declared in TypeScript is never declared. + * No handler runs. + * + * @throws when the request's `dag_ids` is not a list of strings. + */ +export function declareTaskHandlers( + bundle: Bundle, + request: Pick, +): RuntimeTaskHandlerParsingResult { + const requested = new Set(readDagIds(request)); + const declared: [string, TaskHandlerDeclaration[]][] = []; + for (const [dagId, handlers] of getBundleTaskHandlers(bundle)) { + if (!requested.has(dagId)) continue; + declared.push([ + dagId, + [...handlers].map(([taskId, handler]) => declareTaskHandler(taskId, handler)), + ]); + } + return { + type: "TaskHandlerParsingResult", + fileloc: request.file, + // Defined rather than assigned, so a Dag named `__proto__` stays a key. + task_handlers: Object.fromEntries(declared), + }; +} + +/** + * Types are erased, so a handler's `withArgNames` renames are the only names + * the SDK knows. Every other argument reaches the handler by folding, which is + * why the list is open. + */ +function declareTaskHandler(taskId: string, handler: TaskFunction): TaskHandlerDeclaration { + const names = new Set(getArgNames(handler).values()); + return { + task_id: taskId, + binding: "named_open", + params: [...names].map((name) => ({ + name, + value_schema: null, + // A name the call did not pass reads as undefined, and the task still runs. + required: false, + // A rename never falls back to folding. + exact_name: true, + })), + }; +} + +function readDagIds(request: Pick): string[] { + // Checked, or a request this SDK cannot read would be answered with no handlers at all. + const dagIds: unknown = request.dag_ids; + if (!Array.isArray(dagIds) || !dagIds.every((dagId) => typeof dagId === "string")) { + throw new Error( + `TaskHandlerParseRequest.dag_ids must be a list of strings, got ${JSON.stringify(dagIds)}`, + ); + } + return dagIds as string[]; +} diff --git a/ts-sdk/src/sdk/bundle.ts b/ts-sdk/src/sdk/bundle.ts index a3e54f1a23172..31b9863e4376d 100644 --- a/ts-sdk/src/sdk/bundle.ts +++ b/ts-sdk/src/sdk/bundle.ts @@ -226,6 +226,15 @@ export function bundleDags(bundle: Bundle): ReadonlyMap { return dagsOf(bundle); } +/** Internal: the handlers this bundle registered for tasks a Python Dag + * declares, keyed by dag_id and then task_id, both in registration order. + * Excludes the Dags declared in TypeScript. */ +export function getBundleTaskHandlers( + bundle: Bundle, +): ReadonlyMap> { + return taskHandlersOf(bundle); +} + /** Internal: the task IDs this bundle can dispatch, per Dag: the Dags declared * in TypeScript first, then the Dags its task handlers name, each in * registration order. A Dag declared in TypeScript with no tasks is included. */ diff --git a/ts-sdk/tests/coordinator/integration.test.ts b/ts-sdk/tests/coordinator/integration.test.ts index d0b6d53f83fff..35269758ff42a 100644 --- a/ts-sdk/tests/coordinator/integration.test.ts +++ b/ts-sdk/tests/coordinator/integration.test.ts @@ -81,8 +81,8 @@ type Responder = ( body: Record, ) => { body: unknown; error?: unknown } | null; -function frameBytes(id: number, body: unknown, isResponse: boolean): Buffer { - const arr = isResponse ? [id, body, null] : [id, body]; +function frameBytes(id: number, body: unknown, isResponse: boolean, error: unknown = null): Buffer { + const arr = isResponse ? [id, body, error] : [id, body]; const payload = Buffer.from(encode(arr)); const header = Buffer.alloc(4); header.writeUInt32BE(payload.length, 0); @@ -203,7 +203,7 @@ async function driveSupervisor(initialFrame: unknown, responder?: Responder): Pr // runtime never hangs on auto-generated RPCs (e.g. the // return_value XCom push). { body: null }; - commSock.write(frameBytes(f.id, reply.body, true)); + commSock.write(frameBytes(f.id, reply.body, true, reply.error)); continue; } // Arity-3: a response from the runtime. The terminal frame @@ -218,10 +218,12 @@ async function driveSupervisor(initialFrame: unknown, responder?: Responder): Pr // Wait for the runtime to finish AND for the comm socket to deliver // its final bytes (FIN signals that all preceding frames are flushed). const commSocketEnded = new Promise((resolve) => commSock.on("end", () => resolve())); - await Promise.all([runtimeDone, commSocketEnded]); - - comm.server.close(); - logs.server.close(); + try { + await Promise.all([runtimeDone, commSocketEnded]); + } finally { + comm.server.close(); + logs.server.close(); + } const lines = Buffer.concat(logChunks).toString("utf8").split("\n").filter(Boolean); const logRecords = lines.map((l) => JSON.parse(l) as Record); @@ -611,6 +613,56 @@ describe("coordinator runtime integration", () => { ).toBe(true); }); + describe("TaskHandlerParseRequest", () => { + const parseRequest = { + type: "TaskHandlerParseRequest", + file: "/bundles/etl.min.mjs", + dag_ids: ["py_dag"], + bundle_path: "/bundles", + bundle_name: "ts-task-handlers", + }; + + it("sends the declarations as a request once, and runs no handler", async () => { + let called = false; + bundle.register( + new TaskHandler("py_dag", "transform", async () => { + called = true; + }), + ); + + const result = await driveSupervisor(parseRequest); + + expect(result.firstResponse).toBeNull(); + expect(result.runtimeRequests).toEqual([ + { + type: "TaskHandlerParsingResult", + body: { + type: "TaskHandlerParsingResult", + fileloc: "/bundles/etl.min.mjs", + task_handlers: { + py_dag: [{ task_id: "transform", binding: "named_open", params: [] }], + }, + }, + }, + ]); + expect(called).toBe(false); + }); + + it("fails when Airflow answers the result with an error", async () => { + bundle.register(new TaskHandler("py_dag", "transform", async () => undefined)); + + await expect( + driveSupervisor(parseRequest, () => ({ + body: null, + error: { + error: "generic_error", + detail: { message: "A parse result was already received" }, + }, + })), + ).rejects.toThrow(/did not acknowledge the TaskHandlerParsingResult/); + }); + }); + it("aborts the context signal on SIGTERM and reports a thrown task error", async () => { let sawAbort = false; testDag.task("aborted_then_failed", async () => { diff --git a/ts-sdk/tests/coordinator/protocol.test.ts b/ts-sdk/tests/coordinator/protocol.test.ts index 12ba947eeb3a2..e799f10b073fe 100644 --- a/ts-sdk/tests/coordinator/protocol.test.ts +++ b/ts-sdk/tests/coordinator/protocol.test.ts @@ -54,6 +54,17 @@ describe("protocol decode", () => { expect(asMsgFromSupervisor(raw).type).toBe("DagFileParseRequest"); }); + it("accepts TaskHandlerParseRequest", () => { + const raw = { + type: "TaskHandlerParseRequest", + file: "/x.min.mjs", + dag_ids: ["d"], + bundle_path: "/", + bundle_name: "ts-task-handlers", + }; + expect(asMsgFromSupervisor(raw).type).toBe("TaskHandlerParseRequest"); + }); + it("rejects missing type", () => { expect(() => asMsgFromSupervisor({ ti: {} })).toThrow(/missing string 'type'/); }); diff --git a/ts-sdk/tests/coordinator/task-handler-parse.test.ts b/ts-sdk/tests/coordinator/task-handler-parse.test.ts new file mode 100644 index 0000000000000..f25bc1728029e --- /dev/null +++ b/ts-sdk/tests/coordinator/task-handler-parse.test.ts @@ -0,0 +1,117 @@ +/*! + * 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. + */ + +import { describe, expect, it } from "vitest"; + +import { declareTaskHandlers } from "../../src/coordinator/task-handler-parse.js"; +import { withArgNames } from "../../src/sdk/arg-names.js"; +import { Bundle } from "../../src/sdk/bundle.js"; +import { Dag } from "../../src/sdk/dag.js"; +import { TaskHandler } from "../../src/sdk/task-handler.js"; + +const FILE = "/bundles/etl.min.mjs"; + +async function noop(): Promise {} + +function declare(bundle: Bundle, dagIds: unknown) { + return declareTaskHandlers(bundle, { file: FILE, dag_ids: dagIds as string[] }); +} + +describe("declareTaskHandlers", () => { + it("declares each requested Dag's handlers in registration order", () => { + const bundle = new Bundle( + new TaskHandler("etl", "extract", noop), + new TaskHandler("reports", "send", noop), + new TaskHandler("etl", "load", noop), + ); + + const result = declare(bundle, ["reports", "etl"]); + + expect(result).toStrictEqual({ + type: "TaskHandlerParsingResult", + fileloc: FILE, + task_handlers: { + etl: [ + { task_id: "extract", binding: "named_open", params: [] }, + { task_id: "load", binding: "named_open", params: [] }, + ], + reports: [{ task_id: "send", binding: "named_open", params: [] }], + }, + }); + expect(Object.keys(result.task_handlers)).toEqual(["etl", "reports"]); + }); + + it("declares each distinct withArgNames target as an exact, optional param", () => { + interface ReportArgs { + label: string; + title: string; + owner: string; + } + const report = withArgNames( + { label: "run_label", title: "run_label", owner: "team" }, + async ({ label }: ReportArgs) => label, + ); + + const result = declare(new Bundle(new TaskHandler("etl", "report", report)), ["etl"]); + + expect(result.task_handlers).toEqual({ + etl: [ + { + task_id: "report", + binding: "named_open", + params: [ + { name: "run_label", value_schema: null, required: false, exact_name: true }, + { name: "team", value_schema: null, required: false, exact_name: true }, + ], + }, + ], + }); + }); + + it("leaves out unrequested Dags and requested Dags with no handler", () => { + const bundle = new Bundle( + new TaskHandler("etl", "extract", noop), + new TaskHandler("reports", "send", noop), + ); + + expect(Object.keys(declare(bundle, ["etl", "missing"]).task_handlers)).toEqual(["etl"]); + expect(declare(bundle, ["missing"]).task_handlers).toEqual({}); + }); + + it("does not declare a Dag declared in TypeScript", () => { + const native = new Dag("native"); + native.task("run", noop)(); + const bundle = new Bundle(native, new TaskHandler("etl", "extract", noop)); + + expect(Object.keys(declare(bundle, ["native", "etl"]).task_handlers)).toEqual(["etl"]); + }); + + it("keeps a Dag named __proto__ as a key", () => { + const result = declare(new Bundle(new TaskHandler("__proto__", "run", noop)), ["__proto__"]); + + expect(Object.keys(result.task_handlers)).toEqual(["__proto__"]); + expect(Object.getPrototypeOf(result.task_handlers)).toBe(Object.prototype); + }); + + it.each([[undefined], ["etl"], [["etl", 1]]])("rejects dag_ids of %j", (dagIds) => { + const bundle = new Bundle(new TaskHandler("etl", "extract", noop)); + + expect(() => declare(bundle, dagIds)).toThrow(/dag_ids must be a list of strings/); + }); +}); From 8ea0048ce6862fd3965817c3f9b2ac1cbba2b32e Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 16:00:57 +0000 Subject: [PATCH 66/94] Let the Node coordinator probe a bundle's task handlers SDKTaskHandlerProcessorProcess needs a command per coordinator to start the runtime that reports an artifact's task handlers. Unlike the execution command, the Node one does not look up any Dag, since artifacts will stop listing them. The real probe packs ts-sdk/example and asks it through the transport. It needs Node.js 22 and pnpm and installs the SDK's toolchain, so it runs only when AIRFLOW_LANG_SDK_REAL_PROBE_TESTS=1, the switch the Go and Java probes share. --- .../test_task_handler_processor_ts.py | 149 ++++++++++++++++++ .../sdk/coordinators/node/coordinator.py | 6 + .../coordinators/node/test_coordinator.py | 38 +++++ 3 files changed, 193 insertions(+) create mode 100644 airflow-core/tests/unit/dag_processing/test_task_handler_processor_ts.py diff --git a/airflow-core/tests/unit/dag_processing/test_task_handler_processor_ts.py b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_ts.py new file mode 100644 index 0000000000000..1bb7a3863a019 --- /dev/null +++ b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_ts.py @@ -0,0 +1,149 @@ +# +# 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. +"""Probe a bundle packed by the TypeScript SDK through the real Node runtime.""" + +from __future__ import annotations + +import json +import os +import shutil +import subprocess +from typing import TYPE_CHECKING +from unittest import mock + +import pytest +import structlog + +from airflow.dag_processing.processor import ( + TaskHandlerDeclaration, + TaskHandlerParam, + TaskHandlerParsingResult, +) +from airflow.dag_processing.task_handler_processor import LangSDKTaskHandlerProcessorProcess +from airflow.sdk.execution_time import supervisor +from airflow.sdk.execution_time.coordinator import reset_coordinator_manager + +from tests_common.pytest_plugin import AIRFLOW_ROOT_PATH +from tests_common.test_utils.config import conf_vars + +if TYPE_CHECKING: + from collections.abc import Iterator + from pathlib import Path + +TS_SDK_PATH = AIRFLOW_ROOT_PATH / "ts-sdk" +# The SDK's `engines` requirement. +MIN_NODE_MAJOR = 22 + + +def _get_toolchain_problem() -> str | None: + """Return why this machine cannot pack the example bundle, or ``None`` when it can.""" + if not TS_SDK_PATH.is_dir(): + return "the TypeScript SDK sources are absent" + node = shutil.which("node") + if node is None or shutil.which("pnpm") is None: + return "needs node and pnpm" + version = subprocess.run([node, "--version"], capture_output=True, text=True, check=True).stdout + if int(version.strip().lstrip("v").split(".")[0]) < MIN_NODE_MAJOR: + return f"needs Node.js {MIN_NODE_MAJOR} or later, found {version.strip()}" + return None + + +pytestmark = pytest.mark.skipif( + os.environ.get("AIRFLOW_LANG_SDK_REAL_PROBE_TESTS") != "1", + reason="set AIRFLOW_LANG_SDK_REAL_PROBE_TESTS=1 to probe real Lang-SDK artifacts", +) + + +@pytest.fixture(scope="module") +def example_bundle(tmp_path_factory) -> Path: + """Pack ``ts-sdk/example`` from a copy of the SDK sources, so the checkout gets no build output.""" + if (problem := _get_toolchain_problem()) is not None: + pytest.skip(problem) + sdk = tmp_path_factory.mktemp("ts-sdk") / "ts-sdk" + shutil.copytree(TS_SDK_PATH, sdk, ignore=shutil.ignore_patterns("node_modules", "dist", ".pnpm-store")) + env = {**os.environ, "CI": "true", "COREPACK_ENABLE_DOWNLOAD_PROMPT": "0"} + # The example links the SDK, so it is installed again once the SDK is built. + for cwd, args in [ + (sdk, ["install", "--frozen-lockfile"]), + (sdk, ["run", "build"]), + (sdk / "example", ["install"]), + (sdk / "example", ["run", "build"]), + ]: + subprocess.run(["pnpm", *args], cwd=cwd, env=env, check=True) + return sdk / "example" / "dist" / "bundle.min.mjs" + + +@pytest.fixture +def node_coordinator() -> Iterator[None]: + spec = { + "ts": { + "classpath": "airflow.sdk.coordinators.node.NodeCoordinator", + "kwargs": {"node_executable": shutil.which("node")}, + } + } + reset_coordinator_manager() + try: + # The parse child is a bare fork even on macOS, so it sees this config. + with ( + conf_vars({("sdk", "coordinators"): json.dumps(spec)}), + mock.patch.object(supervisor, "_should_use_exec", return_value=False), + ): + yield + finally: + reset_coordinator_manager() + + +def _declare(task_id: str, *renames: str) -> TaskHandlerDeclaration: + return TaskHandlerDeclaration( + task_id=task_id, + binding="named_open", + params=[ + TaskHandlerParam(name=name, value_schema=None, required=False, exact_name=True) + for name in renames + ], + ) + + +@pytest.mark.usefixtures("node_coordinator") +def test_a_packed_bundle_declares_its_task_handlers(example_bundle): + result = LangSDKTaskHandlerProcessorProcess.run( + coordinator="ts", + path=example_bundle, + bundle_path=example_bundle.parent, + bundle_name="ts-task-handlers", + artifact_rel_path=example_bundle.name, + dag_ids=["typescript_taskflow_example", "typescript_example", "not_registered"], + logger=structlog.get_logger(), + ) + + assert result == TaskHandlerParsingResult( + fileloc=os.fspath(example_bundle), + task_handlers={ + "typescript_example": [ + _declare("build_message"), + _declare("read_connection"), + _declare("write_and_delete_variable"), + ], + "typescript_taskflow_example": [ + _declare("summarize"), + _declare("report", "run_label"), + _declare("build_message"), + ], + }, + ) + assert list(result.task_handlers) == ["typescript_example", "typescript_taskflow_example"] diff --git a/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py index 27fc142a3036c..1644cf425082a 100644 --- a/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py @@ -95,6 +95,8 @@ class NodeCoordinator(SubprocessCoordinator): """ Coordinator that launches a Node.js subprocess for task execution. + It also launches a verified ``*.min.mjs`` bundle to report the task handlers it registers. + Configuration is taken from the ``[sdk] coordinators`` entry that constructs this instance:: @@ -123,3 +125,7 @@ def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], roots = self._get_scan_roots() bundle = _Bundle.find(roots, what.dag_id) return [self.node_executable, os.fspath(bundle.path)], bundle.schema_version + + def _build_parse_task_handler_command(self, *, path: pathlib.Path) -> tuple[list[str], str | None]: + metadata = read_bundle(path) + return [self.node_executable, os.fspath(path)], metadata.supervisor_schema_version 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 3dc01be33f297..5c5af582be9ba 100644 --- a/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py @@ -106,6 +106,44 @@ def test_build_execute_task_command_returns_node_bundle_and_schema_version(self, assert schema_version == SCHEMA_VERSION +def _write_plain_file(root: pathlib.Path) -> pathlib.Path: + path = root / BUNDLE_NAME + path.write_bytes(b"export {};\n") + return path + + +def _write_corrupted_bundle(root: pathlib.Path) -> pathlib.Path: + bundle = write_bundle(root, "test_dag") + # The last byte before the trailing newline is in the code region. + mutate_byte(bundle, len(bundle.read_bytes()) - 2) + return bundle + + +class TestNodeCoordinatorParseTaskHandlerCommand: + def test_returns_node_and_the_bundle_schema_version(self, tmp_path): + # A bundle registering no Dag: the command does not depend on the Dags it declares. + bundle = write_bundle(tmp_path) + coordinator = NodeCoordinator(node_executable="/opt/node/bin/node") + + command, schema_version = coordinator._build_parse_task_handler_command(path=bundle) + + assert command == ["/opt/node/bin/node", str(bundle)] + assert schema_version == SCHEMA_VERSION + + @pytest.mark.parametrize( + ("write", "error"), + [ + pytest.param(_write_plain_file, "has no airflow bundle layout", id="not-a-bundle"), + pytest.param(_write_corrupted_bundle, "code SHA-256 mismatch", id="corrupted-code"), + ], + ) + def test_rejects_a_bundle_that_fails_verification(self, tmp_path, write, error): + path = write(tmp_path) + + with pytest.raises(ValueError, match=error): + NodeCoordinator()._build_parse_task_handler_command(path=path) + + class TestBundleFind: @pytest.mark.parametrize( "name", From e12ad5c2bb2d7a0b7a4f481dd15e83bd70473840 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 16:21:48 +0000 Subject: [PATCH 67/94] TS SDK: name the task handler lookup like its siblings The other bundle accessors read as nouns (bundleDags, bundleDagTaskIds), and the parse reply now sits beside the response helper the other first frames use, so a reader needs to see that one sends a request and waits while the other only answers. --- ts-sdk/src/coordinator/runtime.ts | 3 +++ ts-sdk/src/coordinator/task-handler-parse.ts | 4 ++-- ts-sdk/src/sdk/bundle.ts | 2 +- 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/ts-sdk/src/coordinator/runtime.ts b/ts-sdk/src/coordinator/runtime.ts index f80d9bb816e41..3d402c7b9ba27 100644 --- a/ts-sdk/src/coordinator/runtime.ts +++ b/ts-sdk/src/coordinator/runtime.ts @@ -398,6 +398,9 @@ async function sendSupervisorResponse( * Send a parse result as a request and wait for Airflow's acknowledgement, so * the process never exits before the result is recorded. An acknowledgement * carrying an error, including the socket closing first, throws. + * + * Unlike {@link sendSupervisorResponse}, which answers the first frame's id + * and expects no reply, this sends a new request and awaits its answer. */ async function sendParseResult(body: unknown, comm: CommChannel, logs: LogChannel): Promise { const ack = await comm.request(body, { timeoutMs: COORDINATOR_RESPONSE_TIMEOUT_MS }); diff --git a/ts-sdk/src/coordinator/task-handler-parse.ts b/ts-sdk/src/coordinator/task-handler-parse.ts index dd6c0977e3ede..111e0fe0c9ef8 100644 --- a/ts-sdk/src/coordinator/task-handler-parse.ts +++ b/ts-sdk/src/coordinator/task-handler-parse.ts @@ -24,7 +24,7 @@ import type { TaskHandlerDeclaration, TaskHandlerParseRequest } from "../generated/supervisor.js"; import type { RuntimeTaskHandlerParsingResult } from "./protocol.js"; import { getArgNames } from "../sdk/arg-names.js"; -import { getBundleTaskHandlers, type Bundle } from "../sdk/bundle.js"; +import { bundleTaskHandlers, type Bundle } from "../sdk/bundle.js"; import type { TaskFunction } from "../sdk/task.js"; /** @@ -41,7 +41,7 @@ export function declareTaskHandlers( ): RuntimeTaskHandlerParsingResult { const requested = new Set(readDagIds(request)); const declared: [string, TaskHandlerDeclaration[]][] = []; - for (const [dagId, handlers] of getBundleTaskHandlers(bundle)) { + for (const [dagId, handlers] of bundleTaskHandlers(bundle)) { if (!requested.has(dagId)) continue; declared.push([ dagId, diff --git a/ts-sdk/src/sdk/bundle.ts b/ts-sdk/src/sdk/bundle.ts index 31b9863e4376d..4ae41cfe16ed4 100644 --- a/ts-sdk/src/sdk/bundle.ts +++ b/ts-sdk/src/sdk/bundle.ts @@ -229,7 +229,7 @@ export function bundleDags(bundle: Bundle): ReadonlyMap { /** Internal: the handlers this bundle registered for tasks a Python Dag * declares, keyed by dag_id and then task_id, both in registration order. * Excludes the Dags declared in TypeScript. */ -export function getBundleTaskHandlers( +export function bundleTaskHandlers( bundle: Bundle, ): ReadonlyMap> { return taskHandlersOf(bundle); From 57d85ff451561d0e3e305fabb37788be95545b6b Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 16:22:01 +0000 Subject: [PATCH 68/94] Skip the TS real probe when node cannot report a version A node that is on PATH but fails or prints something unexpected made the opt-in probe error instead of saying why it cannot run. --- .../dag_processing/test_task_handler_processor_ts.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/airflow-core/tests/unit/dag_processing/test_task_handler_processor_ts.py b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_ts.py index 1bb7a3863a019..84ea00900c2fe 100644 --- a/airflow-core/tests/unit/dag_processing/test_task_handler_processor_ts.py +++ b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_ts.py @@ -57,9 +57,15 @@ def _get_toolchain_problem() -> str | None: node = shutil.which("node") if node is None or shutil.which("pnpm") is None: return "needs node and pnpm" - version = subprocess.run([node, "--version"], capture_output=True, text=True, check=True).stdout - if int(version.strip().lstrip("v").split(".")[0]) < MIN_NODE_MAJOR: - return f"needs Node.js {MIN_NODE_MAJOR} or later, found {version.strip()}" + try: + version = subprocess.run( + [node, "--version"], capture_output=True, text=True, check=True + ).stdout.strip() + major = int(version.lstrip("v").split(".")[0]) + except (OSError, subprocess.CalledProcessError, ValueError) as e: + return f"cannot read the Node.js version: {e}" + if major < MIN_NODE_MAJOR: + return f"needs Node.js {MIN_NODE_MAJOR} or later, found {version}" return None From 6d0064278c2e67de89133a44ad5f43587077e0d7 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 1 Oct 2026 06:01:46 +0000 Subject: [PATCH 69/94] Patch the TS real probe test with decorators An unspecced _should_use_exec mock accepts any call, and patches hidden in a fixture are harder to see than decorators on the one test that needs them. --- .../test_task_handler_processor_ts.py | 31 ++++++++----------- 1 file changed, 13 insertions(+), 18 deletions(-) diff --git a/airflow-core/tests/unit/dag_processing/test_task_handler_processor_ts.py b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_ts.py index 84ea00900c2fe..a93e639f38081 100644 --- a/airflow-core/tests/unit/dag_processing/test_task_handler_processor_ts.py +++ b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_ts.py @@ -48,6 +48,12 @@ TS_SDK_PATH = AIRFLOW_ROOT_PATH / "ts-sdk" # The SDK's `engines` requirement. MIN_NODE_MAJOR = 22 +COORDINATORS = { + "ts": { + "classpath": "airflow.sdk.coordinators.node.NodeCoordinator", + "kwargs": {"node_executable": shutil.which("node")}, + } +} def _get_toolchain_problem() -> str | None: @@ -95,23 +101,10 @@ def example_bundle(tmp_path_factory) -> Path: @pytest.fixture -def node_coordinator() -> Iterator[None]: - spec = { - "ts": { - "classpath": "airflow.sdk.coordinators.node.NodeCoordinator", - "kwargs": {"node_executable": shutil.which("node")}, - } - } +def fresh_coordinator_manager() -> Iterator[None]: + reset_coordinator_manager() + yield reset_coordinator_manager() - try: - # The parse child is a bare fork even on macOS, so it sees this config. - with ( - conf_vars({("sdk", "coordinators"): json.dumps(spec)}), - mock.patch.object(supervisor, "_should_use_exec", return_value=False), - ): - yield - finally: - reset_coordinator_manager() def _declare(task_id: str, *renames: str) -> TaskHandlerDeclaration: @@ -125,8 +118,10 @@ def _declare(task_id: str, *renames: str) -> TaskHandlerDeclaration: ) -@pytest.mark.usefixtures("node_coordinator") -def test_a_packed_bundle_declares_its_task_handlers(example_bundle): +@pytest.mark.usefixtures("fresh_coordinator_manager") +@conf_vars({("sdk", "coordinators"): json.dumps(COORDINATORS)}) +@mock.patch.object(supervisor, "_should_use_exec", autospec=True, return_value=False) +def test_a_packed_bundle_declares_its_task_handlers(mock_should_use_exec, example_bundle): result = LangSDKTaskHandlerProcessorProcess.run( coordinator="ts", path=example_bundle, From f8acbd27defc9b2a800e47cd29c271685b1077ed Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 1 Oct 2026 07:36:55 +0000 Subject: [PATCH 70/94] TS SDK: answer a task handler parse with every handler The parse request no longer names Dags, and the Dag processor reuses an artifact's answer for every Dag file that uses the artifact, so the answer must not depend on what was asked. --- .../test_task_handler_processor_ts.py | 1 - ts-sdk/src/coordinator/runtime.ts | 5 +-- ts-sdk/src/coordinator/task-handler-parse.ts | 30 ++++------------ ts-sdk/tests/coordinator/integration.test.ts | 1 - ts-sdk/tests/coordinator/protocol.test.ts | 1 - .../coordinator/task-handler-parse.test.ts | 34 ++++++------------- 6 files changed, 18 insertions(+), 54 deletions(-) diff --git a/airflow-core/tests/unit/dag_processing/test_task_handler_processor_ts.py b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_ts.py index a93e639f38081..94541a2573729 100644 --- a/airflow-core/tests/unit/dag_processing/test_task_handler_processor_ts.py +++ b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_ts.py @@ -128,7 +128,6 @@ def test_a_packed_bundle_declares_its_task_handlers(mock_should_use_exec, exampl bundle_path=example_bundle.parent, bundle_name="ts-task-handlers", artifact_rel_path=example_bundle.name, - dag_ids=["typescript_taskflow_example", "typescript_example", "not_registered"], logger=structlog.get_logger(), ) diff --git a/ts-sdk/src/coordinator/runtime.ts b/ts-sdk/src/coordinator/runtime.ts index 3d402c7b9ba27..455cf3ab19b21 100644 --- a/ts-sdk/src/coordinator/runtime.ts +++ b/ts-sdk/src/coordinator/runtime.ts @@ -168,10 +168,7 @@ export async function startCoordinator( const body = asMsgFromSupervisor(firstFrame.body); if (body.type === "TaskHandlerParseRequest") { - runtimeLogs.info("Received task handler parse request", { - file: body.file, - dag_ids: body.dag_ids, - }); + runtimeLogs.info("Received task handler parse request", { file: body.file }); const result = declareTaskHandlers(bundle, body); await sendParseResult(result, comm, runtimeLogs); runtimeLogs.info("Declared task handlers", { diff --git a/ts-sdk/src/coordinator/task-handler-parse.ts b/ts-sdk/src/coordinator/task-handler-parse.ts index 111e0fe0c9ef8..3bc9285f183e0 100644 --- a/ts-sdk/src/coordinator/task-handler-parse.ts +++ b/ts-sdk/src/coordinator/task-handler-parse.ts @@ -17,9 +17,9 @@ * under the License. */ -// The answer to a TaskHandlerParseRequest: the task handlers this bundle -// registers for the Dags a Python file declares, as the Dag processor checks -// them against that file's `@task.stub` calls. +// The answer to a TaskHandlerParseRequest: every task handler this bundle +// registers, as the Dag processor checks them against a Python file's +// `@task.stub` calls. import type { TaskHandlerDeclaration, TaskHandlerParseRequest } from "../generated/supervisor.js"; import type { RuntimeTaskHandlerParsingResult } from "./protocol.js"; @@ -28,21 +28,16 @@ import { bundleTaskHandlers, type Bundle } from "../sdk/bundle.js"; import type { TaskFunction } from "../sdk/task.js"; /** - * Declare the task handlers registered for the requested Dags, each Dag's in - * registration order. A requested Dag with no handler is left out rather than - * mapped to an empty list, and a Dag declared in TypeScript is never declared. - * No handler runs. - * - * @throws when the request's `dag_ids` is not a list of strings. + * Declare every registered task handler, keyed by Dag id, each Dag's in + * registration order. A Dag declared in TypeScript is never declared. No + * handler runs. */ export function declareTaskHandlers( bundle: Bundle, - request: Pick, + request: Pick, ): RuntimeTaskHandlerParsingResult { - const requested = new Set(readDagIds(request)); const declared: [string, TaskHandlerDeclaration[]][] = []; for (const [dagId, handlers] of bundleTaskHandlers(bundle)) { - if (!requested.has(dagId)) continue; declared.push([ dagId, [...handlers].map(([taskId, handler]) => declareTaskHandler(taskId, handler)), @@ -76,14 +71,3 @@ function declareTaskHandler(taskId: string, handler: TaskFunction): TaskHandlerD })), }; } - -function readDagIds(request: Pick): string[] { - // Checked, or a request this SDK cannot read would be answered with no handlers at all. - const dagIds: unknown = request.dag_ids; - if (!Array.isArray(dagIds) || !dagIds.every((dagId) => typeof dagId === "string")) { - throw new Error( - `TaskHandlerParseRequest.dag_ids must be a list of strings, got ${JSON.stringify(dagIds)}`, - ); - } - return dagIds as string[]; -} diff --git a/ts-sdk/tests/coordinator/integration.test.ts b/ts-sdk/tests/coordinator/integration.test.ts index 35269758ff42a..51112847e970f 100644 --- a/ts-sdk/tests/coordinator/integration.test.ts +++ b/ts-sdk/tests/coordinator/integration.test.ts @@ -617,7 +617,6 @@ describe("coordinator runtime integration", () => { const parseRequest = { type: "TaskHandlerParseRequest", file: "/bundles/etl.min.mjs", - dag_ids: ["py_dag"], bundle_path: "/bundles", bundle_name: "ts-task-handlers", }; diff --git a/ts-sdk/tests/coordinator/protocol.test.ts b/ts-sdk/tests/coordinator/protocol.test.ts index e799f10b073fe..0da421eb55b6c 100644 --- a/ts-sdk/tests/coordinator/protocol.test.ts +++ b/ts-sdk/tests/coordinator/protocol.test.ts @@ -58,7 +58,6 @@ describe("protocol decode", () => { const raw = { type: "TaskHandlerParseRequest", file: "/x.min.mjs", - dag_ids: ["d"], bundle_path: "/", bundle_name: "ts-task-handlers", }; diff --git a/ts-sdk/tests/coordinator/task-handler-parse.test.ts b/ts-sdk/tests/coordinator/task-handler-parse.test.ts index f25bc1728029e..029d08f62786c 100644 --- a/ts-sdk/tests/coordinator/task-handler-parse.test.ts +++ b/ts-sdk/tests/coordinator/task-handler-parse.test.ts @@ -29,19 +29,19 @@ const FILE = "/bundles/etl.min.mjs"; async function noop(): Promise {} -function declare(bundle: Bundle, dagIds: unknown) { - return declareTaskHandlers(bundle, { file: FILE, dag_ids: dagIds as string[] }); +function declare(bundle: Bundle) { + return declareTaskHandlers(bundle, { file: FILE }); } describe("declareTaskHandlers", () => { - it("declares each requested Dag's handlers in registration order", () => { + it("declares every Dag's handlers in registration order", () => { const bundle = new Bundle( new TaskHandler("etl", "extract", noop), new TaskHandler("reports", "send", noop), new TaskHandler("etl", "load", noop), ); - const result = declare(bundle, ["reports", "etl"]); + const result = declare(bundle); expect(result).toStrictEqual({ type: "TaskHandlerParsingResult", @@ -68,7 +68,7 @@ describe("declareTaskHandlers", () => { async ({ label }: ReportArgs) => label, ); - const result = declare(new Bundle(new TaskHandler("etl", "report", report)), ["etl"]); + const result = declare(new Bundle(new TaskHandler("etl", "report", report))); expect(result.task_handlers).toEqual({ etl: [ @@ -84,34 +84,20 @@ describe("declareTaskHandlers", () => { }); }); - it("leaves out unrequested Dags and requested Dags with no handler", () => { - const bundle = new Bundle( - new TaskHandler("etl", "extract", noop), - new TaskHandler("reports", "send", noop), - ); - - expect(Object.keys(declare(bundle, ["etl", "missing"]).task_handlers)).toEqual(["etl"]); - expect(declare(bundle, ["missing"]).task_handlers).toEqual({}); - }); - it("does not declare a Dag declared in TypeScript", () => { const native = new Dag("native"); native.task("run", noop)(); - const bundle = new Bundle(native, new TaskHandler("etl", "extract", noop)); - expect(Object.keys(declare(bundle, ["native", "etl"]).task_handlers)).toEqual(["etl"]); + const mixed = new Bundle(native, new TaskHandler("etl", "extract", noop)); + + expect(declare(new Bundle(native)).task_handlers).toEqual({}); + expect(Object.keys(declare(mixed).task_handlers)).toEqual(["etl"]); }); it("keeps a Dag named __proto__ as a key", () => { - const result = declare(new Bundle(new TaskHandler("__proto__", "run", noop)), ["__proto__"]); + const result = declare(new Bundle(new TaskHandler("__proto__", "run", noop))); expect(Object.keys(result.task_handlers)).toEqual(["__proto__"]); expect(Object.getPrototypeOf(result.task_handlers)).toBe(Object.prototype); }); - - it.each([[undefined], ["etl"], [["etl", 1]]])("rejects dag_ids of %j", (dagIds) => { - const bundle = new Bundle(new TaskHandler("etl", "extract", noop)); - - expect(() => declare(bundle, dagIds)).toThrow(/dag_ids must be a list of strings/); - }); }); From 6c3e3e36f83f1d38053c129615c6a6ff3a65490b Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 1 Oct 2026 17:43:00 +0000 Subject: [PATCH 71/94] TS SDK: declare task handlers without listing their params The supervisor schema dropped named_open. Its rename-only params told the Dag processor little, since types are erased and every other argument reaches a handler by folding. A handler now declares named with null params, which says the SDK cannot list them, so the Dag processor checks only that the handler exists. --- .../test_task_handler_processor_ts.py | 19 +++---------- ts-sdk/README.md | 2 +- ts-sdk/src/coordinator/task-handler-parse.ts | 28 ++++--------------- ts-sdk/tests/coordinator/integration.test.ts | 2 +- .../coordinator/task-handler-parse.test.ts | 24 +++++----------- 5 files changed, 18 insertions(+), 57 deletions(-) diff --git a/airflow-core/tests/unit/dag_processing/test_task_handler_processor_ts.py b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_ts.py index 94541a2573729..05cd0a2e4a50a 100644 --- a/airflow-core/tests/unit/dag_processing/test_task_handler_processor_ts.py +++ b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_ts.py @@ -29,11 +29,7 @@ import pytest import structlog -from airflow.dag_processing.processor import ( - TaskHandlerDeclaration, - TaskHandlerParam, - TaskHandlerParsingResult, -) +from airflow.dag_processing.processor import TaskHandlerDeclaration, TaskHandlerParsingResult from airflow.dag_processing.task_handler_processor import LangSDKTaskHandlerProcessorProcess from airflow.sdk.execution_time import supervisor from airflow.sdk.execution_time.coordinator import reset_coordinator_manager @@ -107,15 +103,8 @@ def fresh_coordinator_manager() -> Iterator[None]: reset_coordinator_manager() -def _declare(task_id: str, *renames: str) -> TaskHandlerDeclaration: - return TaskHandlerDeclaration( - task_id=task_id, - binding="named_open", - params=[ - TaskHandlerParam(name=name, value_schema=None, required=False, exact_name=True) - for name in renames - ], - ) +def _declare(task_id: str) -> TaskHandlerDeclaration: + return TaskHandlerDeclaration(task_id=task_id, binding="named", params=None) @pytest.mark.usefixtures("fresh_coordinator_manager") @@ -141,7 +130,7 @@ def test_a_packed_bundle_declares_its_task_handlers(mock_should_use_exec, exampl ], "typescript_taskflow_example": [ _declare("summarize"), - _declare("report", "run_label"), + _declare("report"), _declare("build_message"), ], }, diff --git a/ts-sdk/README.md b/ts-sdk/README.md index 2a7db7cef0391..93fc2db0f8007 100644 --- a/ts-sdk/README.md +++ b/ts-sdk/README.md @@ -289,7 +289,7 @@ Airflow launches the bundled entrypoint with `--comm=host:port` and task startup message, finds the registered handler for the Dag/task pair, and reports the terminal task state back to Airflow. Airflow can also start the bundle to ask which task handlers it registers; `bundle.serve()` then answers with each -handler's task ID and `withArgNames` renames, without running any handler. +handler's task ID, without running any handler. See [`example/`](https://github.com/apache/airflow/tree/main/ts-sdk/example) for a coordinator-runtime example that packs a bundle with `airflow-ts-pack` and diff --git a/ts-sdk/src/coordinator/task-handler-parse.ts b/ts-sdk/src/coordinator/task-handler-parse.ts index 3bc9285f183e0..dfc5320c467db 100644 --- a/ts-sdk/src/coordinator/task-handler-parse.ts +++ b/ts-sdk/src/coordinator/task-handler-parse.ts @@ -23,9 +23,7 @@ import type { TaskHandlerDeclaration, TaskHandlerParseRequest } from "../generated/supervisor.js"; import type { RuntimeTaskHandlerParsingResult } from "./protocol.js"; -import { getArgNames } from "../sdk/arg-names.js"; import { bundleTaskHandlers, type Bundle } from "../sdk/bundle.js"; -import type { TaskFunction } from "../sdk/task.js"; /** * Declare every registered task handler, keyed by Dag id, each Dag's in @@ -38,10 +36,7 @@ export function declareTaskHandlers( ): RuntimeTaskHandlerParsingResult { const declared: [string, TaskHandlerDeclaration[]][] = []; for (const [dagId, handlers] of bundleTaskHandlers(bundle)) { - declared.push([ - dagId, - [...handlers].map(([taskId, handler]) => declareTaskHandler(taskId, handler)), - ]); + declared.push([dagId, [...handlers.keys()].map((taskId) => declareTaskHandler(taskId))]); } return { type: "TaskHandlerParsingResult", @@ -52,22 +47,9 @@ export function declareTaskHandlers( } /** - * Types are erased, so a handler's `withArgNames` renames are the only names - * the SDK knows. Every other argument reaches the handler by folding, which is - * why the list is open. + * Types are erased, so the SDK cannot list the params a handler takes, and the + * Dag processor checks only that the handler exists. */ -function declareTaskHandler(taskId: string, handler: TaskFunction): TaskHandlerDeclaration { - const names = new Set(getArgNames(handler).values()); - return { - task_id: taskId, - binding: "named_open", - params: [...names].map((name) => ({ - name, - value_schema: null, - // A name the call did not pass reads as undefined, and the task still runs. - required: false, - // A rename never falls back to folding. - exact_name: true, - })), - }; +function declareTaskHandler(taskId: string): TaskHandlerDeclaration { + return { task_id: taskId, binding: "named", params: null }; } diff --git a/ts-sdk/tests/coordinator/integration.test.ts b/ts-sdk/tests/coordinator/integration.test.ts index 51112847e970f..5aa1894594436 100644 --- a/ts-sdk/tests/coordinator/integration.test.ts +++ b/ts-sdk/tests/coordinator/integration.test.ts @@ -639,7 +639,7 @@ describe("coordinator runtime integration", () => { type: "TaskHandlerParsingResult", fileloc: "/bundles/etl.min.mjs", task_handlers: { - py_dag: [{ task_id: "transform", binding: "named_open", params: [] }], + py_dag: [{ task_id: "transform", binding: "named", params: null }], }, }, }, diff --git a/ts-sdk/tests/coordinator/task-handler-parse.test.ts b/ts-sdk/tests/coordinator/task-handler-parse.test.ts index 029d08f62786c..f8feab803eee0 100644 --- a/ts-sdk/tests/coordinator/task-handler-parse.test.ts +++ b/ts-sdk/tests/coordinator/task-handler-parse.test.ts @@ -48,39 +48,29 @@ describe("declareTaskHandlers", () => { fileloc: FILE, task_handlers: { etl: [ - { task_id: "extract", binding: "named_open", params: [] }, - { task_id: "load", binding: "named_open", params: [] }, + { task_id: "extract", binding: "named", params: null }, + { task_id: "load", binding: "named", params: null }, ], - reports: [{ task_id: "send", binding: "named_open", params: [] }], + reports: [{ task_id: "send", binding: "named", params: null }], }, }); expect(Object.keys(result.task_handlers)).toEqual(["etl", "reports"]); }); - it("declares each distinct withArgNames target as an exact, optional param", () => { + it("lists no params for a handler with withArgNames renames", () => { interface ReportArgs { label: string; - title: string; owner: string; } const report = withArgNames( - { label: "run_label", title: "run_label", owner: "team" }, + { label: "run_label", owner: "team" }, async ({ label }: ReportArgs) => label, ); const result = declare(new Bundle(new TaskHandler("etl", "report", report))); - expect(result.task_handlers).toEqual({ - etl: [ - { - task_id: "report", - binding: "named_open", - params: [ - { name: "run_label", value_schema: null, required: false, exact_name: true }, - { name: "team", value_schema: null, required: false, exact_name: true }, - ], - }, - ], + expect(result.task_handlers).toStrictEqual({ + etl: [{ task_id: "report", binding: "named", params: null }], }); }); From 0d37b8643a690057210106ca7cfff7a6a50b697d Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 15:56:59 +0000 Subject: [PATCH 72/94] Java SDK: sync the vendored supervisor schema from the monorepo The vendored schema.json was refreshed only when airflowSupervisorSchemaVersion changed, by downloading the published file. The in-progress version is published only at release, so messages added during a release cycle, such as TaskHandlerParseRequest, never reached the Java models. The Gradle task and the prek hook now copy the Task SDK snapshot when it declares the configured version, and download as before otherwise, which keeps standalone source-release builds working. The refreshed file also carries main's changes since the last vendoring (retry_reason, multi_team, more required fields). --- java-sdk/sdk/build.gradle.kts | 18 + java-sdk/sdk/schema/schema.json | 817 ++++++++++++------ .../prek/sync_java_sdk_supervisor_schema.py | 28 +- 3 files changed, 588 insertions(+), 275 deletions(-) diff --git a/java-sdk/sdk/build.gradle.kts b/java-sdk/sdk/build.gradle.kts index 5e9a74a704638..614398800d4bb 100644 --- a/java-sdk/sdk/build.gradle.kts +++ b/java-sdk/sdk/build.gradle.kts @@ -161,6 +161,10 @@ abstract class GenerateDiscriminatorTask : DefaultTask() { } } +// A version is published only when it is released, so the in-progress one +// (which gains messages during a release cycle) is taken from the monorepo +// snapshot whenever that snapshot declares the configured version. Standalone +// (source-release) builds have no snapshot and download the published file. abstract class SyncSupervisorSchemaTask : DefaultTask() { @get:Input abstract val schemaVersion: Property @@ -171,6 +175,9 @@ abstract class SyncSupervisorSchemaTask : DefaultTask() { @get:Internal abstract val schemaFile: RegularFileProperty + @get:Internal + abstract val monorepoFile: RegularFileProperty + private fun apiVersionOf(file: File): String = if (file.exists()) { com.fasterxml.jackson.databind @@ -186,6 +193,16 @@ abstract class SyncSupervisorSchemaTask : DefaultTask() { fun sync() { val file = schemaFile.get().asFile val version = schemaVersion.get() + val monorepo = monorepoFile.get().asFile + if (apiVersionOf(monorepo) == version) { + if (file.exists() && file.readBytes().contentEquals(monorepo.readBytes())) { + logger.lifecycle("Supervisor Schema matches the monorepo snapshot (api_version=$version).") + } else { + logger.lifecycle("Refreshing Supervisor Schema from ${monorepo.path}") + monorepo.copyTo(file, overwrite = true) + } + return + } if (apiVersionOf(file) == version) { logger.lifecycle("Supervisor Schema is up-to-date (api_version=$version).") return @@ -640,6 +657,7 @@ val syncSupervisorSchema by tasks.registering(SyncSupervisorSchemaTask::class) { schemaVersion = airflowSupervisorSchemaVersion baseUrl = schemaBaseUrl schemaFile = schemaInput + monorepoFile = layout.projectDirectory.file("../../task-sdk/src/airflow/sdk/execution_time/schema/schema.json") } tasks.register("generateDiscriminator") { diff --git a/java-sdk/sdk/schema/schema.json b/java-sdk/sdk/schema/schema.json index b671959c50a00..30f196e693cc4 100644 --- a/java-sdk/sdk/schema/schema.json +++ b/java-sdk/sdk/schema/schema.json @@ -34,7 +34,14 @@ "anyOf": [ { "additionalProperties": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": "object" }, @@ -187,34 +194,6 @@ "title": "AssetProfile", "type": "object" }, - "AssetReferenceAssetEventDagRun": { - "additionalProperties": false, - "description": "Schema for AssetModel used in AssetEventDagRunReference.", - "properties": { - "name": { - "title": "Name", - "type": "string" - }, - "uri": { - "title": "Uri", - "type": "string" - }, - "extra": { - "additionalProperties": { - "$ref": "#/$defs/JsonValue" - }, - "title": "Extra", - "type": "object" - } - }, - "required": [ - "name", - "uri", - "extra" - ], - "title": "AssetReferenceAssetEventDagRun", - "type": "object" - }, "AssetResponse": { "description": "Asset schema for responses with fields that are needed for Runtime.", "properties": { @@ -234,7 +213,14 @@ "anyOf": [ { "additionalProperties": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": "object" }, @@ -273,7 +259,14 @@ "anyOf": [ { "additionalProperties": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": "object" }, @@ -304,7 +297,14 @@ "description": "Response to GetAssetStateStore; wraps the generated API response for supervisor to worker comms.", "properties": { "value": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": { "const": "AssetStateStoreResult", @@ -347,17 +347,10 @@ "description": "Park a task instance awaiting human input (Human-in-the-loop), without a trigger.", "properties": { "state": { - "anyOf": [ - { - "const": "awaiting_input", - "type": "string" - }, - { - "type": "null" - } - ], + "const": "awaiting_input", "default": "awaiting_input", - "title": "State" + "title": "State", + "type": "string" }, "timeout": { "anyOf": [ @@ -380,7 +373,14 @@ "anyOf": [ { "additionalProperties": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": "object" }, @@ -513,6 +513,97 @@ "title": "ClearTaskStateStore", "type": "object" }, + "ConnectionResponse": { + "description": "Connection schema for responses with fields that are needed for Runtime.", + "properties": { + "conn_id": { + "title": "Conn Id", + "type": "string" + }, + "conn_type": { + "title": "Conn Type", + "type": "string" + }, + "host": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Host" + }, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Schema" + }, + "login": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Login" + }, + "password": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Password" + }, + "port": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Port" + }, + "extra": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Extra" + } + }, + "required": [ + "conn_id", + "conn_type", + "host", + "schema", + "login", + "password", + "port", + "extra" + ], + "title": "ConnectionResponse", + "type": "object" + }, "ConnectionResult": { "properties": { "conn_id": { @@ -532,7 +623,6 @@ "type": "null" } ], - "default": null, "title": "Host" }, "schema": { @@ -544,7 +634,6 @@ "type": "null" } ], - "default": null, "title": "Schema" }, "login": { @@ -556,7 +645,6 @@ "type": "null" } ], - "default": null, "title": "Login" }, "password": { @@ -568,7 +656,6 @@ "type": "null" } ], - "default": null, "title": "Password" }, "port": { @@ -580,7 +667,6 @@ "type": "null" } ], - "default": null, "title": "Port" }, "extra": { @@ -592,7 +678,6 @@ "type": "null" } ], - "default": null, "title": "Extra" }, "type": { @@ -604,7 +689,13 @@ }, "required": [ "conn_id", - "conn_type" + "conn_type", + "host", + "schema", + "login", + "password", + "port", + "extra" ], "title": "ConnectionResult", "type": "object" @@ -924,6 +1015,20 @@ "default": null, "title": "Import Errors" }, + "parsed_definitions": { + "items": { + "type": "string" + }, + "title": "Parsed Definitions", + "type": "array" + }, + "dag_source_codes": { + "additionalProperties": { + "$ref": "#/$defs/DagSourceCode" + }, + "title": "Dag Source Codes", + "type": "object" + }, "type": { "const": "DagFileParsingResult", "default": "DagFileParsingResult", @@ -957,7 +1062,6 @@ "type": "null" } ], - "default": null, "title": "Bundle Name" }, "bundle_version": { @@ -969,7 +1073,6 @@ "type": "null" } ], - "default": null, "title": "Bundle Version" }, "relative_fileloc": { @@ -981,7 +1084,6 @@ "type": "null" } ], - "default": null, "title": "Relative Fileloc" }, "owners": { @@ -993,7 +1095,6 @@ "type": "null" } ], - "default": null, "title": "Owners" }, "tags": { @@ -1013,7 +1114,6 @@ "type": "null" } ], - "default": null, "title": "Next Dagrun" }, "type": { @@ -1026,7 +1126,12 @@ "required": [ "dag_id", "is_paused", - "tags" + "bundle_name", + "bundle_version", + "relative_fileloc", + "owners", + "tags", + "next_dagrun" ], "title": "DagResult", "type": "object" @@ -1053,13 +1158,19 @@ "type": "null" } ], - "default": null, "title": "Logical Date" }, "start_date": { - "format": "date-time", - "title": "Start Date", - "type": "string" + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Start Date" }, "end_date": { "anyOf": [ @@ -1071,7 +1182,6 @@ "type": "null" } ], - "default": null, "title": "End Date" }, "state": { @@ -1088,7 +1198,6 @@ "type": "null" } ], - "default": null, "title": "Data Interval Start" }, "data_interval_end": { @@ -1101,7 +1210,6 @@ "type": "null" } ], - "default": null, "title": "Data Interval End" }, "partition_key": { @@ -1113,15 +1221,19 @@ "type": "null" } ], - "default": null, "title": "Partition Key" } }, "required": [ "run_id", "dag_id", + "logical_date", "start_date", - "state" + "end_date", + "state", + "data_interval_start", + "data_interval_end", + "partition_key" ], "title": "DagRunAssetReference", "type": "object" @@ -1176,7 +1288,6 @@ "type": "null" } ], - "default": null, "title": "Logical Date" }, "data_interval_start": { @@ -1189,7 +1300,6 @@ "type": "null" } ], - "default": null, "title": "Data Interval Start" }, "data_interval_end": { @@ -1202,7 +1312,6 @@ "type": "null" } ], - "default": null, "title": "Data Interval End" }, "run_after": { @@ -1220,7 +1329,6 @@ "type": "null" } ], - "default": null, "title": "Start Date" }, "end_date": { @@ -1233,7 +1341,6 @@ "type": "null" } ], - "default": null, "title": "End Date" }, "clear_number": { @@ -1295,7 +1402,6 @@ "type": "null" } ], - "default": null, "title": "Partition Key" }, "partition_date": { @@ -1345,10 +1451,16 @@ "required": [ "dag_id", "run_id", + "logical_date", + "data_interval_start", + "data_interval_end", "run_after", + "start_date", + "end_date", "run_type", "state", - "consumed_asset_events" + "consumed_asset_events", + "partition_key" ], "title": "DagRunResult", "type": "object" @@ -1395,22 +1507,34 @@ "title": "DagRunType", "type": "string" }, + "DagSourceCode": { + "description": "Raw source code and its language identifier for a DAG definition.", + "properties": { + "source_code": { + "title": "Source Code", + "type": "string" + }, + "language": { + "title": "Language", + "type": "string" + } + }, + "required": [ + "source_code", + "language" + ], + "title": "DagSourceCode", + "type": "object" + }, "DeferTask": { "additionalProperties": false, "description": "Update a task instance state to deferred.", "properties": { "state": { - "anyOf": [ - { - "const": "deferred", - "type": "string" - }, - { - "type": "null" - } - ], + "const": "deferred", "default": "deferred", - "title": "State" + "title": "State", + "type": "string" }, "classpath": { "title": "Classpath", @@ -1420,7 +1544,14 @@ "anyOf": [ { "additionalProperties": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": "object" }, @@ -1467,7 +1598,14 @@ "anyOf": [ { "additionalProperties": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": "object" }, @@ -3045,7 +3183,17 @@ "title": "InactiveAssetsResult", "type": "object" }, - "JsonValue": {}, + "JsonValue": { + "title": "JsonValue", + "type": [ + "object", + "array", + "string", + "number", + "boolean", + "null" + ] + }, "LazyDeserializedDAG": { "description": "Lazily build information from the serialized DAG structure.\n\nAn object that will present \"enough\" of the DAG like interface to update DAG db models etc, without having\nto deserialize the full DAG and Task hierarchy.", "properties": { @@ -3388,17 +3536,10 @@ "description": "Update a task instance state to reschedule/up_for_reschedule.", "properties": { "state": { - "anyOf": [ - { - "const": "up_for_reschedule", - "type": "string" - }, - { - "type": "null" - } - ], + "const": "up_for_reschedule", "default": "up_for_reschedule", - "title": "State" + "title": "State", + "type": "string" }, "reschedule_date": { "format": "date-time", @@ -3441,17 +3582,10 @@ "description": "Update a task instance state to up_for_retry.", "properties": { "state": { - "anyOf": [ - { - "const": "up_for_retry", - "type": "string" - }, - { - "type": "null" - } - ], + "const": "up_for_retry", "default": "up_for_retry", - "title": "State" + "title": "State", + "type": "string" }, "end_date": { "format": "date-time", @@ -3825,17 +3959,10 @@ "description": "Update a task's state to success. Includes task_outlets and outlet_events for registering asset events.", "properties": { "state": { - "anyOf": [ - { - "const": "success", - "type": "string" - }, - { - "type": "null" - } - ], + "const": "success", "default": "success", - "title": "State" + "title": "State", + "type": "string" }, "end_date": { "format": "date-time", @@ -4029,6 +4156,179 @@ "title": "TaskCallbackRequest", "type": "object" }, + "TaskHandlerDeclaration": { + "description": "A task handler that a Lang-SDK artifact registers for one task.", + "properties": { + "task_id": { + "title": "Task Id", + "type": "string" + }, + "binding": { + "enum": [ + "positional", + "named", + "named_or_whole", + "named_open" + ], + "title": "Binding", + "type": "string" + }, + "params": { + "items": { + "$ref": "#/$defs/TaskHandlerParam" + }, + "title": "Params", + "type": "array" + } + }, + "required": [ + "task_id", + "binding", + "params" + ], + "title": "TaskHandlerDeclaration", + "type": "object" + }, + "TaskHandlerParam": { + "description": "One parameter of a task handler.", + "properties": { + "name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Name" + }, + "value_schema": { + "anyOf": [ + { + "$ref": "#/$defs/ArgValueSchema" + }, + { + "type": "null" + } + ], + "default": null + }, + "required": { + "title": "Required", + "type": "boolean" + }, + "exact_name": { + "default": false, + "title": "Exact Name", + "type": "boolean" + } + }, + "required": [ + "name", + "required" + ], + "title": "TaskHandlerParam", + "type": "object" + }, + "TaskHandlerParseRequest": { + "description": "Request for Task Handler Parsing.\n\nAsks a Lang-SDK runtime which task handlers an artifact registers for the given Dags.", + "properties": { + "file": { + "title": "File", + "type": "string" + }, + "dag_ids": { + "items": { + "type": "string" + }, + "title": "Dag Ids", + "type": "array" + }, + "bundle_path": { + "format": "path", + "title": "Bundle Path", + "type": "string" + }, + "bundle_name": { + "title": "Bundle Name", + "type": "string" + }, + "type": { + "const": "TaskHandlerParseRequest", + "default": "TaskHandlerParseRequest", + "title": "Type", + "type": "string" + } + }, + "required": [ + "file", + "dag_ids", + "bundle_path", + "bundle_name" + ], + "title": "TaskHandlerParseRequest", + "type": "object" + }, + "TaskHandlerParsingResult": { + "description": "Result of Task Handler Parsing.\n\nThe task handlers a Lang-SDK artifact registers, keyed by Dag id.", + "properties": { + "fileloc": { + "title": "Fileloc", + "type": "string" + }, + "task_handlers": { + "additionalProperties": { + "items": { + "$ref": "#/$defs/TaskHandlerDeclaration" + }, + "type": "array" + }, + "title": "Task Handlers", + "type": "object" + }, + "import_errors": { + "anyOf": [ + { + "additionalProperties": { + "type": "string" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Import Errors" + }, + "warnings": { + "anyOf": [ + { + "items": {}, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Warnings" + }, + "type": { + "const": "TaskHandlerParsingResult", + "default": "TaskHandlerParsingResult", + "title": "Type", + "type": "string" + } + }, + "required": [ + "fileloc", + "task_handlers" + ], + "title": "TaskHandlerParsingResult", + "type": "object" + }, "TaskInstanceState": { "description": "All possible states that a Task Instance can be in.\n\nNote that None is also allowed, so always use this in a type hint with Optional.", "enum": [ @@ -4119,6 +4419,18 @@ ], "default": null, "title": "Rendered Map Index" + }, + "retry_reason": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Retry Reason" } }, "required": [ @@ -4132,7 +4444,14 @@ "description": "Response to GetTaskStateStore; wraps the generated API response for supervisor to worker comms.", "properties": { "value": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": { "const": "TaskStateStoreResult", @@ -4267,6 +4586,38 @@ "title": "TriggerDagRun", "type": "object" }, + "UpdateDagRunNote": { + "properties": { + "ti_id": { + "format": "uuid", + "title": "Ti Id", + "type": "string" + }, + "note": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Note" + }, + "type": { + "const": "UpdateDagRunNote", + "default": "UpdateDagRunNote", + "title": "Type", + "type": "string" + } + }, + "required": [ + "ti_id", + "note" + ], + "title": "UpdateDagRunNote", + "type": "object" + }, "UpdateHITLDetail": { "description": "Update the response content part of an existing Human-in-the-loop response.", "properties": { @@ -4357,6 +4708,33 @@ "title": "VariableKeysResult", "type": "object" }, + "VariableResponse": { + "additionalProperties": false, + "description": "Variable schema for responses with fields that are needed for Runtime.", + "properties": { + "key": { + "title": "Key", + "type": "string" + }, + "value": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Value" + } + }, + "required": [ + "key", + "value" + ], + "title": "VariableResponse", + "type": "object" + }, "VariableResult": { "additionalProperties": false, "properties": { @@ -4373,7 +4751,6 @@ "type": "null" } ], - "default": null, "title": "Value" }, "type": { @@ -4384,7 +4761,8 @@ } }, "required": [ - "key" + "key", + "value" ], "title": "VariableResult", "type": "object" @@ -4416,7 +4794,14 @@ "type": "string" }, "value": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": { "const": "XComResult", @@ -4472,97 +4857,6 @@ "title": "XComSequenceSliceResult", "type": "object" }, - "ConnectionResponse": { - "description": "Connection schema for responses with fields that are needed for Runtime.", - "properties": { - "conn_id": { - "title": "Conn Id", - "type": "string" - }, - "conn_type": { - "title": "Conn Type", - "type": "string" - }, - "host": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Host" - }, - "schema": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Schema" - }, - "login": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Login" - }, - "password": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Password" - }, - "port": { - "anyOf": [ - { - "type": "integer" - }, - { - "type": "null" - } - ], - "title": "Port" - }, - "extra": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Extra" - } - }, - "required": [ - "conn_id", - "conn_type", - "host", - "schema", - "login", - "password", - "port", - "extra" - ], - "title": "ConnectionResponse", - "type": "object" - }, "ArgValueSchema": { "additionalProperties": { "$ref": "#/$defs/JsonValue" @@ -4573,11 +4867,6 @@ "LiteralArgBinding": { "description": "One positional stub-task argument carrying an inline literal from the Dag file.", "properties": { - "kind": { - "const": "literal", - "title": "Kind", - "type": "string" - }, "name": { "title": "Name", "type": "string" @@ -4593,6 +4882,11 @@ ], "default": null }, + "kind": { + "const": "literal", + "title": "Kind", + "type": "string" + }, "value": { "anyOf": [ { @@ -4611,8 +4905,8 @@ } }, "required": [ - "kind", - "name" + "name", + "kind" ], "title": "LiteralArgBinding", "type": "object" @@ -4638,11 +4932,6 @@ "XComArgBinding": { "description": "One positional stub-task argument pulled from an upstream task's XCom.", "properties": { - "kind": { - "const": "xcom", - "title": "Kind", - "type": "string" - }, "name": { "title": "Name", "type": "string" @@ -4658,31 +4947,19 @@ ], "default": null }, + "kind": { + "const": "xcom", + "title": "Kind", + "type": "string" + }, "task_id": { "title": "Task Id", "type": "string" - }, - "map_index": { - "default": -1, - "title": "Map Index", - "type": "integer" - }, - "element_index": { - "anyOf": [ - { - "type": "integer" - }, - { - "type": "null" - } - ], - "default": null, - "title": "Element Index" } }, "required": [ - "kind", "name", + "kind", "task_id" ], "title": "XComArgBinding", @@ -4784,6 +5061,34 @@ "title": "AssetEventDagRunReference", "type": "object" }, + "AssetReferenceAssetEventDagRun": { + "additionalProperties": false, + "description": "Schema for AssetModel used in AssetEventDagRunReference.", + "properties": { + "name": { + "title": "Name", + "type": "string" + }, + "uri": { + "title": "Uri", + "type": "string" + }, + "extra": { + "additionalProperties": { + "$ref": "#/$defs/JsonValue" + }, + "title": "Extra", + "type": "object" + } + }, + "required": [ + "name", + "uri", + "extra" + ], + "title": "AssetReferenceAssetEventDagRun", + "type": "object" + }, "DagRun": { "additionalProperties": false, "description": "Schema for DagRun model with minimal required fields needed for Runtime.", @@ -5066,6 +5371,11 @@ ], "default": null, "title": "Arg Bindings" + }, + "multi_team": { + "default": false, + "title": "Multi Team", + "type": "boolean" } }, "required": [ @@ -5150,33 +5460,6 @@ ], "title": "TaskInstance", "type": "object" - }, - "VariableResponse": { - "additionalProperties": false, - "description": "Variable schema for responses with fields that are needed for Runtime.", - "properties": { - "key": { - "title": "Key", - "type": "string" - }, - "value": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Value" - } - }, - "required": [ - "key", - "value" - ], - "title": "VariableResponse", - "type": "object" } } } diff --git a/scripts/ci/prek/sync_java_sdk_supervisor_schema.py b/scripts/ci/prek/sync_java_sdk_supervisor_schema.py index d3f4d271a5f5a..22f6578de1495 100755 --- a/scripts/ci/prek/sync_java_sdk_supervisor_schema.py +++ b/scripts/ci/prek/sync_java_sdk_supervisor_schema.py @@ -18,11 +18,12 @@ """ Keep the Java SDK's bundled Supervisor Schema in sync with ``airflowSupervisorSchemaVersion``. -The Gradle task ``:sdk:syncSupervisorSchema`` downloads a fresh ``schema.json`` when the -``api_version`` inside it differs from the version declared in ``java-sdk/gradle.properties``. +The Gradle task ``:sdk:syncSupervisorSchema`` copies the Task SDK's snapshot when it declares the +configured ``api_version``, and otherwise downloads the published schema when the ``api_version`` +inside ``schema.json`` differs from the version declared in ``java-sdk/gradle.properties``. Starting Gradle for that comparison costs over a minute on a cold CI runner (wrapper download, JVM start, plugin resolution), so this hook does the same comparison in Python first and only -hands over to Gradle when the two versions actually differ. +hands over to Gradle when a download is needed. """ from __future__ import annotations @@ -37,6 +38,9 @@ JAVA_SDK_DIR = REPO_ROOT / "java-sdk" GRADLE_PROPERTIES = JAVA_SDK_DIR / "gradle.properties" SCHEMA_FILE = JAVA_SDK_DIR / "sdk" / "schema" / "schema.json" +MONOREPO_SCHEMA_FILE = ( + REPO_ROOT / "task-sdk" / "src" / "airflow" / "sdk" / "execution_time" / "schema" / "schema.json" +) VERSION_PATTERN = re.compile(r"^airflowSupervisorSchemaVersion\s*=\s*(\S+)\s*$", re.MULTILINE) @@ -49,20 +53,28 @@ def configured_version() -> str: return match.group(1) -def bundled_version() -> str | None: - if not SCHEMA_FILE.exists(): +def api_version_of(path: Path) -> str | None: + if not path.exists(): return None try: - with SCHEMA_FILE.open() as schema: + with path.open() as schema: return json.load(schema).get("api_version") except json.JSONDecodeError: - # A truncated or conflict-marked file counts as drift; Gradle overwrites it. + # A truncated or conflict-marked file counts as drift; it is overwritten. return None def main() -> int: expected = configured_version() - actual = bundled_version() + if api_version_of(MONOREPO_SCHEMA_FILE) == expected: + snapshot = MONOREPO_SCHEMA_FILE.read_bytes() + if SCHEMA_FILE.exists() and SCHEMA_FILE.read_bytes() == snapshot: + print(f"Supervisor Schema matches the Task SDK snapshot (api_version={expected}).") + else: + print(f"Refreshing Supervisor Schema from {MONOREPO_SCHEMA_FILE.relative_to(REPO_ROOT)}.") + SCHEMA_FILE.write_bytes(snapshot) + return 0 + actual = api_version_of(SCHEMA_FILE) if actual == expected: print(f"Supervisor Schema is up-to-date (api_version={expected}).") return 0 From d2539597b1bbfd7b8944d734e526935b2699a78c Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 16:03:18 +0000 Subject: [PATCH 73/94] Java SDK: record handler parameters in generated task classes The Dag processor will ask a Java bundle how each task handler binds its arguments, but Java parameter names do not survive compilation and the generated task class reads its arguments by position only. The annotation processor now gives every class it generates for a @Builder.TaskHandler method a static AIRFLOW_TASK_PARAMS field: each flat data parameter's name and declared type (a TypeRef when it has type arguments, primitives kept primitive), or the TaskInput type. @Builder.Dag task classes are unchanged, because a Java-owned Dag is never declared as task handlers. --- .../apache/airflow/sdk/BuilderProcessor.kt | 52 +++++++-- .../org/apache/airflow/sdk/BuilderTest.kt | 109 ++++++++++++++++++ .../apache/airflow/sdk/internal/TaskParams.kt | 72 ++++++++++++ 3 files changed, 224 insertions(+), 9 deletions(-) create mode 100644 java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/TaskParams.kt diff --git a/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt b/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt index 049aba0d0601a..ee6407380fc55 100644 --- a/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt +++ b/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt @@ -23,6 +23,7 @@ package org.apache.airflow.sdk import com.squareup.javapoet.ClassName import com.squareup.javapoet.CodeBlock +import com.squareup.javapoet.FieldSpec import com.squareup.javapoet.JavaFile import com.squareup.javapoet.MethodSpec import com.squareup.javapoet.ParameterizedTypeName @@ -33,6 +34,7 @@ import org.apache.airflow.sdk.internal.Field import org.apache.airflow.sdk.internal.FieldType import org.apache.airflow.sdk.internal.SchemaFields import org.apache.airflow.sdk.internal.TaskArgs +import org.apache.airflow.sdk.internal.TaskParams import org.apache.airflow.sdk.internal.TypeRef import org.apache.airflow.sdk.internal.foldArgName import org.apache.airflow.sdk.internal.registrarName @@ -129,7 +131,8 @@ class BuilderProcessor : AbstractProcessor() { /** * Generates the registrar for a class of [Builder.TaskHandler] methods: one * [Task] implementation per handler, and a `registerInto` that binds each to - * the Dag and task the annotation names. + * the Dag and task the annotation names. Each implementation carries its + * handler's [TaskParams], so the runtime can report how it binds. * * There is no Dag to build here — the Python Dag file owns it — so this is a * registrar rather than a builder. @@ -162,7 +165,12 @@ class BuilderProcessor : AbstractProcessor() { "@Builder.TaskHandler on '${inner.simpleName}' must name the Dag the Python file declares" } val innerName = inner.simpleName.toString().replaceFirstChar(Char::uppercase) - registrar.addType(buildTask(innerName, inner, el)) + registrar.addType( + buildTask(innerName, inner, el) + .toBuilder() + .addField(taskParamsField(collectDataParams(inner))) + .build(), + ) registerInto.addStatement( $$"bundle.register($S, $S, $L.class)", handler.dag, @@ -415,6 +423,26 @@ class BuilderProcessor : AbstractProcessor() { return params } + /** + * Emits the [TaskParams.FIELD] field recording a handler's data parameters: + * its [TaskInput] type, or each flat parameter's name and declared type. + */ + private fun taskParamsField(dataParams: List): FieldSpec { + val input = dataParams.singleOrNull { it.isTaskInput } + val value = + if (input != null) { + CodeBlock.of($$"$T.input($T.class)", TASK_PARAMS_TYPE, TypeName.get(input.type)) + } else { + val params = + dataParams.map { CodeBlock.of($$"$T.param($S, $L)", TASK_PARAMS_TYPE, it.name, typeToken(TypeName.get(it.type))) } + CodeBlock.of($$"$T.of($L)", TASK_PARAMS_TYPE, CodeBlock.join(params, ", ")) + } + return FieldSpec + .builder(TASK_PARAMS_TYPE, TaskParams.FIELD, Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL) + .initializer(value) + .build() + } + private fun ProcessingEnvironment.isTaskInput(type: TypeMirror): Boolean { val marker = elementUtils.getTypeElement(TASK_INPUT_TYPE.canonicalName()) ?: return false return !type.kind.isPrimitive && typeUtils.isAssignable(type, marker.asType()) @@ -499,6 +527,7 @@ private val CLIENT_TYPE = ClassName.get(Client::class.java) private val CONTEXT_TYPE = ClassName.get(Context::class.java) private val TASK_INPUT_TYPE = ClassName.get(TaskInput::class.java) private val TASK_ARGS_TYPE = ClassName.get(TaskArgs::class.java) +private val TASK_PARAMS_TYPE = ClassName.get(TaskParams::class.java) private val TYPE_REF_TYPE = ClassName.get(TypeRef::class.java) private val ARG_VALUES_TYPE = ClassName.get(ArgValues::class.java) @@ -530,11 +559,16 @@ private fun positionalAccess( ): CodeBlock { val type = TypeName.get(param.type) val reader = if (type.isPrimitive) "require" else "get" - val target = - if (type is ParameterizedTypeName) { - CodeBlock.of($$"new $T<$T>() {}", TYPE_REF_TYPE, type) - } else { - CodeBlock.of($$"$T.class", type.box()) - } - return CodeBlock.of($$"$L.$L($L, $L)", argsLocal, reader, param.position, target) + return CodeBlock.of($$"$L.$L($L, $L)", argsLocal, reader, param.position, typeToken(type.box())) } + +/** + * Emits a runtime token for [type]: a [TypeRef] when it has type arguments, + * which a `Class` literal cannot carry, otherwise its `Class` literal. + */ +private fun typeToken(type: TypeName): CodeBlock = + if (type is ParameterizedTypeName) { + CodeBlock.of($$"new $T<$T>() {}", TYPE_REF_TYPE, type) + } else { + CodeBlock.of($$"$T.class", type) + } diff --git a/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt b/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt index 6fd52e01b75a4..09d0b09eadad5 100644 --- a/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt +++ b/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt @@ -858,6 +858,7 @@ class BuilderTest { import org.apache.airflow.sdk.Context; import org.apache.airflow.sdk.Task; import org.apache.airflow.sdk.internal.TaskArgs; + import org.apache.airflow.sdk.internal.TaskParams; /** * Registers {@link TestExample}'s task handlers against the Dags the Python file owns. @@ -869,6 +870,8 @@ class BuilderTest { } public static final class Score implements Task { + public static final TaskParams AIRFLOW_TASK_PARAMS = TaskParams.of(TaskParams.param("rows", long.class)); + @Override public void execute(Context context, Client client) throws Exception { TaskArgs args = TaskArgs.of(context, client, 1); @@ -878,6 +881,8 @@ class BuilderTest { } public static final class Audit implements Task { + public static final TaskParams AIRFLOW_TASK_PARAMS = TaskParams.of(); + @Override public void execute(Context context, Client client) throws Exception { new TestExample().audit(); @@ -888,6 +893,107 @@ class BuilderTest { ) } + @Test + @DisplayName("record each handler's data parameters for the runtime to declare") + fun generateHandlerTaskParams() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import java.util.List; + import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.Client; + import org.apache.airflow.sdk.Context; + import org.apache.airflow.sdk.TaskInput; + public class TestExample { + public static class ScoreInput implements TaskInput { + public String region; + } + + @Builder.TaskHandler(dag = "etl") + public void flat(Client client, long rows, Integer limit, List regions) {} + + @Builder.TaskHandler(dag = "etl") + public void named(ScoreInput input) {} + + @Builder.TaskHandler(dag = "etl") + public void none(Context context) {} + } + """, + ) + + assertThat(compilation).succeeded() + assertThat(compilation) + .generatedSourceFile("org.apache.airflow.example.TestExampleHandlers") + .hasSourceEquivalentTo( + "org.apache.airflow.example.TestExampleHandlers", + """ + package org.apache.airflow.example; + + import java.lang.Exception; + import java.lang.Integer; + import java.lang.Long; + import java.lang.Override; + import java.lang.String; + import java.util.List; + import org.apache.airflow.sdk.Bundle; + import org.apache.airflow.sdk.Client; + import org.apache.airflow.sdk.Context; + import org.apache.airflow.sdk.Task; + import org.apache.airflow.sdk.internal.ArgValues; + import org.apache.airflow.sdk.internal.TaskArgs; + import org.apache.airflow.sdk.internal.TaskParams; + import org.apache.airflow.sdk.internal.TypeRef; + + /** + * Registers {@link TestExample}'s task handlers against the Dags the Python file owns. + */ + public final class TestExampleHandlers { + public static void registerInto(Bundle bundle) { + bundle.register("etl", "flat", Flat.class); + bundle.register("etl", "named", Named.class); + bundle.register("etl", "none", None.class); + } + + public static final class Flat implements Task { + public static final TaskParams AIRFLOW_TASK_PARAMS = TaskParams.of( + TaskParams.param("rows", long.class), + TaskParams.param("limit", Integer.class), + TaskParams.param("regions", new TypeRef>() {})); + + @Override + public void execute(Context context, Client client) throws Exception { + TaskArgs args = TaskArgs.of(context, client, 3); + long rows = args.require(0, Long.class); + Integer limit = args.get(1, Integer.class); + List regions = args.get(2, new TypeRef>() {}); + new TestExample().flat(client, rows, limit, regions); + } + } + + public static final class Named implements Task { + public static final TaskParams AIRFLOW_TASK_PARAMS = TaskParams.input(TestExample.ScoreInput.class); + + @Override + public void execute(Context context, Client client) throws Exception { + TestExample.ScoreInput input = ArgValues.bindInput(client, TestExample.ScoreInput.class); + new TestExample().named(input); + } + } + + public static final class None implements Task { + public static final TaskParams AIRFLOW_TASK_PARAMS = TaskParams.of(); + + @Override + public void execute(Context context, Client client) throws Exception { + new TestExample().none(context); + } + } + } + """, + ) + } + @Test @DisplayName("reject a handler that names no Dag") fun rejectHandlerWithoutDag() { @@ -941,6 +1047,7 @@ class BuilderTest { import org.apache.airflow.sdk.Context; import org.apache.airflow.sdk.Task; import org.apache.airflow.sdk.internal.TaskArgs; + import org.apache.airflow.sdk.internal.TaskParams; /** * Registers {@link TestExample.Inner}'s task handlers against the Dags the Python file owns. @@ -951,6 +1058,8 @@ class BuilderTest { } public static final class Score implements Task { + public static final TaskParams AIRFLOW_TASK_PARAMS = TaskParams.of(TaskParams.param("rows", long.class)); + @Override public void execute(Context context, Client client) throws Exception { TaskArgs args = TaskArgs.of(context, client, 1); diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/TaskParams.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/TaskParams.kt new file mode 100644 index 0000000000000..7181b13a87e38 --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/TaskParams.kt @@ -0,0 +1,72 @@ +/* + * 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. + */ + +package org.apache.airflow.sdk.internal + +import org.apache.airflow.sdk.TaskInput +import java.lang.reflect.Type + +/** + * @suppress + * + * The data parameters of a task class generated for a `@Builder.TaskHandler` + * method, as the runtime reports them when asked which task handlers a bundle + * registers. The generated class holds it in a static field named [FIELD]. + * + * Java parameter names do not survive compilation, so the annotation + * processor records them here. Public so that generated code can build it; + * not user-facing API. + */ +class TaskParams private constructor( + internal val flat: List, + internal val input: Class?, +) { + /** One flat data parameter: its name in the handler method, and its declared type. */ + class Param internal constructor( + internal val name: String, + internal val type: Type, + ) + + companion object { + /** Name of the static field holding a generated task class's [TaskParams]. */ + const val FIELD = "AIRFLOW_TASK_PARAMS" + + /** Flat data parameters, in declaration order; none for a method that takes no data. */ + @JvmStatic + fun of(vararg params: Param): TaskParams = TaskParams(params.toList(), null) + + /** A single [TaskInput] parameter of [type]. */ + @JvmStatic + fun input(type: Class): TaskParams = TaskParams(emptyList(), type) + + /** A flat parameter whose declared type a `Class` literal expresses; primitives stay primitive. */ + @JvmStatic + fun param( + name: String, + type: Class<*>, + ): Param = Param(name, type) + + /** A flat parameter whose declared type has type arguments. */ + @JvmStatic + fun param( + name: String, + type: TypeRef<*>, + ): Param = Param(name, type.type) + } +} From 078264dd994c06b8381ac63ad61e4a2df3bf3062 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 16:09:05 +0000 Subject: [PATCH 74/94] Java SDK: answer TaskHandlerParseRequest with registered handlers The Dag processor asks a Lang-SDK runtime which task handlers an artifact registers for the Dags a Python file owns, and today a Java bundle exits on that first frame as an unknown message type. Server now replies with one TaskHandlerParsingResult: the task-handler registrations of each requested Dag, in registration order, with a Dag that has none left out. A Dag declared in Java is never reported, since ADR-0011 keeps a Python-owned Dag id out of Java Dag registrations. Flat data parameters bind positional, with their names and required, and a TaskInput binds named, with @ArgName fields matched exactly. A handler that reads no argument binds named with no params, so a stub that passes arguments to it is flagged. value_schema is the declared Java type's schema in pydantic's vocabulary, with null excluded only for primitives. It does not widen to the scalars the decoder also coerces, which is how pydantic describes its lax types too. The reply is a plain map because the generated models cannot hold task_handlers. It waits for the parent's acknowledgement before the process exits. --- java-sdk/README.md | 9 + .../kotlin/org/apache/airflow/sdk/Server.kt | 16 ++ .../airflow/sdk/execution/TaskHandlerParse.kt | 113 +++++++++ .../airflow/sdk/internal/ValueSchema.kt | 154 ++++++++++++ .../org/apache/airflow/sdk/ServerTest.kt | 92 ++++++++ .../sdk/execution/TaskHandlerParseTest.kt | 183 +++++++++++++++ .../airflow/sdk/internal/ValueSchemaTest.kt | 219 ++++++++++++++++++ 7 files changed, 786 insertions(+) create mode 100644 java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParse.kt create mode 100644 java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/ValueSchema.kt create mode 100644 java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParseTest.kt create mode 100644 java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/ValueSchemaTest.kt diff --git a/java-sdk/README.md b/java-sdk/README.md index 752253b6ab826..53f3c037f7303 100644 --- a/java-sdk/README.md +++ b/java-sdk/README.md @@ -674,6 +674,15 @@ launches the JVM application as a subprocess. The flow is: 5. On completion (or exception) the JVM sends a `TaskState` message and closes the socket. The JVM process then exits. +The Dag processor also launches each JAR to ask which task handlers it +registers. The first frame is then a `TaskHandlerParseRequest`, and +`execution/TaskHandlerParse.kt` replies with one `TaskHandlerParsingResult`: +the task-handler registrations of the requested Dags, each with its binding +mode and parameters. Flat data parameters bind `positional` and a `TaskInput` +binds `named`. The annotation processor records a handler's parameter names on +its generated task class (`internal/TaskParams.kt`), because Java does not keep +them at run time. The JVM waits for the reply's acknowledgement, then exits. + Log messages produced by the SDK (not by user code) are forwarded over the `--logs` socket so the supervisor can append them to Airflow's log store. diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt index 4b42eb68d4c39..e3b8bd49a08b7 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt @@ -35,6 +35,8 @@ import org.apache.airflow.sdk.execution.LogSender import org.apache.airflow.sdk.execution.Logger import org.apache.airflow.sdk.execution.comm.ErrorResponse import org.apache.airflow.sdk.execution.comm.StartupDetails +import org.apache.airflow.sdk.execution.comm.TaskHandlerParseRequest +import org.apache.airflow.sdk.execution.parseTaskHandlers import org.apache.airflow.sdk.execution.runTask import kotlin.text.substringAfterLast import kotlin.text.substringBeforeLast @@ -70,6 +72,11 @@ class ApiError( * Connects this JVM process to the Airflow coordinator and dispatches task * execution requests to the registered [Bundle]. * + * The Dag processor starts the same entry point to ask which task handlers + * the bundle registers. That process replies with the handlers registered for + * the requested Dags and how each binds its arguments, runs no task, and + * exits. + * * The typical entry point is: * * ```java @@ -181,6 +188,7 @@ class Server( val frame = coordinator.readMessage() when (val body = frame.body) { is StartupDetails -> runTaskAndReport(bundle, body, coordinator) + is TaskHandlerParseRequest -> parseTaskHandlersAndReport(bundle, body, coordinator) is ErrorResponse -> throw ApiError("[${body.error}] ${body.detail}") else -> throw ApiError("Unexpected initial frame (id=${frame.id})") } @@ -194,4 +202,12 @@ class Server( val result = runTask(bundle, startup, coordinator) coordinator.communicate(result) } + + private suspend fun parseTaskHandlersAndReport( + bundle: Bundle, + request: TaskHandlerParseRequest, + coordinator: CoordinatorComm, + ) { + coordinator.communicate(parseTaskHandlers(bundle, request)) + } } diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParse.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParse.kt new file mode 100644 index 0000000000000..9c177f9fe0fa3 --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParse.kt @@ -0,0 +1,113 @@ +/* + * 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. + */ + +package org.apache.airflow.sdk.execution + +import org.apache.airflow.sdk.Bundle +import org.apache.airflow.sdk.InputTask +import org.apache.airflow.sdk.Task +import org.apache.airflow.sdk.TaskDef +import org.apache.airflow.sdk.TaskInput +import org.apache.airflow.sdk.execution.comm.TaskHandlerDeclaration.Binding +import org.apache.airflow.sdk.execution.comm.TaskHandlerParseRequest +import org.apache.airflow.sdk.internal.TaskParams +import org.apache.airflow.sdk.internal.argNameOf +import org.apache.airflow.sdk.internal.bindableFields +import org.apache.airflow.sdk.internal.buildValueSchema +import org.apache.airflow.sdk.internal.isPinned +import org.apache.airflow.sdk.internal.resolveInputType +import java.lang.reflect.Modifier +import java.lang.reflect.Type + +/** + * Answers a [TaskHandlerParseRequest] with the task handlers [bundle] + * registers for the requested Dags, as a TaskHandlerParsingResult body. + * + * Only task-handler registrations count: a Dag declared in Java is its own, + * not a Python Dag's. A requested Dag with no handlers is left out. The + * generated models cannot hold `task_handlers`, so the body is a map. + */ +internal fun parseTaskHandlers( + bundle: Bundle, + request: TaskHandlerParseRequest, +): Map { + val declarations = linkedMapOf>>() + for (dagId in request.dagIds) { + if (dagId in declarations) continue + val dag = bundle.taskHandlers[dagId] ?: continue + declarations[dagId] = dag.tasks.values.map(::declare) + } + return linkedMapOf( + "type" to "TaskHandlerParsingResult", + "fileloc" to request.file, + "task_handlers" to declarations, + ) +} + +private fun declare(task: TaskDef): Map { + val (binding, params) = bindingOf(task.definition) + return linkedMapOf( + "task_id" to task.id, + "binding" to binding.value(), + "params" to params, + ) +} + +/** + * How [definition] binds the stub call's arguments. Flat parameters bind by + * position; a [TaskInput] binds by name. A task that reads no argument is + * declared as binding by name with no parameters, so an argument passed to it + * stands out as one nothing takes. + */ +private fun bindingOf(definition: Class): Pair>> { + if (InputTask::class.java.isAssignableFrom(definition)) { + return Binding.NAMED to inputParams(resolveInputType(definition)) + } + val declared = taskParamsOf(definition) ?: return Binding.NAMED to emptyList() + declared.input?.let { return Binding.NAMED to inputParams(it) } + if (declared.flat.isEmpty()) return Binding.NAMED to emptyList() + return Binding.POSITIONAL to declared.flat.map { param(it.name, it.type, required = true, exactName = false) } +} + +/** The [TaskParams] the annotation processor recorded on a generated task class, if [definition] is one. */ +private fun taskParamsOf(definition: Class<*>): TaskParams? { + val field = + try { + definition.getField(TaskParams.FIELD) + } catch (_: NoSuchFieldException) { + return null + } + return if (Modifier.isStatic(field.modifiers)) field.get(null) as? TaskParams else null +} + +// A field nothing supplies keeps its Java default and only warns, so no field is required. +private fun inputParams(inputType: Class): List> = + bindableFields(inputType).map { param(argNameOf(it), it.genericType, required = false, exactName = isPinned(it)) } + +private fun param( + name: String, + type: Type, + required: Boolean, + exactName: Boolean, +): Map = + linkedMapOf( + "name" to name, + "required" to required, + "exact_name" to exactName, + ).apply { buildValueSchema(type)?.let { put("value_schema", it) } } diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/ValueSchema.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/ValueSchema.kt new file mode 100644 index 0000000000000..d3470438015ad --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/ValueSchema.kt @@ -0,0 +1,154 @@ +/* + * 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. + */ + +@file:Suppress("PLATFORM_CLASS_MAPPED_TO_KOTLIN") + +package org.apache.airflow.sdk.internal + +import java.lang.reflect.GenericArrayType +import java.lang.reflect.ParameterizedType +import java.lang.reflect.Type +import java.lang.reflect.WildcardType +import java.math.BigDecimal +import java.math.BigInteger +import java.time.Duration +import java.time.Instant +import java.time.LocalDate +import java.time.LocalDateTime +import java.time.LocalTime +import java.time.OffsetDateTime +import java.time.ZonedDateTime +import java.util.UUID + +private typealias Schema = Map + +private val UNCONSTRAINED: Schema = emptyMap() + +/** + * Builds the JSON Schema of the values a parameter or field of [type] + * accepts, in the vocabulary pydantic emits for a Python annotation, or null + * when [type] does not say. + * + * The schema describes the declared type, not every value the decoder + * coerces into it: [ArgValues] also turns `"5"` into a `long`, as pydantic's + * lax mode does for `int`. Only a primitive excludes null, which the decoder + * does enforce. + */ +internal fun buildValueSchema(type: Type): Schema? { + val schema = schemaOf(type) ?: return null + return if (type is Class<*> && type.isPrimitive) schema else nullable(schema) +} + +private fun nullable(schema: Schema): Schema = mapOf("anyOf" to listOf(schema, mapOf("type" to "null"))) + +private fun schemaOf(type: Type): Schema? = + when (type) { + is Class<*> -> classSchema(type) + is ParameterizedType -> parameterizedSchema(type) + is GenericArrayType -> arraySchema(type.genericComponentType) + is WildcardType -> type.upperBounds.firstOrNull()?.let(::schemaOf) + else -> null // A type variable says nothing about its values. + } + +private fun integer(format: String): Schema = mapOf("type" to "integer", "format" to format) + +private fun integerRange( + minimum: Long, + maximum: Long, +): Schema = mapOf("type" to "integer", "minimum" to minimum, "maximum" to maximum) + +private fun number(format: String): Schema = mapOf("type" to "number", "format" to format) + +private fun string(format: String): Schema = mapOf("type" to "string", "format" to format) + +private val SCALARS: Map, Schema> = + buildMap { + fun put( + primitive: Class<*>?, + boxed: Class<*>, + schema: Schema, + ) { + primitive?.let { put(it, schema) } + put(boxed, schema) + } + put(java.lang.Boolean.TYPE, java.lang.Boolean::class.java, mapOf("type" to "boolean")) + put(java.lang.Byte.TYPE, java.lang.Byte::class.java, integerRange(Byte.MIN_VALUE.toLong(), Byte.MAX_VALUE.toLong())) + put(java.lang.Short.TYPE, java.lang.Short::class.java, integerRange(Short.MIN_VALUE.toLong(), Short.MAX_VALUE.toLong())) + put(java.lang.Integer.TYPE, java.lang.Integer::class.java, integer("int32")) + put(java.lang.Long.TYPE, java.lang.Long::class.java, integer("int64")) + put(java.lang.Float.TYPE, java.lang.Float::class.java, number("float")) + put(java.lang.Double.TYPE, java.lang.Double::class.java, number("double")) + put( + java.lang.Character.TYPE, + java.lang.Character::class.java, + mapOf("type" to "string", "minLength" to 1, "maxLength" to 1), + ) + put(null, java.lang.String::class.java, mapOf("type" to "string")) + put(null, BigInteger::class.java, mapOf("type" to "integer")) + put(null, BigDecimal::class.java, mapOf("type" to "number")) + put(null, java.lang.Number::class.java, mapOf("type" to "number")) + put(null, UUID::class.java, string("uuid")) + put(null, OffsetDateTime::class.java, string("date-time")) + put(null, ZonedDateTime::class.java, string("date-time")) + put(null, LocalDateTime::class.java, string("date-time")) + put(null, Instant::class.java, string("date-time")) + put(null, LocalDate::class.java, string("date")) + put(null, LocalTime::class.java, string("time")) + put(null, Duration::class.java, string("duration")) + } + +private fun classSchema(type: Class<*>): Schema? { + SCALARS[type]?.let { return it } + return when { + // A byte array decodes from a base64 string as well as from a list of numbers. + type == ByteArray::class.java -> null + type.isArray -> arraySchema(type.componentType) + type.isEnum -> enumSchema(type) + Collection::class.java.isAssignableFrom(type) -> mapOf("type" to "array", "items" to UNCONSTRAINED) + Map::class.java.isAssignableFrom(type) -> mapOf("type" to "object", "additionalProperties" to true) + // Other platform types (Object, Optional, Date, URI, ...) decode from more + // than one JSON shape, or from none, so they constrain nothing here. + isPlatformType(type) -> null + else -> mapOf("type" to "object") + } +} + +private fun isPlatformType(type: Class<*>): Boolean = + type.name.startsWith("java.") || type.name.startsWith("javax.") || type.name.startsWith("kotlin.") + +// Read from the constant fields rather than enumConstants, which would run the +// enum's static initializer while the bundle is only being described. +private fun enumSchema(type: Class<*>): Schema = + mapOf( + "type" to "string", + "enum" to type.fields.filter { it.isEnumConstant }.map { it.name }, + ) + +private fun arraySchema(component: Type): Schema = mapOf("type" to "array", "items" to (buildValueSchema(component) ?: UNCONSTRAINED)) + +private fun parameterizedSchema(type: ParameterizedType): Schema? { + val raw = type.rawType as? Class<*> ?: return null + val arguments = type.actualTypeArguments + return when { + Collection::class.java.isAssignableFrom(raw) && arguments.size == 1 -> arraySchema(arguments[0]) + Map::class.java.isAssignableFrom(raw) && arguments.size == 2 -> + mapOf("type" to "object", "additionalProperties" to (buildValueSchema(arguments[1]) ?: true)) + else -> classSchema(raw) + } +} diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt index e436dd2fc6160..d5a0db1bfd31f 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt @@ -27,16 +27,25 @@ import kotlinx.coroutines.runBlocking import org.apache.airflow.sdk.execution.CoordinatorComm import org.apache.airflow.sdk.execution.Frame import org.apache.airflow.sdk.execution.IncomingFrame +import org.apache.airflow.sdk.execution.RawFrame import org.apache.airflow.sdk.execution.comm.TaskState import org.junit.jupiter.api.Assertions import org.junit.jupiter.api.DisplayName import org.junit.jupiter.api.Test import org.junit.jupiter.api.Timeout import org.msgpack.core.MessagePack +import org.msgpack.core.buffer.ArrayBufferInput import java.io.ByteArrayOutputStream import java.util.concurrent.ArrayBlockingQueue import java.util.concurrent.TimeUnit +private class ServerHandlerTask : Task { + override fun execute( + context: Context, + client: Client, + ) = Unit +} + class ServerTest { private fun hexToBytes(hex: String): ByteArray = hex @@ -112,6 +121,89 @@ class ServerTest { comm.close() } + @Test + @DisplayName("Should answer a task handler parse request and wait for its acknowledgement") + @Timeout(value = 30, unit = TimeUnit.SECONDS) + fun answersTaskHandlerParseRequest() { + val toServer = ByteChannel(autoFlush = true) + val fromServer = ByteChannel(autoFlush = true) + val comm = CoordinatorComm(toServer, fromServer) + val server = Server(InetSocketAddress("localhost", 0), InetSocketAddress("localhost", 0)) + val bundle = Bundle().register("etl", "audit", ServerHandlerTask::class.java) + + val reported = ArrayBlockingQueue(1) + val supervisor = + Thread { + runBlocking { + toServer.writeFrame(taskHandlerParseRequestFrame(3, "/bundle/etl.jar", listOf("etl"))) + val prefix = fromServer.readByteArray(4) + val raw = Frame.decodeRaw(ArrayBufferInput(fromServer.readByteArray(Frame.parseLengthPrefix(prefix).toInt()))) + reported.put(raw) + toServer.writeFrame(ackFrame(raw.id)) + } + } + supervisor.start() + + runBlocking { server.dispatchTask(bundle, comm) } + supervisor.join() + + val body = reported.take().rawBody as Map<*, *> + Assertions.assertEquals("TaskHandlerParsingResult", body["type"]) + Assertions.assertEquals("/bundle/etl.jar", body["fileloc"]) + Assertions.assertEquals( + mapOf("etl" to listOf(mapOf("task_id" to "audit", "binding" to "named", "params" to emptyList()))), + body["task_handlers"], + ) + comm.close() + } + + @Test + @DisplayName("Should fail a task handler parse whose result is never acknowledged") + @Timeout(value = 30, unit = TimeUnit.SECONDS) + fun failsTaskHandlerParseWithoutAcknowledgement() { + val toServer = ByteChannel(autoFlush = true) + val fromServer = ByteChannel(autoFlush = true) + val comm = CoordinatorComm(toServer, fromServer) + val server = Server(InetSocketAddress("localhost", 0), InetSocketAddress("localhost", 0)) + + val supervisor = + Thread { + runBlocking { + toServer.writeFrame(taskHandlerParseRequestFrame(3, "/bundle/etl.jar", listOf("etl"))) + val prefix = fromServer.readByteArray(4) + fromServer.readByteArray(Frame.parseLengthPrefix(prefix).toInt()) + toServer.close() + } + } + supervisor.start() + + Assertions.assertThrows(ApiError::class.java) { + runBlocking { server.dispatchTask(Bundle(), comm) } + } + supervisor.join() + comm.close() + } + + private fun taskHandlerParseRequestFrame( + id: Int, + file: String, + dagIds: List, + ): ByteArray { + val out = ByteArrayOutputStream() + MessagePack.newDefaultPacker(out).use { packer -> + packer.packArrayHeader(2) + packer.packInt(id) + packer.packMapHeader(5) + packer.packString("type").packString("TaskHandlerParseRequest") + packer.packString("file").packString(file) + packer.packString("dag_ids").packArrayHeader(dagIds.size) + dagIds.forEach { packer.packString(it) } + packer.packString("bundle_path").packString("/bundle") + packer.packString("bundle_name").packString("java-task-handlers") + } + return out.toByteArray() + } + private companion object { // [2, msg, null] with msg coming from // https://github.com/astronomer/airflow/blob/f39c8da8/task-sdk/tests/task_sdk/execution_time/test_comms.py#L73-L108 diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParseTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParseTest.kt new file mode 100644 index 0000000000000..49a769ac9e8bb --- /dev/null +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParseTest.kt @@ -0,0 +1,183 @@ +/* + * 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. + */ + +package org.apache.airflow.sdk.execution + +import org.apache.airflow.sdk.ArgName +import org.apache.airflow.sdk.Bundle +import org.apache.airflow.sdk.Client +import org.apache.airflow.sdk.Context +import org.apache.airflow.sdk.DagDef +import org.apache.airflow.sdk.InputTask +import org.apache.airflow.sdk.Task +import org.apache.airflow.sdk.TaskInput +import org.apache.airflow.sdk.execution.comm.TaskHandlerParseRequest +import org.apache.airflow.sdk.internal.TaskParams +import org.apache.airflow.sdk.internal.TypeRef +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.DisplayName +import org.junit.jupiter.api.Test + +private val NULLABLE_STRING = mapOf("anyOf" to listOf(mapOf("type" to "string"), mapOf("type" to "null"))) + +/** Stands in for a class the annotation processor generates for a handler with flat data parameters. */ +private class GeneratedFlat : Task { + companion object { + @JvmField + val AIRFLOW_TASK_PARAMS: TaskParams = + TaskParams.of( + TaskParams.param("rows", java.lang.Long.TYPE), + TaskParams.param("regions", object : TypeRef>() {}), + ) + } + + override fun execute( + context: Context, + client: Client, + ) = Unit +} + +/** Stands in for a generated class whose handler takes only the injected Client and Context. */ +private class GeneratedNoData : Task { + companion object { + @JvmField + val AIRFLOW_TASK_PARAMS: TaskParams = TaskParams.of() + } + + override fun execute( + context: Context, + client: Client, + ) = Unit +} + +private open class BaseScoreInput : TaskInput { + @JvmField + var threshold: Double = 0.0 +} + +private class ScoreInput : BaseScoreInput() { + @JvmField + @ArgName("region_code") + var region: String? = null +} + +/** Stands in for a generated class whose handler takes a [TaskInput]. */ +private class GeneratedInput : Task { + companion object { + @JvmField + val AIRFLOW_TASK_PARAMS: TaskParams = TaskParams.input(ScoreInput::class.java) + } + + override fun execute( + context: Context, + client: Client, + ) = Unit +} + +private class Score : InputTask { + override fun execute( + context: Context, + client: Client, + input: ScoreInput, + ) = Unit +} + +private class Audit : Task { + override fun execute( + context: Context, + client: Client, + ) = Unit +} + +private fun request(vararg dagIds: String) = + TaskHandlerParseRequest().apply { + file = "/bundles/java/etl.jar" + this.dagIds = dagIds.toList() + bundlePath = "/bundles/java" + bundleName = "java-task-handlers" + } + +private fun param( + name: String, + schema: Map, + required: Boolean = false, + exactName: Boolean = false, +) = mapOf("name" to name, "required" to required, "exact_name" to exactName, "value_schema" to schema) + +private val SCORE_INPUT_PARAMS = + listOf( + param("region_code", NULLABLE_STRING, exactName = true), + param("threshold", mapOf("type" to "number", "format" to "double")), + ) + +class TaskHandlerParseTest { + @Test + @DisplayName("Should declare each requested Dag's task handlers in registration order") + fun declaresRequestedDagsInRegistrationOrder() { + val bundle = + Bundle() + .register("etl", "extract", GeneratedFlat::class.java) + .register("etl", "score", GeneratedInput::class.java) + .register("etl", "summarize", Score::class.java) + .register("etl", "audit", Audit::class.java) + .register("etl", "notify", GeneratedNoData::class.java) + .register("report", "audit", Audit::class.java) + .register(DagDef("native").addTask("audit", Audit::class.java)) + + val result = parseTaskHandlers(bundle, request("etl", "native", "missing", "etl")) + + val regions = + mapOf("anyOf" to listOf(mapOf("type" to "array", "items" to NULLABLE_STRING), mapOf("type" to "null"))) + val expected = + mapOf( + "type" to "TaskHandlerParsingResult", + "fileloc" to "/bundles/java/etl.jar", + "task_handlers" to + mapOf( + "etl" to + listOf( + mapOf( + "task_id" to "extract", + "binding" to "positional", + "params" to + listOf( + param("rows", mapOf("type" to "integer", "format" to "int64"), required = true), + param("regions", regions, required = true), + ), + ), + mapOf("task_id" to "score", "binding" to "named", "params" to SCORE_INPUT_PARAMS), + mapOf("task_id" to "summarize", "binding" to "named", "params" to SCORE_INPUT_PARAMS), + mapOf("task_id" to "audit", "binding" to "named", "params" to emptyList()), + mapOf("task_id" to "notify", "binding" to "named", "params" to emptyList()), + ), + ), + ) + assertEquals(expected, result) + } + + @Test + @DisplayName("Should send an empty map when no requested Dag has task handlers") + fun sendsEmptyTaskHandlersWhenNothingMatches() { + val bundle = Bundle().register("etl", "audit", Audit::class.java) + + val result = parseTaskHandlers(bundle, request("other")) + + assertEquals(emptyMap(), result["task_handlers"]) + } +} diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/ValueSchemaTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/ValueSchemaTest.kt new file mode 100644 index 0000000000000..af7bc525f9b0e --- /dev/null +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/ValueSchemaTest.kt @@ -0,0 +1,219 @@ +/* + * 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. + */ + +@file:Suppress("PLATFORM_CLASS_MAPPED_TO_KOTLIN") + +package org.apache.airflow.sdk.internal + +import org.junit.jupiter.api.Assertions.assertAll +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertFalse +import org.junit.jupiter.api.DisplayName +import org.junit.jupiter.api.Test +import java.math.BigDecimal +import java.math.BigInteger +import java.time.Duration +import java.time.Instant +import java.time.LocalDate +import java.time.LocalTime +import java.time.OffsetDateTime +import java.util.Date +import java.util.Optional +import java.util.UUID + +private fun nullable(schema: Map) = mapOf("anyOf" to listOf(schema, mapOf("type" to "null"))) + +private val INT64 = mapOf("type" to "integer", "format" to "int64") +private val STRING = mapOf("type" to "string") + +private class Payload + +private enum class Color { RED, GREEN } + +private object EnumInitialization { + var ran = false +} + +private enum class Lazy { + ONLY, + ; + + init { + EnumInitialization.ran = true + } +} + +/** Declares one field per case, so each keeps the generic type a Kotlin declaration compiles to. */ +@Suppress("unused") +private class Declared { + @JvmField var flag: Boolean = false + + @JvmField var boxedFlag: Boolean? = null + + @JvmField var tiny: Byte = 0 + + @JvmField var small: Short = 0 + + @JvmField var count: Int = 0 + + @JvmField var total: Long = 0 + + @JvmField var boxedTotal: Long? = null + + @JvmField var ratio: Float = 0f + + @JvmField var score: Double = 0.0 + + @JvmField var letter: Char = 'a' + + @JvmField var text: String? = null + + @JvmField var big: BigInteger? = null + + @JvmField var exact: BigDecimal? = null + + @JvmField var id: UUID? = null + + @JvmField var color: Color? = null + + @JvmField var at: OffsetDateTime? = null + + @JvmField var instant: Instant? = null + + @JvmField var day: LocalDate? = null + + @JvmField var time: LocalTime? = null + + @JvmField var wait: Duration? = null + + @JvmField var names: List? = null + + @JvmField var anything: List? = null + + @JvmField var words: Array? = null + + @JvmField var totals: LongArray? = null + + @JvmField var bytes: ByteArray? = null + + @JvmField var counts: Map? = null + + @JvmField var payload: Payload? = null + + @JvmField var any: Any? = null + + @JvmField var date: Date? = null + + @JvmField var maybe: Optional? = null + + // Kotlin compiles a non-final type argument of a parameter to a wildcard, `List`. + fun acceptNumbers(numbers: List) = numbers +} + +private fun schemaOf(field: String) = buildValueSchema(Declared::class.java.getField(field).genericType) + +class ValueSchemaTest { + @Test + @DisplayName("Should describe scalars in pydantic's vocabulary, excluding null only for primitives") + fun describesScalars() { + assertAll( + { assertEquals(mapOf("type" to "boolean"), schemaOf("flag")) }, + { assertEquals(nullable(mapOf("type" to "boolean")), schemaOf("boxedFlag")) }, + { assertEquals(mapOf("type" to "integer", "minimum" to -128L, "maximum" to 127L), schemaOf("tiny")) }, + { assertEquals(mapOf("type" to "integer", "minimum" to -32768L, "maximum" to 32767L), schemaOf("small")) }, + { assertEquals(mapOf("type" to "integer", "format" to "int32"), schemaOf("count")) }, + { assertEquals(INT64, schemaOf("total")) }, + { assertEquals(nullable(INT64), schemaOf("boxedTotal")) }, + { assertEquals(mapOf("type" to "number", "format" to "float"), schemaOf("ratio")) }, + { assertEquals(mapOf("type" to "number", "format" to "double"), schemaOf("score")) }, + { assertEquals(mapOf("type" to "string", "minLength" to 1, "maxLength" to 1), schemaOf("letter")) }, + { assertEquals(nullable(STRING), schemaOf("text")) }, + { assertEquals(nullable(mapOf("type" to "integer")), schemaOf("big")) }, + { assertEquals(nullable(mapOf("type" to "number")), schemaOf("exact")) }, + { assertEquals(nullable(mapOf("type" to "string", "format" to "uuid")), schemaOf("id")) }, + { assertEquals(nullable(mapOf("type" to "string", "enum" to listOf("RED", "GREEN"))), schemaOf("color")) }, + ) + } + + @Test + @DisplayName("Should describe temporal types by their ISO-8601 string format") + fun describesTemporalTypes() { + assertAll( + { assertEquals(nullable(mapOf("type" to "string", "format" to "date-time")), schemaOf("at")) }, + { assertEquals(nullable(mapOf("type" to "string", "format" to "date-time")), schemaOf("instant")) }, + { assertEquals(nullable(mapOf("type" to "string", "format" to "date")), schemaOf("day")) }, + { assertEquals(nullable(mapOf("type" to "string", "format" to "time")), schemaOf("time")) }, + { assertEquals(nullable(mapOf("type" to "string", "format" to "duration")), schemaOf("wait")) }, + ) + } + + @Test + @DisplayName("Should describe containers with the schema of what they hold") + fun describesContainers() { + assertAll( + { assertEquals(nullable(mapOf("type" to "array", "items" to nullable(STRING))), schemaOf("names")) }, + { assertEquals(nullable(mapOf("type" to "array", "items" to emptyMap())), schemaOf("anything")) }, + { + assertEquals( + nullable(mapOf("type" to "array", "items" to nullable(mapOf("type" to "number")))), + buildValueSchema(Declared::class.java.getMethod("acceptNumbers", List::class.java).genericParameterTypes[0]), + ) + }, + { assertEquals(nullable(mapOf("type" to "array", "items" to nullable(STRING))), schemaOf("words")) }, + { assertEquals(nullable(mapOf("type" to "array", "items" to INT64)), schemaOf("totals")) }, + { + assertEquals( + nullable(mapOf("type" to "object", "additionalProperties" to nullable(INT64))), + schemaOf("counts"), + ) + }, + { + assertEquals( + nullable(mapOf("type" to "array", "items" to emptyMap())), + buildValueSchema(java.util.List::class.java), + ) + }, + { + assertEquals( + nullable(mapOf("type" to "object", "additionalProperties" to true)), + buildValueSchema(java.util.Map::class.java), + ) + }, + { assertEquals(nullable(mapOf("type" to "object")), schemaOf("payload")) }, + ) + } + + @Test + @DisplayName("Should say nothing for a type that accepts more than one shape, or none") + fun describesNothingForOpenTypes() { + assertAll( + { assertEquals(null, schemaOf("bytes")) }, + { assertEquals(null, schemaOf("any")) }, + { assertEquals(null, schemaOf("date")) }, + { assertEquals(null, schemaOf("maybe")) }, + ) + } + + @Test + @DisplayName("Should list an enum's constants without initializing it") + fun listsEnumConstantsWithoutInitializing() { + assertEquals(nullable(mapOf("type" to "string", "enum" to listOf("ONLY"))), buildValueSchema(Lazy::class.java)) + assertFalse(EnumInitialization.ran) + } +} From 8d360b8b17cdbb8a4337a6831d48798f1c6f1960 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 16:14:50 +0000 Subject: [PATCH 75/94] Java SDK: stamp a cache digest on the bundle JAR manifest The Dag processor will skip re-probing a JAR whose cache digest it already knows (ADR-0014), so a Java bundle needs one that stays the same across rebuilds that change nothing and differs otherwise. The JAR's own hash does not: Gradle keeps file timestamps in archives by default. The Gradle plugin now writes Airflow-Cache-Digest on the bundle JAR, the shadow JAR or the thin jar. It is a SHA-256 over the main class, the packaging mode, the main source set's files, and every runtime dependency, which is read from each archive's central directory without decompressing anything. This is how the reproducible-build hash is realised without making the archive itself byte-reproducible. The thin jar task now takes the runtime classpath as an input, so a dependency change rebuilds it. --- java-sdk/plugin/build.gradle.kts | 6 + .../airflow/sdk/plugin/AirflowSdkPlugin.kt | 43 +++++-- .../apache/airflow/sdk/plugin/BundleDigest.kt | 113 ++++++++++++++++++ .../airflow/sdk/plugin/BundleDigestTest.kt | 98 +++++++++++++++ 4 files changed, 251 insertions(+), 9 deletions(-) create mode 100644 java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/BundleDigest.kt create mode 100644 java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/BundleDigestTest.kt diff --git a/java-sdk/plugin/build.gradle.kts b/java-sdk/plugin/build.gradle.kts index edf930bb6fc84..7efcdc36e612c 100644 --- a/java-sdk/plugin/build.gradle.kts +++ b/java-sdk/plugin/build.gradle.kts @@ -25,6 +25,12 @@ plugins { dependencies { implementation("com.gradleup.shadow:shadow-gradle-plugin:9.1.0") // Last supporting Java 11. + + testImplementation(kotlin("test")) +} + +tasks.withType { + useJUnitPlatform() } gradlePlugin { diff --git a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt index 3756ef688e174..817f7c0a32176 100644 --- a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt +++ b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt @@ -23,6 +23,7 @@ import com.github.jengelman.gradle.plugins.shadow.tasks.ShadowJar import org.gradle.api.Plugin import org.gradle.api.Project import org.gradle.api.provider.Property +import org.gradle.api.tasks.ClasspathNormalizer import org.gradle.api.tasks.Copy import org.gradle.api.tasks.Input import org.gradle.api.tasks.SourceSetContainer @@ -98,6 +99,10 @@ abstract class AirflowBundleExtension { * directory instead. In this mode, `Airflow-Supervisor-Schema-Version` lives in * the `airflow-sdk` JAR instead. The bundle JAR still contains `Main-Class`. * + * In both modes the bundle JAR also carries `Airflow-Cache-Digest`, a SHA-256 + * over the bundle's classes, resources, and runtime dependencies. It stays the + * same when a rebuild changes none of them, so Airflow can tell that a JAR it + * has already inspected is unchanged. */ class AirflowSdkPlugin : Plugin { override fun apply(project: Project) { @@ -115,15 +120,25 @@ class AirflowSdkPlugin : Plugin { } } - val classFiles = - project.objects.fileCollection().from( - project.extensions - .getByType(SourceSetContainer::class.java) - .getByName("main") - .output - .classesDirs, - project.configurations.getByName("runtimeClasspath"), - ) + val mainOutput = + project.extensions + .getByType(SourceSetContainer::class.java) + .getByName("main") + .output + val runtimeClasspath = project.configurations.getByName("runtimeClasspath") + val classFiles = project.objects.fileCollection().from(mainOutput.classesDirs, runtimeClasspath) + + fun stampCacheDigest(task: Jar) { + // Plain file collections and values only, which the configuration cache can store. + val outputFiles = project.objects.fileCollection().from(mainOutput) + val classpathFiles = project.objects.fileCollection().from(runtimeClasspath) + val mainClass = ext.mainClass + val fatJar = ext.fatJar.get() + task.doFirst { + val digest = BundleDigest.compute(mainClass.orNull, fatJar, outputFiles, classpathFiles) + task.manifest.attributes(mapOf(BundleDigest.ATTRIBUTE to digest)) + } + } val verifyTask = project.tasks.register("verifyBundleMainClass") { task -> @@ -195,6 +210,7 @@ class AirflowSdkPlugin : Plugin { mapOf("Airflow-Supervisor-Schema-Version" to schemaVersionProvider.get()), ) } + stampCacheDigest(task) } project.tasks.register("bundle", Copy::class.java) { task -> @@ -205,6 +221,15 @@ class AirflowSdkPlugin : Plugin { task.into(project.layout.buildDirectory.dir("bundle")) } } else { + project.tasks.named("jar", Jar::class.java).configure { task -> + // The digest covers the dependencies copied next to the thin JAR, so + // a dependency change has to rebuild it. + task.inputs + .files(runtimeClasspath) + .withPropertyName("airflowCacheDigestClasspath") + .withNormalizer(ClasspathNormalizer::class.java) + stampCacheDigest(task) + } // bundle copies the thin JAR and all runtime dependency JARs into // build/bundle/, mirroring what installDist puts in lib/. project.tasks.register("bundle", Copy::class.java) { task -> diff --git a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/BundleDigest.kt b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/BundleDigest.kt new file mode 100644 index 0000000000000..e7c09decb90da --- /dev/null +++ b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/BundleDigest.kt @@ -0,0 +1,113 @@ +/* + * 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. + */ + +package org.apache.airflow.sdk.plugin + +import java.io.File +import java.nio.ByteBuffer +import java.security.MessageDigest +import java.util.zip.ZipException +import java.util.zip.ZipFile + +/** + * The `Airflow-Cache-Digest` a bundle JAR carries: a SHA-256 over what the + * bundle runs, which Airflow compares to tell whether a JAR changed since it + * last asked it for its task handlers. + * + * It hashes content, never archive metadata, so a rebuild that changes + * nothing gives the same value although the JAR's bytes differ (Gradle keeps + * file timestamps in archives by default). It covers the main class, the + * packaging mode, every file of the main source set's output, and every + * runtime dependency. A dependency JAR contributes its entries' names, sizes + * and CRC-32 values from the central directory, so no entry is decompressed. + */ +internal object BundleDigest { + const val ATTRIBUTE = "Airflow-Cache-Digest" + + fun compute( + mainClass: String?, + fatJar: Boolean, + outputs: Iterable, + classpath: Iterable, + ): String { + val writer = Writer() + writer.write("airflow-cache-digest", 1, mainClass.orEmpty(), if (fatJar) "fat" else "thin") + writer.writeTree(outputs) + for (entry in classpath) { + when { + entry.isDirectory -> writer.writeTree(listOf(entry)) + entry.isFile -> writer.writeArchive(entry) + } + } + return writer.hex() + } + + private class Writer { + private val digest = MessageDigest.getInstance("SHA-256") + + // Each token is length-prefixed, so no two token sequences hash alike. + fun write(vararg tokens: Any) { + for (token in tokens) { + val bytes = token.toString().toByteArray(Charsets.UTF_8) + digest.update(ByteBuffer.allocate(Int.SIZE_BYTES).putInt(bytes.size).array()) + digest.update(bytes) + } + } + + fun writeTree(roots: Iterable) { + roots + .filter { it.isDirectory } + .flatMap { root -> + root.walkTopDown().filter { it.isFile }.map { root.toPath().relativize(it.toPath()).joinToString("/") to it } + }.sortedBy { it.first } + .forEach { (path, file) -> write("file", path, sha256(file)) } + } + + fun writeArchive(file: File) { + try { + ZipFile(file).use { zip -> + write("archive", file.name) + zip + .entries() + .asSequence() + .filterNot { it.isDirectory } + .sortedBy { it.name } + .forEach { write("entry", it.name, it.size, it.crc) } + } + } catch (_: ZipException) { + write("file", file.name, sha256(file)) + } + } + + fun hex(): String = digest.digest().joinToString("") { "%02x".format(it) } + + private fun sha256(file: File): String { + val fileDigest = MessageDigest.getInstance("SHA-256") + file.inputStream().use { input -> + val buffer = ByteArray(DEFAULT_BUFFER_SIZE) + while (true) { + val read = input.read(buffer) + if (read < 0) break + fileDigest.update(buffer, 0, read) + } + } + return fileDigest.digest().joinToString("") { "%02x".format(it) } + } + } +} diff --git a/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/BundleDigestTest.kt b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/BundleDigestTest.kt new file mode 100644 index 0000000000000..272b5eac1cb90 --- /dev/null +++ b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/BundleDigestTest.kt @@ -0,0 +1,98 @@ +/* + * 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. + */ + +package org.apache.airflow.sdk.plugin + +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertNotEquals +import org.junit.jupiter.api.Assertions.assertTrue +import org.junit.jupiter.api.DisplayName +import org.junit.jupiter.api.Test +import org.junit.jupiter.api.io.TempDir +import java.io.File +import java.util.zip.ZipEntry +import java.util.zip.ZipOutputStream + +class BundleDigestTest { + @TempDir + lateinit var tmp: File + + private fun classesDir( + name: String, + files: Map, + ): File = + tmp.resolve(name).also { dir -> + files.forEach { (path, content) -> dir.resolve(path).apply { parentFile.mkdirs() }.writeText(content) } + } + + private fun jar( + name: String, + entries: List>, + time: Long = 0L, + ): File = + tmp.resolve(name).also { file -> + file.parentFile.mkdirs() + ZipOutputStream(file.outputStream()).use { zip -> + entries.forEach { (path, content) -> + zip.putNextEntry(ZipEntry(path).apply { this.time = time }) + zip.write(content.toByteArray()) + zip.closeEntry() + } + } + } + + private val classes = mapOf("com/example/Main.class" to "main", "com/example/Handler.class" to "handler") + private val depEntries = listOf("dep/A.class" to "a", "dep/B.class" to "b") + + private fun digest( + mainClass: String = "com.example.Main", + fatJar: Boolean = true, + outputs: List = listOf(classesDir("classes", classes)), + classpath: List = listOf(jar("lib/dep.jar", depEntries)), + ) = BundleDigest.compute(mainClass, fatJar, outputs, classpath) + + @Test + @DisplayName("Should give the same SHA-256 for the same content") + fun sameContentSameDigest() { + val first = digest() + assertTrue(Regex("[0-9a-f]{64}").matches(first)) + assertEquals(first, digest()) + } + + @Test + @DisplayName("Should ignore dependency archive timestamps and entry order") + fun ignoresArchiveMetadata() { + val first = digest() + val rezipped = jar("lib/dep.jar", depEntries.reversed(), time = 1_700_000_000_000L) + + assertEquals(first, digest(classpath = listOf(rezipped))) + } + + @Test + @DisplayName("Should change when anything the bundle runs changes") + fun changesWithContent() { + val first = digest() + + assertNotEquals(first, digest(outputs = listOf(classesDir("changed", classes + ("com/example/Handler.class" to "HANDLER"))))) + assertNotEquals(first, digest(classpath = listOf(jar("lib/dep.jar", depEntries), jar("lib/extra.jar", depEntries)))) + assertNotEquals(first, digest(classpath = listOf(jar("lib/dep.jar", listOf("dep/A.class" to "a", "dep/B.class" to "bb"))))) + assertNotEquals(first, digest(mainClass = "com.example.Other")) + assertNotEquals(first, digest(fatJar = false)) + } +} From e4adc58a92a68a4321336c9d3da9101058e8f832 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 16:29:46 +0000 Subject: [PATCH 76/94] Java SDK: register stub Dags' Java tasks as task handlers The example, the e2e test bundle, the Scala Spark example and the k8s example registered the Java tasks of Dags that a Python file owns as a Dag declared in Java. ADR-0011 keeps a Python-owned Dag id out of Java Dag registrations, the parse probe reports only task-handler registrations, and native Dag parsing would treat such a DagDef as a second definition of the same Dag. They now register those tasks as task handlers. Execution already falls back to task handlers, so the runs do not change. --- .../apache/airflow/e2e/TestBundleBuilder.java | 21 +++++++------------ .../airflow/example/ExampleBundleBuilder.java | 3 +-- .../example/InterfaceExampleBuilder.java | 12 ++++++----- .../airflow/example/ScalaSparkExample.scala | 21 ++++++++----------- .../airflow/k8sexample/CombinedExample.java | 12 +++++------ .../airflow/k8sexample/K8sBundleBuilder.java | 15 ++----------- 6 files changed, 32 insertions(+), 52 deletions(-) diff --git a/airflow-e2e-tests/java-test-bundle/src/java/org/apache/airflow/e2e/TestBundleBuilder.java b/airflow-e2e-tests/java-test-bundle/src/java/org/apache/airflow/e2e/TestBundleBuilder.java index ad33a15708a46..1c15b425932df 100644 --- a/airflow-e2e-tests/java-test-bundle/src/java/org/apache/airflow/e2e/TestBundleBuilder.java +++ b/airflow-e2e-tests/java-test-bundle/src/java/org/apache/airflow/e2e/TestBundleBuilder.java @@ -19,7 +19,6 @@ package org.apache.airflow.e2e; -import java.util.List; import org.apache.airflow.sdk.*; import org.jetbrains.annotations.NotNull; @@ -27,7 +26,7 @@ * Bundle for the runner-behaviour E2E tests: deliberately broken task classes that exercise * instantiation failures, and a task that round-trips Airflow Variables through the supervisor. */ -public class TestBundleBuilder implements BundleBuilder { +public class TestBundleBuilder { public static class MissingNoArgConstructor implements Task { public MissingNoArgConstructor(String unused) {} @@ -60,19 +59,15 @@ public void execute(@NotNull Context context, Client client) { } } - @NotNull - @Override - public Iterable getDags() { - var uninstantiable = new DagDef("java_uninstantiable"); - uninstantiable.addTask("missing_no_arg_constructor", MissingNoArgConstructor.class); - uninstantiable.addTask("non_static_inner", NonStaticInner.class); - var variableWrite = new DagDef("java_variable_write"); - variableWrite.addTask("write_and_delete", WriteAndDeleteVariable.class); - return List.of(uninstantiable, variableWrite); + public static Bundle build() { + return new Bundle() + .register( + "java_uninstantiable", "missing_no_arg_constructor", MissingNoArgConstructor.class) + .register("java_uninstantiable", "non_static_inner", NonStaticInner.class) + .register("java_variable_write", "write_and_delete", WriteAndDeleteVariable.class); } public static void main(String[] args) { - var bundle = new TestBundleBuilder().build(); - Server.create(args).serve(bundle); + Server.create(args).serve(build()); } } diff --git a/java-sdk/example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java b/java-sdk/example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java index d75f5762ba12d..7ce2fd4de9200 100644 --- a/java-sdk/example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java +++ b/java-sdk/example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java @@ -23,8 +23,7 @@ public class ExampleBundleBuilder { public static Bundle build() { - return new Bundle() - .register(InterfaceExampleBuilder.build()) + return InterfaceExampleBuilder.registerInto(new Bundle()) .register(AnnotationExample.class) .register(XComCastingExample.class) .register(org.apache.airflow.example.nativedag.InterfaceExample.build()); diff --git a/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java b/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java index 4d0a494d6dcaa..0a07a20ccdc22 100644 --- a/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java +++ b/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java @@ -89,10 +89,12 @@ public void execute(@NotNull Context context, Client client, SummarizeInput inpu } } - public static DagDef build() { - return new DagDef("java_interface_example") - .addTask("extract", Extract.class) - .addTask("transform", Transform.class) - .addTask("summarize", Summarize.class); + // The Python Dag file owns java_interface_example, so its tasks register as + // task handlers rather than as a Dag declared in Java. + public static Bundle registerInto(Bundle bundle) { + return bundle + .register("java_interface_example", "extract", Extract.class) + .register("java_interface_example", "transform", Transform.class) + .register("java_interface_example", "summarize", Summarize.class); } } diff --git a/java-sdk/scala_spark_example/src/main/scala/org/apache/airflow/example/ScalaSparkExample.scala b/java-sdk/scala_spark_example/src/main/scala/org/apache/airflow/example/ScalaSparkExample.scala index f7f32bad8e296..8c6b3cf28b277 100644 --- a/java-sdk/scala_spark_example/src/main/scala/org/apache/airflow/example/ScalaSparkExample.scala +++ b/java-sdk/scala_spark_example/src/main/scala/org/apache/airflow/example/ScalaSparkExample.scala @@ -19,7 +19,7 @@ package org.apache.airflow.example -import org.apache.airflow.sdk.{Bundle, BundleBuilder, Client, Context, DagDef, Server, Task} +import org.apache.airflow.sdk.{Bundle, Client, Context, Server, Task} import org.apache.logging.log4j.{LogManager, Logger} import org.apache.spark.sql.SparkSession import org.apache.spark.sql.functions.sum @@ -132,18 +132,15 @@ class SparkLoad extends Task { } } -object ScalaSparkExample { - def build(): DagDef = - new DagDef(SparkEtl.DagId) - .addTask(SparkEtl.ExtractTaskId, classOf[SparkExtract]) - .addTask(SparkEtl.TransformTaskId, classOf[SparkTransform]) - .addTask(SparkEtl.LoadTaskId, classOf[SparkLoad]) -} - /** Bundle entry point served to Airflow's Java coordinator. */ -object ScalaSparkBundleBuilder extends BundleBuilder { - override def getDags(): java.lang.Iterable[DagDef] = java.util.List.of(ScalaSparkExample.build()) +object ScalaSparkBundleBuilder { + // The Python Dag file owns the Dag, so its tasks register as task handlers. + def build(): Bundle = + new Bundle() + .register(SparkEtl.DagId, SparkEtl.ExtractTaskId, classOf[SparkExtract]) + .register(SparkEtl.DagId, SparkEtl.TransformTaskId, classOf[SparkTransform]) + .register(SparkEtl.DagId, SparkEtl.LoadTaskId, classOf[SparkLoad]) def main(args: Array[String]): Unit = - Server.create(args).serve(new Bundle(getDags())) + Server.create(args).serve(build()) } diff --git a/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/CombinedExample.java b/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/CombinedExample.java index 63aaa5b6f051d..9e44edca50cfa 100644 --- a/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/CombinedExample.java +++ b/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/CombinedExample.java @@ -25,21 +25,19 @@ import org.slf4j.LoggerFactory; // Java half of the KubernetesExecutor lang-SDK system test bundle. Registers the -// Java tasks of the shared "lang_sdk_combined" Dag; the Go tasks of the same -// dag_id live in ../go_example and the Python stub Dag in ../dags. The -// coordinator locates this jar by dag_id, so only the Java tasks are registered -// here. -@Builder.Dag(id = "lang_sdk_combined") +// Java task handlers of the shared "lang_sdk_combined" Dag; the Go tasks of the +// same dag_id live in ../go_example and the Python stub Dag in ../dags. The +// Python file owns the Dag, so only the Java tasks are registered here. public class CombinedExample { private static final Logger logger = LoggerFactory.getLogger(CombinedExample.class); - @Builder.Task(id = "java_extract") + @Builder.TaskHandler(dag = "lang_sdk_combined", task = "java_extract") public long extract(Client client) { logger.info("java_extract running"); return new Date().getTime(); } - @Builder.Task(id = "java_transform") + @Builder.TaskHandler(dag = "lang_sdk_combined", task = "java_transform") public void transform(Client client) { var variable = client.getVariable("my_variable"); logger.info("java_transform obtained variable {}", variable); diff --git a/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/K8sBundleBuilder.java b/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/K8sBundleBuilder.java index c3b511b898ced..064dfebd3566f 100644 --- a/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/K8sBundleBuilder.java +++ b/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/K8sBundleBuilder.java @@ -19,21 +19,10 @@ package org.apache.airflow.k8sexample; -import java.util.List; import org.apache.airflow.sdk.*; -import org.jetbrains.annotations.NotNull; - -// CombinedExampleBuilder is generated by the annotation processor from the -// @Builder.Dag class CombinedExample. -public class K8sBundleBuilder implements BundleBuilder { - @NotNull - @Override - public Iterable getDags() { - return List.of(CombinedExampleBuilder.build()); - } +public class K8sBundleBuilder { public static void main(String[] args) { - var bundle = new K8sBundleBuilder().build(); - Server.create(args).serve(bundle); + Server.create(args).serve(new Bundle().register(CombinedExample.class)); } } From f72d37f25bfadcda53f6366ee386858833e09cde Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 16:30:40 +0000 Subject: [PATCH 77/94] Let the Java coordinator probe a JAR's task handlers The Dag-file parse will ask each candidate JAR which task handlers it registers. JavaCoordinator now builds that command: the probed JAR's own Main-Class with the whole bundle on the classpath, so the probe never picks among executable JARs (#71134). The schema version comes from the JAR, or for a thin JAR from the airflow-sdk JAR beside it, as for execution. Manifests are now parsed as the JAR spec defines them: a value longer than a 72-byte line continues on the next line, which the email parser kept folded, so a long Main-Class came out broken. The real-JAR test builds the example as a thin bundle and probes it through a JVM. It runs only with AIRFLOW_LANG_SDK_REAL_PROBE_TESTS=1 and a JDK. --- .../test_task_handler_processor_java.py | 212 ++++++++++++++++++ .../sdk/coordinators/java/coordinator.py | 54 ++++- .../coordinators/java/test_coordinator.py | 81 +++++++ 3 files changed, 341 insertions(+), 6 deletions(-) create mode 100644 airflow-core/tests/unit/dag_processing/test_task_handler_processor_java.py diff --git a/airflow-core/tests/unit/dag_processing/test_task_handler_processor_java.py b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_java.py new file mode 100644 index 0000000000000..9dd27999e1c69 --- /dev/null +++ b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_java.py @@ -0,0 +1,212 @@ +# +# 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. +"""Probe a bundle built with the Java SDK's Gradle plugin through a real JVM.""" + +from __future__ import annotations + +import json +import os +import re +import shutil +import subprocess +import zipfile +from typing import TYPE_CHECKING +from unittest import mock + +import pytest +import structlog + +from airflow.dag_processing.processor import ( + TaskHandlerDeclaration, + TaskHandlerParam, + TaskHandlerParsingResult, +) +from airflow.dag_processing.task_handler_processor import LangSDKTaskHandlerProcessorProcess +from airflow.sdk.coordinators.java.coordinator import _parse_manifest +from airflow.sdk.execution_time import supervisor +from airflow.sdk.execution_time.coordinator import reset_coordinator_manager + +from tests_common.pytest_plugin import AIRFLOW_ROOT_PATH +from tests_common.test_utils.config import conf_vars + +if TYPE_CHECKING: + from collections.abc import Iterator + from pathlib import Path + +JAVA_SDK_PATH = AIRFLOW_ROOT_PATH / "java-sdk" + +# The example's sources, built as a thin bundle against the SDK as an included build, so nothing is +# published. The plugin's fat mode looks the SDK up as a published airflow-sdk artifact. +_SETTINGS = """ +pluginManagement { + includeBuild("../java-sdk") +} +includeBuild("../java-sdk") +rootProject.name = "probe-example" +""" + +_BUILD = """ +plugins { + id("org.apache.airflow.sdk") +} +repositories { + mavenCentral() +} +dependencies { + annotationProcessor("org.apache.airflow:processor") + implementation("org.apache.airflow:sdk") + implementation("org.apache.airflow:jpl") +} +sourceSets { + main { + java.srcDir("../java-sdk/example/src/java") + } +} +airflowBundle { + mainClass = "org.apache.airflow.example.ExampleBundleBuilder" + fatJar = false +} +""" + +pytestmark = pytest.mark.skipif( + os.environ.get("AIRFLOW_LANG_SDK_REAL_PROBE_TESTS") != "1", + reason="set AIRFLOW_LANG_SDK_REAL_PROBE_TESTS=1 to probe real Lang-SDK artifacts", +) + + +@pytest.fixture(scope="module") +def example_bundle(tmp_path_factory) -> Path: + """Build ``java-sdk/example`` from a copy of the SDK sources, so the checkout gets no build output.""" + if not JAVA_SDK_PATH.is_dir(): + pytest.skip("the Java SDK sources are absent") + if shutil.which("java") is None or shutil.which("javac") is None: + pytest.skip("needs a JDK") + root = tmp_path_factory.mktemp("java-sdk") + sdk = root / "java-sdk" + shutil.copytree(JAVA_SDK_PATH, sdk, ignore=shutil.ignore_patterns("build", ".gradle", ".kotlin")) + project = root / "probe-example" + project.mkdir() + (project / "settings.gradle").write_text(_SETTINGS) + (project / "build.gradle").write_text(_BUILD) + subprocess.run([sdk / "gradlew", "--no-daemon", "--quiet", "bundle"], cwd=project, check=True) + return project / "build" / "bundle" + + +@pytest.fixture +def java_coordinator() -> Iterator[None]: + spec = {"java": {"classpath": "airflow.sdk.coordinators.java.JavaCoordinator"}} + reset_coordinator_manager() + try: + # The parse child is a bare fork even on macOS, so it sees this config. + with ( + conf_vars({("sdk", "coordinators"): json.dumps(spec)}), + mock.patch.object(supervisor, "_should_use_exec", return_value=False), + ): + yield + finally: + reset_coordinator_manager() + + +def _nullable(schema: dict) -> dict: + return {"anyOf": [schema, {"type": "null"}]} + + +INT32 = {"type": "integer", "format": "int32"} +INT64 = {"type": "integer", "format": "int64"} +FLOAT = {"type": "number", "format": "float"} +DOUBLE = {"type": "number", "format": "double"} +STRING = {"type": "string"} + + +def _positional(task_id: str, *params: tuple[str, dict]) -> TaskHandlerDeclaration: + return TaskHandlerDeclaration( + task_id=task_id, + binding="positional", + params=[TaskHandlerParam(name=name, value_schema=schema, required=True) for name, schema in params], + ) + + +def _named(task_id: str, *params: tuple[str, dict, bool]) -> TaskHandlerDeclaration: + return TaskHandlerDeclaration( + task_id=task_id, + binding="named", + params=[ + TaskHandlerParam(name=name, value_schema=schema, required=False, exact_name=exact) + for name, schema, exact in params + ], + ) + + +@pytest.mark.usefixtures("java_coordinator") +def test_a_built_bundle_declares_its_task_handlers(example_bundle): + jar = example_bundle / "probe-example.jar" + + result = LangSDKTaskHandlerProcessorProcess.run( + coordinator="java", + path=jar, + bundle_path=example_bundle, + bundle_name="java-task-handlers", + artifact_rel_path=jar.name, + dag_ids=[ + "java_xcom_casting_example", + "java_interface_example", + "java_annotation_example", + "java_native_interface_example", + "not_registered", + ], + logger=structlog.get_logger(), + ) + + assert result == TaskHandlerParsingResult( + fileloc=os.fspath(jar), + task_handlers={ + "java_xcom_casting_example": [ + _named("produce_number"), + _positional("widen_to_long", ("value", INT64)), + _positional("widen_to_double", ("value", DOUBLE)), + _named("produce_nothing"), + _positional("consume_nullable", ("value", _nullable(INT32))), + _named("produce_fraction"), + _positional("consume_float", ("value", FLOAT)), + _positional( + "consume_double_list", + ("values", _nullable({"type": "array", "items": _nullable(DOUBLE)})), + ), + ], + "java_interface_example": [ + _named("extract"), + _named("transform", ("extracted", INT64, False)), + _named("summarize", ("region_code", _nullable(STRING), True), ("transformed", INT64, False)), + ], + "java_annotation_example": [ + _named("extract"), + _positional("transform", ("extracted", INT64)), + _positional("load", ("transformed", INT64)), + _named("report", ("runLabel", _nullable(STRING), False), ("transformed", INT64, False)), + _named("concurrent"), + ], + }, + ) + assert list(result.task_handlers) == [ + "java_xcom_casting_example", + "java_interface_example", + "java_annotation_example", + ] + with zipfile.ZipFile(jar) as zf: + manifest = _parse_manifest(zf.read("META-INF/MANIFEST.MF")) + assert re.fullmatch(r"[0-9a-f]{64}", manifest["airflow-cache-digest"]) diff --git a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py index 5d483110580f2..168ed2691d144 100644 --- a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py @@ -19,9 +19,9 @@ from __future__ import annotations -import email import os import pathlib +import re import stat import zipfile from typing import TYPE_CHECKING @@ -87,10 +87,33 @@ def _calculate_classpath(roots: Sequence[pathlib.Path]) -> str: return os.pathsep.join(sorted(jars)) # Keep output deterministic. +def _parse_manifest(data: bytes) -> dict[str, str]: + """ + Return the main-section attributes of a JAR manifest, keyed by lowercased name. + + A line holds at most 72 bytes, so a longer value continues on lines starting with one space, + which are joined without it. + """ + attributes: dict[str, str] = {} + name: str | None = None + for line in re.split(r"\r\n|\r|\n", data.decode("utf-8", errors="replace")): + if not line: + break # The main section ends at the first blank line. + if line.startswith(" "): + if name is not None: + attributes[name] += line[1:] + continue + key, sep, value = line.partition(":") + name = key.lower() if sep else None + if name is not None: + attributes[name] = value.removeprefix(" ") + return attributes + + @attrs.define class _JarMetadata: - main_class: str - schema_version: str + main_class: str | None + schema_version: str | None @classmethod def from_jar(cls, path: pathlib.Path) -> Self | None: @@ -101,9 +124,8 @@ def from_jar(cls, path: pathlib.Path) -> Self | None: except KeyError: log.debug("JAR does not contain META-INF/MANIFEST.MF; ignored", path=path) return None - with zf.open(manifest_info) as f: - manifest = email.message_from_binary_file(f) - return cls(manifest["Main-Class"], manifest["Airflow-Supervisor-Schema-Version"]) + manifest = _parse_manifest(zf.read(manifest_info)) + return cls(manifest.get("main-class"), manifest.get("airflow-supervisor-schema-version")) except zipfile.BadZipFile: log.exception("Cannot read JAR; ignored", path=path) return None @@ -189,6 +211,10 @@ class JavaCoordinator(SubprocessCoordinator): up being executed, so set *main_class* when more than one JAR in the bundle declares Main-Class. + To report the task handlers a JAR registers, the coordinator runs that JAR's + own Main-Class with the whole bundle on the classpath. *main_class* applies + to task execution only. + A JAR containing metadata *Airflow-Supervisor-Schema-Version* should also be available to specify the wire schema version. The JAR containing the Java SDK automatically sets this, so you don't generally need to do anything if @@ -220,3 +246,19 @@ def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], jar.main_class, ] return command, jar.schema_version + + def _build_parse_task_handler_command(self, *, path: pathlib.Path) -> tuple[list[str], str | None]: + metadata = _JarMetadata.from_jar(path) + if metadata is None or not metadata.main_class: + raise ValueError(f"{path} is not an executable JAR: its manifest sets no Main-Class") + roots = self._get_scan_roots() + # A thin JAR leaves the version to the airflow-sdk JAR beside it, where execution finds it too. + schema_version = metadata.schema_version or _JarInfo.find(roots, metadata.main_class).schema_version + command = [ + self.java_executable, + "-classpath", + _calculate_classpath(roots), + *self.jvm_args, + metadata.main_class, + ] + return command, schema_version 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 56e21d970e612..e41b7afec9f14 100644 --- a/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py @@ -35,6 +35,7 @@ JavaCoordinator, _calculate_classpath, _JarInfo, + _parse_manifest, _walk_jars, ) from airflow.sdk.execution_time.coordinator import BaseCoordinator @@ -251,6 +252,86 @@ def test_build_command_scans_passed_roots_in_colocated_mode(self, tmp_path): assert schema_version == "2026-06-16" +def test_parse_manifest_joins_continued_values_and_stops_at_the_main_section(): + manifest = ( + b"Manifest-Version: 1.0\r\n" + b"Airflow-Cache-Digest: 6a7cf9d4ddd454e97952b07038a8ebc6ff37c3a52ae176e921\r\n" + b" bc03065d6404cd\r\n" + b"Main-Class: com.example.Main\r\n" + b"\r\n" + b"Name: com/example/\r\n" + b"Sealed: true\r\n" + ) + + assert _parse_manifest(manifest) == { + "manifest-version": "1.0", + "airflow-cache-digest": "6a7cf9d4ddd454e97952b07038a8ebc6ff37c3a52ae176e921bc03065d6404cd", + "main-class": "com.example.Main", + } + + +class TestJavaCoordinatorParseTaskHandlerCommand: + @pytest.fixture + def bundle(self, tmp_path): + _make_jar(tmp_path / "app.jar", main_class="com.example.App", schema_version="2026-06-16") + _make_jar(tmp_path / "other.jar", main_class="com.example.Other", schema_version="2026-06-16") + (tmp_path / "lib").mkdir() + _make_jar(tmp_path / "lib" / "dep.jar", main_class=None) + return tmp_path + + @pytest.mark.parametrize( + ("jar", "main_class"), [("app.jar", "com.example.App"), ("other.jar", "com.example.Other")] + ) + def test_runs_the_probed_jars_own_main_class_with_the_bundle_on_the_classpath( + self, bundle, jar, main_class + ): + coordinator = JavaCoordinator(jvm_args=["-Xmx1g"], main_class="com.example.Configured") + with coordinator._set_scan_roots([bundle]): + command, schema_version = coordinator._build_parse_task_handler_command(path=bundle / jar) + + classpath = os.pathsep.join( + (bundle / name).as_posix() for name in ("app.jar", "lib/dep.jar", "other.jar") + ) + assert command == ["java", "-classpath", classpath, "-Xmx1g", main_class] + assert schema_version == "2026-06-16" + + def test_a_thin_jar_takes_the_schema_version_from_the_sdk_jar(self, tmp_path): + _make_jar(tmp_path / "app.jar", main_class="com.example.App") + _make_jar(tmp_path / "airflow-sdk.jar", main_class=None, schema_version="2026-10-30") + coordinator = JavaCoordinator() + with coordinator._set_scan_roots([tmp_path]): + command, schema_version = coordinator._build_parse_task_handler_command(path=tmp_path / "app.jar") + + assert command[-1] == "com.example.App" + assert schema_version == "2026-10-30" + + def test_reads_a_main_class_continued_on_the_next_manifest_line(self, tmp_path): + main_class = "org.example.a.very.long.package.name.that.needs.folding.BundleMain" + manifest = f"Manifest-Version: 1.0\r\nMain-Class: {main_class[:50]}\r\n {main_class[50:]}\r\n" + manifest += "Airflow-Supervisor-Schema-Version: 2026-06-16\r\n\r\n" + with zipfile.ZipFile(tmp_path / "app.jar", "w") as zf: + zf.writestr("META-INF/MANIFEST.MF", manifest) + coordinator = JavaCoordinator() + with coordinator._set_scan_roots([tmp_path]): + command, _ = coordinator._build_parse_task_handler_command(path=tmp_path / "app.jar") + + assert command[-1] == main_class + + @pytest.mark.parametrize("content", ["no-main-class", "not-a-zip"]) + def test_a_jar_that_is_not_executable_is_rejected(self, tmp_path, content): + path = tmp_path / "app.jar" + if content == "no-main-class": + _make_jar(path, main_class=None, schema_version="2026-06-16") + else: + path.write_text("not a zip") + coordinator = JavaCoordinator() + with ( + coordinator._set_scan_roots([tmp_path]), + pytest.raises(ValueError, match="is not an executable JAR"), + ): + coordinator._build_parse_task_handler_command(path=path) + + @pytest.fixture def jars_dir(tmp_path): _make_jar(tmp_path.joinpath("app.jar"), main_class="com.example.TaskRunner", schema_version="2026-06-16") From 11fb9964dd1ce7980301629b9996cdc2032aff19 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 17:05:45 +0000 Subject: [PATCH 78/94] Java SDK: list a value schema's enum constants in a stable order Class.getFields has no defined order, so the enum list in a parameter's value_schema could differ between JVMs. Declarations are stored and compared across parses, so the names are now sorted. --- .../org/apache/airflow/sdk/internal/ValueSchema.kt | 10 ++++++++-- .../org/apache/airflow/sdk/internal/ValueSchemaTest.kt | 2 +- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/ValueSchema.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/ValueSchema.kt index d3470438015ad..1ffd087998271 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/ValueSchema.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/ValueSchema.kt @@ -125,6 +125,7 @@ private fun classSchema(type: Class<*>): Schema? { // Other platform types (Object, Optional, Date, URI, ...) decode from more // than one JSON shape, or from none, so they constrain nothing here. isPlatformType(type) -> null + // A POJO's fields are not described: the schema says only that it takes an object. else -> mapOf("type" to "object") } } @@ -133,11 +134,16 @@ private fun isPlatformType(type: Class<*>): Boolean = type.name.startsWith("java.") || type.name.startsWith("javax.") || type.name.startsWith("kotlin.") // Read from the constant fields rather than enumConstants, which would run the -// enum's static initializer while the bundle is only being described. +// enum's static initializer while the bundle is only being described. getFields +// has no defined order, so the names are sorted. private fun enumSchema(type: Class<*>): Schema = mapOf( "type" to "string", - "enum" to type.fields.filter { it.isEnumConstant }.map { it.name }, + "enum" to + type.fields + .filter { it.isEnumConstant } + .map { it.name } + .sorted(), ) private fun arraySchema(component: Type): Schema = mapOf("type" to "array", "items" to (buildValueSchema(component) ?: UNCONSTRAINED)) diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/ValueSchemaTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/ValueSchemaTest.kt index af7bc525f9b0e..7f82c7cddb101 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/ValueSchemaTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/ValueSchemaTest.kt @@ -147,7 +147,7 @@ class ValueSchemaTest { { assertEquals(nullable(mapOf("type" to "integer")), schemaOf("big")) }, { assertEquals(nullable(mapOf("type" to "number")), schemaOf("exact")) }, { assertEquals(nullable(mapOf("type" to "string", "format" to "uuid")), schemaOf("id")) }, - { assertEquals(nullable(mapOf("type" to "string", "enum" to listOf("RED", "GREEN"))), schemaOf("color")) }, + { assertEquals(nullable(mapOf("type" to "string", "enum" to listOf("GREEN", "RED"))), schemaOf("color")) }, ) } From e4c081cb7189704f4dd685ce4763b3d1f396deaf Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 17:08:42 +0000 Subject: [PATCH 79/94] Java SDK: check the task handler reply against the supervisor schema The runtime builds the TaskHandlerParsingResult body as a map, and the only test that decodes it on the Python side needs a JVM and is opt-in. A renamed or missing key would pass the Kotlin assertions and fail only at parse time. The vendored supervisor schema is generated from the Python models, so the reply is now checked against its declared and required fields and its binding values in the Java SDK's own tests. --- .../sdk/execution/TaskHandlerParseTest.kt | 49 +++++++++++++++++++ 1 file changed, 49 insertions(+) diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParseTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParseTest.kt index 49a769ac9e8bb..fa1c167bc8c8e 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParseTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParseTest.kt @@ -19,6 +19,8 @@ package org.apache.airflow.sdk.execution +import com.fasterxml.jackson.databind.JsonNode +import com.fasterxml.jackson.databind.ObjectMapper import org.apache.airflow.sdk.ArgName import org.apache.airflow.sdk.Bundle import org.apache.airflow.sdk.Client @@ -31,8 +33,10 @@ import org.apache.airflow.sdk.execution.comm.TaskHandlerParseRequest import org.apache.airflow.sdk.internal.TaskParams import org.apache.airflow.sdk.internal.TypeRef import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertTrue import org.junit.jupiter.api.DisplayName import org.junit.jupiter.api.Test +import java.io.File private val NULLABLE_STRING = mapOf("anyOf" to listOf(mapOf("type" to "string"), mapOf("type" to "null"))) @@ -126,6 +130,25 @@ private val SCORE_INPUT_PARAMS = param("threshold", mapOf("type" to "number", "format" to "double")), ) +// Generated from the Python models, so it is what the Dag processor decodes the reply with. +private val SUPERVISOR_SCHEMA: JsonNode = ObjectMapper().readTree(File("schema/schema.json")).path("${'$'}defs") + +private fun assertConforms( + definition: String, + body: Map<*, *>, +) { + val schema = SUPERVISOR_SCHEMA.path(definition) + val properties = + schema + .path("properties") + .fieldNames() + .asSequence() + .toSet() + val required = schema.path("required").map { it.asText() }.toSet() + assertTrue(properties.containsAll(body.keys), "$definition has no ${body.keys - properties}") + assertTrue(body.keys.containsAll(required), "$definition misses ${required - body.keys}") +} + class TaskHandlerParseTest { @Test @DisplayName("Should declare each requested Dag's task handlers in registration order") @@ -171,6 +194,32 @@ class TaskHandlerParseTest { assertEquals(expected, result) } + @Test + @DisplayName("Should send only the fields the supervisor schema declares, and every required one") + fun conformsToSupervisorSchema() { + val bundle = + Bundle() + .register("etl", "extract", GeneratedFlat::class.java) + .register("etl", "score", GeneratedInput::class.java) + .register("etl", "audit", Audit::class.java) + + val result = parseTaskHandlers(bundle, request("etl")) + + assertConforms("TaskHandlerParsingResult", result) + val bindings = + SUPERVISOR_SCHEMA + .path("TaskHandlerDeclaration") + .path("properties") + .path("binding") + .path("enum") + for (declaration in (result["task_handlers"] as Map<*, *>).values.flatMap { it as List<*> }) { + declaration as Map<*, *> + assertConforms("TaskHandlerDeclaration", declaration) + assertTrue(bindings.any { it.asText() == declaration["binding"] }, "unknown binding ${declaration["binding"]}") + (declaration["params"] as List<*>).forEach { assertConforms("TaskHandlerParam", it as Map<*, *>) } + } + } + @Test @DisplayName("Should send an empty map when no requested Dag has task handlers") fun sendsEmptyTaskHandlersWhenNothingMatches() { From 238799ed5d8673ea9593c3867b90ec1f7d4ad5a6 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 17:11:08 +0000 Subject: [PATCH 80/94] Say why the Java coordinator rejects a JAR it cannot probe An unreadable file or a JAR without a manifest was reported as having no Main-Class, which points the Dag author at the wrong fix. The two cases now have their own messages. Manifest values are now decoded after their continuation lines are joined, because the 72-byte fold can split a multi-byte character. Execution and the probe share one command builder, so the two cannot drift. --- .../sdk/coordinators/java/coordinator.py | 42 ++++++++----------- .../coordinators/java/test_coordinator.py | 22 ++++++++-- 2 files changed, 36 insertions(+), 28 deletions(-) diff --git a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py index 168ed2691d144..05cbad27b828d 100644 --- a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py @@ -92,22 +92,23 @@ def _parse_manifest(data: bytes) -> dict[str, str]: Return the main-section attributes of a JAR manifest, keyed by lowercased name. A line holds at most 72 bytes, so a longer value continues on lines starting with one space, - which are joined without it. + which are joined without it. The fold can split a multi-byte character, so values are decoded + only once joined. """ - attributes: dict[str, str] = {} + attributes: dict[str, bytes] = {} name: str | None = None - for line in re.split(r"\r\n|\r|\n", data.decode("utf-8", errors="replace")): + for line in re.split(rb"\r\n|\r|\n", data): if not line: break # The main section ends at the first blank line. - if line.startswith(" "): + if line.startswith(b" "): if name is not None: attributes[name] += line[1:] continue - key, sep, value = line.partition(":") - name = key.lower() if sep else None + key, sep, value = line.partition(b":") + name = key.decode("utf-8", errors="replace").lower() if sep else None if name is not None: - attributes[name] = value.removeprefix(" ") - return attributes + attributes[name] = value.removeprefix(b" ") + return {key: value.decode("utf-8", errors="replace") for key, value in attributes.items()} @attrs.define @@ -238,27 +239,18 @@ def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], # 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(roots), - *self.jvm_args, - jar.main_class, - ] - return command, jar.schema_version + return self._build_java_command(roots, jar.main_class), jar.schema_version def _build_parse_task_handler_command(self, *, path: pathlib.Path) -> tuple[list[str], str | None]: metadata = _JarMetadata.from_jar(path) - if metadata is None or not metadata.main_class: + if metadata is None: + raise ValueError(f"{path} is not an executable JAR: it has no readable manifest") + if not metadata.main_class: raise ValueError(f"{path} is not an executable JAR: its manifest sets no Main-Class") roots = self._get_scan_roots() # A thin JAR leaves the version to the airflow-sdk JAR beside it, where execution finds it too. schema_version = metadata.schema_version or _JarInfo.find(roots, metadata.main_class).schema_version - command = [ - self.java_executable, - "-classpath", - _calculate_classpath(roots), - *self.jvm_args, - metadata.main_class, - ] - return command, schema_version + return self._build_java_command(roots, metadata.main_class), schema_version + + def _build_java_command(self, roots: Sequence[pathlib.Path], main_class: str) -> list[str]: + return [self.java_executable, "-classpath", _calculate_classpath(roots), *self.jvm_args, main_class] 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 e41b7afec9f14..bae667d635b8e 100644 --- a/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py @@ -270,6 +270,12 @@ def test_parse_manifest_joins_continued_values_and_stops_at_the_main_section(): } +def test_parse_manifest_decodes_a_character_split_across_a_fold(): + manifest = b"Manifest-Version: 1.0\r\nImplementation-Vendor: Z\xc3\r\n \xbcrich\r\n\r\n" + + assert _parse_manifest(manifest)["implementation-vendor"] == "Z\u00fcrich" + + class TestJavaCoordinatorParseTaskHandlerCommand: @pytest.fixture def bundle(self, tmp_path): @@ -317,17 +323,27 @@ def test_reads_a_main_class_continued_on_the_next_manifest_line(self, tmp_path): assert command[-1] == main_class - @pytest.mark.parametrize("content", ["no-main-class", "not-a-zip"]) - def test_a_jar_that_is_not_executable_is_rejected(self, tmp_path, content): + @pytest.mark.parametrize( + ("content", "reason"), + [ + ("no-main-class", "its manifest sets no Main-Class"), + ("no-manifest", "it has no readable manifest"), + ("not-a-zip", "it has no readable manifest"), + ], + ) + def test_a_jar_that_is_not_executable_is_rejected(self, tmp_path, content, reason): path = tmp_path / "app.jar" if content == "no-main-class": _make_jar(path, main_class=None, schema_version="2026-06-16") + elif content == "no-manifest": + with zipfile.ZipFile(path, "w") as zf: + zf.writestr("com/example/App.class", b"") else: path.write_text("not a zip") coordinator = JavaCoordinator() with ( coordinator._set_scan_roots([tmp_path]), - pytest.raises(ValueError, match="is not an executable JAR"), + pytest.raises(ValueError, match=f"is not an executable JAR: {reason}$"), ): coordinator._build_parse_task_handler_command(path=path) From 5826da213a6e8ff9aecc8768addc39ec589b8811 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 1 Oct 2026 06:10:22 +0000 Subject: [PATCH 81/94] Test the Java SDK schema sync hook's snapshot copy The branch that copies the Task SDK snapshot had only a manual run behind it, and its version reader had a noun name that reads as an attribute. --- .../prek/sync_java_sdk_supervisor_schema.py | 6 +- .../test_sync_java_sdk_supervisor_schema.py | 90 +++++++++++++++++++ 2 files changed, 93 insertions(+), 3 deletions(-) create mode 100644 scripts/tests/ci/prek/test_sync_java_sdk_supervisor_schema.py diff --git a/scripts/ci/prek/sync_java_sdk_supervisor_schema.py b/scripts/ci/prek/sync_java_sdk_supervisor_schema.py index 22f6578de1495..c270920c09ac2 100755 --- a/scripts/ci/prek/sync_java_sdk_supervisor_schema.py +++ b/scripts/ci/prek/sync_java_sdk_supervisor_schema.py @@ -53,7 +53,7 @@ def configured_version() -> str: return match.group(1) -def api_version_of(path: Path) -> str | None: +def read_api_version(path: Path) -> str | None: if not path.exists(): return None try: @@ -66,7 +66,7 @@ def api_version_of(path: Path) -> str | None: def main() -> int: expected = configured_version() - if api_version_of(MONOREPO_SCHEMA_FILE) == expected: + if read_api_version(MONOREPO_SCHEMA_FILE) == expected: snapshot = MONOREPO_SCHEMA_FILE.read_bytes() if SCHEMA_FILE.exists() and SCHEMA_FILE.read_bytes() == snapshot: print(f"Supervisor Schema matches the Task SDK snapshot (api_version={expected}).") @@ -74,7 +74,7 @@ def main() -> int: print(f"Refreshing Supervisor Schema from {MONOREPO_SCHEMA_FILE.relative_to(REPO_ROOT)}.") SCHEMA_FILE.write_bytes(snapshot) return 0 - actual = api_version_of(SCHEMA_FILE) + actual = read_api_version(SCHEMA_FILE) if actual == expected: print(f"Supervisor Schema is up-to-date (api_version={expected}).") return 0 diff --git a/scripts/tests/ci/prek/test_sync_java_sdk_supervisor_schema.py b/scripts/tests/ci/prek/test_sync_java_sdk_supervisor_schema.py new file mode 100644 index 0000000000000..ec60dd043734f --- /dev/null +++ b/scripts/tests/ci/prek/test_sync_java_sdk_supervisor_schema.py @@ -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. +from __future__ import annotations + +import json +import os +import subprocess +from typing import TYPE_CHECKING +from unittest import mock + +import pytest +from ci.prek import sync_java_sdk_supervisor_schema as sync + +if TYPE_CHECKING: + from pathlib import Path + +CONFIGURED_VERSION = "2026-06-16" + + +def _write_schema(path: Path, api_version: str, **body) -> None: + path.write_text(json.dumps({"api_version": api_version, **body})) + + +@pytest.fixture +def schemas(tmp_path, monkeypatch): + """Point the hook at a temp gradle.properties, vendored schema and Task SDK snapshot.""" + gradle_properties = tmp_path / "gradle.properties" + gradle_properties.write_text(f"airflowSupervisorSchemaVersion={CONFIGURED_VERSION}\n") + vendored = tmp_path / "java-sdk" / "schema.json" + snapshot = tmp_path / "task-sdk" / "schema.json" + vendored.parent.mkdir() + snapshot.parent.mkdir() + monkeypatch.setattr(sync, "REPO_ROOT", tmp_path) + monkeypatch.setattr(sync, "GRADLE_PROPERTIES", gradle_properties) + monkeypatch.setattr(sync, "SCHEMA_FILE", vendored) + monkeypatch.setattr(sync, "MONOREPO_SCHEMA_FILE", snapshot) + return vendored, snapshot + + +@mock.patch.object(sync.subprocess, "run", autospec=True) +class TestMain: + def test_leaves_a_vendored_schema_that_matches_the_snapshot(self, mock_run, schemas, capsys): + vendored, snapshot = schemas + _write_schema(snapshot, CONFIGURED_VERSION) + _write_schema(vendored, CONFIGURED_VERSION) + os.utime(vendored, ns=(0, 0)) + + assert sync.main() == 0 + + assert "matches the Task SDK snapshot" in capsys.readouterr().out + assert vendored.stat().st_mtime_ns == 0 + mock_run.assert_not_called() + + def test_copies_the_snapshot_over_a_differing_vendored_schema(self, mock_run, schemas): + vendored, snapshot = schemas + _write_schema(snapshot, CONFIGURED_VERSION, definitions={"New": {}}) + _write_schema(vendored, CONFIGURED_VERSION) + + assert sync.main() == 0 + + assert vendored.read_bytes() == snapshot.read_bytes() + mock_run.assert_not_called() + + def test_runs_gradle_when_the_snapshot_declares_another_version(self, mock_run, schemas): + vendored, snapshot = schemas + _write_schema(snapshot, "2026-01-01") + _write_schema(vendored, "2025-01-01") + mock_run.return_value = subprocess.CompletedProcess(args=[], returncode=0) + + assert sync.main() == 0 + + mock_run.assert_called_once_with( + [str(sync.JAVA_SDK_DIR / "gradlew"), "-p", str(sync.JAVA_SDK_DIR), ":sdk:syncSupervisorSchema"], + check=False, + ) + assert json.loads(vendored.read_text()) == {"api_version": "2025-01-01"} From 8f942ec422deaf08a18b8d23ec35e3deac5435bf Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 1 Oct 2026 06:10:35 +0000 Subject: [PATCH 82/94] Tidy names and mocks in the Java probe test An unspecced patch of _should_use_exec accepts any call, and the noun or adjective helper names read as attributes. --- .../test_task_handler_processor_java.py | 46 ++++++++++--------- 1 file changed, 25 insertions(+), 21 deletions(-) diff --git a/airflow-core/tests/unit/dag_processing/test_task_handler_processor_java.py b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_java.py index 9dd27999e1c69..66f52d978e9ba 100644 --- a/airflow-core/tests/unit/dag_processing/test_task_handler_processor_java.py +++ b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_java.py @@ -115,14 +115,14 @@ def java_coordinator() -> Iterator[None]: # The parse child is a bare fork even on macOS, so it sees this config. with ( conf_vars({("sdk", "coordinators"): json.dumps(spec)}), - mock.patch.object(supervisor, "_should_use_exec", return_value=False), + mock.patch.object(supervisor, "_should_use_exec", autospec=True, return_value=False), ): yield finally: reset_coordinator_manager() -def _nullable(schema: dict) -> dict: +def _make_nullable(schema: dict) -> dict: return {"anyOf": [schema, {"type": "null"}]} @@ -133,7 +133,7 @@ def _nullable(schema: dict) -> dict: STRING = {"type": "string"} -def _positional(task_id: str, *params: tuple[str, dict]) -> TaskHandlerDeclaration: +def _declare_positional(task_id: str, *params: tuple[str, dict]) -> TaskHandlerDeclaration: return TaskHandlerDeclaration( task_id=task_id, binding="positional", @@ -141,7 +141,7 @@ def _positional(task_id: str, *params: tuple[str, dict]) -> TaskHandlerDeclarati ) -def _named(task_id: str, *params: tuple[str, dict, bool]) -> TaskHandlerDeclaration: +def _declare_named(task_id: str, *params: tuple[str, dict, bool]) -> TaskHandlerDeclaration: return TaskHandlerDeclaration( task_id=task_id, binding="named", @@ -176,29 +176,33 @@ def test_a_built_bundle_declares_its_task_handlers(example_bundle): fileloc=os.fspath(jar), task_handlers={ "java_xcom_casting_example": [ - _named("produce_number"), - _positional("widen_to_long", ("value", INT64)), - _positional("widen_to_double", ("value", DOUBLE)), - _named("produce_nothing"), - _positional("consume_nullable", ("value", _nullable(INT32))), - _named("produce_fraction"), - _positional("consume_float", ("value", FLOAT)), - _positional( + _declare_named("produce_number"), + _declare_positional("widen_to_long", ("value", INT64)), + _declare_positional("widen_to_double", ("value", DOUBLE)), + _declare_named("produce_nothing"), + _declare_positional("consume_nullable", ("value", _make_nullable(INT32))), + _declare_named("produce_fraction"), + _declare_positional("consume_float", ("value", FLOAT)), + _declare_positional( "consume_double_list", - ("values", _nullable({"type": "array", "items": _nullable(DOUBLE)})), + ("values", _make_nullable({"type": "array", "items": _make_nullable(DOUBLE)})), ), ], "java_interface_example": [ - _named("extract"), - _named("transform", ("extracted", INT64, False)), - _named("summarize", ("region_code", _nullable(STRING), True), ("transformed", INT64, False)), + _declare_named("extract"), + _declare_named("transform", ("extracted", INT64, False)), + _declare_named( + "summarize", ("region_code", _make_nullable(STRING), True), ("transformed", INT64, False) + ), ], "java_annotation_example": [ - _named("extract"), - _positional("transform", ("extracted", INT64)), - _positional("load", ("transformed", INT64)), - _named("report", ("runLabel", _nullable(STRING), False), ("transformed", INT64, False)), - _named("concurrent"), + _declare_named("extract"), + _declare_positional("transform", ("extracted", INT64)), + _declare_positional("load", ("transformed", INT64)), + _declare_named( + "report", ("runLabel", _make_nullable(STRING), False), ("transformed", INT64, False) + ), + _declare_named("concurrent"), ], }, ) From 18958abe00f88466027faf9bdbbf97b39c1c33c0 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 1 Oct 2026 06:11:28 +0000 Subject: [PATCH 83/94] Test the cache digest of classpath directories and non-zip files Composite and included builds put directories on the runtime classpath, and a classpath file can be something other than a JAR. Neither path was tested, so a change there could leave a stale digest unnoticed. --- .../airflow/sdk/plugin/BundleDigestTest.kt | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/BundleDigestTest.kt b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/BundleDigestTest.kt index 272b5eac1cb90..b87dcdb0c6316 100644 --- a/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/BundleDigestTest.kt +++ b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/BundleDigestTest.kt @@ -95,4 +95,27 @@ class BundleDigestTest { assertNotEquals(first, digest(mainClass = "com.example.Other")) assertNotEquals(first, digest(fatJar = false)) } + + @Test + @DisplayName("Should change when a classpath directory's content changes") + fun changesWithClasspathDirectoryContent() { + val included = classesDir("included", classes) + val first = digest(classpath = listOf(included)) + + included.resolve("com/example/Handler.class").writeText("HANDLER") + + assertNotEquals(first, digest(classpath = listOf(included))) + } + + @Test + @DisplayName("Should change when a non-zip classpath file's content changes") + fun changesWithNonZipClasspathFileContent() { + val native = tmp.resolve("lib/native.so").apply { parentFile.mkdirs() } + native.writeText("native") + val first = digest(classpath = listOf(native)) + + native.writeText("NATIVE") + + assertNotEquals(first, digest(classpath = listOf(native))) + } } From 628e134066d0e2b50e48fe20b4badc3b602003e3 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 1 Oct 2026 07:40:39 +0000 Subject: [PATCH 84/94] Java SDK: answer a task handler parse with every handler The parse request no longer names Dags, and the Dag processor reuses an artifact's answer for every Dag file that uses the artifact, so the answer must not depend on what was asked. The vendored schema is resynced in the same commit because the request model generated from it no longer has dag_ids. --- .../test_task_handler_processor_java.py | 9 +------- java-sdk/README.md | 3 +-- java-sdk/sdk/schema/schema.json | 12 ++-------- .../kotlin/org/apache/airflow/sdk/Server.kt | 5 ++--- .../airflow/sdk/execution/TaskHandlerParse.kt | 22 +++++++------------ .../org/apache/airflow/sdk/ServerTest.kt | 9 +++----- .../sdk/execution/TaskHandlerParseTest.kt | 21 +++++++++--------- 7 files changed, 28 insertions(+), 53 deletions(-) diff --git a/airflow-core/tests/unit/dag_processing/test_task_handler_processor_java.py b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_java.py index 66f52d978e9ba..6b3ba66acd632 100644 --- a/airflow-core/tests/unit/dag_processing/test_task_handler_processor_java.py +++ b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_java.py @@ -162,13 +162,6 @@ def test_a_built_bundle_declares_its_task_handlers(example_bundle): bundle_path=example_bundle, bundle_name="java-task-handlers", artifact_rel_path=jar.name, - dag_ids=[ - "java_xcom_casting_example", - "java_interface_example", - "java_annotation_example", - "java_native_interface_example", - "not_registered", - ], logger=structlog.get_logger(), ) @@ -207,9 +200,9 @@ def test_a_built_bundle_declares_its_task_handlers(example_bundle): }, ) assert list(result.task_handlers) == [ - "java_xcom_casting_example", "java_interface_example", "java_annotation_example", + "java_xcom_casting_example", ] with zipfile.ZipFile(jar) as zf: manifest = _parse_manifest(zf.read("META-INF/MANIFEST.MF")) diff --git a/java-sdk/README.md b/java-sdk/README.md index 53f3c037f7303..f757301b7e1ce 100644 --- a/java-sdk/README.md +++ b/java-sdk/README.md @@ -677,8 +677,7 @@ launches the JVM application as a subprocess. The flow is: The Dag processor also launches each JAR to ask which task handlers it registers. The first frame is then a `TaskHandlerParseRequest`, and `execution/TaskHandlerParse.kt` replies with one `TaskHandlerParsingResult`: -the task-handler registrations of the requested Dags, each with its binding -mode and parameters. Flat data parameters bind `positional` and a `TaskInput` +every task-handler registration, each with its binding mode and parameters. Flat data parameters bind `positional` and a `TaskInput` binds `named`. The annotation processor records a handler's parameter names on its generated task class (`internal/TaskParams.kt`), because Java does not keep them at run time. The JVM waits for the reply's acknowledgement, then exits. diff --git a/java-sdk/sdk/schema/schema.json b/java-sdk/sdk/schema/schema.json index 30f196e693cc4..6ca8e046cbdd3 100644 --- a/java-sdk/sdk/schema/schema.json +++ b/java-sdk/sdk/schema/schema.json @@ -4232,19 +4232,12 @@ "type": "object" }, "TaskHandlerParseRequest": { - "description": "Request for Task Handler Parsing.\n\nAsks a Lang-SDK runtime which task handlers an artifact registers for the given Dags.", + "description": "Request for Task Handler Parsing.\n\nAsks a Lang-SDK runtime for every task handler an artifact registers.", "properties": { "file": { "title": "File", "type": "string" }, - "dag_ids": { - "items": { - "type": "string" - }, - "title": "Dag Ids", - "type": "array" - }, "bundle_path": { "format": "path", "title": "Bundle Path", @@ -4263,7 +4256,6 @@ }, "required": [ "file", - "dag_ids", "bundle_path", "bundle_name" ], @@ -4271,7 +4263,7 @@ "type": "object" }, "TaskHandlerParsingResult": { - "description": "Result of Task Handler Parsing.\n\nThe task handlers a Lang-SDK artifact registers, keyed by Dag id.", + "description": "Result of Task Handler Parsing.\n\nEvery task handler a Lang-SDK artifact registers, keyed by Dag id.\n\nThe answer depends only on the artifact, never on the request.", "properties": { "fileloc": { "title": "Fileloc", diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt index e3b8bd49a08b7..e039f4e8671c6 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt @@ -73,9 +73,8 @@ class ApiError( * execution requests to the registered [Bundle]. * * The Dag processor starts the same entry point to ask which task handlers - * the bundle registers. That process replies with the handlers registered for - * the requested Dags and how each binds its arguments, runs no task, and - * exits. + * the bundle registers. That process replies with every registered handler + * and how each binds its arguments, runs no task, and exits. * * The typical entry point is: * diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParse.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParse.kt index 9c177f9fe0fa3..c4573998d0cdb 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParse.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParse.kt @@ -36,29 +36,23 @@ import java.lang.reflect.Modifier import java.lang.reflect.Type /** - * Answers a [TaskHandlerParseRequest] with the task handlers [bundle] - * registers for the requested Dags, as a TaskHandlerParsingResult body. + * Answers a [TaskHandlerParseRequest] with every task handler [bundle] + * registers, keyed by Dag id in registration order, as a + * TaskHandlerParsingResult body. * * Only task-handler registrations count: a Dag declared in Java is its own, - * not a Python Dag's. A requested Dag with no handlers is left out. The - * generated models cannot hold `task_handlers`, so the body is a map. + * not a Python Dag's. The generated models cannot hold `task_handlers`, so + * the body is a map. */ internal fun parseTaskHandlers( bundle: Bundle, request: TaskHandlerParseRequest, -): Map { - val declarations = linkedMapOf>>() - for (dagId in request.dagIds) { - if (dagId in declarations) continue - val dag = bundle.taskHandlers[dagId] ?: continue - declarations[dagId] = dag.tasks.values.map(::declare) - } - return linkedMapOf( +): Map = + linkedMapOf( "type" to "TaskHandlerParsingResult", "fileloc" to request.file, - "task_handlers" to declarations, + "task_handlers" to bundle.taskHandlers.mapValues { (_, dag) -> dag.tasks.values.map(::declare) }, ) -} private fun declare(task: TaskDef): Map { val (binding, params) = bindingOf(task.definition) diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt index d5a0db1bfd31f..65e062b2dedcd 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt @@ -135,7 +135,7 @@ class ServerTest { val supervisor = Thread { runBlocking { - toServer.writeFrame(taskHandlerParseRequestFrame(3, "/bundle/etl.jar", listOf("etl"))) + toServer.writeFrame(taskHandlerParseRequestFrame(3, "/bundle/etl.jar")) val prefix = fromServer.readByteArray(4) val raw = Frame.decodeRaw(ArrayBufferInput(fromServer.readByteArray(Frame.parseLengthPrefix(prefix).toInt()))) reported.put(raw) @@ -169,7 +169,7 @@ class ServerTest { val supervisor = Thread { runBlocking { - toServer.writeFrame(taskHandlerParseRequestFrame(3, "/bundle/etl.jar", listOf("etl"))) + toServer.writeFrame(taskHandlerParseRequestFrame(3, "/bundle/etl.jar")) val prefix = fromServer.readByteArray(4) fromServer.readByteArray(Frame.parseLengthPrefix(prefix).toInt()) toServer.close() @@ -187,17 +187,14 @@ class ServerTest { private fun taskHandlerParseRequestFrame( id: Int, file: String, - dagIds: List, ): ByteArray { val out = ByteArrayOutputStream() MessagePack.newDefaultPacker(out).use { packer -> packer.packArrayHeader(2) packer.packInt(id) - packer.packMapHeader(5) + packer.packMapHeader(4) packer.packString("type").packString("TaskHandlerParseRequest") packer.packString("file").packString(file) - packer.packString("dag_ids").packArrayHeader(dagIds.size) - dagIds.forEach { packer.packString(it) } packer.packString("bundle_path").packString("/bundle") packer.packString("bundle_name").packString("java-task-handlers") } diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParseTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParseTest.kt index fa1c167bc8c8e..5e72621a2cc5f 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParseTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParseTest.kt @@ -109,10 +109,9 @@ private class Audit : Task { ) = Unit } -private fun request(vararg dagIds: String) = +private fun request() = TaskHandlerParseRequest().apply { file = "/bundles/java/etl.jar" - this.dagIds = dagIds.toList() bundlePath = "/bundles/java" bundleName = "java-task-handlers" } @@ -151,8 +150,8 @@ private fun assertConforms( class TaskHandlerParseTest { @Test - @DisplayName("Should declare each requested Dag's task handlers in registration order") - fun declaresRequestedDagsInRegistrationOrder() { + @DisplayName("Should declare every task handler in registration order") + fun declaresEveryTaskHandlerInRegistrationOrder() { val bundle = Bundle() .register("etl", "extract", GeneratedFlat::class.java) @@ -163,7 +162,7 @@ class TaskHandlerParseTest { .register("report", "audit", Audit::class.java) .register(DagDef("native").addTask("audit", Audit::class.java)) - val result = parseTaskHandlers(bundle, request("etl", "native", "missing", "etl")) + val result = parseTaskHandlers(bundle, request()) val regions = mapOf("anyOf" to listOf(mapOf("type" to "array", "items" to NULLABLE_STRING), mapOf("type" to "null"))) @@ -189,9 +188,11 @@ class TaskHandlerParseTest { mapOf("task_id" to "audit", "binding" to "named", "params" to emptyList()), mapOf("task_id" to "notify", "binding" to "named", "params" to emptyList()), ), + "report" to listOf(mapOf("task_id" to "audit", "binding" to "named", "params" to emptyList())), ), ) assertEquals(expected, result) + assertEquals(listOf("etl", "report"), (result["task_handlers"] as Map<*, *>).keys.toList()) } @Test @@ -203,7 +204,7 @@ class TaskHandlerParseTest { .register("etl", "score", GeneratedInput::class.java) .register("etl", "audit", Audit::class.java) - val result = parseTaskHandlers(bundle, request("etl")) + val result = parseTaskHandlers(bundle, request()) assertConforms("TaskHandlerParsingResult", result) val bindings = @@ -221,11 +222,11 @@ class TaskHandlerParseTest { } @Test - @DisplayName("Should send an empty map when no requested Dag has task handlers") - fun sendsEmptyTaskHandlersWhenNothingMatches() { - val bundle = Bundle().register("etl", "audit", Audit::class.java) + @DisplayName("Should send an empty map when the bundle registers no task handler") + fun sendsEmptyTaskHandlersWithoutHandlers() { + val bundle = Bundle().register(DagDef("native").addTask("audit", Audit::class.java)) - val result = parseTaskHandlers(bundle, request("other")) + val result = parseTaskHandlers(bundle, request()) assertEquals(emptyMap(), result["task_handlers"]) } From f599ee696049add9e286499ce579e0ad0008e67c Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 1 Oct 2026 17:46:04 +0000 Subject: [PATCH 85/94] Java SDK: declare task handler params without required The supervisor schema dropped TaskHandlerParam.required, so a reply that still sends it no longer matches the schema the Dag processor decodes with. Flat params stay positional and a TaskInput stays named, which now means a field nothing fills only warns. The vendored schema is synced from the monorepo snapshot with sync-java-sdk-supervisor-schema. --- .../test_task_handler_processor_java.py | 4 +-- java-sdk/sdk/schema/schema.json | 28 +++++++++---------- .../airflow/sdk/execution/TaskHandlerParse.kt | 11 +++----- .../sdk/execution/TaskHandlerParseTest.kt | 7 ++--- 4 files changed, 23 insertions(+), 27 deletions(-) diff --git a/airflow-core/tests/unit/dag_processing/test_task_handler_processor_java.py b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_java.py index 6b3ba66acd632..a76ae128c17d8 100644 --- a/airflow-core/tests/unit/dag_processing/test_task_handler_processor_java.py +++ b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_java.py @@ -137,7 +137,7 @@ def _declare_positional(task_id: str, *params: tuple[str, dict]) -> TaskHandlerD return TaskHandlerDeclaration( task_id=task_id, binding="positional", - params=[TaskHandlerParam(name=name, value_schema=schema, required=True) for name, schema in params], + params=[TaskHandlerParam(name=name, value_schema=schema) for name, schema in params], ) @@ -146,7 +146,7 @@ def _declare_named(task_id: str, *params: tuple[str, dict, bool]) -> TaskHandler task_id=task_id, binding="named", params=[ - TaskHandlerParam(name=name, value_schema=schema, required=False, exact_name=exact) + TaskHandlerParam(name=name, value_schema=schema, exact_name=exact) for name, schema, exact in params ], ) diff --git a/java-sdk/sdk/schema/schema.json b/java-sdk/sdk/schema/schema.json index 6ca8e046cbdd3..98b9be109d26f 100644 --- a/java-sdk/sdk/schema/schema.json +++ b/java-sdk/sdk/schema/schema.json @@ -4166,19 +4166,24 @@ "binding": { "enum": [ "positional", - "named", - "named_or_whole", - "named_open" + "named" ], "title": "Binding", "type": "string" }, "params": { - "items": { - "$ref": "#/$defs/TaskHandlerParam" - }, - "title": "Params", - "type": "array" + "anyOf": [ + { + "items": { + "$ref": "#/$defs/TaskHandlerParam" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Task Handler Params" } }, "required": [ @@ -4214,10 +4219,6 @@ ], "default": null }, - "required": { - "title": "Required", - "type": "boolean" - }, "exact_name": { "default": false, "title": "Exact Name", @@ -4225,8 +4226,7 @@ } }, "required": [ - "name", - "required" + "name" ], "title": "TaskHandlerParam", "type": "object" diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParse.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParse.kt index c4573998d0cdb..12ac86159d071 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParse.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParse.kt @@ -66,8 +66,8 @@ private fun declare(task: TaskDef): Map { /** * How [definition] binds the stub call's arguments. Flat parameters bind by * position; a [TaskInput] binds by name. A task that reads no argument is - * declared as binding by name with no parameters, so an argument passed to it - * stands out as one nothing takes. + * declared as binding by name with no parameters: it ignores an argument + * passed to it, so the Dag processor only warns about one. */ private fun bindingOf(definition: Class): Pair>> { if (InputTask::class.java.isAssignableFrom(definition)) { @@ -76,7 +76,7 @@ private fun bindingOf(definition: Class): Pair): TaskParams? { return if (Modifier.isStatic(field.modifiers)) field.get(null) as? TaskParams else null } -// A field nothing supplies keeps its Java default and only warns, so no field is required. private fun inputParams(inputType: Class): List> = - bindableFields(inputType).map { param(argNameOf(it), it.genericType, required = false, exactName = isPinned(it)) } + bindableFields(inputType).map { param(argNameOf(it), it.genericType, exactName = isPinned(it)) } private fun param( name: String, type: Type, - required: Boolean, exactName: Boolean, ): Map = linkedMapOf( "name" to name, - "required" to required, "exact_name" to exactName, ).apply { buildValueSchema(type)?.let { put("value_schema", it) } } diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParseTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParseTest.kt index 5e72621a2cc5f..02ae55c96a518 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParseTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParseTest.kt @@ -119,9 +119,8 @@ private fun request() = private fun param( name: String, schema: Map, - required: Boolean = false, exactName: Boolean = false, -) = mapOf("name" to name, "required" to required, "exact_name" to exactName, "value_schema" to schema) +) = mapOf("name" to name, "exact_name" to exactName, "value_schema" to schema) private val SCORE_INPUT_PARAMS = listOf( @@ -179,8 +178,8 @@ class TaskHandlerParseTest { "binding" to "positional", "params" to listOf( - param("rows", mapOf("type" to "integer", "format" to "int64"), required = true), - param("regions", regions, required = true), + param("rows", mapOf("type" to "integer", "format" to "int64")), + param("regions", regions), ), ), mapOf("task_id" to "score", "binding" to "named", "params" to SCORE_INPUT_PARAMS), From 580f68ab98e828c3780f2403642c73fe3692ec62 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 18:16:57 +0000 Subject: [PATCH 86/94] Let the coordinator registry list its task-handler bundles The Dag processor manager must know which bundles can hold task-handler artifacts before it dispatches a parse, and it must not build coordinators to find out, since that imports and constructs every runtime. The validated specs already carry the answer. Only the coordinators a queue routes to are listed, since only they run stub tasks. --- .../airflow/sdk/execution_time/coordinator.py | 15 ++++++ .../execution_time/test_coordinator.py | 50 +++++++++++++++++++ 2 files changed, 65 insertions(+) diff --git a/task-sdk/src/airflow/sdk/execution_time/coordinator.py b/task-sdk/src/airflow/sdk/execution_time/coordinator.py index 904faf5068de3..f3d37858503b9 100644 --- a/task-sdk/src/airflow/sdk/execution_time/coordinator.py +++ b/task-sdk/src/airflow/sdk/execution_time/coordinator.py @@ -357,6 +357,21 @@ def extra_for_queue(self, queue: str) -> dict[str, Any] | None: return None return spec.extra + def get_task_handler_bundle_names(self) -> dict[str, str | None]: + """ + Return the ``task_handler_bundle_name`` of each coordinator a queue routes to, keyed by coordinator. + + ``None`` means the coordinator reads task handlers from the task's own Dag bundle. + A coordinator no queue routes to runs no stub task, so it is left out and its bundle is not read. + Only the declarative specs are read; no coordinator is instantiated. + """ + routed_keys = set(self._queue_to_coordinator.values()) + return { + key: spec.kwargs.get("task_handler_bundle_name") + for key, spec in self._coordinator_specs.items() + if key in routed_keys + } + @functools.cache def get_coordinator_manager() -> CoordinatorManager: diff --git a/task-sdk/tests/task_sdk/execution_time/test_coordinator.py b/task-sdk/tests/task_sdk/execution_time/test_coordinator.py index cc8a81410a04f..cfca1e95baaba 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_coordinator.py +++ b/task-sdk/tests/task_sdk/execution_time/test_coordinator.py @@ -376,6 +376,56 @@ def test_extra_for_queue_does_not_instantiate_coordinator(self, sdk_config): } assert manager._created_coordinators == {} + def test_get_task_handler_bundle_names_does_not_instantiate_coordinator(self, sdk_config): + sdk_config( + coordinators=json.dumps( + { + "named": { + "classpath": f"{_ExplodingCoordinator.__module__}._ExplodingCoordinator", + "kwargs": {"task_handler_bundle_name": "java-task-handlers"}, + }, + "own-bundle": {"classpath": f"{_ExplodingCoordinator.__module__}._ExplodingCoordinator"}, + } + ), + queue_to_coordinator=json.dumps({"queue-named": "named", "queue-own-bundle": "own-bundle"}), + ) + bundles = [ + { + "name": "java-task-handlers", + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {}, + } + ] + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(bundles)}): + manager = CoordinatorManager.from_config() + + assert manager.get_task_handler_bundle_names() == { + "named": "java-task-handlers", + "own-bundle": None, + } + assert manager._created_coordinators == {} + + @pytest.mark.parametrize( + "kwargs", + [{"task_handler_bundle_name": "ghost"}, {"task_handler_bundle_name": ["ghost"]}, {}], + ids=["unconfigured", "not-a-string", "own-bundle"], + ) + def test_get_task_handler_bundle_names_leaves_out_unrouted_coordinator(self, sdk_config, kwargs): + sdk_config( + coordinators=json.dumps( + { + "routed": {"classpath": f"{_ExplodingCoordinator.__module__}._ExplodingCoordinator"}, + "unrouted": { + "classpath": f"{_ExplodingCoordinator.__module__}._ExplodingCoordinator", + "kwargs": kwargs, + }, + } + ), + queue_to_coordinator=json.dumps({"queue-routed": "routed"}), + ) + + assert CoordinatorManager.from_config().get_task_handler_bundle_names() == {"routed": None} + class TestConfigYamlCoordinatorsExample: """Guard the ``[sdk] coordinators`` example in ``config.yml`` against drift. From 657bdc5ac84e3b93c111634a69ee0b3527b484c5 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 18:20:47 +0000 Subject: [PATCH 87/94] Add known task-handler artifacts to the Dag parse request The Dag-parsing child has no database, so the artifacts the fast path compares against must arrive with the parse request. Each carries its whole probe answer next to its fingerprint: a match then lets the child validate its stub tasks against that answer without starting a runtime, and an answer is only ever trusted under the fingerprint it was probed with. DagFileParseRequest already existed at 2026-06-16, so a runtime pinned there must not receive the new field. --- .../src/airflow/dag_processing/processor.py | 67 +++++++++++----- .../tests/unit/dag_processing/test_manager.py | 70 +++++++++++++++- .../unit/dag_processing/test_processor.py | 79 +++++++++++++++++++ .../sdk/execution_time/schema/schema.json | 56 +++++++++++++ .../schema/versions/__init__.py | 2 + .../schema/versions/v2026_10_30.py | 12 ++- .../execution_time/schema/test_migrator.py | 75 +++++++++++++++++- 7 files changed, 335 insertions(+), 26 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/processor.py b/airflow-core/src/airflow/dag_processing/processor.py index b3b177f451447..408d108869437 100644 --- a/airflow-core/src/airflow/dag_processing/processor.py +++ b/airflow-core/src/airflow/dag_processing/processor.py @@ -97,26 +97,6 @@ from airflow.typing_compat import Self -class DagFileParseRequest(BaseModel): - """ - Request for DAG File Parsing. - - This is the request that the manager will send to the DAG parser with the dag file and - any other necessary metadata. - """ - - file: str - - bundle_path: Path - """Passing bundle path around lets us figure out relative file path.""" - - bundle_name: str - """Bundle name for team-specific executor validation.""" - - callback_requests: list[CallbackRequest] = Field(default_factory=list) - type: Literal["DagFileParseRequest"] = "DagFileParseRequest" - - class DagFileParsingResult(BaseModel): """ Result of DAG File Parsing. @@ -180,6 +160,47 @@ class TaskHandlerDeclaration(BaseModel): """ +class TaskHandlerArtifact(BaseModel): + """A Lang-SDK artifact and every task handler it registers.""" + + bundle_name: str + + relative_fileloc: str = Field(max_length=2000) + """Path of the artifact within its bundle.""" + + size_bytes: int + + cache_digest: str | None = Field(max_length=128) + """Opaque content fingerprint defined by the coordinator; ``None`` when the artifact stores none, so it is always probed.""" + + task_handlers: dict[str, list[TaskHandlerDeclaration]] + """Every Dag id the artifact registers a task handler for, with those handlers.""" + + +class DagFileParseRequest(BaseModel): + """ + Request for DAG File Parsing. + + This is the request that the manager will send to the DAG parser with the dag file and + any other necessary metadata. + """ + + file: str + + bundle_path: Path + """Passing bundle path around lets us figure out relative file path.""" + + bundle_name: str + """Bundle name for team-specific executor validation.""" + + callback_requests: list[CallbackRequest] = Field(default_factory=list) + + known_artifacts: list[TaskHandlerArtifact] = Field(default_factory=list) + """The recorded task-handler artifacts, with their answers, that this file's stub tasks may resolve against.""" + + type: Literal["DagFileParseRequest"] = "DagFileParseRequest" + + class TaskHandlerParseRequest(BaseModel): """ Request for Task Handler Parsing. @@ -787,6 +808,7 @@ def start( # type: ignore[override] bundle_name: str, dag_file_rel_path: str, callbacks: list[CallbackRequest], + known_artifacts: Sequence[TaskHandlerArtifact] = (), target: Callable[[], None] = _parse_file_entrypoint, client: Client, **kwargs, @@ -814,7 +836,7 @@ def start( # type: ignore[override] **kwargs, ) proc.had_callbacks = bool(callbacks) # Track if this process had callbacks - proc._on_child_started(callbacks, path, bundle_path, bundle_name) + proc._on_child_started(callbacks, path, bundle_path, bundle_name, known_artifacts=known_artifacts) return proc def _on_child_started( @@ -823,11 +845,14 @@ def _on_child_started( path: str | os.PathLike[str], bundle_path: Path, bundle_name: str, + *, + known_artifacts: Sequence[TaskHandlerArtifact] = (), ) -> None: msg = DagFileParseRequest( file=os.fspath(path), bundle_path=bundle_path, bundle_name=bundle_name, callback_requests=callbacks, + known_artifacts=list(known_artifacts), ) self.send_msg(msg, request_id=0) diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index a3624e1d4a143..a8d21d3990f60 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -61,6 +61,9 @@ DagFileParseRequest, DagFileParsingResult, DagFileProcessorProcess, + TaskHandlerArtifact, + TaskHandlerDeclaration, + TaskHandlerParam, _parse_file, ) from airflow.models import DagModel, DbCallbackRequest @@ -1859,9 +1862,10 @@ def test_collect_results_tolerates_stale_file_handle_on_close(self): @pytest.mark.usefixtures("testing_dag_bundle") @pytest.mark.parametrize( - ("callbacks", "path", "expected_body"), + ("callbacks", "known_artifacts", "path", "expected_body"), [ pytest.param( + [], [], "/opt/airflow/dags/test_dag.py", { @@ -1869,8 +1873,62 @@ def test_collect_results_tolerates_stale_file_handle_on_close(self): "bundle_path": "/opt/airflow/dags", "bundle_name": "testing", "callback_requests": [], + "known_artifacts": [], + "type": "DagFileParseRequest", + }, + ), + pytest.param( + [], + [ + TaskHandlerArtifact( + bundle_name="java-task-handlers", + relative_fileloc="etl.jar", + size_bytes=1024, + cache_digest="ab12", + task_handlers={ + "etl": [ + TaskHandlerDeclaration( + task_id="extract", + binding="positional", + params=[TaskHandlerParam(name=None, required=True)], + ) + ] + }, + ) + ], + "/opt/airflow/dags/etl.py", + { + "file": "/opt/airflow/dags/etl.py", + "bundle_path": "/opt/airflow/dags", + "bundle_name": "testing", + "callback_requests": [], + "known_artifacts": [ + { + "bundle_name": "java-task-handlers", + "relative_fileloc": "etl.jar", + "size_bytes": 1024, + "cache_digest": "ab12", + "task_handlers": { + "etl": [ + { + "task_id": "extract", + "binding": "positional", + "params": [ + { + "name": None, + "value_schema": None, + "required": True, + "exact_name": False, + } + ], + } + ] + }, + } + ], "type": "DagFileParseRequest", }, + id="known-artifacts", ), pytest.param( [ @@ -1884,6 +1942,7 @@ def test_collect_results_tolerates_stale_file_handle_on_close(self): is_failure_callback=False, ) ], + [], "/opt/airflow/dags/dag_callback_dag.py", { "file": "/opt/airflow/dags/dag_callback_dag.py", @@ -1903,17 +1962,22 @@ def test_collect_results_tolerates_stale_file_handle_on_close(self): "type": "DagCallbackRequest", } ], + "known_artifacts": [], "type": "DagFileParseRequest", }, ), ], ) - def test_serialize_callback_requests(self, callbacks, path, expected_body): + def test_serialize_parse_request(self, callbacks, known_artifacts, path, expected_body): from airflow.sdk.execution_time.comms import _ResponseFrame processor, read_socket = self.mock_processor() processor._on_child_started( - callbacks, path, bundle_path=Path("/opt/airflow/dags"), bundle_name="testing" + callbacks, + path, + bundle_path=Path("/opt/airflow/dags"), + bundle_name="testing", + known_artifacts=known_artifacts, ) read_socket.settimeout(0.1) diff --git a/airflow-core/tests/unit/dag_processing/test_processor.py b/airflow-core/tests/unit/dag_processing/test_processor.py index 8887440ecb527..1e1e710948994 100644 --- a/airflow-core/tests/unit/dag_processing/test_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_processor.py @@ -58,6 +58,7 @@ DagFileParseRequest, DagFileParsingResult, DagFileProcessorProcess, + TaskHandlerArtifact, TaskHandlerDeclaration, TaskHandlerParam, TaskHandlerParseRequest, @@ -621,6 +622,84 @@ def test_parse_file_entrypoint_parses_dag_callbacks(mocker): ] +def _make_known_artifact_body(**overrides) -> dict: + return { + "bundle_name": "java-task-handlers", + "relative_fileloc": "etl.jar", + "size_bytes": 1024, + "cache_digest": "ab12", + "task_handlers": { + "etl": [ + {"task_id": "extract", "binding": "positional", "params": [{"name": None, "required": True}]} + ] + }, + **overrides, + } + + +def test_parse_file_entrypoint_decodes_known_artifacts(): + r, w = socketpair() + frame = comms._ResponseFrame( + id=0, + body={ + "file": "/files/dags/etl.py", + "bundle_path": "/files/dags", + "bundle_name": "testing", + "known_artifacts": [ + _make_known_artifact_body(), + _make_known_artifact_body(relative_fileloc="empty.jar", cache_digest=None, task_handlers={}), + ], + "type": "DagFileParseRequest", + }, + ) + w.sendall(frame.as_bytes()) + + # The same decoder _parse_file_entrypoint builds. + decoder = comms.CommsDecoder[ToDagProcessor, ToManager]( + socket=r, + body_decoder=TypeAdapter[ToDagProcessor](ToDagProcessor), + ) + + msg = decoder._get_response() + assert isinstance(msg, DagFileParseRequest) + assert msg.known_artifacts == [ + TaskHandlerArtifact( + bundle_name="java-task-handlers", + relative_fileloc="etl.jar", + size_bytes=1024, + cache_digest="ab12", + task_handlers={ + "etl": [ + TaskHandlerDeclaration( + task_id="extract", + binding="positional", + params=[TaskHandlerParam(name=None, required=True)], + ) + ] + }, + ), + TaskHandlerArtifact( + bundle_name="java-task-handlers", + relative_fileloc="empty.jar", + size_bytes=1024, + cache_digest=None, + task_handlers={}, + ), + ] + + +@pytest.mark.parametrize( + ("field", "value"), + [ + pytest.param("cache_digest", "a" * 129, id="cache-digest"), + pytest.param("relative_fileloc", "a" * 2001, id="rel-path"), + ], +) +def test_known_artifact_rejects_values_wider_than_their_column(field, value): + with pytest.raises(ValidationError, match=field): + TaskHandlerArtifact.model_validate(_make_known_artifact_body(**{field: value})) + + def test_parse_file_with_dag_callbacks(spy_agency): from airflow import DAG diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index 98b9be109d26f..e7a512de7b574 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -958,6 +958,13 @@ "title": "Callback Requests", "type": "array" }, + "known_artifacts": { + "items": { + "$ref": "#/$defs/TaskHandlerArtifact" + }, + "title": "Known Artifacts", + "type": "array" + }, "type": { "const": "DagFileParseRequest", "default": "DagFileParseRequest", @@ -4156,6 +4163,55 @@ "title": "TaskCallbackRequest", "type": "object" }, + "TaskHandlerArtifact": { + "description": "A Lang-SDK artifact and every task handler it registers.", + "properties": { + "bundle_name": { + "title": "Bundle Name", + "type": "string" + }, + "relative_fileloc": { + "maxLength": 2000, + "title": "Relative Fileloc", + "type": "string" + }, + "size_bytes": { + "title": "Size Bytes", + "type": "integer" + }, + "cache_digest": { + "anyOf": [ + { + "maxLength": 128, + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Cache Digest" + }, + "task_handlers": { + "additionalProperties": { + "items": { + "$ref": "#/$defs/TaskHandlerDeclaration" + }, + "type": "array" + }, + "title": "Task Handlers", + "type": "object" + } + }, + "required": [ + "bundle_name", + "relative_fileloc", + "size_bytes", + "cache_digest", + "task_handlers" + ], + "title": "TaskHandlerArtifact", + "type": "object" + }, "TaskHandlerDeclaration": { "description": "A task handler that a Lang-SDK artifact registers for one task.", "properties": { diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py index c974f797bb0ad..aec4e0874b728 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py @@ -40,6 +40,7 @@ def get_bundle() -> VersionBundle: from airflow.sdk.execution_time.schema.versions.v2026_10_30 import ( AddArgBindingsToSupervisorTIRunContext, AddDagDefinitionsToDagFileParsingResult, + AddKnownArtifactsToDagFileParseRequest, AddRetryReasonToTaskState, ) @@ -50,6 +51,7 @@ def get_bundle() -> VersionBundle: AddArgBindingsToSupervisorTIRunContext, AddRetryReasonToTaskState, AddDagDefinitionsToDagFileParsingResult, + AddKnownArtifactsToDagFileParseRequest, ), Version("2026-06-16"), ) diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py index 4aca18cf50cc0..3b810cbaaa254 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py @@ -28,7 +28,7 @@ from cadwyn import VersionChange, schema -from airflow.dag_processing.processor import DagFileParsingResult # noqa: SDK002 +from airflow.dag_processing.processor import DagFileParseRequest, DagFileParsingResult # noqa: SDK002 from airflow.sdk.api.datamodels._generated import TIRunContext from airflow.sdk.execution_time.comms import TaskState @@ -64,3 +64,13 @@ class AddDagDefinitionsToDagFileParsingResult(VersionChange): schema(DagFileParsingResult).field("parsed_definitions").didnt_exist, schema(DagFileParsingResult).field("dag_source_codes").didnt_exist, ) + + +class AddKnownArtifactsToDagFileParseRequest(VersionChange): + """Add ``known_artifacts``, the recorded task-handler artifacts and their answers, to ``DagFileParseRequest``.""" + + description = __doc__ + + instructions_to_migrate_to_previous_version = ( + schema(DagFileParseRequest).field("known_artifacts").didnt_exist, + ) diff --git a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py index 204ffe7335593..ae698712ea966 100644 --- a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py +++ b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py @@ -30,6 +30,7 @@ import subprocess import sys +from pathlib import Path from typing import Literal import pytest @@ -48,7 +49,13 @@ _SupervisorResponse, ) -from airflow.dag_processing.processor import TaskHandlerParseRequest +from airflow.dag_processing.processor import ( + DagFileParseRequest, + TaskHandlerArtifact, + TaskHandlerDeclaration, + TaskHandlerParam, + TaskHandlerParseRequest, +) from airflow.sdk import TaskInstanceState from airflow.sdk.execution_time.comms import TaskState from airflow.sdk.execution_time.schema import ( @@ -507,3 +514,69 @@ def test_upgrade_fills_missing_retry_reason_with_none(self, real_migrator): body = {"type": "TaskState", "state": "failed", "end_date": None, "rendered_map_index": None} out = real_migrator.upgrade(body, TaskState, "2026-06-16") assert out["retry_reason"] is None + + +class TestRealBundleKnownArtifacts: + """ + Drive the *real* supervisor bundle through the ``known_artifacts`` migration. + + ``DagFileParseRequest`` flows supervisor -> runtime, so ``downgrade`` is the direction a + pinned runtime travels. + """ + + @pytest.fixture + def request_with_known_artifacts(self) -> DagFileParseRequest: + return DagFileParseRequest( + file="/files/dags/etl.py", + bundle_path=Path("/files/dags"), + bundle_name="dags-folder", + known_artifacts=[ + TaskHandlerArtifact( + bundle_name="java-task-handlers", + relative_fileloc="etl.jar", + size_bytes=1024, + cache_digest="ab12", + task_handlers={ + "etl": [ + TaskHandlerDeclaration( + task_id="extract", + binding="positional", + params=[TaskHandlerParam(name=None, required=True)], + ) + ] + }, + ) + ], + ) + + @pytest.fixture + def real_migrator(self) -> SchemaVersionMigrator: + return get_schema_version_migrator() + + def test_downgrade_strips_known_artifacts_for_previous_version( + self, real_migrator, request_with_known_artifacts + ): + out = real_migrator.downgrade(request_with_known_artifacts, "2026-06-16").model_dump() + assert "known_artifacts" not in out + + def test_downgrade_keeps_known_artifacts_at_head(self, real_migrator, request_with_known_artifacts): + out = real_migrator.downgrade(request_with_known_artifacts, "2026-10-30").model_dump() + assert out["known_artifacts"] == [ + { + "bundle_name": "java-task-handlers", + "relative_fileloc": "etl.jar", + "size_bytes": 1024, + "cache_digest": "ab12", + "task_handlers": { + "etl": [ + { + "task_id": "extract", + "binding": "positional", + "params": [ + {"name": None, "value_schema": None, "required": True, "exact_name": False} + ], + } + ] + }, + } + ] From 90cfe6adaf3cbbf7e29c1094432f2f51149fed4a Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 18:22:42 +0000 Subject: [PATCH 88/94] Regenerate TS SDK supervisor types for known artifacts check-ts-sdk-supervisor-schema fails whenever supervisor.ts lags the snapshot, and CI runs it over all files on any ts-sdk change. The shifted numeric suffixes are the generator's own renumbering; no hand-written TS code refers to them. --- ts-sdk/src/generated/supervisor.ts | 148 +++++++++++++++++------------ 1 file changed, 85 insertions(+), 63 deletions(-) diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index 4567d65a39835..af0c907a5ed05 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -269,6 +269,16 @@ export type Msg2 = string | null; export type EmailType = "failure" | "retry"; export type Type14 = "EmailRequest"; export type CallbackRequests = (DagCallbackRequest | TaskCallbackRequest | EmailRequest)[]; +export type BundleName4 = string; +export type RelativeFileloc = string; +export type SizeBytes = number; +export type CacheDigest = string | null; +export type TaskId2 = string; +export type Binding = "positional" | "named"; +export type TaskHandlerParams = TaskHandlerParam[] | null; +export type Name10 = string | null; +export type ExactName = boolean; +export type KnownArtifacts = TaskHandlerArtifact[]; export type Type15 = "DagFileParseRequest"; export type Fileloc = string; export type LastLoaded = string | null; @@ -283,9 +293,9 @@ export type Language = string; export type Type16 = "DagFileParsingResult"; export type DagId4 = string; export type IsPaused = boolean; -export type BundleName4 = string | null; +export type BundleName5 = string | null; export type BundleVersion3 = string | null; -export type RelativeFileloc = string | null; +export type RelativeFileloc1 = string | null; export type Owners = string | null; export type Tags = string[]; export type NextDagrun = string | null; @@ -326,7 +336,7 @@ export type NextKwargs2 = { } | null; export type RenderedMapIndex1 = string | null; export type Type20 = "DeferTask"; -export type Name10 = string; +export type Name11 = string; export type Key1 = string; export type Type21 = "DeleteAssetStateStoreByName"; export type Uri5 = string; @@ -340,7 +350,7 @@ export type Type24 = "DeleteVariable"; export type Key5 = string; export type DagId6 = string; export type RunId5 = string; -export type TaskId2 = string; +export type TaskId3 = string; export type MapIndex1 = number | null; export type Type25 = "DeleteXCom"; /** @@ -378,11 +388,11 @@ export type ErrorType1 = | "PERMISSION_DENIED" | "GENERIC_ERROR" | "API_SERVER_ERROR"; -export type Name11 = string; +export type Name12 = string; export type Type27 = "GetAssetByName"; export type Uri6 = string; export type Type28 = "GetAssetByUri"; -export type Name12 = string | null; +export type Name13 = string | null; export type Uri7 = string | null; export type After = string | null; export type Before = string | null; @@ -405,7 +415,7 @@ export type Extra8 = { [k: string]: string; } | null; export type Type30 = "GetAssetEventByAssetAlias"; -export type Name13 = string; +export type Name14 = string; export type Key6 = string; export type Type31 = "GetAssetStateStoreByName"; export type Uri8 = string; @@ -437,7 +447,7 @@ export type LogicalDate3 = string; export type State3 = string | null; export type Type41 = "GetPreviousDagRun"; export type DagId12 = string; -export type TaskId3 = string; +export type TaskId4 = string; export type LogicalDate4 = string | null; export type MapIndex2 = number; export type Type42 = "GetPreviousTI"; @@ -474,25 +484,25 @@ export type Type49 = "GetVariableKeys"; export type Key10 = string; export type DagId16 = string; export type RunId9 = string; -export type TaskId4 = string; +export type TaskId5 = string; export type MapIndex5 = number | null; export type IncludePriorDates = boolean; export type Type50 = "GetXCom"; export type Key11 = string; export type DagId17 = string; export type RunId10 = string; -export type TaskId5 = string; +export type TaskId6 = string; export type Type51 = "GetXComCount"; export type Key12 = string; export type DagId18 = string; export type RunId11 = string; -export type TaskId6 = string; +export type TaskId7 = string; export type Offset1 = number; export type Type52 = "GetXComSequenceItem"; export type Key13 = string; export type DagId19 = string; export type RunId12 = string; -export type TaskId7 = string; +export type TaskId8 = string; export type Start = number | null; export type Stop = number | null; export type Step = number | null; @@ -514,7 +524,7 @@ export type AssignedUsers1 = HITLUser[] | null; export type Type54 = "HITLDetailRequestResult"; export type InactiveAssets = AssetProfile[] | null; export type Type55 = "InactiveAssetsResult"; -export type Name14 = string | null; +export type Name15 = string | null; export type Type56 = "MaskSecret"; export type Ok = boolean; export type Type57 = "OKResponse"; @@ -524,7 +534,7 @@ export type StartDate4 = string | null; export type EndDate3 = string | null; export type Type58 = "PrevSuccessfulDagRunResult"; export type Type59 = "PreviousDagRunResult"; -export type TaskId8 = string; +export type TaskId9 = string; export type DagId20 = string; export type RunId13 = string; export type LogicalDate5 = string | null; @@ -552,7 +562,7 @@ export type RetryReason = string | null; export type Type64 = "RetryTask"; export type Type65 = "SentFDs"; export type Fds = number[]; -export type Name15 = string; +export type Name16 = string; export type Key15 = string; export type Type66 = "SetAssetStateStoreByName"; export type Uri9 = string; @@ -568,7 +578,7 @@ export type Type70 = "SetTaskStateStore"; export type Key18 = string; export type DagId21 = string; export type RunId14 = string; -export type TaskId9 = string; +export type TaskId10 = string; export type MapIndex7 = number | null; export type DagResult1 = boolean; export type MappedLength = number | null; @@ -595,14 +605,9 @@ export type Breadcrumbs = { [k: string]: unknown; }[]; export type Type76 = "TaskBreadcrumbsResult"; -export type TaskId10 = string; -export type Binding = "positional" | "named"; -export type TaskHandlerParams = TaskHandlerParam[] | null; -export type Name16 = string | null; -export type ExactName = boolean; export type File1 = string; export type BundlePath1 = string; -export type BundleName5 = string; +export type BundleName6 = string; export type Type77 = "TaskHandlerParseRequest"; export type Fileloc1 = string; export type ImportErrors1 = { @@ -1015,6 +1020,7 @@ export interface DagFileParseRequest { bundle_path: BundlePath; bundle_name: BundleName1; callback_requests?: CallbackRequests; + known_artifacts?: KnownArtifacts; type?: Type15; } /** @@ -1116,6 +1122,44 @@ export interface EmailRequest { context_from_server: TIRunContext; type?: Type14; } +/** + * A Lang-SDK artifact and every task handler it registers. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TaskHandlerArtifact". + */ +export interface TaskHandlerArtifact { + bundle_name: BundleName4; + relative_fileloc: RelativeFileloc; + size_bytes: SizeBytes; + cache_digest: CacheDigest; + task_handlers: TaskHandlers; +} +export interface TaskHandlers { + [k: string]: TaskHandlerDeclaration[]; +} +/** + * A task handler that a Lang-SDK artifact registers for one task. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TaskHandlerDeclaration". + */ +export interface TaskHandlerDeclaration { + task_id: TaskId2; + binding: Binding; + params: TaskHandlerParams; +} +/** + * One parameter of a task handler. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TaskHandlerParam". + */ +export interface TaskHandlerParam { + name: Name10; + value_schema?: ArgValueSchema | null; + exact_name?: ExactName; +} /** * Result of DAG File Parsing. * @@ -1170,9 +1214,9 @@ export interface DagSourceCode { export interface DagResult { dag_id: DagId4; is_paused: IsPaused; - bundle_name: BundleName4; + bundle_name: BundleName5; bundle_version: BundleVersion3; - relative_fileloc: RelativeFileloc; + relative_fileloc: RelativeFileloc1; owners: Owners; tags: Tags; next_dagrun: NextDagrun; @@ -1233,7 +1277,7 @@ export interface DeferTask { * via the `definition` "DeleteAssetStateStoreByName". */ export interface DeleteAssetStateStoreByName { - name: Name10; + name: Name11; key: Key1; type?: Type21; } @@ -1271,7 +1315,7 @@ export interface DeleteXCom { key: Key5; dag_id: DagId6; run_id: RunId5; - task_id: TaskId2; + task_id: TaskId3; map_index?: MapIndex1; type?: Type25; } @@ -1289,7 +1333,7 @@ export interface ErrorResponse { * via the `definition` "GetAssetByName". */ export interface GetAssetByName { - name: Name11; + name: Name12; type?: Type27; } /** @@ -1305,7 +1349,7 @@ export interface GetAssetByUri { * via the `definition` "GetAssetEventByAsset". */ export interface GetAssetEventByAsset { - name: Name12; + name: Name13; uri: Uri7; after?: After; before?: Before; @@ -1336,7 +1380,7 @@ export interface GetAssetEventByAssetAlias { * via the `definition` "GetAssetStateStoreByName". */ export interface GetAssetStateStoreByName { - name: Name13; + name: Name14; key: Key6; type?: Type31; } @@ -1438,7 +1482,7 @@ export interface GetPreviousDagRun { */ export interface GetPreviousTI { dag_id: DagId12; - task_id: TaskId3; + task_id: TaskId4; logical_date?: LogicalDate4; map_index?: MapIndex2; state?: TaskInstanceState | null; @@ -1524,7 +1568,7 @@ export interface GetXCom { key: Key10; dag_id: DagId16; run_id: RunId9; - task_id: TaskId4; + task_id: TaskId5; map_index?: MapIndex5; include_prior_dates?: IncludePriorDates; type?: Type50; @@ -1539,7 +1583,7 @@ export interface GetXComCount { key: Key11; dag_id: DagId17; run_id: RunId10; - task_id: TaskId5; + task_id: TaskId6; type?: Type51; } /** @@ -1550,7 +1594,7 @@ export interface GetXComSequenceItem { key: Key12; dag_id: DagId18; run_id: RunId11; - task_id: TaskId6; + task_id: TaskId7; offset: Offset1; type?: Type52; } @@ -1562,7 +1606,7 @@ export interface GetXComSequenceSlice { key: Key13; dag_id: DagId19; run_id: RunId12; - task_id: TaskId7; + task_id: TaskId8; start: Start; stop: Stop; step: Step; @@ -1604,7 +1648,7 @@ export interface InactiveAssetsResult { */ export interface MaskSecret { value: JsonValue; - name?: Name14; + name?: Name15; type?: Type56; } /** @@ -1643,7 +1687,7 @@ export interface PreviousDagRunResult { * via the `definition` "PreviousTIResponse". */ export interface PreviousTIResponse { - task_id: TaskId8; + task_id: TaskId9; dag_id: DagId20; run_id: RunId13; logical_date?: LogicalDate5; @@ -1720,7 +1764,7 @@ export interface SentFDs { * via the `definition` "SetAssetStateStoreByName". */ export interface SetAssetStateStoreByName { - name: Name15; + name: Name16; key: Key15; value: JsonValue; type?: Type66; @@ -1778,7 +1822,7 @@ export interface SetXCom { value: JsonValue; dag_id: DagId21; run_id: RunId14; - task_id: TaskId9; + task_id: TaskId10; map_index?: MapIndex7; dag_result?: DagResult1; mapped_length?: MappedLength; @@ -1839,28 +1883,6 @@ export interface TaskBreadcrumbsResult { breadcrumbs: Breadcrumbs; type?: Type76; } -/** - * A task handler that a Lang-SDK artifact registers for one task. - * - * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema - * via the `definition` "TaskHandlerDeclaration". - */ -export interface TaskHandlerDeclaration { - task_id: TaskId10; - binding: Binding; - params: TaskHandlerParams; -} -/** - * One parameter of a task handler. - * - * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema - * via the `definition` "TaskHandlerParam". - */ -export interface TaskHandlerParam { - name: Name16; - value_schema?: ArgValueSchema | null; - exact_name?: ExactName; -} /** * Request for Task Handler Parsing. * @@ -1872,7 +1894,7 @@ export interface TaskHandlerParam { export interface TaskHandlerParseRequest { file: File1; bundle_path: BundlePath1; - bundle_name: BundleName5; + bundle_name: BundleName6; type?: Type77; } /** @@ -1887,12 +1909,12 @@ export interface TaskHandlerParseRequest { */ export interface TaskHandlerParsingResult { fileloc: Fileloc1; - task_handlers: TaskHandlers; + task_handlers: TaskHandlers1; import_errors?: ImportErrors1; warnings?: Warnings1; type?: Type78; } -export interface TaskHandlers { +export interface TaskHandlers1 { [k: string]: TaskHandlerDeclaration[]; } /** From f858597052c73d7fe6dc5ca4b67e85c8b81b8411 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 18:23:28 +0000 Subject: [PATCH 89/94] Sync the Java SDK supervisor schema for known artifacts The vendored copy tracks the monorepo snapshot while both declare the same api_version, and sync-java-sdk-supervisor-schema rewrites it on the next Java SDK change otherwise. --- java-sdk/sdk/schema/schema.json | 56 +++++++++++++++++++++++++++++++++ 1 file changed, 56 insertions(+) diff --git a/java-sdk/sdk/schema/schema.json b/java-sdk/sdk/schema/schema.json index 98b9be109d26f..e7a512de7b574 100644 --- a/java-sdk/sdk/schema/schema.json +++ b/java-sdk/sdk/schema/schema.json @@ -958,6 +958,13 @@ "title": "Callback Requests", "type": "array" }, + "known_artifacts": { + "items": { + "$ref": "#/$defs/TaskHandlerArtifact" + }, + "title": "Known Artifacts", + "type": "array" + }, "type": { "const": "DagFileParseRequest", "default": "DagFileParseRequest", @@ -4156,6 +4163,55 @@ "title": "TaskCallbackRequest", "type": "object" }, + "TaskHandlerArtifact": { + "description": "A Lang-SDK artifact and every task handler it registers.", + "properties": { + "bundle_name": { + "title": "Bundle Name", + "type": "string" + }, + "relative_fileloc": { + "maxLength": 2000, + "title": "Relative Fileloc", + "type": "string" + }, + "size_bytes": { + "title": "Size Bytes", + "type": "integer" + }, + "cache_digest": { + "anyOf": [ + { + "maxLength": 128, + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Cache Digest" + }, + "task_handlers": { + "additionalProperties": { + "items": { + "$ref": "#/$defs/TaskHandlerDeclaration" + }, + "type": "array" + }, + "title": "Task Handlers", + "type": "object" + } + }, + "required": [ + "bundle_name", + "relative_fileloc", + "size_bytes", + "cache_digest", + "task_handlers" + ], + "title": "TaskHandlerArtifact", + "type": "object" + }, "TaskHandlerDeclaration": { "description": "A task handler that a Lang-SDK artifact registers for one task.", "properties": { From 703d64f022254ab56b000b04934d8c44316617f2 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Fri, 2 Oct 2026 02:19:08 +0000 Subject: [PATCH 90/94] Regenerate Go SDK models for known artifacts check-go-sdk-generated-drift regenerates the Go SDK's coordinator-protocol models from the supervisor schema snapshot and fails on any difference, so DagFileParseRequest must carry known_artifacts on the Go side too. go-jsonschema moves some type blocks, so the diff is larger than the new field. --- go-sdk/pkg/execution/genmodels/models.gen.go | 339 ++++++++++--------- 1 file changed, 180 insertions(+), 159 deletions(-) diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go b/go-sdk/pkg/execution/genmodels/models.gen.go index 0b0b4eed57d51..a6b4542757690 100644 --- a/go-sdk/pkg/execution/genmodels/models.gen.go +++ b/go-sdk/pkg/execution/genmodels/models.gen.go @@ -397,6 +397,9 @@ type DagFileParseRequest struct { // File corresponds to the JSON schema field "file". File string `msgpack:"file"` + // KnownArtifacts corresponds to the JSON schema field "known_artifacts". + KnownArtifacts []TaskHandlerArtifact `msgpack:"known_artifacts,omitempty"` + // Type corresponds to the JSON schema field "type". Type string `msgpack:"type,omitempty"` } @@ -888,86 +891,6 @@ type GetAssetEventByAssetAlias struct { Type string `msgpack:"type,omitempty"` } -// Variable schema for responses with fields that are needed for Runtime. -type VariableResponse struct { - // Key corresponds to the JSON schema field "key". - Key string `msgpack:"key"` - - // Value corresponds to the JSON schema field "value". - Value interface{} `msgpack:"value"` -} - -type TaskArgBinding interface{} - -type TriggerKwargs map[string]interface{} - -type Warnings []interface{} - -// Lazily build information from the serialized DAG structure. -// -// An object that will present "enough" of the DAG like interface to update DAG db -// models etc, without having -// to deserialize the full DAG and Task hierarchy. -type LazyDeserializedDAG struct { - // Data corresponds to the JSON schema field "data". - Data Data `msgpack:"data"` - - // LastLoaded corresponds to the JSON schema field "last_loaded". - LastLoaded interface{} `msgpack:"last_loaded,omitempty"` -} - -type ImportErrors map[string]string - -// Schema for TaskInstance model with minimal required fields needed for Runtime. -type TaskInstance struct { - // ContextCarrier corresponds to the JSON schema field "context_carrier". - ContextCarrier *ContextCarrier `msgpack:"context_carrier,omitempty"` - - // DagID corresponds to the JSON schema field "dag_id". - DagID string `msgpack:"dag_id"` - - // DagVersionID corresponds to the JSON schema field "dag_version_id". - DagVersionID string `msgpack:"dag_version_id"` - - // Hostname corresponds to the JSON schema field "hostname". - Hostname interface{} `msgpack:"hostname,omitempty"` - - // ID corresponds to the JSON schema field "id". - ID string `msgpack:"id"` - - // MapIndex corresponds to the JSON schema field "map_index". - MapIndex *int `msgpack:"map_index,omitempty"` - - // Queue corresponds to the JSON schema field "queue". - Queue string `msgpack:"queue,omitempty"` - - // RunID corresponds to the JSON schema field "run_id". - RunID string `msgpack:"run_id"` - - // TaskID corresponds to the JSON schema field "task_id". - TaskID string `msgpack:"task_id"` - - // TryNumber corresponds to the JSON schema field "try_number". - TryNumber int `msgpack:"try_number"` -} - -type Params map[string]interface{} - -// Schema for a Human-in-the-loop users. -type HITLUser struct { - // ID corresponds to the JSON schema field "id". - ID string `msgpack:"id"` - - // Name corresponds to the JSON schema field "name". - Name string `msgpack:"name"` -} - -type VersionData map[string]interface{} - -type NextKwargs map[string]interface{} - -type JsonValue interface{} - type GetAssetStateStoreByName struct { // Key corresponds to the JSON schema field "key". Key string `msgpack:"key"` @@ -1006,12 +929,6 @@ type GetConnection struct { Type string `msgpack:"type,omitempty"` } -type LogicalDates []time.Time - -type RunIds []string - -type States []string - type GetDRCount struct { // DagID corresponds to the JSON schema field "dag_id". DagID string `msgpack:"dag_id"` @@ -1090,22 +1007,6 @@ type GetPreviousDagRun struct { Type string `msgpack:"type,omitempty"` } -type TaskInstanceState string - -const TaskInstanceStateRemoved TaskInstanceState = "removed" -const TaskInstanceStateScheduled TaskInstanceState = "scheduled" -const TaskInstanceStateQueued TaskInstanceState = "queued" -const TaskInstanceStateRunning TaskInstanceState = "running" -const TaskInstanceStateSuccess TaskInstanceState = "success" -const TaskInstanceStateRestarting TaskInstanceState = "restarting" -const TaskInstanceStateFailed TaskInstanceState = "failed" -const TaskInstanceStateUpForRetry TaskInstanceState = "up_for_retry" -const TaskInstanceStateUpForReschedule TaskInstanceState = "up_for_reschedule" -const TaskInstanceStateUpstreamFailed TaskInstanceState = "upstream_failed" -const TaskInstanceStateSkipped TaskInstanceState = "skipped" -const TaskInstanceStateDeferred TaskInstanceState = "deferred" -const TaskInstanceStateAwaitingInput TaskInstanceState = "awaiting_input" - // Request to get previous task instance. type GetPreviousTI struct { // DagID corresponds to the JSON schema field "dag_id". @@ -1127,8 +1028,6 @@ type GetPreviousTI struct { Type string `msgpack:"type,omitempty"` } -type TaskIds []string - type GetTICount struct { // DagID corresponds to the JSON schema field "dag_id". DagID string `msgpack:"dag_id"` @@ -1353,6 +1252,17 @@ type HITLDetailRequestResult struct { Type string `msgpack:"type,omitempty"` } +// Schema for a Human-in-the-loop users. +type HITLUser struct { + // ID corresponds to the JSON schema field "id". + ID string `msgpack:"id"` + + // Name corresponds to the JSON schema field "name". + Name string `msgpack:"name"` +} + +type ImportErrors map[string]string + type InactiveAssets []AssetProfile // Response of InactiveAssets requests. @@ -1364,44 +1274,19 @@ type InactiveAssetsResult struct { Type string `msgpack:"type,omitempty"` } -// Response schema for TaskInstance run context. -type TIRunContext struct { - // ArgBindings corresponds to the JSON schema field "arg_bindings". - ArgBindings *ArgBindings `msgpack:"arg_bindings,omitempty"` - - // Connections corresponds to the JSON schema field "connections". - Connections []ConnectionResponse `msgpack:"connections,omitempty"` - - // DagRun corresponds to the JSON schema field "dag_run". - DagRun DagRun `msgpack:"dag_run"` - - // MaxTries corresponds to the JSON schema field "max_tries". - MaxTries int `msgpack:"max_tries"` - - // MultiTeam corresponds to the JSON schema field "multi_team". - MultiTeam bool `msgpack:"multi_team,omitempty"` - - // NextKwargs corresponds to the JSON schema field "next_kwargs". - NextKwargs *NextKwargs `msgpack:"next_kwargs,omitempty"` - - // NextMethod corresponds to the JSON schema field "next_method". - NextMethod interface{} `msgpack:"next_method,omitempty"` - - // ShouldRetry corresponds to the JSON schema field "should_retry". - ShouldRetry bool `msgpack:"should_retry,omitempty"` - - // StartDate corresponds to the JSON schema field "start_date". - StartDate interface{} `msgpack:"start_date,omitempty"` - - // TaskRescheduleCount corresponds to the JSON schema field - // "task_reschedule_count". - TaskRescheduleCount int `msgpack:"task_reschedule_count,omitempty"` +type JsonValue interface{} - // Variables corresponds to the JSON schema field "variables". - Variables []VariableResponse `msgpack:"variables,omitempty"` +// Lazily build information from the serialized DAG structure. +// +// An object that will present "enough" of the DAG like interface to update DAG db +// models etc, without having +// to deserialize the full DAG and Task hierarchy. +type LazyDeserializedDAG struct { + // Data corresponds to the JSON schema field "data". + Data Data `msgpack:"data"` - // XcomKeysToClear corresponds to the JSON schema field "xcom_keys_to_clear". - XcomKeysToClear []string `msgpack:"xcom_keys_to_clear,omitempty"` + // LastLoaded corresponds to the JSON schema field "last_loaded". + LastLoaded interface{} `msgpack:"last_loaded,omitempty"` } // One positional stub-task argument carrying an inline literal from the Dag file. @@ -1422,6 +1307,8 @@ type LiteralArgBinding struct { ValueSchema *ArgValueSchema `msgpack:"value_schema,omitempty"` } +type LogicalDates []time.Time + // Add a new value to be redacted in task logs. type MaskSecret struct { // Name corresponds to the JSON schema field "name". @@ -1434,6 +1321,8 @@ type MaskSecret struct { Value JsonValue `msgpack:"value"` } +type NextKwargs map[string]interface{} + type OKResponse struct { // Ok corresponds to the JSON schema field "ok". Ok bool `msgpack:"ok"` @@ -1442,6 +1331,12 @@ type OKResponse struct { Type string `msgpack:"type,omitempty"` } +type OutletEvents []map[string]interface{} + +type Params map[string]interface{} + +type ParamsInput map[string]interface{} + type PrevSuccessfulDagRunResult struct { // DataIntervalEnd corresponds to the JSON schema field "data_interval_end". DataIntervalEnd interface{} `msgpack:"data_interval_end,omitempty"` @@ -1524,6 +1419,8 @@ type PutVariable struct { Value interface{} `msgpack:"value"` } +type RenderedFields map[string]JsonValue + // Update a task instance state to reschedule/up_for_reschedule. type RescheduleTask struct { // EndDate corresponds to the JSON schema field "end_date". @@ -1565,6 +1462,8 @@ type RetryTask struct { Type string `msgpack:"type,omitempty"` } +type RunIds []string + type SentFDs struct { // Fds corresponds to the JSON schema field "fds". Fds []int `msgpack:"fds"` @@ -1601,8 +1500,6 @@ type SetAssetStateStoreByURI struct { Value JsonValue `msgpack:"value"` } -type RenderedFields map[string]JsonValue - // Payload for setting RTIF for a task instance. type SetRenderedFields struct { // RenderedFields corresponds to the JSON schema field "rendered_fields". @@ -1700,9 +1597,7 @@ type StartupDetails struct { Type string `msgpack:"type,omitempty"` } -type OutletEvents []map[string]interface{} - -type TaskOutlets []AssetProfile +type States []string // Update a task's state to success. Includes task_outlets and outlet_events for // registering asset events. @@ -1735,7 +1630,47 @@ type TICount struct { Type string `msgpack:"type,omitempty"` } -type TaskBreadcrumbsResultBreadcrumbsElem map[string]interface{} +// Response schema for TaskInstance run context. +type TIRunContext struct { + // ArgBindings corresponds to the JSON schema field "arg_bindings". + ArgBindings *ArgBindings `msgpack:"arg_bindings,omitempty"` + + // Connections corresponds to the JSON schema field "connections". + Connections []ConnectionResponse `msgpack:"connections,omitempty"` + + // DagRun corresponds to the JSON schema field "dag_run". + DagRun DagRun `msgpack:"dag_run"` + + // MaxTries corresponds to the JSON schema field "max_tries". + MaxTries int `msgpack:"max_tries"` + + // MultiTeam corresponds to the JSON schema field "multi_team". + MultiTeam bool `msgpack:"multi_team,omitempty"` + + // NextKwargs corresponds to the JSON schema field "next_kwargs". + NextKwargs *NextKwargs `msgpack:"next_kwargs,omitempty"` + + // NextMethod corresponds to the JSON schema field "next_method". + NextMethod interface{} `msgpack:"next_method,omitempty"` + + // ShouldRetry corresponds to the JSON schema field "should_retry". + ShouldRetry bool `msgpack:"should_retry,omitempty"` + + // StartDate corresponds to the JSON schema field "start_date". + StartDate interface{} `msgpack:"start_date,omitempty"` + + // TaskRescheduleCount corresponds to the JSON schema field + // "task_reschedule_count". + TaskRescheduleCount int `msgpack:"task_reschedule_count,omitempty"` + + // Variables corresponds to the JSON schema field "variables". + Variables []VariableResponse `msgpack:"variables,omitempty"` + + // XcomKeysToClear corresponds to the JSON schema field "xcom_keys_to_clear". + XcomKeysToClear []string `msgpack:"xcom_keys_to_clear,omitempty"` +} + +type TaskArgBinding interface{} type TaskBreadcrumbsResult struct { // Breadcrumbs corresponds to the JSON schema field "breadcrumbs". @@ -1745,6 +1680,8 @@ type TaskBreadcrumbsResult struct { Type string `msgpack:"type,omitempty"` } +type TaskBreadcrumbsResultBreadcrumbsElem map[string]interface{} + // Task callback status information. // // A Class with information about the success/failure TI callback to be executed. @@ -1780,10 +1717,43 @@ type TaskCallbackRequest struct { VersionData *VersionData `msgpack:"version_data,omitempty"` } -type TaskHandlerDeclarationBinding string +type Warnings []interface{} -const TaskHandlerDeclarationBindingPositional TaskHandlerDeclarationBinding = "positional" -const TaskHandlerDeclarationBindingNamed TaskHandlerDeclarationBinding = "named" +type TriggerKwargs map[string]interface{} + +// A Lang-SDK artifact and every task handler it registers. +type TaskHandlerArtifact struct { + // BundleName corresponds to the JSON schema field "bundle_name". + BundleName string `msgpack:"bundle_name"` + + // CacheDigest corresponds to the JSON schema field "cache_digest". + CacheDigest interface{} `msgpack:"cache_digest"` + + // RelativeFileloc corresponds to the JSON schema field "relative_fileloc". + RelativeFileloc string `msgpack:"relative_fileloc"` + + // SizeBytes corresponds to the JSON schema field "size_bytes". + SizeBytes int `msgpack:"size_bytes"` + + // TaskHandlers corresponds to the JSON schema field "task_handlers". + TaskHandlers TaskHandlers `msgpack:"task_handlers"` +} + +type TaskHandlers map[string][]TaskHandlerDeclaration + +// A task handler that a Lang-SDK artifact registers for one task. +type TaskHandlerDeclaration struct { + // Binding corresponds to the JSON schema field "binding". + Binding TaskHandlerDeclarationBinding `msgpack:"binding"` + + // Params corresponds to the JSON schema field "params". + Params *TaskHandlerParams `msgpack:"params"` + + // TaskID corresponds to the JSON schema field "task_id". + TaskID string `msgpack:"task_id"` +} + +type TaskHandlerParams []TaskHandlerParam // One parameter of a task handler. type TaskHandlerParam struct { @@ -1797,18 +1767,71 @@ type TaskHandlerParam struct { ValueSchema *ArgValueSchema `msgpack:"value_schema,omitempty"` } -type TaskHandlerParams []TaskHandlerParam +type TaskIds []string -// A task handler that a Lang-SDK artifact registers for one task. -type TaskHandlerDeclaration struct { - // Binding corresponds to the JSON schema field "binding". - Binding TaskHandlerDeclarationBinding `msgpack:"binding"` +const TaskInstanceStateAwaitingInput TaskInstanceState = "awaiting_input" +const TaskInstanceStateDeferred TaskInstanceState = "deferred" +const TaskInstanceStateSkipped TaskInstanceState = "skipped" +const TaskInstanceStateUpstreamFailed TaskInstanceState = "upstream_failed" +const TaskInstanceStateUpForReschedule TaskInstanceState = "up_for_reschedule" +const TaskInstanceStateUpForRetry TaskInstanceState = "up_for_retry" +const TaskInstanceStateFailed TaskInstanceState = "failed" +const TaskInstanceStateRestarting TaskInstanceState = "restarting" +const TaskInstanceStateSuccess TaskInstanceState = "success" +const TaskInstanceStateRunning TaskInstanceState = "running" +const TaskInstanceStateQueued TaskInstanceState = "queued" +const TaskInstanceStateScheduled TaskInstanceState = "scheduled" +const TaskInstanceStateRemoved TaskInstanceState = "removed" - // Params corresponds to the JSON schema field "params". - Params *TaskHandlerParams `msgpack:"params"` +type TaskInstanceState string + +// Variable schema for responses with fields that are needed for Runtime. +type VariableResponse struct { + // Key corresponds to the JSON schema field "key". + Key string `msgpack:"key"` + + // Value corresponds to the JSON schema field "value". + Value interface{} `msgpack:"value"` +} + +const TaskHandlerDeclarationBindingNamed TaskHandlerDeclarationBinding = "named" +const TaskHandlerDeclarationBindingPositional TaskHandlerDeclarationBinding = "positional" + +type TaskHandlerDeclarationBinding string + +type TaskOutlets []AssetProfile + +// Schema for TaskInstance model with minimal required fields needed for Runtime. +type TaskInstance struct { + // ContextCarrier corresponds to the JSON schema field "context_carrier". + ContextCarrier *ContextCarrier `msgpack:"context_carrier,omitempty"` + + // DagID corresponds to the JSON schema field "dag_id". + DagID string `msgpack:"dag_id"` + + // DagVersionID corresponds to the JSON schema field "dag_version_id". + DagVersionID string `msgpack:"dag_version_id"` + + // Hostname corresponds to the JSON schema field "hostname". + Hostname interface{} `msgpack:"hostname,omitempty"` + + // ID corresponds to the JSON schema field "id". + ID string `msgpack:"id"` + + // MapIndex corresponds to the JSON schema field "map_index". + MapIndex *int `msgpack:"map_index,omitempty"` + + // Queue corresponds to the JSON schema field "queue". + Queue string `msgpack:"queue,omitempty"` + + // RunID corresponds to the JSON schema field "run_id". + RunID string `msgpack:"run_id"` // TaskID corresponds to the JSON schema field "task_id". TaskID string `msgpack:"task_id"` + + // TryNumber corresponds to the JSON schema field "try_number". + TryNumber int `msgpack:"try_number"` } // Request for Task Handler Parsing. @@ -1828,8 +1851,6 @@ type TaskHandlerParseRequest struct { Type string `msgpack:"type,omitempty"` } -type TaskHandlers map[string][]TaskHandlerDeclaration - // Result of Task Handler Parsing. // // Every task handler a Lang-SDK artifact registers, keyed by Dag id. @@ -1950,7 +1971,7 @@ type UpdateDagRunNote struct { Type string `msgpack:"type,omitempty"` } -type ParamsInput map[string]interface{} +type VersionData map[string]interface{} // Update the response content part of an existing Human-in-the-loop response. type UpdateHITLDetail struct { From 27370d053fc27ff6db74ff1819d60e55c849225a Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 18:29:40 +0000 Subject: [PATCH 91/94] Send the Dag-parsing child its known task-handler artifacts The fast path can only skip a probe if the child knows what the Dag processor already recorded, answers included. Reading once per loop, not once per pass, lets children started later see rows persisted by earlier ones, so a cold start does not re-probe every artifact for every file. An answer one file recorded is trusted by every file that reads it, so a child only gets the named bundles of its own team, plus its own Dag bundle when a coordinator falls back to it; never another Dag bundle, which keeps the request small. A broken [sdk] coordinators value must not stop Python-only parsing. --- .../src/airflow/dag_processing/manager.py | 122 ++++++- .../tests/unit/dag_processing/test_manager.py | 317 ++++++++++++++++++ 2 files changed, 436 insertions(+), 3 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index 8c4fd0588eb87..fa8705493215c 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -39,6 +39,7 @@ import attrs import structlog +from pydantic import ValidationError from sqlalchemy import select, update from sqlalchemy.exc import OperationalError from sqlalchemy.orm import load_only @@ -55,7 +56,11 @@ ) from airflow.dag_processing.bundles.manager import DagBundlesManager from airflow.dag_processing.collection import update_dag_parsing_results_in_db -from airflow.dag_processing.processor import DagFileParsingResult, DagFileProcessorProcess +from airflow.dag_processing.processor import ( + DagFileParsingResult, + DagFileProcessorProcess, + TaskHandlerArtifact, +) from airflow.models.asset import remove_references_to_deleted_dags from airflow.models.dag import DagModel from airflow.models.dagbag import DagPriorityParsingRequest @@ -63,8 +68,10 @@ from airflow.models.dagwarning import DagWarning from airflow.models.db_callback_request import DbCallbackRequest from airflow.models.errors import ParseImportError +from airflow.models.lang_sdk_task_handler import LangSDKTaskHandlerArtifact from airflow.observability.metrics import stats_utils from airflow.sdk import SecretCache +from airflow.sdk.execution_time.coordinator import get_coordinator_manager # noqa: SDK001 from airflow.sdk.log import init_log_file, logging_processors from airflow.typing_compat import assert_never from airflow.utils.file import list_py_file_paths, might_contain_dag @@ -151,6 +158,16 @@ def normalized_file_path_for_stats(self) -> str: return normalize_name_for_stats(str(self.rel_path), log_warning=False) +class _TaskHandlerBundles(NamedTuple): + """Where the coordinators that queues route to read task-handler artifacts from.""" + + named: frozenset[str] + """Bundles named by a coordinator's ``task_handler_bundle_name``.""" + + include_dag_bundles: bool + """Whether a coordinator without ``task_handler_bundle_name`` reads the task's own Dag bundle.""" + + def _config_int_factory(section: str, key: str): return functools.partial(conf.getint, section, key) @@ -1444,7 +1461,97 @@ def client(self) -> Client: client.base_url = "http://in-process.invalid./" return client - def _create_process(self, dag_file: DagFileInfo) -> DagFileProcessorProcess: + @functools.cached_property + def _task_handler_bundles(self) -> _TaskHandlerBundles: + try: + bundle_names = get_coordinator_manager().get_task_handler_bundle_names().values() + except Exception: + # A Lang-SDK configuration error must not stop Python-only parsing. + self.log.exception( + "Cannot read [sdk] coordinators; Dag files are parsed without known task-handler artifacts" + ) + return _TaskHandlerBundles(named=frozenset(), include_dag_bundles=False) + return _TaskHandlerBundles( + named=frozenset(name for name in bundle_names if name is not None), + include_dag_bundles=None in bundle_names, + ) + + def get_known_task_handler_artifacts( + self, bundle_names: Collection[str] + ) -> dict[str, list[TaskHandlerArtifact]]: + """ + Return the recorded task-handler artifacts in *bundle_names*, keyed by bundle name. + + Default implementation reads from the metadata DB; override to source them from an API. + """ + return self._get_known_task_handler_artifacts_from_db(bundle_names) + + @provide_session + def _get_known_task_handler_artifacts_from_db( + self, bundle_names: Collection[str], *, session: Session = NEW_SESSION + ) -> dict[str, list[TaskHandlerArtifact]]: + rows = session.execute( + select( + LangSDKTaskHandlerArtifact.bundle_name, + LangSDKTaskHandlerArtifact.relative_fileloc, + LangSDKTaskHandlerArtifact.size_bytes, + LangSDKTaskHandlerArtifact.cache_digest, + LangSDKTaskHandlerArtifact.task_handlers, + ) + .where(LangSDKTaskHandlerArtifact.bundle_name.in_(bundle_names)) + .order_by(LangSDKTaskHandlerArtifact.bundle_name, LangSDKTaskHandlerArtifact.relative_fileloc) + ) + known: dict[str, list[TaskHandlerArtifact]] = defaultdict(list) + for row in rows: + try: + artifact = TaskHandlerArtifact.model_validate(row._asdict()) + except ValidationError as exc: + # Left out, the artifact is probed again and its new answer replaces the row. + self.log.warning( + "Ignoring a recorded task-handler artifact that fails validation", + bundle_name=row.bundle_name, + relative_fileloc=row.relative_fileloc, + error=str(exc), + ) + continue + known[row.bundle_name].append(artifact) + return dict(known) + + def _query_known_task_handler_artifacts(self) -> dict[str, list[TaskHandlerArtifact]]: + """Return the recorded artifacts of every bundle that can hold task handlers for this manager's Dag files.""" + bundle_names = set(self._task_handler_bundles.named) + if self._task_handler_bundles.include_dag_bundles: + bundle_names.update(bundle.name for bundle in self._dag_bundles) + if not bundle_names: + return {} + try: + return self.get_known_task_handler_artifacts(bundle_names) + except Exception: + # Without known artifacts the children probe their candidates, so parsing keeps running. + self.log.exception( + "Cannot read the recorded task-handler artifacts; Dag files are parsed without them" + ) + return {} + + def _get_task_handler_artifact_bundle_names(self, dag_bundle_name: str) -> frozenset[str]: + """Return the bundles whose task-handler artifacts a Dag file in *dag_bundle_name* may read and record.""" + bundles = self._task_handler_bundles + teams = self._get_team_names(bundles.named | {dag_bundle_name}) + scope = {name for name in bundles.named if teams.get(name) == teams.get(dag_bundle_name)} + if bundles.include_dag_bundles: + scope.add(dag_bundle_name) + return frozenset(scope) + + def _select_known_task_handler_artifacts( + self, known_artifacts: dict[str, list[TaskHandlerArtifact]], dag_bundle_name: str + ) -> list[TaskHandlerArtifact]: + """Return the artifacts a Dag file in *dag_bundle_name* may resolve its stub tasks against.""" + bundle_names = self._get_task_handler_artifact_bundle_names(dag_bundle_name) + return [artifact for name in sorted(bundle_names) for artifact in known_artifacts.get(name, ())] + + def _create_process( + self, dag_file: DagFileInfo, *, known_artifacts: Sequence[TaskHandlerArtifact] = () + ) -> DagFileProcessorProcess: id = uuid7() callback_to_execute_for_file = self._callback_to_execute.pop(dag_file, []) @@ -1457,6 +1564,7 @@ def _create_process(self, dag_file: DagFileInfo) -> DagFileProcessorProcess: bundle_name=dag_file.bundle_name, dag_file_rel_path=str(dag_file.rel_path), callbacks=callback_to_execute_for_file, + known_artifacts=known_artifacts, selector=self.selector, logger=logger, logger_filehandle=logger_filehandle, @@ -1467,6 +1575,9 @@ def _create_process(self, dag_file: DagFileInfo) -> DagFileProcessorProcess: def _start_new_processes(self): """Start more processors if we have enough slots and files to process.""" bundle_to_team = self._get_team_names({file.bundle_name for file in self._file_queue}) + # Read once per loop, and only when a child starts: rows persisted by earlier children are + # then visible to later ones, which a cache held across a whole pass would hide. + known_artifacts: dict[str, list[TaskHandlerArtifact]] | None = None while self._parallelism > len(self._processors) and self._file_queue: file, _ = self._file_queue.popitem(last=False) @@ -1474,7 +1585,12 @@ def _start_new_processes(self): if file in self._processors: continue - processor = self._create_process(file) + if known_artifacts is None: + known_artifacts = self._query_known_task_handler_artifacts() + processor = self._create_process( + file, + known_artifacts=self._select_known_task_handler_artifacts(known_artifacts, file.bundle_name), + ) stats.incr( "dag_processing.processes", tags=prune_dict( diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index a8d21d3990f60..1cd036d9fadfa 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -71,10 +71,12 @@ from airflow.models.dag_version import DagVersion from airflow.models.dagbundle import DagBundleModel from airflow.models.dagcode import DagCode +from airflow.models.lang_sdk_task_handler import LangSDKTaskHandlerArtifact from airflow.models.serialized_dag import SerializedDagModel from airflow.models.team import Team from airflow.providers.standard.operators.empty import EmptyOperator from airflow.sdk import DAG as SdkDAG +from airflow.sdk.execution_time.coordinator import reset_coordinator_manager from airflow.sdk.importers import DagSourceCode from airflow.serialization.serialized_objects import LazyDeserializedDAG from airflow.utils.net import get_hostname @@ -2687,6 +2689,7 @@ def test_callback_queue(self, mock_get_logger, configure_testing_dag_bundle): bundle_name="testing", dag_file_rel_path=str(dag2_path.rel_path), callbacks=[dag2_req1], + known_artifacts=[], selector=mock.ANY, logger=mock_logger, logger_filehandle=mock_filehandle, @@ -2700,6 +2703,7 @@ def test_callback_queue(self, mock_get_logger, configure_testing_dag_bundle): bundle_name="testing", dag_file_rel_path=str(dag1_path.rel_path), callbacks=[dag1_req1, dag1_req2], + known_artifacts=[], selector=mock.ANY, logger=mock_logger, logger_filehandle=mock_filehandle, @@ -4166,6 +4170,319 @@ def test_a_sweep_pays_the_fixed_cost_once_per_call(self, session, testing_dag_bu ) +JAVA_TASK_HANDLERS = "java-task-handlers" + + +def _coordinator(task_handler_bundle_name: str | None = None) -> dict: + kwargs = ( + {} if task_handler_bundle_name is None else {"task_handler_bundle_name": task_handler_bundle_name} + ) + return {"classpath": "airflow.sdk.coordinators.java.JavaCoordinator", "kwargs": kwargs} + + +def _routed_coordinators(coordinators: dict[str, dict]) -> dict[tuple[str, str], str]: + """Return the ``[sdk]`` options that configure *coordinators* and route a queue to each.""" + return { + ("sdk", "coordinators"): json.dumps(coordinators), + ("sdk", "queue_to_coordinator"): json.dumps({f"queue-{key}": key for key in coordinators}), + } + + +def _known_artifact(bundle_name: str) -> TaskHandlerArtifact: + return TaskHandlerArtifact( + bundle_name=bundle_name, + relative_fileloc=f"{bundle_name}.jar", + size_bytes=1024, + cache_digest="ab12", + task_handlers={ + "etl": [ + TaskHandlerDeclaration( + task_id="extract", + binding="positional", + params=[TaskHandlerParam(name=None, required=True)], + ) + ] + }, + ) + + +GO_TASK_HANDLERS = "go-task-handlers" +TS_TASK_HANDLERS = "ts-task-handlers" +TEAM_COORDINATORS = { + "java": _coordinator(JAVA_TASK_HANDLERS), + "go": _coordinator(GO_TASK_HANDLERS), + "ts": _coordinator(TS_TASK_HANDLERS), +} + + +class TestKnownTaskHandlerArtifacts: + @pytest.fixture(autouse=True) + def _reset_coordinator_manager(self): + reset_coordinator_manager() + yield + reset_coordinator_manager() + + @pytest.fixture + def manager(self, configure_dag_bundles, tmp_path): + with ( + conf_vars({("core", "load_examples"): "False"}), + configure_dag_bundles({"dags-a": tmp_path, "dags-b": tmp_path, JAVA_TASK_HANDLERS: tmp_path}), + ): + manager = DagFileProcessorManager(max_runs=1) + manager._dag_bundles = list(DagBundlesManager().get_all_dag_bundles()) + yield manager + + @conf_vars(_routed_coordinators({"java": _coordinator(JAVA_TASK_HANDLERS), "go": _coordinator()})) + @mock.patch.object(DagFileProcessorProcess, "start") + @mock.patch.object(DagFileProcessorManager, "get_known_task_handler_artifacts", autospec=True) + def test_start_new_processes_reads_known_artifacts_once_per_loop(self, get_known, start, manager): + known = {name: [_known_artifact(name)] for name in ("dags-a", "dags-b", JAVA_TASK_HANDLERS)} + get_known.return_value = known + files = [ + DagFileInfo(bundle_name=bundle_name, rel_path=Path(rel_path), bundle_path=TEST_DAGS_FOLDER) + for bundle_name, rel_path in (("dags-a", "one.py"), ("dags-a", "two.py"), ("dags-b", "three.py")) + ] + manager._parallelism = 3 + manager._file_queue = OrderedDict.fromkeys(files) + + manager._start_new_processes() + + get_known.assert_called_once() + assert { + call.kwargs["dag_file_rel_path"]: call.kwargs["known_artifacts"] for call in start.call_args_list + } == { + "one.py": [*known["dags-a"], *known[JAVA_TASK_HANDLERS]], + "two.py": [*known["dags-a"], *known[JAVA_TASK_HANDLERS]], + "three.py": [*known["dags-b"], *known[JAVA_TASK_HANDLERS]], + } + + manager._processors.clear() + manager._file_queue = OrderedDict.fromkeys(files[:1]) + manager._start_new_processes() + + assert get_known.call_count == 2 + + @pytest.mark.parametrize( + ("coordinators", "teams", "expected_bundle_names"), + [ + pytest.param( + TEAM_COORDINATORS, + {}, + {JAVA_TASK_HANDLERS, GO_TASK_HANDLERS, TS_TASK_HANDLERS}, + id="multi-team-off", + ), + pytest.param( + TEAM_COORDINATORS, + {"dags-a": "team-a", JAVA_TASK_HANDLERS: "team-a", GO_TASK_HANDLERS: "team-b"}, + {JAVA_TASK_HANDLERS}, + id="same-team", + ), + pytest.param( + TEAM_COORDINATORS, + {JAVA_TASK_HANDLERS: "team-a", GO_TASK_HANDLERS: "team-b"}, + {TS_TASK_HANDLERS}, + id="no-team", + ), + pytest.param( + {**TEAM_COORDINATORS, "fallback": _coordinator()}, + {"dags-a": "team-a", JAVA_TASK_HANDLERS: "team-a", GO_TASK_HANDLERS: "team-b"}, + {JAVA_TASK_HANDLERS, "dags-a"}, + id="own-bundle-fallback", + ), + ], + ) + def test_known_artifacts_are_scoped_by_team( + self, configure_dag_bundles, tmp_path, coordinators, teams, expected_bundle_names + ): + bundle_names = ("dags-a", JAVA_TASK_HANDLERS, GO_TASK_HANDLERS, TS_TASK_HANDLERS) + with ( + conf_vars({("core", "load_examples"): "False", **_routed_coordinators(coordinators)}), + configure_dag_bundles(dict.fromkeys(bundle_names, tmp_path)), + mock.patch.object(DagFileProcessorManager, "_get_team_names", autospec=True, return_value=teams), + ): + selected = DagFileProcessorManager(max_runs=1)._select_known_task_handler_artifacts( + {name: [_known_artifact(name)] for name in bundle_names}, "dags-a" + ) + + assert {artifact.bundle_name for artifact in selected} == expected_bundle_names + + @conf_vars(_routed_coordinators({"java": _coordinator(JAVA_TASK_HANDLERS)})) + @mock.patch.object(DagFileProcessorManager, "get_known_task_handler_artifacts", autospec=True) + def test_start_new_processes_reads_nothing_without_a_child_to_start(self, get_known, manager): + manager._start_new_processes() + + get_known.assert_not_called() + + @pytest.mark.parametrize( + ("coordinators", "expected_bundle_names"), + [ + pytest.param({"java": _coordinator(JAVA_TASK_HANDLERS)}, {JAVA_TASK_HANDLERS}, id="named-only"), + pytest.param( + {"java": _coordinator(JAVA_TASK_HANDLERS), "go": _coordinator()}, + {JAVA_TASK_HANDLERS, "dags-a", "dags-b"}, + id="own-bundle-fallback", + ), + ], + ) + @mock.patch.object(DagFileProcessorManager, "get_known_task_handler_artifacts", autospec=True) + def test_query_asks_only_for_bundles_that_can_hold_artifacts( + self, get_known, manager, coordinators, expected_bundle_names + ): + with conf_vars(_routed_coordinators(coordinators)): + manager._query_known_task_handler_artifacts() + + get_known.assert_called_once_with(manager, expected_bundle_names) + + @pytest.mark.parametrize( + "unrouted", + [ + _coordinator("ghost"), + { + "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", + "kwargs": {"task_handler_bundle_name": ["ghost"]}, + }, + _coordinator(), + ], + ids=["unconfigured", "not-a-string", "own-bundle"], + ) + @mock.patch.object(DagFileProcessorManager, "get_known_task_handler_artifacts", autospec=True) + def test_query_leaves_out_the_bundle_of_an_unrouted_coordinator(self, get_known, manager, unrouted): + with conf_vars( + { + ("sdk", "coordinators"): json.dumps( + {"java": _coordinator(JAVA_TASK_HANDLERS), "unrouted": unrouted} + ), + ("sdk", "queue_to_coordinator"): json.dumps({"queue-java": "java"}), + } + ): + manager._query_known_task_handler_artifacts() + + get_known.assert_called_once_with(manager, {JAVA_TASK_HANDLERS}) + + @mock.patch.object(DagFileProcessorManager, "get_known_task_handler_artifacts", autospec=True) + def test_query_skips_the_read_without_coordinators(self, get_known, manager): + assert manager._query_known_task_handler_artifacts() == {} + + get_known.assert_not_called() + + @conf_vars(_routed_coordinators({"java": _coordinator("ghost")})) + @mock.patch.object(DagFileProcessorManager, "get_known_task_handler_artifacts", autospec=True) + def test_invalid_coordinator_config_is_logged_once_and_sends_nothing(self, get_known, manager, caplog): + assert manager._query_known_task_handler_artifacts() == {} + assert manager._query_known_task_handler_artifacts() == {} + + get_known.assert_not_called() + event = "Cannot read [sdk] coordinators; Dag files are parsed without known task-handler artifacts" + assert [entry["log_level"] for entry in caplog.entries if entry["event"] == event] == ["error"] + + @conf_vars(_routed_coordinators({"java": _coordinator(JAVA_TASK_HANDLERS)})) + @mock.patch.object(DagFileProcessorProcess, "start") + @mock.patch.object(DagFileProcessorManager, "get_known_task_handler_artifacts", autospec=True) + def test_start_new_processes_sends_no_known_artifacts_when_the_read_fails( + self, get_known, start, manager, caplog + ): + get_known.side_effect = OperationalError("SELECT", {}, Exception("connection lost")) + manager._parallelism = 2 + manager._file_queue = OrderedDict.fromkeys( + DagFileInfo(bundle_name="dags-a", rel_path=Path(rel_path), bundle_path=TEST_DAGS_FOLDER) + for rel_path in ("one.py", "two.py") + ) + + manager._start_new_processes() + + get_known.assert_called_once() + assert [call.kwargs["known_artifacts"] for call in start.call_args_list] == [[], []] + assert { + "event": "Cannot read the recorded task-handler artifacts; Dag files are parsed without them", + "log_level": "error", + } in caplog + + def test_get_known_task_handler_artifacts_from_db_skips_a_row_that_fails_validation( + self, manager, session, caplog + ): + valid = _known_artifact(JAVA_TASK_HANDLERS) + session.add_all( + [ + LangSDKTaskHandlerArtifact(**valid.model_dump(mode="json")), + LangSDKTaskHandlerArtifact( + bundle_name=JAVA_TASK_HANDLERS, + relative_fileloc="broken.jar", + size_bytes=1024, + cache_digest="ab12", + task_handlers={"etl": [{"task_id": "extract", "binding": "unknown"}]}, + ), + ] + ) + session.flush() + + known = manager._get_known_task_handler_artifacts_from_db({JAVA_TASK_HANDLERS}, session=session) + + assert known == {JAVA_TASK_HANDLERS: [valid]} + assert { + "event": "Ignoring a recorded task-handler artifact that fails validation", + "bundle_name": JAVA_TASK_HANDLERS, + "relative_fileloc": "broken.jar", + "log_level": "warning", + } in caplog + + def test_get_known_task_handler_artifacts_from_db_orders_by_bundle_and_path(self, manager, session): + session.add_all( + LangSDKTaskHandlerArtifact( + bundle_name=bundle_name, + relative_fileloc=rel_path, + size_bytes=1024, + cache_digest=None, + task_handlers={}, + ) + for bundle_name, rel_path in ( + (JAVA_TASK_HANDLERS, "c.jar"), + ("dags-a", "b.jar"), + (JAVA_TASK_HANDLERS, "a.jar"), + (JAVA_TASK_HANDLERS, "b.jar"), + ) + ) + session.flush() + + known = manager._get_known_task_handler_artifacts_from_db( + {JAVA_TASK_HANDLERS, "dags-a"}, session=session + ) + + assert [ + (bundle_name, [artifact.relative_fileloc for artifact in artifacts]) + for bundle_name, artifacts in known.items() + ] == [("dags-a", ["b.jar"]), (JAVA_TASK_HANDLERS, ["a.jar", "b.jar", "c.jar"])] + + def test_get_known_task_handler_artifacts_from_db_reads_only_the_given_bundles(self, manager, session): + session.add_all( + LangSDKTaskHandlerArtifact( + bundle_name=name, + relative_fileloc=f"{name}.jar", + size_bytes=1024, + cache_digest="ab12", + task_handlers={ + "etl": [ + { + "task_id": "extract", + "binding": "positional", + "params": [{"name": None, "required": True}], + } + ] + }, + ) + for name in (JAVA_TASK_HANDLERS, "dags-a", "unconfigured") + ) + session.flush() + + known = manager._get_known_task_handler_artifacts_from_db( + {JAVA_TASK_HANDLERS, "dags-a"}, session=session + ) + + assert known == { + JAVA_TASK_HANDLERS: [_known_artifact(JAVA_TASK_HANDLERS)], + "dags-a": [_known_artifact("dags-a")], + } + + class TestMultiTeamMetrics: """Tests for team_name tag on dag processing metrics in multi-team mode.""" From f78db1cb06a94bad7786899851b225a13f1f30e8 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 18:30:22 +0000 Subject: [PATCH 92/94] Describe how known artifacts are scoped per Dag-parsing child ADR-0013 predates the fallback to the task's own Dag bundle, so it said the manager reads only the named bundles and sends every child the same list. A fallback coordinator makes Dag bundles hold artifacts too, and no child can use another Dag bundle's rows. Each row now carries the artifact's whole answer, which every file reading it trusts, so a child only sees the named bundles of its own team, and the ADR's probe sketch still asked for Dag ids the runtimes no longer filter by. The message is named TaskHandlerArtifact, like the other task-handler parse messages. --- .../0013-persisted-task-handler-bindings.md | 36 ++++++++++--------- 1 file changed, 19 insertions(+), 17 deletions(-) diff --git a/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md b/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md index 6a62c732497c3..e38c4c3c3cb58 100644 --- a/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md +++ b/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md @@ -210,15 +210,16 @@ fingerprint beside it. A binding row is then only a mapping from a stub task to already knows about. The parse child processor subprocesses run in the client context without a database connection ([`_parse_file_entrypoint`](https://github.com/apache/airflow/blob/79991cd4db0c9346a28b23c453377f6df0c6b4ed/airflow-core/src/airflow/dag_processing/processor.py#L208-L232)), so the prior cache state must be pushed down from the manager. The alternative is a dedicated -Execution API for the child processor process to retrieve `KnownSDKTaskHandlerArtifact` itself, which +Execution API for the child processor process to retrieve `TaskHandlerArtifact` itself, which was rejected on blast radius. ```python -class KnownSDKTaskHandlerArtifact(BaseModel): +class TaskHandlerArtifact(BaseModel): bundle_name: str relative_fileloc: str size_bytes: int - cache_digest: str + cache_digest: str | None # None: the artifact stores no fingerprint, so it is always probed + task_handlers: dict[str, list[TaskHandlerDeclaration]] # the probe answer: dag_id -> declarations class DagFileParseRequest(BaseModel): @@ -226,14 +227,11 @@ class DagFileParseRequest(BaseModel): bundle_path: Path bundle_name: str callback_requests: list[CallbackRequest] - known_artifacts: list[KnownSDKTaskHandlerArtifact] = [] # new + known_artifacts: list[TaskHandlerArtifact] = [] # new type: Literal["DagFileParseRequest"] ``` -`known_artifacts` is scoped by *artifact* bundle, not by Dag file, so the manager reads it **once per -parsing loop** for every configured `task_handler_bundle_name` and pushes the same list to every -child. Ten Dag files resolving against one twenty-jar bundle therefore probe that bundle once in -total, not once each. +`known_artifacts` is scoped by *artifact* bundle, not by Dag file, so the manager reads it **once per parsing loop** for every bundle that can hold task handlers: the `task_handler_bundle_name` of each coordinator a queue routes to, plus the Dag bundles it parses when one of them sets none and so reads the task's own Dag bundle. A coordinator no queue routes to runs no stub task, so its bundle is not read. Each child gets the rows of the named bundles whose team matches its Dag bundle's team, and of its own Dag bundle when a coordinator falls back to it, never those of another Dag bundle, which no coordinator reads for it. The team rule exists because an answer one file's parse records is trusted by every file that reads it. Teams match only when equal: without `[core] multi_team` every named bundle is in scope, and with it a team-less Dag bundle sees only team-less named bundles. One recorded answer serves every file, so ten Dag files resolving against one bundle probe a changed artifact once in total, apart from the children already started when it changed. **Dag-parsing child → coordinator subprocess.** Introduced here. The child spawns the runtime and forwards bytes in both directions, decoding nothing; the process that spawned the parse decodes the @@ -243,7 +241,6 @@ one member, and `ToManager` gains `SDKTaskHandlerParsingResult`. ```python class SDKTaskHandlerParseRequest(BaseModel): # parent -> runtime, on ToSDKTaskHandlerProcessor file: str # the candidate artifact being probed - dag_ids: list[str] # every Dag in this file with stub tasks routed here bundle_path: Path bundle_name: str type: Literal["SDKTaskHandlerParseRequest"] @@ -269,9 +266,9 @@ class TaskHandlerParam(BaseModel): exact_name: bool = False # match as spelled, not case-insensitively with underscores ignored ``` -A `dag_id` the artifact registers nothing for is **omitted** from `task_handlers` rather than returned -empty, so a probe that matches nothing is distinguishable from a probe that matched a Dag with zero -tasks. +`task_handlers` maps every `dag_id` the artifact registers a handler for, and is `{}` when it registers +none. The answer must not depend on the request, because the manager records it once and serves it to +every Dag file that resolves against the artifact. **Dag-parsing child → manager.** [`DagFileParsingResult`](https://github.com/apache/airflow/blob/79991cd4db0c9346a28b23c453377f6df0c6b4ed/airflow-core/src/airflow/dag_processing/processor.py#L133-L145) gains the resolved bindings. @@ -350,12 +347,17 @@ stub Dag is validated against bindings that have not been written yet. ``` DagProcessorManager [reads DB] │ - │ once per parsing loop, per configured task_handler_bundle_name: - │ SELECT bundle_name, relative_fileloc, size_bytes, cache_digest + │ once per parsing loop, for every bundle that can hold task handlers + │ (each routed coordinator's task_handler_bundle_name, plus the Dag + │ bundles when one sets none and reads the task's own Dag bundle): + │ SELECT bundle_name, relative_fileloc, size_bytes, cache_digest, + │ task_handlers │ FROM lang_sdk_task_handler_artifact - │ WHERE bundle_name IN (:configured bundles) ──▶ known_artifacts + │ WHERE bundle_name IN (:those bundles) ──▶ known_artifacts │ - │ per file about to be dispatched: + │ per file about to be dispatched: the rows of the bundles in its scope + │ (named bundles of its team, plus its own Dag bundle when a routed + │ coordinator sets none) │ (the child re-reads nothing; it has no DB) │ ├── DagFileParseRequest(file=etl.py, bundle_*, known_artifacts=[...]) @@ -394,7 +396,7 @@ DagFileProcessorProcess(etl.py) [no DB — client con │ target=_start_task_handler_runtime_entrypoint, │ coordinator=JavaCoordinator("jdk-17"), │ path=) - │ ──SDKTaskHandlerParseRequest(file=…, dag_ids=["etl"])──▶ runtime + │ ──SDKTaskHandlerParseRequest(file=…)──▶ runtime │ ◀─SDKTaskHandlerParsingResult(task_handlers={"etl": […]})── runtime │ Get* from the runtime is relayed up ToManager unchanged │ From 411c2092e11bf13b31098951921908007b3e7888 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 18:47:32 +0000 Subject: [PATCH 93/94] Test that a Dag file gets only named bundles without a fallback The other tests that start children all configure a coordinator that falls back to the Dag's own bundle, so dropping that condition on the way to the child would have gone unnoticed, and an override of get_known_task_handler_artifacts may return bundles it was not asked for. --- .../tests/unit/dag_processing/test_manager.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 1cd036d9fadfa..ecc99a49c461a 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -4306,6 +4306,20 @@ def test_known_artifacts_are_scoped_by_team( assert {artifact.bundle_name for artifact in selected} == expected_bundle_names + @conf_vars(_routed_coordinators({"java": _coordinator(JAVA_TASK_HANDLERS)})) + @mock.patch.object(DagFileProcessorProcess, "start") + @mock.patch.object(DagFileProcessorManager, "get_known_task_handler_artifacts", autospec=True) + def test_start_new_processes_keeps_own_bundle_out_without_a_fallback(self, get_known, start, manager): + # An override may return bundles it was not asked for. + get_known.return_value = {name: [_known_artifact(name)] for name in ("dags-a", JAVA_TASK_HANDLERS)} + manager._file_queue = OrderedDict.fromkeys( + [DagFileInfo(bundle_name="dags-a", rel_path=Path("one.py"), bundle_path=TEST_DAGS_FOLDER)] + ) + + manager._start_new_processes() + + assert start.call_args.kwargs["known_artifacts"] == [_known_artifact(JAVA_TASK_HANDLERS)] + @conf_vars(_routed_coordinators({"java": _coordinator(JAVA_TASK_HANDLERS)})) @mock.patch.object(DagFileProcessorManager, "get_known_task_handler_artifacts", autospec=True) def test_start_new_processes_reads_nothing_without_a_child_to_start(self, get_known, manager): From 00586dbfdb89979ed3727e691b5904162a0fd9c8 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 1 Oct 2026 17:49:17 +0000 Subject: [PATCH 94/94] Build the known-artifact test params without required TaskHandlerParam no longer has required. A model dump no longer carries it, so the expected request bodies must drop it, and the constructors would otherwise pass a keyword pydantic only ignores. --- airflow-core/tests/unit/dag_processing/test_manager.py | 7 +++---- airflow-core/tests/unit/dag_processing/test_processor.py | 6 ++---- .../tests/task_sdk/execution_time/schema/test_migrator.py | 6 ++---- 3 files changed, 7 insertions(+), 12 deletions(-) diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index ecc99a49c461a..1f0a66cbf303d 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -1892,7 +1892,7 @@ def test_collect_results_tolerates_stale_file_handle_on_close(self): TaskHandlerDeclaration( task_id="extract", binding="positional", - params=[TaskHandlerParam(name=None, required=True)], + params=[TaskHandlerParam(name=None)], ) ] }, @@ -1919,7 +1919,6 @@ def test_collect_results_tolerates_stale_file_handle_on_close(self): { "name": None, "value_schema": None, - "required": True, "exact_name": False, } ], @@ -4199,7 +4198,7 @@ def _known_artifact(bundle_name: str) -> TaskHandlerArtifact: TaskHandlerDeclaration( task_id="extract", binding="positional", - params=[TaskHandlerParam(name=None, required=True)], + params=[TaskHandlerParam(name=None)], ) ] }, @@ -4478,7 +4477,7 @@ def test_get_known_task_handler_artifacts_from_db_reads_only_the_given_bundles(s { "task_id": "extract", "binding": "positional", - "params": [{"name": None, "required": True}], + "params": [{"name": None}], } ] }, diff --git a/airflow-core/tests/unit/dag_processing/test_processor.py b/airflow-core/tests/unit/dag_processing/test_processor.py index 1e1e710948994..2de2b89aa313b 100644 --- a/airflow-core/tests/unit/dag_processing/test_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_processor.py @@ -629,9 +629,7 @@ def _make_known_artifact_body(**overrides) -> dict: "size_bytes": 1024, "cache_digest": "ab12", "task_handlers": { - "etl": [ - {"task_id": "extract", "binding": "positional", "params": [{"name": None, "required": True}]} - ] + "etl": [{"task_id": "extract", "binding": "positional", "params": [{"name": None}]}] }, **overrides, } @@ -673,7 +671,7 @@ def test_parse_file_entrypoint_decodes_known_artifacts(): TaskHandlerDeclaration( task_id="extract", binding="positional", - params=[TaskHandlerParam(name=None, required=True)], + params=[TaskHandlerParam(name=None)], ) ] }, diff --git a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py index ae698712ea966..d837181c6f6d4 100644 --- a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py +++ b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py @@ -541,7 +541,7 @@ def request_with_known_artifacts(self) -> DagFileParseRequest: TaskHandlerDeclaration( task_id="extract", binding="positional", - params=[TaskHandlerParam(name=None, required=True)], + params=[TaskHandlerParam(name=None)], ) ] }, @@ -572,9 +572,7 @@ def test_downgrade_keeps_known_artifacts_at_head(self, real_migrator, request_wi { "task_id": "extract", "binding": "positional", - "params": [ - {"name": None, "value_schema": None, "required": True, "exact_name": False} - ], + "params": [{"name": None, "value_schema": None, "exact_name": False}], } ] },