Repository: apache/airflow

Issue:
[AIRFLOW-6361] Run LocalTaskJob directly in Celery task
Hello,

The executor runs multiple processes to perform one task. Many processes have a very short life cycle, so the process of starting it is a significant overhead.

Firstly, the Celery executor trigger Celery tasks - app.task. This task runs the CLI  command (first process), which contains LocalTaskJob. LocalTaskJob runs the separate command (second process) that executes user-code. This level of isolation is redundant because LocalTaskJob doesn't execute unsafe code. The first command is run by a new process creation, not by a fork, so this is an expensive operation. I suggest running code from the first process as part of the celery task to reduce the need to create new processes.

The code currently uses CLIFactory to run the LocalTaskJob It is better to do this without unnecessary dependence on CLI, but it is a big change and I plan to do it in a separate PR.
WIP PR: https://github.com/mik-laj/incubator-airflow/pull/10 (Travis green :-D )

Performance benchmark: 
===================
Example DAG from Airflow with unneeded sleep instructions deleted.
```python
"""Example DAG demonstrating the usage of the BashOperator."""

from datetime import timedelta

import airflow
from airflow.models import DAG
from airflow.operators.bash_operator import BashOperator
from airflow.operators.dummy_operator import DummyOperator

args = {
    'owner': 'airflow',
    'start_date': airflow.utils.dates.days_ago(2),
}

dag = DAG(
    dag_id='example_bash_operator',
    default_args=args,
    schedule_interval='0 0 * * *',
    dagrun_timeout=timedelta(minutes=60),
)

run_this_last = DummyOperator(
    task_id='run_this_last',
    dag=dag,
)

# [START howto_operator_bash]
run_this = BashOperator(
    task_id='run_after_loop',
    bash_command='echo 1',
    dag=dag,
)
# [END howto_operator_bash]

run_this >> run_this_last

for i in range(3):
    task = BashOperator(
        task_id='runme_' + str(i),
        bash_command='echo "{{ task_instance_key_str }}",
        dag=dag,
    )
    task >> run_this

# [START howto_operator_bash_template]
also_run_this = BashOperator(
    task_id='also_run_this',
    bash_command='echo "run_id={{ run_id }} | dag_run={{ dag_run }}"',
    dag=dag,
)
# [END howto_operator_bash_template]
also_run_this >> run_this_last

if __name__ == "__main__":
    dag.cli()

```
```python
import airflow
from airflow import DAG
from airflow.models import DagBag

dagbag = airflow.models.DagBag()
dag: DAG = dagbag.get_dag("example_bash_operator")

dag.clear()
dag.run()
```
Environment: Brreze
```
unset AIRFLOW__CORE__DAGS_FOLDER
unset AIRFLOW__CORE__UNIT_TEST_MODE
chmod -R 777 /root
sudo -E su airflow
export AIRFLOW__CORE__EXECUTOR="CeleryExecutor"
export AIRFLOW__CELERY__BROKER_URL="redis://redis:6379/0"
export AIRFLOW__CELERY__WORKER_CONCURRENCY=8
seq 1 10 | xargs -n 1 -I {} bash -c "time python /files/benchmark_speed.py > /dev/null 2>&1" | grep '^(real\|user\|sys)';
```

Result: 

|Fn.     | After | Before | Change|
|--------|-------|--------|-------|
|AVERAGE | 56.48 | 38.32  | -32%  |
|VAR     | 23.60 | 0.04   | -98%  |
|MAX     | 68.29 | 38.68  | -43%  |
|MIN     | 53.26 | 38.08  | -28%  |
|STDEV   | 4.86  | 0.19   | -96%. |

Raw data
After:
```
real 0m38.394s
user 0m4.340s
sys 0m1.600s

real 0m38.355s
user 0m4.700s
sys 0m1.340s

real 0m38.675s
user 0m4.760s
sys 0m1.530s

real 0m38.488s
user 0m4.770s
sys 0m1.280s

real 0m38.434s
user 0m4.600s
sys 0m1.390s

real 0m38.378s
user 0m4.500s
sys 0m1.270s

real 0m38.106s
user 0m4.200s
sys 0m1.100s

real 0m38.082s
user 0m4.170s
sys 0m1.030s

real 0m38.173s
user 0m4.290s
sys 0m1.340s

real 0m38.161s
user 0m4.460s
sys 0m1.370s
```

