Deprecate and move airflow.utils.task_group to SDK (#53450)

ephraimbuddy · web-flow · commit b58f5d30dcd2 · 2025-07-19T14:23:27.000+02:00
* Deprecate and move `airflow.utils.task_group` to SDK

Some part of the module is already moved to SDK. This completes the
move

* fixup! Deprecate and move `airflow.utils.task_group` to SDK

* fixup! fixup! Deprecate and move `airflow.utils.task_group` to SDK
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/grid.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/grid.py
@@ -52,7 +52,7 @@
 from airflow.models.dagrun import DagRun
 from airflow.models.serialized_dag import SerializedDagModel
 from airflow.models.taskinstance import TaskInstance
-from airflow.utils.task_group import (
+from airflow.sdk.definitions.taskgroup import (
     get_task_group_children_getter,
     task_group_to_dict_grid,
 )
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/structure.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/structure.py
@@ -32,8 +32,8 @@
 )
 from airflow.models.dag_version import DagVersion
 from airflow.models.serialized_dag import SerializedDagModel
+from airflow.sdk.definitions.taskgroup import task_group_to_dict
 from airflow.utils.dag_edges import dag_edges
-from airflow.utils.task_group import task_group_to_dict
 
 structure_router = AirflowRouter(tags=["Structure"], prefix="/structure")
 
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/services/ui/grid.py b/airflow-core/src/airflow/api_fastapi/core_api/services/ui/grid.py
@@ -25,9 +25,8 @@
 from airflow.api_fastapi.common.parameters import state_priority
 from airflow.models.taskmap import TaskMap
 from airflow.sdk.definitions.mappedoperator import MappedOperator
-from airflow.sdk.definitions.taskgroup import MappedTaskGroup, TaskGroup
+from airflow.sdk.definitions.taskgroup import MappedTaskGroup, TaskGroup, get_task_group_children_getter
 from airflow.serialization.serialized_objects import SerializedBaseOperator
-from airflow.utils.task_group import get_task_group_children_getter
 
 log = structlog.get_logger(logger_name=__name__)
 
diff --git a/airflow-core/src/airflow/example_dags/example_setup_teardown.py b/airflow-core/src/airflow/example_dags/example_setup_teardown.py
@@ -23,7 +23,7 @@
 
 from airflow.providers.standard.operators.bash import BashOperator
 from airflow.sdk import DAG
-from airflow.utils.task_group import TaskGroup
+from airflow.sdk.definitions.taskgroup import TaskGroup
 
 with DAG(
     dag_id="example_setup_teardown",
diff --git a/airflow-core/src/airflow/example_dags/example_task_group.py b/airflow-core/src/airflow/example_dags/example_task_group.py
@@ -24,7 +24,7 @@
 from airflow.providers.standard.operators.bash import BashOperator
 from airflow.providers.standard.operators.empty import EmptyOperator
 from airflow.sdk import DAG
-from airflow.utils.task_group import TaskGroup
+from airflow.sdk.definitions.taskgroup import TaskGroup
 
 # [START howto_task_group]
 with DAG(
diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py
@@ -124,11 +124,10 @@
     from airflow.sdk.definitions.asset import AssetNameRef, AssetUniqueKey, AssetUriRef
     from airflow.sdk.definitions.dag import DAG
     from airflow.sdk.definitions.mappedoperator import MappedOperator
-    from airflow.sdk.definitions.taskgroup import MappedTaskGroup
+    from airflow.sdk.definitions.taskgroup import MappedTaskGroup, TaskGroup
     from airflow.sdk.types import RuntimeTaskInstanceProtocol
     from airflow.serialization.serialized_objects import SerializedBaseOperator
     from airflow.utils.context import Context
-    from airflow.utils.task_group import TaskGroup
 
     Operator: TypeAlias = BaseOperator | MappedOperator
 
diff --git a/airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py b/airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py
@@ -26,9 +26,9 @@
 from sqlalchemy import and_, func, or_, select
 
 from airflow.models.taskinstance import PAST_DEPENDS_MET
+from airflow.sdk.definitions.taskgroup import MappedTaskGroup
 from airflow.ti_deps.deps.base_ti_dep import BaseTIDep
 from airflow.utils.state import TaskInstanceState
-from airflow.utils.task_group import MappedTaskGroup
 from airflow.utils.trigger_rule import TriggerRule as TR
 
 if TYPE_CHECKING:
diff --git a/airflow-core/src/airflow/utils/__init__.py b/airflow-core/src/airflow/utils/__init__.py
@@ -28,6 +28,12 @@
     "xcom": {
         "XCOM_RETURN_KEY": "airflow.models.xcom.XCOM_RETURN_KEY",
     },
+    "task_group": {
+        "TaskGroup": "airflow.sdk.definitions.taskgroup.TaskGroup",
+        "MappedTaskGroup": "airflow.sdk.definitions.taskgroup.MappedOperator",
+        "get_task_group_children_getter": "airflow.sdk.definitions.taskgroup.get_task_group_children_getter",
+        "task_group_to_dict": "airflow.sdk.definitions.taskgroup.task_group_to_dict",
+    },
 }
 
 add_deprecated_classes(__deprecated_classes, __name__)
diff --git a/airflow-core/src/airflow/utils/dot_renderer.py b/airflow-core/src/airflow/utils/dot_renderer.py
@@ -26,10 +26,10 @@
 from airflow.exceptions import AirflowException
 from airflow.sdk import BaseOperator
 from airflow.sdk.definitions.mappedoperator import MappedOperator
+from airflow.sdk.definitions.taskgroup import TaskGroup
 from airflow.serialization.serialized_objects import SerializedBaseOperator
 from airflow.utils.dag_edges import dag_edges
 from airflow.utils.state import State
-from airflow.utils.task_group import TaskGroup
 
 if TYPE_CHECKING:
     import graphviz
diff --git a/airflow-core/src/airflow/utils/task_group.py b/airflow-core/src/airflow/utils/task_group.py
diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_grid.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_grid.py
@@ -29,10 +29,10 @@
 from airflow.models.taskinstance import TaskInstance
 from airflow.providers.standard.operators.empty import EmptyOperator
 from airflow.sdk import task_group
+from airflow.sdk.definitions.taskgroup import TaskGroup
 from airflow.utils import timezone
 from airflow.utils.session import provide_session
 from airflow.utils.state import DagRunState, TaskInstanceState
-from airflow.utils.task_group import TaskGroup
 from airflow.utils.types import DagRunTriggeredByType, DagRunType
 
 from tests_common.test_utils.db import clear_db_assets, clear_db_dags, clear_db_runs, clear_db_serialized_dags
diff --git a/airflow-core/tests/unit/models/test_dagrun.py b/airflow-core/tests/unit/models/test_dagrun.py
@@ -2774,7 +2774,7 @@ def test_teardown_and_fail_fast(dag_maker):
     in this case, the second teardown skips because its setup skips.
     """
     from airflow.sdk import task as task_decorator
-    from airflow.utils.task_group import TaskGroup
+    from airflow.sdk.definitions.taskgroup import TaskGroup
 
     with dag_maker(fail_fast=True) as dag:
         for num in (1, 2):
diff --git a/airflow-core/tests/unit/models/test_mappedoperator.py b/airflow-core/tests/unit/models/test_mappedoperator.py
@@ -34,8 +34,8 @@
 from airflow.models.taskmap import TaskMap
 from airflow.providers.standard.operators.python import PythonOperator
 from airflow.sdk import setup, task, task_group, teardown
+from airflow.sdk.definitions.taskgroup import TaskGroup
 from airflow.utils.state import TaskInstanceState
-from airflow.utils.task_group import TaskGroup
 from airflow.utils.trigger_rule import TriggerRule
 
 from tests_common.test_utils.mapping import expand_mapped_task
diff --git a/airflow-core/tests/unit/models/test_taskinstance.py b/airflow-core/tests/unit/models/test_taskinstance.py
@@ -67,6 +67,7 @@
 from airflow.sdk.api.datamodels._generated import AssetEventResponse, AssetResponse
 from airflow.sdk.definitions.asset import Asset, AssetAlias
 from airflow.sdk.definitions.param import process_params
+from airflow.sdk.definitions.taskgroup import TaskGroup
 from airflow.sdk.execution_time.comms import (
     AssetEventsResult,
 )
@@ -83,7 +84,6 @@
 from airflow.utils.session import create_session, provide_session
 from airflow.utils.span_status import SpanStatus
 from airflow.utils.state import DagRunState, State, TaskInstanceState
-from airflow.utils.task_group import TaskGroup
 from airflow.utils.types import DagRunTriggeredByType, DagRunType
 
 from tests_common.test_utils import db
diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py b/airflow-core/tests/unit/serialization/test_dag_serialization.py
@@ -68,6 +68,7 @@
 from airflow.sdk.definitions._internal.expandinput import EXPAND_INPUT_EMPTY
 from airflow.sdk.definitions.asset import Asset, AssetUniqueKey
 from airflow.sdk.definitions.param import Param, ParamsDict
+from airflow.sdk.definitions.taskgroup import TaskGroup
 from airflow.security import permissions
 from airflow.serialization.enums import Encoding
 from airflow.serialization.json_schema import load_dag_schema_dict
@@ -84,7 +85,6 @@
 from airflow.utils import timezone
 from airflow.utils.module_loading import qualname
 from airflow.utils.operator_resources import Resources
-from airflow.utils.task_group import TaskGroup
 
 from tests_common.test_utils.config import conf_vars
 from tests_common.test_utils.markers import skip_if_force_lowest_dependencies_marker
diff --git a/airflow-core/tests/unit/serialization/test_serialized_objects.py b/airflow-core/tests/unit/serialization/test_serialized_objects.py
@@ -59,6 +59,7 @@
 from airflow.sdk.definitions.deadline import DeadlineAlert, DeadlineAlertFields, DeadlineReference
 from airflow.sdk.definitions.decorators import task
 from airflow.sdk.definitions.param import Param
+from airflow.sdk.definitions.taskgroup import TaskGroup
 from airflow.sdk.execution_time.context import OutletEventAccessor, OutletEventAccessors
 from airflow.serialization.enums import DagAttributeTypes as DAT, Encoding
 from airflow.serialization.serialized_objects import BaseSerialization, LazyDeserializedDAG, SerializedDAG
@@ -68,7 +69,6 @@
 from airflow.utils.db import LazySelectSequence
 from airflow.utils.operator_resources import Resources
 from airflow.utils.state import DagRunState, State
-from airflow.utils.task_group import TaskGroup
 from airflow.utils.types import DagRunType
 
 from unit.models import DEFAULT_DATE
diff --git a/airflow-core/tests/unit/utils/test_dag_cycle.py b/airflow-core/tests/unit/utils/test_dag_cycle.py
@@ -22,8 +22,8 @@
 from airflow.models.dag import DAG
 from airflow.providers.standard.operators.empty import EmptyOperator
 from airflow.sdk import Label
+from airflow.sdk.definitions.taskgroup import TaskGroup
 from airflow.utils.dag_cycle_tester import check_cycle
-from airflow.utils.task_group import TaskGroup
 
 from unit.models import DEFAULT_DATE
 
diff --git a/airflow-core/tests/unit/utils/test_dot_renderer.py b/airflow-core/tests/unit/utils/test_dot_renderer.py
@@ -24,10 +24,10 @@
 
 from airflow.models.dag import DAG
 from airflow.providers.standard.operators.empty import EmptyOperator
+from airflow.sdk.definitions.taskgroup import TaskGroup
 from airflow.serialization.dag_dependency import DagDependency
 from airflow.utils import dot_renderer, timezone
 from airflow.utils.state import State
-from airflow.utils.task_group import TaskGroup
 
 from tests_common.test_utils.compat import BashOperator
 from tests_common.test_utils.db import clear_db_dags
diff --git a/airflow-core/tests/unit/utils/test_edgemodifier.py b/airflow-core/tests/unit/utils/test_edgemodifier.py
@@ -25,8 +25,8 @@
 from airflow.providers.standard.operators.empty import EmptyOperator
 from airflow.providers.standard.operators.python import PythonOperator
 from airflow.sdk import Label
+from airflow.sdk.definitions.taskgroup import TaskGroup
 from airflow.utils.dag_edges import dag_edges
-from airflow.utils.task_group import TaskGroup
 
 DEFAULT_ARGS = {
     "owner": "test",
diff --git a/airflow-core/tests/unit/utils/test_task_group.py b/airflow-core/tests/unit/utils/test_task_group.py
@@ -34,8 +34,8 @@
     task_group as task_group_decorator,
     teardown,
 )
+from airflow.sdk.definitions.taskgroup import TaskGroup, task_group_to_dict
 from airflow.utils.dag_edges import dag_edges
-from airflow.utils.task_group import TaskGroup, task_group_to_dict
 
 from tests_common.test_utils.compat import BashOperator, PythonOperator
 from unit.models import DEFAULT_DATE
diff --git a/providers/standard/tests/unit/standard/decorators/test_python.py b/providers/standard/tests/unit/standard/decorators/test_python.py
@@ -37,6 +37,7 @@
     from airflow.sdk.bases.decorator import DecoratedMappedOperator
     from airflow.sdk.definitions._internal.expandinput import DictOfListsExpandInput
     from airflow.sdk.definitions.mappedoperator import MappedOperator
+
 else:
     from airflow.decorators import setup, task as task_decorator, teardown
     from airflow.decorators.base import DecoratedMappedOperator  # type: ignore[no-redef]
@@ -45,7 +46,7 @@
     from airflow.models.expandinput import DictOfListsExpandInput
     from airflow.models.mappedoperator import MappedOperator
     from airflow.models.xcom_arg import XComArg
-    from airflow.utils.task_group import TaskGroup
+    from airflow.utils.task_group import TaskGroup  # type: ignore[no-redef]
 
 pytestmark = pytest.mark.db_test
 
diff --git a/providers/standard/tests/unit/standard/operators/test_branch_operator.py b/providers/standard/tests/unit/standard/operators/test_branch_operator.py
@@ -28,9 +28,14 @@
 from airflow.timetables.base import DataInterval
 from airflow.utils import timezone
 from airflow.utils.state import State
-from airflow.utils.task_group import TaskGroup
 from airflow.utils.types import DagRunType
 
+try:
+    from airflow.sdk.definitions.taskgroup import TaskGroup
+except ImportError:
+    # Fallback for Airflow < 3.1
+    from airflow.utils.task_group import TaskGroup  # type: ignore[no-redef]
+
 from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_1, AIRFLOW_V_3_0_PLUS
 
 if AIRFLOW_V_3_0_PLUS:
diff --git a/providers/standard/tests/unit/standard/sensors/test_external_task_sensor.py b/providers/standard/tests/unit/standard/sensors/test_external_task_sensor.py
diff --git a/providers/standard/tests/unit/standard/utils/test_sensor_helper.py b/providers/standard/tests/unit/standard/utils/test_sensor_helper.py
diff --git a/task-sdk/src/airflow/sdk/definitions/taskgroup.py b/task-sdk/src/airflow/sdk/definitions/taskgroup.py

Original file line number	Diff line number	Diff line change
`@@ -52,7 +52,7 @@`
`52`	`52`	`from airflow.models.dagrun import DagRun`
`53`	`53`	`from airflow.models.serialized_dag import SerializedDagModel`
`54`	`54`	`from airflow.models.taskinstance import TaskInstance`
`55`		`-from airflow.utils.task_group import (`
	`55`	`+from airflow.sdk.definitions.taskgroup import (`
`56`	`56`	`get_task_group_children_getter,`
`57`	`57`	`task_group_to_dict_grid,`
`58`	`58`	`)`
Original file line number	Diff line number	Diff line change
`@@ -32,8 +32,8 @@`
`32`	`32`	`)`
`33`	`33`	`from airflow.models.dag_version import DagVersion`
`34`	`34`	`from airflow.models.serialized_dag import SerializedDagModel`
	`35`	`+from airflow.sdk.definitions.taskgroup import task_group_to_dict`
`35`	`36`	`from airflow.utils.dag_edges import dag_edges`
`36`		`-from airflow.utils.task_group import task_group_to_dict`
`37`	`37`
`38`	`38`	`structure_router = AirflowRouter(tags=["Structure"], prefix="/structure")`
`39`	`39`