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 335fd71f58f Stop AWS Glue job run when a deferred task is cleared
(#73702)
335fd71f58f is described below
commit 335fd71f58f0cefff3e5c41b0191c751fdc256c5
Author: leon.jeon <[email protected]>
AuthorDate: Sat Nov 7 01:42:50 2026 +0900
Stop AWS Glue job run when a deferred task is cleared (#73702)
stop_job_run_on_kill was only honoured by the operator's synchronous
on_kill, which never runs for a deferred task, so clearing, failing or
succeeding a deferred Glue task left its job run going.
The trigger's on_kill fires on every user action and cannot tell a clear,
which re-runs the task, from mark failed or success. A cleared task's new
try would otherwise reconnect (durable) to the run that on_kill is
stopping and then fail on STOPPED. So when stop_job_run_on_kill is set,
the retry stops the previous run itself and waits for it to end before
submitting a new one: a STOPPING run still counts toward
MaxConcurrentRuns, so submitting earlier would fail.
---
.../src/airflow/providers/amazon/aws/exceptions.py | 4 +
.../airflow/providers/amazon/aws/operators/glue.py | 51 +++++++-
.../airflow/providers/amazon/aws/triggers/glue.py | 30 ++++-
.../tests/unit/amazon/aws/operators/test_glue.py | 141 ++++++++++++++++++++-
.../tests/unit/amazon/aws/triggers/test_glue.py | 41 ++++++
5 files changed, 262 insertions(+), 5 deletions(-)
diff --git a/providers/amazon/src/airflow/providers/amazon/aws/exceptions.py
b/providers/amazon/src/airflow/providers/amazon/aws/exceptions.py
index 7a59aa0fe2f..4c88b2e6a24 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/exceptions.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/exceptions.py
@@ -86,6 +86,10 @@ class GlueJobRunStoppedError(AirflowException):
"""Raised when a Glue job run finishes in a state that is not a real
success."""
+class GlueJobRunStopTimeoutError(AirflowException):
+ """Raised when a previous Glue job run does not reach a terminal state
after being stopped."""
+
+
class DataSyncTaskNotFoundError(AirflowException):
"""Raised when a DataSync task could not be identified or created for the
requested locations."""
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py
b/providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py
index 94e4d7b93cc..fcc63c788be 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py
@@ -18,6 +18,7 @@
from __future__ import annotations
import os
+import time
import urllib.parse
import warnings
from collections.abc import Sequence
@@ -26,7 +27,7 @@ from typing import TYPE_CHECKING, Any, cast
from botocore.exceptions import BotoCoreError, ClientError
from airflow.exceptions import AirflowProviderDeprecationWarning
-from airflow.providers.amazon.aws.exceptions import GlueJobRunStoppedError
+from airflow.providers.amazon.aws.exceptions import GlueJobRunStoppedError,
GlueJobRunStopTimeoutError
from airflow.providers.amazon.aws.hooks.glue import GlueDataQualityHook,
GlueJobHook
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
from airflow.providers.amazon.aws.links.glue import GlueJobRunDetailsLink
@@ -60,6 +61,8 @@ JOB_RUN_SUCCESS_STATES = ("SUCCEEDED",)
JOB_RUN_TERMINAL_STATES = (*JOB_RUN_SUCCESS_STATES, "STOPPED", "FAILED",
"TIMEOUT", "ERROR", "EXPIRED")
# Synthetic state for a run id Glue no longer knows about, so a retry submits
fresh rather than failing.
NOT_FOUND_STATE = "NOT_FOUND"
+# Upper bound on waiting for a previous run to leave STOPPING before a fresh
run is submitted.
+PREVIOUS_JOB_RUN_STOP_TIMEOUT_SECONDS = 600
class GlueJobOperator(ResumableJobMixin, AwsBaseOperator[GlueJobHook]):
@@ -94,7 +97,9 @@ class GlueJobOperator(ResumableJobMixin,
AwsBaseOperator[GlueJobHook]):
:param verbose: If True, Glue Job Run logs show in the Airflow Task Logs.
(default: False)
:param update_config: If True, Operator will update job configuration.
(default: False)
:param replace_script_file: If True, the script file will be replaced in
S3. (default: False)
- :param stop_job_run_on_kill: If True, Operator will stop the job run when
task is killed.
+ :param stop_job_run_on_kill: If True, Operator will stop the job run when
task is killed. In
+ deferrable mode this requires Airflow 3.3+, and with ``durable`` a
retried or cleared task stops
+ its previous run and waits for it to end before starting a new one,
instead of reconnecting to it.
:param sleep_before_return: time in seconds to wait before returning final
status. This is meaningful in case
of limiting concurrency, Glue needs 5-10 seconds to clean up resources.
Thus if status is returned immediately it might end up in case of more
than 1 concurrent run.
@@ -342,6 +347,7 @@ class GlueJobOperator(ResumableJobMixin,
AwsBaseOperator[GlueJobHook]):
region_name=self.region_name,
verify=self.verify,
botocore_config=self.botocore_config,
+ stop_job_run_on_kill=self.stop_job_run_on_kill,
),
method_name="execute_complete",
)
@@ -451,13 +457,26 @@ class GlueJobOperator(ResumableJobMixin,
AwsBaseOperator[GlueJobHook]):
def submit_job(self, context: Context) -> str:
"""Start a Glue job run and return its run id, or reconnect to one
this task already started."""
script_args = self._build_script_args(context)
+ if (
+ AIRFLOW_V_3_3_PLUS
+ and self.deferrable
+ and self.durable
+ and self.stop_job_run_on_kill
+ and context["ti"].try_number > 1
+ ):
+ # The trigger's on_kill stops the previous run on any user action,
and it cannot tell a
+ # clear (which re-runs the task) from mark failed/success.
Reattaching would wait on a
+ # run that on_kill is stopping, so stop it here and start fresh
instead.
+ previous_job_run_id = self._find_previous_job_run(context,
script_args[self.TASK_UUID_ARG])
+ if previous_job_run_id:
+ self._stop_job_run_and_wait(previous_job_run_id)
# Scan only when there's nothing else to go on: first attempt, no
store, or a store that
# never recorded this key. The store check itself is skipped on the
deferrable path, where
# nothing is ever written to it -- a store error there would otherwise
block the scan.
# Deferrable retries could reconnect via task_state_store too,
avoiding this scan entirely,
# if ResumableJobMixin exposed its reconnect decision apart from its
polling loop;
# tracked at https://github.com/apache/airflow/issues/71485.
- if (
+ elif (
self.durable
and (self.deferrable or not AIRFLOW_V_3_3_PLUS)
and context["ti"].try_number > 1
@@ -480,6 +499,32 @@ class GlueJobOperator(ResumableJobMixin,
AwsBaseOperator[GlueJobHook]):
self._set_job_run_id(context, glue_job_run["JobRunId"])
return glue_job_run["JobRunId"]
+ def _stop_job_run_and_wait(self, job_run_id: str) -> None:
+ """
+ Stop a job run and wait until it is terminal.
+
+ A STOPPING run still counts toward MaxConcurrentRuns, so submitting
before then would fail.
+ """
+ self.log.info(
+ "Stopping previous AWS Glue job %s run %s before starting a new
one.", self.job_name, job_run_id
+ )
+ response = self.hook.conn.batch_stop_job_run(JobName=self.job_name,
JobRunIds=[job_run_id])
+ if response.get("Errors"):
+ # The trigger's on_kill may already have stopped it, or it may
have just finished.
+ self.log.info("Stop request for run %s returned errors: %s",
job_run_id, response["Errors"])
+ deadline = time.monotonic() + PREVIOUS_JOB_RUN_STOP_TIMEOUT_SECONDS
+ while True:
+ state = self.hook.get_job_state(self.job_name, job_run_id)
+ if state in JOB_RUN_TERMINAL_STATES:
+ self.log.info("Previous AWS Glue job run %s reached %s.",
job_run_id, state)
+ return
+ if time.monotonic() >= deadline:
+ raise GlueJobRunStopTimeoutError(
+ f"Previous AWS Glue job {self.job_name} run {job_run_id}
is still {state} "
+ f"{PREVIOUS_JOB_RUN_STOP_TIMEOUT_SECONDS}s after being
stopped; not starting a new run."
+ )
+ time.sleep(self.job_poll_interval)
+
def get_job_status(self, external_id: JsonValue, context: Context) -> str:
"""Query the raw job run state; a run id Glue no longer knows about
degrades to NOT_FOUND."""
job_run_id = cast("str", external_id)
diff --git a/providers/amazon/src/airflow/providers/amazon/aws/triggers/glue.py
b/providers/amazon/src/airflow/providers/amazon/aws/triggers/glue.py
index 42c0e501dec..026626ff522 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/triggers/glue.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/glue.py
@@ -50,6 +50,8 @@ class GlueJobCompleteTrigger(AwsBaseWaiterTrigger):
if not specified.
:param verify: Whether or not to verify SSL certificates.
:param botocore_config: Configuration dictionary (key-values) for botocore
client.
+ :param stop_job_run_on_kill: If True, stop the Glue job run when the
deferred task is killed
+ (for example, cleared while running). Defaults to False.
"""
aws_hook_class = GlueJobHook
@@ -65,9 +67,15 @@ class GlueJobCompleteTrigger(AwsBaseWaiterTrigger):
region_name: str | None = None,
verify: bool | str | None = None,
botocore_config: dict | None = None,
+ stop_job_run_on_kill: bool = False,
):
super().__init__(
- serialized_fields={"job_name": job_name, "run_id": run_id,
"verbose": verbose},
+ serialized_fields={
+ "job_name": job_name,
+ "run_id": run_id,
+ "verbose": verbose,
+ "stop_job_run_on_kill": stop_job_run_on_kill,
+ },
waiter_name="job_complete",
waiter_args={"JobName": job_name, "RunId": run_id},
failure_message="AWS Glue job failed.",
@@ -85,6 +93,26 @@ class GlueJobCompleteTrigger(AwsBaseWaiterTrigger):
self.job_name = job_name
self.run_id = run_id
self.verbose = verbose
+ self.stop_job_run_on_kill = stop_job_run_on_kill
+
+ async def on_kill(self) -> None:
+ """
+ Stop the Glue job run when the user acts on the deferred task.
+
+ Requires Airflow 3.3+, where the triggerer calls
``BaseTrigger.on_kill()``; on older
+ versions this hook is inert. Released 3.3.x versions also call it when
the trigger is
+ reassigned to another triggerer (fixed in #73454), which stops a run
the new triggerer
+ is still watching.
+ """
+ if not self.stop_job_run_on_kill or not self.run_id:
+ return
+ self.log.info("Stopping AWS Glue job %s run %s.", self.job_name,
self.run_id)
+ async with await self.hook().get_async_conn() as client:
+ response = await client.batch_stop_job_run(JobName=self.job_name,
JobRunIds=[self.run_id])
+ if response.get("Errors"):
+ self.log.error(
+ "Failed to stop AWS Glue job %s run %s: %s", self.job_name,
self.run_id, response["Errors"]
+ )
async def run(self) -> AsyncIterator[TriggerEvent]:
if not self.verbose:
diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_glue.py
b/providers/amazon/tests/unit/amazon/aws/operators/test_glue.py
index 56c0f89a1a7..1e2da4a8c57 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_glue.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_glue.py
@@ -30,7 +30,7 @@ from moto import mock_aws
from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.models.dag import DAG
-from airflow.providers.amazon.aws.exceptions import GlueJobRunStoppedError
+from airflow.providers.amazon.aws.exceptions import GlueJobRunStoppedError,
GlueJobRunStopTimeoutError
from airflow.providers.amazon.aws.hooks.glue import GlueDataQualityHook,
GlueJobHook
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
from airflow.providers.amazon.aws.links.glue import GlueJobRunDetailsLink
@@ -233,6 +233,28 @@ class TestGlueJobOperator:
assert serialized["verify"] is False
assert serialized["botocore_config"] == botocore_config
+ @mock.patch.object(GlueJobHook, "initialize_job")
+ @mock.patch.object(GlueJobHook, "get_conn")
+ def test_execute_deferrable_forwards_stop_job_run_on_kill_to_trigger(self,
_, mock_initialize_job):
+ """stop_job_run_on_kill must reach the trigger so a cleared deferred
task can stop the run."""
+ glue = GlueJobOperator(
+ durable=False,
+ task_id=TASK_ID,
+ job_name=JOB_NAME,
+ script_location="s3://folder/file",
+ s3_bucket="some_bucket",
+ iam_role_name="my_test_role",
+ deferrable=True,
+ stop_job_run_on_kill=True,
+ )
+ mock_initialize_job.return_value = {"JobRunState": "RUNNING",
"JobRunId": JOB_RUN_ID}
+
+ with pytest.raises(TaskDeferred) as defer:
+ glue.execute({"ti": mock.MagicMock(try_number=1)})
+
+ _, serialized = defer.value.trigger.serialize()
+ assert serialized["stop_job_run_on_kill"] is True
+
@mock.patch.object(GlueJobHook, "conn", new_callable=mock.PropertyMock)
@mock.patch.object(GlueJobHook, "initialize_job")
@mock.patch.object(GlueJobHook, "get_conn")
@@ -405,6 +427,123 @@ class TestGlueJobOperator:
mock_initialize_job.assert_not_called()
+ @staticmethod
+ def _make_retry_with_running_previous_run(glue: GlueJobOperator) ->
mock.MagicMock:
+ mock_ti = mock.MagicMock()
+ mock_ti.dag_id = DAG_ID
+ mock_ti.task_id = TASK_ID
+ mock_ti.run_id = "manual__2024-01-01T00:00:00+00:00"
+ mock_ti.map_index = -1
+ mock_ti.try_number = 2
+ mock_ti.xcom_pull.return_value = None
+ task_uuid =
f"{mock_ti.dag_id}:{mock_ti.task_id}:{mock_ti.run_id}:{mock_ti.map_index}"
+ glue.hook.conn = mock.MagicMock()
+ glue.hook.conn.get_job_runs.return_value = {
+ "JobRuns": [
+ {
+ "Id": JOB_RUN_ID,
+ "Arguments": {GlueJobOperator.TASK_UUID_ARG: task_uuid},
+ "JobRunState": "RUNNING",
+ }
+ ]
+ }
+ return mock_ti
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="the trigger's on_kill
only runs on Airflow 3.3+")
+ @pytest.mark.parametrize(
+ ("stop_response", "states_after_stop"),
+ [
+ pytest.param(
+ {"SuccessfulSubmissions": [{"JobRunId": JOB_RUN_ID}]},
+ ["STOPPING", "STOPPED"],
+ id="stopped_here",
+ ),
+ pytest.param(
+ {"Errors": [{"JobRunId": JOB_RUN_ID}]}, ["STOPPED"],
id="already_stopped_by_on_kill"
+ ),
+ ],
+ )
+ @mock.patch("airflow.providers.amazon.aws.operators.glue.time.sleep")
+ @mock.patch.object(GlueJobHook, "get_job_state")
+ @mock.patch.object(GlueJobHook, "initialize_job")
+ @mock.patch.object(GlueJobHook, "get_conn")
+ def
test_deferrable_retry_with_stop_job_run_on_kill_stops_previous_run_and_starts_fresh(
+ self,
+ mock_get_conn,
+ mock_initialize_job,
+ mock_get_job_state,
+ mock_sleep,
+ stop_response,
+ states_after_stop,
+ ):
+ glue = GlueJobOperator(
+ task_id=TASK_ID,
+ job_name=JOB_NAME,
+ script_location="s3://folder/file",
+ deferrable=True,
+ durable=True,
+ stop_job_run_on_kill=True,
+ )
+ mock_ti = self._make_retry_with_running_previous_run(glue)
+ glue.hook.conn.batch_stop_job_run.return_value = stop_response
+ mock_get_job_state.side_effect = states_after_stop
+ mock_initialize_job.return_value = {"JobRunState": "RUNNING",
"JobRunId": "22222"}
+
+ with pytest.raises(TaskDeferred) as defer:
+ glue.execute({"ti": mock_ti})
+
+
glue.hook.conn.batch_stop_job_run.assert_called_once_with(JobName=JOB_NAME,
JobRunIds=[JOB_RUN_ID])
+ assert mock_get_job_state.call_count == len(states_after_stop)
+ mock_initialize_job.assert_called_once()
+ assert defer.value.trigger.run_id == "22222"
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="the trigger's on_kill
only runs on Airflow 3.3+")
+
@mock.patch("airflow.providers.amazon.aws.operators.glue.PREVIOUS_JOB_RUN_STOP_TIMEOUT_SECONDS",
0)
+ @mock.patch.object(GlueJobHook, "get_job_state", return_value="STOPPING")
+ @mock.patch.object(GlueJobHook, "initialize_job")
+ @mock.patch.object(GlueJobHook, "get_conn")
+ def
test_deferrable_retry_with_stop_job_run_on_kill_raises_when_previous_run_does_not_stop(
+ self, mock_get_conn, mock_initialize_job, mock_get_job_state
+ ):
+ glue = GlueJobOperator(
+ task_id=TASK_ID,
+ job_name=JOB_NAME,
+ script_location="s3://folder/file",
+ deferrable=True,
+ stop_job_run_on_kill=True,
+ )
+ mock_ti = self._make_retry_with_running_previous_run(glue)
+ glue.hook.conn.batch_stop_job_run.return_value = {}
+
+ with pytest.raises(GlueJobRunStopTimeoutError, match="is still
STOPPING"):
+ glue.execute({"ti": mock_ti})
+
+ mock_initialize_job.assert_not_called()
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="the trigger's on_kill
only runs on Airflow 3.3+")
+ @mock.patch.object(GlueJobHook, "initialize_job")
+ @mock.patch.object(GlueJobHook, "get_conn")
+ def
test_deferrable_retry_with_stop_job_run_on_kill_without_durable_starts_fresh(
+ self, mock_get_conn, mock_initialize_job
+ ):
+ glue = GlueJobOperator(
+ task_id=TASK_ID,
+ job_name=JOB_NAME,
+ script_location="s3://folder/file",
+ deferrable=True,
+ durable=False,
+ stop_job_run_on_kill=True,
+ )
+ mock_ti = self._make_retry_with_running_previous_run(glue)
+ mock_initialize_job.return_value = {"JobRunState": "RUNNING",
"JobRunId": "22222"}
+
+ with pytest.raises(TaskDeferred) as defer:
+ glue.execute({"ti": mock_ti})
+
+ glue.hook.conn.batch_stop_job_run.assert_not_called()
+ assert GlueJobOperator.TASK_UUID_ARG not in
mock_initialize_job.call_args[0][0]
+ assert defer.value.trigger.run_id == "22222"
+
@pytest.mark.skipif(
not AIRFLOW_V_3_3_PLUS,
reason="task_state_store only exists as an execute() context key on
Airflow 3.3+",
diff --git a/providers/amazon/tests/unit/amazon/aws/triggers/test_glue.py
b/providers/amazon/tests/unit/amazon/aws/triggers/test_glue.py
index e4de3b7e994..add8dbb0b6e 100644
--- a/providers/amazon/tests/unit/amazon/aws/triggers/test_glue.py
+++ b/providers/amazon/tests/unit/amazon/aws/triggers/test_glue.py
@@ -108,6 +108,7 @@ class TestGlueJobTrigger:
"job_name": "job_name",
"run_id": "JobRunId",
"verbose": False,
+ "stop_job_run_on_kill": False,
"aws_conn_id": "aws_conn_id",
"waiter_max_attempts": 3,
"waiter_delay": 10,
@@ -357,6 +358,46 @@ class TestGlueJobTrigger:
log_output = mock_log_info.call_args[0][0]
assert "No new log from the Glue Job in /aws-glue/python-jobs/output"
in log_output
+ def test_serialization_includes_stop_job_run_on_kill(self):
+ trigger = GlueJobCompleteTrigger(job_name="job_name",
run_id="JobRunId", stop_job_run_on_kill=True)
+ _, kwargs = trigger.serialize()
+ assert kwargs["stop_job_run_on_kill"] is True
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize(
+ ("response", "expect_error_log"),
+ [
+ (
+ {"SuccessfulSubmissions": [{"JobName": "job_name", "JobRunId":
"JobRunId"}], "Errors": []},
+ False,
+ ),
+ (
+ {"SuccessfulSubmissions": [], "Errors": [{"JobName":
"job_name", "JobRunId": "JobRunId"}]},
+ True,
+ ),
+ ],
+ )
+ @mock.patch.object(GlueJobHook, "get_async_conn")
+ async def test_on_kill_stops_job_run_when_enabled(self, mock_glue_conn,
response, expect_error_log):
+ glue_client = AsyncMock()
+ glue_client.batch_stop_job_run = AsyncMock(return_value=response)
+ mock_glue_conn.return_value.__aenter__ =
AsyncMock(return_value=glue_client)
+ mock_glue_conn.return_value.__aexit__ = AsyncMock(return_value=False)
+ trigger = GlueJobCompleteTrigger(job_name="job_name",
run_id="JobRunId", stop_job_run_on_kill=True)
+
+ with mock.patch.object(trigger.log, "error") as mock_log_error:
+ await trigger.on_kill()
+
+
glue_client.batch_stop_job_run.assert_awaited_once_with(JobName="job_name",
JobRunIds=["JobRunId"])
+ assert mock_log_error.called is expect_error_log
+
+ @pytest.mark.asyncio
+ @mock.patch.object(GlueJobHook, "get_async_conn")
+ async def test_on_kill_does_not_stop_when_disabled(self, mock_glue_conn):
+ trigger = GlueJobCompleteTrigger(job_name="job_name",
run_id="JobRunId", stop_job_run_on_kill=False)
+ await trigger.on_kill()
+ mock_glue_conn.assert_not_called()
+
class TestGlueCatalogPartitionSensorTrigger:
@pytest.mark.asyncio
