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