ferruzzi commented on code in PR #63657:
URL: https://github.com/apache/airflow/pull/63657#discussion_r3222883123
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/utils.py:
##########
@@ -35,7 +35,10 @@
from airflow.utils.state import State
if TYPE_CHECKING:
- from airflow.models.taskinstance import TaskInstanceKey
+ from airflow.providers.amazon.version_compat import AIRFLOW_V_3_3_PLUS
+
+ if AIRFLOW_V_3_3_PLUS:
+ from airflow.executors.workloads.types import WorkloadKey
Review Comment:
```suggestion
if AIRFLOW_V_3_3_PLUS:
from airflow.executors.workloads.types import WorkloadKey
else:
from airflow.models.taskinstance import TaskInstanceKey as
WorkloadKey # type: ignore[assignment]
```
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py:
##########
@@ -134,27 +140,39 @@ def __init__(self, *args, **kwargs):
def queue_workload(self, workload: workloads.All, session: Session | None)
-> None:
from airflow.executors import workloads
- if not isinstance(workload, workloads.ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type
{type(workload)}")
- ti = workload.ti
- self.queued_tasks[ti.key] = workload
+ if isinstance(workload, workloads.ExecuteTask):
+ ti = workload.ti
+ self.queued_tasks[ti.key] = workload
+ return
+ if AIRFLOW_V_3_3_PLUS and isinstance(workload,
workloads.ExecuteCallback):
+ self.queued_callbacks[workload.callback.id] = workload
+ return
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type
{type(workload)}")
+
+ def _process_workloads(self, workload_items: Sequence[workloads.All]) ->
None:
Review Comment:
PR#65392 should have made most of these ignores in the following block
unnecessary. Have you confirmed that mypy still needs them?
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py:
##########
@@ -329,13 +353,13 @@ def __log_container_failures(self, task_arn: str):
"The ECS task failed due to the following containers
failing:\n%s", "\n".join(reasons)
)
- def __handle_failed_task(self, task_arn: str, reason: str):
+ def __handle_failed_workload(self, task_arn: str, reason: str):
"""
If an API failure occurs, the task is rescheduled.
Review Comment:
```suggestion
If an API failure occurs, the workload is rescheduled.
```
--
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.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
