Deprecate and move airflow.utils.task_group to SDK

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,
 )

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")
 

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__)
 

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",

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(

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
 

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:

airflow-core/src/airflow/utils/__init__.py

@@ -46,5 +46,11 @@ def __getattr__(name: str):
     "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__)

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

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

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):

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

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

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

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

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
 

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

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",

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

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
 

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: