Add ExecutableCoordinator for native self-contained Dag bundles #67161

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft

jason810496 wants to merge 137 commits into apache:main from jason810496:refactor/go-sdk/executable-coordinator

.pre-commit-config.yaml

-Original file line number
+Diff line change
@@ Expand Up / @@ -1147,6 +1147,20 @@ repos: @@
             pass_filenames: true
             files: ^airflow-core/src/airflow/api_fastapi/execution_api/(datamodels|versions)/.*\.py$
             require_serial: true
+          - id: generate-supervisor-schemas-snapshot
+            name: Regenerate supervisor schema snapshot
+            entry: ./scripts/ci/prek/generate_supervisor_schemas_snapshot.py
+            language: python
+            pass_filenames: false
+            files: ^(task-sdk/src/airflow/sdk/execution_time/(comms\.py|schema/.*\.py)|airflow-core/src/airflow/dag_processing/processor\.py)$
+            require_serial: true
+          - id: check-supervisor-schemas-versions
+            name: Check supervisor schema changes have corresponding version updates
+            entry: ./scripts/ci/prek/check_supervisor_schemas_versions.py
+            language: python
+            pass_filenames: true
+            files: ^(task-sdk/src/airflow/sdk/execution_time/(comms\.py|schema/.*\.py)|airflow-core/src/airflow/dag_processing/processor\.py)$
+            require_serial: true
           - id: generate-tasksdk-datamodels
             name: Generate Datamodels for TaskSDK client
             language: python
@@ Expand Down @@

airflow-core/src/airflow/config_templates/config.yml

-Original file line number
+Diff line change
@@ Expand Up / @@ -2004,6 +2004,43 @@ workers: @@
           type: integer
           example: ~
           default: "60"
+    sdk:
+      description: Settings for non-Python SDK runtime coordination
+      options:
+        coordinators:
+          description: |
+            JSON object mapping of coordinator keys to coordinator definitions.
+            Each value is an object with ``classpath`` and optional ``kwargs``.
+            ``classpath`` is resolved via ``import_string`` and constructed with
+            ``kwargs`` on first use.  Entries are
+            independent instances, so the same ``classpath`` can be configured
+            multiple times under different names with different ``kwargs`` (for
+            example, two ``JavaCoordinator`` instances pinned to different JDK
+            versions).
+          version_added: 3.3.0
+          type: string
+          example: |
+            {
+              "jdk-17": {
+                "classpath": "airflow.sdk.coordinators.java.JavaCoordinator",
+                "kwargs": {"java_executable": "/usr/lib/jvm/java-17-openjdk/bin/java", "jvm_args": ["-Xmx1024m"]}
+              }
+            }
+          default: ~
+        queue_to_coordinator:
+          description: |
+            JSON mapping of queue names to a coordinator key from
+            ``[sdk] coordinators``.
+            When a task's ``language`` field is not set, this mapping is checked
+            to route the task to a configured coordinator instance based on its
+            queue.  This is useful when queues are used as environment or
+            isolation identifiers (e.g. ``legacy-java``, ``modern-java``).
+          version_added: 3.3.0
+          type: string
+          example: '{"legacy-java": "jdk-11", "modern-java": "jdk-17"}'
+          default: ~
     api_auth:
       description: Settings relating to authentication on the Airflow APIs
       options:
@@ Expand Down @@

