Skip to content

Use Task SDK BaseHook in Amazon executors #62090

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
dylanschw wants to merge 3 commits into from
Closed

Use Task SDK BaseHook in Amazon executors #62090

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
9ea0ed7
Use Task SDK BaseHook in Amazon executors
dylanschw Feb 17, 2026
53c03da
Merge branch 'main' into fix-58205-ecs-command-as-list
dylanschw Feb 17, 2026
71ccf1f
Merge branch 'main' into fix-58205-ecs-command-as-list
dylanschw Feb 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 14 additions & 1 deletion providers/amazon/src/airflow/providers/amazon/aws/executors/batch/batch_executor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down Expand Up @@ -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,
Expand Down
17 changes: 15 additions & 2 deletions providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,19 @@


class AwsEcsExecutor(BaseExecutor):
"""Executor that runs Airflow tasks on AWS ECS."""

def _ti_to_execute_workload_cmd(self, ti):

@SameerMesiah97 SameerMesiah97 Feb 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe this functionality can be included in the executor/utils folder to avoid duplication.. I do not think this is a massive issue as it is only duplicated once but I think it should be considered. I would wait for the codeowner to weigh in here.

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,
]

@SameerMesiah97 SameerMesiah97 Feb 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this compatible with airflow 2.x? It seems like this will not work for versions before Airflow 3.0 as airflow.sdk was introduced then. I think you should have a look at PR #58207. There are issues with that but I think the implementation in that PR was more backwards compatible.

@dylanschw dylanschw Feb 18, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for calling this out. As written, using airflow.sdk.execution_time.execute_workload is not compatible with Airflow 2.x since airflow.sdk is only available in Airflow 3.x. I will update the executor command construction to branch on Airflow version, using the SDK entrypoint for Airflow 3.x and the legacy task runner entrypoint for Airflow 2.x. I will add unit tests that assert the exact command built for both Airflow >= 3.0 and Airflow < 3.0, including the entrypoint and arguments like --json-string. Since the regression occurred there, I will take another look at test_try_adopt_task_instances for ECS and Batch. If unskipping is not feasible due to environment assumptions, I will add a targeted test to cover the same behavior with mocking.

On factoring this into executor/utils, that sounds reasonable. I can do it in this PR

"""
Executes the provided Airflow command on an ECS instance.

Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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,
)
Expand Down