Before:
```
real 0m53.488s
user 0m5.140s
sys 0m1.700s

real 1m8.288s
user 0m6.430s
sys 0m2.200s

real 0m53.371s
user 0m5.330s
sys 0m1.630s

real 0m58.939s
user 0m6.470s
sys 0m1.730s

real 0m53.255s
user 0m4.950s
sys 0m1.640s

real 0m58.802s
user 0m5.970s
sys 0m1.790s

real 0m58.449s
user 0m5.380s
sys 0m1.580s

real 0m53.308s
user 0m5.120s
sys 0m1.430s

real 0m53.485s
user 0m5.220s
sys 0m1.290s

real 0m53.387s
user 0m5.020s
sys 0m1.590s
```

---
Link to JIRA issue: https://issues.apache.org/jira/browse/AIRFLOW-6361

- [x] Description above provides context of the change
- [x] Commit message starts with `[AIRFLOW-NNNN]`, where AIRFLOW-NNNN = JIRA ID*
- [x] Unit tests coverage for changes (not needed for documentation changes)
- [x] Commits follow "[How to write a good git commit message](http://chris.beams.io/posts/git-commit/)"
- [x] Relevant documentation is updated including usage instructions.
- [x] I will engage committers as explained in [Contribution Workflow Example](https://github.com/apache/airflow/blob/master/CONTRIBUTING.rst#contribution-workflow-example).

(*) For document-only changes, no JIRA issue is needed. Commit message starts `[AIRFLOW-XXXX]`.

---
In case of fundamental code change, Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvements+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
In case of backwards incompatible changes please leave a note in [UPDATING.md](https://github.com/apache/airflow/blob/master/UPDATING.md).
Read the [Pull Request Guidelines](https://github.com/apache/airflow/blob/master/CONTRIBUTING.rst#pull-request-guidelines) for more information.



Discussion:
# [Codecov](https://codecov.io/gh/apache/airflow/pull/6905?src=pr&el=h1) Report
> Merging [#6905](https://codecov.io/gh/apache/airflow/pull/6905?src=pr&el=desc) into [master](https://codecov.io/gh/apache/airflow/commit/7d97d16c47ba181839fd583b8653d2872dc4290b?src=pr&el=desc) will **increase** coverage by `0.03%`.
> The diff coverage is `70%`.

[![Impacted file tree graph](https://codecov.io/gh/apache/airflow/pull/6905/graphs/tree.svg?width=650&token=WdLKlKHOAU&height=150&src=pr)](https://codecov.io/gh/apache/airflow/pull/6905?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master    #6905      +/-   ##
==========================================
+ Coverage    84.7%   84.74%   +0.03%     
==========================================
  Files         680      680              
  Lines       38556    38755     +199     
==========================================
+ Hits        32659    32841     +182     
- Misses       5897     5914      +17
```


| [Impacted Files](https://codecov.io/gh/apache/airflow/pull/6905?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [airflow/task/task\_runner/standard\_task\_runner.py](https://codecov.io/gh/apache/airflow/pull/6905/diff?src=pr&el=tree#diff-YWlyZmxvdy90YXNrL3Rhc2tfcnVubmVyL3N0YW5kYXJkX3Rhc2tfcnVubmVyLnB5) | `73.58% <0%> (+6.91%)` | :arrow_up: |
| [airflow/executors/celery\_executor.py](https://codecov.io/gh/apache/airflow/pull/6905/diff?src=pr&el=tree#diff-YWlyZmxvdy9leGVjdXRvcnMvY2VsZXJ5X2V4ZWN1dG9yLnB5) | `88.11% <100%> (-0.33%)` | :arrow_down: |
| [airflow/bin/cli.py](https://codecov.io/gh/apache/airflow/pull/6905/diff?src=pr&el=tree#diff-YWlyZmxvdy9iaW4vY2xpLnB5) | `92.78% <73.33%> (-1.27%)` | :arrow_down: |
| [airflow/kubernetes/volume\_mount.py](https://codecov.io/gh/apache/airflow/pull/6905/diff?src=pr&el=tree#diff-YWlyZmxvdy9rdWJlcm5ldGVzL3ZvbHVtZV9tb3VudC5weQ==) | `44.44% <0%> (-55.56%)` | :arrow_down: |
| [airflow/kubernetes/volume.py](https://codecov.io/gh/apache/airflow/pull/6905/diff?src=pr&el=tree#diff-YWlyZmxvdy9rdWJlcm5ldGVzL3ZvbHVtZS5weQ==) | `52.94% <0%> (-47.06%)` | :arrow_down: |
| [airflow/kubernetes/pod\_launcher.py](https://codecov.io/gh/apache/airflow/pull/6905/diff?src=pr&el=tree#diff-YWlyZmxvdy9rdWJlcm5ldGVzL3BvZF9sYXVuY2hlci5weQ==) | `45.25% <0%> (-46.72%)` | :arrow_down: |
| [airflow/kubernetes/refresh\_config.py](https://codecov.io/gh/apache/airflow/pull/6905/diff?src=pr&el=tree#diff-YWlyZmxvdy9rdWJlcm5ldGVzL3JlZnJlc2hfY29uZmlnLnB5) | `50.98% <0%> (-23.53%)` | :arrow_down: |
| [...rflow/contrib/operators/kubernetes\_pod\_operator.py](https://codecov.io/gh/apache/airflow/pull/6905/diff?src=pr&el=tree#diff-YWlyZmxvdy9jb250cmliL29wZXJhdG9ycy9rdWJlcm5ldGVzX3BvZF9vcGVyYXRvci5weQ==) | `78.75% <0%> (-20%)` | :arrow_down: |
| [airflow/utils/file.py](https://codecov.io/gh/apache/airflow/pull/6905/diff?src=pr&el=tree#diff-YWlyZmxvdy91dGlscy9maWxlLnB5) | `86.3% <0%> (-5.17%)` | :arrow_down: |
| [airflow/utils/cli.py](https://codecov.io/gh/apache/airflow/pull/6905/diff?src=pr&el=tree#diff-YWlyZmxvdy91dGlscy9jbGkucHk=) | `68.75% <0%> (-2.95%)` | :arrow_down: |
| ... and [90 more](https://codecov.io/gh/apache/airflow/pull/6905/diff?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/apache/airflow/pull/6905?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/apache/airflow/pull/6905?src=pr&el=footer). Last update [7d97d16...530125d](https://codecov.io/gh/apache/airflow/pull/6905?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

Really nice! This looks really nice - unless there is some hidden side effect in the CeleryExecutor, this looks like really nice improvement. LocalTaskJob uses task runner underneath that provides appropriate isolation between tasks already so indeed that looks like an unnecessary overhead to start yet another process. I'd love others to take a look but this one looks great for me!
The improvements sound attractive. One question, does this mean every time a change is made to the tasks, the worker needs to be restarted to pick up those changes?
No. Behaviour in this situation is completely identical. DAG was earlier and is still reloaded before running LocalTaskJob and before executing a raw task.  However I think, this can be optimized so that LocalTaskJob doesn't need to have access to DAG. it has all the necessary information in the database, but this will be part of another PR.
To show the process of execution and communication between processes, I drew two sequence diagrams. I hope it will be helpful.
![image](https://user-images.githubusercontent.com/12058428/71550255-f55f8800-29cb-11ea-9f4e-09821ef801eb.png)

![image](https://user-images.githubusercontent.com/12058428/71550253-ebd62000-29cb-11ea-827f-d6d66cf25256.png)
![image](https://user-images.githubusercontent.com/12058428/71550259-145e1a00-29cc-11ea-892a-51cc86cba276.png)


@mik-laj - should we add this graph to the documentation once merged? It's SUUUPER useful! And maybe we can add the first graph to the documentation in v1-10-test - so that it can be released with 1.10.8 docs ?
Yes. I wanna to add it to documentation. It's only early draft now. I working on it yet. 
I created a separate PR with a sequence diagram.  https://github.com/apache/airflow/pull/6961
Broadly sounds good, however as it stands this PR would load the dag in to the long running Celery worker process. Which, when coupled with my fork change _might_ mean that this stops tasks automatically picking up dag file changes!
@ashb I tested locally and this does not happen because a new Dagbag is created each time CLI command is called. I would also like to make LocalTaskJob independent of DAG and operator and only use information from TaskInstance.  It seems to me that there is no need to load the entire DAG in LocalTaskJob. This can further increase performance.
> @ashb I tested locally and this does not happen because a new Dagbag is created each time CLI command is called.

Oh yes, we recreate a new dag bag. The only thing that would be preloaded then is modules or plugins. This is okay, but should be written down somewhere. (I'm not sure where - we don't have anything like this in the docs right now)

> I would also like to make LocalTaskJob independent of DAG and operator and only use information from TaskInstance. It seems to me that there is no need to load the entire DAG in LocalTaskJob. This can further increase performance.

Yes totally! 
> @ashb I tested locally and this does not happen because a new Dagbag is created each time CLI command is called.

Wait, do we? In on my phone so checking the code is hard right now, but we used to have a single module-global dagbag object in cli.py. Where in the code does the dagbag get (re)created now?
@ashb 
1. https://github.com/apache/airflow/blob/master/airflow/cli/commands/task_command.py#L142
2. https://github.com/apache/airflow/blob/master/airflow/utils/cli.py#L155
@ashb I added note about plugins in `celery.rst` file.  Although I will try to remove this limitation in the future.  It doesn't seem difficult. The next step is to decouple the executor from the CLI. 
This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.

do we want to include this in 2.0?
@ryw  It is not necessary to be included in Airflow 2.0. We can do it in any version. After that, I realized it would be a problem for my users. They use Git/GCS sync to sync plugins, and this can be problematic after making this change. For this reason, if we would like to add this change, it should be activated by the feature gate.

If you are interested you can take over this change as it is not currently in my task queue.

Resolving patch:
diff --git a/airflow/executors/celery_executor.py b/airflow/executors/celery_executor.py
--- a/airflow/executors/celery_executor.py
+++ b/airflow/executors/celery_executor.py
@@ -37,7 +37,9 @@
 from celery.backends.base import BaseKeyValueStoreBackend
 from celery.backends.database import DatabaseBackend, Task as TaskDb, session_cleanup
 from celery.result import AsyncResult
+from setproctitle import setproctitle  # pylint: disable=no-name-in-module
 
+import airflow.settings as settings
 from airflow.config_templates.default_celery import DEFAULT_CELERY_CONFIG
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
@@ -78,6 +80,45 @@ def execute_command(command_to_exec: CommandType) -> None:
     """Executes command."""
     BaseExecutor.validate_command(command_to_exec)
     log.info("Executing command in Celery: %s", command_to_exec)
+
+    if settings.EXECUTE_TASKS_NEW_PYTHON_INTERPRETER:
+        _execute_in_subprocees(command_to_exec)
+    else:
+        _execute_in_fork(command_to_exec)
+
+
+def _execute_in_fork(command_to_exec: CommandType) -> None:
+    pid = os.fork()
+    if pid:
+        # In parent, wait for the child
+        pid, ret = os.waitpid(pid, 0)
+        if ret == 0:
+            return
+
+        raise AirflowException('Celery command failed on host: ' + get_hostname())
+
+    from airflow.sentry import Sentry
+
+    ret = 1
+    try:
+        from airflow.cli.cli_parser import get_parser
+        parser = get_parser()
+        # [1:] - remove "airflow" from the start of the command
+        args = parser.parse_args(command_to_exec[1:])
+
+        setproctitle(f"airflow task supervisor: {command_to_exec}")
+
+        args.func(args)
+        ret = 0
+    except Exception as e:  # pylint: disable=broad-except
+        log.error("Failed to execute task %s.", str(e))
+        ret = 1
+    finally:
+        Sentry.flush()
+        os._exit(ret)  # pylint: disable=protected-access
+
+
+def _execute_in_subprocees(command_to_exec: CommandType) -> None:
     env = os.environ.copy()
     try:
         # pylint: disable=unexpected-keyword-arg


Test patch:
diff --git a/tests/executors/test_celery_executor.py b/tests/executors/test_celery_executor.py
--- a/tests/executors/test_celery_executor.py
+++ b/tests/executors/test_celery_executor.py
@@ -225,18 +225,20 @@ def test_gauge_executor_metrics(self, mock_stats_gauge, mock_trigger_tasks, mock
         [['airflow', 'version'], ValueError],
         [['airflow', 'tasks', 'run'], None]
     ))
-    @mock.patch('subprocess.check_output')
-    def test_command_validation(self, command, expected_exception, mock_check_output):
+    def test_command_validation(self, command, expected_exception):
         # Check that we validate _on the receiving_ side, not just sending side
-        if expected_exception:
-            with pytest.raises(expected_exception):
+        with mock.patch('airflow.executors.celery_executor._execute_in_subprocees') as mock_subproc, \
+             mock.patch('airflow.executors.celery_executor._execute_in_fork') as mock_fork:
+            if expected_exception:
+                with pytest.raises(expected_exception):
+                    celery_executor.execute_command(command)
+                mock_subproc.assert_not_called()
+                mock_fork.assert_not_called()
+            else:
                 celery_executor.execute_command(command)
-            mock_check_output.assert_not_called()
-        else:
-            celery_executor.execute_command(command)
-            mock_check_output.assert_called_once_with(
-                command, stderr=mock.ANY, close_fds=mock.ANY, env=mock.ANY,
-            )
+                # One of these should be called.
+                assert mock_subproc.call_args == ((command,),) or \
+                    mock_fork.call_args == ((command,),)
 
     @pytest.mark.backend("mysql", "postgres")
     def test_try_adopt_task_instances_none(self):