check transform job status before deferring SageMakerTransformOperator

airflow/providers/amazon/aws/operators/sagemaker.py

@@ -700,8 +700,24 @@ def execute(self, context: Context) -> dict:
             raise AirflowException(f"Sagemaker transform Job creation failed: {response}")
 
         if self.deferrable and self.wait_for_completion:
+            response = self.hook.describe_transform_job(transform_config["TransformJobName"])
+            status = response["TransformJobStatus"]
+            if status in self.hook.failed_states:
+                raise AirflowException(f"SageMaker job failed because {response['FailureReason']}")
+
+            if status == "Completed":
+                self.log.info("%s completed successfully.", self.task_id)
+                return {
+                    "Model": serialize(self.hook.describe_model(transform_config["ModelName"])),
+                    "Transform": serialize(response),
+                }
+
+            timeout = self.execution_timeout
+            if self.max_ingestion_time:
+                timeout = datetime.timedelta(seconds=self.max_ingestion_time)
+
             self.defer(
-                timeout=self.execution_timeout,
+                timeout=timeout,
                 trigger=SageMakerTrigger(
                     job_name=transform_config["TransformJobName"],
                     job_type="Transform",
@@ -712,17 +728,18 @@ def execute(self, context: Context) -> dict:
                 method_name="execute_complete",
             )
 
-        self.serialized_model = serialize(self.hook.describe_model(transform_config["ModelName"]))
-        self.serialized_transform = serialize(
-            self.hook.describe_transform_job(transform_config["TransformJobName"])
-        )
-        return {"Model": self.serialized_model, "Transform": self.serialized_transform}
+        return self.serialize_result()
 
-    def execute_complete(self, context, event=None):
-        if event["status"] != "success":
-            raise AirflowException(f"Error while running job: {event}")
-        else:
-            self.log.info(event["message"])
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, dict]:
+        if event is None:
+            err_msg = "Trigger error: event is None"
+            self.log.error(err_msg)
+            raise AirflowException(err_msg)
+
+        self.log.info(event["message"])
+        return self.serialize_result()
+
+    def serialize_result(self) -> dict[str, dict]:
         transform_config = self.config.get("Transform", self.config)
         self.serialized_model = serialize(self.hook.describe_model(transform_config["ModelName"]))
         self.serialized_transform = serialize(

tests/providers/amazon/aws/operators/test_sagemaker_transform.py

@@ -167,18 +167,66 @@ def test_execute_without_check_if_job_exists(self, _, __, ___, mock_transform, _
             max_ingestion_time=None,
         )
 
+    @mock.patch("airflow.providers.amazon.aws.operators.sagemaker.SageMakerTransformOperator.defer")
+    @mock.patch.object(
+        SageMakerHook,
+        "describe_transform_job",
+        return_value={"TransformJobStatus": "Failed", "FailureReason": "it failed"},
+    )
     @mock.patch.object(SageMakerHook, "create_transform_job")
     @mock.patch.object(SageMakerHook, "create_model")
-    def test_operator_defer(self, _, mock_transform):
+    def test_operator_failed_before_defer(self, _, mock_transform, mock_describe_transform_job, mock_defer):
         mock_transform.return_value = {
             "TransformJobArn": "test_arn",
             "ResponseMetadata": {"HTTPStatusCode": 200},
         }
         self.sagemaker.deferrable = True
         self.sagemaker.wait_for_completion = True
         self.sagemaker.check_if_job_exists = False
+
+        with pytest.raises(AirflowException):
+            self.sagemaker.execute(context=None)
+        assert not mock_defer.called
+
+    @mock.patch("airflow.providers.amazon.aws.operators.sagemaker.SageMakerTransformOperator.defer")
+    @mock.patch.object(SageMakerHook, "describe_model")
+    @mock.patch.object(
+        SageMakerHook, "describe_transform_job", return_value={"TransformJobStatus": "Completed"}
+    )
+    @mock.patch.object(SageMakerHook, "create_transform_job")
+    @mock.patch.object(SageMakerHook, "create_model")
+    def test_operator_complete_before_defer(
+        self, _, mock_transform, mock_describe_transform_job, mock_describe_model, mock_defer
+    ):
+        mock_transform.return_value = {
+            "TransformJobArn": "test_arn",
+            "ResponseMetadata": {"HTTPStatusCode": 200},
+        }
+        mock_describe_model.return_value = {"PrimaryContainer": {"ModelPackageName": "package-name"}}
+        self.sagemaker.deferrable = True
+        self.sagemaker.wait_for_completion = True
+        self.sagemaker.check_if_job_exists = False
+
+        self.sagemaker.execute(context=None)
+        assert not mock_defer.called
+
+    @mock.patch.object(
+        SageMakerHook, "describe_transform_job", return_value={"TransformJobStatus": "InProgress"}
+    )
+    @mock.patch.object(SageMakerHook, "create_transform_job")
+    @mock.patch.object(SageMakerHook, "create_model")
+    def test_operator_defer(self, _, mock_transform, mock_describe_transform_job):
+        mock_transform.return_value = {
+            "TransformJobArn": "test_arn",
+            "ResponseMetadata": {"HTTPStatusCode": 200},
+        }
+        self.sagemaker.deferrable = True
+        self.sagemaker.wait_for_completion = True
+        self.sagemaker.check_if_job_exists = False
+
         with pytest.raises(TaskDeferred) as exc:
             self.sagemaker.execute(context=None)
+
         assert isinstance(exc.value.trigger, SageMakerTrigger), "Trigger is not a SagemakerTrigger"
 
     @mock.patch.object(SageMakerHook, "describe_transform_job")