From 024958f76536a54987da8c528f9691a222185e12 Mon Sep 17 00:00:00 2001 From: Kaxil Naik Date: Thu, 11 Jun 2026 14:34:23 +0100 Subject: [PATCH 1/7] Version the worker-bound TaskInstance fields in the execution API schema The supervisor's task routing (added in #65958) reads queue, pool_slots and priority_weight from the TaskInstance it receives, but those fields lived in a TaskInstanceDTO duplicated between airflow-core and task-sdk and kept in sync by an AST-comparison prek hook, outside the versioned execution API schema the Task SDK datamodels are generated from. Add the three fields to the execution API TaskInstance schema (with a Cadwyn version change) and regenerate the SDK datamodels, so the generated TaskInstance carries everything the supervisor needs. Re-point StartupDetails and the coordinator interfaces at the generated model, fold the core TaskInstanceDTO into a subclass of the schema model that only adds the executor-side fields, and drop the duplicated task-sdk DTO and the sync hook. The unused parent_context_carrier field (no readers or writers) is removed. --- .pre-commit-config.yaml | 6 - .../docs/core-concepts/executor/index.rst | 7 +- .../execution_api/datamodels/taskinstance.py | 5 + .../execution_api/versions/__init__.py | 2 + .../execution_api/versions/v2026_06_30.py | 17 +- .../src/airflow/executors/workloads/task.py | 35 +-- .../tests/unit/executors/test_workloads.py | 33 +++ .../src/tests_common/pytest_plugin.py | 16 +- .../edge3/worker_api/v2-edge-generated.yaml | 52 ++-- .../ci/prek/check_task_instance_dto_sync.py | 125 --------- .../airflow/sdk/api/datamodels/_generated.py | 3 + .../airflow/sdk/coordinators/_subprocess.py | 9 +- .../coordinators/executable/coordinator.py | 4 +- .../sdk/coordinators/java/coordinator.py | 4 +- .../src/airflow/sdk/execution_time/comms.py | 7 +- .../airflow/sdk/execution_time/coordinator.py | 6 +- .../sdk/execution_time/schema/schema.json | 246 +++++++----------- .../schema/versions/__init__.py | 3 + .../airflow/sdk/execution_time/supervisor.py | 7 +- .../sdk/execution_time/workloads/__init__.py | 23 -- .../sdk/execution_time/workloads/task.py | 53 ---- .../executable/test_coordinator.py | 6 +- .../coordinators/java/test_coordinator.py | 6 +- .../task_sdk/coordinators/test_subprocess.py | 6 +- .../execution_time/test_supervisor.py | 38 +-- .../execution_time/test_task_runner.py | 33 ++- 26 files changed, 254 insertions(+), 498 deletions(-) delete mode 100755 scripts/ci/prek/check_task_instance_dto_sync.py delete mode 100644 task-sdk/src/airflow/sdk/execution_time/workloads/__init__.py delete mode 100644 task-sdk/src/airflow/sdk/execution_time/workloads/task.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 43a488e2bccf1..3da2fb0135b00 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -532,12 +532,6 @@ repos: language: python pass_filenames: false files: ^dev/registry/registry_tools/types\.py$|^registry/src/_data/types\.json$ - - id: check-task-instance-dto-sync - name: Check BaseTaskInstanceDTO duplicate is in sync between core and task-sdk - entry: ./scripts/ci/prek/check_task_instance_dto_sync.py - language: python - pass_filenames: false - files: ^airflow-core/src/airflow/executors/workloads/task\.py$|^task-sdk/src/airflow/sdk/execution_time/workloads/task\.py$ - id: ruff name: Run 'ruff' for extremely fast Python linting description: "Run 'ruff' for extremely fast Python linting" diff --git a/airflow-core/docs/core-concepts/executor/index.rst b/airflow-core/docs/core-concepts/executor/index.rst index 9420c55d84e26..0097c1239f4dd 100644 --- a/airflow-core/docs/core-concepts/executor/index.rst +++ b/airflow-core/docs/core-concepts/executor/index.rst @@ -271,20 +271,17 @@ Example: ExecuteTask( token="mock", - ti=TaskInstance( + ti=TaskInstanceDTO( id=UUID("4d828a62-a417-4936-a7a6-2b3fabacecab"), task_id="mock", dag_id="mock", run_id="mock", try_number=1, + dag_version_id=UUID("4d828a62-a417-4936-a7a6-2b3fabacecab"), map_index=-1, pool_slots=1, queue="default", priority_weight=1, - executor_config=None, - parent_context_carrier=None, - context_carrier=None, - queued_dttm=None, ), dag_rel_path=PurePosixPath("mock.py"), bundle_info=BundleInfo(name="n/a", version="no matter"), diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py index c7f6b1ee9f8a9..f8d45123e94a3 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py @@ -286,6 +286,11 @@ class TaskInstance(BaseModel): map_index: int = -1 hostname: str | None = None context_carrier: dict | None = None + # Defaults keep hand-built instances (tests, dry runs) valid; the executor + # workload always sends the real values. + queue: str = "default" + pool_slots: int = 1 + priority_weight: int = 1 class AssetReferenceAssetEventDagRun(StrictBaseModel): diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py index 9e4d486aa3028..b536066f0cea8 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py @@ -49,6 +49,7 @@ from airflow.api_fastapi.execution_api.versions.v2026_06_30 import ( AddAwaitingInputStatePayload, AddConnectionTestEndpoint, + AddTaskInstanceQueueFields, AddVariableKeysEndpoint, ) @@ -59,6 +60,7 @@ AddVariableKeysEndpoint, AddConnectionTestEndpoint, AddAwaitingInputStatePayload, + AddTaskInstanceQueueFields, ), Version( "2026-06-16", diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_06_30.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_06_30.py index f9c22f6cd8281..31f227c6cb6e7 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_06_30.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_06_30.py @@ -19,7 +19,10 @@ from cadwyn import VersionChange, endpoint, schema -from airflow.api_fastapi.execution_api.datamodels.taskinstance import TIAwaitingInputStatePayload +from airflow.api_fastapi.execution_api.datamodels.taskinstance import ( + TaskInstance, + TIAwaitingInputStatePayload, +) class AddVariableKeysEndpoint(VersionChange): @@ -41,6 +44,18 @@ class AddConnectionTestEndpoint(VersionChange): ) +class AddTaskInstanceQueueFields(VersionChange): + """Add the `queue`, `pool_slots` and `priority_weight` fields to the TaskInstance model.""" + + description = __doc__ + + instructions_to_migrate_to_previous_version = ( + schema(TaskInstance).field("queue").didnt_exist, + schema(TaskInstance).field("pool_slots").didnt_exist, + schema(TaskInstance).field("priority_weight").didnt_exist, + ) + + class AddAwaitingInputStatePayload(VersionChange): """Add the awaiting_input task instance state transition payload (Human-in-the-loop, no trigger).""" diff --git a/airflow-core/src/airflow/executors/workloads/task.py b/airflow-core/src/airflow/executors/workloads/task.py index 9af3f33c10efd..584dfbb10d058 100644 --- a/airflow-core/src/airflow/executors/workloads/task.py +++ b/airflow-core/src/airflow/executors/workloads/task.py @@ -18,12 +18,12 @@ from __future__ import annotations -import uuid from pathlib import Path from typing import TYPE_CHECKING, Literal -from pydantic import BaseModel, Field +from pydantic import Field +from airflow.api_fastapi.execution_api.datamodels.taskinstance import TaskInstance from airflow.executors.workloads.base import BaseDagBundleWorkload, BundleInfo from airflow.utils.state import TaskInstanceState @@ -33,36 +33,17 @@ from airflow.models.taskinstancekey import TaskInstanceKey -class BaseTaskInstanceDTO(BaseModel): +class TaskInstanceDTO(TaskInstance): """ - Base schema for TaskInstance with the minimal fields shared by Executors and the Task SDK. + The versioned execution API ``TaskInstance`` schema plus executor-only fields. - This definition is duplicated in :mod:`airflow.sdk.execution_time.workloads.task` - and the two are kept in sync by the ``check-task-instance-dto-sync`` prek - hook. Update both files together. + The base class is the single source of truth for the fields a worker + receives; the fields added here are excluded from serialization and never + leave the executor. """ - id: uuid.UUID - dag_version_id: uuid.UUID - task_id: str - dag_id: str - run_id: str - try_number: int - map_index: int = -1 - - pool_slots: int - queue: str - priority_weight: int - executor_config: dict | None = Field(default=None, exclude=True) - - parent_context_carrier: dict | None = None - context_carrier: dict | None = None - - -class TaskInstanceDTO(BaseTaskInstanceDTO): - """TaskInstanceDTO with executor-specific ``external_executor_id`` field and ``key`` property.""" - external_executor_id: str | None = Field(default=None, exclude=True) + executor_config: dict | None = Field(default=None, exclude=True) # TODO: Task-SDK: Can we replace TaskInstanceKey with just the uuid across the codebase? @property diff --git a/airflow-core/tests/unit/executors/test_workloads.py b/airflow-core/tests/unit/executors/test_workloads.py index 6f027a4d3be6c..3522572701604 100644 --- a/airflow-core/tests/unit/executors/test_workloads.py +++ b/airflow-core/tests/unit/executors/test_workloads.py @@ -32,6 +32,7 @@ from airflow.executors.workloads.task import ExecuteTask from airflow.executors.workloads.types import state_class_for_key from airflow.models.callback import CallbackKey +from airflow.sdk.api.datamodels._generated import TaskInstance as GeneratedTaskInstance def test_task_instance_alias_keeps_backwards_compat(): @@ -134,3 +135,35 @@ def test_callback_dto_key_returns_callback_key_instance(): assert isinstance(key, CallbackKey) assert key.id == cid assert str(key) == cid + + +def test_workload_ti_round_trips_through_sdk_generated_model(): + """ + The executor-side DTO and the SDK's generated TaskInstance share the + execution API schema; the serialized workload must carry the routing + fields and exclude the executor-only ones. + """ + ti = TaskInstanceDTO( + id=uuid4(), + dag_version_id=uuid4(), + task_id="test_task", + dag_id="test_dag", + run_id="test_run", + try_number=2, + map_index=3, + pool_slots=4, + queue="jdk-17", + priority_weight=5, + external_executor_id="celery-id", + executor_config={"KubernetesExecutor": {"image": "custom"}}, + ) + + dumped = ti.model_dump(mode="json") + assert "external_executor_id" not in dumped + assert "executor_config" not in dumped + + received = GeneratedTaskInstance.model_validate(dumped) + assert received.queue == "jdk-17" + assert received.pool_slots == 4 + assert received.priority_weight == 5 + assert received.map_index == 3 diff --git a/devel-common/src/tests_common/pytest_plugin.py b/devel-common/src/tests_common/pytest_plugin.py index 0958dbd663a1a..1d3bd04768788 100644 --- a/devel-common/src/tests_common/pytest_plugin.py +++ b/devel-common/src/tests_common/pytest_plugin.py @@ -2584,16 +2584,12 @@ 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.sdk.api.datamodels._generated import ( + DagRun, + DagRunState, + TaskInstance as TaskInstanceDTO, + TIRunContext, + ) from airflow.utils.types import DagRunType if isinstance(logical_date, str): diff --git a/providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml b/providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml index 408cb14078381..8d29e60cc7f19 100644 --- a/providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml +++ b/providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml @@ -1277,10 +1277,6 @@ components: type: string format: uuid title: Id - dag_version_id: - type: string - format: uuid - title: Dag Version Id task_id: type: string title: Task Id @@ -1293,45 +1289,55 @@ components: try_number: type: integer title: Try Number + dag_version_id: + type: string + format: uuid + title: Dag Version Id map_index: type: integer title: Map Index default: -1 - pool_slots: - type: integer - title: Pool Slots - queue: - type: string - title: Queue - priority_weight: - type: integer - title: Priority Weight - parent_context_carrier: + hostname: anyOf: - - additionalProperties: true - type: object + - type: string - type: 'null' - title: Parent Context Carrier + title: Hostname context_carrier: anyOf: - additionalProperties: true type: object - type: 'null' title: Context Carrier + queue: + type: string + title: Queue + default: default + pool_slots: + type: integer + title: Pool Slots + default: 1 + priority_weight: + type: integer + title: Priority Weight + default: 1 type: object required: - id - - dag_version_id - task_id - dag_id - run_id - try_number - - pool_slots - - queue - - priority_weight + - dag_version_id title: TaskInstanceDTO - description: TaskInstanceDTO with executor-specific ``external_executor_id`` - field and ``key`` property. + description: 'The versioned execution API ``TaskInstance`` schema plus executor-only + fields. + + + The base class is the single source of truth for the fields a worker + + receives; the fields added here are excluded from serialization and never + + leave the executor.' TaskInstanceState: type: string enum: diff --git a/scripts/ci/prek/check_task_instance_dto_sync.py b/scripts/ci/prek/check_task_instance_dto_sync.py deleted file mode 100755 index 689d35a4d15e3..0000000000000 --- a/scripts/ci/prek/check_task_instance_dto_sync.py +++ /dev/null @@ -1,125 +0,0 @@ -#!/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. -""" -Verify that the duplicate ``BaseTaskInstanceDTO`` definitions in airflow-core -and task-sdk stay structurally identical. - -``BaseTaskInstanceDTO`` is duplicated (not shared) in: - -- ``airflow-core/src/airflow/executors/workloads/task.py`` -- ``task-sdk/src/airflow/sdk/execution_time/workloads/task.py`` - -This hook compares the *fields* (annotated assignments) and bases of both -``BaseTaskInstanceDTO`` classes. The concrete ``TaskInstanceDTO`` subclasses -in each file are allowed to differ (airflow-core adds an executor-specific -``key`` property that depends on ``airflow.models``, which the Task SDK -does not have access to). -""" - -from __future__ import annotations - -import ast -import sys -from pathlib import Path - -AIRFLOW_ROOT = Path(__file__).parents[3].resolve() -CORE_FILE = AIRFLOW_ROOT / "airflow-core" / "src" / "airflow" / "executors" / "workloads" / "task.py" -SDK_FILE = AIRFLOW_ROOT / "task-sdk" / "src" / "airflow" / "sdk" / "execution_time" / "workloads" / "task.py" -CLASS_NAME = "BaseTaskInstanceDTO" - - -def _find_class(tree: ast.AST, class_name: str) -> ast.ClassDef | None: - for node in ast.walk(tree): - if isinstance(node, ast.ClassDef) and node.name == class_name: - return node - return None - - -def _field_signature(class_node: ast.ClassDef) -> list[tuple[str, str, str | None]]: - """Return a normalized list of ``(name, annotation, default)`` for each field.""" - fields: list[tuple[str, str, str | None]] = [] - for stmt in class_node.body: - if isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name): - name = stmt.target.id - annotation = ast.unparse(stmt.annotation) - default = ast.unparse(stmt.value) if stmt.value is not None else None - fields.append((name, annotation, default)) - return fields - - -def _bases(class_node: ast.ClassDef) -> list[str]: - return [ast.unparse(base) for base in class_node.bases] - - -def _extract(file_path: Path) -> tuple[list[str], list[tuple[str, str, str | None]]]: - source = file_path.read_text() - tree = ast.parse(source, filename=str(file_path)) - class_node = _find_class(tree, CLASS_NAME) - if class_node is None: - print(f"ERROR: Could not find class {CLASS_NAME} in {file_path}", file=sys.stderr) - sys.exit(1) - return _bases(class_node), _field_signature(class_node) - - -def main() -> None: - core_bases, core_fields = _extract(CORE_FILE) - sdk_bases, sdk_fields = _extract(SDK_FILE) - - if core_bases == sdk_bases and core_fields == sdk_fields: - sys.exit(0) - - print( - f"\nERROR: {CLASS_NAME} definitions in airflow-core and task-sdk are out of sync!", - file=sys.stderr, - ) - print(f"\n airflow-core: {CORE_FILE.relative_to(AIRFLOW_ROOT)}", file=sys.stderr) - print(f" task-sdk: {SDK_FILE.relative_to(AIRFLOW_ROOT)}", file=sys.stderr) - - if core_bases != sdk_bases: - print("\nClass bases differ:", file=sys.stderr) - print(f" airflow-core: {core_bases}", file=sys.stderr) - print(f" task-sdk: {sdk_bases}", file=sys.stderr) - - if core_fields != sdk_fields: - core_set = {f[0]: f for f in core_fields} - sdk_set = {f[0]: f for f in sdk_fields} - only_in_core = sorted(set(core_set) - set(sdk_set)) - only_in_sdk = sorted(set(sdk_set) - set(core_set)) - differing = sorted(name for name in set(core_set) & set(sdk_set) if core_set[name] != sdk_set[name]) - if only_in_core: - print(f"\n Fields only in airflow-core: {only_in_core}", file=sys.stderr) - if only_in_sdk: - print(f"\n Fields only in task-sdk: {only_in_sdk}", file=sys.stderr) - for name in differing: - print( - f"\n Field {name!r} differs:" - f"\n airflow-core: {core_set[name]}" - f"\n task-sdk: {sdk_set[name]}", - file=sys.stderr, - ) - - print( - f"\nUpdate both files together so the two {CLASS_NAME} definitions stay in sync.", - file=sys.stderr, - ) - sys.exit(1) - - -if __name__ == "__main__": - main() diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 9bdd8ac4fba61..4823c05db94ba 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -560,6 +560,9 @@ class TaskInstance(BaseModel): map_index: Annotated[int | None, Field(title="Map Index")] = -1 hostname: Annotated[str | None, Field(title="Hostname")] = None context_carrier: Annotated[dict[str, Any] | None, Field(title="Context Carrier")] = None + queue: Annotated[str | None, Field(title="Queue")] = "default" + pool_slots: Annotated[int | None, Field(title="Pool Slots")] = 1 + priority_weight: Annotated[int | None, Field(title="Priority Weight")] = 1 class BundleInfo(BaseModel): diff --git a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py index 49bea634ac545..bac9f828b78d8 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py +++ b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py @@ -51,8 +51,7 @@ from typing_extensions import Self from airflow.sdk.api.client import Client - from airflow.sdk.api.datamodels._generated import BundleInfo - from airflow.sdk.execution_time.workloads.task import TaskInstanceDTO + from airflow.sdk.api.datamodels._generated import BundleInfo, TaskInstance Tracked = TypeVar("Tracked", socket.socket, subprocess.Popen) @@ -279,7 +278,7 @@ class _PopenActivitySubprocess(ActivitySubprocess): def start( # type: ignore[override] cls, *, - what: TaskInstanceDTO, + what: TaskInstance, dag_rel_path: str | os.PathLike[str], bundle_info, logger: FilteringBoundLogger | None = None, @@ -369,7 +368,7 @@ class SubprocessCoordinator(BaseCoordinator): task_startup_timeout: float = 10.0 - def _build_execute_task_command(self, *, what: TaskInstanceDTO) -> tuple[list[str], str | None]: + def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], str | None]: """ Build the subprocess command and resolve its supervisor wire-schema version for *what*. @@ -385,7 +384,7 @@ def _build_execute_task_command(self, *, what: TaskInstanceDTO) -> tuple[list[st def execute_task( self, *, - what: TaskInstanceDTO, + what: TaskInstance, dag_rel_path: str | os.PathLike[str], bundle_info: BundleInfo, client: Client, diff --git a/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py index 7b5290be021db..ed11f97165d1a 100644 --- a/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py @@ -40,7 +40,7 @@ from structlog.typing import FilteringBoundLogger from typing_extensions import Self - from airflow.sdk.execution_time.workloads.task import TaskInstanceDTO + from airflow.sdk.api.datamodels._generated import TaskInstance log: FilteringBoundLogger = structlog.get_logger(logger_name="coordinators.executable") @@ -390,6 +390,6 @@ class ExecutableCoordinator(SubprocessCoordinator): validator=attrs.validators.min_len(1), ) - def _build_execute_task_command(self, *, what: TaskInstanceDTO) -> tuple[list[str], str | None]: + def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], str | None]: bundle = _Bundle.find(self.executables_root, what.dag_id) return [str(bundle.path)], bundle.schema_version diff --git a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py index d6aebe707f59f..cb74c64dfd48f 100644 --- a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py @@ -38,7 +38,7 @@ from structlog.typing import FilteringBoundLogger from typing_extensions import Self - from airflow.sdk.execution_time.workloads.task import TaskInstanceDTO + from airflow.sdk.api.datamodels._generated import TaskInstance log: FilteringBoundLogger = structlog.get_logger(logger_name="coordinators.java") @@ -219,7 +219,7 @@ class JavaCoordinator(SubprocessCoordinator): ) main_class: str = "" - def _build_execute_task_command(self, *, what: TaskInstanceDTO) -> tuple[list[str], str | None]: + def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], str | None]: jar = _JarInfo.find(self.jars_root, self.main_class) command = [ self.java_executable, diff --git a/task-sdk/src/airflow/sdk/execution_time/comms.py b/task-sdk/src/airflow/sdk/execution_time/comms.py index 330fe934c7c45..c2c4fd6951617 100644 --- a/task-sdk/src/airflow/sdk/execution_time/comms.py +++ b/task-sdk/src/airflow/sdk/execution_time/comms.py @@ -80,6 +80,7 @@ PreviousTIResponse, PrevSuccessfulDagRunResponse, TaskBreadcrumbsResponse, + TaskInstance, TaskInstanceState, TaskStatesResponse, TaskStoreResponse, @@ -98,10 +99,6 @@ XComSequenceSliceResponse, ) from airflow.sdk.exceptions import ErrorType -from airflow.sdk.execution_time.workloads.task import ( - # Pydantic needs this at runtime since we don't model_rebuild() StartupDetails. - TaskInstanceDTO, # noqa: TC001 -) try: from socket import recv_fds @@ -338,7 +335,7 @@ def _get_response(self) -> ReceiveMsgType | None: class StartupDetails(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) - ti: TaskInstanceDTO + ti: TaskInstance dag_rel_path: str bundle_info: BundleInfo start_date: datetime diff --git a/task-sdk/src/airflow/sdk/execution_time/coordinator.py b/task-sdk/src/airflow/sdk/execution_time/coordinator.py index 072e9c0b6c739..e233ab8e4ae2b 100644 --- a/task-sdk/src/airflow/sdk/execution_time/coordinator.py +++ b/task-sdk/src/airflow/sdk/execution_time/coordinator.py @@ -57,7 +57,7 @@ from typing_extensions import Self from airflow.sdk.api.client import Client - from airflow.sdk.execution_time.workloads.task import TaskInstanceDTO + from airflow.sdk.api.datamodels._generated import TaskInstance __all__ = [ "BaseCoordinator", @@ -89,7 +89,7 @@ class ExecutionResult: def execute_task( self, *, - what: TaskInstanceDTO, + what: TaskInstance, dag_rel_path: str | PathLike[str], bundle_info, client: Client, @@ -122,7 +122,7 @@ class _PythonCoordinator(BaseCoordinator): def execute_task( self, *, - what: TaskInstanceDTO, + what: TaskInstance, dag_rel_path: str | PathLike[str], bundle_info, client: Client, diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index c02aa76791ed4..a2dde7b9998c9 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -3656,7 +3656,7 @@ "StartupDetails": { "properties": { "ti": { - "$ref": "#/$defs/TaskInstanceDTO" + "$ref": "#/$defs/TaskInstance" }, "dag_rel_path": { "title": "Dag Rel Path", @@ -3891,164 +3891,6 @@ "title": "TaskCallbackRequest", "type": "object" }, - "TaskInstance": { - "description": "Schema for TaskInstance model with minimal required fields needed for Runtime.", - "properties": { - "id": { - "format": "uuid", - "title": "Id", - "type": "string" - }, - "task_id": { - "title": "Task Id", - "type": "string" - }, - "dag_id": { - "title": "Dag Id", - "type": "string" - }, - "run_id": { - "title": "Run Id", - "type": "string" - }, - "try_number": { - "title": "Try Number", - "type": "integer" - }, - "dag_version_id": { - "format": "uuid", - "title": "Dag Version Id", - "type": "string" - }, - "map_index": { - "default": -1, - "title": "Map Index", - "type": "integer" - }, - "hostname": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "default": null, - "title": "Hostname" - }, - "context_carrier": { - "anyOf": [ - { - "additionalProperties": true, - "type": "object" - }, - { - "type": "null" - } - ], - "default": null, - "title": "Context Carrier" - } - }, - "required": [ - "id", - "task_id", - "dag_id", - "run_id", - "try_number", - "dag_version_id" - ], - "title": "TaskInstance", - "type": "object" - }, - "TaskInstanceDTO": { - "description": "Task SDK TaskInstanceDTO.", - "properties": { - "id": { - "format": "uuid", - "title": "Id", - "type": "string" - }, - "dag_version_id": { - "format": "uuid", - "title": "Dag Version Id", - "type": "string" - }, - "task_id": { - "title": "Task Id", - "type": "string" - }, - "dag_id": { - "title": "Dag Id", - "type": "string" - }, - "run_id": { - "title": "Run Id", - "type": "string" - }, - "try_number": { - "title": "Try Number", - "type": "integer" - }, - "map_index": { - "default": -1, - "title": "Map Index", - "type": "integer" - }, - "pool_slots": { - "title": "Pool Slots", - "type": "integer" - }, - "queue": { - "title": "Queue", - "type": "string" - }, - "priority_weight": { - "title": "Priority Weight", - "type": "integer" - }, - "parent_context_carrier": { - "anyOf": [ - { - "additionalProperties": true, - "type": "object" - }, - { - "type": "null" - } - ], - "default": null, - "title": "Parent Context Carrier" - }, - "context_carrier": { - "anyOf": [ - { - "additionalProperties": true, - "type": "object" - }, - { - "type": "null" - } - ], - "default": null, - "title": "Context Carrier" - } - }, - "required": [ - "id", - "dag_version_id", - "task_id", - "dag_id", - "run_id", - "try_number", - "pool_slots", - "queue", - "priority_weight" - ], - "title": "TaskInstanceDTO", - "type": "object" - }, "TaskInstanceState": { "description": "All possible states that a Task Instance can be in.\n\nNote that None is also allowed, so always use this in a type hint with Optional.", "enum": [ @@ -4930,6 +4772,92 @@ "title": "TIRunContext", "type": "object" }, + "TaskInstance": { + "description": "Schema for TaskInstance model with minimal required fields needed for Runtime.", + "properties": { + "id": { + "format": "uuid", + "title": "Id", + "type": "string" + }, + "task_id": { + "title": "Task Id", + "type": "string" + }, + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "run_id": { + "title": "Run Id", + "type": "string" + }, + "try_number": { + "title": "Try Number", + "type": "integer" + }, + "dag_version_id": { + "format": "uuid", + "title": "Dag Version Id", + "type": "string" + }, + "map_index": { + "default": -1, + "title": "Map Index", + "type": "integer" + }, + "hostname": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Hostname" + }, + "context_carrier": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Context Carrier" + }, + "queue": { + "default": "default", + "title": "Queue", + "type": "string" + }, + "pool_slots": { + "default": 1, + "title": "Pool Slots", + "type": "integer" + }, + "priority_weight": { + "default": 1, + "title": "Priority Weight", + "type": "integer" + } + }, + "required": [ + "id", + "task_id", + "dag_id", + "run_id", + "try_number", + "dag_version_id" + ], + "title": "TaskInstance", + "type": "object" + }, "VariableResponse": { "additionalProperties": false, "description": "Variable schema for responses with fields that are needed for Runtime.", diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py index d8ae4e4580ed5..6acf83994ec68 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py @@ -21,5 +21,8 @@ bundle = VersionBundle( HeadVersion(), + # The initial, not-yet-released version of the contract. Shape changes to + # registered models land directly in the head snapshot until this ships; + # no VersionChange entries are possible against the bundle floor. Version("2026-06-16"), ) diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 6527e651041b8..fc21104c50251 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -172,7 +172,6 @@ from airflow.executors.workloads import BundleInfo from airflow.sdk.bases.secrets_backend import BaseSecretsBackend from airflow.sdk.definitions.connection import Connection - from airflow.sdk.execution_time.workloads.task import TaskInstanceDTO from airflow.sdk.types import RuntimeTaskInstanceProtocol as RuntimeTI __all__ = ["ActivitySubprocess", "WatchedSubprocess", "supervise", "supervise_task"] @@ -1317,7 +1316,7 @@ class ActivitySubprocess(WatchedSubprocess): def start( # type: ignore[override] cls, *, - what: TaskInstanceDTO, + what: TaskInstance, dag_rel_path: str | os.PathLike[str], bundle_info, client: Client, @@ -1346,7 +1345,7 @@ def start( # type: ignore[override] def _on_child_started( self, *, - ti: TaskInstanceDTO, + ti: TaskInstance, dag_rel_path: str | os.PathLike[str], bundle_info, sentry_integration: str, @@ -2412,7 +2411,7 @@ def _configure_logging(log_path: str, client: Client) -> tuple[FilteringBoundLog def supervise_task( *, - ti: TaskInstanceDTO, + ti: TaskInstance, bundle_info: BundleInfo, dag_rel_path: str | os.PathLike[str], token: str, diff --git a/task-sdk/src/airflow/sdk/execution_time/workloads/__init__.py b/task-sdk/src/airflow/sdk/execution_time/workloads/__init__.py deleted file mode 100644 index cdf955e742dc0..0000000000000 --- a/task-sdk/src/airflow/sdk/execution_time/workloads/__init__.py +++ /dev/null @@ -1,23 +0,0 @@ -# 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. -"""Workload schemas for Task SDK execution-time communication.""" - -from __future__ import annotations - -from airflow.sdk.execution_time.workloads.task import TaskInstanceDTO - -__all__ = ["TaskInstanceDTO"] diff --git a/task-sdk/src/airflow/sdk/execution_time/workloads/task.py b/task-sdk/src/airflow/sdk/execution_time/workloads/task.py deleted file mode 100644 index ceff200856f06..0000000000000 --- a/task-sdk/src/airflow/sdk/execution_time/workloads/task.py +++ /dev/null @@ -1,53 +0,0 @@ -# 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 workload schemas for Task SDK execution-time communication.""" - -from __future__ import annotations - -import uuid - -from pydantic import BaseModel, Field - - -class BaseTaskInstanceDTO(BaseModel): - """ - Base schema for TaskInstance with the minimal fields shared by Executors and the Task SDK. - - This class is duplicated in :mod:`airflow.executors.workloads.task` and the - two definitions are kept in sync by the ``check-task-instance-dto-sync`` - prek hook. Update both files together. - """ - - id: uuid.UUID - dag_version_id: uuid.UUID - task_id: str - dag_id: str - run_id: str - try_number: int - map_index: int = -1 - - pool_slots: int - queue: str - priority_weight: int - executor_config: dict | None = Field(default=None, exclude=True) - - parent_context_carrier: dict | None = None - context_carrier: dict | None = None - - -class TaskInstanceDTO(BaseTaskInstanceDTO): - """Task SDK TaskInstanceDTO.""" diff --git a/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py index 23b45e4f1446b..dc8081694975b 100644 --- a/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py @@ -31,6 +31,7 @@ import yaml from uuid6 import uuid7 +from airflow.sdk.api.datamodels._generated import TaskInstance from airflow.sdk.coordinators.executable.coordinator import ( FOOTER_MAGIC, FOOTER_SIZE, @@ -41,7 +42,6 @@ ) from airflow.sdk.execution_time.coordinator import BaseCoordinator from airflow.sdk.execution_time.supervisor import ActivitySubprocess -from airflow.sdk.execution_time.workloads.task import TaskInstanceDTO from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS @@ -111,8 +111,8 @@ def _make_executable(path: Path) -> Path: return path -def _make_ti(dag_id: str = "tutorial_dag", queue: str = "executable") -> TaskInstanceDTO: - return TaskInstanceDTO( +def _make_ti(dag_id: str = "tutorial_dag", queue: str = "executable") -> TaskInstance: + return TaskInstance( id=uuid7(), dag_version_id=uuid7(), task_id="task_1", diff --git a/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py index 8670a0e895c4a..816135fdaf590 100644 --- a/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py @@ -29,6 +29,7 @@ import pytest from uuid6 import uuid7 +from airflow.sdk.api.datamodels._generated import TaskInstance from airflow.sdk.coordinators.java.coordinator import ( JavaCoordinator, _calculate_classpath, @@ -37,7 +38,6 @@ ) from airflow.sdk.execution_time.coordinator import BaseCoordinator from airflow.sdk.execution_time.supervisor import ActivitySubprocess -from airflow.sdk.execution_time.workloads.task import TaskInstanceDTO from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS @@ -45,8 +45,8 @@ pytest.skip("Coordinator is only compatible with Airflow >= 3.3.0", allow_module_level=True) -def _make_ti(dag_id: str = "test_dag", queue: str = "java") -> TaskInstanceDTO: - return TaskInstanceDTO( +def _make_ti(dag_id: str = "test_dag", queue: str = "java") -> TaskInstance: + return TaskInstance( id=uuid7(), dag_version_id=uuid7(), task_id="task_1", diff --git a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py index 2c043526a4e2a..34a1d0ec3d40f 100644 --- a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py +++ b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py @@ -32,6 +32,7 @@ from uuid6 import uuid7 from airflow.sdk.api.client import Client, TaskInstanceOperations +from airflow.sdk.api.datamodels._generated import TaskInstance from airflow.sdk.coordinators._subprocess import ( SubprocessCoordinator, _accept_connections, @@ -43,7 +44,6 @@ ) from airflow.sdk.execution_time.coordinator import BaseCoordinator from airflow.sdk.execution_time.supervisor import ActivitySubprocess -from airflow.sdk.execution_time.workloads.task import TaskInstanceDTO from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS @@ -51,8 +51,8 @@ pytest.skip("Coordinator is only compatible with Airflow >= 3.3.0", allow_module_level=True) -def _make_ti(dag_id: str = "tutorial_dag", queue: str = "socket") -> TaskInstanceDTO: - return TaskInstanceDTO( +def _make_ti(dag_id: str = "tutorial_dag", queue: str = "socket") -> TaskInstance: + return TaskInstance( id=uuid7(), dag_version_id=uuid7(), task_id="task_1", diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py index 62a9d00fdf286..eb126f0e81743 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -65,6 +65,7 @@ DagRunState, DagRunType, PreviousTIResponse, + TaskInstance, TaskInstanceState, ) from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType, TaskAlreadyRunningError @@ -172,7 +173,6 @@ supervise_task, ) from airflow.sdk.execution_time.task_runner import run -from airflow.sdk.execution_time.workloads.task import TaskInstanceDTO from tests_common.test_utils.config import conf_vars @@ -235,7 +235,7 @@ def test_supervise( """ Test that the supervisor validates server URL and dry_run parameter combinations correctly. """ - ti = TaskInstanceDTO( + ti = TaskInstance( id=uuid7(), task_id="async", dag_id="super_basic_deferred_run", @@ -271,7 +271,7 @@ def test_on_kill_hook_called_when_supervisor_receives_sigterm( client_with_ti_start, ): """SIGTERM to the supervisor process is forwarded to the task subprocess.""" - ti = TaskInstanceDTO( + ti = TaskInstance( id=uuid7(), task_id="signal_task", dag_id="signal_forward_test", @@ -379,7 +379,7 @@ def subprocess_main(): proc = ActivitySubprocess.start( dag_rel_path=os.devnull, bundle_info=FAKE_BUNDLE, - what=TaskInstanceDTO( + what=TaskInstance( id="4d828a62-a417-4936-a7a6-2b3fabacecab", task_id="b", dag_id="c", @@ -457,7 +457,7 @@ def subprocess_main(): proc = ActivitySubprocess.start( dag_rel_path=os.devnull, bundle_info=FAKE_BUNDLE, - what=TaskInstanceDTO( + what=TaskInstance( id="4d828a62-a417-4936-a7a6-2b3fabacecab", task_id="b", dag_id="c", @@ -557,7 +557,7 @@ def on_kill(self) -> None: proc = ActivitySubprocess.start( dag_rel_path=os.devnull, bundle_info=FAKE_BUNDLE, - what=TaskInstanceDTO( + what=TaskInstance( id=ti_id, task_id="b", dag_id="c", @@ -589,7 +589,7 @@ def subprocess_main(): proc = ActivitySubprocess.start( dag_rel_path=os.devnull, bundle_info=FAKE_BUNDLE, - what=TaskInstanceDTO( + what=TaskInstance( id="4d828a62-a417-4936-a7a6-2b3fabacecab", task_id="b", dag_id="c", @@ -627,7 +627,7 @@ def subprocess_main(): proc = ActivitySubprocess.start( dag_rel_path=os.devnull, bundle_info=FAKE_BUNDLE, - what=TaskInstanceDTO( + what=TaskInstance( id=uuid7(), task_id="b", dag_id="c", @@ -674,7 +674,7 @@ def test_resume_start_date_from_context(self, mocker, make_ti_context, start_dat proc = ActivitySubprocess.start( dag_rel_path=os.devnull, bundle_info=FAKE_BUNDLE, - what=TaskInstanceDTO( + what=TaskInstance( id=uuid7(), task_id="b", dag_id="c", @@ -720,7 +720,7 @@ def subprocess_main(): proc = ActivitySubprocess.start( dag_rel_path=os.devnull, bundle_info=FAKE_BUNDLE, - what=TaskInstanceDTO( + what=TaskInstance( id=ti_id, task_id="b", dag_id="c", @@ -765,7 +765,7 @@ def _on_child_started(self, *args, **kwargs): proc = ActivitySubprocess.start( dag_rel_path=os.devnull, bundle_info=FAKE_BUNDLE, - what=TaskInstanceDTO( + what=TaskInstance( id=ti_id, task_id="b", dag_id="c", @@ -789,7 +789,7 @@ def test_run_simple_dag(self, test_dags_dir, captured_logs, time_machine, mocker time_machine.move_to(instant, tick=False) dagfile_path = test_dags_dir - ti = TaskInstanceDTO( + ti = TaskInstance( id=uuid7(), task_id="hello", dag_id="super_basic_run", @@ -833,7 +833,7 @@ def test_supervise_handles_deferred_task( """ instant = timezone.datetime(2024, 11, 7, 12, 34, 56, 0) - ti = TaskInstanceDTO( + ti = TaskInstance( id=uuid7(), task_id="async", dag_id="super_basic_deferred_run", @@ -963,7 +963,7 @@ def handle_request(request: httpx.Request) -> httpx.Response: proc = ActivitySubprocess.start( dag_rel_path=os.devnull, - what=TaskInstanceDTO( + what=TaskInstance( id=ti_id, task_id="b", dag_id="c", @@ -1049,7 +1049,7 @@ def subprocess_main(): ActivitySubprocess.start( dag_rel_path=os.devnull, bundle_info=FAKE_BUNDLE, - what=TaskInstanceDTO( + what=TaskInstance( id=ti_id, task_id="b", dag_id="c", @@ -1262,7 +1262,7 @@ def subprocess_main(): proc = ActivitySubprocess.start( dag_rel_path=os.devnull, bundle_info=FAKE_BUNDLE, - what=TaskInstanceDTO( + what=TaskInstance( id="4d828a62-a417-4936-a7a6-2b3fabacecab", task_id="b", dag_id="c", @@ -1425,7 +1425,7 @@ def _handler(sig, frame): proc = ActivitySubprocess.start( dag_rel_path=os.devnull, bundle_info=FAKE_BUNDLE, - what=TaskInstanceDTO( + what=TaskInstance( id=ti_id, task_id="b", dag_id="c", @@ -3454,7 +3454,7 @@ def execute(self, context: Context): task.dag = DAG(dag_id="test_dag") # Create a simple TaskInstance datamodel to pass to the supervisor - ti = TaskInstanceDTO( + ti = TaskInstance( id=uuid7(), dag_version_id=uuid7(), dag_id="test_dag", @@ -3880,7 +3880,7 @@ def subprocess_main(): proc = ActivitySubprocess.start( dag_rel_path=os.devnull, bundle_info=FAKE_BUNDLE, - what=TaskInstanceDTO( + what=TaskInstance( id="4d828a62-a417-4936-a7a6-2b3fabacecab", task_id="b", dag_id="c", diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py index 471541a6429d5..18b4da9018847 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py +++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py @@ -176,7 +176,6 @@ run, startup, ) -from airflow.sdk.execution_time.workloads.task import TaskInstanceDTO from airflow.sdk.execution_time.xcom import XCom from airflow.sdk.serde import deserialize from airflow.triggers.base import BaseEventTrigger, BaseTrigger, TriggerEvent @@ -208,7 +207,7 @@ def execute(self, context): def test_parse(test_dags_dir: Path, make_ti_context): """Test that checks parsing of a basic dag with an un-mocked parse.""" what = StartupDetails( - ti=TaskInstanceDTO( + ti=TaskInstance( id=uuid7(), task_id="a", dag_id="super_basic", @@ -263,7 +262,7 @@ def test_parse_dag_bag(mock_dagbag, test_dags_dir: Path, make_ti_context): mock_dag.tasks = [mock_task] what = StartupDetails( - ti=TaskInstanceDTO( + ti=TaskInstance( id=uuid7(), task_id="a", dag_id="super_basic", @@ -326,7 +325,7 @@ def test_parse_dag_bag(mock_dagbag, test_dags_dir: Path, make_ti_context): def test_parse_not_found(test_dags_dir: Path, make_ti_context, dag_id, task_id, expected_error): """Check for nice error messages on dag not found.""" what = StartupDetails( - ti=TaskInstanceDTO( + ti=TaskInstance( id=uuid7(), task_id=task_id, dag_id=dag_id, @@ -375,7 +374,7 @@ def test_parse_not_found_does_not_reschedule_when_max_attempts_reached(test_dags and should surface as a hard failure (SystemExit in the task runner process). """ what = StartupDetails( - ti=TaskInstanceDTO( + ti=TaskInstance( id=uuid7(), task_id="a", dag_id="madeup_dag_id", @@ -439,7 +438,7 @@ def test_main_sends_reschedule_task_when_startup_reschedules( mock_comms_instance.socket = None mock_comms_decoder_cls.__getitem__.return_value.return_value = mock_comms_instance what = StartupDetails( - ti=TaskInstanceDTO( + ti=TaskInstance( id=uuid7(), task_id="my_task", dag_id="test_dag", @@ -615,7 +614,7 @@ def test_task_span_is_child_of_dag_run_span(make_ti_context): # Step 3: build StartupDetails with ti.context_carrier = ti_carrier. what = StartupDetails( - ti=TaskInstanceDTO( + ti=TaskInstance( id=uuid7(), task_id="my_task", dag_id="test_dag", @@ -686,7 +685,7 @@ def test_task_span_no_parent_when_no_context_carrier(make_ti_context): provider.add_span_processor(SimpleSpanProcessor(in_mem_exporter)) what = StartupDetails( - ti=TaskInstanceDTO( + ti=TaskInstance( id=uuid7(), task_id="standalone_task", dag_id="test_dag", @@ -730,7 +729,7 @@ def test_parse_module_in_bundle_root(tmp_path: Path, make_ti_context): dag1_path.write_text(textwrap.dedent(dag1_code)) what = StartupDetails( - ti=TaskInstanceDTO( + ti=TaskInstance( id=uuid7(), task_id="a", dag_id="dag_name", @@ -1180,7 +1179,7 @@ def test_basic_templated_dag(mocked_parse, make_ti_context, mock_supervisor_comm ) what = StartupDetails( - ti=TaskInstanceDTO( + ti=TaskInstance( id=uuid7(), task_id="templated_task", dag_id="basic_templated_dag", @@ -1299,7 +1298,7 @@ def execute(self, context): instant = timezone.datetime(2024, 12, 3, 10, 0) what = StartupDetails( - ti=TaskInstanceDTO( + ti=TaskInstance( id=uuid7(), task_id="templated_task", dag_id="basic_dag", @@ -1350,7 +1349,7 @@ def execute(self, context): instant = timezone.datetime(2024, 12, 3, 10, 0) what = StartupDetails( - ti=TaskInstanceDTO( + ti=TaskInstance( id=uuid7(), task_id="impersonation_task", dag_id="basic_dag", @@ -1401,7 +1400,7 @@ def execute(self, context): instant = timezone.datetime(2024, 12, 3, 10, 0) what = StartupDetails( - ti=TaskInstanceDTO( + ti=TaskInstance( id=uuid7(), task_id="impersonation_task", dag_id="basic_dag", @@ -1444,7 +1443,7 @@ def execute(self, context): instant = timezone.datetime(2024, 12, 3, 10, 0) what = StartupDetails( - ti=TaskInstanceDTO( + ti=TaskInstance( id=uuid7(), task_id="impersonation_task", dag_id="basic_dag", @@ -1620,7 +1619,7 @@ def test_dag_parsing_context(make_ti_context, mock_supervisor_comms, monkeypatch task_id = "conditional_task" what = StartupDetails( - ti=TaskInstanceDTO( + ti=TaskInstance( id=uuid7(), task_id=task_id, dag_id=dag_id, @@ -2269,7 +2268,7 @@ def execute(self, context): task = CustomOperator(task_id=test_task_id) # In case of the specific map_index we should check it is passed to TI. - # ``None`` is not a valid TaskInstanceDTO.map_index value, but xcom_pull's + # ``None`` is not a valid TaskInstance.map_index value, but xcom_pull's # behaviour with ``map_indexes=None`` is independent of the TI's own map_index. extra_for_ti = {"map_index": map_indexes} if isinstance(map_indexes, int) else {} runtime_ti = create_runtime_ti(task=task, **extra_for_ti) @@ -4166,7 +4165,7 @@ def execute(self, context): task_id="test_task_runner_calls_listeners", do_xcom_push=True, multiple_outputs=True ) what = StartupDetails( - ti=TaskInstanceDTO( + ti=TaskInstance( id=uuid7(), task_id="templated_task", dag_id="basic_dag", From 78a6b821503f9496fadbf31f70a7900210245131 Mon Sep 17 00:00:00 2001 From: Kaxil Naik Date: Thu, 11 Jun 2026 15:45:30 +0100 Subject: [PATCH 2/7] Keep pool_slots and priority_weight executor-side; only queue joins the schema The supervisor only routes on queue; pool_slots and priority_weight are executor concerns (queued-workload priority ordering and edge concurrency slot accounting) that the worker never reads. Keep them on the executor DTO, serialized as before, and add only queue to the worker-facing schema. --- .../execution_api/datamodels/taskinstance.py | 7 +++---- .../api_fastapi/execution_api/versions/__init__.py | 4 ++-- .../api_fastapi/execution_api/versions/v2026_06_30.py | 10 +++------- airflow-core/src/airflow/executors/workloads/task.py | 9 ++++++--- airflow-core/tests/unit/executors/test_workloads.py | 8 ++++++-- .../providers/edge3/worker_api/v2-edge-generated.yaml | 10 +++++----- task-sdk/src/airflow/sdk/api/datamodels/_generated.py | 2 -- .../src/airflow/sdk/execution_time/schema/schema.json | 10 ---------- 8 files changed, 25 insertions(+), 35 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py index f8d45123e94a3..fe7f5a5ce050d 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py @@ -286,11 +286,10 @@ class TaskInstance(BaseModel): map_index: int = -1 hostname: str | None = None context_carrier: dict | None = None - # Defaults keep hand-built instances (tests, dry runs) valid; the executor - # workload always sends the real values. + # The supervisor routes tasks to a coordinator by queue. The default keeps + # hand-built instances (tests, dry runs) valid; the executor workload + # always sends the real value. queue: str = "default" - pool_slots: int = 1 - priority_weight: int = 1 class AssetReferenceAssetEventDagRun(StrictBaseModel): diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py index b536066f0cea8..656bb8dce107f 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py @@ -49,7 +49,7 @@ from airflow.api_fastapi.execution_api.versions.v2026_06_30 import ( AddAwaitingInputStatePayload, AddConnectionTestEndpoint, - AddTaskInstanceQueueFields, + AddTaskInstanceQueueField, AddVariableKeysEndpoint, ) @@ -60,7 +60,7 @@ AddVariableKeysEndpoint, AddConnectionTestEndpoint, AddAwaitingInputStatePayload, - AddTaskInstanceQueueFields, + AddTaskInstanceQueueField, ), Version( "2026-06-16", diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_06_30.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_06_30.py index 31f227c6cb6e7..cfa5b616396bf 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_06_30.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_06_30.py @@ -44,16 +44,12 @@ class AddConnectionTestEndpoint(VersionChange): ) -class AddTaskInstanceQueueFields(VersionChange): - """Add the `queue`, `pool_slots` and `priority_weight` fields to the TaskInstance model.""" +class AddTaskInstanceQueueField(VersionChange): + """Add the `queue` field to the TaskInstance model.""" description = __doc__ - instructions_to_migrate_to_previous_version = ( - schema(TaskInstance).field("queue").didnt_exist, - schema(TaskInstance).field("pool_slots").didnt_exist, - schema(TaskInstance).field("priority_weight").didnt_exist, - ) + instructions_to_migrate_to_previous_version = (schema(TaskInstance).field("queue").didnt_exist,) class AddAwaitingInputStatePayload(VersionChange): diff --git a/airflow-core/src/airflow/executors/workloads/task.py b/airflow-core/src/airflow/executors/workloads/task.py index 584dfbb10d058..b4bf02ea47b8d 100644 --- a/airflow-core/src/airflow/executors/workloads/task.py +++ b/airflow-core/src/airflow/executors/workloads/task.py @@ -37,11 +37,14 @@ class TaskInstanceDTO(TaskInstance): """ The versioned execution API ``TaskInstance`` schema plus executor-only fields. - The base class is the single source of truth for the fields a worker - receives; the fields added here are excluded from serialization and never - leave the executor. + The base class is the single source of truth for the fields a worker needs; + the fields added here are executor concerns (queueing order and pool + accounting) the worker never reads. """ + pool_slots: int + priority_weight: int + external_executor_id: str | None = Field(default=None, exclude=True) executor_config: dict | None = Field(default=None, exclude=True) diff --git a/airflow-core/tests/unit/executors/test_workloads.py b/airflow-core/tests/unit/executors/test_workloads.py index 3522572701604..a063b91140dfb 100644 --- a/airflow-core/tests/unit/executors/test_workloads.py +++ b/airflow-core/tests/unit/executors/test_workloads.py @@ -161,9 +161,13 @@ def test_workload_ti_round_trips_through_sdk_generated_model(): dumped = ti.model_dump(mode="json") assert "external_executor_id" not in dumped assert "executor_config" not in dumped + # Executor-side scheduling fields stay on the workload wire (older workers + # deserialize the workload with a model that requires them) but are not + # part of the worker-facing schema. + assert dumped["pool_slots"] == 4 + assert dumped["priority_weight"] == 5 received = GeneratedTaskInstance.model_validate(dumped) assert received.queue == "jdk-17" - assert received.pool_slots == 4 - assert received.priority_weight == 5 assert received.map_index == 3 + assert not hasattr(received, "pool_slots") diff --git a/providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml b/providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml index 8d29e60cc7f19..7a7fb194fb03d 100644 --- a/providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml +++ b/providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml @@ -1315,11 +1315,9 @@ components: pool_slots: type: integer title: Pool Slots - default: 1 priority_weight: type: integer title: Priority Weight - default: 1 type: object required: - id @@ -1328,16 +1326,18 @@ components: - run_id - try_number - dag_version_id + - pool_slots + - priority_weight title: TaskInstanceDTO description: 'The versioned execution API ``TaskInstance`` schema plus executor-only fields. - The base class is the single source of truth for the fields a worker + The base class is the single source of truth for the fields a worker needs; - receives; the fields added here are excluded from serialization and never + the fields added here are executor concerns (queueing order and pool - leave the executor.' + accounting) the worker never reads.' TaskInstanceState: type: string enum: diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 4823c05db94ba..99aedae25b03b 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -561,8 +561,6 @@ class TaskInstance(BaseModel): hostname: Annotated[str | None, Field(title="Hostname")] = None context_carrier: Annotated[dict[str, Any] | None, Field(title="Context Carrier")] = None queue: Annotated[str | None, Field(title="Queue")] = "default" - pool_slots: Annotated[int | None, Field(title="Pool Slots")] = 1 - priority_weight: Annotated[int | None, Field(title="Priority Weight")] = 1 class BundleInfo(BaseModel): diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index a2dde7b9998c9..0d1a3dc61c703 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -4835,16 +4835,6 @@ "default": "default", "title": "Queue", "type": "string" - }, - "pool_slots": { - "default": 1, - "title": "Pool Slots", - "type": "integer" - }, - "priority_weight": { - "default": 1, - "title": "Priority Weight", - "type": "integer" } }, "required": [ From 749d874c1aa12e0765eac2d6cf5e8981b0045c24 Mon Sep 17 00:00:00 2001 From: Kaxil Naik Date: Thu, 11 Jun 2026 16:00:09 +0100 Subject: [PATCH 3/7] Update Java SDK test for the renamed supervisor schema TaskInstance model --- .../test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt index 1bf91220f9342..0542516228de8 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt @@ -29,7 +29,7 @@ import org.apache.airflow.sdk.execution.comm.DagRun import org.apache.airflow.sdk.execution.comm.StartupDetails import org.apache.airflow.sdk.execution.comm.SucceedTask import org.apache.airflow.sdk.execution.comm.TIRunContext -import org.apache.airflow.sdk.execution.comm.TaskInstanceDTO +import org.apache.airflow.sdk.execution.comm.TaskInstance import org.apache.airflow.sdk.execution.comm.TaskState import org.junit.jupiter.api.Assertions import org.junit.jupiter.api.DisplayName @@ -76,7 +76,7 @@ class TaskTest { private fun startupDetails(taskId: String): StartupDetails = StartupDetails().also { it.ti = - TaskInstanceDTO().also { o -> + TaskInstance().also { o -> o.id = UUID.randomUUID() o.taskId = taskId o.dagId = "test_dag" From ecd7836ef6526f74d884a6ec790a52a10acb09bb Mon Sep 17 00:00:00 2001 From: Kaxil Naik Date: Thu, 11 Jun 2026 16:59:12 +0100 Subject: [PATCH 4/7] Replace confusing version-bundle comment with a module docstring --- .../airflow/sdk/execution_time/schema/versions/__init__.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py index 6acf83994ec68..10b529a9825b8 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py @@ -14,6 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +"""Supervisor wire-schema version bundle. See the AGENTS.md in the parent package for how to evolve it.""" from __future__ import annotations @@ -21,8 +22,5 @@ bundle = VersionBundle( HeadVersion(), - # The initial, not-yet-released version of the contract. Shape changes to - # registered models land directly in the head snapshot until this ships; - # no VersionChange entries are possible against the bundle floor. Version("2026-06-16"), ) From f1231d99ea483edb8b80b2f7cd593c64a8e030af Mon Sep 17 00:00:00 2001 From: Kaxil Naik Date: Thu, 11 Jun 2026 17:04:28 +0100 Subject: [PATCH 5/7] Remove the version-bundle comment --- .../src/airflow/sdk/execution_time/schema/versions/__init__.py | 1 - 1 file changed, 1 deletion(-) diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py index 10b529a9825b8..d8ae4e4580ed5 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py @@ -14,7 +14,6 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""Supervisor wire-schema version bundle. See the AGENTS.md in the parent package for how to evolve it.""" from __future__ import annotations From 3d48d40bc94c8719e57c056eba5d0dcabf3cb3ca Mon Sep 17 00:00:00 2001 From: Kaxil Naik Date: Thu, 11 Jun 2026 17:08:17 +0100 Subject: [PATCH 6/7] Restore executor_config in the executor doc example --- airflow-core/docs/core-concepts/executor/index.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/airflow-core/docs/core-concepts/executor/index.rst b/airflow-core/docs/core-concepts/executor/index.rst index 0097c1239f4dd..cf21daafd3ffd 100644 --- a/airflow-core/docs/core-concepts/executor/index.rst +++ b/airflow-core/docs/core-concepts/executor/index.rst @@ -282,6 +282,7 @@ Example: pool_slots=1, queue="default", priority_weight=1, + executor_config=None, ), dag_rel_path=PurePosixPath("mock.py"), bundle_info=BundleInfo(name="n/a", version="no matter"), From ea2e92db7ffb37fe936d07e84fdc758f1ff38e69 Mon Sep 17 00:00:00 2001 From: Kaxil Naik Date: Thu, 11 Jun 2026 18:12:17 +0100 Subject: [PATCH 7/7] Fix mypy errors and ECS executor test after TaskInstance schema change The generated TaskInstance no longer accepts pool_slots/priority_weight kwargs, its queue is Optional in the generated typing, and the DTO now validates the inherited hostname field that the ECS adoption test's mock did not set. --- .../src/tests_common/pytest_plugin.py | 2 -- .../aws/executors/ecs/test_ecs_executor.py | 1 + .../airflow/sdk/execution_time/supervisor.py | 2 +- .../executable/test_coordinator.py | 2 -- .../coordinators/java/test_coordinator.py | 2 -- .../task_sdk/coordinators/test_subprocess.py | 2 -- .../execution_time/test_supervisor.py | 36 ------------------- .../execution_time/test_task_runner.py | 30 ---------------- 8 files changed, 2 insertions(+), 75 deletions(-) diff --git a/devel-common/src/tests_common/pytest_plugin.py b/devel-common/src/tests_common/pytest_plugin.py index 1d3bd04768788..eaac4e37fe01c 100644 --- a/devel-common/src/tests_common/pytest_plugin.py +++ b/devel-common/src/tests_common/pytest_plugin.py @@ -2674,9 +2674,7 @@ def _create_task_instance( try_number=try_number, 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"), diff --git a/providers/amazon/tests/unit/amazon/aws/executors/ecs/test_ecs_executor.py b/providers/amazon/tests/unit/amazon/aws/executors/ecs/test_ecs_executor.py index f350c88498124..ca6fab54fa2c3 100644 --- a/providers/amazon/tests/unit/amazon/aws/executors/ecs/test_ecs_executor.py +++ b/providers/amazon/tests/unit/amazon/aws/executors/ecs/test_ecs_executor.py @@ -1303,6 +1303,7 @@ def test_try_adopt_task_instances(self, mock_executor): task.pool_slots = 1 task.priority_weight = 1 task.context_carrier = {} + task.hostname = None task.queued_dttm = dt.datetime.now() # Set up nested attributes for BundleInfo task.dag_model = mock.Mock() diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index fc21104c50251..4ea82c4532748 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -2481,7 +2481,7 @@ def supervise_task( raise ValueError("dag_path is required") try: - coordinator = get_coordinator_manager().for_queue(ti.queue) + coordinator = get_coordinator_manager().for_queue(ti.queue or "default") except: log.exception( "Failed to initialize coordinator for task", diff --git a/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py index dc8081694975b..ec426a09ed741 100644 --- a/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py @@ -120,9 +120,7 @@ def _make_ti(dag_id: str = "tutorial_dag", queue: str = "executable") -> TaskIns run_id="run_1", try_number=1, map_index=-1, - pool_slots=1, queue=queue, - priority_weight=1, ) diff --git a/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py index 816135fdaf590..c30ecff35b9d2 100644 --- a/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py @@ -54,9 +54,7 @@ def _make_ti(dag_id: str = "test_dag", queue: str = "java") -> TaskInstance: run_id="run_1", try_number=1, map_index=-1, - pool_slots=1, queue=queue, - priority_weight=1, ) diff --git a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py index 34a1d0ec3d40f..8bdd905c6b059 100644 --- a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py +++ b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py @@ -60,9 +60,7 @@ def _make_ti(dag_id: str = "tutorial_dag", queue: str = "socket") -> TaskInstanc run_id="run_1", try_number=1, map_index=-1, - pool_slots=1, queue=queue, - priority_weight=1, ) diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py index eb126f0e81743..8ae91b8c15e8d 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -242,9 +242,7 @@ def test_supervise( run_id="d", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, ) bundle_info = BundleInfo(name="my-bundle", version=None) @@ -278,9 +276,7 @@ def test_on_kill_hook_called_when_supervisor_receives_sigterm( run_id="r", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, ) bundle_info = BundleInfo(name="my-bundle", version=None) @@ -386,9 +382,7 @@ def subprocess_main(): run_id="d", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, ), client=client_with_ti_start, target=subprocess_main, @@ -464,9 +458,7 @@ def subprocess_main(): run_id="d", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, ), client=client_with_ti_start, target=subprocess_main, @@ -564,9 +556,7 @@ def on_kill(self) -> None: run_id="d", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, ), client=make_client(transport=httpx.MockTransport(handle_request)), target=subprocess_main, @@ -596,9 +586,7 @@ def subprocess_main(): run_id="d", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, ), client=client_with_ti_start, target=subprocess_main, @@ -634,9 +622,7 @@ def subprocess_main(): run_id="d", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, ), client=mock_client, target=subprocess_main, @@ -681,9 +667,7 @@ def test_resume_start_date_from_context(self, mocker, make_ti_context, start_dat run_id="d", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, ), client=mock_client, target=lambda: None, @@ -727,9 +711,7 @@ def subprocess_main(): run_id="d", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, ), client=sdk_client.Client(base_url="", dry_run=True, token=""), target=subprocess_main, @@ -772,9 +754,7 @@ def _on_child_started(self, *args, **kwargs): run_id="d", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, ), client=sdk_client.Client(base_url="", dry_run=True, token=""), target=subprocess_main, @@ -796,9 +776,7 @@ def test_run_simple_dag(self, test_dags_dir, captured_logs, time_machine, mocker run_id="c", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, ) bundle_info = BundleInfo(name="my-bundle", version=None) @@ -840,9 +818,7 @@ def test_supervise_handles_deferred_task( run_id="d", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, ) # Create a mock client to assert calls to the client @@ -970,9 +946,7 @@ def handle_request(request: httpx.Request) -> httpx.Response: run_id="d", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, ), client=make_client(transport=httpx.MockTransport(handle_request)), target=subprocess_main, @@ -1056,9 +1030,7 @@ def subprocess_main(): run_id="d", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, ), client=make_client(transport=httpx.MockTransport(handle_request)), target=subprocess_main, @@ -1269,9 +1241,7 @@ def subprocess_main(): run_id="d", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, ), client=client_with_ti_start, target=subprocess_main, @@ -1432,9 +1402,7 @@ def _handler(sig, frame): run_id="d", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, ), client=client_with_ti_start, target=subprocess_main, @@ -3461,9 +3429,7 @@ def execute(self, context: Context): task_id=task.task_id, run_id="r", try_number=1, - pool_slots=1, queue="default", - priority_weight=1, ) # Patch the API client used by InProcessTestSupervisor to return a predictable TI context @@ -3887,9 +3853,7 @@ def subprocess_main(): run_id="d", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, ), client=client_with_ti_start, target=subprocess_main, diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py index 18b4da9018847..48b5139046963 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py +++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py @@ -214,9 +214,7 @@ def test_parse(test_dags_dir: Path, make_ti_context): run_id="c", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, ), dag_rel_path="super_basic.py", bundle_info=BundleInfo(name="my-bundle", version=None), @@ -269,9 +267,7 @@ def test_parse_dag_bag(mock_dagbag, test_dags_dir: Path, make_ti_context): run_id="c", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, ), dag_rel_path="super_basic.py", bundle_info=BundleInfo(name="my-bundle", version=None), @@ -332,9 +328,7 @@ def test_parse_not_found(test_dags_dir: Path, make_ti_context, dag_id, task_id, run_id="c", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, ), dag_rel_path="super_basic.py", bundle_info=BundleInfo(name="my-bundle", version=None), @@ -381,9 +375,7 @@ def test_parse_not_found_does_not_reschedule_when_max_attempts_reached(test_dags run_id="c", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, ), dag_rel_path="super_basic.py", bundle_info=BundleInfo(name="my-bundle", version=None), @@ -445,9 +437,7 @@ def test_main_sends_reschedule_task_when_startup_reschedules( run_id="test_run", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, context_carrier={}, ), dag_rel_path="", @@ -621,9 +611,7 @@ def test_task_span_is_child_of_dag_run_span(make_ti_context): run_id="test_run", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, context_carrier=ti_carrier, ), dag_rel_path="", @@ -692,9 +680,7 @@ def test_task_span_no_parent_when_no_context_carrier(make_ti_context): run_id="test_run", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, context_carrier=None, ), dag_rel_path="", @@ -736,9 +722,7 @@ def test_parse_module_in_bundle_root(tmp_path: Path, make_ti_context): run_id="c", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, ), dag_rel_path="path_test.py", bundle_info=BundleInfo(name="my-bundle", version=None), @@ -1186,9 +1170,7 @@ def test_basic_templated_dag(mocked_parse, make_ti_context, mock_supervisor_comm run_id="c", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, ), bundle_info=FAKE_BUNDLE, dag_rel_path="", @@ -1305,9 +1287,7 @@ def execute(self, context): run_id="c", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, ), dag_rel_path="", bundle_info=FAKE_BUNDLE, @@ -1356,9 +1336,7 @@ def execute(self, context): run_id="c", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, ), dag_rel_path="", bundle_info=FAKE_BUNDLE, @@ -1407,9 +1385,7 @@ def execute(self, context): run_id="c", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, ), dag_rel_path="", bundle_info=FAKE_BUNDLE, @@ -1450,9 +1426,7 @@ def execute(self, context): run_id="c", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, ), dag_rel_path="", bundle_info=FAKE_BUNDLE, @@ -1626,9 +1600,7 @@ def test_dag_parsing_context(make_ti_context, mock_supervisor_comms, monkeypatch run_id="c", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, ), dag_rel_path="dag_parsing_context.py", bundle_info=BundleInfo(name="my-bundle", version=None), @@ -4172,9 +4144,7 @@ def execute(self, context): run_id="c", try_number=1, dag_version_id=uuid7(), - pool_slots=1, queue="default", - priority_weight=1, ), dag_rel_path="", bundle_info=FAKE_BUNDLE,