Type TaskInstance.task to Operator and call unmap() when needed

airflow/cli/commands/task_command.py

@@ -113,9 +113,9 @@ def _get_ti(
 ) -> TaskInstance:
     """Get the task instance through DagRun.run_id, if that fails, get the TI the old way"""
     if task.is_mapped:
-        if map_index == -1:
+        if map_index < 0:
             raise RuntimeError("No map_index passed to mapped task")
-    elif map_index != -1:
+    elif map_index >= 0:
         raise RuntimeError("map_index passed to non-mapped task")
     dag_run = _get_dag_run(
         dag=task.dag,
@@ -224,8 +224,7 @@ def _run_task_by_local_task_job(args, ti):
 
 def _run_raw_task(args, ti: TaskInstance) -> None:
     """Runs the main task handling code"""
-    if ti.task.is_mapped:
-        ti.task = ti.task.unmap()
+    ti.task = ti.task.unmap()
     ti._run_raw_task(
         mark_success=args.mark_success,
         job_id=args.job_id,
@@ -531,6 +530,7 @@ def task_render(args):
     dag = get_dag(args.subdir, args.dag_id)
     task = dag.get_task(task_id=args.task_id)
     ti = _get_ti(task, args.execution_date_or_run_id, args.map_index, create_if_necessary=True)
+    ti.task = ti.task.unmap()
     ti.render_templates()
     for attr in task.__class__.template_fields:
         print(

airflow/dag_processing/processor.py

@@ -445,7 +445,7 @@ def manage_slas(self, dag: DAG, session: Session = None) -> None:
             blocking_tis: List[TI] = []
             for ti in fetched_tis:
                 if ti.task_id in dag.task_ids:
-                    ti.task = dag.get_task(ti.task_id)  # type: ignore[assignment]  # TODO: Fix task: Operator
+                    ti.task = dag.get_task(ti.task_id)
                     blocking_tis.append(ti)
                 else:
                     session.delete(ti)

airflow/decorators/base.py

@@ -43,7 +43,7 @@
 from airflow.compat.functools import cache, cached_property
 from airflow.exceptions import AirflowException
 from airflow.models.abstractoperator import DEFAULT_RETRIES, DEFAULT_RETRY_DELAY
-from airflow.models.baseoperator import BaseOperator, coerce_retry_delay, parse_retries
+from airflow.models.baseoperator import BaseOperator, coerce_resources, coerce_retry_delay, parse_retries
 from airflow.models.dag import DAG, DagContext
 from airflow.models.mappedoperator import MappedOperator
 from airflow.models.pool import Pool
@@ -308,6 +308,7 @@ def map(self, **kwargs) -> XComArg:
         partial_kwargs["retry_delay"] = coerce_retry_delay(
             partial_kwargs.get("retry_delay", DEFAULT_RETRY_DELAY),
         )
+        partial_kwargs["resources"] = coerce_resources(partial_kwargs.get("resources"))
         partial_kwargs.setdefault("executor_config", {})
 
         # Mypy does not work well with a subclassed attrs class :(

airflow/executors/debug_executor.py

@@ -76,8 +76,7 @@ def _run_task(self, ti: TaskInstance) -> bool:
         key = ti.key
         try:
             params = self.tasks_params.pop(ti.key, {})
-            if ti.task.is_mapped:
-                ti.task = ti.task.unmap()
+            ti.task = ti.task.unmap()
             ti._run_raw_task(job_id=ti.job_id, **params)
             self.change_state(key, State.SUCCESS)
             ti._run_finished_callback()

airflow/jobs/local_task_job.py

@@ -107,8 +107,7 @@ def signal_handler(signum, frame):
             # Unmap the task _after_ it has forked/execed. (This is a bit of a kludge, but if we unmap before
             # fork, then the "run_raw_task" command will see the mapping index and an Non-mapped task and
             # fail)
-            if self.task_instance.task.is_mapped:
-                self.task_instance.task = self.task_instance.task.unmap()
+            self.task_instance.task = self.task_instance.task.unmap()
 
             heartbeat_time_limit = conf.getint('scheduler', 'scheduler_zombie_task_threshold')
 

airflow/models/baseoperator.py

@@ -48,7 +48,6 @@
 )
 
 import attr
-import jinja2
 import pendulum
 from dateutil.relativedelta import relativedelta
 from sqlalchemy.orm import Session
@@ -91,6 +90,8 @@
 from airflow.utils.weight_rule import WeightRule
 
 if TYPE_CHECKING:
+    import jinja2  # Slow import.
+
     from airflow.models.dag import DAG
     from airflow.utils.task_group import TaskGroup
 
@@ -122,6 +123,12 @@ def coerce_retry_delay(retry_delay: Union[float, timedelta]) -> timedelta:
     return timedelta(seconds=retry_delay)
 
 
+def coerce_resources(resources: Optional[Dict[str, Any]]) -> Optional[Resources]:
+    if resources is None:
+        return None
+    return Resources(**resources)
+
+
 def _get_dag_defaults(dag: Optional["DAG"], task_group: Optional["TaskGroup"]) -> Tuple[dict, ParamsDict]:
     if not dag:
         return {}, ParamsDict()
@@ -175,6 +182,7 @@ def partial(
     owner: str = DEFAULT_OWNER,
     email: Union[None, str, Iterable[str]] = None,
     params: Optional[dict] = None,
+    resources: Optional[Dict[str, Any]] = None,
     trigger_rule: str = DEFAULT_TRIGGER_RULE,
     depends_on_past: bool = False,
     wait_for_downstream: bool = False,
@@ -251,6 +259,7 @@ def partial(
     partial_kwargs.setdefault("executor_config", executor_config)
     partial_kwargs.setdefault("inlets", inlets)
     partial_kwargs.setdefault("outlets", outlets)
+    partial_kwargs.setdefault("resources", resources)
 
     # Post-process arguments. Should be kept in sync with _TaskDecorator.map().
     if "task_concurrency" in kwargs:  # Reject deprecated option.
@@ -264,6 +273,7 @@ def partial(
     partial_kwargs["retries"] = parse_retries(partial_kwargs["retries"])
     partial_kwargs["retry_delay"] = coerce_retry_delay(partial_kwargs["retry_delay"])
     partial_kwargs["executor_config"] = partial_kwargs["executor_config"] or {}
+    partial_kwargs["resources"] = coerce_resources(partial_kwargs["resources"])
 
     return OperatorPartial(operator_class=operator_class, kwargs=partial_kwargs)
 
@@ -696,7 +706,7 @@ def __init__(
         pre_execute: Optional[TaskPreExecuteHook] = None,
         post_execute: Optional[TaskPostExecuteHook] = None,
         trigger_rule: str = DEFAULT_TRIGGER_RULE,
-        resources: Optional[Dict] = None,
+        resources: Optional[Dict[str, Any]] = None,
         run_as_user: Optional[str] = None,
         task_concurrency: Optional[int] = None,
         max_active_tis_per_dag: Optional[int] = None,
@@ -828,7 +838,7 @@ def __init__(
                 f"received '{weight_rule}'."
             )
         self.weight_rule = weight_rule
-        self.resources: Optional[Resources] = Resources(**resources) if resources else None
+        self.resources = coerce_resources(resources)
         if task_concurrency and not max_active_tis_per_dag:
             # TODO: Remove in Airflow 3.0
             warnings.warn(
@@ -1140,7 +1150,7 @@ def __setstate__(self, state):
         self._log = logging.getLogger("airflow.task.operators")
 
     def render_template_fields(
-        self, context: Context, jinja_env: Optional[jinja2.Environment] = None
+        self, context: Context, jinja_env: Optional["jinja2.Environment"] = None
     ) -> None:
         """
         Template all attributes listed in template_fields. Note this operation is irreversible.
@@ -1158,7 +1168,7 @@ def _do_render_template_fields(
         parent: Any,
         template_fields: Iterable[str],
         context: Context,
-        jinja_env: jinja2.Environment,
+        jinja_env: "jinja2.Environment",
         seen_oids: Set,
     ) -> None:
         for attr_name in template_fields:
@@ -1178,7 +1188,7 @@ def render_template(
         self,
         content: Any,
         context: Context,
-        jinja_env: Optional[jinja2.Environment] = None,
+        jinja_env: Optional["jinja2.Environment"] = None,
         seen_oids: Optional[Set] = None,
     ) -> Any:
         """
@@ -1236,7 +1246,7 @@ def render_template(
             return content
 
     def _render_nested_template_fields(
-        self, content: Any, context: Context, jinja_env: jinja2.Environment, seen_oids: Set
+        self, content: Any, context: Context, jinja_env: "jinja2.Environment", seen_oids: Set
     ) -> None:
         if id(content) not in seen_oids:
             seen_oids.add(id(content))
@@ -1553,8 +1563,7 @@ def defer(
 
     def unmap(self) -> "BaseOperator":
         """:meta private:"""
-        # Exists to make typing easier
-        raise TypeError("Internal code error: Do not call unmap on BaseOperator!")
+        return self
 
 
 # TODO: Deprecate for Airflow 3.0

airflow/models/dagrun.py

@@ -624,7 +624,7 @@ def task_instance_scheduling_decisions(self, session: Session = NEW_SESSION) ->
         dag = self.get_dag()
         for ti in tis:
             try:
-                ti.task = dag.get_task(ti.task_id)  # type: ignore[assignment]  # TODO: Fix task: Operator
+                ti.task = dag.get_task(ti.task_id)
             except TaskNotFound:
                 self.log.warning(
                     "Failed to get task '%s' for dag '%s'. Marking it as removed.", ti, ti.dag_id

airflow/models/mappedoperator.py

@@ -59,6 +59,7 @@
 from airflow.serialization.enums import DagAttributeTypes
 from airflow.ti_deps.deps.base_ti_dep import BaseTIDep
 from airflow.ti_deps.deps.mapped_task_expanded import MappedTaskIsExpanded
+from airflow.utils.operator_resources import Resources
 from airflow.utils.session import NEW_SESSION
 from airflow.utils.state import State, TaskInstanceState
 from airflow.utils.task_group import TaskGroup
@@ -339,6 +340,10 @@ def sla(self) -> Optional[datetime.timedelta]:
     def max_active_tis_per_dag(self) -> Optional[int]:
         return self.partial_kwargs.get("max_active_tis_per_dag")
 
+    @property
+    def resources(self) -> Optional[Resources]:
+        return self.partial_kwargs.get("resources")
+
     @property
     def on_execute_callback(self) -> Optional[TaskStateChangeCallback]:
         return self.partial_kwargs.get("on_execute_callback")

airflow/models/taskinstance.py

@@ -122,7 +122,6 @@
 
 
 if TYPE_CHECKING:
-    from airflow.models.baseoperator import BaseOperator
     from airflow.models.dag import DAG, DagModel
     from airflow.models.dagrun import DagRun
     from airflow.models.operator import Operator
@@ -423,11 +422,7 @@ class TaskInstance(Base, LoggingMixin):
 
     execution_date = association_proxy("dag_run", "execution_date")
 
-    # TODO: Investigate TaskInstance's lifecycle and call stack to determine
-    # when self.task is guaranteed to be unmapped, and when it may be mapped.
-    # Add 'assert isinstance(self.task, BaseOperator) to methods accordingly,
-    # and change this to 'task: Operator' instead.
-    task: "BaseOperator"  # Not always set...
+    task: "Operator"  # Not always set...
 
     def __init__(
         self,
@@ -806,7 +801,7 @@ def refresh_from_task(self, task: "Operator", pool_override=None):
         :param task: The task object to copy from
         :param pool_override: Use the pool_override instead of task's pool
         """
-        self.task = task  # type: ignore[assignment]  # TODO: Fix task: Operator
+        self.task = task
         self.queue = task.queue
         self.pool = pool_override or task.pool
         self.pool_slots = task.pool_slots
@@ -1321,8 +1316,9 @@ def _run_raw_task(
                 f'task property of {self.task_id!r} was still a MappedOperator -- it should have been '
                 'expanded already!'
             )
+        task = self.task.unmap()
         self.test_mode = test_mode
-        self.refresh_from_task(self.task, pool_override=pool)
+        self.refresh_from_task(task, pool_override=pool)
         self.refresh_from_db(session=session)
         self.job_id = job_id
         self.hostname = get_hostname()
@@ -1334,7 +1330,7 @@ def _run_raw_task(
         Stats.incr(f'ti.start.{self.task.dag_id}.{self.task.task_id}')
         try:
             if not mark_success:
-                self.task = self.task.prepare_for_execution()
+                self.task = task.prepare_for_execution()
                 context = self.get_template_context(ignore_param_exceptions=False)
                 self._execute_task_with_callbacks(context)
             if not test_mode:
@@ -1397,7 +1393,7 @@ def _run_raw_task(
             session.commit()
             raise
         finally:
-            Stats.incr(f'ti.finish.{self.task.dag_id}.{self.task.task_id}.{self.state}')
+            Stats.incr(f'ti.finish.{task.dag_id}.{task.task_id}.{self.state}')
 
         # Recording SKIPPED or SUCCESS
         self.clear_next_method_args()
@@ -1660,8 +1656,7 @@ def run(
 
     def dry_run(self):
         """Only Renders Templates for the TI"""
-        task = self.task
-        task_copy = task.prepare_for_execution()
+        task_copy = self.task.unmap().prepare_for_execution()
         self.task = task_copy
 
         self.render_templates()
@@ -1742,9 +1737,7 @@ def handle_failure(
         if not test_mode:
             self.refresh_from_db(session)
 
-        task = self.task
-        if task.is_mapped:
-            task = task.unmap()
+        task = self.task.unmap()
         self.end_date = timezone.utcnow()
         self.set_duration()
         Stats.incr(f'operator_failures_{task.task_type}', 1, 1)
@@ -2032,10 +2025,14 @@ def overwrite_params_with_dag_run_conf(self, params, dag_run):
 
     def render_templates(self, context: Optional[Context] = None) -> None:
         """Render templates in the operator fields."""
+        if self.task.is_mapped:
+            raise RuntimeError(
+                f'task property of {self.task_id!r} was still a MappedOperator -- it should have been '
+                'expanded already!'
+            )
         if not context:
             context = self.get_template_context()
-
-        self.task.render_template_fields(context)
+        self.task.unmap().render_template_fields(context)
 
     def render_k8s_pod_yaml(self) -> Optional[dict]:
         """Render k8s pod yaml"""

airflow/providers/elasticsearch/log/es_task_handler.py

@@ -104,8 +104,10 @@ def __init__(
 
     def _render_log_id(self, ti: TaskInstance, try_number: int) -> str:
         dag_run = ti.get_dagrun()
+        dag = ti.task.dag
+        assert dag is not None  # For Mypy.
         try:
-            data_interval: Tuple[datetime, datetime] = ti.task.dag.get_run_data_interval(dag_run)
+            data_interval: Tuple[datetime, datetime] = dag.get_run_data_interval(dag_run)
         except AttributeError:  # ti.task is not always set.
             data_interval = (dag_run.data_interval_start, dag_run.data_interval_end)
 

airflow/utils/log/file_task_handler.py

@@ -84,8 +84,10 @@ def _render_filename(self, ti: "TaskInstance", try_number: int) -> str:
             return render_template_to_string(self.filename_jinja_template, context)
         elif self.filename_template:
             dag_run = ti.get_dagrun()
+            dag = ti.task.dag
+            assert dag is not None  # For Mypy.
             try:
-                data_interval: Tuple[datetime, datetime] = ti.task.dag.get_run_data_interval(dag_run)
+                data_interval: Tuple[datetime, datetime] = dag.get_run_data_interval(dag_run)
             except AttributeError:  # ti.task is not always set.
                 data_interval = (dag_run.data_interval_start, dag_run.data_interval_end)
             if data_interval[0]:

airflow/www/views.py

@@ -1205,9 +1205,9 @@ def rendered_templates(self, session):
         root = request.args.get('root', '')
 
         logging.info("Retrieving rendered templates.")
-        dag = current_app.dag_bag.get_dag(dag_id)
+        dag: DAG = current_app.dag_bag.get_dag(dag_id)
         dag_run = dag.get_dagrun(execution_date=dttm, session=session)
-        task = copy.copy(dag.get_task(task_id))
+        task = copy.copy(dag.get_task(task_id).unmap())
 
         if dag_run is None:
             # No DAG run matching given logical date. This usually means this