@@ -63,9 +63,10 @@ def _timestamp_now() -> Timestamp:
6363 return ts
6464
6565
66- def _get_execution_mode (job_type ):
66+ def _get_execution_mode (utask_module , job_type ):
6767 """Determines whether this task in executed on swarming or batch."""
68- if swarming .is_swarming_task (job_type ):
68+ command = task_utils .get_command_from_module (utask_module .__name__ )
69+ if swarming .is_swarming_task (command , job_type ):
6970 return Mode .SWARMING
7071 return Mode .BATCH
7172
@@ -409,7 +410,7 @@ def tworker_preprocess(utask_module, task_argument, job_type, uworker_env):
409410 signed download URL for the uworker's input and the (unsigned) download URL
410411 for its output."""
411412 with _MetricRecorder (_Subtask .PREPROCESS ) as recorder :
412- execution_mode = _get_execution_mode (job_type )
413+ execution_mode = _get_execution_mode (utask_module , job_type )
413414 uworker_input = _preprocess (utask_module , task_argument , job_type ,
414415 uworker_env , recorder , execution_mode )
415416 if not uworker_input :
@@ -500,7 +501,8 @@ def tworker_postprocess(output_download_url) -> None:
500501 task_utils .reset_task_stage_env ()
501502
502503 utask_module = get_utask_module (uworker_output .uworker_input .module_name )
503- execution_mode = _get_execution_mode (uworker_output .uworker_input .job_type )
504+ execution_mode = _get_execution_mode (utask_module ,
505+ uworker_output .uworker_input .job_type )
504506 recorder .set_task_details (
505507 utask_module ,
506508 uworker_output .uworker_input .job_type ,
0 commit comments