airflow-core/src/airflow/executors/base_executor.py

            
                      Original file line number
                      Diff line number
                      Diff line change
                  
    @@ -652,11 +652,10 @@ def run_workload(
  
            if isinstance(workload, ExecuteTask):

                from airflow.sdk.execution_time.supervisor import supervise_task

                from airflow.sdk.execution_time.workloads.task import TaskInstanceDTO as SDKTaskInstanceDTO

                # workload.ti is a TaskInstanceDTO which duck-types as TaskInstance.

                # TODO: Create a protocol for this.

                return supervise_task(

                    ti=workload.ti,  # type: ignore[arg-type]

                    ti=SDKTaskInstanceDTO.model_validate(workload.ti, from_attributes=True),

                    bundle_info=workload.bundle_info,

                    dag_rel_path=workload.dag_rel_path,

                    token=workload.token,

airflow-core/src/airflow/executors/workloads/task.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -28,6 +28,8 @@ @@
     from airflow.utils.state import TaskInstanceState
     if TYPE_CHECKING:
+        import uuid
         from airflow.api_fastapi.auth.tokens import JWTGenerator
         from airflow.models.taskinstance import TaskInstance as TIModel
         from airflow.models.taskinstancekey import TaskInstanceKey
@@ Expand Down @@

devel-common/src/tests_common/pytest_plugin.py

            
                      Original file line number
                      Diff line number
                      Diff line change
                  
    @@ -2519,7 +2519,6 @@ def execute(self, context):
  
        from uuid6 import uuid7

        from airflow.sdk import DAG

        from airflow.sdk.api.datamodels._generated import TaskInstance

        from airflow.sdk.execution_time.comms import BundleInfo, StartupDetails

        from airflow.timetables.base import TimeRestriction

    @@ -2547,6 +2546,15 @@ def _create_task_instance(
  
            should_retry: bool | None = None,

            max_tries: int | None = None,

        ) -> RuntimeTaskInstance:

            from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS

            if AIRFLOW_V_3_3_PLUS:

                from airflow.sdk.execution_time.workloads.task import TaskInstanceDTO

            else:

                from airflow.sdk.api.datamodels._generated import (  # type: ignore[no-redef,assignment]

                    TaskInstance as TaskInstanceDTO,

                )

            from airflow.sdk.api.datamodels._generated import DagRun, DagRunState, TIRunContext

            from airflow.utils.types import DagRunType

    @@ -2624,14 +2632,17 @@ def _create_task_instance(
  
            }

            startup_details = StartupDetails(

                ti=TaskInstance(

                ti=TaskInstanceDTO(

                    id=ti_id,

                    task_id=task.task_id,

                    dag_id=dag_id,

                    run_id=run_id,

                    try_number=try_number,

                    map_index=map_index,

                    map_index=map_index,  # type: ignore[arg-type]

                    dag_version_id=uuid7(),

                    pool_slots=1,

                    queue="default",

                    priority_weight=1,

                ),

                dag_rel_path="",

                bundle_info=BundleInfo(name="anything", version="any"),

docs/spelling_wordlist.txt

-Original file line number
+Diff line change
@@ Expand Up / @@ -860,6 +860,7 @@ iTerm @@
     iterm
     itertools
     Jarek
+    JavaCoordinator
     javascript
     jaydebeapi
     Jdbc
@@ Expand Down Expand Up / @@ -897,6 +898,7 @@ jsonl @@
     juli
     Jupyter
     jupyter
+    jvm
     jwks
     JWT
     jwt
@@ Expand Down Expand Up / @@ -1130,6 +1132,7 @@ openai @@
     openapi
     openfaas
     OpenID
+    openjdk
     openlineage
     OpenSearch
     opensearch
@@ Expand Down Expand Up / @@ -1860,6 +1863,7 @@ XComs @@
     Xiaodong
     xlarge
     xml
+    Xmx
     xpath
     XSS
     xyz
@@ Expand Down @@

providers/standard/src/airflow/providers/standard/decorators/stub.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -85,7 +85,6 @@ def stub( @@
         Stub tasks exist in the Dag graph only, but the execution must happen in an external
         environment via the Task Execution Interface.
         """
         return task_decorator_factory(
             decorated_operator_class=_StubOperator,
@@ Expand Down @@

scripts/ci/prek/check_supervisor_schemas_versions.py

-Original file line number
+Diff line change
@@ -0,0 +1,204 @@
+    #!/usr/bin/env python
+    # 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.
+    # /// script
+    # requires-python = ">=3.10,<3.11"
+    # dependencies = [
+    #   "rich>=13.6.0",
+    # ]
+    # ///
+    """
+    Fail when a supervisor schema has changed without a matching
+    ``VersionChange`` entry under
+    ``task-sdk/src/airflow/sdk/execution_time/schema/versions/``.
+    Mirrors :mod:`scripts.ci.prek.check_execution_api_versions` for the
+    supervisor bundle. The check is per-commit: every PR that mutates a
+    registered supervisor schema must add an instruction to the in-progress head
+    ``v<YYYY>_<MM>_<DD>.py`` file. The release-time version-file bump itself
+    is one-per-release; this hook is what keeps the in-progress file
+    honest between releases.
+    The comparison is done by dumping the snapshot JSON in this worktree
+    and in a temporary worktree of the upstream target branch, then
+    diffing them. Both sides invoke the sibling ``dump_supervisor_schemas.py``
+    script so the comparison is dump-version stable.
+    """
+    from __future__ import annotations
+    import os
+    import subprocess
+    import sys
+    import tempfile
+    from pathlib import Path
+    from common_prek_utils import console, get_remote_for_main
+    SUPERVISOR_SCHEMAS_PREFIX = "task-sdk/src/airflow/sdk/execution_time/schema/"
+    VERSIONS_PREFIX = SUPERVISOR_SCHEMAS_PREFIX + "versions/"
+    TASK_SDK_COMMS_PATH = "task-sdk/src/airflow/sdk/execution_time/comms.py"
+    CORE_PROCESSOR_PATH = "airflow-core/src/airflow/dag_processing/processor.py"
+    DUMP_SCRIPT = Path(__file__).parent / "dump_supervisor_schemas.py"
+    # TODO: We should consolidte the common logic with check_execution_api_versions.py into common_prek_utils
+    def get_target_branch() -> str:
+        """Branch to compare against. GITHUB_BASE_REF for PRs, DEFAULT_BRANCH in CI, else main."""
+        return os.environ.get("GITHUB_BASE_REF") or os.environ.get("DEFAULT_BRANCH") or "main"
+    def get_changed_files(filenames: list[str]) -> list[str]:
+        """Get changed files. Uses filenames from prek when provided, else staged files for local runs."""
+        if filenames:
+            return filenames
+        result = subprocess.run(
+            ["git", "diff", "--cached", "--name-only"],
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+        return [f for f in result.stdout.strip().splitlines() if f]
+    def dump_snapshot(cwd: Path) -> str:
+        """Run ``dump_supervisor_schemas.py`` in *cwd* and return its stdout."""
+        result = subprocess.run(
+            [
+                "uv",
+                "run",
+                "-p",
+                "3.12",
+                "--no-progress",
+                "--project",
+                "task-sdk",
+                "-s",
+                str(DUMP_SCRIPT),
+            ],
+            cwd=cwd,
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+        if result.returncode != 0:
+            raise RuntimeError(f"Snapshot dump failed: {result.stderr}")
+        return result.stdout
+    def _upstream_has_schema() -> bool:
+        """Return True if the target branch carries the schema package."""
+        target_branch = get_target_branch()
+        remote = get_remote_for_main()
+        ref = f"{remote}/{target_branch}"
+        subprocess.run(["git", "fetch", remote, target_branch], capture_output=True, check=False)
+        # ``git cat-file -e`` exits zero iff the path exists at the ref.
+        result = subprocess.run(
+            ["git", "cat-file", "-e", f"{ref}:{VERSIONS_PREFIX}__init__.py"],
+            capture_output=True,
+            check=False,
+        )
+        return result.returncode == 0
+    def dump_snapshot_from_main() -> str:
+        """Dump snapshot from target branch using a temporary worktree."""
+        target_branch = get_target_branch()
+        remote = get_remote_for_main()
+        ref = f"{remote}/{target_branch}"
+        worktree_path = Path(tempfile.mkdtemp()) / "airflow-main"
+        subprocess.run(["git", "fetch", remote, target_branch], capture_output=True, check=False)
+        subprocess.run(["git", "worktree", "add", str(worktree_path), ref], capture_output=True, check=True)
+        try:
+            return dump_snapshot(worktree_path)
+        finally:
+            subprocess.run(
+                ["git", "worktree", "remove", "--force", str(worktree_path)],
+                capture_output=True,
+                check=False,
+            )
+    def main() -> int:
+        changed_files = get_changed_files(sys.argv[1:])
+        # Files under schema/ that reference the bundle's
+        # registered models. Schema changes in those models' homes
+        # (``comms.py``, ``processor.py``) trigger this hook too because
+        # the snapshot embeds their head shape.
+        schema_source_files = [
+            f
+            for f in changed_files
+            if f.startswith(SUPERVISOR_SCHEMAS_PREFIX) or f == TASK_SDK_COMMS_PATH or f == CORE_PROCESSOR_PATH
+        ]
+        version_files = [f for f in changed_files if f.startswith(VERSIONS_PREFIX)]
+        if not schema_source_files:
+            return 0
+        if version_files:
+            # Contributor added a version-change entry: trust them.
+            return 0
+        if not _upstream_has_schema():
+            # The package is being introduced in this PR -- nothing on the
+            # target branch to compare against. The check will start firing
+            # normally once the package is on the target branch.
+            console.print(
+                "[yellow]Skipping supervisor-schemas version check:[/] target branch "
+                "has no schema package yet. The check activates once "
+                "this PR merges."
+            )
+            return 0
+        try:
+            main_snapshot = dump_snapshot_from_main()
+        except Exception as e:
+            console.print(f"[bold red]ERROR:[/] Failed to generate upstream snapshot for comparison: {e}")
+            return 1
+        try:
+            current_snapshot = dump_snapshot(Path.cwd())
+        except Exception as e:
+            console.print(f"[bold red]ERROR:[/] Failed to generate current snapshot: {e}")
+            return 1
+        if current_snapshot != main_snapshot:
+            console.print("[bold red]ERROR:[/] Supervisor schema has changed but no version file was updated.")
+            console.print("")
+            console.print("The following files were changed:")
+            for f in schema_source_files:
+                console.print(f"  - [magenta]{f}[/]")
+            console.print("")
+            remote = get_remote_for_main()
+            target_branch = get_target_branch()
+            console.print(
+                f"Snapshot diff against [cyan]{remote}/{target_branch}[/] detected differences.\n"
+                "\n"
+                "Append a ``VersionChange`` subclass to the in-progress head "
+                "``v<YYYY>_<MM>_<DD>.py`` file under:\n"
+                f"  [cyan]{VERSIONS_PREFIX}[/]\n"
+                "\n"
+                "See [cyan]task-sdk/src/airflow/sdk/execution_time/schema/AGENTS.md[/]."
+            )
+            return 1
+        console.print("[green]Snapshot unchanged:[/] Source changes do not affect the supervisor schema.")
+        return 0
+    if __name__ == "__main__":
+        sys.exit(main())

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add ExecutableCoordinator for native self-contained Dag bundles #67161

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!

Add ExecutableCoordinator for native self-contained Dag bundles #67161

Are you sure you want to change the base?

Uh oh!

Add ExecutableCoordinator for native self-contained Dag bundles #67161

Uh oh!

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!