fix: WaiterError on failed pipeline execution. results() (#1337)

Co-authored-by: Namrata Madan <[email protected]>

Author: 2 people

src/sagemaker/workflow/pipeline.py

@@ -22,13 +22,14 @@
 import attr
 import botocore
 import pytz
-from botocore.exceptions import ClientError
+from botocore.exceptions import ClientError, WaiterError
 
 from sagemaker import s3
 from sagemaker._studio import _append_project_tags
 from sagemaker.config import PIPELINE_ROLE_ARN_PATH, PIPELINE_TAGS_PATH
 from sagemaker.remote_function.core.serialization import deserialize_obj_from_s3
 from sagemaker.remote_function.core.stored_function import RESULTS_FOLDER
+from sagemaker.remote_function.errors import RemoteFunctionError
 from sagemaker.remote_function.job import JOBS_CONTAINER_ENTRYPOINT
 from sagemaker.s3_utils import s3_path_join
 from sagemaker.session import Session
@@ -965,8 +966,13 @@ def result(self, step_name: str):
 
         Raises:
               ValueError if the provided step is not a ``@step`` decorated function.
+              RemoteFunctionError if the provided step is not in "Completed" status
         """
-        self.wait()
+        try:
+            self.wait()
+        except WaiterError as e:
+            if "Waiter encountered a terminal failure state" in str(e):
+                pass
         step = next(filter(lambda x: x["StepName"] == step_name, self.list_steps()), None)
         if not step:
             raise ValueError(f"Invalid step name {step_name}")
@@ -986,15 +992,22 @@ def result(self, step_name: str):
         ]
         if container_args != JOBS_CONTAINER_ENTRYPOINT:
             raise ValueError(
-                "This method can only be used on pipeline steps created using @step" " decorator."
+                "This method can only be used on pipeline steps created using @step decorator."
             )
-
         s3_output_path = describe_training_job_response["OutputDataConfig"]["S3OutputPath"]
-        return deserialize_obj_from_s3(
-            sagemaker_session=self.sagemaker_session,
-            s3_uri=s3_path_join(s3_output_path, self.arn.split("/")[-1], step_name, RESULTS_FOLDER),
-            hmac_key=describe_training_job_response["Environment"]["REMOTE_FUNCTION_SECRET_KEY"],
-        )
+
+        job_status = describe_training_job_response["TrainingJobStatus"]
+        if job_status == "Completed":
+            return deserialize_obj_from_s3(
+                sagemaker_session=self.sagemaker_session,
+                s3_uri=s3_path_join(
+                    s3_output_path, self.arn.split("/")[-1], step_name, RESULTS_FOLDER
+                ),
+                hmac_key=describe_training_job_response["Environment"][
+                    "REMOTE_FUNCTION_SECRET_KEY"
+                ],
+            )
+        raise RemoteFunctionError(f"Pipeline step {step_name} is in {job_status} status.")
 
 
 class PipelineGraph:

tests/integ/sagemaker/workflow/helpers.py

@@ -71,7 +71,7 @@ def create_and_execute_pipeline(
     if step_result_value:
         result = execution.result(execution_steps[0]["StepName"])
         assert result == step_result_value, f"Expected {step_result_value}, instead found {result}"
-    return execution_steps
+    return execution, execution_steps
 
 
 def validate_scheduled_pipeline_execution(

tests/integ/sagemaker/workflow/test_step_decorator.py

@@ -22,6 +22,7 @@
 from sagemaker import get_execution_role, utils
 from sagemaker.config import load_sagemaker_config
 from sagemaker.processing import ProcessingInput
+from sagemaker.remote_function.errors import RemoteFunctionError
 from sagemaker.sklearn import SKLearnProcessor
 from sagemaker.remote_function.core.serialization import CloudpickleSerializer
 from sagemaker.s3 import S3Uploader
@@ -577,7 +578,7 @@ def divide(x, y):
     )
 
     try:
-        create_and_execute_pipeline(
+        execution, execution_steps = create_and_execute_pipeline(
             pipeline=pipeline,
             pipeline_name=pipeline_name,
             region_name=region_name,
@@ -587,6 +588,11 @@ def divide(x, y):
             execution_parameters=dict(),
             step_status="Failed",
         )
+
+        step_name = execution_steps[0]["StepName"]
+        with pytest.raises(RemoteFunctionError) as e:
+            execution.result(step_name)
+            assert f"Pipeline step {step_name} is in Failed status." in str(e)
     finally:
         try:
             pipeline.delete()
@@ -720,7 +726,7 @@ def func(var: int):
     )
 
     try:
-        execution_steps = create_and_execute_pipeline(
+        _, execution_steps = create_and_execute_pipeline(
             pipeline=pipeline,
             pipeline_name=pipeline_name,
             region_name=region_name,

tests/unit/sagemaker/workflow/test_pipeline.py

@@ -287,6 +287,7 @@ def test_pipeline_update(
                 "/opt/ml/input/data/sagemaker_remote_function_bootstrap/job_driver.sh",
             ]
         },
+        "TrainingJobStatus": "Completed",
         "OutputDataConfig": {"S3OutputPath": "s3:/my-bucket/my-key"},
         "Environment": {"REMOTE_FUNCTION_SECRET_KEY": "abcdefg"},
     }