Report a native Dag's task that reaches a Python worker

A Python worker handed a task from a Dag file that a coordinator's Dag
importer claims now exits with an error naming the queue to route. It
checks before building a Dag bag, so it never starts the runtime. As in
the Dag processor, a .py, .pyc or .zip file is never native.

Author: jason810496

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

@@ -60,6 +60,7 @@
 from airflow.sdk.bases.skipmixin import XCOM_SKIPMIXIN_KEY
 from airflow.sdk.bases.xcom import BaseXCom
 from airflow.sdk.configuration import conf
+from airflow.sdk.coordinators._dag_importer import CoordinatorDagImporter
 from airflow.sdk.definitions._internal.dag_parsing_context import _airflow_parsing_context_manager
 from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet, is_arg_set
 from airflow.sdk.definitions.asset import (
@@ -149,6 +150,7 @@
 from airflow.sdk.execution_time.sentry import Sentry
 from airflow.sdk.execution_time.tracing import detail_span
 from airflow.sdk.execution_time.xcom import XCom
+from airflow.sdk.importers import get_file_suffix, get_importer_registry
 from airflow.sdk.listener import get_listener_manager
 from airflow.sdk.observability.metrics import stats_utils
 from airflow.sdk.serde import allow_class, iter_pydantic_models
@@ -994,6 +996,22 @@ def _register_deserialization_allowed_classes(dag, log: Logger) -> None:
                     )
 
 
+# The Dag processor always parses these with the legacy importer, whatever importer is configured.
+_LEGACY_DAG_FILE_SUFFIXES = frozenset({".py", ".pyc", ".zip"})
+
+
+def _is_native_dag_file(path: str, bundle_name: str) -> bool:
+    """Return whether a coordinator's Dag importer claims *path*, so that a Lang-SDK runtime parses it."""
+    if get_file_suffix(path) in _LEGACY_DAG_FILE_SUFFIXES:
+        return False
+    try:
+        importer = get_importer_registry(bundle_name).get_importer(path)
+    except Exception:
+        # Building the Dag bag reports a broken importer configuration.
+        return False
+    return isinstance(importer, CoordinatorDagImporter)
+
+
 @detail_span("parse")
 def parse(what: StartupDetails, log: Logger) -> RuntimeTaskInstance:
     # TODO: Task-SDK:
@@ -1006,6 +1024,17 @@ def parse(what: StartupDetails, log: Logger) -> RuntimeTaskInstance:
     bundle_prepare_ms = int((time.monotonic() - bundle_prepare_start) * 1000)
 
     dag_absolute_path = os.fspath(Path(bundle_instance.path, what.dag_rel_path))
+    if _is_native_dag_file(dag_absolute_path, bundle_info.name):
+        log.error(
+            "A task of a native Lang-SDK Dag cannot run in Python. Route its queue to the coordinator "
+            "that parses the Dag, with [sdk] queue_to_coordinator",
+            dag_id=what.ti.dag_id,
+            task_id=what.ti.task_id,
+            queue=what.ti.queue,
+            path=what.dag_rel_path,
+        )
+        sys.exit(1)
+
     dag_file_parse_start = time.monotonic()
     bag = BundleDagBag(
         dag_folder=dag_absolute_path,

task-sdk/tests/task_sdk/execution_time/test_task_runner.py

@@ -31,6 +31,7 @@
 from unittest import mock
 from unittest.mock import call, patch
 
+import attrs
 import pandas as pd
 import pytest
 import structlog
@@ -78,6 +79,8 @@
 )
 from airflow.sdk.bases.operator import ExecutorSafeguard
 from airflow.sdk.bases.xcom import BaseXCom
+from airflow.sdk.coordinators._dag_importer import CoordinatorDagImporter
+from airflow.sdk.coordinators._subprocess import SubprocessCoordinator
 from airflow.sdk.definitions._internal.types import NOTSET, SET_DURING_EXECUTION, is_arg_set
 from airflow.sdk.definitions.asset import Asset, AssetAlias, AssetUniqueKey, AssetUriRef, Dataset, Model
 from airflow.sdk.definitions.param import DagParam
