Serialize mapped tasks and task groups

[ashb]

This is needed for the scheduler to be able to expand the task instances at runtime

[uranusjr]

I don’t get what this implies.

[ashb]

Delete the task_id from the kwargs (so it isn't stored in the mapped data structure) -- so del kwargs['task_id'] but without the try/except for when it's not there? I'll double check if I can just do this with a del.

[uranusjr]

I guess my question is more like why do we need to treat task_id specially (and why we didn’t need to previously, but only when we add operator mapping)

[ashb]

Task ID was getting stored in partial_kwargs which I hadn't noticed until I looked at the serialized representation of a MappedOperator.

[uranusjr]

How about storing task_id separately instead and not in partial_kwargs?

[ashb]

This is dealing with it in partial_kwargs --

            # Store the args passed to init -- we need them to support task.map serialzation!
            kwargs.pop('task_id', None)
            self._BaseOperator__init_kwargs.update(kwargs)  # type: ignore

(The only use of __init_kwargs is for building partial_kwargs when mapping a task)

[ashb]

This is the other option:

diff --git a/airflow/models/baseoperator.py b/airflow/models/baseoperator.py
index 752b8a6cf..71170320d 100644
--- a/airflow/models/baseoperator.py
+++ b/airflow/models/baseoperator.py
@@ -136,6 +136,7 @@ class BaseOperatorMeta(abc.ABCMeta):
         non_optional_args = {
             name for (name, param) in non_variadic_params.items() if param.default == param.empty
         }
+        non_optional_args -= {'task_id'}
 
         class autostacklevel_warn:
             def __init__(self):
@@ -158,7 +159,7 @@ class BaseOperatorMeta(abc.ABCMeta):
             func.__globals__['warnings'] = autostacklevel_warn()
 
         @functools.wraps(func)
-        def apply_defaults(self: "BaseOperator", *args: Any, **kwargs: Any) -> Any:
+        def apply_defaults(self: "BaseOperator", *args: Any, task_id: str, **kwargs: Any) -> Any:
             from airflow.models.dag import DagContext
             from airflow.utils.task_group import TaskGroupContext
 
@@ -201,15 +202,15 @@ class BaseOperatorMeta(abc.ABCMeta):
 
             hook = getattr(self, '_hook_apply_defaults', None)
             if hook:
-                args, kwargs = hook(**kwargs, default_args=default_args)
+                args, kwargs = hook(task_id=task_id, **kwargs, default_args=default_args)
+                task_id = kwargs.pop('task_id')
                 default_args = kwargs.pop('default_args', {})
 
             if not hasattr(self, '_BaseOperator__init_kwargs'):
                 self._BaseOperator__init_kwargs = {}
 
-            result = func(self, **kwargs, default_args=default_args)
+            result = func(self, **kwargs, task_id=task_id, default_args=default_args)
             # Store the args passed to init -- we need them to support task.map serialzation!
-            kwargs.pop('task_id', None)
             self._BaseOperator__init_kwargs.update(kwargs)  # type: ignore
 
             # Here we set upstream task defined by XComArgs passed to template fields of the operator

Which do you think is better @uranusjr?

[uranusjr]

The “other” one looks better to me, it feels like we are intentionally treating task_id different from all other arguments.

[ashb]

Done in 9439378 -- and a bit of a justification/reason in the commit message

[uranusjr]

Turns out this has an unfortunate side effect:

>>> FooOperator()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: apply_defaults() missing 1 required keyword-only argument: 'task_id'

and error message seems to be hard-coded in the interpreter and cannot be patched. I’m trying yet another approach (see new commits)

[kaxil]

Shouldn't partial_kwargs={} be a default on the MappedOperator itself?

[ashb]

No, cos in all other situations where this is constructor we need to pass a value.

[kaxil]

Same for these three, shouldn't we have a default on MappedOperator itself?

[ashb]

I'm now saying this is "good enough" -- we can make changes later if we want to

[uranusjr]

I’m assuming many of these defaults (this, _is_dummy, and template_fields, I believe) don’t need to consider when operator_class is a str because in that case these values would’ve been supplied explicitly instead. (It’s kind of bad it’s designed this way but I guess that can be said for many things regarding the current serialisation implementation...)

[ashb]

Yes, exactly that. I've already started thinking in my head about how to refactor/rearchitect the serialization and deserialization.

[github-actions]

The PR most likely needs to run full matrix of tests because it modifies parts of the core of Airflow. However, committers might decide to merge it quickly and take the risk. If they don't merge it quickly - please rebase it to the latest main at your convenience, or amend the last commit of the PR, and push it with --force-with-lease.

[uranusjr]

Seems to be working. If the latest two commits look good to you @ashb this is ready to go in.

[ashb]

Looks good -- thanks for picking up those fiixes TP.

[uranusjr]

I’ll start working on attribute parity between BaseOperator and MappedOperator.

[ashb]

I'm going to work on adding --map-index to task execution side.