chore(components): SageMaker integ tests, fix for unbound variable (#…

…4595)
kubeflow · Oct 9, 2020 · e87d74f · e87d74f
1 parent 2a1fe45
commit e87d74f
Showing 1 changed file with 18 additions and 13 deletions.
diff --git a/...nents/aws/sagemaker/tests/integration_tests/component_tests/test_groundtruth_component.py b/...nents/aws/sagemaker/tests/integration_tests/component_tests/test_groundtruth_component.py
@@ -51,6 +51,7 @@ def test_groundtruth_labeling_job(
         )
     )
 
+    workteam_arn = ground_truth_train_job_name = None
     # Verify the GroundTruthJob was created in SageMaker and is InProgress.
     # TODO: Add a bot to complete the labeling job and check for completion instead.
     try:
@@ -110,20 +111,24 @@ def test_groundtruth_labeling_job(
         )
         assert response["LabelingJobStatus"] in ["Stopping", "Stopped"]
     finally:
-        # Check if terminate failed, and stop the labeling job
-        labeling_jobs = sagemaker_utils.list_labeling_jobs_for_workteam(
-            sagemaker_client, workteam_arn
+        print(
+            f"Clean up workteam: {workteam_arn} and GT job: {ground_truth_train_job_name}"
         )
-        if len(labeling_jobs["LabelingJobSummaryList"]) > 0:
-            sagemaker_utils.stop_labeling_job(
-                sagemaker_client, ground_truth_train_job_name
-            )
-
-        # Cleanup the workteam
-        workteams = sagemaker_utils.list_workteams(sagemaker_client)["Workteams"]
-        workteam_names = list(map((lambda x: x["WorkteamName"]), workteams))
-        if workteam_name in workteam_names:
-            sagemaker_utils.delete_workteam(sagemaker_client, workteam_name)
+        if workteam_arn:
+            if ground_truth_train_job_name:
+                # Check if terminate failed, and stop the labeling job
+                labeling_jobs = sagemaker_utils.list_labeling_jobs_for_workteam(
+                    sagemaker_client, workteam_arn
+                )
+                if len(labeling_jobs["LabelingJobSummaryList"]) > 0:
+                    sagemaker_utils.stop_labeling_job(
+                        sagemaker_client, ground_truth_train_job_name
+                    )
+            # Cleanup the workteam
+            workteams = sagemaker_utils.list_workteams(sagemaker_client)["Workteams"]
+            workteam_names = list(map((lambda x: x["WorkteamName"]), workteams))
+            if workteam_name in workteam_names:
+                sagemaker_utils.delete_workteam(sagemaker_client, workteam_name)
 
     # Delete generated files
     utils.remove_dir(download_dir)