amoghrajesh opened a new pull request, #71211:
URL: https://github.com/apache/airflow/pull/71211
<!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
<!--
Thank you for contributing!
Please provide above a brief description of the changes made in this pull
request.
Write a good git commit message following this guide:
https://chris.beams.io/posts/git-commit/
Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the
documentation.
For user-facing UI changes, please attach before/after screenshots (or a
short
screen recording) so reviewers can assess the visual impact.
Feel free to ping (in general) for the review if you do not see reaction for
a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we
sometimes miss notifications.
In case of an existing issue, reference it using one of the following:
* closes: #ISSUE
* related: #ISSUE
-->
---
##### Was generative AI tooling used to co-author this PR?
<!--
If generative AI tooling has been used in the process of authoring this PR,
please
change below checkbox to `[X]` followed by the name of the tool, uncomment
the "Generated-by".
-->
- [x] Yes - claude sonnet
<!--
Generated-by: [Tool Name] following [the
guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
-->
## Summary
This PR intends to extend GlueJobOperator for durable execution (AIP-103):
on Airflow 3.3+, the job
run id is persisted to task state store immediately after submission, so a
worker crash and retry reconnects to the run already executing in Glue instead
of starting a new one.
## How it works today
GlueJobOperator's resume_glue_job_on_retry mechanism has two tiers: an XCom
cached run id (fast path), falling back to a paginated scan of the job's run
history tagged by task UUID. I reproduced twice against real AWS Glue that the
fast path never fires on Airflow 3.x: the server clears a task instance's XComs
before every non deferral attempt, so the cached id is always gone by the time
the next retry looks for it. Every retry was paying for the full scan
regardless, and the scan itself returned the first UUID match regardless of
state, so an earlier terminated run could mask a active one underneath it.
## Why
`resume_glue_job_on_retry` already tries to do this via XCom, falling back to
a scan of the job's run history when XCom is empty. Confirmed against real
AWS Glue (killed the worker mid-poll, twice, under two different kill
methods) that the XCom fast path never fires on Airflow 3.x: the server
clears a task instance's XComs before every non-deferral attempt
(`task_instances.py`'s `xcom_keys_to_clear`), so by the time the next retry
looks for the cached id it's already gone. Every retry pays for the full
paginated scan regardless of timing, not just the narrow crash-before-push
window the mechanism's original PR (#62560) was written to cover.
The scan itself also has a latent bug: it returns the first task-UUID match
regardless of state, so an earlier terminated run can mask a still-active
one further down the page. Not fixed in this PR -- see Scope below.
This matters more for Glue than for most `ResumableJobMixin` ports: a Glue
job's `concurrent_run_limit` defaults to `1`, so a duplicate submission
doesn't just waste compute, it fails outright with
`ConcurrentRunsExceededException` and the task keeps retrying against a run
it can never observe.
## What changed
- `GlueJobOperator` now inherits `ResumableJobMixin`. On Airflow 3.3+, the
mixin checks task state store before doing anything else: an active run
reconnects directly (one `get_job_state` call), a succeeded or externally
stopped run returns without resubmitting, a terminal run resubmits fresh.
- The **existing** XCom followed by scan mechanism is kept exactly as it
was, relocated into `submit_job` (the mixin only calls `submit_job` when task
state has nothing usable). This is what Airflow <3.3 runs on entirely, and what
3.3+ falls back to as a one-time bootstrap for a run started before this
upgrade, or before task state store had anything persisted.
- `resume_glue_job_on_retry` is deprecated in favour of `durable`
(`AirflowProviderDeprecationWarning`, gated to fire only on Airflow 3.3+);
passing it still works and maps onto `durable`.
- A run `STOPPED` outside Airflow (e.g. cancelled manually) is treated as a
success rather than resubmitted -- the work is genuinely finished, just not the
way the task expected -- with a warning logged. Matches how
`GlueJobHook.job_completion` and `GlueJobCompleteTrigger` already treat
`STOPPED`.
- Deferrable takes precedence over durable, unchanged in spirit from today:
the Triggerer already tracks the run across the wait, so the run id isn't
persisted to task state on that path.
## Behaviour change
`durable` is owned by the mixin and defaults to `True`.
`resume_glue_job_on_retry` defaulted to `False`. A Dag that never touched the
flag goes from "always resubmit on retry" to "reattach on retry" with no code
change.
Documented in `providers/amazon/docs/changelog.rst`
## Scope
Deliberately left alone, all pre existing, none introduced by this PR:
- The scan's first match regardless of state bug described above in the PR.
- `GlueJobHook._handle_state` treats only `FAILED`/`TIMEOUT` as failure, so
AWS's terminal `ERROR`/`EXPIRED` states fall through to "keep polling" and the
sync poll loop never exits on them.
- `ResumableJobMixin.execute_resumable`'s reconnect path returns
`poll_until_complete`'s `None` instead of calling `get_job_result`, unlike the
other two paths. Not touched here (no task-sdk change) -- every method instead
sets `self._job_run_id`, and `execute()` returns that directly.
## Why `task_state_store` is a better crash recovery mechanism
**1. It survives retries while xcom does not.**
Airflow clears all xcoms before each retry attempt. `task_state_store` is
built specifically to persist across retries.
**2. One call vs. a full scan.**
Store hit -> single `get_job_state` call. Scan -> paginates the entire job
run history, 50 at a time.
**3. No ambiguity.**
Store returns the exact id you saved. The scan matches by tag and can pick a
stale run if an older one shares the tag.
## Testing
Uploaded this script on AWS Glue and created an _ETL JOB_ there:
```python
import sys
import time
def get_arg(name, default=None):
flag = f"--{name}"
if flag in sys.argv:
i = sys.argv.index(flag)
if i + 1 < len(sys.argv):
return sys.argv[i + 1]
return default
task_uuid = get_arg("airflow_task_uuid", "<NOT INJECTED>")
sleep_total = int(get_arg("sleep_seconds", "600"))
print(f"argv: {sys.argv}", flush=True)
print(f"airflow_task_uuid = {task_uuid}", flush=True)
print(f"sleeping for {sleep_total}s", flush=True)
elapsed = 0
while elapsed < sleep_total:
time.sleep(15)
elapsed += 15
print(f"alive {elapsed}/{sleep_total}s uuid={task_uuid}", flush=True)
print("done", flush=True)
```
<img width="2559" height="1148" alt="image"
src="https://github.com/user-attachments/assets/e55604b8-064d-4d7b-832a-718f42e7dba5";
/>
Using this DAG:
```python
from datetime import datetime, timedelta
from airflow.providers.amazon.aws.operators.glue import GlueJobOperator
from airflow.sdk import DAG
with DAG(
dag_id="glue_job_trial",
schedule=None,
start_date=datetime(2024, 1, 1),
catchup=False,
) as dag:
run_glue = GlueJobOperator(
task_id="run_glue",
job_name="airflow_durable_test",
aws_conn_id="aws_default",
region_name="us-east-1",
durable=True,
script_args={"--sleep_seconds": "900"},
wait_for_completion=True,
deferrable=False,
verbose=True,
retries=3,
retry_delay=timedelta(seconds=20),
)
```
### Case 1: Crash Recovery after worker crash
Launched the DAG:
<img width="2559" height="1148" alt="image"
src="https://github.com/user-attachments/assets/2e377d09-ba24-4410-9dbd-138b9259c4cf";
/>
Glue Job:
<img width="2559" height="1148" alt="image"
src="https://github.com/user-attachments/assets/e707d43a-0112-4234-b24e-b27f9ff45e34";
/>
State Store:
<img width="2559" height="1148" alt="image"
src="https://github.com/user-attachments/assets/1fe0f914-0cd8-44ac-8169-661825230af5";
/>
Brought the worker down, see that the airflow task instance goes to retry
state
<img width="2559" height="1148" alt="image"
src="https://github.com/user-attachments/assets/121885f1-6005-4496-ba8d-fc51c118f7fd";
/>
Reconnects to same job
<img width="2559" height="1148" alt="image"
src="https://github.com/user-attachments/assets/f498ad7d-bae2-4469-941c-f06fb0d5800e";
/>
### Case 2: Stop a run which was submitted from Glue Console, there should
be no resubmit
<img width="2559" height="1148" alt="image"
src="https://github.com/user-attachments/assets/af71dbcd-612b-4884-9e23-3f8c7cd5a959";
/>
<img width="2559" height="1148" alt="image"
src="https://github.com/user-attachments/assets/781c7ea9-5a75-433e-881d-9de2d4f38d2f";
/>
<img width="2559" height="1191" alt="image"
src="https://github.com/user-attachments/assets/d28b19bb-fd57-4d9f-92f1-ebc92fa538a0";
/>
### Case 3: `wait_for_completion=False` persists to state store
Update task to:
```python
run_glue = GlueJobOperator(
task_id="run_glue",
job_name="airflow_durable_test",
aws_conn_id="aws_default",
region_name="us-east-1",
durable=True,
script_args={"--sleep_seconds": "900"},
wait_for_completion=False,
deferrable=False,
verbose=True,
retries=3,
retry_delay=timedelta(seconds=20),
)
```
<img width="2559" height="1191" alt="image"
src="https://github.com/user-attachments/assets/906035d1-0f5e-47a1-83be-76453e5b90af";
/>
<img width="2559" height="1191" alt="image"
src="https://github.com/user-attachments/assets/b36f16da-0498-4cb7-bcfc-1b2a81fd66f5";
/>
---
* Read the **[Pull Request
Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)**
for more information. Note: commit author/co-author name and email in commits
become permanently public when merged.
* For fundamental code changes, an Airflow Improvement Proposal
([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals))
is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License
Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment:
`{pr_number}.significant.rst`, in
[airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments).
You can add this file in a follow-up commit after the PR is created so you
know the PR number.
--
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]
