Repository: apache/airflow
Pull request: Add support for custom task runner
https://github.com/apache/airflow/issues/8067
 I'm not sure if anyone will use it, but it can stimulate community creativity.

---
Make sure to mark the boxes below before creating PR: [x]

- [X] Description above provides context of the change
- [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).

---
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.


Problem statement:
Allow Custom Task Runner via plugins
**Description**
Airflow provides a plugin manager to integrate external features. However, there is no configuration available to have custom task runner as part of plugins.

**Use case / motivation**
 Enabling custom task runner integration with Airflow will help users in configuring runners as per environment requirement and dependencies.

Hints from issue discussion:
It looks trivial. Do you want to work on it? I think this should not be configured by plugins, but the configuration option. Similar to the executor.
Please look at https://github.com/apache/airflow/pull/7563

How would you like to use it? Do you have any special task runner already implemented?  I wonder if it would be a better idea to contribute it to the community.

Commit message: Add support for custom task runner

Patch to review:
diff --git a/airflow/config_templates/config.yml b/airflow/config_templates/config.yml
index 5033c342e503d..bcfbf3f15e507 100644
--- a/airflow/config_templates/config.yml
+++ b/airflow/config_templates/config.yml
@@ -54,6 +54,7 @@
       description: |
         The executor class that airflow should use. Choices include
         SequentialExecutor, LocalExecutor, CeleryExecutor, DaskExecutor, KubernetesExecutor
+        or the full import path to the class when using a custom executor.
       version_added: ~
       type: string
       example: ~
@@ -237,7 +238,9 @@
       default: "50"
     - name: task_runner
       description: |
-        The class to use for running task instances in a subprocess
+        The class to use for running task instances in a subprocess.
+        Choices include StandardTaskRunner, CgroupTaskRunner or the full import path to the class
+        when using a custom task runner.
       version_added: ~
       type: string
       example: ~
diff --git a/airflow/config_templates/default_airflow.cfg b/airflow/config_templates/default_airflow.cfg
index dfd779056fffe..f822f640efc2d 100644
--- a/airflow/config_templates/default_airflow.cfg
+++ b/airflow/config_templates/default_airflow.cfg
@@ -49,6 +49,7 @@ default_timezone = utc
 
 # The executor class that airflow should use. Choices include
 # SequentialExecutor, LocalExecutor, CeleryExecutor, DaskExecutor, KubernetesExecutor
+# or the full import path to the class when using a custom executor.
 executor = SequentialExecutor
 
 # The SqlAlchemy connection string to the metadata database.
@@ -144,7 +145,9 @@ dagbag_import_timeout = 30
 # How long before timing out a DagFileProcessor, which processes a dag file
 dag_file_processor_timeout = 50
 
-# The class to use for running task instances in a subprocess
+# The class to use for running task instances in a subprocess.
+# Choices include StandardTaskRunner, CgroupTaskRunner or the full import path to the class
+# when using a custom task runner.
 task_runner = StandardTaskRunner
 
 # If set, tasks without a ``run_as_user`` argument will be run with this user
diff --git a/airflow/task/task_runner/__init__.py b/airflow/task/task_runner/__init__.py
index ef7686ee80f16..52c3052e8362f 100644
--- a/airflow/task/task_runner/__init__.py
+++ b/airflow/task/task_runner/__init__.py
@@ -17,12 +17,23 @@
 # under the License.
 
 # pylint: disable=missing-docstring
+import logging
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
-from airflow.task.task_runner.standard_task_runner import StandardTaskRunner
+from airflow.utils.module_loading import import_string
 
-_TASK_RUNNER = conf.get('core', 'TASK_RUNNER')
+log = logging.getLogger(__name__)
+
+_TASK_RUNNER_NAME = conf.get('core', 'TASK_RUNNER')
+
+STANDARD_TASK_RUNNER = "StandardTaskRunner"
+
+CGROUP_TASK_RUNNER = "CgroupTaskRunner"
+
+CORE_TASK_RUNNERS = {
+    STANDARD_TASK_RUNNER: "airflow.task.task_runner.standard_task_runner.StandardTaskRunner",
+    CGROUP_TASK_RUNNER: "airflow.task.task_runner.cgroup_task_runner.CgroupTaskRunner",
+}
 
 
 def get_task_runner(local_task_job):
