amoghrajesh commented on code in PR #71211:
URL: https://github.com/apache/airflow/pull/71211#discussion_r3756745715
##########
providers/amazon/docs/operators/glue.rst:
##########
@@ -139,6 +139,79 @@ To submit a new AWS Glue job you can use
:class:`~airflow.providers.amazon.aws.o
The same AWS IAM role used for the crawler can be used here as well, but it
will need
policies to provide access to the output location for result data.
+Durable execution
+==================
+
+``GlueJobOperator`` submits a job run and then polls it to completion on the
worker. By default
+the operator runs in a *durable* mode that makes this crash-safe: the Glue job
run id is
+persisted to :doc:`task state store
<apache-airflow:core-concepts/task-state-store>` before
+polling begins, so if the worker crashes or is preempted and the task is
retried, the operator
+reconnects to the run that is already executing in Glue instead of starting a
new one.
+
+This matters more for Glue because a Glue job's ``concurrent_run_limit``
defaults to ``1``, so
+submitting a second run while the first is still active does not create a
harmless duplicate, it
+fails outright with ``ConcurrentRunsExceededException`` and the task keeps
retrying against a run
+it can never see. Durable execution turns that retry into a normal reconnect.
+
+On retry the operator checks the prior run's state:
+
+* if it is still starting, running, waiting for capacity, or being stopped,
the operator
+ reconnects and continues polling
+* if it already succeeded, the operator returns immediately without
resubmitting
+* if it stopped, failed terminally, or its id has expired and is no longer
found, the operator
+ submits the job fresh
+
+A stopped run resubmits rather than being treated as a success. Glue's API has
no way to tell a
+run that was cancelled manually (for example, in the AWS console) apart from
one this operator's
+own
:meth:`~airflow.providers.amazon.aws.operators.glue.GlueJobOperator.on_kill`
stopped, which
+happens whenever ``stop_job_run_on_kill=True`` and the task is killed -- on
SIGTERM, on
+``execution_timeout``, or when the task is cleared while running. Since the
two cases can't be
+told apart, a stopped run always resubmits, so a self-inflicted stop never
gets silently reported
+as a false success.
+
+This protection also applies when ``wait_for_completion=False`` -- even though
that task attempt
+never polls at all, a retry after a successful submission still reconnects
rather than
+resubmitting, since the run id is persisted immediately after submission
regardless of whether the
+task waits for it to finish.
+
+Durable execution requires Airflow 3.3 or newer for the task state store
lookup above. On earlier
+Airflow versions, or if the task state store is unavailable at runtime,
``durable=True`` still
+recovers a prior run, just via an older mechanism: the operator checks XCom
for a cached run id
+first, then falls back to scanning the job's run history for a run tagged with
this task
+instance's identity, and reconnects if it finds one that is still active.
+
+Like the persisted state itself, the stored run id isn't deleted
automatically, that only happens
+when someone runs ``airflow state-store clean``. If a task's ``retry_delay``
is longer than
+``[state_store] default_retention_days`` (30 days by default) and cleanup runs
in between, the run
+id won't be there for the next retry, and the operator falls back to the
XCom/scan mechanism
+above rather than reconnecting via task state store. Avoid running cleanup on
a schedule shorter
+than your longest ``retry_delay``.
+
+Clearing a task is treated the same as a retry, which matters specifically for
a task whose job
+already succeeded: clearing does not delete the stored run id, so the next
attempt reads it back
+and returns immediately without submitting anything to Glue. See
+:doc:`apache-airflow:core-concepts/resumable-tasks` for why, and for the
+``[state_store] clear_on_success`` setting that restores "clearing always
resubmits."
+
+To opt out and always start a fresh run on retry, set ``durable=False``:
+
+.. code-block:: python
+
+ glue_job = GlueJobOperator(
+ task_id="glue_job",
+ job_name="my_glue_job",
+ script_location="s3://glue-examples/glue-scripts/sample_aws_glue_job.py",
+ durable=False,
+ )
+
+Durable execution applies to the synchronous path. When ``deferrable=True`` is
set, the Triggerer
+already tracks the run across the wait, so deferrable mode takes precedence
and ``durable`` has no
Review Comment:
Handled in 118dacfd80
##########
providers/amazon/docs/operators/glue.rst:
##########
@@ -139,6 +139,79 @@ To submit a new AWS Glue job you can use
:class:`~airflow.providers.amazon.aws.o
The same AWS IAM role used for the crawler can be used here as well, but it
will need
policies to provide access to the output location for result data.
+Durable execution
+==================
+
+``GlueJobOperator`` submits a job run and then polls it to completion on the
worker. By default
+the operator runs in a *durable* mode that makes this crash-safe: the Glue job
run id is
+persisted to :doc:`task state store
<apache-airflow:core-concepts/task-state-store>` before
+polling begins, so if the worker crashes or is preempted and the task is
retried, the operator
+reconnects to the run that is already executing in Glue instead of starting a
new one.
+
+This matters more for Glue because a Glue job's ``concurrent_run_limit``
defaults to ``1``, so
+submitting a second run while the first is still active does not create a
harmless duplicate, it
+fails outright with ``ConcurrentRunsExceededException`` and the task keeps
retrying against a run
+it can never see. Durable execution turns that retry into a normal reconnect.
+
+On retry the operator checks the prior run's state:
+
+* if it is still starting, running, waiting for capacity, or being stopped,
the operator
+ reconnects and continues polling
+* if it already succeeded, the operator returns immediately without
resubmitting
+* if it stopped, failed terminally, or its id has expired and is no longer
found, the operator
+ submits the job fresh
+
+A stopped run resubmits rather than being treated as a success. Glue's API has
no way to tell a
+run that was cancelled manually (for example, in the AWS console) apart from
one this operator's
+own
:meth:`~airflow.providers.amazon.aws.operators.glue.GlueJobOperator.on_kill`
stopped, which
+happens whenever ``stop_job_run_on_kill=True`` and the task is killed -- on
SIGTERM, on
+``execution_timeout``, or when the task is cleared while running. Since the
two cases can't be
+told apart, a stopped run always resubmits, so a self-inflicted stop never
gets silently reported
+as a false success.
+
+This protection also applies when ``wait_for_completion=False`` -- even though
that task attempt
+never polls at all, a retry after a successful submission still reconnects
rather than
+resubmitting, since the run id is persisted immediately after submission
regardless of whether the
+task waits for it to finish.
+
+Durable execution requires Airflow 3.3 or newer for the task state store
lookup above. On earlier
+Airflow versions, or if the task state store is unavailable at runtime,
``durable=True`` still
Review Comment:
Handled in 118dacfd80
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py:
##########
@@ -384,6 +366,150 @@ def on_kill(self):
if not response["SuccessfulSubmissions"]:
self.log.error("Failed to stop AWS Glue Job: %s. Run Id: %s",
self.job_name, self._job_run_id)
+ def _set_job_run_id(self, context: Context, job_run_id: str) -> None:
+ """
+ Record the run id and surface its console link.
+
+ Called from every path that learns a run id (fresh submit, bootstrap
search, reconnect) so
+ ``on_kill`` can stop the run and the link is available before polling
starts. Guarded on the
+ id changing, so a single attempt logs the link once.
+ """
+ if self._job_run_id == job_run_id:
+ return
+ self._job_run_id = job_run_id
+ GlueJobRunDetailsLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ job_name=urllib.parse.quote(self.job_name, safe=""),
+ job_run_id=job_run_id,
+ )
+ self.log.info(
+ "You can monitor this Glue Job run at: %s",
+ GlueJobRunDetailsLink.format_str.format(
+
aws_domain=GlueJobRunDetailsLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ job_name=urllib.parse.quote(self.job_name, safe=""),
+ job_run_id=job_run_id,
+ ),
+ )
+
+ def _build_script_args(self, context: Context) -> dict:
+ script_args = dict(self.script_args)
+ if self.openlineage_inject_parent_job_info:
+ self.log.info("Injecting OpenLineage parent job information into
Glue job arguments.")
+ script_args =
inject_parent_job_information_into_glue_arguments(script_args, context)
+ if self.openlineage_inject_transport_info:
+ self.log.info("Injecting OpenLineage transport information into
Glue job arguments.")
+ script_args =
inject_transport_information_into_glue_arguments(script_args, context)
+ if self.durable and not self.deferrable:
+ script_args, _ = self._prepare_script_args_with_task_uuid(context,
base_args=script_args)
+ return script_args
+
+ def _find_previous_job_run(self, context: Context, task_uuid: str) -> str
| None:
+ """
+ Look for a Glue job run this task instance already started.
+
+ Checks XCom for a cached run id first; falls back to a task-UUID scan
of the job's run
+ history when XCom has nothing (e.g. the cached id expired, or this is
the first retry to
+ run under a provider version that persists it).
+ """
+ ti = context["ti"]
+ previous_job_run_id = ti.xcom_pull(key="glue_job_run_id",
task_ids=ti.task_id)
+ if previous_job_run_id:
+ try:
+ job_run = self.hook.conn.get_job_run(JobName=self.job_name,
RunId=previous_job_run_id)
+ state = job_run.get("JobRun", {}).get("JobRunState")
+ self.log.info("Previous Glue job_run_id: %s, state: %s",
previous_job_run_id, state)
+ if state in ("RUNNING", "STARTING"):
+ return previous_job_run_id
+ except Exception:
+ self.log.warning("Failed to get previous Glue job run state",
exc_info=True)
+ else:
+ try:
+ existing = self._find_job_run_id_by_task_uuid(task_uuid)
+ if existing:
+ existing_job_run_id, existing_job_run_state = existing
+ self.log.info(
+ "Found Glue job_run_id by task UUID: %s, state: %s",
+ existing_job_run_id,
+ existing_job_run_state,
+ )
+ if existing_job_run_state in ("RUNNING", "STARTING"):
+ ti.xcom_push(key="glue_job_run_id",
value=existing_job_run_id)
+ return existing_job_run_id
+ except Exception:
+ self.log.warning("Failed to find previous Glue job run by task
UUID", exc_info=True)
+ return None
+
+ 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)
+ # The --airflow_task_uuid tag is still written on every nondeferred
attempt,
+ # including the first: a later retry's scan can only find it if it was
there from the
+ # start. The lookup itself is skipped on the first attempt, since
nothing tagged with this
+ # task's UUID could possibly exist yet -- paying for a full
run-history scan there would
+ # be pure waste.
+ if self.durable and not self.deferrable and context["ti"].try_number >
1:
+ existing_job_run_id = self._find_previous_job_run(context,
script_args[self.TASK_UUID_ARG])
+ if existing_job_run_id:
+ self._set_job_run_id(context, existing_job_run_id)
+ return existing_job_run_id
+ self.log.info(
+ "Initializing AWS Glue Job: %s. Wait for completion: %s",
+ self.job_name,
+ self.wait_for_completion,
+ )
+ glue_job_run = self.hook.initialize_job(script_args,
self.run_job_kwargs)
+ # Set immediately (before any polling) so on_kill can stop the run
even if the worker dies
+ # before poll_until_complete runs.
+ self._set_job_run_id(context, glue_job_run["JobRunId"])
+ # Pushed unconditionally, matching pre-durable-execution behavior:
downstream tasks read
+ # this key directly, and it's also the only writer that makes
_find_previous_job_run's
+ # XCom tier reachable on the next retry.
+ context["ti"].xcom_push(key="glue_job_run_id",
value=glue_job_run["JobRunId"])
+ return glue_job_run["JobRunId"]
+
+ 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)
+ # This is the first place a reconnecting attempt learns the run id.
+ self._set_job_run_id(context, job_run_id)
Review Comment:
Handled in 118dacfd80
##########
providers/amazon/tests/unit/amazon/aws/operators/test_glue.py:
##########
@@ -781,6 +784,295 @@ def test_inject_parent_job_info_with_resume_on_retry(
assert GlueJobOperator.TASK_UUID_ARG in call_args
+class TestGlueJobOperatorDeprecation:
+ @pytest.mark.parametrize("resume_value", [True, False])
+ def test_warns_on_3_3_plus_and_maps_to_durable(self, resume_value):
+ with
mock.patch("airflow.providers.amazon.aws.operators.glue.AIRFLOW_V_3_3_PLUS",
True):
+ with pytest.warns(AirflowProviderDeprecationWarning,
match="resume_glue_job_on_retry"):
+ glue = GlueJobOperator(
+ task_id=TASK_ID, job_name=JOB_NAME,
resume_glue_job_on_retry=resume_value
+ )
+ assert glue.durable is resume_value
+
+ def test_silent_below_3_3(self):
+ with
mock.patch("airflow.providers.amazon.aws.operators.glue.AIRFLOW_V_3_3_PLUS",
False):
+ with warnings.catch_warnings(record=True) as caught:
+ warnings.simplefilter("always")
+ glue = GlueJobOperator(task_id=TASK_ID, job_name=JOB_NAME,
resume_glue_job_on_retry=True)
+ deprecation_warnings = [
+ w for w in caught if issubclass(w.category,
AirflowProviderDeprecationWarning)
+ ]
+ assert deprecation_warnings == []
+ assert glue.durable is True
+
+ def test_both_flags_passed_durable_wins(self):
+ with
mock.patch("airflow.providers.amazon.aws.operators.glue.AIRFLOW_V_3_3_PLUS",
True):
+ with pytest.warns(AirflowProviderDeprecationWarning):
+ glue = GlueJobOperator(
+ task_id=TASK_ID,
+ job_name=JOB_NAME,
+ durable=True,
+ resume_glue_job_on_retry=False,
+ )
+ assert glue.durable is True
+
+ @pytest.mark.skipif(
+ not AIRFLOW_V_3_3_PLUS,
+ reason="The <3.3 compat stub's __init__ isn't decorated with
BaseOperatorMeta._apply_defaults, "
+ "so default_args injection for durable only works on the real
ResumableJobMixin.",
+ )
+ def test_default_args_durable_reaches_operator(self):
+ with DAG(
+ dag_id="test_glue_durable_default_args",
+ schedule=None,
+ start_date=datetime(2024, 1, 1),
+ default_args={"durable": False},
+ ):
+ glue = GlueJobOperator(task_id=TASK_ID, job_name=JOB_NAME)
+ assert glue.durable is False
+
+
+class FakeTaskStateStore:
+ """In-memory task state store for tests."""
+
+ def __init__(self, stored: dict[str, str] | None = None):
+ self._store: dict[str, str] = dict(stored or {})
+
+ def get(self, key: str) -> str | None:
+ return self._store.get(key)
+
+ def set(self, key: str, value: str) -> None:
+ self._store[key] = value
+
+
[email protected](
+ not AIRFLOW_V_3_3_PLUS,
+ reason="ResumableJobMixin reconnect requires task_state_store, available
in Airflow 3.3+",
+)
+class TestGlueJobOperatorDurableExecution:
+ def _build(self, **kwargs):
+ return GlueJobOperator(task_id=TASK_ID, job_name=JOB_NAME, **kwargs)
+
+ def _stub_empty_scan(self, glue):
+ # submit_job scans for a task UUID tagged run whenever durable is set,
regardless of why it
+ # was called. Stub it to return no matches so tests don't depend on
that fallback mechanism.
+ glue.hook.conn = mock.MagicMock()
+ glue.hook.conn.get_job_runs.return_value = {"JobRuns": []}
+
+ def _context(self, store=None, try_number=2):
+ ti = mock.MagicMock()
+ ti.try_number = try_number
+ ti.xcom_pull.return_value = None
+ ctx = {"ti": ti}
+ if store is not None:
+ ctx["task_state_store"] = store
+ return ctx
+
+ @mock.patch.object(GlueJobHook, "job_completion")
+ @mock.patch.object(GlueJobHook, "initialize_job")
+ @mock.patch.object(GlueJobHook, "get_conn")
+ def test_fresh_submit_persists_before_polling(
+ self, mock_get_conn, mock_initialize_job, mock_job_completion
+ ):
+ glue = self._build(durable=True)
+ self._stub_empty_scan(glue)
+ mock_initialize_job.return_value = {"JobRunId": "jr_new"}
+ store = FakeTaskStateStore()
+ persisted_before_poll = []
+ mock_job_completion.side_effect = lambda *a, **k: (
+ persisted_before_poll.append(store.get("glue_job_run_id")) or
{"JobRunState": "SUCCEEDED"}
+ )
+
+ job_run_id = glue.execute(self._context(store))
+
+ assert job_run_id == "jr_new"
+ assert store.get("glue_job_run_id") == "jr_new"
+ assert persisted_before_poll == ["jr_new"]
+ mock_initialize_job.assert_called_once()
+
+ @mock.patch.object(GlueJobHook, "job_completion")
+ @mock.patch.object(GlueJobHook, "get_job_state")
+ @mock.patch.object(GlueJobHook, "initialize_job")
+ @mock.patch.object(GlueJobHook, "get_conn")
+ def test_reconnect_when_stored_run_is_running(
+ self, mock_get_conn, mock_initialize_job, mock_get_job_state,
mock_job_completion
+ ):
+ glue = self._build(durable=True)
+ mock_get_job_state.return_value = "RUNNING"
+ mock_job_completion.return_value = {"JobRunState": "SUCCEEDED"}
+ store = FakeTaskStateStore({"glue_job_run_id": "jr_old"})
+
+ job_run_id = glue.execute(self._context(store))
+
+ assert job_run_id == "jr_old"
+ mock_initialize_job.assert_not_called()
+ mock_job_completion.assert_called_once_with(JOB_NAME, "jr_old", False,
0)
+
+ @pytest.mark.parametrize("status", ["STARTING", "RUNNING", "WAITING",
"STOPPING"])
+ @mock.patch.object(GlueJobHook, "job_completion")
+ @mock.patch.object(GlueJobHook, "get_job_state")
+ @mock.patch.object(GlueJobHook, "initialize_job")
+ @mock.patch.object(GlueJobHook, "get_conn")
+ def test_reconnect_from_every_active_state(
+ self, mock_get_conn, mock_initialize_job, mock_get_job_state,
mock_job_completion, status
+ ):
+ glue = self._build(durable=True)
+ mock_get_job_state.return_value = status
+ mock_job_completion.return_value = {"JobRunState": "SUCCEEDED"}
+ store = FakeTaskStateStore({"glue_job_run_id": "jr_old"})
+
+ glue.execute(self._context(store))
+
+ mock_initialize_job.assert_not_called()
+
+ @mock.patch.object(GlueJobHook, "job_completion")
+ @mock.patch.object(GlueJobHook, "get_job_state")
+ @mock.patch.object(GlueJobHook, "initialize_job")
+ @mock.patch.object(GlueJobHook, "get_conn")
+ def test_already_succeeded_returns_without_resubmit(
+ self, mock_get_conn, mock_initialize_job, mock_get_job_state,
mock_job_completion
+ ):
+ glue = self._build(durable=True)
+ mock_get_job_state.return_value = "SUCCEEDED"
+ store = FakeTaskStateStore({"glue_job_run_id": "jr_old"})
+
+ job_run_id = glue.execute(self._context(store))
+
+ assert job_run_id == "jr_old"
+ mock_initialize_job.assert_not_called()
+ mock_job_completion.assert_not_called()
+
+ @pytest.mark.parametrize("status", ["FAILED", "TIMEOUT", "STOPPED"])
+ @mock.patch.object(GlueJobHook, "job_completion")
+ @mock.patch.object(GlueJobHook, "get_job_state")
+ @mock.patch.object(GlueJobHook, "initialize_job")
+ @mock.patch.object(GlueJobHook, "get_conn")
+ def test_terminal_failure_resubmits_fresh(
+ self, mock_get_conn, mock_initialize_job, mock_get_job_state,
mock_job_completion, status
+ ):
+ """STOPPED resubmits like any other terminal state: Glue's API can't
tell a console
+ cancellation from a run this operator's own on_kill stopped, so
treating it as success
+ would risk silently reporting a self-inflicted stop as done."""
+ glue = self._build(durable=True)
+ self._stub_empty_scan(glue)
+ mock_get_job_state.return_value = status
+ mock_initialize_job.return_value = {"JobRunId": "jr_new"}
+ mock_job_completion.return_value = {"JobRunState": "SUCCEEDED"}
+ store = FakeTaskStateStore({"glue_job_run_id": "jr_old"})
+
+ job_run_id = glue.execute(self._context(store))
+
+ assert job_run_id == "jr_new"
+ assert store.get("glue_job_run_id") == "jr_new"
+ mock_initialize_job.assert_called_once()
+
+ @mock.patch.object(GlueJobHook, "job_completion")
+ @mock.patch.object(GlueJobHook, "get_job_state")
+ @mock.patch.object(GlueJobHook, "initialize_job")
+ @mock.patch.object(GlueJobHook, "get_conn")
+ def test_not_found_resubmits_fresh(
+ self, mock_get_conn, mock_initialize_job, mock_get_job_state,
mock_job_completion
+ ):
+ glue = self._build(durable=True)
+ self._stub_empty_scan(glue)
+ mock_get_job_state.side_effect = ClientError(
+ {"Error": {"Code": "EntityNotFoundException", "Message": "gone"}},
"GetJobRun"
+ )
+ mock_initialize_job.return_value = {"JobRunId": "jr_new"}
+ mock_job_completion.return_value = {"JobRunState": "SUCCEEDED"}
+ store = FakeTaskStateStore({"glue_job_run_id": "jr_old"})
+
+ job_run_id = glue.execute(self._context(store))
+
+ assert job_run_id == "jr_new"
+ mock_initialize_job.assert_called_once()
+
+ @mock.patch.object(GlueJobHook, "job_completion")
+ @mock.patch.object(GlueJobHook, "initialize_job")
+ @mock.patch.object(GlueJobHook, "get_conn")
+ def test_durable_false_never_touches_store(self, mock_get_conn,
mock_initialize_job, mock_job_completion):
+ glue = self._build(durable=False)
+ mock_initialize_job.return_value = {"JobRunId": "jr_new"}
+ mock_job_completion.return_value = {"JobRunState": "SUCCEEDED"}
+ store = FakeTaskStateStore({"glue_job_run_id": "jr_old"})
+
+ job_run_id = glue.execute(self._context(store))
+
+ assert job_run_id == "jr_new"
+ assert store.get("glue_job_run_id") == "jr_old", "store must be left
untouched"
+ mock_initialize_job.assert_called_once()
+
+ @mock.patch.object(GlueJobHook, "conn", new_callable=mock.PropertyMock)
+ @mock.patch.object(GlueJobHook, "job_completion")
+ @mock.patch.object(GlueJobHook, "initialize_job")
+ @mock.patch.object(GlueJobHook, "get_conn")
+ def test_first_attempt_skips_the_retry_lookup_entirely(
+ self, mock_get_conn, mock_initialize_job, mock_job_completion,
mock_conn
+ ):
+ glue = self._build(durable=True)
+ mock_initialize_job.return_value = {"JobRunId": "jr_new"}
+ mock_job_completion.return_value = {"JobRunState": "SUCCEEDED"}
+ store = FakeTaskStateStore()
+
+ job_run_id = glue.execute(self._context(store, try_number=1))
+
+ assert job_run_id == "jr_new"
+ mock_initialize_job.assert_called_once()
+ mock_conn.return_value.get_job_run.assert_not_called()
+ mock_conn.return_value.get_job_runs.assert_not_called()
Review Comment:
Handled in 118dacfd80
--
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]
