From 2a062ce05ebce56517ab14a7f7b6171f4126acf7 Mon Sep 17 00:00:00 2001
From: Akshay Shah <akshayshah0803@gmail.com>
Date: Fri, 9 Oct 2026 20:57:43 -0500
Subject: [PATCH] fix(mlops): stop PipelineExecution.wait() on a Stopped
 execution

The PipelineExecutionComplete waiter only had acceptors for Succeeded and
Failed. Stopped is also a terminal pipeline execution status, so waiting on
a stopped execution kept polling until max_attempts (30 minutes by default)
and then failed with 'Max attempts exceeded'. result() swallows terminal
failures but not that error, so it raised too instead of reporting the step.

Treat Stopped as a terminal failure, like Failed.
---
 .../src/sagemaker/mlops/workflow/pipeline.py  |  6 +++++
 .../tests/unit/workflow/test_pipeline.py      | 27 +++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/pipeline.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/pipeline.py
index 05622e9ca0..e37ff9a1b0 100644
--- a/sagemaker-mlops/src/sagemaker/mlops/workflow/pipeline.py
+++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/pipeline.py
@@ -1127,6 +1127,12 @@ def wait(self, delay=30, max_attempts=60):
                                 "state": "failure",
                                 "argument": "PipelineExecutionStatus",
                             },
+                            {
+                                "expected": "Stopped",
+                                "matcher": "path",
+                                "state": "failure",
+                                "argument": "PipelineExecutionStatus",
+                            },
                         ],
                     }
                 },
diff --git a/sagemaker-mlops/tests/unit/workflow/test_pipeline.py b/sagemaker-mlops/tests/unit/workflow/test_pipeline.py
index 759b906c02..bc137afc5b 100644
--- a/sagemaker-mlops/tests/unit/workflow/test_pipeline.py
+++ b/sagemaker-mlops/tests/unit/workflow/test_pipeline.py
@@ -518,6 +518,33 @@ def test_pipeline_execution_wait(mock_session):
         mock_waiter.return_value.wait.assert_called_once()
 
 
+@pytest.mark.parametrize("status", ["Failed", "Stopped"])
+def test_pipeline_execution_wait_stops_on_terminal_failure(mock_session, status):
+    import boto3
+    from botocore.exceptions import WaiterError
+    from botocore.stub import Stubber
+    from sagemaker.mlops.workflow.pipeline import PipelineExecution
+
+    client = boto3.client(
+        "sagemaker",
+        region_name="us-west-2",
+        aws_access_key_id="testing",
+        aws_secret_access_key="testing",
+    )
+    mock_session.sagemaker_client = client
+    execution = PipelineExecution(arn="arn", sagemaker_session=mock_session)
+
+    with Stubber(client) as stubber:
+        stubber.add_response(
+            "describe_pipeline_execution",
+            {"PipelineExecutionStatus": status},
+            {"PipelineExecutionArn": "arn"},
+        )
+        with pytest.raises(WaiterError, match="terminal failure state"):
+            execution.wait(delay=0, max_attempts=3)
+        stubber.assert_no_pending_responses()
+
+
 def test_get_function_step_result_invalid_step(mock_session):
     from sagemaker.mlops.workflow.pipeline import get_function_step_result