Repository: apache/airflow

Issue:
task_instance_mutation_hook not using my custom implementation
<!--

Welcome to Apache Airflow!  For a smooth issue process, try to answer the following questions.
Don't worry if they're not all applicable; just try to include what you can :-)

If you need to include code snippets or logs, please put them in fenced code
blocks.  If they're super-long, please use the details tag like
<details><summary>super-long log</summary> lots of stuff </details>

Please delete these comment blocks before submitting the issue.

-->

<!--

IMPORTANT!!!

PLEASE CHECK "SIMILAR TO X EXISTING ISSUES" OPTION IF VISIBLE
NEXT TO "SUBMIT NEW ISSUE" BUTTON!!!

PLEASE CHECK IF THIS ISSUE HAS BEEN REPORTED PREVIOUSLY USING SEARCH!!!

Please complete the next sections or the issue will be closed.
This questions are the first thing we need to know to understand the context.

-->

**Apache Airflow version**: v1.10.11


**Kubernetes version (if you are using kubernetes)** (use `kubectl version`): NO

**Environment**: airflow 1.10.11 running on ubuntu 18.04 (used pip for installation) 

- **Cloud provider or hardware configuration**: bare metal 
- **OS** (e.g. from /etc/os-release): ubuntu 18.04
- **Kernel** (e.g. `uname -a`): Linux dev1 4.15.0-108-generic 109-Ubuntu SMP Fri Jun 19 11:33:10 UTC 2020 x86_64 x86_64 x86_64 GNU/Linux
- **Install tools**: pip
- **Others**:

**What happened**:
I'm not able to get task_instance_mutation_hook (https://github.com/apache/airflow/pull/8852) to load my custom implementation from `$AIRFLOW_HOME/config/airflow_local_settings.py`.
<!-- (please include exact error messages if you can) -->

**What you expected to happen**: Airflow to use task_instance_mutation_hook provided in `$AIRFLOW_HOME/config/airflow_local_settings.py`.

<!-- What do you think went wrong? -->

**How to reproduce it**: make sure your running v1.10.11 and create a custom `$AIRFLOW_HOME/config/airflow_local_settings.py` where you define a dummy `task_instance_mutation_hook` as next
```python
def task_instance_mutation_hook(ti):
   # every new or up to retry task will be added to non existent queue, so no task will pass through
   if ti.state in [State.NONE, State.UP_FOR_RETRY]:
      ti.queue = 'dummy-queue'
```
if any of the airflow example dags run successfully then, the previous hook didn't work.
<!---

As minimally and precisely as possible. Keep in mind we do not have access to your cluster or dags.

If you are using kubernetes, please attempt to recreate the issue using minikube or kind.

## Install minikube/kind

- Minikube https://minikube.sigs.k8s.io/docs/start/
- Kind https://kind.sigs.k8s.io/docs/user/quick-start/

If this is a UI bug, please provide a screenshot of the bug or a link to a youtube video of the bug in action

