kaxil commented on code in PR #75405:
URL: https://github.com/apache/airflow/pull/75405#discussion_r4350621804
##########
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)
Review Comment:
The supervisor picked the coordinator with `for_queue(ti.queue or
"default")` (supervisor.py:3148), but this looks it up again with
`ti.task.queue`. Those differ whenever the TI's queue was changed after
parsing, e.g. by `task_instance_mutation_hook`, and then the handler is
resolved through a different coordinator (or the bundle-less fallback) than the
one running the task. Using `ti.queue or "default"` here keeps the two in step.
##########
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)
+ if declaration is None or declaration.handler 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."
+ )
+ handler = declaration.handler
+
+ log.info("Running Python task handler", dag_id=ti.dag_id,
task_id=ti.task_id)
+
+ bindings = ctx.arg_bindings or () if (ctx :=
ti._ti_context_from_server) else ()
+ return call_handler(
+ handler,
+ cast("list[ArgBinding]", [b.root.model_dump() for b in bindings]),
Review Comment:
`_bind_arguments` calls the resolver as `xcom_resolver(task_id,
"return_value")`, but the second positional parameter of
`RuntimeTaskInstance.xcom_pull` is `dag_id`, not `key`. So every `{$x:
upstream}` argument is pulled from a Dag called `"return_value"` and the
handler gets `None`, with no error. Passing `lambda task_id, key:
ti.xcom_pull(task_ids=task_id, key=key)` here (or having the binder call by
keyword) fixes it. The fake in `test_run_task_handler.py` defines
`xcom_pull(task_ids, key)`, which is why the tests pass;
`create_autospec(RuntimeTaskInstance.xcom_pull)` would have caught it.
##########
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:
+ """
+ Import the handler modules in the specified Dag bundle.
+
+ This function does nothing if *bundle_name* is falsy (no handler bundle
+ configured) or already loaded. Resolution goes through
``DagBundlesManager``
+ so the bundle inherits download/refresh/versioning like any other Dag
+ bundle. Imports are idempotent and at most once per-process, same as normal
+ Python imports.
+ """
+ if not bundle_name or bundle_name in _loaded_bundles:
+ return
+
+ # Make this lazy so the outer task_handlers package (which re-exports
+ # load_handler_bundle) imports cleanly without pulling in Core.
+ from airflow.dag_processing.bundles.manager import DagBundlesManager #
noqa: SDK002
+
+ bundle = DagBundlesManager().get_bundle(bundle_name)
+ bundle.initialize()
+ _import_python_files(pathlib.Path(bundle.path))
+ _loaded_bundles.add(bundle_name)
+ log.debug("Loaded task-handler bundle", bundle=bundle_name,
path=str(bundle.path))
+
+
+def _import_python_files(root: pathlib.Path) -> None:
+ """Import every .py under *root* so module-level decorators run."""
+ for path in sorted(root.rglob("*.py")):
+ if "__pycache__" in path.parts:
+ continue
+ if (module_name := get_unique_dag_module_name(str(path))) in
sys.modules:
+ continue
+ if (spec := importlib.util.spec_from_file_location(module_name, path))
is None or spec.loader is None:
+ continue
+ spec.loader.exec_module(module :=
importlib.util.module_from_spec(spec))
Review Comment:
The module goes into `sys.modules` only after `exec_module`, which breaks
any handler file that uses `@dataclass` with `from __future__ import
annotations` (`AttributeError: 'NoneType' object has no attribute '__dict__'`;
reproduced with this exact load order). The Python Dag importer registers the
module first (python_importer.py:264-266). Separately, the bundle root is never
put on `sys.path` and each file is loaded as an isolated `unusual_prefix_*`
module, so `from helpers import x` or a relative import inside a handler bundle
fails. Every `.py` in the bundle (tests, conftest, scripts) also gets executed,
regardless of `.airflowignore`. Would importing by dotted name with the bundle
root on `sys.path` work better here?
##########
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:
+ """
+ Import the handler modules in the specified Dag bundle.
+
+ This function does nothing if *bundle_name* is falsy (no handler bundle
+ configured) or already loaded. Resolution goes through
``DagBundlesManager``
+ so the bundle inherits download/refresh/versioning like any other Dag
+ bundle. Imports are idempotent and at most once per-process, same as normal
+ Python imports.
+ """
+ if not bundle_name or bundle_name in _loaded_bundles:
+ return
+
+ # Make this lazy so the outer task_handlers package (which re-exports
+ # load_handler_bundle) imports cleanly without pulling in Core.
+ from airflow.dag_processing.bundles.manager import DagBundlesManager #
noqa: SDK002
+
+ bundle = DagBundlesManager().get_bundle(bundle_name)
Review Comment:
`get_bundle(bundle_name)` has no version here, so this is the shared
tracking checkout, and for a `GitDagBundle` `initialize()` runs `refresh()`
(fetch, then hard reset). `_loaded_bundles` doesn't help much since both
callers are short-lived forks: each handler task and each parse of a YAML file
with a `run:` task does its own fetch. The bundle lock is dropped before
`_import_python_files` runs too, so a refresh from another task can reset the
tree mid-import. `SubprocessCoordinator` avoids this for the same
`task_handler_bundle_name` with `_initialize_pinned_bundle` plus
`BundleVersionLock`. Can the loader reuse that?
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -1666,7 +1667,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)
Review Comment:
`_execute_task(context=context, ti=ti, log=log)` used the `log` passed into
`run()`, but `_run_operator` now uses the module-level `log`, so the
`::endgroup::` marker, the hook runners and the `on_execute_callback` error no
longer go to the caller's logger. That's what fails six tests in the Task SDK
job (`test_run_emits_endgroup_before_execute`,
`test_task_runner_not_fail_on_failed_callback[...]`, and others). Threading
`log` through `_TaskRunner.run(ti, context, log)` restores the old contract.
The core CLI jobs fail separately because `test_dag_command.py:1302` still
patches `task_runner._execute_task`, and its `call_args_list[0].kwargs["ti"]`
assertion needs to read the positional arg now.
##########
task-sdk/src/airflow/sdk/task_handlers/declaration.py:
##########
@@ -0,0 +1,116 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+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 types
+import typing
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from inspect import Parameter, signature
+
+from pydantic import TypeAdapter
+from pydantic.json_schema import GenerateJsonSchema
+
+from airflow.sdk.definitions.context import KNOWN_CONTEXT_KEYS
+
+__all__ = ["TaskHandlerParam", "TaskHandlerDeclaration", "declare_handler",
"infer_value_schema"]
+
+
+@dataclass(frozen=True)
+class TaskHandlerParam:
+ """One declared handler parameter (ADR-0012). Bindings are positional;
name is carried too."""
+
+ name: str
+ value_schema: dict[str, typing.Any] | None = None
+ required: bool = True # the handler declares no default
+
+
+@dataclass
+class TaskHandlerDeclaration:
+ """What a handler declares for one task_id, to compare against the stub's
arg_bindings."""
+
+ task_id: str
+ params: list[TaskHandlerParam] = field(default_factory=list)
+ handler: Callable[..., typing.Any] | None = field(default=None,
compare=False, repr=False)
+
+
+def declare_handler(
+ task_id: str,
+ fn: Callable,
+ *,
+ schema_fn: Callable[[typing.Any], dict[str, typing.Any] | None] | None =
None,
+) -> TaskHandlerDeclaration:
+ """Introspect ``fn`` into a ``TaskHandlerDeclaration`` (ordered params +
optional schemas)."""
+ schema_fn = schema_fn or infer_value_schema
+ sig = signature(fn)
+
+ try:
+ hints = typing.get_type_hints(fn)
+ except Exception:
+ # Unresolvable annotations (forward refs to unavailable names) ->
treat as untyped.
+ hints = {}
+
+ params: list[TaskHandlerParam] = []
+ for name, p in sig.parameters.items():
+ if p.kind in (Parameter.VAR_POSITIONAL, Parameter.VAR_KEYWORD):
Review Comment:
Positional-only parameters (`def h(region, /)`) pass this check and the
builder's name comparison, but `call_handler` invokes `fn(**kwargs)`, so the
task always fails with a `TypeError` at run time. Rejecting
`Parameter.POSITIONAL_ONLY` here too would surface it at parse time. The
`TaskHandlerParam` docstring also says bindings are positional, while binding
is by name.
##########
task-sdk/src/airflow/sdk/task_handlers/declaration.py:
##########
@@ -0,0 +1,116 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+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 types
+import typing
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from inspect import Parameter, signature
+
+from pydantic import TypeAdapter
+from pydantic.json_schema import GenerateJsonSchema
+
+from airflow.sdk.definitions.context import KNOWN_CONTEXT_KEYS
+
+__all__ = ["TaskHandlerParam", "TaskHandlerDeclaration", "declare_handler",
"infer_value_schema"]
+
+
+@dataclass(frozen=True)
+class TaskHandlerParam:
+ """One declared handler parameter (ADR-0012). Bindings are positional;
name is carried too."""
+
+ name: str
+ value_schema: dict[str, typing.Any] | None = None
+ required: bool = True # the handler declares no default
+
+
+@dataclass
+class TaskHandlerDeclaration:
+ """What a handler declares for one task_id, to compare against the stub's
arg_bindings."""
+
+ task_id: str
+ params: list[TaskHandlerParam] = field(default_factory=list)
+ handler: Callable[..., typing.Any] | None = field(default=None,
compare=False, repr=False)
+
+
+def declare_handler(
+ task_id: str,
+ fn: Callable,
+ *,
+ schema_fn: Callable[[typing.Any], dict[str, typing.Any] | None] | None =
None,
+) -> TaskHandlerDeclaration:
+ """Introspect ``fn`` into a ``TaskHandlerDeclaration`` (ordered params +
optional schemas)."""
+ schema_fn = schema_fn or infer_value_schema
+ sig = signature(fn)
+
+ try:
+ hints = typing.get_type_hints(fn)
+ except Exception:
+ # Unresolvable annotations (forward refs to unavailable names) ->
treat as untyped.
+ hints = {}
+
+ params: list[TaskHandlerParam] = []
+ for name, p in sig.parameters.items():
+ if p.kind in (Parameter.VAR_POSITIONAL, Parameter.VAR_KEYWORD):
+ raise ValueError(f"task handler {task_id!r} may not use varadic
*args or **kwargs")
+ if name in KNOWN_CONTEXT_KEYS:
+ raise ValueError(f"task handler {task_id!r} argument {name!r} is a
reserved context key")
+ params.append(
+ TaskHandlerParam(
+ name=name,
+ value_schema=schema_fn(hints.get(name, p.annotation)),
+ required=p.default is Parameter.empty,
+ )
+ )
+ return TaskHandlerDeclaration(task_id=task_id, params=params, handler=fn)
+
+
+class _ValueSchemaGenerator(GenerateJsonSchema):
+ # OpenAPI fixed-width numeric formats, matching core so a foreign runtime
decodes
+ # into the right machine types.
+ def int_schema(self, schema):
+ return {**super().int_schema(schema), "format": "int64"}
+
+ def float_schema(self, schema):
+ return {**super().float_schema(schema), "format": "double"}
+
+
+def infer_value_schema(annotation: typing.Any) -> dict[str, typing.Any] | None:
+ """
+ Default ``value_schema`` generator: pydantic-based, mirroring core's
``_infer_value_schema``.
+
+ Returns ``None`` for unconstrained annotations, and also when pydantic is
unavailable, so
Review Comment:
This duplicates core's `_infer_value_schema`/`_ValueSchemaGenerator` in
`serialization/stub_arg_bindings.py`, minus its `@cache` and narrower excepts,
so the two will drift. The docstring says it returns None when pydantic is
unavailable, but pydantic is imported at module top and is a hard dependency.
Nothing reads `value_schema` yet either: the builder uses `name`/`required` and
`_run_handler` only `.handler`, so every task run builds TypeAdapters it then
discards.
##########
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.
+
+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")``.
+
+Deploying handlers
+------------------
+
+A handler must be importable in two places: the Dag processor (to validate a
+stub's arguments against the handler's signature at parse time) and the worker
+(to run it). The handler modules should be made available by placing them in a
+:doc:`Dag bundle <apache-airflow:administration-and-deployment/dag-bundles>`
and
+pointing a :class:`~airflow.sdk.coordinators.python.PythonCoordinator` at it.
+The coordinator should be configured in ``[sdk] coordinators`` with an
+appropriate ``task_handler_bundle_name``, and have the task's queue routed to
it
+with ``[sdk] queue_to_coordinator``:
+
+.. code-block:: ini
+
+ [dag_processor]
+ # A Dag bundle holding the @task_handler modules, registered like any
other bundle.
+ dag_bundle_config_list = [
+ {"name": "python-handlers", "classpath": "...", "kwargs": {...}},
+ ]
+
+ [sdk]
+ coordinators = {
+ "python": {
+ "classpath": "airflow.sdk.coordinators.python.PythonCoordinator",
+ "kwargs": {"task_handler_bundle_name": "python-handlers"}
+ }
+ }
+ queue_to_coordinator = {"python-workers": "python"}
+
+With this, a stub task whose ``queue`` is ``python-workers`` is routed to the
+``PythonCoordinator``, which imports ``python-handlers`` and runs the
registered
+handler. ``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
+-----
+
+* A handler bundle is imported once per worker process.
Review Comment:
This isn't what happens: the bundle is loaded in the task process, which is
forked per task, and in the per-file parse process, so it's imported again for
every handler task and every YAML parse (the loader docstring makes the same
assumption). It's also worth saying that every `.py` in the handler bundle gets
imported, and that a handler bundle in `dag_bundle_config_list` is parsed for
Dags too, so an `.airflowignore` tip like the one on the Go page would help.
##########
task-sdk/src/airflow/sdk/task_handlers/registry.py:
##########
@@ -0,0 +1,58 @@
+# 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: ...
+
+ :meta private:
Review Comment:
`task_handler` is in `airflow.sdk.__all__` and documented as the authoring
API, but `:meta private:` hides it from the reference, which is what fails
`test_docs_inventory_matches_public_api` in the Task SDK job. The same job
fails `test_airflow_sdk_no_unexpected_exports` because the lazy import leaves
`task_handlers` on the `airflow.sdk` namespace. Should
`airflow.sdk.task_handlers` export only the decorator too? Right now
`call_handler`, `declare_handler`, `load_handler_bundle` and friends are in its
`__all__`, which makes them compat surface.
##########
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:
These TypedDicts restate the generated `LiteralArgBinding`/`XComArgBinding`
models (minus `from_default`, which the `match` below reads), and
`_run_handler` dumps already-validated models to dicts and `cast`s them just to
get here. The missing name/kind/task_id branches can't fire on that path.
Matching on the generated model classes directly would drop the TypedDicts, the
cast and the dead branches.
##########
task-sdk/src/airflow/sdk/coordinators/python.py:
##########
@@ -0,0 +1,162 @@
+#
+# 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 Generator
+ from os import PathLike
+
+ from structlog.typing import FilteringBoundLogger
+
+ 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")
+
+
[email protected]
+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)
+
+
[email protected](kw_only=True)
+class PythonCoordinator(BaseCoordinator):
+ """
+ Coordinator that executes Python tasks in an in-process Task SDK
subprocess.
Review Comment:
Nit: "in-process Task SDK subprocess" contradicts itself, and line 101 has a
stray "in" ("configured explicitly in to carry").
##########
task-sdk/tests/task_sdk/task_handlers/test_loader.py:
##########
@@ -0,0 +1,118 @@
+# 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 declare_handler, get_handler,
load_handler_bundle, loader as loader_mod
+from airflow.sdk.task_handlers.registry import _REGISTRY
+
+
[email protected](autouse=True)
+def _clean():
+ _REGISTRY.clear()
+ loader_mod._loaded_bundles.clear()
+ stale = [name for name in sys.modules if
name.startswith(UNUSUAL_MODULE_PREFIX)]
+ for name in stale:
+ del sys.modules[name]
+ yield
+ _REGISTRY.clear()
+ loader_mod._loaded_bundles.clear()
+
+
+def _fake_bundle(monkeypatch, name, path):
+ from airflow.dag_processing.bundles import manager as mgr
+
+ bundle = SimpleNamespace(name=name, path=path, initialize=lambda: None)
+ monkeypatch.setattr(mgr.DagBundlesManager, "get_bundle", lambda self,
bundle_name, *a, **k: bundle)
+
+
+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
+ """
+ )
+ )
+
+
+class TestLoadHandlerBundle:
+ def test_imports_modules_and_populates_registry(self, tmp_path,
monkeypatch):
+ _write_handler(tmp_path / "handlers.py", "events", "transform")
+ _fake_bundle(monkeypatch, "b", tmp_path)
+
+ load_handler_bundle("b")
+
+ assert get_handler("events", "transform") is not None
+
+ def test_none_is_a_noop(self):
+ load_handler_bundle(None)
+ assert get_handler("x", "y") is None
+
+ def test_idempotent_second_call_does_not_reimport(self, tmp_path,
monkeypatch):
Review Comment:
This stays green with the `_loaded_bundles` guard removed, since
`_import_python_files` also skips names already in `sys.modules` (deterministic
per path). Asserting `get_bundle`/`initialize` is called once would pin the
guard. Also, `_clean` purges `unusual_prefix_*` modules on setup but not on
teardown, so this file's handler modules leak into later tests, and
`TestDeclarationCarriesHandler` below tests `declaration.py` and probably
belongs in test_task_handlers.py.
##########
task-sdk/tests/task_sdk/importers/yaml_importer/test_builder_validation.py:
##########
@@ -0,0 +1,91 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import textwrap
+from types import SimpleNamespace
+
+import pytest
+
+from airflow.sdk.importers.base import FilesystemDagDefinition
+from airflow.sdk.importers.yaml_importer.importer import YamlDagImporter
+from airflow.sdk.task_handlers import TaskHandlerDeclaration, TaskHandlerParam
+
+HEAD_SCHEMA = "https://airflow.apache.org/schemas/dag/2026-10-30.json";
+
+
[email protected]
+def mock_bundle(tmp_path):
+ bundle_dir = tmp_path / "bundle"
+ bundle_dir.mkdir(parents=True, exist_ok=True)
+ return SimpleNamespace(name="test_bundle", 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(describe_task_handler=lambda dag_id,
task_id: 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):
Review Comment:
Every declaration in this file has a single `required=True` param, so
dropping `if p.required` in the builder keeps all five tests green. A case with
`TaskHandlerParam(name="limit", required=False)` left out of `run:` would pin
it. `_patch_declaration`'s `for_queue=lambda queue: coordinator` ignores the
queue, so routing by the task's `queue:` isn't tested either, and
`test_none_declaration_skips` covers the same branch as
`test_no_declaration_available_skips_and_builds`.
##########
task-sdk/tests/task_sdk/task_handlers/test_task_handlers.py:
##########
@@ -0,0 +1,186 @@
+# 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.
+
+# Licensed to the ASF under Apache-2.0.
Review Comment:
Stray second license line under the full ASF header. Also in this file,
`test_call_handler_end_to_end` repeats `test_literal_and_xcom_bind_to_kwargs`,
and `test_xcom_missing_task_id_errors` loops over its inputs where
`@pytest.mark.parametrize` would report which one failed.
##########
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):
Review Comment:
This path never goes through `_run_execute_callable`, so `execution_timeout`
is ignored for handler tasks: no `timeout()` in the task process and no
`SetExecutionTimeout` to the supervisor. A YAML `run:` task accepts
`execution_timeout` at task level (it's in `_BASE_TASK_ARGS` and is passed to
`task_decorators.stub`), and `[core] default_task_execution_timeout` applies
too, so a hung handler holds the worker slot indefinitely. Could this return
`_run_execute_callable(context, lambda context: call_handler(...), ti.task)`?
That would also keep `on_kill` and the `task.execute` span consistent with
`_run_operator`.
##########
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:
Review Comment:
ADR-0011 (still Proposed) puts this check in `_parse_file` above every
importer ("No importer knows about coordinators"), treats a missing handler as
an import error, and compares params by name and order, and ADR-0012 defines
the coordinator hook as `parse_task_handler`. This does it in the YAML builder,
skips a missing handler, compares names only, and adds a separate
`describe_task_handler(dag_id, task_id)` hook. Is `describe_task_handler` meant
to replace `parse_task_handler` or sit next to it? Worth settling (or updating
the ADR) before `BaseCoordinator.describe_task_handler` ships as an override
point, since changing its signature later breaks out-of-tree coordinators.
##########
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:
Review Comment:
The docstring above says validation "does not fail parsing, so building a
Dag never depends on handler availability", but nothing here catches exceptions
from `describe_task_handler`. A git failure while initializing the handler
bundle, an ImportError in any file in it, or `declare_handler` rejecting
`**kwargs` all propagate out of `build()` and become an import error on the
YAML file, for every YAML Dag routed to that coordinator. Either catch
availability failures here and log them, or update this docstring and the note
in task-handler.rst to say it fails loudly.
##########
task-sdk/src/airflow/sdk/coordinators/python.py:
##########
@@ -0,0 +1,162 @@
+#
+# 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 Generator
+ from os import PathLike
+
+ from structlog.typing import FilteringBoundLogger
+
+ 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")
+
+
[email protected]
+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)
+
+
[email protected](kw_only=True)
+class PythonCoordinator(BaseCoordinator):
+ """
+ Coordinator that executes Python tasks in an in-process 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 to
+ carry ``task_handler_bundle_name`` via ``[sdk] coordinators`` for 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 describe_task_handler(self, dag_id: str, task_id: str) ->
TaskHandlerDeclaration | None:
+ """
+ Describe the registered handler for the specified task.
+
+ Since Python task handlers are registered in-process, they are simply
+ imported and introspected.
+
+ The returned declaration carries both the signature (for validation)
and
+ the resolved callable (for execution).
+ """
+ from airflow.sdk.task_handlers import declare_handler, get_handler,
load_handler_bundle
+
+ load_handler_bundle(self.task_handler_bundle_name)
Review Comment:
On `SubprocessCoordinator` an unset `task_handler_bundle_name` means "use
the task's own Dag bundle" (config.yml documents it that way too), but here
unset loads nothing. That includes the default fallback from
`_build_python_coordinator`, so a `run:` task on an unmapped queue, or with its
handler next to the YAML file, parses cleanly (validation silently skips) and
fails every run with "No task handler". Should unset fall back to the task's
bundle here as well? If the difference is intended, task-handler.rst should say
the kwarg is required for Python handlers.
##########
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.
+
+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
Review Comment:
Worth listing the restrictions `declare_handler` enforces: no
`*args`/`**kwargs`, and no parameter named after a context key (`ti`, `ds`,
...). Someone coming from `@task` will write `**context` and get a YAML import
error the docs never mentioned.
##########
task-sdk/tests/task_sdk/execution_time/test_run_task_handler.py:
##########
@@ -0,0 +1,127 @@
+# 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.execution_time.task_runner import _TaskRunner
+from airflow.sdk.task_handlers import TaskHandlerDeclaration
+
+
+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",
+ task=SimpleNamespace(queue=queue, is_stub=True),
+ _ti_context_from_server=SimpleNamespace(arg_bindings=arg_bindings),
+ xcom_pull=xcom_pull,
+ )
+ ti.pulled = pulled
+ return ti
+
+
+def _run(ti):
+ # A stub ti always routes through the public entry to the handler path;
context is unused there.
+ return _TaskRunner.run(ti, mock.MagicMock())
+
+
+def _patch_coordinator(monkeypatch, declaration):
+ coordinator = SimpleNamespace(describe_task_handler=lambda dag_id,
task_id: declaration)
+ 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,
TaskHandlerDeclaration(task_id="transform", handler=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, TaskHandlerDeclaration(task_id="extract",
handler=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_declaration_without_callable_raises(monkeypatch):
+ # e.g. a declaration that came from a foreign runtime over the wire
carries no callable.
+ _patch_coordinator(monkeypatch,
TaskHandlerDeclaration(task_id="transform", handler=None))
+ ti = _fake_ti(arg_bindings=None)
+
+ with pytest.raises(RuntimeError, match="No task handler"):
+ _run(ti)
+
+
+def test_run_dispatches_is_stub_to_handler(monkeypatch):
Review Comment:
This runs the same path as `test_binds_literal_and_xcom_then_calls_handler`
(every fake TI has `is_stub=True`), and nothing here pushes or checks an XCom
despite the comment. A test for the other branch, a non-stub TI going to
`_run_operator`, would add more.
--
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]
