Implement Python task handler and support it in YAML Dags by uranusjr · Pull Request #75405 · apache/airflow

uranusjr · 2026-10-22T12:51:35Z

This implements a Python equivalent to Language SDK’s ADR-0011 interface as a @task_handler(...) decorator function. The general design matches exactly to the ADR; a task handler is registered (with the decorator) under a dag_id, task_id key, and a coordinator is specified to find the handler.

This information is used twice, once at parse-time to do a best-effort signature validation (as specified in ADR-0011), and again on execution through the same AIP-108 interface that routes a Python stub to non-Python task handlers.

While the feature as it stands now solely serves the YAML Dag format, it also is not specific to it. We now have the Python interface to allow Dags in any language to create a stub to resolve into a Python function at runtime—although no Dag formats currently have such a stub interface (except Python, but you don’t need this to call Python from Python in the first place).

Close #75273.

This is the Python equivalent to the Language SDK ADR-0007 interface. The new interface allows a Dag author to write a Dag in YAML (other languages will also automatically work as long as they implement the stub interface), and wire up a PythonCoordinator in the configuration plus the correct queue mapping to have it execute a Python function at runtime.

This implements the Language SDK ADR-0011 interface between Python task handlers and the YAML Dag parser. This allows the parser to use the interface to validate 'run:' arguments against a matched task handler if it exists. As specified in ADR-0011, a missing task handler at parse-time does not fail parsing; the error is only surfaced at execution time.

kaxil · 2026-10-22T16:02:33Z

+        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]),


_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.

kaxil · 2026-10-22T16:02:33Z

-    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):


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.

kaxil · 2026-10-22T16:02:33Z

+        from airflow.sdk.task_handlers import call_handler

-    return result
+        coordinator = get_coordinator_manager().for_queue(ti.task.queue)


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.

kaxil · 2026-10-22T16:02:34Z


            try:
-                result = _execute_task(context=context, ti=ti, log=log)
+                result = _TaskRunner.run(ti, context)


_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.

kaxil · 2026-10-22T16:02:34Z

+    # 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)


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?

kaxil · 2026-10-22T16:02:34Z

+    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


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.

kaxil · 2026-10-22T16:02:34Z

+    assert getattr(result.dags[0].get_task("t"), "is_stub", False) is True
+
+
+def test_matching_arguments_pass(mock_bundle, monkeypatch):


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.

kaxil · 2026-10-22T16:02:34Z

+        _run(ti)
+
+
+def test_run_dispatches_is_stub_to_handler(monkeypatch):


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.

kaxil · 2026-10-22T16:02:35Z

+        load_handler_bundle(None)
+        assert get_handler("x", "y") is None
+
+    def test_idempotent_second_call_does_not_reimport(self, tmp_path, monkeypatch):


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.

kaxil · 2026-10-22T16:02:35Z

+# specific language governing permissions and limitations
+# under the License.
+
+# Licensed to the ASF under Apache-2.0.


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.

jason810496

Thanks for the filling up the gap.

jason810496 · 2026-10-23T02:07:04Z

