This is an automated email from the ASF dual-hosted git repository.
vincbeck 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 80e870d3dfa Fix SageMakerTransformOperator succeeding on a failed
deferred job (#69042)
80e870d3dfa is described below
commit 80e870d3dfa3a04a2dfa5693872f2a5e3dcec8d2
Author: Steve Ahn <[email protected]>
AuthorDate: Mon Jun 29 06:07:25 2026 -0700
Fix SageMakerTransformOperator succeeding on a failed deferred job (#69042)
A deferred SageMaker transform job that fails during the wait was reported
as a
successful task: execute_complete did not check the trigger event status,
so the
failed job's description was pushed as the result and downstream tasks ran
against
missing or invalid transform output with no error surfaced. Since the
trigger was
migrated to AwsBaseWaiterTrigger it yields a {"status": "error"} event on
failure
instead of raising, so the status guard that every sibling SageMaker
operator has
became load-bearing here too.
---
.../providers/amazon/aws/operators/sagemaker.py | 3 +++
.../aws/operators/test_sagemaker_transform.py | 27 ++++++++++++++++++++++
2 files changed, 30 insertions(+)
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py
b/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py
index e23db37ae2d..21fdaa6b988 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py
@@ -944,6 +944,9 @@ class SageMakerTransformOperator(SageMakerBaseOperator):
def execute_complete(self, context: Context, event: dict[str, Any] | None
= None) -> dict[str, dict]:
validated_event = validate_execute_complete_event(event)
+ if validated_event["status"] != "success":
+ raise RuntimeError(f"Error while running transform job:
{validated_event}")
+
self.log.info("SageMaker job %s completed.",
validated_event["job_name"])
return self.serialize_result(validated_event["job_name"])
diff --git
a/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_transform.py
b/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_transform.py
index 40c86115c48..de8ed7841e8 100644
---
a/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_transform.py
+++
b/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_transform.py
@@ -341,6 +341,33 @@ class TestSageMakerTransformOperator:
self.sagemaker.execute(context=None)
assert not mock_defer.called
+ @mock.patch.object(sagemaker, "serialize", return_value="")
+ @mock.patch.object(SageMakerHook, "describe_model",
return_value={"ModelName": "model_name"})
+ @mock.patch.object(
+ SageMakerHook,
+ "describe_transform_job",
+ return_value={
+ "ModelName": "model_name",
+ "TransformJobStatus": "Failed",
+ "FailureReason": "it failed",
+ },
+ )
+ def test_execute_complete_raises_when_job_failed_during_deferred_wait(
+ self, mock_describe_transform_job, mock_describe_model, mock_serialize
+ ):
+ # When the transform job fails during the deferred wait, the trigger
(an
+ # AwsBaseWaiterTrigger) yields {"status": "error", ...} instead of
raising, so
+ # execute_complete must reject a non-success status rather than report
the task as
+ # successful — matching every other SageMaker operator's
execute_complete.
+ event = {
+ "status": "error",
+ "message": "Error while waiting for transform job: terminal
failure",
+ "job_name": "job_name",
+ }
+
+ with pytest.raises(RuntimeError, match="Error while running transform
job"):
+ self.sagemaker.execute_complete(context=None, event=event)
+
@mock.patch("airflow.providers.amazon.aws.operators.sagemaker.SageMakerTransformOperator.defer")
@mock.patch.object(SageMakerHook, "describe_model")
@mock.patch.object(
