rabsr commented on a change in pull request #8085: Add support for custom task
runner
URL: https://github.com/apache/airflow/pull/8085#discussion_r402944882
##########
File path: airflow/task/task_runner/__init__.py
##########
@@ -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)
Review comment:
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)
```
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
