SageMakerProcessingOperator would not terminate if Sagemaker processing job ends in the Stopped state

[dym-ok]

Under which category would you file this issue?

Providers

Apache Airflow version

3.2.1

What happened and how to reproduce it?

When using the SageMakerProcessingOperator with deferrable=True and wait_for_completion=True if the job ends in a 'Stopped' state the AirflowException(f"SageMaker job failed because {response['FailureReason']}") will not be thrown because execute() checks SageMakerHook.failed_states and not SageMakerHook.processing_job_failed_states, which includes the Stopped state.

It keeps looping even if the job is not running anymore.

What you think should happen instead?

The task should end when the SageMaker processing job ends in a "Stopped" state with the corresponding AirflowException thrown.

Operating System

No response

Deployment

Astronomer

Apache Airflow Provider(s)

amazon

Versions of Apache Airflow Providers

apache-airflow-providers-amazon==9.25.0
apache-airflow-providers-common-compat==1.14.3
apache-airflow-providers-common-io==1.7.2
apache-airflow-providers-common-sql==1.34.0
apache-airflow-providers-github==2.11.2
apache-airflow-providers-google==21.1.0
apache-airflow-providers-http==6.0.2
apache-airflow-providers-openlineage==2.14.0
apache-airflow-providers-salesforce==5.14.0
apache-airflow-providers-sftp==5.7.3
apache-airflow-providers-slack==9.10.0
apache-airflow-providers-smtp==2.4.5
apache-airflow-providers-snowflake==6.12.1
apache-airflow-providers-ssh==5.0.0
apache-airflow-providers-standard==1.12.3

Official Helm Chart version

Not Applicable

Kubernetes Version

No response

Helm Chart configuration

No response

Docker Image customizations

No response

Anything else?

No response

Are you willing to submit PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct

[Codingaditya17]

Hi, I would like to work on this issue.

From the description, it looks like SageMakerProcessingOperator is checking the generic SageMakerHook.failed_states instead of SageMakerHook.processing_job_failed_states, so the Stopped state is not treated as a terminal failed state for processing jobs.

I can look into updating the operator logic and adding/adjusting tests to cover the Stopped state when deferrable=True and wait_for_completion=True.

Thanks!

[Codingaditya17]

Thanks! I have updated the branch with the latest upstream/main. The focused local test could not be completed because the uv environment attempted to build jpype1 and failed due to missing Java runtime on my machine, but I added a targeted unit test for the Stopped processing job state.