You can include images using the .md sytle of
![alt text](http://url/to/img.png)

To record a screencast, mac users can use QuickTime and then create an unlisted youtube video with the resulting .mov file.

--->


**Anything else we need to know**: task_instance_mutation_hook is a new feature introduced in v1.10.11 

<!--

How often does this problem occur? Once? Every time etc?

Any relevant logs to include? Put them here in side a detail tag:
<details><summary>x.log</summary> lots of stuff </details>

-->



Discussion:

Resolving patch:
diff --git a/airflow/operators/python_operator.py b/airflow/operators/python_operator.py
--- a/airflow/operators/python_operator.py
+++ b/airflow/operators/python_operator.py
@@ -330,13 +330,19 @@ def _write_string_args(self, filename):
 
     def _write_args(self, input_filename):
         # serialize args to file
+        if self.use_dill:
+            serializer = dill
+        else:
+            serializer = pickle
+        # some items from context can't be loaded in virtual env
+        # see pr https://github.com/apache/airflow/pull/8256
+        not_serializable = {'dag', 'task', 'ti', 'macros', 'task_instance', 'var'}
         if self._pass_op_args():
+            kwargs = {key: value for key, value in self.op_kwargs.items()
+                      if key not in not_serializable}
             with open(input_filename, 'wb') as f:
-                arg_dict = ({'args': self.op_args, 'kwargs': self.op_kwargs})
-                if self.use_dill:
-                    dill.dump(arg_dict, f)
-                else:
-                    pickle.dump(arg_dict, f)
+                arg_dict = ({'args': self.op_args, 'kwargs': kwargs})
+                serializer.dump(arg_dict, f)
 
     def _read_result(self, output_filename):
         if os.stat(output_filename).st_size == 0:


Test patch:
diff --git a/tests/operators/test_python_operator.py b/tests/operators/test_python_operator.py
--- a/tests/operators/test_python_operator.py
+++ b/tests/operators/test_python_operator.py
@@ -35,9 +35,11 @@
 from airflow.models.taskinstance import clear_task_instances
 from airflow.operators.dummy_operator import DummyOperator
 from airflow.operators.python_operator import PythonOperator, BranchPythonOperator
+from airflow.operators.python_operator import PythonVirtualenvOperator
 from airflow.operators.python_operator import ShortCircuitOperator
 from airflow.settings import Session
 from airflow.utils import timezone
+from tests.test_utils.db import clear_db_runs, clear_db_dags
 from airflow.utils.db import create_session
 from airflow.utils.state import State
 
@@ -341,6 +343,75 @@ def test_echo_env_variables(self):
         t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
 
 
+class TestPythonVirtualenvOperator(TestPythonBase):
+    @classmethod
+    def setUpClass(cls):
+        super(TestPythonVirtualenvOperator, cls).setUpClass()
+        clear_db_runs()
+
+    def setUp(self):
+        super(TestPythonVirtualenvOperator, self).setUp()
+
+        def del_env(key):
+            try:
+                del os.environ[key]
+            except KeyError:
+                pass
+
+        del_env('AIRFLOW_CTX_DAG_ID')
+        del_env('AIRFLOW_CTX_TASK_ID')
+        del_env('AIRFLOW_CTX_EXECUTION_DATE')
+        del_env('AIRFLOW_CTX_DAG_RUN_ID')
+        self.dag = DAG(
+            'test_dag',
+            default_args={
+                'owner': 'airflow',
+                'start_date': DEFAULT_DATE})
+        self.addCleanup(self.dag.clear)
+        self.clear_run()
+        self.addCleanup(self.clear_run)
+
+    def tearDown(self):
+        super(TestPythonVirtualenvOperator, self).tearDown()
+        clear_db_runs()
+        clear_db_dags()
+
+        for var in TI_CONTEXT_ENV_VARS:
+            if var in os.environ:
+                del os.environ[var]
+
+    def clear_run(self):
+        self.run = False
+
+    def do_run(self):
+        self.run = True
+
+    def is_run(self):
+        return self.run
+
+    def test_config_context(self):
+        """
+        This test ensures we can use dag_run from the context
+        to access the configuration at run time that's being
+        passed from the UI, CLI, and REST API.
+        """
+        self.dag.create_dagrun(
+            run_id='manual__' + DEFAULT_DATE.isoformat(),
+            execution_date=DEFAULT_DATE,
+            start_date=DEFAULT_DATE,
+            state=State.RUNNING,
+            external_trigger=False,
+        )
+
+        def pass_function(**kwargs):
+            kwargs['dag_run'].conf
+
+        t = PythonVirtualenvOperator(task_id='config_dag_run', dag=self.dag,
+                                     provide_context=True,
+                                     python_callable=pass_function)
+        t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
+
+
 class BranchOperatorTest(unittest.TestCase):
     @classmethod
     def setUpClass(cls):