@@ -448,7 +449,9 @@ implemented in a future release.
  arguments`_).
 * **Dynamic task mapping.** There is no equivalent of ``.expand()``.


We probably need to rename as "Mapped Tasks" as Ask called out in -- #71757 (comment)

jason810496 · 2026-10-23T02:12:09Z

+(``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.
+


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.

jason810496 · 2026-10-23T03:29:06Z


-    return result
+        coordinator = get_coordinator_manager().for_queue(ti.task.queue)
+        declaration = coordinator.describe_task_handler(ti.dag_id, ti.task_id)


Would it be better to make the describe_task_handler method PythonCoordinator specific instead of changing the base interface of coordinator or remove it entirely?

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.

airflow/task-sdk/src/airflow/sdk/coordinators/_subprocess.py

Lines 764 to 792 in 7ac92c9

def execute_task(

self,

*,

what: TaskInstance,

dag_rel_path: str | os.PathLike[str],

bundle_info: BundleInfo,

client: Client,

logger: FilteringBoundLogger | None = None,

sentry_integration: str = "",

subprocess_logs_to_stdout: bool,

**kwargs,

) -> BaseCoordinator.ExecutionResult:

"""

Run *what*.

A task of a native Dag runs its own Dag file, *dag_rel_path* in the Dag bundle of

*bundle_info* at the version of the run, and never another Dag file. Its runtime may still

read other files of the bundle, such as the JARs on a Java classpath. Any other task runs the

artifact found by scanning the bundle that holds the artifacts.

:raises TaskLaunchError: when a native Dag file cannot run, before the runtime starts.

"""

task_logger = logger or log

rel_path = os.fspath(dag_rel_path)

is_native = self._is_native_dag_file(rel_path, bundle_info.name)

if is_native:

bundle = _resolve_task_bundle(bundle_info, task_logger)

else:

bundle = self._resolve_artifact_bundle(bundle_info, task_logger)

IIUC, the necessary info for importing a handler are task_handler_bundle_path, dag_id and task_id, so it's about how we resolve the task_handler_bundle_path.
There're two other approaches in my mind:

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.

passing the task_handler_bundle_path that we resolved at supervisor as _AIRFLOW__TASK_HANDLER_BUNDLE_PATH env to the Python task runner.

jason810496 · 2026-10-23T03:47:46Z

+_loaded_bundles: set[str] = set()
+
+
+def load_handler_bundle(bundle_name: str | None) -> None:


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.

airflow/task-sdk/src/airflow/sdk/coordinators/_subprocess.py

Lines 789 to 806 in 7ac92c9

if is_native:

bundle = _resolve_task_bundle(bundle_info, task_logger)

else:

bundle = self._resolve_artifact_bundle(bundle_info, task_logger)

# Hold the version lock across start()/wait() so bundle cleanup cannot

# rmtree a version this task is still reading from, mirroring

# task_runner.main() for the Python task path.

with (

BundleVersionLock(bundle_name=bundle.name, bundle_version=bundle.version),

self._set_scan_roots([bundle.path]),

):

if is_native:

command, subprocess_schema_version = self._build_native_task_command(

what=what, bundle=bundle, rel_path=rel_path

)

else:

command, subprocess_schema_version = self._build_execute_task_command(what=what)

process = _PopenActivitySubprocess.start(

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.

airflow/task-sdk/src/airflow/sdk/execution_time/task_runner.py

Lines 2593 to 2597 in 7ac92c9

with BundleVersionLock(

bundle_name=ti.bundle_instance.name,

bundle_version=ti.bundle_instance.version,

):

state, _, error = run(ti, context, log)

jason810496 · 2026-10-23T04:00:26Z

+            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"


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 being binded to a argument in list type

jason810496 · 2026-10-23T04:02:24Z

+__all__ = ["ArgBinding", "call_handler"]
+
+
+class LiteralArgBinding(TypedDict):


Yes, we can reuse the _generated datamodels.

uranusjr added 4 commits October 22, 2026 12:39

Plumbing for Python task handler

ab12337

Document Python task handler support for YAML

b43c499

uranusjr added this to the Airflow 3.5.0 milestone Oct 22, 2026

uranusjr self-assigned this Oct 22, 2026

uranusjr added the area:task-sdk label Oct 22, 2026

uranusjr requested review from amoghrajesh, ashb and jason810496 as code owners October 22, 2026 12:51

uranusjr added the area:coordinator Coordinator: The interface to spawn Lang-SDK subprocesses label Oct 22, 2026

boring-cyborg Bot added the kind:documentation label Oct 22, 2026

uranusjr requested a review from kaxil October 22, 2026 14:45

kaxil reviewed Oct 22, 2026

View reviewed changes

jason810496 reviewed Oct 23, 2026

View reviewed changes

		assert getattr(result.dags[0].get_task("t"), "is_stub", False) is True


		def test_matching_arguments_pass(mock_bundle, monkeypatch):

		_run(ti)


		def test_run_dispatches_is_stub_to_handler(monkeypatch):

		@@ -448,7 +449,9 @@ implemented in a future release.
		arguments`_).
		* Dynamic task mapping. There is no equivalent of ``.expand()``.

	def execute_task(
	self,
	*,
	what: TaskInstance,
	dag_rel_path: str \| os.PathLike[str],
	bundle_info: BundleInfo,
	client: Client,
	logger: FilteringBoundLogger \| None = None,
	sentry_integration: str = "",
	subprocess_logs_to_stdout: bool,
	**kwargs,
	) -> BaseCoordinator.ExecutionResult:
	"""
	Run what.

	A task of a native Dag runs its own Dag file, dag_rel_path in the Dag bundle of
	bundle_info at the version of the run, and never another Dag file. Its runtime may still
	read other files of the bundle, such as the JARs on a Java classpath. Any other task runs the
	artifact found by scanning the bundle that holds the artifacts.

	:raises TaskLaunchError: when a native Dag file cannot run, before the runtime starts.
	"""
	task_logger = logger or log
	rel_path = os.fspath(dag_rel_path)
	is_native = self._is_native_dag_file(rel_path, bundle_info.name)
	if is_native:
	bundle = _resolve_task_bundle(bundle_info, task_logger)
	else:
	bundle = self._resolve_artifact_bundle(bundle_info, task_logger)

		_loaded_bundles: set[str] = set()


		def load_handler_bundle(bundle_name: str \| None) -> None:

	with BundleVersionLock(
	bundle_name=ti.bundle_instance.name,
	bundle_version=ti.bundle_instance.version,
	):
	state, _, error = run(ti, context, log)

		__all__ = ["ArgBinding", "call_handler"]


		class LiteralArgBinding(TypedDict):

Conversation

uranusjr commented Oct 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jason810496 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jason810496 Oct 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jason810496 Oct 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

uranusjr commented Oct 22, 2026 •

edited

Loading

jason810496 Oct 23, 2026 •

edited

Loading

jason810496 Oct 23, 2026 •

edited

Loading