This is an automated email from the ASF dual-hosted git repository.
o-nikolas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 7c34d3097eb Prevent a broken job from stalling AWS Batch executor sync
(#71380)
7c34d3097eb is described below
commit 7c34d3097eba77910ec7ea4b121741704b890e50
Author: Alexander Oreshkevich <[email protected]>
AuthorDate: Tue Aug 18 06:19:16 2026 +0700
Prevent a broken job from stalling AWS Batch executor sync (#71380)
A deterministic per-job error in sync_running_jobs() aborts every
heartbeat before attempt_submit_jobs(), halting all task submission
until the scheduler is restarted. Seen in production: a job whose
bookkeeping was partially cleaned made pop_by_id() raise KeyError
forever, wedging the whole executor.
---
.../amazon/aws/executors/batch/batch_executor.py | 37 ++++-
.../providers/amazon/aws/executors/batch/utils.py | 14 ++
.../aws/executors/batch/test_batch_executor.py | 168 +++++++++++++++++++++
3 files changed, 214 insertions(+), 5 deletions(-)
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/executors/batch/batch_executor.py
b/providers/amazon/src/airflow/providers/amazon/aws/executors/batch/batch_executor.py
index 00c7dcb7881..454ba38c0a4 100644
---
a/providers/amazon/src/airflow/providers/amazon/aws/executors/batch/batch_executor.py
+++
b/providers/amazon/src/airflow/providers/amazon/aws/executors/batch/batch_executor.py
@@ -268,11 +268,38 @@ class AwsBatchExecutor(BaseExecutor):
self.log.debug("Active Workers: %s", describe_job_response)
for job in describe_job_response:
- if job.get_job_state() == State.FAILED:
- self._handle_failed_job(job)
- elif job.get_job_state() == State.SUCCESS:
- workload_key = self.active_workers.pop_by_id(job.job_id)
- self.success(workload_key)
+ # snapshot the key before handling the job: if the error strikes
after
+ # pop_by_id() already removed it, the collection can no longer
tell us
+ # whose workload it was, and the workload would be left with no
terminal state
+ workload_key = self.active_workers.id_to_key.get(job.job_id)
+ try:
+ if job.get_job_state() == State.FAILED:
+ self._handle_failed_job(job)
+ elif job.get_job_state() == State.SUCCESS:
+ self.success(self.active_workers.pop_by_id(job.job_id))
+ except (ClientError, NoCredentialsError):
+ # credential problems are executor-wide, not job-specific: let
sync() handle them
+ raise
+ except Exception:
+ self.log.exception(
+ "Evicting Batch job %s after an unexpected error while
syncing it.", job.job_id
+ )
+ self._evict_job(job.job_id, workload_key)
+
+ def _evict_job(self, job_id: str, workload_key: BatchJobWorkloadKey | None
= None) -> None:
+ """
+ Remove a job from the collection and fail its workload, tolerating
corrupted bookkeeping.
+
+ workload_key is the caller's snapshot of the mapping taken before the
error;
+ it is the fallback when the job was already removed from the
collection.
+ """
+ workload_key = self.active_workers.remove_job(job_id) or workload_key
+ if workload_key is None:
+ return
+ try:
+ self.fail(workload_key)
+ except Exception:
+ self.log.exception("Failed to fail workload %s of evicted Batch
job %s", workload_key, job_id)
def _handle_failed_job(self, job):
"""
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/executors/batch/utils.py
b/providers/amazon/src/airflow/providers/amazon/aws/executors/batch/utils.py
index 64685184902..d47fc93ff5b 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/executors/batch/utils.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/executors/batch/utils.py
@@ -128,6 +128,20 @@ class BatchJobCollection:
del self.id_to_failure_counts[job_id]
return workload_key
+ def remove_job(self, job_id: str) -> BatchJobWorkloadKey | None:
+ """
+ Remove a job from the collection, tolerating partially cleaned
bookkeeping.
+
+ Unlike pop_by_id, this never raises: every mapping is popped
defensively.
+ Returns the workload key if the job was still tracked, otherwise None.
+ """
+ workload_key = self.id_to_key.pop(job_id, None)
+ self.id_to_failure_counts.pop(job_id, None)
+ self.id_to_job_info.pop(job_id, None)
+ if workload_key is not None:
+ self.key_to_id.pop(workload_key, None)
+ return workload_key
+
def failure_count_by_id(self, job_id: str) -> int:
"""Get the number of times a job has failed given a Batch Job Id."""
return self.id_to_failure_counts[job_id]
diff --git
a/providers/amazon/tests/unit/amazon/aws/executors/batch/test_batch_executor.py
b/providers/amazon/tests/unit/amazon/aws/executors/batch/test_batch_executor.py
index 4bfd7c716ef..50a90757eae 100644
---
a/providers/amazon/tests/unit/amazon/aws/executors/batch/test_batch_executor.py
+++
b/providers/amazon/tests/unit/amazon/aws/executors/batch/test_batch_executor.py
@@ -146,6 +146,21 @@ class TestBatchJobCollection:
assert len(self.collection) == 1
assert self.collection.get_all_jobs() == [self.second_job_id]
+ def test_remove_job(self):
+ """Test remove_job() clears every mapping and returns the workload
key"""
+ assert self.collection.remove_job(self.first_job_id) is
self.first_airflow_key
+ assert self.collection.get_all_jobs() == [self.second_job_id]
+ assert self.first_job_id not in self.collection.id_to_failure_counts
+ assert self.first_job_id not in self.collection.id_to_job_info
+ assert self.first_airflow_key not in self.collection.key_to_id
+
+ def test_remove_job_tolerates_partially_cleaned_state(self):
+ """remove_job() must not raise where pop_by_id() would"""
+ del self.collection.key_to_id[self.first_airflow_key]
+ assert self.collection.remove_job(self.first_job_id) is
self.first_airflow_key
+ assert self.first_job_id not in self.collection.id_to_key
+ assert self.collection.remove_job("untracked-job-id") is None
+
class TestBatchJob:
"""Tests the BatchJob DTO"""
@@ -644,6 +659,159 @@ class TestAwsBatchExecutor:
fail_mock.assert_called_once()
assert success_mock.call_count == 0
+ @staticmethod
+ def _describe_jobs_response(*job_id_status_pairs: tuple[str, str]) -> dict:
+ return {
+ "jobs": [
+ {
+ "jobName": "some-job-name",
+ "jobId": job_id,
+ "jobQueue": "some-job-queue",
+ "status": status,
+ "statusReason": "",
+ "createdAt": dt.datetime.now().timestamp(),
+ "jobDefinition": "some-job-def",
+ }
+ for job_id, status in job_id_status_pairs
+ ]
+ }
+
+ @mock.patch.object(BaseExecutor, "fail")
+ @mock.patch.object(BaseExecutor, "success")
+ def test_sync_evicts_job_that_cannot_be_synced(self, success_mock,
fail_mock, mock_executor):
+ """
+ An error while handling one job must evict that job and fail its
workload,
+ without blocking the handling of the remaining jobs in the same cycle.
+ """
+ poisoned_key = mock.Mock(spec=TaskInstanceKey)
+ healthy_key = mock.Mock(spec=TaskInstanceKey)
+ for job_id, key in (("001", poisoned_key), ("002", healthy_key)):
+ mock_executor.active_workers.add_job(
+ job_id=job_id,
+ airflow_workload_key=key,
+ airflow_cmd="airflow_cmd",
+ queue="queue",
+ exec_config={},
+ attempt_number=1,
+ )
+ # Partially cleaned bookkeeping (e.g. left behind by a duplicate
submission for
+ # the same workload key) makes pop_by_id raise KeyError for this job
forever.
+ del mock_executor.active_workers.key_to_id[poisoned_key]
+ mock_executor.batch.describe_jobs.return_value =
self._describe_jobs_response(
+ ("001", "SUCCEEDED"), ("002", "SUCCEEDED")
+ )
+
+ mock_executor.sync_running_jobs()
+
+ fail_mock.assert_called_once_with(poisoned_key)
+ success_mock.assert_called_once_with(healthy_key)
+ assert len(mock_executor.active_workers) == 0
+ assert mock_executor.active_workers.get_all_jobs() == []
+
+ @mock.patch.object(BaseExecutor, "fail")
+ def test_sync_eviction_survives_fail_raising(self, fail_mock,
mock_executor, mock_airflow_key):
+ """Even fail() blowing up during eviction must not leave the job
tracked."""
+ airflow_key = mock_airflow_key()
+ mock_executor.active_workers.add_job(
+ job_id="001",
+ airflow_workload_key=airflow_key,
+ airflow_cmd="airflow_cmd",
+ queue="queue",
+ exec_config={},
+ attempt_number=1,
+ )
+ del mock_executor.active_workers.key_to_id[airflow_key]
+ mock_executor.batch.describe_jobs.return_value =
self._describe_jobs_response(("001", "SUCCEEDED"))
+ fail_mock.side_effect = RuntimeError("fail() also broken")
+
+ mock_executor.sync_running_jobs()
+
+ assert mock_executor.active_workers.get_all_jobs() == []
+ assert "001" not in mock_executor.active_workers.id_to_key
+
+ @mock.patch.object(BaseExecutor, "fail")
+ @mock.patch.object(BaseExecutor, "success")
+ def test_sync_eviction_fails_workload_already_removed_from_collection(
+ self, success_mock, fail_mock, mock_executor, mock_airflow_key
+ ):
+ """
+ When the error strikes after pop_by_id() already removed the job (here
success()
+ itself raising), the collection can no longer resolve the workload
key, so the
+ key snapshotted before handling must be used to fail the workload
instead of
+ leaving it with no terminal state.
+ """
+ airflow_key = mock_airflow_key()
+ mock_executor.active_workers.add_job(
+ job_id="001",
+ airflow_workload_key=airflow_key,
+ airflow_cmd="airflow_cmd",
+ queue="queue",
+ exec_config={},
+ attempt_number=1,
+ )
+ mock_executor.batch.describe_jobs.return_value =
self._describe_jobs_response(("001", "SUCCEEDED"))
+ success_mock.side_effect = RuntimeError("state change failed")
+
+ mock_executor.sync_running_jobs()
+
+ fail_mock.assert_called_once_with(airflow_key)
+ assert mock_executor.active_workers.get_all_jobs() == []
+ assert airflow_key not in mock_executor.active_workers.key_to_id
+
+ @mock.patch.object(BaseExecutor, "fail")
+ @mock.patch.object(BaseExecutor, "success")
+ def test_sync_reraises_credential_errors_instead_of_evicting(
+ self, success_mock, fail_mock, mock_executor, mock_airflow_key
+ ):
+ """
+ Credential problems are executor-wide, not job-specific: they must
propagate to
+ sync()'s connection handling rather than evict the job they happened
to surface on.
+ """
+ airflow_key = mock_airflow_key()
+ mock_executor.active_workers.add_job(
+ job_id="001",
+ airflow_workload_key=airflow_key,
+ airflow_cmd="airflow_cmd",
+ queue="queue",
+ exec_config={},
+ attempt_number=1,
+ )
+ mock_executor.batch.describe_jobs.return_value =
self._describe_jobs_response(("001", "SUCCEEDED"))
+ success_mock.side_effect = ClientError(
+ {"Error": {"Code": "ExpiredTokenException", "Message": "token
expired"}}, "SomeBotoOperation"
+ )
+
+ with pytest.raises(ClientError):
+ mock_executor.sync_running_jobs()
+
+ fail_mock.assert_not_called()
+
+ @mock.patch.object(AwsBatchExecutor, "attempt_submit_jobs")
+ @mock.patch.object(BaseExecutor, "fail")
+ def test_sync_submits_pending_jobs_despite_poisoned_job(
+ self, fail_mock, attempt_submit_jobs_mock, mock_executor,
mock_airflow_key
+ ):
+ """
+ The reported production impact: a job that deterministically errors
during sync
+ used to abort sync() before attempt_submit_jobs(), halting all task
submission
+ until a scheduler restart.
+ """
+ airflow_key = mock_airflow_key()
+ mock_executor.active_workers.add_job(
+ job_id="001",
+ airflow_workload_key=airflow_key,
+ airflow_cmd="airflow_cmd",
+ queue="queue",
+ exec_config={},
+ attempt_number=1,
+ )
+ del mock_executor.active_workers.key_to_id[airflow_key]
+ mock_executor.batch.describe_jobs.return_value =
self._describe_jobs_response(("001", "SUCCEEDED"))
+
+ mock_executor.sync()
+
+ attempt_submit_jobs_mock.assert_called_once()
+
def test_start_failure_with_invalid_permissions(self, set_env_vars):
executor = AwsBatchExecutor()
