diff --git a/providers/amazon/src/airflow/providers/amazon/aws/executors/batch/batch_executor.py b/providers/amazon/src/airflow/providers/amazon/aws/executors/batch/batch_executor.py index 01bed883ddc92..eb65427cd76e8 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/executors/batch/batch_executor.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/executors/batch/batch_executor.py @@ -67,6 +67,19 @@ class AwsBatchExecutor(BaseExecutor): + """Executor that runs Airflow tasks on AWS Batch.""" + + def _ti_to_execute_workload_cmd(self, ti): + from airflow.executors.workloads import ExecuteTask + + workload_json = ExecuteTask.make(ti).model_dump_json() + return [ + "python3", + "-m", + "airflow.sdk.execution_time.execute_workload", + "--json-string", + workload_json, + ] """ The Airflow Scheduler creates a shell command, and passes it to the executor. @@ -501,7 +514,7 @@ def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[Task self.active_workers.add_job( job_id=batch_job.job_id, airflow_task_key=ti.key, - airflow_cmd=ti.command_as_list(), + airflow_cmd=self._ti_to_execute_workload_cmd(ti), queue=ti.queue, exec_config=ti.executor_config, attempt_number=ti.try_number, diff --git a/providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py b/providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py index 3605c803cd0d6..ffb1718051ed6 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py @@ -69,6 +69,19 @@ class AwsEcsExecutor(BaseExecutor): + """Executor that runs Airflow tasks on AWS ECS.""" + + def _ti_to_execute_workload_cmd(self, ti): + from airflow.executors.workloads import ExecuteTask + + workload_json = ExecuteTask.make(ti).model_dump_json() + return [ + "python3", + "-m", + "airflow.sdk.execution_time.execute_workload", + "--json-string", + workload_json, + ] """ Executes the provided Airflow command on an ECS instance. @@ -159,7 +172,7 @@ def _process_workloads(self, workloads: Sequence[workloads.All]) -> None: def start(self): """Call this when the Executor is run for the first time by the scheduler.""" check_health = self.conf.getboolean( - CONFIG_GROUP_NAME, AllEcsConfigKeys.CHECK_HEALTH_ON_STARTUP, fallback=False + CONFIG_GROUP_NAME, AllEcsConfigKeys.CHECK_HEALTH_ON_STARTUP, fallback=True ) if not check_health: @@ -590,7 +603,7 @@ def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[Task task, ti.key, ti.queue, - ti.command_as_list(), + self._ti_to_execute_workload_cmd(ti), ti.executor_config, ti.try_number, )