diff --git a/airflow-core/docs/authoring-and-scheduling/yaml-dags.rst b/airflow-core/docs/authoring-and-scheduling/yaml-dags.rst index ceef4e9d4b1ef..e0d9f4caf53f6 100644 --- a/airflow-core/docs/authoring-and-scheduling/yaml-dags.rst +++ b/airflow-core/docs/authoring-and-scheduling/yaml-dags.rst @@ -26,9 +26,9 @@ Writing Dags in YAML Airflow can build a Dag from a YAML document. One YAML document describes one Dag, and the YAML Dag importer turns it into the same Dag structure that Airflow -runs like any other source. The tasks can be ordinary Airflow operators, -or a pointer to a task handler implemented in a Language SDK. (Support for the -Python Task SDK is planned for a future release.) +runs like any other source. The tasks can be ordinary Airflow operators, or a +pointer to a task handler implemented in any language with a Task SDK, including +Python, Java, and more. A document is declarative. Everything in it is a value, and the shape of each value decides what it means. @@ -423,11 +423,12 @@ which holds the code's arguments directly: source: {$x: extract} factor: 2 -A ``run:`` task is executed by a Task SDK runtime, so it must be routed to a -coordinator. In practice, a ``run:`` task always needs a task-level ``queue:`` -mapped to a coordinator under ``[sdk] queue_to_coordinator``. See -:doc:`language-sdks/index` for writing the handler and configuring the -coordinator. +A ``run:`` task must be accompanied by a task handler implemented by a Task SDK +that will be executed at runtime. How the handler is implemented is not named by +the YAML file; it follows the task-level ``queue:``, and then is mapped to the +appropriate coordinator through ``[sdk] queue_to_coordinator``. For more on +writing task handlers, see :doc:`task-sdk:task-handler` for Python, +and :doc:`language-sdks/index` for other languages. Each argument in ``run:`` is either a JSON-serializable literal, or an upstream task's output, in the ``{$x: task}`` form, which also wires the dependency. @@ -446,9 +447,11 @@ implemented in a future release. be scheduled on an asset (see `Scheduling`_). * **Task groups.** Tasks cannot be nested into groups (see `Reserved arguments`_). -* **Dynamic task mapping.** There is no equivalent of ``.expand()``. +* **Mapped tasks.** There is no equivalent of ``.expand()``. * **Arguments that must be a Python object.** This particularly affects - callbacks (which need a Python function). + callbacks (which need a Python function). Note that using ``PythonOperator`` + is generally not needed in a YAML Dag; implement a Python task handler to use + with a ``run:`` task instead (see `Code tasks`_). A rejected key is reported as an error rather than silently ignored, as is any other argument the target operator does not accept. diff --git a/airflow-core/tests/unit/cli/commands/test_dag_command.py b/airflow-core/tests/unit/cli/commands/test_dag_command.py index ff97b2fe7169a..8ff16d01de07c 100644 --- a/airflow-core/tests/unit/cli/commands/test_dag_command.py +++ b/airflow-core/tests/unit/cli/commands/test_dag_command.py @@ -1299,8 +1299,8 @@ def execute(self, context, event=None): tis = dr.get_task_instances() assert next(x for x in tis if x.task_id == "abc").state == "success" - @mock.patch("airflow.sdk.execution_time.task_runner._execute_task") - def test_dag_test_with_mark_success(self, mock__execute_task): + @mock.patch("airflow.sdk.execution_time.task_runner._TaskRunner.run") + def test_dag_test_with_mark_success(self, mock_run): """ option `--mark-success-pattern` should mark matching tasks as success without executing them. """ @@ -1317,8 +1317,8 @@ def test_dag_test_with_mark_success(self, mock__execute_task): dag_command.dag_test(cli_args) # only second operator was actually executed, first one was marked as success - assert len(mock__execute_task.call_args_list) == 1 - assert mock__execute_task.call_args_list[0].kwargs["ti"].task_id == "dummy_operator" + assert len(mock_run.call_args_list) == 1 + assert mock_run.call_args_list[0].args[0].task_id == "dummy_operator" @conf_vars({("core", "load_examples"): "false"}) def test_get_dag_excludes_examples_with_bundle(self, configure_testing_dag_bundle): diff --git a/task-sdk/docs/api.rst b/task-sdk/docs/api.rst index 824fde2638e48..e2855a28ff2a9 100644 --- a/task-sdk/docs/api.rst +++ b/task-sdk/docs/api.rst @@ -67,6 +67,8 @@ Task Decorators: .. autoapifunction:: airflow.sdk.task_group +.. autoapifunction:: airflow.sdk.task_handler + .. autoapifunction:: airflow.sdk.setup .. autoapifunction:: airflow.sdk.teardown diff --git a/task-sdk/docs/index.rst b/task-sdk/docs/index.rst index c0653f6ee3fcf..35cce71109005 100644 --- a/task-sdk/docs/index.rst +++ b/task-sdk/docs/index.rst @@ -180,6 +180,7 @@ For the full public API reference, see the :doc:`api` page. resumable-job-mixin api concepts + task-handler lang-sdk executable-bundle-spec ts-bundle-spec diff --git a/task-sdk/docs/task-handler.rst b/task-sdk/docs/task-handler.rst new file mode 100644 index 0000000000000..e5509bfc946d2 --- /dev/null +++ b/task-sdk/docs/task-handler.rst @@ -0,0 +1,125 @@ + .. 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. + + +Task handlers +============= + +A task may be declared in one place and have its body implemented in another. +The *stub task* interface is implemented by the Python (``@task.stub``) and YAML +(``run:``) Dag formats. A stub has its body supplied by a *task handler*. +Handlers can be written in any language with a Task SDK that implements the +task handler interface. + +Writing a handler +----------------- + +Decorate a plain function with the ``(dag_id, task_id)`` it implements: + +.. code-block:: python + + from airflow.sdk import task_handler + + + @task_handler(dag_id="events_pipeline", task_id="transform") + def transform(region: str, events: list[dict]) -> dict: + return {"region": region, "count": len(events)} + +* The function's parameters are the task's arguments. They are bound from the + stub's captured arguments by name, and may accept either a literal value, or + a reference to a return value XCom. +* The return value becomes the task's return value XCom, so downstream tasks + consume it like any other task's output. +* ``task_id`` defaults to the function name, so the above could be written + equivalently as simply ``@task_handler(dag_id="events_pipeline")``. +* Arguments are bound to parameters by name, so ``*args``, ``**kwargs``, and + positional-only parameters are not allowed. No parameter may be named after a + task-context key (``ti``, ``ds``, etc.). + +Deploying handlers +------------------ + +A stub task's handler must be available in two places: the Dag processor, which +validates the stub's arguments against the handler's signature when the Dag is +parsed, and the worker that runs the task. + +In both cases, the handler is loaded from a Dag bundle, so the most important +thing to decide is which bundle should hold it. + +In the stub's own bundle +~~~~~~~~~~~~~~~~~~~~~~~~ + +A handler is discovered automatically in the stub task's own +:doc:`Dag bundle `. +If you place the handler in the same bundle as the Dag that uses it, it is +imported with no configuration at all. This is the recommended layout. + +Keep the bundle lean. Everything in a handler's bundle is loaded on every parse +and every task run to discover its handlers. Airflow cannot tell which +files are handlers, so it loads them all. A bundle stuffed with unrelated +definitions (Python Dag files, for example) can therefore be slow to load. The +leaner the bundle, the faster it loads; prefer a focused bundle of stub Dags +together with the handlers that implement them. + +In a separate bundle +~~~~~~~~~~~~~~~~~~~~ + +Alternatively, the stub's coordinator can be pointed at the handler's bundle +with the ``task_handler_bundle_name`` keyword argument. Register the handler +bundle in ``[dag_processor] dag_bundle_config_list`` like any other bundle, set +``task_handler_bundle_name`` on the coordinator in ``[sdk] coordinators``, and +route the stub's ``queue`` to that coordinator in ``[sdk] queue_to_coordinator`` +like this: + +.. code-block:: ini + + [dag_processor] + # The bundle holding the handlers, registered like any other bundle. + dag_bundle_config_list = [ + {"name": "shared-handlers", "classpath": "...", "kwargs": {...}}, + ] + + [sdk] + coordinators = { + "handlers": { + "classpath": "airflow.sdk.coordinators.python.PythonCoordinator", + "kwargs": {"task_handler_bundle_name": "shared-handlers"} + } + } + queue_to_coordinator = {"workers": "handlers"} + +A stub whose ``queue`` is ``workers`` is then routed to that coordinator, which +loads ``shared-handlers`` instead of the stub's own bundle. +``task_handler_bundle_name`` is validated against ``dag_bundle_config_list`` +when configuration loads, so a typo fails fast. + +It is considered best practice to give handler tasks dedicated queues for +management purposes. However, you can have a queue that's shared by regular +tasks and stub task handlers, and it would work just fine as long as the +coordinator routed to the queue is capable of executing them. + +Notes +----- + +* For the Python SDK, a handler bundle is imported fresh in each Dag parse and + task run; every non-ignored ``.py`` file in it is imported to discover + handlers (``.airflowignore`` is honored). +* A stub's arguments must satisfy cross-language requirements. An argument must + be either JSON-compatible, or reference an upstream task's return value XCom. +* A stub's arguments are validated against the handler's signature when the Dag + is parsed. If a handler cannot be resolved at parse time, argument validation + is deferred; a genuinely missing handler is reported when the task runs. diff --git a/task-sdk/src/airflow/sdk/__init__.py b/task-sdk/src/airflow/sdk/__init__.py index 46418ceeaf125..2e2f08feb0f46 100644 --- a/task-sdk/src/airflow/sdk/__init__.py +++ b/task-sdk/src/airflow/sdk/__init__.py @@ -120,6 +120,7 @@ "setup", "task", "task_group", + "task_handler", "teardown", ] @@ -234,6 +235,7 @@ from airflow.sdk.execution_time import macros from airflow.sdk.execution_time.context import NEVER_EXPIRE from airflow.sdk.io.path import ObjectStoragePath + from airflow.sdk.task_handlers import task_handler from airflow.sdk.types import TaskInstance conf: AirflowSDKConfigParser @@ -342,6 +344,7 @@ "setup": (".definitions.decorators", "setup"), "task": (".definitions.decorators", "task"), "task_group": (".definitions.decorators", "task_group"), + "task_handler": (".task_handlers", "task_handler"), "teardown": (".definitions.decorators", "teardown"), } diff --git a/task-sdk/src/airflow/sdk/__init__.pyi b/task-sdk/src/airflow/sdk/__init__.pyi index 221d3eb48a6c6..da47ece08bdb7 100644 --- a/task-sdk/src/airflow/sdk/__init__.pyi +++ b/task-sdk/src/airflow/sdk/__init__.pyi @@ -136,6 +136,7 @@ from airflow.sdk.execution_time import macros as macros from airflow.sdk.execution_time.cache import SecretCache as SecretCache from airflow.sdk.execution_time.context import NEVER_EXPIRE from airflow.sdk.io.path import ObjectStoragePath as ObjectStoragePath +from airflow.sdk.task_handlers import task_handler as task_handler conf: AirflowSDKConfigParser @@ -236,6 +237,7 @@ __all__ = [ "setup", "task", "task_group", + "task_handler", "teardown", "conf", ] diff --git a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py index aa30b834ef64d..21403aa82d212 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py +++ b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py @@ -47,12 +47,12 @@ import psutil import structlog -from airflow.dag_processing.bundles.base import BundleVersionLock, unpack_bundle_version # noqa: SDK002 +from airflow.dag_processing.bundles.base import BundleVersionLock # noqa: SDK002 from airflow.sdk._shared.module_loading import import_string from airflow.sdk.api.datamodels._generated import BundleInfo, TaskInstanceState from airflow.sdk.configuration import conf from airflow.sdk.coordinators._dag_importer import find_claiming_importer -from airflow.sdk.execution_time.bundles import initialize_ti_bundle +from airflow.sdk.execution_time.bundles import initialize_pinned_bundle from airflow.sdk.execution_time.comms import RetryTask, TaskState from airflow.sdk.execution_time.coordinator import BaseCoordinator, TaskLaunchError from airflow.sdk.execution_time.schema import get_schema_version_migrator @@ -605,29 +605,6 @@ def _handle_set_task_state_store( return super()._handle_set_task_state_store(msg, log, req_id) -def _initialize_pinned_bundle(target: BundleInfo, logger: FilteringBoundLogger) -> BaseDagBundle: - """ - Materialize *target* at a concrete version, so the tree handed to the subprocess is lockable. - - A bundle resolved without a version points at the bundle's shared, mutable - checkout: another task refreshing the same bundle resets it underneath a running - subprocess, and ``BundleVersionLock`` cannot protect it because a version-less - lock is a no-op. Re-resolving at the version current now yields a private - ``versions/`` tree that the lock does cover. - - Bundles that do not track versions have nothing to pin and keep their single path. - """ - bundle = initialize_ti_bundle(target) - if bundle.version is not None: - return bundle - - version, version_data = unpack_bundle_version(bundle.get_current_version(), bundle) - if version is None: - return bundle - logger.debug("Pinning Dag bundle to its current version", bundle=target.name, version=version) - return initialize_ti_bundle(BundleInfo(name=target.name, version=version, version_data=version_data)) - - def _resolve_bundle(bundle_info: BundleInfo, logger: FilteringBoundLogger) -> BaseDagBundle: """ Materialize the Dag bundle of *bundle_info*, pinned to a concrete version. @@ -638,7 +615,7 @@ def _resolve_bundle(bundle_info: BundleInfo, logger: FilteringBoundLogger) -> Ba """ cannot_read = f"Dag bundle {bundle_info.name!r} cannot be read" try: - bundle = _initialize_pinned_bundle(bundle_info, logger) + bundle = initialize_pinned_bundle(bundle_info, logger) except Exception as e: raise TaskLaunchError(f"{cannot_read}: {e}") from e try: diff --git a/task-sdk/src/airflow/sdk/coordinators/python.py b/task-sdk/src/airflow/sdk/coordinators/python.py new file mode 100644 index 0000000000000..79f3b3134700f --- /dev/null +++ b/task-sdk/src/airflow/sdk/coordinators/python.py @@ -0,0 +1,192 @@ +# +# 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. + +""" +Python runtime coordinator. + +Unlike foreign language coordinators, this starts the ordinary Task SDK +:class:`~airflow.sdk.execution_time.supervisor.ActivitySubprocess`. It is used +by default when a queue is not mapped to any other coordinator. + +The coordinator can also be configured explicitly in ``[sdk] coordinators`` and +``[sdk] queue_to_coordinator`` to carry ``task_handler_bundle_name``. This is +useful for YAML Dags, where ``run:`` tasks can point to a task handler defined +in Python. The value is read by the parse and execution paths to import handler +modules; see :func:`~airflow.sdk.task_handlers.task_handler`. +""" + +from __future__ import annotations + +import contextlib +import os +import signal +from typing import TYPE_CHECKING, Any + +import attrs +import structlog + +from airflow.sdk.execution_time.coordinator import BaseCoordinator + +if TYPE_CHECKING: + from collections.abc import Callable, Generator + from os import PathLike + + from structlog.typing import FilteringBoundLogger + + from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: SDK002 + from airflow.sdk.api.client import Client + from airflow.sdk.api.datamodels._generated import TaskInstance + from airflow.sdk.task_handlers import TaskHandlerDeclaration + +log: FilteringBoundLogger = structlog.get_logger(logger_name="coordinators.python") + + +@contextlib.contextmanager +def _warm_shutdown_signals() -> Generator[None, None, None]: + """ + Install SIGTERM/SIGINT warm-shutdown handlers during task supervision. + + While supervising a task the supervisor must not be torn down by a + termination signal; instead it keeps running so the task can finish (or be + shut down gracefully) and its terminal state and logs are reported. The + handlers are installed around BOTH ``start()`` (which transitions the TI to + RUNNING) and ``wait()`` (which runs the task and then reports the terminal + state / uploads logs), so there is no window where Python's default SIGTERM + disposition could kill the supervisor and tear the just-started task down + with it. + + The previous dispositions are restored on exit so a long-lived supervisor + process (e.g. a reused Celery prefork worker) does not leak the handler into + later tasks or clobber the worker's own signal handling. + """ + + def _warm_shutdown(signum, frame): + log.info( + "Received signal; warm shutdown in progress, waiting for the running task to complete.", + signal=signal.Signals(signum).name, + pid=os.getpid(), + ) + + prev_sigterm = signal.getsignal(signal.SIGTERM) + prev_sigint = signal.getsignal(signal.SIGINT) + signal.signal(signal.SIGTERM, _warm_shutdown) + signal.signal(signal.SIGINT, _warm_shutdown) + try: + yield + finally: + signal.signal(signal.SIGTERM, prev_sigterm) + signal.signal(signal.SIGINT, prev_sigint) + + +@attrs.define(kw_only=True) +class PythonCoordinator(BaseCoordinator): + """ + Coordinator that executes Python tasks by launching a Task SDK subprocess. + + This is automatically used as the default fallback when a queue is not + mapped to another coordinator. It may also be configured explicitly in + ``[sdk] coordinators`` to carry a ``task_handler_bundle_name``, then mapped + to specific queues in ``[sdk] queue_to_coordinator``. + + :param task_handler_bundle_name: Name of the Dag bundle holding the Python + ``@task_handler`` modules for stub tasks routed to the coordinator. This + is validated against ``[dag_processor] dag_bundle_config_list`` when + configuration loads. + """ + + task_handler_bundle_name: str | None = None + + def execute_task( + self, + *, + what: TaskInstance, + dag_rel_path: str | PathLike[str], + bundle_info, + client: Client, + logger: FilteringBoundLogger | None = None, + sentry_integration: str = "", + subprocess_logs_to_stdout: bool, + **kwargs: Any, + ) -> BaseCoordinator.ExecutionResult: + # TODO: Importing this at the top causes circular imports. + # ActivitySubprocess and WatchedSubprocess should be moved out of the + # supervisor, and maybe with additional refactoring to abstract out + # process handling. + from airflow.sdk.execution_time.supervisor import ActivitySubprocess + + # Keep the warm-shutdown handlers installed across both start() (which + # transitions the TI to RUNNING) and wait() (which runs the task and + # reports its terminal state / uploads logs) so a SIGTERM at any point + # in this window can't kill the supervisor and tear the task down. + with _warm_shutdown_signals(): + process = ActivitySubprocess.start( + dag_rel_path=dag_rel_path, + what=what, + client=client, + logger=logger, + bundle_info=bundle_info, + subprocess_logs_to_stdout=subprocess_logs_to_stdout, + sentry_integration=sentry_integration, + ) + exit_code = process.wait() + return self.ExecutionResult(exit_code, process.final_state) + + def parse_task_handler( + self, + *, + dag_id: str, + task_id: str, + bundle: BaseDagBundle, + ) -> TaskHandlerDeclaration | None: + """ + Parse the in-process Python handler into its declared signature. + + Python task handlers are registered in the current process, so an + in-process import and inspection is enough. The returned declaration + carries the ordered parameters for validation. + """ + from airflow.sdk.task_handlers import TaskHandlerDeclaration + + if (handler := self.resolve_handler(dag_id=dag_id, task_id=task_id, bundle=bundle)) is None: + return None + try: + return TaskHandlerDeclaration.from_handler(handler) + except ValueError as e: + raise ValueError(f"task handler {task_id!r} {e}") from e + + def resolve_handler( + self, + dag_id: str, + task_id: str, + *, + bundle: BaseDagBundle, + ) -> Callable[..., Any] | None: + """ + Resolve the Python callable registered for the task. + + This is the execution-time counterpart to :meth:`parse_task_handler` for + the Python coordinator, which executes the handler in-process. + + :meta private: + """ + from airflow.sdk.task_handlers import get_handler, import_handler_modules, resolve_handler_bundle + + if self.task_handler_bundle_name is not None: + bundle = resolve_handler_bundle(self.task_handler_bundle_name) + import_handler_modules(bundle) + return get_handler(dag_id, task_id) diff --git a/task-sdk/src/airflow/sdk/execution_time/bundles.py b/task-sdk/src/airflow/sdk/execution_time/bundles.py index 8ab131beb64f3..e41077497d71f 100644 --- a/task-sdk/src/airflow/sdk/execution_time/bundles.py +++ b/task-sdk/src/airflow/sdk/execution_time/bundles.py @@ -28,15 +28,40 @@ from getpass import getuser from typing import TYPE_CHECKING +from airflow.dag_processing.bundles.base import 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.exceptions import AirflowException from airflow.sdk.execution_time.tracing import detail_span if TYPE_CHECKING: + from structlog.typing import FilteringBoundLogger + from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: SDK002 - from airflow.sdk.api.datamodels._generated import BundleInfo -__all__ = ["initialize_ti_bundle", "verify_bundle_access"] +__all__ = ["initialize_pinned_bundle", "initialize_ti_bundle", "verify_bundle_access"] + + +def initialize_pinned_bundle(target: BundleInfo, logger: FilteringBoundLogger) -> BaseDagBundle: + """ + Materialize *target* at a concrete version, so the tree handed to the caller is lockable. + + A bundle resolved without a version points at the bundle's shared, mutable + checkout: another task refreshing the same bundle resets it underneath a caller + still reading it, and ``BundleVersionLock`` cannot protect it because a + version-less lock is a no-op. Re-resolving at the version current now yields a + private ``versions/`` tree that the lock does cover. + + Bundles that do not track versions have nothing to pin and keep their single path. + """ + bundle = initialize_ti_bundle(target) + if bundle.version is not None: + return bundle + version, version_data = unpack_bundle_version(bundle.get_current_version(), bundle) + if version is None: + return bundle + logger.debug("Pinning Dag bundle to its current version", bundle=target.name, version=version) + return initialize_ti_bundle(BundleInfo(name=target.name, version=version, version_data=version_data)) def initialize_ti_bundle(bundle_info: BundleInfo) -> BaseDagBundle: diff --git a/task-sdk/src/airflow/sdk/execution_time/coordinator.py b/task-sdk/src/airflow/sdk/execution_time/coordinator.py index c17f031a83a0d..cdd792ef4be94 100644 --- a/task-sdk/src/airflow/sdk/execution_time/coordinator.py +++ b/task-sdk/src/airflow/sdk/execution_time/coordinator.py @@ -32,8 +32,6 @@ import contextlib import functools -import os -import signal from typing import TYPE_CHECKING, Any import attrs @@ -46,14 +44,16 @@ from airflow.sdk.exceptions import AirflowConfigException if TYPE_CHECKING: - from collections.abc import Generator, Mapping + from collections.abc import Mapping from os import PathLike from typing import Self from structlog.typing import FilteringBoundLogger + from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: SDK002 from airflow.sdk.api.client import Client from airflow.sdk.api.datamodels._generated import TaskInstance + from airflow.sdk.task_handlers import TaskHandlerDeclaration __all__ = [ "BaseCoordinator", @@ -84,10 +84,10 @@ def __init__(self, message: str, *, retryable: bool = True) -> None: class BaseCoordinator: """ - Base coordinator for runtime-specific DAG file processing and task execution. + Base coordinator for runtime-specific Dag processing and execution. Coordinators are instantiated from the ``[sdk] coordinators`` configuration - (see :class:`CoordinatorManager`) — each entry's ``classpath`` is resolved + (see :class:`CoordinatorManager`). Each entry's ``classpath`` is resolved via :func:`~airflow.sdk._shared.module_loading.import_string` and constructed with the entry's ``kwargs``. """ @@ -122,6 +122,33 @@ def execute_task( """ raise NotImplementedError + def parse_task_handler( + self, + *, + dag_id: str, + task_id: str, + bundle: BaseDagBundle, + ) -> TaskHandlerDeclaration | None: + """ + Parse the handler for the specified task into its declared signature. + + This is the cross-language, parse-time validation hook defined by + Language SDK ADR-0011 and ADR-0012. A stub task's captured arguments are + compared against the returned signature the same way for every runtime. + + *bundle* is the stub's own Dag bundle, where a coordinator looks for + a handler colocated with the stub when no dedicated handler bundle is + configured. + + If this coordinator has no handler for the *dag_id* and *task_id* + combination, *None* is returned and validation is deferred to run time. + This hook is for validation only and never resolves a callable to run. + + The default implementation returns *None* to indicate this coordinator + declares nothing. + """ + return None + class _CoordinatorSpec(pydantic.BaseModel): classpath: str @@ -131,90 +158,11 @@ class _CoordinatorSpec(pydantic.BaseModel): extra: dict[str, Any] | None = None -@contextlib.contextmanager -def _warm_shutdown_signals() -> Generator[None, None, None]: - """ - Install SIGTERM/SIGINT warm-shutdown handlers for the duration of task supervision. - - While supervising a task the supervisor must not be torn down by a - termination signal; instead it keeps running so the task can finish (or be - shut down gracefully) and its terminal state and logs are reported. The - handlers are installed around BOTH ``start()`` (which transitions the TI to - RUNNING) and ``wait()`` (which runs the task and then reports the terminal - state / uploads logs), so there is no window where Python's default SIGTERM - disposition could kill the supervisor and tear the just-started task down - with it. - - The previous dispositions are restored on exit so a long-lived supervisor - process (e.g. a reused Celery prefork worker) does not leak the handler into - later tasks or clobber the worker's own signal handling. - """ - - def _warm_shutdown(signum, frame): - log.info( - "Received signal; warm shutdown in progress, waiting for the running task to complete.", - signal=signal.Signals(signum).name, - pid=os.getpid(), - ) - - prev_sigterm = signal.getsignal(signal.SIGTERM) - prev_sigint = signal.getsignal(signal.SIGINT) - signal.signal(signal.SIGTERM, _warm_shutdown) - signal.signal(signal.SIGINT, _warm_shutdown) - try: - yield - finally: - signal.signal(signal.SIGTERM, prev_sigterm) - signal.signal(signal.SIGINT, prev_sigint) - - -class _PythonCoordinator(BaseCoordinator): - """ - Coordinator implementation to execute Python tasks. - - This is not supposed to be specified by users directly, but the fallback - used by default when nothing is specified. - """ - - def execute_task( - self, - *, - what: TaskInstance, - dag_rel_path: str | PathLike[str], - bundle_info, - client: Client, - logger: FilteringBoundLogger | None = None, - sentry_integration: str = "", - subprocess_logs_to_stdout: bool, - **kwargs, - ) -> BaseCoordinator.ExecutionResult: - # TODO: Importing this at the top causes circular imports. - # ActivitySubprocess and WatchedSubprocess should be moved out of the - # supervisor, and maybe with additional refactoring to abstract out - # process handling. - from airflow.sdk.execution_time.supervisor import ActivitySubprocess - - # Keep the warm-shutdown handlers installed across both start() (which - # transitions the TI to RUNNING) and wait() (which runs the task and - # reports its terminal state / uploads logs) so a SIGTERM at any point - # in this window can't kill the supervisor and tear the task down. - with _warm_shutdown_signals(): - process = ActivitySubprocess.start( - dag_rel_path=dag_rel_path, - what=what, - client=client, - logger=logger, - bundle_info=bundle_info, - subprocess_logs_to_stdout=subprocess_logs_to_stdout, - sentry_integration=sentry_integration, - ) - exit_code = process.wait() - return self.ExecutionResult(exit_code, process.final_state) - - @functools.cache -def _build_python_coordinator() -> _PythonCoordinator: - return _PythonCoordinator() +def _build_python_coordinator() -> BaseCoordinator: + from airflow.sdk.coordinators.python import PythonCoordinator # Lazy to avoid circular import. + + return PythonCoordinator() class InvalidCoordinatorError(ValueError): diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 05810f0e45326..516b30595982f 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -3135,6 +3135,7 @@ def supervise_task( _make_process_nondumpable() from airflow.sdk._shared.secrets_masker import reset_secrets_masker + from airflow.sdk.definitions._internal.abstractoperator import DEFAULT_QUEUE if not client: if dry_run and server: @@ -3172,7 +3173,7 @@ def supervise_task( raise ValueError("dag_path is required") try: - coordinator = get_coordinator_manager().for_queue(ti.queue or "default") + coordinator = get_coordinator_manager().for_queue(ti.queue or DEFAULT_QUEUE) except: log.exception( "Failed to initialize coordinator for task", diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py index 3160f2d01db86..193427d5d3fc2 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -1670,7 +1670,7 @@ def _on_term(signum, frame): return state, msg, error try: - result = _execute_task(context=context, ti=ti, log=log) + result = _TaskRunner.run(ti, context, log) log.info("::group::Post Execute") except Exception: import jinja2 @@ -2361,52 +2361,108 @@ def _run_execute_callable( return result -@detail_span("_execute_task") -def _execute_task(context: Context, ti: RuntimeTaskInstance, log: Logger): - """Execute Task (optionally with a Timeout) and push Xcom results.""" - task = ti.task - execute = task.execute +class _TaskRunner: + """ + Run a task instance and return its result. - if ti._ti_context_from_server and (next_method := ti._ti_context_from_server.next_method): - from airflow.sdk.serde import deserialize + Only :meth:`run` is meant to be called from outside. It selects how the task + runs and delegates. A stub task goes to its registered handler; a concrete + operator to its ``execute`` lifecycle (:meth:`_run_operator`). + + .. note:: + + Ideally, the stub-or-not dispatching should live with the coordinator + dispatch. A coordinator is picked by queue, then either runs the handler + (if stub) or the (non-stub) task. This requires a larger refactor and is + deferred to a future PR. + + .. seealso:: Language SDK ADR-0002. + """ + + @classmethod + @detail_span("_execute_task") + def run(cls, ti: RuntimeTaskInstance, context: Context, log: Logger): + """Run the task specified by *ti* with *context*.""" + if ti.task.is_stub: + return cls._run_handler(ti, context, log) + return cls._run_operator(ti, context, log) + + @classmethod + def _run_operator(cls, ti: RuntimeTaskInstance, context: Context, log: Logger): + """Execute a concrete operator.""" + task = ti.task + execute = task.execute + + if ti._ti_context_from_server and (next_method := ti._ti_context_from_server.next_method): + from airflow.sdk.serde import deserialize + + next_kwargs_data = ti._ti_context_from_server.next_kwargs or {} + try: + if TYPE_CHECKING: + assert isinstance(next_kwargs_data, dict) + kwargs = deserialize(next_kwargs_data) + except (ImportError, KeyError, AttributeError, TypeError): + from airflow.serialization.serialized_objects import BaseSerialization + + kwargs = BaseSerialization.deserialize(next_kwargs_data) - next_kwargs_data = ti._ti_context_from_server.next_kwargs or {} - try: if TYPE_CHECKING: - assert isinstance(next_kwargs_data, dict) - kwargs = deserialize(next_kwargs_data) - except (ImportError, KeyError, AttributeError, TypeError): - from airflow.serialization.serialized_objects import BaseSerialization + assert isinstance(kwargs, dict) + execute = functools.partial(task.resume_execution, next_method=next_method, next_kwargs=kwargs) - kwargs = BaseSerialization.deserialize(next_kwargs_data) + # Export context in os.environ to make it available for operators to use. + airflow_context_vars = context_to_airflow_vars(context, in_env_var_format=True) + os.environ.update(airflow_context_vars) - if TYPE_CHECKING: - assert isinstance(kwargs, dict) - execute = functools.partial(task.resume_execution, next_method=next_method, next_kwargs=kwargs) + outlet_events = context_get_outlet_events(context) - # Export context in os.environ to make it available for operators to use. - airflow_context_vars = context_to_airflow_vars(context, in_env_var_format=True) - os.environ.update(airflow_context_vars) + if (pre_execute_hook := task._pre_execute_hook) is not None: + create_executable_runner(pre_execute_hook, outlet_events, logger=log).run(context) + if getattr(pre_execute_hook := task.pre_execute, "__func__", None) is not BaseOperator.pre_execute: + create_executable_runner(pre_execute_hook, outlet_events, logger=log).run(context) - outlet_events = context_get_outlet_events(context) + _run_task_state_change_callbacks(task, "on_execute_callback", context, log) - if (pre_execute_hook := task._pre_execute_hook) is not None: - create_executable_runner(pre_execute_hook, outlet_events, logger=log).run(context) - if getattr(pre_execute_hook := task.pre_execute, "__func__", None) is not BaseOperator.pre_execute: - create_executable_runner(pre_execute_hook, outlet_events, logger=log).run(context) + log.info("::endgroup::") - _run_task_state_change_callbacks(task, "on_execute_callback", context, log) + result = _run_execute_callable(context, execute, task) - log.info("::endgroup::") + if (post_execute_hook := task._post_execute_hook) is not None: + create_executable_runner(post_execute_hook, outlet_events, logger=log).run(context, result) + if getattr(post_execute_hook := task.post_execute, "__func__", None) is not BaseOperator.post_execute: + create_executable_runner(post_execute_hook, outlet_events, logger=log).run(context) - result = _run_execute_callable(context, execute, task) + return result - if (post_execute_hook := task._post_execute_hook) is not None: - create_executable_runner(post_execute_hook, outlet_events, logger=log).run(context, result) - if getattr(post_execute_hook := task.post_execute, "__func__", None) is not BaseOperator.post_execute: - create_executable_runner(post_execute_hook, outlet_events, logger=log).run(context) + @classmethod + def _run_handler(cls, ti: RuntimeTaskInstance, context: Context, log: Logger): + """Run a stub task's registered Python task handler.""" + from airflow.sdk.coordinators.python import PythonCoordinator + from airflow.sdk.definitions._internal.abstractoperator import DEFAULT_QUEUE + from airflow.sdk.execution_time.coordinator import get_coordinator_manager + from airflow.sdk.task_handlers import call_handler - return result + coordinator = get_coordinator_manager().for_queue(ti.queue or DEFAULT_QUEUE) + if not isinstance(coordinator, PythonCoordinator): + raise RuntimeError( + f"No task handler for ({ti.dag_id!r}, {ti.task_id!r}); " + f"{type(coordinator).__name__} does not run Python task handlers." + ) + if (handler := coordinator.resolve_handler(ti.dag_id, ti.task_id, bundle=ti.bundle_instance)) is None: + raise RuntimeError( + f"No task handler for ({ti.dag_id!r}, {ti.task_id!r}); check the " + "coordinator's task_handler_bundle_name bundle." + ) + + def _invoke(context: Context) -> Any: + return call_handler( + handler, + ctx.arg_bindings if (ctx := ti._ti_context_from_server) else None, + lambda task_id, key: ti.xcom_pull(task_ids=task_id, key=key), + ) + + log.info("Running Python task handler", dag_id=ti.dag_id, task_id=ti.task_id) + return _run_execute_callable(context, _invoke, ti.task) def _render_map_index(context: Context, ti: RuntimeTaskInstance, log: Logger) -> str | None: diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py index 207ef0dc1f6c9..2b7aa25748796 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py @@ -45,6 +45,7 @@ import attrs import pendulum +import structlog from dateutil.relativedelta import relativedelta from airflow.sdk import DAG, BaseOperator, XComArg, literal, task as task_decorators @@ -64,7 +65,7 @@ ) if TYPE_CHECKING: - from collections.abc import Callable + from collections.abc import Callable, Iterator from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: SDK002 from airflow.sdk.importers.base import DagDefinition @@ -79,6 +80,8 @@ class YamlDagBuildError(ValueError): _UNSUPPORTED_TASK_ARGS = frozenset({"task_id", "dag", "task_group", "default_args"}) _MARKER_KEYS = frozenset({*XCOM_KEYS, *TEMPLATE_KEYS, *FILE_KEYS, CONST_KEY}) +log = structlog.get_logger(__name__) + @attrs.define(kw_only=True) class DagBuilder: @@ -105,11 +108,64 @@ def build(self, doc: DagDocument) -> DAG: for task in doc.tasks: for upstream in sorted(task.get_dependency_ids()): self._operators[upstream] >> self._operators[task.id_] + self._validate_stub_handlers(doc) dag.bundle_name = self._bundle.name dag.fileloc = repr(self._definition) dag.relative_fileloc = self._definition.get_relative_loc(self._bundle.path) return dag + def _validate_stub_handlers(self, doc: DagDocument) -> None: + """ + Check each ``run:`` task's arguments against its handler. + + This implements the interface defined in AIP-108 and Language SDK + ADR-0011. The routed coordinator is queried to obtain the handler's + signature and compared against the ``run:`` arguments. Mismatches are + collected and raised together as a :class:`YamlDagBuildError`, + eventually recorded as an import error. + + Validation is skipped when no declaration is available (no handler + bundle or an unregistered handler), or resolving one fails (an + uninitializable bundle, an ImportError in it, or a malformed handler). + + This never fails parsing, so building a Dag does not depend on handler + availability. A real problem surfaces when the task runs. + """ + if not (stub_tasks := [task for task in doc.tasks if self._operators[task.id_].is_stub]): + return + if problems := "; ".join(self._iter_stub_handler_collect_problems(doc, stub_tasks)): + raise YamlDagBuildError(problems) + + def _iter_stub_handler_collect_problems(self, doc: DagDocument, stub_tasks: list[Task]) -> Iterator[str]: + from airflow.sdk.execution_time.coordinator import get_coordinator_manager + + manager = get_coordinator_manager() + for task in stub_tasks: + coordinator = manager.for_queue(self._operators[task.id_].queue) + try: + declaration = coordinator.parse_task_handler( + dag_id=doc.dag_id, + task_id=task.id_, + bundle=self._bundle, + ) + except Exception: + log.warning( + "Could not resolve task handler; skipping parse-time validation", + dag_id=doc.dag_id, + task_id=task.id_, + exc_info=True, + ) + continue + if declaration is None: + continue # Best-effort: nothing to match against; defer to execution. + run_names = set(task.run) if task.run else set() + if missing := sorted({p.name for p in declaration.params if p.required} - run_names): + yield f"task {task.id_!r}: missing required handler argument(s) {missing}" + if unknown := sorted(run_names - {p.name for p in declaration.params}): + yield f"task {task.id_!r}: run argument(s) {unknown} are not parameters of its task handler" + # Only names and arity are checked here. Type validation will follow + # once Language SDK ADR-0011 is settled. + def _build_stub_task(self, task: Task, args: dict[str, Value]) -> BaseOperator: """ Build a stub from a "run:" task. diff --git a/task-sdk/src/airflow/sdk/task_handlers/__init__.py b/task-sdk/src/airflow/sdk/task_handlers/__init__.py new file mode 100644 index 0000000000000..6fa98d6eb7b66 --- /dev/null +++ b/task-sdk/src/airflow/sdk/task_handlers/__init__.py @@ -0,0 +1,42 @@ +# 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. +""" +Python task handlers: Python's equivalent to Language SDK task handlers. + +A function decorated by ``@task_handler(dag_id, task_id)`` is the Python body of +a stub task. +""" + +from __future__ import annotations + +from airflow.sdk.task_handlers.binding import call_handler +from airflow.sdk.task_handlers.declaration import ( + TaskHandlerDeclaration, + TaskHandlerParam, +) +from airflow.sdk.task_handlers.loader import import_handler_modules, resolve_handler_bundle +from airflow.sdk.task_handlers.registry import get_handler, task_handler + +__all__ = [ + "import_handler_modules", + "resolve_handler_bundle", + "TaskHandlerDeclaration", + "TaskHandlerParam", + "call_handler", + "get_handler", + "task_handler", +] diff --git a/task-sdk/src/airflow/sdk/task_handlers/binding.py b/task-sdk/src/airflow/sdk/task_handlers/binding.py new file mode 100644 index 0000000000000..357dfb2eea644 --- /dev/null +++ b/task-sdk/src/airflow/sdk/task_handlers/binding.py @@ -0,0 +1,79 @@ +# 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. + +""" +Runtime argument binder. + +At execution time, a stub task's inputs arrive as ``TIRunContext.arg_bindings``. +This converts it into keyword arguments, and uses them to call the registered +task handler. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from airflow.sdk.api.datamodels._generated import LiteralArgBinding, XComArgBinding +from airflow.sdk.bases.xcom import BaseXCom + +if TYPE_CHECKING: + from collections.abc import Callable, Iterator + + from airflow.sdk.api.datamodels._generated import TaskArgBinding + + XComResolver = Callable[[str, str], Any] + +__all__ = ["call_handler"] + + +def _bind_arguments( + arg_bindings: list[TaskArgBinding] | None, + xcom_resolver: XComResolver, + *, + skip_defaults: bool = True, +) -> Iterator[tuple[str, Any]]: + """ + Resolve an ordered arg-binding spec into handler keyword arguments. + + * a ``LiteralArgBinding`` -> its inline ``value``. + * an ``XComArgBinding`` -> ``xcom_resolver(task_id, "return_value")``. + + ``skip_defaults``: when true (default), literal entries flagged ``from_default`` are omitted + so the handler's own default applies -- matching ADR-0007's "defaults are explicit" so a + keyword consumer can leave them unclaimed. Set false to bind them explicitly. + + ``None`` (an argless call) yields nothing. + """ + for binding in arg_bindings or (): + match binding.root: + case LiteralArgBinding(from_default=True) if skip_defaults: + continue + case LiteralArgBinding(name=name, value=value): + yield name, value + case XComArgBinding(name=name, task_id=task_id): + yield name, xcom_resolver(task_id, BaseXCom.XCOM_RETURN_KEY) + + +def call_handler( + fn: Callable[..., Any], + arg_bindings: list[TaskArgBinding] | None, + xcom_resolver: XComResolver, + *, + skip_defaults: bool = True, +) -> Any: + """Bind ``arg_bindings`` and invoke ``fn``, returning its result (the task's return_value XCom).""" + return fn(**dict(_bind_arguments(arg_bindings, xcom_resolver, skip_defaults=skip_defaults))) diff --git a/task-sdk/src/airflow/sdk/task_handlers/declaration.py b/task-sdk/src/airflow/sdk/task_handlers/declaration.py new file mode 100644 index 0000000000000..91572a3bbc621 --- /dev/null +++ b/task-sdk/src/airflow/sdk/task_handlers/declaration.py @@ -0,0 +1,75 @@ +# 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. +""" +Introspect a Python handler for mixed-language validation. + +This is conceptually the inverse of +:func:`airflow.serialization.stub_arg_bindings.build_arg_bindings`. That +function builds the stub side from the caller; this builds the handler side. +""" + +from __future__ import annotations + +import dataclasses +from inspect import Parameter, signature +from typing import TYPE_CHECKING + +from airflow.sdk.definitions.context import KNOWN_CONTEXT_KEYS + +if TYPE_CHECKING: + from collections.abc import Callable + +__all__ = ["TaskHandlerParam", "TaskHandlerDeclaration"] + + +@dataclasses.dataclass(frozen=True, kw_only=True) +class TaskHandlerParam: + """ + One declared handler parameter (ADR-0012). + + Arguments are bound to the handler by name. + """ + + name: str + required: bool + + @classmethod + def from_parameter(cls, parameter: Parameter) -> TaskHandlerParam: + if parameter.kind in (Parameter.VAR_POSITIONAL, Parameter.VAR_KEYWORD): + raise ValueError("may not use variadic *args or **kwargs") + if parameter.kind is Parameter.POSITIONAL_ONLY: + raise ValueError(f"may not use positional-only parameter {parameter.name!r}") + if (name := parameter.name) in KNOWN_CONTEXT_KEYS: + raise ValueError(f"argument {name!r} is a reserved context key") + return cls(name=name, required=parameter.default is Parameter.empty) + + +@dataclasses.dataclass +class TaskHandlerDeclaration: + """ + The signature a handler declares. + + This is the parse-time argument specification compared against a stub's + captured arguments. It carries only the parameters: the caller holds the + dag_id/task_id it queried, and the callable to run is resolved separately. + """ + + params: list[TaskHandlerParam] + + @classmethod + def from_handler(cls, fn: Callable) -> TaskHandlerDeclaration: + return cls([TaskHandlerParam.from_parameter(p) for p in signature(fn).parameters.values()]) diff --git a/task-sdk/src/airflow/sdk/task_handlers/loader.py b/task-sdk/src/airflow/sdk/task_handlers/loader.py new file mode 100644 index 0000000000000..fabf499a73c6e --- /dev/null +++ b/task-sdk/src/airflow/sdk/task_handlers/loader.py @@ -0,0 +1,111 @@ +# 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 the Python modules that register ``@task_handler`` functions. + +A non-Python Dag defines a stub task (e.g. YAML's ``run:`` task) but imports no +Python; the ``@task_handler`` bodies live in Python modules in a Dag bundle. +Importing those modules runs the decorators, which register each handler in the +in-memory registry keyed by ``(dag_id, task_id)``. The bundle is either named +(the coordinator's ``task_handler_bundle_name``) or, when unset, the stub's own +Dag bundle. + +Resolving a named bundle and importing its modules are separate steps: +:func:`resolve_handler_bundle` turns a name into a pinned bundle, and +:func:`import_handler_modules` imports any already-materialized bundle. Both are +done at most once per process, but Dag processing and task workers are +short-lived, so each process does its own. Every non-ignored Python file under +the bundle is imported, so a handler bundle should not also hold Python Dag files +(see the task-handler documentation). +""" + +from __future__ import annotations + +import contextlib +import importlib.util +import pathlib +import sys +from typing import TYPE_CHECKING + +import structlog + +from airflow.sdk._shared.module_loading.dag_file import get_unique_dag_module_name +from airflow.sdk._shared.module_loading.file_discovery import find_path_from_directory + +if TYPE_CHECKING: + from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: SDK002 + +log = structlog.get_logger(logger_name="task_handlers.loader") + +_resolved_bundles: dict[str, BaseDagBundle] = {} +_imported_bundles: set[str] = set() + + +def resolve_handler_bundle(name: str) -> BaseDagBundle: + """ + Materialize the named handler bundle at a pinned version. + + Resolves *name* through ``DagBundlesManager`` and pins it to a concrete + version, once per process. + """ + with contextlib.suppress(KeyError): + return _resolved_bundles[name] + + from airflow.sdk.api.datamodels._generated import BundleInfo + from airflow.sdk.execution_time.bundles import initialize_pinned_bundle + + _resolved_bundles[name] = bundle = initialize_pinned_bundle(BundleInfo(name=name), log) + return bundle + + +def import_handler_modules(bundle: BaseDagBundle) -> None: + """ + Import *bundle*'s handler modules so ``@task_handler`` decorators register. + + Imports every non-ignored Python file in *bundle*. The bundle must already + be materialized. Importing is idempotent per process by bundle name. + """ + from airflow.dag_processing.bundles.base import BundleVersionLock # noqa: SDK002 + + if bundle.name in _imported_bundles: + return + with BundleVersionLock(bundle_name=bundle.name, bundle_version=bundle.version): + _import_python_files(bundle.path) + _imported_bundles.add(bundle.name) + log.debug("Imported task-handler bundle", bundle=bundle.name, path=bundle.path) + + +def _import_python_files(root: pathlib.Path) -> None: + """Import every non-ignored ``.py`` under *root*.""" + for found in find_path_from_directory(root, ".airflowignore"): + path = pathlib.Path(found) + if path.suffix != ".py" or "__pycache__" in path.parts: + continue + if (module_name := get_unique_dag_module_name(found)) in sys.modules: + continue + if (spec := importlib.util.spec_from_file_location(module_name, path)) is None or spec.loader is None: + continue + module = importlib.util.module_from_spec(spec) + # Register before exec so module-level code that looks itself up in + # sys.modules works (e.g. @dataclass). Drop it on failure. + sys.modules[module_name] = module + try: + spec.loader.exec_module(module) + except BaseException: + sys.modules.pop(module_name, None) + raise diff --git a/task-sdk/src/airflow/sdk/task_handlers/registry.py b/task-sdk/src/airflow/sdk/task_handlers/registry.py new file mode 100644 index 0000000000000..661ef098915ea --- /dev/null +++ b/task-sdk/src/airflow/sdk/task_handlers/registry.py @@ -0,0 +1,56 @@ +# 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 typing import TypeVar + +TaskHandlerF = TypeVar("TaskHandlerF", bound=Callable) + +_REGISTRY: dict[tuple[str, str], Callable] = {} + + +def get_handler(dag_id: str, task_id: str) -> Callable | None: + """Get the handler registered for the task, or *None* if none is.""" + return _REGISTRY.get((dag_id, task_id)) + + +def task_handler(dag_id: str, task_id: str | None = None) -> Callable[[TaskHandlerF], TaskHandlerF]: + """ + Register the decorated function as the body of a stub task. + + If not specified, *task_id* defaults to the function name. The decorated + function is returned unchanged. + + Example: + + .. code-block:: python + + @task_handler("events_pipeline") + def transform(region: str, events: list[dict]) -> dict: ... + """ + + def _task_handler(fn: TaskHandlerF) -> TaskHandlerF: + if not (resolved_task_id := task_id or getattr(fn, "__name__", None)): + raise ValueError("task_id could not be derived from the handler; pass task_id explicitly") + if (key := (dag_id, resolved_task_id)) in _REGISTRY: + raise ValueError(f"handler already registered for dag_id={dag_id!r} task_id={resolved_task_id!r}") + _REGISTRY[key] = fn + return fn + + return _task_handler 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 37ac3fe9f73d4..0de0aa5c627d7 100644 --- a/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py @@ -714,7 +714,7 @@ def test_runs_a_bundle_found_by_its_bundle_relative_name(self, tmp_path, monkeyp try: with ( conf_vars(coordinators), - patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle", autospec=True) as init, + patch("airflow.sdk.execution_time.bundles.initialize_ti_bundle", autospec=True) as init, patch("airflow.sdk.coordinators._subprocess.BundleVersionLock", autospec=True), patch.object(_PopenActivitySubprocess, "start", autospec=True) as mock_start, patch.object( 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 36e7616089bc7..15ebfb75426d0 100644 --- a/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py @@ -497,7 +497,7 @@ def mock_start(self, tmp_path): bundle = MagicMock(path=tmp_path, version="v1") bundle.name = "dags-folder" with ( - patch("airflow.sdk.coordinators._subprocess._initialize_pinned_bundle", return_value=bundle), + patch("airflow.sdk.coordinators._subprocess.initialize_pinned_bundle", return_value=bundle), patch("airflow.sdk.coordinators._subprocess.BundleVersionLock"), patch.object(_PopenActivitySubprocess, "start") as mock_start, ): 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 2d2e4ba6b82f9..9208b3d3fe819 100644 --- a/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py @@ -138,7 +138,7 @@ def mock_start(self, tmp_path): bundle = MagicMock(path=tmp_path, version="v1") bundle.name = "dags-folder" with ( - patch("airflow.sdk.coordinators._subprocess._initialize_pinned_bundle", return_value=bundle), + patch("airflow.sdk.coordinators._subprocess.initialize_pinned_bundle", return_value=bundle), patch("airflow.sdk.coordinators._subprocess.BundleVersionLock"), patch.object(_PopenActivitySubprocess, "start") as mock_start, ): diff --git a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py index b9f16aeeb9a84..94aaea7ec31d1 100644 --- a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py +++ b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py @@ -1367,7 +1367,7 @@ def test_arms_execution_timeout_from_run_context( class TestResolveArtifactBundle: """Execute-time resolution of the Dag bundle holding the artifacts.""" - @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle", autospec=True) + @patch("airflow.sdk.execution_time.bundles.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 @@ -1377,7 +1377,7 @@ def test_unset_name_uses_the_task_bundle(self, mock_initialize, tmp_path): 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", autospec=True) + @patch("airflow.sdk.execution_time.bundles.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. @@ -1402,7 +1402,7 @@ def test_version_less_bundle_is_rematerialized_at_its_current_version(self, mock call(BundleInfo(name="artifacts", version="sha-abc", version_data={"k": "v"})), ] - @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle", autospec=True) + @patch("airflow.sdk.execution_time.bundles.initialize_ti_bundle", autospec=True) def test_unversioned_bundle_keeps_its_single_path(self, mock_initialize, tmp_path): resolved = _make_bundle(tmp_path) resolved.get_current_version.return_value = None @@ -1413,7 +1413,7 @@ def test_unversioned_bundle_keeps_its_single_path(self, mock_initialize, tmp_pat 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", autospec=True) + @patch("airflow.sdk.execution_time.bundles.initialize_ti_bundle", autospec=True) def test_missing_resolved_path_raises(self, mock_initialize, tmp_path): mock_initialize.return_value = _make_bundle(tmp_path / "nope", version="v1") coordinator = _StubSubprocessCoordinator(command=["x"], task_handler_bundle_name="artifacts") @@ -1474,7 +1474,7 @@ def test_fails_without_a_retry_for_a_coordinator_bug(self, mock_client): assert isinstance(raised.value.__cause__, RuntimeError) @patch( - "airflow.sdk.coordinators._subprocess.initialize_ti_bundle", + "airflow.sdk.execution_time.bundles.initialize_ti_bundle", autospec=True, side_effect=RuntimeError("Error cloning repository"), ) @@ -1663,7 +1663,7 @@ class TestExecuteTaskBundleWiring: """execute_task passes the task bundle, forwards resolved roots, and holds the version lock.""" @patch("airflow.sdk.coordinators._subprocess.BundleVersionLock", autospec=True) - @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle", autospec=True) + @patch("airflow.sdk.execution_time.bundles.initialize_ti_bundle", autospec=True) @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 @@ -1689,7 +1689,7 @@ def test_task_bundle_mode_binds_forwards_roots_and_locks( mock_lock.return_value.__exit__.assert_called_once() @patch("airflow.sdk.coordinators._subprocess.BundleVersionLock", autospec=True) - @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle", autospec=True) + @patch("airflow.sdk.execution_time.bundles.initialize_ti_bundle", autospec=True) @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 @@ -1804,7 +1804,7 @@ def bundle(self, tmp_path): @pytest.fixture def mock_initialize(self, bundle): with patch( - "airflow.sdk.coordinators._subprocess.initialize_ti_bundle", autospec=True, return_value=bundle + "airflow.sdk.execution_time.bundles.initialize_ti_bundle", autospec=True, return_value=bundle ) as m: yield m diff --git a/task-sdk/tests/task_sdk/docs/test_public_api.py b/task-sdk/tests/task_sdk/docs/test_public_api.py index 0b00b32bc2809..51cdb6ca78e50 100644 --- a/task-sdk/tests/task_sdk/docs/test_public_api.py +++ b/task-sdk/tests/task_sdk/docs/test_public_api.py @@ -77,6 +77,7 @@ def test_airflow_sdk_no_unexpected_exports(): "timezone", "types", "state", + "task_handlers", "yaml", } unexpected = actual - public - ignore 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 6ade4fb1e1233..41ba25a27d991 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_coordinator.py +++ b/task-sdk/tests/task_sdk/execution_time/test_coordinator.py @@ -27,12 +27,11 @@ from airflow.sdk._shared.module_loading import import_string from airflow.sdk.configuration import conf, retrieve_configuration_description +from airflow.sdk.coordinators.python import PythonCoordinator, _warm_shutdown_signals from airflow.sdk.execution_time.coordinator import ( BaseCoordinator, CoordinatorManager, InvalidCoordinatorError, - _PythonCoordinator, - _warm_shutdown_signals, get_coordinator_manager, reset_coordinator_manager, ) @@ -143,7 +142,7 @@ def test_from_config_loads_specs_and_resolves_instances(self, sdk_config): assert manager._created_coordinators == {"alpha": coordinator_for_queue_a} coordinator_for_queue_missing = manager.for_queue("queue-1") - assert isinstance(coordinator_for_queue_missing, _PythonCoordinator) + assert isinstance(coordinator_for_queue_missing, PythonCoordinator) assert manager.for_queue("queue-1") is coordinator_for_queue_missing assert manager._created_coordinators == {"alpha": coordinator_for_queue_a} @@ -780,7 +779,7 @@ def _fake_start(*args, **kwargs): monkeypatch.setattr(supervisor_mod.ActivitySubprocess, "start", staticmethod(_fake_start)) - coordinator = _PythonCoordinator() + coordinator = PythonCoordinator() result = coordinator.execute_task( what=mock.MagicMock(), dag_rel_path="some_dag.py", diff --git a/task-sdk/tests/task_sdk/execution_time/test_python_coordinator.py b/task-sdk/tests/task_sdk/execution_time/test_python_coordinator.py new file mode 100644 index 0000000000000..eaaa3564525b9 --- /dev/null +++ b/task-sdk/tests/task_sdk/execution_time/test_python_coordinator.py @@ -0,0 +1,166 @@ +# 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 re +import sys +import textwrap +from types import SimpleNamespace + +import pytest + +from airflow.sdk._shared.module_loading.dag_file import UNUSUAL_MODULE_PREFIX +from airflow.sdk.coordinators.python import PythonCoordinator +from airflow.sdk.execution_time.coordinator import BaseCoordinator, CoordinatorManager +from airflow.sdk.task_handlers import loader as loader_mod +from airflow.sdk.task_handlers.registry import _REGISTRY + + +def _purge_handler_modules(): + for name in [n for n in sys.modules if n.startswith(UNUSUAL_MODULE_PREFIX)]: + del sys.modules[name] + + +@pytest.fixture(autouse=True) +def _clean_registry(): + _REGISTRY.clear() + loader_mod._imported_bundles.clear() + loader_mod._resolved_bundles.clear() + _purge_handler_modules() + yield + _REGISTRY.clear() + loader_mod._imported_bundles.clear() + loader_mod._resolved_bundles.clear() + _purge_handler_modules() + + +def _bundle(path): + """A stand-in for the stub's own, already-materialized Dag bundle.""" + return SimpleNamespace(name="own", version=None, path=path) + + +def _write_handler(path, body): + (path / "handlers.py").write_text(textwrap.dedent(body)) + + +class TestParseTaskHandler: + def test_parses_handler_from_the_tasks_own_bundle(self, tmp_path): + _write_handler( + tmp_path, + """ + from airflow.sdk import task_handler + + @task_handler("d", "t") + def t(region, events): + return len(events) + """, + ) + + coordinator = PythonCoordinator(task_handler_bundle_name=None) + decl = coordinator.parse_task_handler(dag_id="d", task_id="t", bundle=_bundle(tmp_path)) + + assert decl is not None + assert [p.name for p in decl.params] == ["region", "events"] + + def test_returns_none_when_no_handler_registered(self, tmp_path): + coordinator = PythonCoordinator(task_handler_bundle_name=None) + assert coordinator.parse_task_handler(dag_id="d", task_id="x", bundle=_bundle(tmp_path)) is None + + def test_base_coordinator_default_returns_none(self, tmp_path): + # The default (foreign / not-adopted) coordinator declares nothing; validation defers. + base = BaseCoordinator() + assert base.parse_task_handler(dag_id="d", task_id="t", bundle=_bundle(tmp_path)) is None + + def test_invalid_signature_is_reported_against_the_task(self, tmp_path): + _write_handler( + tmp_path, + """ + from airflow.sdk import task_handler + + @task_handler("d", "t") + def t(*args): ... + """, + ) + + # from_handler raises a bare signature error; parse_task_handler adds the task identity. + coordinator = PythonCoordinator(task_handler_bundle_name=None) + with pytest.raises( + ValueError, match=re.escape("task handler 't' may not use variadic *args or **kwargs") + ): + coordinator.parse_task_handler(dag_id="d", task_id="t", bundle=_bundle(tmp_path)) + + +class TestResolveHandler: + def test_resolves_the_handler_from_the_tasks_own_bundle(self, tmp_path): + _write_handler( + tmp_path, + """ + from airflow.sdk import task_handler + + @task_handler("d", "t") + def t(region): + return region + """, + ) + coordinator = PythonCoordinator(task_handler_bundle_name=None) + handler = coordinator.resolve_handler(dag_id="d", task_id="t", bundle=_bundle(tmp_path)) + + assert handler is not None + assert handler.__name__ == "t" + + def test_returns_none_when_no_handler_registered(self, tmp_path): + coordinator = PythonCoordinator(task_handler_bundle_name=None) + assert coordinator.resolve_handler(dag_id="d", task_id="x", bundle=_bundle(tmp_path)) is None + + def test_is_not_a_base_coordinator_hook(self): + # Execution resolution is Python-specific; it is deliberately not on BaseCoordinator. + assert not hasattr(BaseCoordinator(), "resolve_handler") + + +class TestConfigurable: + def test_configured_in_sdk_coordinators_carries_bundle_name(self, monkeypatch): + from airflow.dag_processing.bundles import manager as mgr + + monkeypatch.setattr(mgr.DagBundlesManager, "is_bundle_configured", lambda name: True) + monkeypatch.setenv( + "AIRFLOW__SDK__COORDINATORS", + json.dumps( + { + "py": { + "classpath": "airflow.sdk.coordinators.python.PythonCoordinator", + "kwargs": {"task_handler_bundle_name": "handlers"}, + } + } + ), + ) + monkeypatch.setenv("AIRFLOW__SDK__QUEUE_TO_COORDINATOR", json.dumps({"myqueue": "py"})) + + coordinator = CoordinatorManager.from_config().for_queue("myqueue") + + assert isinstance(coordinator, PythonCoordinator) + assert coordinator.task_handler_bundle_name == "handlers" + + def test_unmapped_queue_falls_back_to_bundleless_python(self, monkeypatch): + monkeypatch.setenv("AIRFLOW__SDK__COORDINATORS", "{}") + monkeypatch.setenv("AIRFLOW__SDK__QUEUE_TO_COORDINATOR", "{}") + + coordinator = CoordinatorManager.from_config().for_queue("whatever") + + assert isinstance(coordinator, PythonCoordinator) + assert coordinator.task_handler_bundle_name is None diff --git a/task-sdk/tests/task_sdk/execution_time/test_run_task_handler.py b/task-sdk/tests/task_sdk/execution_time/test_run_task_handler.py new file mode 100644 index 0000000000000..dfe5fa8a4c2de --- /dev/null +++ b/task-sdk/tests/task_sdk/execution_time/test_run_task_handler.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. + +from __future__ import annotations + +from types import SimpleNamespace +from unittest import mock + +import pytest + +from airflow.sdk.api.datamodels._generated import ( + LiteralArgBinding, + TaskArgBinding, + XComArgBinding, +) +from airflow.sdk.coordinators.python import PythonCoordinator +from airflow.sdk.execution_time.task_runner import _TaskRunner + + +def _literal(name, value): + return TaskArgBinding(root=LiteralArgBinding(name=name, kind="literal", value=value)) + + +def _xcom(name, task_id): + return TaskArgBinding(root=XComArgBinding(name=name, kind="xcom", task_id=task_id)) + + +def _fake_ti(*, arg_bindings, xcom_values=None, queue="default"): + xcom_values = xcom_values or {} + pulled = [] + + def xcom_pull(*, task_ids, key): + pulled.append((task_ids, key)) + return xcom_values.get(task_ids) + + ti = SimpleNamespace( + dag_id="events_pipeline", + task_id="transform", + queue=queue, + task=SimpleNamespace(queue=queue, is_stub=True, execution_timeout=None), + bundle_instance=SimpleNamespace(name="b", version=None, path="/tmp/does-not-matter"), + _ti_context_from_server=SimpleNamespace(arg_bindings=arg_bindings), + xcom_pull=xcom_pull, + ) + ti.pulled = pulled + return ti + + +def _run(ti): + return _TaskRunner.run(ti, mock.MagicMock(), mock.MagicMock()) + + +def _patch_coordinator(monkeypatch, handler): + class _Coordinator(PythonCoordinator): + def resolve_handler(self, dag_id, task_id, *, bundle): + return handler + + coordinator = _Coordinator() + manager = SimpleNamespace(for_queue=lambda queue: coordinator) + monkeypatch.setattr("airflow.sdk.execution_time.coordinator.get_coordinator_manager", lambda: manager) + + +def test_binds_literal_and_xcom_then_calls_handler(monkeypatch): + def transform(region, events): + return f"{region}:{len(events)}" + + _patch_coordinator(monkeypatch, transform) + ti = _fake_ti( + arg_bindings=[_literal("region", "eu"), _xcom("events", "extract")], + xcom_values={"extract": [1, 2, 3]}, + ) + + result = _run(ti) + + assert result == "eu:3" + assert ti.pulled == [("extract", "return_value")] + + +def test_no_arg_bindings_calls_handler_with_nothing(monkeypatch): + def extract(): + return [1, 2] + + _patch_coordinator(monkeypatch, extract) + ti = _fake_ti(arg_bindings=None) + + assert _run(ti) == [1, 2] + + +def test_missing_declaration_raises(monkeypatch): + _patch_coordinator(monkeypatch, None) + ti = _fake_ti(arg_bindings=None) + + with pytest.raises(RuntimeError, match="No task handler"): + _run(ti) + + +def test_run_dispatches_on_is_stub(monkeypatch): + seen = [] + monkeypatch.setattr( + _TaskRunner, "_run_handler", classmethod(lambda cls, ti, context, log: seen.append("handler")) + ) + monkeypatch.setattr( + _TaskRunner, "_run_operator", classmethod(lambda cls, ti, context, log: seen.append("operator")) + ) + + _run(_fake_ti(arg_bindings=None)) + + operator_ti = _fake_ti(arg_bindings=None) + operator_ti.task.is_stub = False + _run(operator_ti) + + assert seen == ["handler", "operator"] + + +def test_handler_runs_through_execute_callable(monkeypatch): + # The handler must go through _run_execute_callable so execution_timeout, on_kill and the + # task.execute span apply to it exactly as to a concrete operator. + seen = {} + + def fake_run_execute_callable(context, execute, task): + seen["task"] = task + return execute(context=context) + + monkeypatch.setattr( + "airflow.sdk.execution_time.task_runner._run_execute_callable", fake_run_execute_callable + ) + + def transform(region): + return region + + _patch_coordinator(monkeypatch, transform) + ti = _fake_ti(arg_bindings=[_literal("region", "eu")]) + + assert _run(ti) == "eu" + assert seen["task"] is ti.task diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py index 387543a0bff71..11d9d295df0eb 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py +++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py @@ -185,13 +185,13 @@ TaskRunnerMarker, _axcom_push, _defer_task, - _execute_task, _find_native_dag_importer, _make_task_span, _push_xcom_if_needed, _register_deserialization_allowed_classes, _run_execute_callable, _serialize_outlet_events, + _TaskRunner, _xcom_push, check_xcom_backend_supports_loops, finalize, @@ -2113,7 +2113,7 @@ def sleep_and_catch_other_exceptions(): ti = create_runtime_ti(task=op, dag_id="dag_execution_timeout") with pytest.raises(AirflowTaskTimeout): - _execute_task(context=ti.get_template_context(), ti=ti, log=mock.MagicMock()) + _TaskRunner.run(ti, ti.get_template_context(), mock.MagicMock()) def test_basic_templated_dag(mocked_parse, make_ti_context, mock_supervisor_comms, spy_agency): diff --git a/task-sdk/tests/task_sdk/importers/yaml_importer/test_builder_validation.py b/task-sdk/tests/task_sdk/importers/yaml_importer/test_builder_validation.py new file mode 100644 index 0000000000000..942aeac50b2f2 --- /dev/null +++ b/task-sdk/tests/task_sdk/importers/yaml_importer/test_builder_validation.py @@ -0,0 +1,169 @@ +# 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 sys +import textwrap +from types import SimpleNamespace + +import pytest + +from airflow.sdk._shared.module_loading.dag_file import UNUSUAL_MODULE_PREFIX +from airflow.sdk.importers.base import FilesystemDagDefinition +from airflow.sdk.importers.yaml_importer.importer import YamlDagImporter +from airflow.sdk.task_handlers import TaskHandlerDeclaration, TaskHandlerParam, loader as loader_mod +from airflow.sdk.task_handlers.registry import _REGISTRY + +HEAD_SCHEMA = "https://airflow.apache.org/schemas/dag/2026-10-30.json" + + +@pytest.fixture(autouse=True) +def _clean_handlers(): + def _reset(): + _REGISTRY.clear() + loader_mod._imported_bundles.clear() + loader_mod._resolved_bundles.clear() + for name in [n for n in sys.modules if n.startswith(UNUSUAL_MODULE_PREFIX)]: + del sys.modules[name] + + _reset() + yield + _reset() + + +@pytest.fixture +def mock_bundle(tmp_path): + bundle_dir = tmp_path / "bundle" + bundle_dir.mkdir(parents=True, exist_ok=True) + return SimpleNamespace(name="test_bundle", version=None, path=bundle_dir) + + +def _build(mock_bundle, tasks: str): + body = f"$schema: {HEAD_SCHEMA}\ndag_id: d\ntasks:\n{textwrap.indent(textwrap.dedent(tasks), ' ')}" + path = mock_bundle.path / "dag.yaml" + path.write_text(body) + definition = FilesystemDagDefinition(path=path) + return YamlDagImporter().import_definition(definition, bundle=mock_bundle) + + +def _patch_declaration(monkeypatch, declaration): + coordinator = SimpleNamespace(parse_task_handler=lambda dag_id, task_id, bundle: declaration) + manager = SimpleNamespace(for_queue=lambda queue: coordinator) + monkeypatch.setattr("airflow.sdk.execution_time.coordinator.get_coordinator_manager", lambda: manager) + + +def test_no_declaration_available_skips_and_builds(mock_bundle): + result = _build(mock_bundle, "- {id: t, run: {region: eu}}") + assert result.success, result.errors + assert getattr(result.dags[0].get_task("t"), "is_stub", False) is True + + +def test_matching_arguments_pass(mock_bundle, monkeypatch): + _patch_declaration( + monkeypatch, + TaskHandlerDeclaration([TaskHandlerParam(name="region", required=True)]), + ) + result = _build(mock_bundle, "- {id: t, run: {region: eu}}") + assert result.success, result.errors + + +def test_missing_required_argument_errors(mock_bundle, monkeypatch): + _patch_declaration( + monkeypatch, + TaskHandlerDeclaration([TaskHandlerParam(name="region", required=True)]), + ) + result = _build(mock_bundle, "- {id: t, run: {}}") + assert not result.dags + assert any("missing required handler argument(s) ['region']" in e.message for e in result.errors) + + +def test_unknown_run_argument_errors(mock_bundle, monkeypatch): + _patch_declaration( + monkeypatch, + TaskHandlerDeclaration([TaskHandlerParam(name="region", required=True)]), + ) + result = _build(mock_bundle, "- {id: t, run: {region: eu, bogus: 1}}") + assert not result.dags + assert any("run argument(s) ['bogus'] are not parameters" in e.message for e in result.errors) + + +def _patch_raising_coordinator(monkeypatch, exc): + def describe(dag_id, task_id, bundle): + raise exc + + coordinator = SimpleNamespace(parse_task_handler=describe) + manager = SimpleNamespace(for_queue=lambda queue: coordinator) + monkeypatch.setattr("airflow.sdk.execution_time.coordinator.get_coordinator_manager", lambda: manager) + + +def test_handler_resolution_failure_skips_and_builds(mock_bundle, monkeypatch): + _patch_raising_coordinator(monkeypatch, RuntimeError("bundle fetch failed")) + result = _build(mock_bundle, "- {id: t, run: {region: eu}}") + assert result.success, result.errors + assert getattr(result.dags[0].get_task("t"), "is_stub", False) is True + + +def test_optional_argument_may_be_omitted(mock_bundle, monkeypatch): + _patch_declaration( + monkeypatch, + TaskHandlerDeclaration( + params=[ + TaskHandlerParam(name="region", required=True), + TaskHandlerParam(name="limit", required=False), + ], + ), + ) + result = _build(mock_bundle, "- {id: t, run: {region: eu}}") + assert result.success, result.errors + + +def test_coordinator_resolved_by_task_queue(mock_bundle, monkeypatch): + queues_seen = [] + declaration = TaskHandlerDeclaration([TaskHandlerParam(name="region", required=True)]) + coordinator = SimpleNamespace(parse_task_handler=lambda dag_id, task_id, bundle: declaration) + + def for_queue(queue): + queues_seen.append(queue) + return coordinator + + monkeypatch.setattr( + "airflow.sdk.execution_time.coordinator.get_coordinator_manager", + lambda: SimpleNamespace(for_queue=for_queue), + ) + result = _build(mock_bundle, "- {id: t, run: {region: eu}}") + assert result.success, result.errors + assert queues_seen == [result.dags[0].get_task("t").queue] + + +def test_colocated_handler_is_validated(mock_bundle): + # With no configured handler bundle, a handler in the stub's own bundle is + # auto-discovered and validated at parse time -- no coordinator patch here. + (mock_bundle.path / "handlers.py").write_text( + textwrap.dedent( + """ + from airflow.sdk import task_handler + + @task_handler("d", "t") + def t(region): + return region + """ + ) + ) + result = _build(mock_bundle, "- {id: t, run: {region: eu, bogus: 1}}") + assert not result.dags + assert any("run argument(s) ['bogus'] are not parameters" in e.message for e in result.errors) diff --git a/task-sdk/tests/task_sdk/task_handlers/test_loader.py b/task-sdk/tests/task_sdk/task_handlers/test_loader.py new file mode 100644 index 0000000000000..cd96edadbba9e --- /dev/null +++ b/task-sdk/tests/task_sdk/task_handlers/test_loader.py @@ -0,0 +1,165 @@ +# 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 sys +import textwrap +from types import SimpleNamespace + +import pytest + +from airflow.sdk._shared.module_loading.dag_file import UNUSUAL_MODULE_PREFIX +from airflow.sdk.task_handlers import ( + get_handler, + import_handler_modules, + loader as loader_mod, + resolve_handler_bundle, +) +from airflow.sdk.task_handlers.registry import _REGISTRY + + +def _purge_handler_modules(): + for name in [n for n in sys.modules if n.startswith(UNUSUAL_MODULE_PREFIX)]: + del sys.modules[name] + + +@pytest.fixture(autouse=True) +def _clean(): + _REGISTRY.clear() + loader_mod._imported_bundles.clear() + loader_mod._resolved_bundles.clear() + _purge_handler_modules() + yield + _REGISTRY.clear() + loader_mod._imported_bundles.clear() + loader_mod._resolved_bundles.clear() + _purge_handler_modules() + + +def _write_handler(path, dag_id, task_id, params="region, events"): + path.write_text( + textwrap.dedent( + f""" + from airflow.sdk import task_handler + + @task_handler({dag_id!r}, {task_id!r}) + def {task_id}({params}): + return 1 + """ + ) + ) + + +def _materialized(name, path): + """Stand-in for an already-materialized bundle (name + path).""" + return SimpleNamespace(name=name, version=None, path=path) + + +class TestImportHandlerModules: + def test_imports_modules_and_populates_registry(self, tmp_path): + _write_handler(tmp_path / "handlers.py", "events", "transform") + + import_handler_modules(_materialized("b", tmp_path)) + + assert get_handler("events", "transform") is not None + + def test_handler_module_can_use_dataclass_with_future_annotations(self, tmp_path): + (tmp_path / "handlers.py").write_text( + textwrap.dedent( + """ + from __future__ import annotations + + from dataclasses import dataclass + + from airflow.sdk import task_handler + + @dataclass + class Config: + factor: int + + @task_handler("d", "t") + def t(region): + return Config(factor=2) + """ + ) + ) + + import_handler_modules(_materialized("b", tmp_path)) + + assert get_handler("d", "t") is not None + + def test_airflowignore_skips_ignored_files(self, tmp_path): + _write_handler(tmp_path / "handlers.py", "d", "keep", params="a") + _write_handler(tmp_path / "scratch.py", "d", "skip", params="a") + (tmp_path / ".airflowignore").write_text("scratch.py\n") + + import_handler_modules(_materialized("b", tmp_path)) + + assert get_handler("d", "keep") is not None + assert get_handler("d", "skip") is None + + def test_idempotent_by_bundle_name(self, tmp_path): + # A re-import would re-run @task_handler and hit the duplicate-registration guard. + _write_handler(tmp_path / "handlers.py", "d", "t", params="a") + bundle = _materialized("b", tmp_path) + + import_handler_modules(bundle) + import_handler_modules(bundle) + + assert get_handler("d", "t") is not None + + def test_two_bundles_with_same_filename_do_not_collide(self, tmp_path): + b1, b2 = tmp_path / "b1", tmp_path / "b2" + b1.mkdir() + b2.mkdir() + _write_handler(b1 / "handlers.py", "dag1", "a", params="x") + _write_handler(b2 / "handlers.py", "dag2", "b", params="y") + + import_handler_modules(_materialized("one", b1)) + import_handler_modules(_materialized("two", b2)) + + assert get_handler("dag1", "a") is not None + assert get_handler("dag2", "b") is not None + + +class TestResolveHandlerBundle: + def test_materializes_and_pins_once(self, monkeypatch): + calls = {"pin": 0} + pinned = SimpleNamespace(name="b", version="v1", path="/does-not-matter") + + def _pin(target, logger): + calls["pin"] += 1 + return pinned + + monkeypatch.setattr("airflow.sdk.execution_time.bundles.initialize_pinned_bundle", _pin) + + assert resolve_handler_bundle("b") is pinned + assert resolve_handler_bundle("b") is pinned # memoized: materialized once per process + assert calls["pin"] == 1 + + def test_resolve_then_import_populates_registry(self, tmp_path, monkeypatch): + _write_handler(tmp_path / "handlers.py", "d", "t", params="a") + pinned = _materialized("b", tmp_path) + monkeypatch.setattr( + "airflow.sdk.execution_time.bundles.initialize_pinned_bundle", + lambda target, logger: pinned, + ) + + import_handler_modules(resolve_handler_bundle("b")) + + assert get_handler("d", "t") is not None diff --git a/task-sdk/tests/task_sdk/task_handlers/test_task_handlers.py b/task-sdk/tests/task_sdk/task_handlers/test_task_handlers.py new file mode 100644 index 0000000000000..b61cd5fd2b2da --- /dev/null +++ b/task-sdk/tests/task_sdk/task_handlers/test_task_handlers.py @@ -0,0 +1,132 @@ +# 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. + +"""Unit tests for airflow.sdk.task_handlers (registry, declaration, runtime binder).""" + +from __future__ import annotations + +import re + +import pytest + +from airflow.sdk import task_handler +from airflow.sdk.api.datamodels._generated import LiteralArgBinding, TaskArgBinding, XComArgBinding +from airflow.sdk.task_handlers import ( + TaskHandlerDeclaration, + call_handler, + get_handler, +) +from airflow.sdk.task_handlers.registry import _REGISTRY + + +@pytest.fixture(autouse=True) +def _clean_registry(): + _REGISTRY.clear() + yield + _REGISTRY.clear() + + +def _bound(bindings, resolver=lambda *_: None): + """Return the kwargs ``call_handler`` would pass to a handler (identity handler).""" + return call_handler(lambda **kwargs: kwargs, bindings, resolver) + + +def _lit(name, value, from_default=False): + root = LiteralArgBinding(name=name, kind="literal", value=value, from_default=from_default) + return TaskArgBinding(root=root) + + +def _xc(name, task_id): + return TaskArgBinding(root=XComArgBinding(name=name, kind="xcom", task_id=task_id)) + + +class TestRegistry: + def test_registers_and_looks_up_by_dag_and_task(self): + def transform(region, events): + return len(events) + + transform_h = task_handler("events_pipeline", "transform")(transform) + assert transform is transform_h, "@task_handler should return the original function" + assert get_handler("events_pipeline", "transform") is transform + + def test_task_id_defaults_to_function_name(self): + @task_handler("d") + def extract(): + return 1 + + assert get_handler("d", "extract") is extract + + def test_unknown_lookup_returns_none(self): + assert get_handler("nope", "nope") is None + + def test_duplicate_registration_rejected(self): + @task_handler("d", "t") + def a(): ... + + with pytest.raises( + ValueError, match=re.escape("handler already registered for dag_id='d' task_id='t'") + ): + + @task_handler("d", "t") + def b(): ... + + +class TestDeclaration: + def test_params_ordered_with_required_flag(self): + def transform(region: str, events: list, limit: int = 10): ... + + decl = TaskHandlerDeclaration.from_handler(transform) + assert [p.name for p in decl.params] == ["region", "events", "limit"] + assert [p.required for p in decl.params] == [True, True, False] + + def test_varargs_rejected(self): + def f(*args): ... + + with pytest.raises(ValueError, match=re.escape("may not use variadic *args or **kwargs")): + TaskHandlerDeclaration.from_handler(f) + + def test_positional_only_rejected(self): + def f(region, /): ... + + with pytest.raises(ValueError, match=re.escape("may not use positional-only parameter 'region'")): + TaskHandlerDeclaration.from_handler(f) + + def test_context_key_rejected(self): + def f(ti, region: str): ... + + with pytest.raises(ValueError, match=re.escape("argument 'ti' is a reserved context key")): + TaskHandlerDeclaration.from_handler(f) + + +class TestBinding: + def test_literal_and_xcom_bind_to_kwargs(self): + pulled = [] + + def resolver(task_id, key): + pulled.append((task_id, key)) + return [{"row": 1}] + + kwargs = _bound([_lit("region", "eu"), _xc("events", "extract")], resolver) + assert kwargs == {"region": "eu", "events": [{"row": 1}]} + assert pulled == [("extract", "return_value")] + + def test_from_default_literal_is_skipped(self): + bindings = [_lit("region", "eu"), _lit("limit", 10, from_default=True)] + assert _bound(bindings) == {"region": "eu"} + + def test_none_bindings_is_empty(self): + assert _bound(None) == {}