From 898842486f79a80e88712614bd981ba2475b5f82 Mon Sep 17 00:00:00 2001
From: Akshay Shah <akshayshah0803@gmail.com>
Date: Mon, 12 Oct 2026 11:59:35 -0500
Subject: [PATCH] fix(train): stop EvaluationPipelineExecution.wait() on a
 Stopped execution

wait() returned on the target status and raised on statuses containing
'failed', but did not treat Stopped as terminal. With the default
timeout=None, waiting on a stopped evaluation (execution.stop() or a stop
from the console) polled forever; with a timeout it eventually raised
TimeoutExceededError saying the job is still running.

Raise FailedStatusError for Stopped (unless it is the target status), in
both the Jupyter and the terminal branches, as trainer_wait already does
for Stopped training jobs.
---
 .../src/sagemaker/train/evaluate/execution.py | 22 +++++++++++++++++++
 .../unit/train/evaluate/test_execution.py     | 22 +++++++++++++++++++
 2 files changed, 44 insertions(+)

diff --git a/sagemaker-train/src/sagemaker/train/evaluate/execution.py b/sagemaker-train/src/sagemaker/train/evaluate/execution.py
index 04810ac97d..b1938c125a 100644
--- a/sagemaker-train/src/sagemaker/train/evaluate/execution.py
+++ b/sagemaker-train/src/sagemaker/train/evaluate/execution.py
@@ -1321,6 +1321,17 @@ def get_cached_mlflow_url():
                         reason=self.status.failure_reason,
                     )
 
+                # Stopped is terminal too; without this a stopped execution is polled forever
+                # (or until the timeout, reported as "still running").
+                if current_status == "Stopped":
+                    from sagemaker.core.utils.exceptions import FailedStatusError
+
+                    raise FailedStatusError(
+                        resource_type="PipelineExecution",
+                        status=current_status,
+                        reason=self.status.failure_reason,
+                    )
+
                 if timeout is not None and time.time() - start_time >= timeout:
                     from sagemaker.core.utils.exceptions import TimeoutExceededError
 
@@ -1441,6 +1452,17 @@ def get_cached_mlflow_url():
                         reason=self.status.failure_reason,
                     )
 
+                # Stopped is terminal too; without this a stopped execution is polled forever
+                # (or until the timeout, reported as "still running").
+                if current_status == "Stopped":
+                    from sagemaker.core.utils.exceptions import FailedStatusError
+
+                    raise FailedStatusError(
+                        resource_type="PipelineExecution",
+                        status=current_status,
+                        reason=self.status.failure_reason,
+                    )
+
                 if timeout is not None and time.time() - start_time >= timeout:
                     from sagemaker.core.utils.exceptions import TimeoutExceededError
 
diff --git a/sagemaker-train/tests/unit/train/evaluate/test_execution.py b/sagemaker-train/tests/unit/train/evaluate/test_execution.py
index 0ccfaa7f7c..d834dcb95a 100644
--- a/sagemaker-train/tests/unit/train/evaluate/test_execution.py
+++ b/sagemaker-train/tests/unit/train/evaluate/test_execution.py
@@ -1112,6 +1112,28 @@ def test_wait_fails_on_failed_status(self):
         with pytest.raises(FailedStatusError):
             execution.wait(target_status="Succeeded", poll=1)
 
+    @patch("sagemaker.train.evaluate.execution.time.sleep")
+    def test_wait_stops_on_stopped_status(self, mock_sleep):
+        """A Stopped execution is terminal: wait() must not poll it forever."""
+        mock_sleep.side_effect = AssertionError("wait() kept polling a Stopped execution")
+        execution = EvaluationPipelineExecution(
+            name="test", status=PipelineExecutionStatus(overall_status="Executing")
+        )
+
+        mock_pe = MagicMock()
+        mock_pe.pipeline_execution_status = "Stopped"
+        mock_pe.failure_reason = None
+        mock_pe.last_modified_time = datetime.now()
+        mock_pe.get_all_steps.return_value = iter([])
+
+        execution._pipeline_execution = mock_pe
+
+        with pytest.raises(FailedStatusError):
+            execution.wait(target_status="Succeeded", poll=1)
+
+        mock_pe.get_all_steps.return_value = iter([])
+        execution.wait(target_status="Stopped", poll=1)
+
     @patch("sagemaker.train.evaluate.execution.time")
     def test_wait_timeout_exceeded(self, mock_time):
         """Test wait raises exception on timeout."""