jason810496 commented on code in PR #75405: URL: https://github.com/apache/airflow/pull/75405#discussion_r4355418620
########## airflow-core/docs/authoring-and-scheduling/yaml-dags.rst: ########## @@ -448,7 +449,9 @@ implemented in a future release. arguments`_). * **Dynamic task mapping.** There is no equivalent of ``.expand()``. Review Comment: We probably need to rename as "Mapped Tasks" as Ask called out in -- https://github.com/apache/airflow/pull/71757#discussion_r4347923371 ########## task-sdk/docs/task-handler.rst: ########## @@ -0,0 +1,97 @@ + .. 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. + Review Comment: How about mentioning the minimal YAML Dag that matches the handler here to to show the proper usage? Since the code example we show in YAML Dag section is using go-worker and the task signature doesn't match as well. ########## task-sdk/src/airflow/sdk/execution_time/task_runner.py: ########## @@ -2355,52 +2356,102 @@ 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): + """Run the task specified by *ti* with *context*.""" + if ti.task.is_stub: + return cls._run_handler(ti) + return cls._run_operator(ti, context) + + @classmethod + def _run_operator(cls, ti: RuntimeTaskInstance, context: Context): + """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): + """Run a stub task's registered Python task handler.""" + 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.task.queue) + declaration = coordinator.describe_task_handler(ti.dag_id, ti.task_id) Review Comment: Would it be better to make the `describe_task_handler` method `PythonCoordinator` specific instead of changing the base interface of coordinator? For the non-python SDK, we actually encapsulate the logic of handling execution path for the task handler and for the Dag definition within the `execute_task` method. https://github.com/apache/airflow/blob/7ac92c98ce63201de7b4bf9a59f0364088e7c447/task-sdk/src/airflow/sdk/coordinators/_subprocess.py#L764-L792 IIUC, the necessary info for importing a handler are `task_handler_bundle_path`, dag_id and task_id. There're two other approaches in my mind: 1) define a public attribute for accessing `task_handler_bundle_path` -- this means the child need to reparse the config and init the coordinator again just for resolving the `task_handler_bundle_path`, but change less. 2) passing the `task_handler_bundle_path` that we resolved at supervisor as `_AIRFLOW__TASK_HANDLER_BUNDLE_PATH` env to the Python task runner. ########## task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py: ########## @@ -105,11 +105,46 @@ 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, e.g. no handler + bundle or an unregistered handler. This does not fail parsing, so + building a Dag never depends on handler availability. Errors may happen + at execution time if the signatures are incompatible. + """ + 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) + if (declaration := coordinator.describe_task_handler(doc.dag_id, task.id_)) is None: + continue # Best-effort: If nothing can be matched, allow the stub to exist. + 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" Review Comment: Do we need to validate against the type like https://github.com/apache/airflow/pull/74067/changes#diff-a62ed98ec4693a9a9f43c7807c68805f64c92e233f808ce26121dcb70d03b81cR184 ? e.g. reject a int literal binding to a argument in list type ########## task-sdk/src/airflow/sdk/task_handlers/loader.py: ########## @@ -0,0 +1,82 @@ +# 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 a separate module. The Dag bundle +that contains a handler is resolved by name, and its Python files are imported +so the decorators run and populate the in-memory handler registry. + +The same code path is used by the parse path for validation, and the execution +path to run the handler. + +Import is per-process and once per bundle. A redeployed bundle is therefore only +picked up on the next fresh processing loop. +""" + +from __future__ import annotations + +import importlib.util +import pathlib +import sys + +import structlog + +from airflow.sdk._shared.module_loading.dag_file import get_unique_dag_module_name + +log = structlog.get_logger(logger_name="task_handlers.loader") + +_loaded_bundles: set[str] = set() + + +def load_handler_bundle(bundle_name: str | None) -> None: Review Comment: How about moving the initialization and for the task handler bundle with lock to the coordinator.execute_task? Just like the non-python coordinator, we lock the bundle before launch the task subprocess. https://github.com/apache/airflow/blob/7ac92c98ce63201de7b4bf9a59f0364088e7c447/task-sdk/src/airflow/sdk/coordinators/_subprocess.py#L789-L806 I know for typical Python task, we lock the bundle in the task_runner level. So another direction might be lock the task handler bundle at the similar place before the task run. https://github.com/apache/airflow/blob/7ac92c98ce63201de7b4bf9a59f0364088e7c447/task-sdk/src/airflow/sdk/execution_time/task_runner.py#L2593-L2597 ########## task-sdk/src/airflow/sdk/task_handlers/binding.py: ########## @@ -0,0 +1,104 @@ +# 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, TypedDict + +from airflow.sdk.bases.xcom import BaseXCom + +if TYPE_CHECKING: + from collections.abc import Callable, Iterator + from typing import Literal, NotRequired, TypeAlias + + XComResolver = Callable[[str, str], Any] + +__all__ = ["ArgBinding", "call_handler"] + + +class LiteralArgBinding(TypedDict): Review Comment: Yes, we can reuse the _generated datamodels. -- 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]
