ferruzzi commented on code in PR #62984:
URL: https://github.com/apache/airflow/pull/62984#discussion_r3114005158
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/batch/batch_executor.py:
##########
@@ -88,14 +97,15 @@ class AwsBatchExecutor(BaseExecutor):
"""
supports_multi_team: bool = True
+ supports_callbacks: bool = True
Review Comment:
This should be gated
```suggestion
if AIRFLOW_V_3_3_PLUS:
supports_callbacks: bool = True
```
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/batch/batch_executor.py:
##########
@@ -124,29 +134,48 @@ def __init__(self, *args, **kwargs):
fallback=CONFIG_DEFAULTS[AllBatchConfigKeys.MAX_SUBMIT_JOB_ATTEMPTS],
)
- def queue_workload(self, workload: workloads.All, session: Session | None)
-> None:
+ def queue_workload(self, workload: WorkloadTypes, session: Session | None)
-> None:
Review Comment:
Double check, but you should be able to drop this entirely and rely on
BaseExecutor's version
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/batch/utils.py:
##########
@@ -43,7 +52,7 @@
class BatchQueuedJob:
"""Represents a Batch job that is queued. The job will be run in the next
heartbeat."""
- key: TaskInstanceKey
+ key: BatchJobWorkloadKey
command: CommandType
queue: str
Review Comment:
Double check me here, but `_process_workloads` can set queue = None so this
should be
```suggestion
queue: str | None
```
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/batch/batch_executor.py:
##########
@@ -124,29 +134,48 @@ def __init__(self, *args, **kwargs):
fallback=CONFIG_DEFAULTS[AllBatchConfigKeys.MAX_SUBMIT_JOB_ATTEMPTS],
)
- def queue_workload(self, workload: workloads.All, session: Session | None)
-> None:
+ def queue_workload(self, workload: WorkloadTypes, session: Session | None)
-> None:
from airflow.executors import workloads
- if not isinstance(workload, workloads.ExecuteTask):
+ if AIRFLOW_V_3_3_PLUS and isinstance(workload,
workloads.ExecuteCallback):
+ self.queued_callbacks[workload.callback.id] = workload
+ elif isinstance(workload, workloads.ExecuteTask):
+ ti = workload.ti
+ self.queued_tasks[ti.key] = workload
+ else:
raise RuntimeError(f"{type(self)} cannot handle workloads of type
{type(workload)}")
- ti = workload.ti
- self.queued_tasks[ti.key] = workload
- def _process_workloads(self, workloads: Sequence[workloads.All]) -> None:
- from airflow.executors.workloads import ExecuteTask
+ def _process_workloads(self, workloads: Sequence[WorkloadTypes]) -> None:
+ from airflow.executors import workloads as wl
Review Comment:
Same comment as https://github.com/apache/airflow/pull/63657/:
```suggestion
def _process_workloads(self, workload_items: Sequence[workloads.All])
-> None:
from airflow.executors import workloads
```
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/batch/batch_executor.py:
##########
@@ -124,29 +134,48 @@ def __init__(self, *args, **kwargs):
fallback=CONFIG_DEFAULTS[AllBatchConfigKeys.MAX_SUBMIT_JOB_ATTEMPTS],
)
- def queue_workload(self, workload: workloads.All, session: Session | None)
-> None:
+ def queue_workload(self, workload: WorkloadTypes, session: Session | None)
-> None:
from airflow.executors import workloads
- if not isinstance(workload, workloads.ExecuteTask):
+ if AIRFLOW_V_3_3_PLUS and isinstance(workload,
workloads.ExecuteCallback):
+ self.queued_callbacks[workload.callback.id] = workload
+ elif isinstance(workload, workloads.ExecuteTask):
+ ti = workload.ti
+ self.queued_tasks[ti.key] = workload
+ else:
raise RuntimeError(f"{type(self)} cannot handle workloads of type
{type(workload)}")
- ti = workload.ti
- self.queued_tasks[ti.key] = workload
- def _process_workloads(self, workloads: Sequence[workloads.All]) -> None:
- from airflow.executors.workloads import ExecuteTask
+ def _process_workloads(self, workloads: Sequence[WorkloadTypes]) -> None:
+ from airflow.executors import workloads as wl
# Airflow V3 version
Review Comment:
We can drop this comment while you are in there please.
##########
providers/amazon/tests/unit/amazon/aws/executors/batch/test_batch_executor.py:
##########
@@ -269,6 +270,98 @@ def test_task_sdk(self, running_state_mock,
mock_airflow_key, mock_executor, moc
assert job_id == ARN1
running_state_mock.assert_called_once_with(workload.ti.key, ARN1)
+ @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="Test requires Airflow
3.2+")
Review Comment:
Here and below, update the version in the reason
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/batch/batch_executor.py:
##########
@@ -337,18 +357,18 @@ def attempt_submit_jobs(self):
self.log.error(
(
"This job has been unsuccessfully attempted too
many times (%s). "
- "Dropping the task. Reason: %s"
+ "Dropping the workload. Reason: %s"
),
attempt_number,
failure_reason,
)
self.log_task_event(
event="batch job submit failure",
extra=f"This job has been unsuccessfully attempted too
many times ({attempt_number}). "
- f"Dropping the task. Reason: {failure_reason}",
- ti_key=key,
+ f"Dropping the workload. Reason: {failure_reason}",
+ ti_key=workload_key,
)
- self.fail(key=key)
+ self.fail(key=workload_key)
Review Comment:
Yup, rebasing this PR should resolve this. It looks like it's pretty far
behind main and needs a rebase anyway.
--
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]
