Remove shared workloads dependency and refactor TaskInstanceDTO usage

- Removed the shared workloads dependency from pyproject.toml and related files.
- Deleted the workloads directory and its references in the codebase.
- Refactored imports of TaskInstanceDTO to point to the new location in execution_time.workloads.task.
- Introduced new files for TaskInstanceDTO and its base class in the execution_time module.
- Updated tests to reflect the changes in TaskInstanceDTO imports.

Author: jason810496

.pre-commit-config.yaml

@@ -463,6 +463,12 @@ 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"

airflow-core/pyproject.toml

@@ -249,7 +249,6 @@ exclude = [
 "../shared/serialization/src/airflow_shared/serialization" = "src/airflow/_shared/serialization"
 "../shared/state/src/airflow_shared/state" = "src/airflow/_shared/state"
 "../shared/timezones/src/airflow_shared/timezones" = "src/airflow/_shared/timezones"
-"../shared/workloads/src/airflow_shared/workloads" = "src/airflow/_shared/workloads"
 "../shared/listeners/src/airflow_shared/listeners" = "src/airflow/_shared/listeners"
 "../shared/plugins_manager/src/airflow_shared/plugins_manager" = "src/airflow/_shared/plugins_manager"
 "../shared/providers_discovery/src/airflow_shared/providers_discovery" = "src/airflow/_shared/providers_discovery"
@@ -338,7 +337,6 @@ shared_distributions = [
     "apache-airflow-shared-serialization",
     "apache-airflow-shared-state",
     "apache-airflow-shared-timezones",
-    "apache-airflow-shared-workloads",
     "apache-airflow-shared-plugins-manager",
     "apache-airflow-shared-providers-discovery",
 ]

airflow-core/src/airflow/_shared/workloads

@@ -648,9 +648,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
 
             return supervise_task(
-                ti=workload.ti,
+                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/base_executor.py

@@ -22,9 +22,8 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Literal
 
-from pydantic import Field
+from pydantic import BaseModel, Field
 
-from airflow._shared.workloads import TaskInstanceDTO as _BaseTaskInstanceDTO
 from airflow.executors.workloads.base import BaseDagBundleWorkload, BundleInfo
 from airflow.utils.state import TaskInstanceState
 
@@ -34,12 +33,13 @@
     from airflow.models.taskinstancekey import TaskInstanceKey
 
 
-class TaskInstanceDTO(_BaseTaskInstanceDTO):
+class BaseTaskInstanceDTO(BaseModel):
     """
-    TaskInstanceDTO with executor-specific ``key`` property.
+    Base schema for TaskInstance with the minimal fields shared by Executors and the Task SDK.
 
-    Extends the shared :class:`~airflow._shared.workloads.TaskInstanceDTO`
-    to add the :attr:`key` property used by executors for workload tracking.
+    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.
     """
 
     id: uuid.UUID
@@ -54,11 +54,16 @@ class TaskInstanceDTO(_BaseTaskInstanceDTO):
     queue: str
     priority_weight: int
     executor_config: dict | None = Field(default=None, exclude=True)
-    external_executor_id: str | 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)
+
     # TODO: Task-SDK: Can we replace TaskInstanceKey with just the uuid across the codebase?
     @property
     def key(self) -> TaskInstanceKey:

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

@@ -2540,7 +2540,7 @@ def _create_task_instance(
         from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS
 
         if AIRFLOW_V_3_3_PLUS:
-            from airflow.sdk._shared.workloads import TaskInstanceDTO
+            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,

devel-common/src/tests_common/pytest_plugin.py

@@ -1259,12 +1259,8 @@ components:
       - queue
       - priority_weight
       title: TaskInstanceDTO
-      description: 'TaskInstanceDTO with executor-specific ``key`` property.
-
-
-        Extends the shared :class:`~airflow._shared.workloads.TaskInstanceDTO`
-
-        to add the :attr:`key` property used by executors for workload tracking.'
+      description: TaskInstanceDTO with executor-specific ``external_executor_id``
+        field and ``key`` property.
     TaskInstanceState:
       type: string
       enum:

providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml

@@ -29,8 +29,8 @@
 from airflow.sdk.execution_time.coordinator import BaseCoordinator
 
 if TYPE_CHECKING:
-    from airflow.sdk._shared.workloads import TaskInstanceDTO
     from airflow.sdk.api.datamodels._generated import BundleInfo
+    from airflow.sdk.execution_time.workloads.task import TaskInstanceDTO
 
 
 class JavaCoordinator(BaseCoordinator):

providers/sdk/java/src/airflow/providers/sdk/java/coordinator.py

@@ -25,7 +25,6 @@
 import pytest
 import yaml
 
-from airflow._shared.workloads import TaskInstanceDTO
 from airflow.providers.sdk.java.bundle_scanner import (
     MAIN_CLASS_MANIFEST_KEY,
     MANIFEST_PATH,
@@ -34,6 +33,7 @@
 )
 from airflow.providers.sdk.java.coordinator import JavaCoordinator
 from airflow.sdk.api.datamodels._generated import BundleInfo
+from airflow.sdk.execution_time.workloads.task import TaskInstanceDTO
 
 from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS
 

providers/sdk/java/tests/unit/sdk/java/test_coordinator.py

@@ -1335,7 +1335,6 @@ dev = [
     "apache-airflow-shared-state",
     "apache-airflow-shared-template-rendering",
     "apache-airflow-shared-timezones",
-    "apache-airflow-shared-workloads",
 ]
 
 # To build docs:
@@ -1504,7 +1503,6 @@ apache-airflow-shared-serialization = false
 apache-airflow-shared-state = false
 apache-airflow-shared-template-rendering = false
 apache-airflow-shared-timezones = false
-apache-airflow-shared-workloads = false
 apache-airflow-task-sdk = false
 apache-airflow-task-sdk-integration-tests = false
 apache-aurflow-docker-stack = false
@@ -1657,7 +1655,6 @@ apache-airflow-shared-serialization = false
 apache-airflow-shared-state = false
 apache-airflow-shared-template-rendering = false
 apache-airflow-shared-timezones = false
-apache-airflow-shared-workloads = false
 apache-airflow-task-sdk = false
 apache-airflow-task-sdk-integration-tests = false
 apache-aurflow-docker-stack = false
@@ -1703,7 +1700,6 @@ apache-airflow-shared-serialization = { workspace = true }
 apache-airflow-shared-state = { workspace = true }
 apache-airflow-shared-template-rendering = { workspace = true }
 apache-airflow-shared-timezones = { workspace = true }
-apache-airflow-shared-workloads = { workspace = true }
 # Automatically generated provider workspace items (update_airflow_pyproject_toml.py)
 apache-airflow-providers-airbyte = { workspace = true }
 apache-airflow-providers-akeyless = { workspace = true }
@@ -1842,7 +1838,6 @@ members = [
     "shared/state",
     "shared/template_rendering",
     "shared/timezones",
-    "shared/workloads",
     # Automatically generated provider workspace members (update_airflow_pyproject_toml.py)
     "providers/airbyte",
     "providers/akeyless",