Repository: apache/airflow

Issue:
Task preceeding PythonVirtualenvOperator fails: "cannot pickle 'module' object"
**Apache Airflow version**

13faa6912f7cd927737a1dc15630d3bbaf2f5d4d

**Environment**

- **Configuration**: Local Executor
- **OS** (e.g. from /etc/os-release): Mac OS 11.3
- **Kernel**: Darwin Kernel Version 20.4.0
- **Install tools**: `pip install -e .`

**The DAG**

```python
def callable():
    print("hi")

with DAG(dag_id="two_virtualenv") as dag:

    a = PythonOperator(
        task_id="a",
        python_callable=callable,
    )

    # b = PythonOperator(          # works
    b = PythonVirtualenvOperator(  # doesn't work
        task_id="b",
        python_callable=callable,
    )

    a >> b
```

**What happened**:
Failure somewhere between first task and second:
```
INFO - Marking task as SUCCESS. dag_id=two_virtualenv, task_id=a
ERROR - Failed to execute task: cannot pickle 'module' object.
Traceback (most recent call last):
  File "/Users/matt/src/airflow/airflow/executors/debug_executor.py", line 79, in _run_task
    ti._run_raw_task(job_id=ti.job_id, **params)  # pylint: disable=protected-access
  File "/Users/matt/src/airflow/airflow/utils/session.py", line 70, in wrapper
    return func(*args, session=session, **kwargs)
  File "/Users/matt/src/airflow/airflow/models/taskinstance.py", line 1201, in _run_raw_task
    self._run_mini_scheduler_on_child_tasks(session)
  File "/Users/matt/src/airflow/airflow/utils/session.py", line 67, in wrapper
    return func(*args, **kwargs)
  File "/Users/matt/src/airflow/airflow/models/taskinstance.py", line 1223, in _run_mini_scheduler_on_child_tasks
    partial_dag = self.task.dag.partial_subset(
  File "/Users/matt/src/airflow/airflow/models/dag.py", line 1490, in partial_subset
    dag.task_dict = {
  File "/Users/matt/src/airflow/airflow/models/dag.py", line 1491, in <dictcomp>
    t.task_id: copy.deepcopy(t, {id(t.dag): dag})  # type: ignore
  File "/usr/local/Cellar/python@3.9/3.9.4/Frameworks/Python.framework/Versions/3.9/lib/python3.9/copy.py", line 153, in deepcopy
    y = copier(memo)
  File "/Users/matt/src/airflow/airflow/models/baseoperator.py", line 961, in __deepcopy__
    setattr(result, k, copy.deepcopy(v, memo))  # noqa
  File "/usr/local/Cellar/python@3.9/3.9.4/Frameworks/Python.framework/Versions/3.9/lib/python3.9/copy.py", line 161, in deepcopy
    rv = reductor(4)
TypeError: cannot pickle 'module' object
ERROR - Task instance <TaskInstance: two_virtualenv.a 2021-05-11 00:00:00+00:00 [failed]> failed
```

**What you expected to happen**:

Both tasks say "hi" and succeed

**To Replicate**

The DAG and output above are shortened for brevity.  A more complete story: https://gist.github.com/MatrixManAtYrService/6b27378776470491eb20b60e01cfb675

Ran it like this:
```
 $ airflow dags test two_virtualenv $(date "+%Y-%m-%d")
```


Discussion:
I hate pickle :/

Turns out it's nothing really to do with pickle, just to do with trying to copy _any_ module object.

And the horrifying thing? This has been broken since 2.0.0.

Resolving patch:
diff --git a/airflow/models/dag.py b/airflow/models/dag.py
--- a/airflow/models/dag.py
+++ b/airflow/models/dag.py
@@ -1463,13 +1463,8 @@ def partial_subset(
         """
         # deep-copying self.task_dict and self._task_group takes a long time, and we don't want all
         # the tasks anyway, so we copy the tasks manually later
-        task_dict = self.task_dict
-        task_group = self._task_group
-        self.task_dict = {}
-        self._task_group = None  # type: ignore
-        dag = copy.deepcopy(self)
-        self.task_dict = task_dict
-        self._task_group = task_group
+        memo = {id(self.task_dict): None, id(self._task_group): None}
+        dag = copy.deepcopy(self, memo)  # type: ignore
 
         if isinstance(task_ids_or_regex, (str, RePatternType)):
             matched_tasks = [t for t in self.tasks if re.findall(task_ids_or_regex, t.task_id)]
diff --git a/airflow/operators/python.py b/airflow/operators/python.py
--- a/airflow/operators/python.py
+++ b/airflow/operators/python.py
@@ -432,6 +432,11 @@ def _read_result(self, filename):
                 )
                 raise
 
+    def __deepcopy__(self, memo):
+        # module objects can't be copied _at all__
+        memo[id(self.pickling_library)] = self.pickling_library
+        return super().__deepcopy__(memo)
+
 
 def get_current_context() -> Dict[str, Any]:
     """


Test patch:
diff --git a/tests/operators/test_python.py b/tests/operators/test_python.py
--- a/tests/operators/test_python.py
+++ b/tests/operators/test_python.py
@@ -1025,6 +1025,19 @@ def f(
 
         self._run_as_operator(f, use_dill=True, system_site_packages=False, requirements=None)
 
+    def test_deepcopy(self):
+        """Test that PythonVirtualenvOperator are deep-copyable."""
+
+        def f():
+            return 1
+
+        task = PythonVirtualenvOperator(
+            python_callable=f,
+            task_id='task',
+            dag=self.dag,
+        )
+        copy.deepcopy(task)
+
 
 DEFAULT_ARGS = {
     "owner": "test",