amoghrajesh opened a new pull request, #76036:
URL: https://github.com/apache/airflow/pull/76036
<!-- 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
-->
`SageMakerTrainingOperator` submits a new training job on every retry. When
a worker dies mid-training, attempt 2 finds the job name taken, appends a
timestamp via `_get_unique_job_name`, and submits a duplicate. The original job
keeps running and keeps billing, unwatched.
This brings durable execution (AIP-103) to the operator by porting it to
`ResumableJobMixin`, joining `GlueJobOperator`, `RedshiftDataOperator`,
`DatabricksSubmitRunOperator` and the others already on it. The submitted job
name
goes into the task state store before polling, and a retry reads it back and
asks SageMaker what happened to it: still running means reconnect and keep
polling, already `Completed` means return the result without resubmitting,
terminal failure
means submit fresh.
`execute()` now splits into the mixin's six methods. `submit_job` is
everything that used to run before polling; `poll_until_complete` is the wait
that used to live inside `hook.create_training_job`, which is now called as a
pure submit. The deferrable path moved verbatim into `_execute_deferrable` and
is otherwise unchanged.
### On the timestamp rename
It stays, and it needed no special casing. The rename lives inside
`submit_job`, and
the mixin only calls `submit_job` once it has decided to submit:
- `durable=False` runs submit, poll, result with the task state store
untouched, so
the rename behaves exactly as it does today.
- `durable=True` reconnecting never calls `submit_job`, so the rename never
runs.
- `durable=True` after a terminal job does rename, which is what you want,
since the
old job still holds that name.
`CheckpointConfig.S3Uri` is set by the user and is independent of the job
name, so a
renamed resubmit still reads the same checkpoint folder. That is the case
that matters
for managed spot training.
### Scope
Durable execution applies to the synchronous path only. A deferrable run
behaves as before and does not reconnect, matching what `ResumableJobMixin`
documents and what `GlueJobOperator` does.
### Behaviour change
`durable` defaults to `True`, following the mixin and the other ported
operators. On upgrade, a retry reconnects instead of submitting a duplicate.
Set `durable=False` for the old behaviour. Noted in the provider changelog.
Below Airflow 3.3 the compat stand-in applies and the operator behaves as
with `durable=False`.
### Testing
Used this Dag:
```python
with DAG(
dag_id="test_sagemaker_durable",
schedule=None,
start_date=datetime(2026, 1, 1),
catchup=False,
tags=["manual-test"],
) as dag:
create_bucket = S3CreateBucketOperator(task_id="create_bucket",
bucket_name=BUCKET, region_name=REGION)
upload_data = S3CreateObjectOperator(
task_id="upload_data",
s3_bucket=BUCKET,
s3_key="train/train.csv",
data=TRAIN_CSV,
replace=True,
)
# SageMaker rejects a training image from a different region than the
client, so the region is
# pinned here rather than left to the connection's default.
train_model = SageMakerTrainingOperator(
task_id="train_model",
config=TRAINING_CONFIG,
region_name=REGION,
print_log=False,
check_interval=10,
durable=True,
retries=1,
retry_delay=timedelta(seconds=30),
execution_timeout=timedelta(minutes=3),
)
delete_bucket = S3DeleteBucketOperator(
task_id="delete_bucket",
bucket_name=BUCKET,
force_delete=True,
trigger_rule=TriggerRule.ALL_DONE,
)
chain(create_bucket, upload_data, train_model, delete_bucket)
```
Ran the dag and killed the worker process with `kill -9 <pid>`
Try 1:
<img width="1723" height="972" alt="image"
src="https://github.com/user-attachments/assets/04107098-02c0-4454-9c1b-2331132c90f9";
/>
Try 2:
<img width="1723" height="972" alt="image"
src="https://github.com/user-attachments/assets/92439d25-26b4-4ccf-8a12-02ee520382ae";
/>
Only one job submitted
<img width="1723" height="972" alt="image"
src="https://github.com/user-attachments/assets/5a2f0617-28da-4d5b-bfa0-595edd47531d";
/>
<img width="1723" height="972" alt="image"
src="https://github.com/user-attachments/assets/84f6aec5-a01b-4598-a7e7-b6a294935510";
/>
---
##### 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".
-->
- [ ] Yes (please specify the tool below)
<!--
Generated-by: [Tool Name] following [the
guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
-->
---
* 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]
