chore(components): AWS SageMaker - Fix leaking Workteam(GroundTruth) resources (#4536)

This commit is contained in:
Meghna Baijal 2020-09-30 13:08:54 -07:00 committed by GitHub
parent 58584e1d1f
commit 237795539f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 10 additions and 3 deletions

View File

@ -110,12 +110,19 @@ def test_groundtruth_labeling_job(
)
assert response["LabelingJobStatus"] in ["Stopping", "Stopped"]
finally:
# Cleanup the SageMaker Resources
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
)
if workteam_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