@@ -35,10 +46,12 @@ def get_task_runner(local_task_job):
     :return: The task runner to use to run the task.
     :rtype: airflow.task.task_runner.base_task_runner.BaseTaskRunner
     """
-    if _TASK_RUNNER == "StandardTaskRunner":
-        return StandardTaskRunner(local_task_job)
-    elif _TASK_RUNNER == "CgroupTaskRunner":
-        from airflow.task.task_runner.cgroup_task_runner import CgroupTaskRunner
-        return CgroupTaskRunner(local_task_job)
-    else:
-        raise AirflowException("Unknown task runner type {}".format(_TASK_RUNNER))
+    if _TASK_RUNNER_NAME in CORE_TASK_RUNNERS:
+        log.debug("Loading core task runner: %s", _TASK_RUNNER_NAME)
+        task_runner_class = import_string(CORE_TASK_RUNNERS[_TASK_RUNNER_NAME])
+        return task_runner_class(local_task_job)
+
+    log.debug("Loading task runner from custom path: %s", _TASK_RUNNER_NAME)
+    task_runner_class = import_string(_TASK_RUNNER_NAME)
+    task_runner = task_runner_class(local_task_job)
+    return task_runner
diff --git a/tests/task/task_runner/test_task_runner.py b/tests/task/task_runner/test_task_runner.py
new file mode 100644
index 0000000000000..ff6b2375d4643
--- /dev/null
+++ b/tests/task/task_runner/test_task_runner.py
@@ -0,0 +1,52 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import unittest
+from unittest import mock
+
+from airflow.task.task_runner import get_task_runner
+
+custom_task_runner = mock.MagicMock()
+
+
+class GetTaskRunner(unittest.TestCase):
+
+    @mock.patch(
+        'airflow.task.task_runner.base_task_runner.subprocess'
+    )
+    @mock.patch('airflow.task.task_runner._TASK_RUNNER_NAME', "StandardTaskRunner")
+    def test_should_support_core_task_runner(self, mock_subprocess):
+        local_task_job = mock.MagicMock(
+            **{'task_instance.get_template_context.return_value': {"ti": mock.MagicMock()}}
+        )
+        task_runner = get_task_runner(local_task_job)
+
+        self.assertEqual("StandardTaskRunner", task_runner.__class__.__name__)
+
+    @mock.patch(
+        'airflow.task.task_runner._TASK_RUNNER_NAME',
+        "tests.task.task_runner.test_task_runner.custom_task_runner"
+    )
+    def test_should_support_custom_task_runner(self):
+        local_task_job = mock.MagicMock(
+            **{'task_instance.get_template_context.return_value': {"ti": mock.MagicMock()}}
+        )
+        custom_task_runner.reset_mock()
+
+        task_runner = get_task_runner(local_task_job)
+
+        custom_task_runner.assert_called_once_with(local_task_job)


Code review comments:
diff_hunk: @@ -35,10 +46,12 @@ def get_task_runner(local_task_job):
     :return: The task runner to use to run the task.
     :rtype: airflow.task.task_runner.base_task_runner.BaseTaskRunner
     """
-    if _TASK_RUNNER == "StandardTaskRunner":
-        return StandardTaskRunner(local_task_job)
-    elif _TASK_RUNNER == "CgroupTaskRunner":
-        from airflow.task.task_runner.cgroup_task_runner import CgroupTaskRunner
-        return CgroupTaskRunner(local_task_job)
-    else:
-        raise AirflowException("Unknown task runner type {}".format(_TASK_RUNNER))
+    if _TASK_RUNNER_NAME in CORE_TASK_RUNNERS:
+        log.debug("Loading core task runner: %s", _TASK_RUNNER_NAME)
+        task_runner_class = import_string(CORE_TASK_RUNNERS[_TASK_RUNNER_NAME])
+        return task_runner_class(local_task_job)
+
+    log.debug("Loading task runner from custom path: %s", _TASK_RUNNER_NAME)
path: airflow/task/task_runner/__init__.py
text: @user1:
Please remove code duplication

@author:
We don't have code duplication. These are separate ways to load modules. 

There is only one common fragment.
``
task_runner_class(local_task_job)
``
However, pulling out this common part will not reduce the amount of code.

@user1:
The only difference I see is that the parameter pass to `import_string`. Rest code can be taken out in a method passing the string parameter required by `import_string` to new method. 


@user1:
Something like this 
```
def get_runner(runner_class, local_task_job):
    log.debug("Loading core task runner: %s", _TASK_RUNNER_NAME)
    task_runner_class = import_string(runner_class)
    return task_runner_class(local_task_job)


def get_task_runner(local_task_job):
    if _TASK_RUNNER_NAME in CORE_TASK_RUNNERS:
        return get_runner(CORE_TASK_RUNNERS[_TASK_RUNNER_NAME], local_task_job)
    return get_runner(_TASK_RUNNER_NAME, local_task_job)

```