Implement Python task handler and support it in YAML Dags

airflow-core/docs/authoring-and-scheduling/yaml-dags.rst

@@ -26,9 +26,9 @@ Writing Dags in YAML
 
 Airflow can build a Dag from a YAML document. One YAML document describes one
 Dag, and the YAML Dag importer turns it into the same Dag structure that Airflow
-runs like any other source. The tasks can be ordinary Airflow operators,
-or a pointer to a task handler implemented in a Language SDK. (Support for the
-Python Task SDK is planned for a future release.)
+runs like any other source. The tasks can be ordinary Airflow operators, or a
+pointer to a task handler implemented in any language with a Task SDK, including
+Python, Java, and more.
 
 A document is declarative. Everything in it is a value, and the shape of each
 value decides what it means.
@@ -423,11 +423,12 @@ which holds the code's arguments directly:
         source: {$x: extract}
         factor: 2
 
-A ``run:`` task is executed by a Task SDK runtime, so it must be routed to a
-coordinator. In practice, a ``run:`` task always needs a task-level ``queue:``
-mapped to a coordinator under ``[sdk] queue_to_coordinator``. See
-:doc:`language-sdks/index` for writing the handler and configuring the
-coordinator.
+A ``run:`` task must be accompanied by a task handler implemented by a Task SDK
+that will be executed at runtime. How the handler is implemented is not named by
+the YAML file; it follows the task-level ``queue:``, and then is mapped to the
+appropriate coordinator through ``[sdk] queue_to_coordinator``. For more on
+writing task handlers, see :doc:`task-sdk:task-handler` for Python,
+and :doc:`language-sdks/index` for other languages.
 
 Each argument in ``run:`` is either a JSON-serializable literal, or an upstream
 task's output, in the ``{$x: task}`` form, which also wires the dependency.
@@ -448,7 +449,9 @@ implemented in a future release.
   arguments`_).
 * **Dynamic task mapping.** There is no equivalent of ``.expand()``.
 * **Arguments that must be a Python object.** This particularly affects
-  callbacks (which need a Python function).
+  callbacks (which need a Python function). Note that using ``PythonOperator``
+  is generally not needed in a YAML Dag; implement a Python task handler to use
+  with a ``run:`` task instead (see `Code tasks`_).
 
 A rejected key is reported as an error rather than silently ignored, as is any
 other argument the target operator does not accept.

task-sdk/docs/index.rst

@@ -180,6 +180,7 @@ For the full public API reference, see the :doc:`api` page.
   resumable-job-mixin
   api
   concepts
+  task-handler
   lang-sdk
   executable-bundle-spec
   ts-bundle-spec

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.
+* A stub's arguments must satisfy cross-language requirements. An argument must
+  be either JSON-compatible, or reference an upstream task's return value XCom.
+* When a handler cannot be resolved at parse time (no bundle configured yet),
+  argument validation is deferred; a genuinely missing handler is reported when
+  the task runs.

task-sdk/src/airflow/sdk/__init__.py

@@ -120,6 +120,7 @@
     "setup",
     "task",
     "task_group",
+    "task_handler",
     "teardown",
 ]
 
@@ -233,6 +234,7 @@
     from airflow.sdk.execution_time import macros
     from airflow.sdk.execution_time.context import NEVER_EXPIRE
     from airflow.sdk.io.path import ObjectStoragePath
+    from airflow.sdk.task_handlers import task_handler
     from airflow.sdk.types import TaskInstance
 
     conf: AirflowSDKConfigParser
@@ -339,6 +341,7 @@
     "setup": ".definitions.decorators",
     "task": ".definitions.decorators",
     "task_group": ".definitions.decorators",
+    "task_handler": ".task_handlers",
     "teardown": ".definitions.decorators",
 }
 

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")
+
+
+@contextlib.contextmanager
+def _warm_shutdown_signals() -> Generator[None, None, None]:
+    """
+    Install SIGTERM/SIGINT warm-shutdown handlers during task supervision.
+
+    While supervising a task the supervisor must not be torn down by a
+    termination signal; instead it keeps running so the task can finish (or be
+    shut down gracefully) and its terminal state and logs are reported. The
+    handlers are installed around BOTH ``start()`` (which transitions the TI to
+    RUNNING) and ``wait()`` (which runs the task and then reports the terminal
+    state / uploads logs), so there is no window where Python's default SIGTERM
+    disposition could kill the supervisor and tear the just-started task down
+    with it.
+
+    The previous dispositions are restored on exit so a long-lived supervisor
+    process (e.g. a reused Celery prefork worker) does not leak the handler into
+    later tasks or clobber the worker's own signal handling.
+    """
+
+    def _warm_shutdown(signum, frame):
+        log.info(
+            "Received signal; warm shutdown in progress, waiting for the running task to complete.",
+            signal=signal.Signals(signum).name,
+            pid=os.getpid(),
+        )
+
+    prev_sigterm = signal.getsignal(signal.SIGTERM)
+    prev_sigint = signal.getsignal(signal.SIGINT)
+    signal.signal(signal.SIGTERM, _warm_shutdown)
+    signal.signal(signal.SIGINT, _warm_shutdown)
+    try:
+        yield
+    finally:
+        signal.signal(signal.SIGTERM, prev_sigterm)
+        signal.signal(signal.SIGINT, prev_sigint)
+
+
+@attrs.define(kw_only=True)
+class PythonCoordinator(BaseCoordinator):
+    """
+    Coordinator that executes Python tasks 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)
+        handler = get_handler(dag_id, task_id)
+        return None if handler is None else declare_handler(task_id, handler)