@@ -175,6 +178,7 @@
     TaskRunnerMarker,
     _defer_task,
     _execute_task,
+    _is_native_dag_file,
     _make_task_span,
     _push_xcom_if_needed,
     _register_deserialization_allowed_classes,
@@ -912,6 +916,94 @@ def test_parse_dag_from_bundle_importer(tmp_path: Path, make_ti_context):
     assert ti.task.dag.relative_fileloc == "listed.dags"
 
 
+class NativeDagImporter(CoordinatorDagImporter):
+    artifact_suffix = ".native"
+    supported_extensions = [".native"]
+
+    def get_source_code(self, definition):
+        return DagSourceCode(source_code=definition.read_text(), language="native")
+
+
+@attrs.define(kw_only=True)
+class NativeCoordinator(SubprocessCoordinator):
+    """A coordinator whose Dag importer claims ``.native`` files in every bundle."""
+
+    def get_dag_importer(self):
+        return NativeDagImporter(coordinator=self)
+
+
+@patch("airflow.dag_processing.dagbag.BundleDagBag", autospec=True)
+def test_parse_rejects_a_task_of_a_native_dag(mock_bag, tmp_path: Path, make_ti_context):
+    tmp_path.joinpath("dag.native").write_text("{}")
+    what = StartupDetails(
+        ti=TaskInstance(
+            id=uuid7(),
+            task_id="a",
+            dag_id="native_dag",
+            run_id="c",
+            try_number=1,
+            dag_version_id=uuid7(),
+            queue="default",
+        ),
+        dag_rel_path="dag.native",
+        bundle_info=BundleInfo(name="my-bundle", version=None),
+        ti_context=make_ti_context(),
+        start_date=timezone.utcnow(),
+        sentry_integration="",
+    )
+    bundle_config = [
+        {
+            "name": "my-bundle",
+            "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle",
+            "kwargs": {"path": str(tmp_path), "refresh_interval": 1},
+        }
+    ]
+    coordinators = {"native": {"classpath": f"{__name__}.NativeCoordinator", "kwargs": {}}}
+    log = mock.Mock()
+
+    reset_importer_registry()
+    try:
+        with (
+            patch.dict(
+                os.environ,
+                {
+                    "AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(bundle_config),
+                    "AIRFLOW__SDK__COORDINATORS": json.dumps(coordinators),
+                },
+            ),
+            pytest.raises(SystemExit, match="1"),
+        ):
+            parse(what, log)
+    finally:
+        reset_importer_registry()
+
+    mock_bag.assert_not_called()
+    log.error.assert_called_once_with(
+        "A task of a native Lang-SDK Dag cannot run in Python. Route its queue to the coordinator "
+        "that parses the Dag, with [sdk] queue_to_coordinator",
+        dag_id="native_dag",
+        task_id="a",
+        queue="default",
+        path="dag.native",
+    )
+
+
+@pytest.mark.parametrize(
+    ("file_name", "expected"),
+    [("dag.native", True), ("dag.py", False), ("dag.pyc", False), ("dags.zip", False)],
+)
+@patch.object(NativeDagImporter, "supported_extensions", [".native", ".py", ".pyc", ".zip"])
+def test_a_python_dag_file_is_never_native(file_name, expected):
+    coordinators = {"native": {"classpath": f"{__name__}.NativeCoordinator", "kwargs": {}}}
+
+    reset_importer_registry()
+    try:
+        with patch.dict(os.environ, {"AIRFLOW__SDK__COORDINATORS": json.dumps(coordinators)}):
+            assert _is_native_dag_file(f"/bundle/{file_name}", "my-bundle") is expected
+    finally:
+        reset_importer_registry()
+
+
 @pytest.mark.parametrize("use_queues", [False, True])
 def test_run_deferred_basic(time_machine, create_runtime_ti, mock_supervisor_comms, use_queues: bool):
     """Test that a task can transition to a deferred state."""