fix(train): stop EvaluationPipelineExecution.wait() on a Stopped execution

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
 

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."""