ferruzzi commented on code in PR #63657:
URL: https://github.com/apache/airflow/pull/63657#discussion_r3113903586
##########
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:
I believe you can remove most/all of these ignores now. Try removing them
and see which ones mypy insists on keeping.
##########
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:
+ from airflow.executors import workloads
+
+ for workload in workload_items:
+ if isinstance(workload, workloads.ExecuteTask):
+ command = [workload]
+ key = workload.ti.key
+ queue = workload.ti.queue
+ executor_config = workload.ti.executor_config or {}
- def _process_workloads(self, workloads: Sequence[workloads.All]) -> None:
- from airflow.executors.workloads import ExecuteTask
+ del self.queued_tasks[key]
+ self.execute_async(key=key, command=command, queue=queue,
executor_config=executor_config) # type: ignore[arg-type]
+ self.running.add(key)
- # Airflow V3 version
- for w in workloads:
- if not isinstance(w, ExecuteTask):
- 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]
Review Comment:
Here and elsewhere: it's semantically the same thing, but use
workload.callback.key as a key please, it will makes things more uniform
##########
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:
Review Comment:
Double check me, but we should be able to drop this entire `queue_workload`
method now and fall back on BaseExecutor's version.
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py:
##########
@@ -103,7 +107,9 @@ class AwsEcsExecutor(BaseExecutor):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.active_workers: EcsTaskCollection = EcsTaskCollection()
- self.pending_tasks: deque = deque()
+ self.pending_workloads: deque = deque()
+ if AIRFLOW_V_3_3_PLUS:
+ self.queued_callbacks: dict[str, workloads.ExecuteCallback] = {}
Review Comment:
This is redundant since BaseExecutor should be initializing that now
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py:
##########
@@ -261,11 +279,11 @@ def sync(self):
# up and kill the scheduler process
self.log.exception("Failed to sync %s", self.__class__.__name__)
- def sync_running_tasks(self):
- """Check and update state on all running tasks."""
+ def sync_running_workloads(self):
Review Comment:
Why no shims? Pretty sure we need them for back-compat, no? Just something
like
```python
@property
def pending_tasks(self) -> deque:
"""Deprecated: use pending_workloads."""
return self.pending_workloads
def sync_running_tasks(self):
"""Deprecated: use sync_running_workloads."""
return self.sync_running_workloads()
def attempt_task_runs(self):
"""Deprecated: use attempt_workload_runs."""
return self.attempt_workload_runs()
```
--
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]
