Repository: apache/airflow

Issue:
AIP-45 breaks follow-on mini scheduler for mapped tasks
I've just noticed that this causes a problem for the follow-on mini scheduler for mapped tasks. I guess that code path wasn't sufficiently unit tested.

DAG

```python
import csv
import io
import os
import json
from datetime import datetime

from airflow import DAG
from airflow.decorators import task
from airflow.models.xcom_arg import XComArg
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
from airflow.providers.amazon.aws.operators.s3 import S3ListOperator


with DAG(dag_id='mapped_s3', start_date=datetime(2022, 5, 19)) as dag:
    files = S3ListOperator(
        task_id="get_inputs",
        bucket="airflow-summit-2022",
        prefix="data_provider_a/{{ data_interval_end | ds }}/",
        delimiter='/',
        do_xcom_push=True,
    )

    @task
    def csv_to_json(aws_conn_id, input_bucket, key, output_bucket):
        hook = S3Hook(aws_conn_id=aws_conn_id)

        csv_data = hook.read_key(key, input_bucket)
        reader = csv.DictReader(io.StringIO(csv_data))

        output = io.BytesIO()

        for row in reader:
            output.write(json.dumps(row, indent=None).encode('utf-8'))
            output.write(b"\n")

        output.seek(0, os.SEEK_SET)
        hook.load_file_obj(output, key=key.replace(".csv", ".json"), bucket_name=output_bucket)

    csv_to_json.partial(
        aws_conn_id="aws_default", input_bucket=files.bucket, output_bucket="airflow-summit-2022-processed"
    ).expand(key=XComArg(files))
```

Error:

```
  File "/home/ash/code/airflow/airflow/airflow/jobs/local_task_job.py", line 253, in _run_mini_scheduler_on_child_tasks
    info = dag_run.task_instance_scheduling_decisions(session)
  File "/home/ash/code/airflow/airflow/airflow/utils/session.py", line 68, in wrapper
    return func(*args, **kwargs)
  File "/home/ash/code/airflow/airflow/airflow/models/dagrun.py", line 658, in task_instance_scheduling_decisions
    schedulable_tis, changed_tis, expansion_happened = self._get_ready_tis(
  File "/home/ash/code/airflow/airflow/airflow/models/dagrun.py", line 714, in _get_ready_tis
    expanded_tis, _ = schedulable.task.expand_mapped_task(self.run_id, session=session)
  File "/home/ash/code/airflow/airflow/airflow/models/mappedoperator.py", line 609, in expand_mapped_task
    operator.mul, self._resolve_map_lengths(run_id, session=session).values()
  File "/home/ash/code/airflow/airflow/airflow/models/mappedoperator.py", line 591, in _resolve_map_lengths
    expansion_kwargs = self._get_expansion_kwargs()
  File "/home/ash/code/airflow/airflow/airflow/models/mappedoperator.py", line 526, in _get_expansion_kwargs
    return getattr(self, self._expansion_kwargs_attr)
AttributeError: 'MappedOperator' object has no attribute 'mapped_op_kwargs'
```

_Originally posted by @ashb in https://github.com/apache/airflow/issues/21877#issuecomment-1133409500_


Discussion:
FYI @pingzh 

I can look at this later, right now I've got to write my talk for the Summit :) 
> FYI @pingzh
> 
> I can look at this later, right now I've got to write my talk for the Summit :)

thanks @ashb . i can also take a look. looking forward to your talk.

--- 

it is strange that the scheduler can scheduler this task, it means `dagrun.task_instance_scheduling_decisions`  is called based on the deserialized dag, but it failed in the `_run_mini_scheduler_on_child_tasks`
update: 

it only happens to the new DagRun. If you try to clear an existing DagRun, it does not throw that error
That might be because the task is already expanded?.

I've not looked at this at all - just noticed main was broken but (thankfully for my talk) 2.3 is fine
> I've not looked at this at all - just noticed main was broken but (thankfully for my talk) 2.3 is fine

Your talk was main reason why we waited with merging it until we branched off to 2.3 :). 
@ashb do you mind taking a look at this? thanks
@pingzh Yeah, I or @uranusjr will take a look before we start closing out 2.4
The most immediate problem is we’re trying to expand a _deserialised_ mapped operator. This does not work; only a “real” mapped operator can be properly expanded. We should fix this, _and_ figure out a way to prevent this from happening again.

Resolving patch:
diff --git a/airflow/models/mappedoperator.py b/airflow/models/mappedoperator.py
--- a/airflow/models/mappedoperator.py
+++ b/airflow/models/mappedoperator.py
@@ -240,7 +240,17 @@ def _expand(self, **mapped_kwargs: "Mappable") -> "MappedOperator":
         return op
 
 
-@attr.define(kw_only=True)
+@attr.define(
+    kw_only=True,
+    # Disable custom __getstate__ and __setstate__ generation since it interacts
+    # badly with Airflow's DAG serialization and pickling. When a mapped task is
+    # deserialized, subclasses are coerced into MappedOperator, but when it goes
+    # through DAG pickling, all attributes defined in the subclasses are dropped
+    # by attrs's custom state management. Since attrs does not do anything too
+    # special here (the logic is only important for slots=True), we use Python's
+    # built-in implementation, which works (as proven by good old BaseOperator).
+    getstate_setstate=False,
+)
 class MappedOperator(AbstractOperator):
     """Object representing a mapped operator in a DAG."""
 


Test patch:
diff --git a/tests/serialization/test_dag_serialization.py b/tests/serialization/test_dag_serialization.py
--- a/tests/serialization/test_dag_serialization.py
+++ b/tests/serialization/test_dag_serialization.py
@@ -24,6 +24,7 @@
 import json
 import multiprocessing
 import os
+import pickle
 from datetime import datetime, timedelta
 from glob import glob
 from unittest import mock
@@ -1862,6 +1863,21 @@ def x(arg1, arg2, arg3):
         "retry_delay": timedelta(seconds=30),
     }
 
+    # Ensure the serialized operator can also be correctly pickled, to ensure
+    # correct interaction between DAG pickling and serialization. This is done
+    # here so we don't need to duplicate tests between pickled and non-pickled
+    # DAGs everywhere else.
+    pickled = pickle.loads(pickle.dumps(deserialized))
+    assert pickled.mapped_op_kwargs == {
+        "arg2": {"a": 1, "b": 2},
+        "arg3": _XComRef("op1", XCOM_RETURN_KEY),
+    }
+    assert pickled.partial_kwargs == {
+        "op_args": [],
+        "op_kwargs": {"arg1": [1, 2, {"a": "b"}]},
+        "retry_delay": timedelta(seconds=30),
+    }
+
 
 @pytest.mark.filterwarnings("ignore::DeprecationWarning")
 @pytest.mark.parametrize(