Merge pull request #3 from aurora-opensource/patches/handle-stop-state

fix(component): Handle Stop State for Sagemaker Jobs Fixes #kubeflow#7117
burrito-party · Mar 7, 2022 · 40c347e · 40c347e
2 parents 87c29ac + 0167ec4
commit 40c347e
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git a/components/aws/sagemaker/tests/unit_tests/tests/train/test_train_component.py b/components/aws/sagemaker/tests/unit_tests/tests/train/test_train_component.py
@@ -120,6 +120,19 @@ def test_get_job_status(self):
             SageMakerJobStatus(is_completed=True, raw_status="Completed"),
         )
 
+        self.component._get_debug_rule_status = MagicMock(
+            return_value=SageMakerJobStatus(
+                is_completed=True, has_error=False, raw_status="Stopped"
+            )
+        )
+        self.component._sm_client.describe_training_job.return_value = {
+            "TrainingJobStatus": "Stopped"
+        }
+        self.assertEqual(
+            self.component._get_job_status(),
+            SageMakerJobStatus(is_completed=True, raw_status="Stopped"),
+        )
+
         self.component._sm_client.describe_training_job.return_value = {
             "TrainingJobStatus": "Failed",
             "FailureReason": "lolidk",

diff --git a/components/aws/sagemaker/train/src/sagemaker_training_component.py b/components/aws/sagemaker/train/src/sagemaker_training_component.py
@@ -54,7 +54,7 @@ def _get_job_status(self) -> SageMakerJobStatus:
         )
         status = response["TrainingJobStatus"]
 
-        if status == "Completed":
+        if status == "Completed" or status == "Stopped":
             return self._get_debug_rule_status()
         if status == "Failed":
             message = response["FailureReason"]