jason810496 commented on code in PR #75405: URL: https://github.com/apache/airflow/pull/75405#discussion_r4439031438
########## 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}") Review Comment: Verified with `dags reserialize`: the Dag bundle's parse imports handlers, then handlers.py is re-imported as a Dag file and this raises. Make re-registering the same function a no-op or maybe list out as limitation in the user facing doc? ########## task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py: ########## @@ -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: Review Comment: `**kwargs` and positional-only handlers raise ValueError here and get swallowed, so the documented restrictions never surface at parse time. Only availability failures should be skipped. ```suggestion except ValueError as e: raise YamlDagBuildError(str(e)) from e # Malformed handler, not availability. except Exception: ``` ########## 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 <apache-airflow:administration-and-deployment/dag-bundles>`. +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. Review Comment: Verified locally: A named handler bundle never reaches `sys.path`, unlike the stub's own bundle via BundleDagBag, so `from common import helper` inside a shared handler bundle fails. Maybe no need to fix in this PR, but track it as follow-up. ########## 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] + + [email protected](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() Review Comment: nit: Duplicated fixtures from `task-sdk/tests/task_sdk/execution_time/test_python_coordinator.py`. Would it be better to place in cross module fixtures? ########## 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", +] Review Comment: IIUC, only the decorator needs exporting. ```suggestion __all__ = ["task_handler"] ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
