Add ExecuteCallback support to AWS ECS Executor by shivaam · Pull Request #63657 · apache/airflow

shivaam · 2026-03-15T14:44:46Z

Implements executor callback support for the AWS ECS Executor.

Sibling to merged Lambda PR #63035; builds on #62645 (merged) and #65392 (merged).

related: #62887

Changes

supports_callbacks = True gated on AIRFLOW_V_3_3_PLUS so base executor routes ExecuteCallback to ECS.
Widens key types from TaskInstanceKey to WorkloadKey in ecs_executor.py, EcsQueuedTask, EcsTaskCollection.
queue_workload() override branches on ExecuteTask vs ExecuteCallback (kept until min Airflow version reaches 3.3, after which BaseExecutor.queue_workload routes both natively; mirrors merged Lambda). No tracking issue yet — happy to open one for the override removal as a follow-up.
_process_workloads() dispatches both workload types via execute_async().
Renames methods/attrs to reflect workloads, with AirflowProviderDeprecationWarning shims for the public ones (matches merged Lambda):
- public + shim: sync_running_tasks → sync_running_workloads, attempt_task_runs → attempt_workload_runs, pending_tasks → pending_workloads
- private (no shim): __update_running_task → __update_running_workload, __handle_failed_task → __handle_failed_workload
Relies on Widen BaseExecutor method signatures to accept WorkloadKey #65392 widening BaseExecutor signatures — no cast() needed, log_task_event() skips callback keys internally.

Tests

CI green: ruff, mypy, full test suite, plus compat suites on Airflow 2.11.1 / 3.0.6 / 3.1.8 / 3.2.0.

End-to-end re-verified 2026-05-02 against rebased main on real ECS Fargate. Each DAG runs a regular ExecuteTask (slow_task 120s) followed by a deadline ExecuteCallback:

DAG	Callback pattern	Task result	Callback result
`ecs_callback_simple`	Plain function from S3 bundle	SUCCESS ~129s	SUCCESS in 3.9s
`ecs_callback_context`	Function reading dag_run + deadline context	SUCCESS	SUCCESS in 3.8s
`ecs_callback_failure`	Function raises `RuntimeError`	SUCCESS	FAILED after 3 retries
`ecs_callback_notifier_class`	Callable class (notifier `__call__` pattern)	SUCCESS	SUCCESS in 4.0s

Verified: PENDING → QUEUED → RUNNING → SUCCESS/FAILED for both workload types; __handle_failed_workload + pending_workloads.append(...) reschedule callbacks correctly; log_task_event and self.fail() with string callback keys handled correctly by #65392's base class. Callback UUIDs and ECS task ARNs from the run are available on request.

Backwards compatibility

Existing task path unchanged. Callback paths gated behind AIRFLOW_V_3_3_PLUS. Public renamed methods/attrs preserved as deprecation shims so subclassers continue to work.

Was generative AI tooling used to co-author this PR?

Yes — Claude Code (Opus 4.7) for code drafting, sibling-PR review comparison, and end-to-end test orchestration

Generated-by: Claude Code (Opus 4.7) following the guidelines

shivaam · 2026-03-15T14:51:57Z

@ferruzzi Draft PR

ferruzzi

Once #65392 merges (hopefully today) you'll need to rebase and make some small tweaks. That may also let you clean up some of those #ignore directives in _process_workloads, hopefully.

ferruzzi · 2026-04-17T20:01:10Z

                failure_count,
            )
-            self.fail(task_key)
+            self.fail(cast("TaskInstanceKey", task_key))


Once #65392 merges I think you can drop this and the other places you used cast(). Same for the places where you used isinstance(task_key, tuple)

shivaam · 2026-04-19T01:39:52Z

@ferruzzi Thanks for the review! I have updated the Pr and addressed all your comments.

Dropped the cast() wrappers and isinstance(tuple) guards now that Widen BaseExecutor method signatures to accept WorkloadKey #65392 is merged — BaseExecutor handles WorkloadKey natively and skips callback keys in log_task_event internally. Removed the unused cast import and one test whose coverage moved to airflow-core.
Renamed the _process_workloads param to workload_items so the module import no longer needs as wl.
Renamed the task-named methods and pending_tasks to their workload equivalents. Straight rename, no compat shim

ferruzzi · 2026-04-20T21:52:19Z

+        if AIRFLOW_V_3_3_PLUS:
+            self.queued_callbacks: dict[str, workloads.ExecuteCallback] = {}


This is redundant since BaseExecutor should be initializing that now

ferruzzi · 2026-04-20T21:57:47Z

-                raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+            elif AIRFLOW_V_3_3_PLUS and isinstance(workload, workloads.ExecuteCallback):
+                command = [workload]  # type: ignore[list-item]
+                key = workload.callback.id  # type: ignore[assignment]


Here and elsewhere: it's semantically the same thing, but use workload.callback.key as a key please, it will makes things more uniform

ferruzzi · 2026-04-20T22:57:17Z

    from sqlalchemy.orm import Session

    from airflow.executors import workloads
+    from airflow.executors.workloads.types import WorkloadKey


This will need to be wrapped in a version check. The user could be running latest provider package with core 3.0 and this will blow up. It'll be something like

if TYPE_CHECKING: if AIRFLOW_V_3_3_PLUS: from airflow.executors.workloads.types import WorkloadKey else: WorkloadKey: TypeAlias = TaskInstanceKey # type: ignore[no-redef, misc]

Enables the ECS executor to dispatch ExecuteCallback workloads (deadline alerts) alongside regular ExecuteTask workloads. Builds on apache#65392 which widened BaseExecutor signatures to accept WorkloadKey. - supports_callbacks = True (gated on AIRFLOW_V_3_3_PLUS) - Widen key types to WorkloadKey throughout EcsQueuedTask / EcsTaskCollection - Branch _process_workloads on ExecuteTask vs ExecuteCallback - Add AIRFLOW_V_3_3_PLUS to version_compat.py - Unit tests for queueing, processing, serialization, sync, mixed keys

Renames (mirrors the merged Lambda callback PR — straight rename, no shim, executor-internal surface): sync_running_tasks -> sync_running_workloads attempt_task_runs -> attempt_workload_runs pending_tasks (attr) -> pending_workloads __update_running_task -> __update_running_workload __handle_failed_task -> __handle_failed_workload Fix CI on older Airflow compat tests: - Restore queue_workload() override. Airflow 3.3+ BaseExecutor routes ExecuteCallback natively, but pre-3.3 raises ValueError for anything not ExecuteTask. Override works across versions. - Import AIRFLOW_V_3_3_PLUS from tests_common (main bumped to 3.3). check-airflow-v-imports-in-tests hook disallows provider-internal version_compat imports from test files.

vincbeck · 2026-05-04T17:43:56Z


    supports_multi_team: bool = True

+    if AIRFLOW_V_3_3_PLUS:


You should not need this? You should be able to declare this variable regardless o the Airflow version?

boring-cyborg Bot added area:providers provider:amazon AWS/Amazon - related issues labels Mar 15, 2026

shivaam force-pushed the feature/ecs-executor-callback branch from f3156ea to 4728b66 Compare March 15, 2026 14:50

shivaam force-pushed the feature/ecs-executor-callback branch 5 times, most recently from 38ddfea to 6ac2634 Compare March 15, 2026 18:44

shivaam marked this pull request as ready for review March 15, 2026 21:26

shivaam requested a review from o-nikolas as a code owner March 15, 2026 21:26

potiuk added the ready for maintainer review Set after triaging when all criteria pass. label Mar 16, 2026

ferruzzi mentioned this pull request Mar 16, 2026

Executor Callback support #62887

Open

2 tasks

shivaam force-pushed the feature/ecs-executor-callback branch 2 times, most recently from 0372a19 to a93048e Compare April 16, 2026 05:15

ferruzzi reviewed Apr 17, 2026

View reviewed changes

shivaam force-pushed the feature/ecs-executor-callback branch 2 times, most recently from 6a1c191 to a543ec3 Compare April 19, 2026 01:35

shivaam force-pushed the feature/ecs-executor-callback branch from ea660cf to 28b791f Compare April 19, 2026 02:44

ferruzzi reviewed Apr 20, 2026

View reviewed changes

ferruzzi mentioned this pull request Apr 20, 2026

feat: add callback support to aws batch executor #62984

Open

1 task

ferruzzi reviewed Apr 20, 2026

View reviewed changes

ferruzzi mentioned this pull request Apr 24, 2026

Small param rename in BaseExecutor and CeleryExecutor #65814

Closed

shivaam force-pushed the feature/ecs-executor-callback branch 3 times, most recently from 08bb01a to 726fbd1 Compare May 3, 2026 00:24

shivaam force-pushed the feature/ecs-executor-callback branch from 726fbd1 to 233ec6b Compare May 3, 2026 01:19

eladkal requested a review from vincbeck May 4, 2026 16:56

vincbeck reviewed May 4, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add ExecuteCallback support to AWS ECS Executor#63657

Add ExecuteCallback support to AWS ECS Executor#63657
shivaam wants to merge 2 commits intoapache:mainfrom
shivaam:feature/ecs-executor-callback

shivaam commented Mar 15, 2026 •

edited

Loading

Uh oh!

shivaam commented Mar 15, 2026

Uh oh!

ferruzzi left a comment

Uh oh!

Uh oh!

Uh oh!

ferruzzi Apr 17, 2026

Uh oh!

shivaam commented Apr 19, 2026 •

edited

Loading

Uh oh!

ferruzzi Apr 20, 2026

Uh oh!

Uh oh!

ferruzzi Apr 20, 2026

Uh oh!

Uh oh!

Uh oh!

ferruzzi Apr 20, 2026

Uh oh!

vincbeck May 4, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

		if AIRFLOW_V_3_3_PLUS:
		self.queued_callbacks: dict[str, workloads.ExecuteCallback] = {}

Conversation

shivaam commented Mar 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Was generative AI tooling used to co-author this PR?

Uh oh!

shivaam commented Mar 15, 2026

Uh oh!

ferruzzi left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

ferruzzi Apr 17, 2026

Choose a reason for hiding this comment

Uh oh!

shivaam commented Apr 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ferruzzi Apr 20, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ferruzzi Apr 20, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

ferruzzi Apr 20, 2026

Choose a reason for hiding this comment

Uh oh!

vincbeck May 4, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

shivaam commented Mar 15, 2026 •

edited

Loading

shivaam commented Apr 19, 2026 •

edited

Loading