def test_terminate_trainingjob(kfp_client, experiment_id, region,
sagemaker_client):
test_file_dir = "resources/config/simple-mnist-training"
download_dir = utils.mkdir(
os.path.join(test_file_dir + "/generated_test_terminate"))
test_params = utils.load_params(
utils.replace_placeholders(
os.path.join(test_file_dir, "config.yaml"),
os.path.join(download_dir, "config.yaml"),
))
input_job_name = test_params["Arguments"]["job_name"] = (
utils.generate_random_string(4) + "-terminate-job")
run_id, _, workflow_json = kfp_client_utils.compile_run_monitor_pipeline(
kfp_client,
experiment_id,
test_params["PipelineDefinition"],
test_params["Arguments"],
download_dir,
test_params["TestName"],
60,
"running",
)
print(
f"Terminating run: {run_id} where Training job_name: {input_job_name}")
kfp_client_utils.terminate_run(kfp_client, run_id)
response = sagemaker_utils.describe_training_job(sagemaker_client,
input_job_name)
assert response["TrainingJobStatus"] in ["Stopping", "Stopped"]
utils.remove_dir(download_dir)
def test_groundtruth_labeling_job(kfp_client, experiment_id, region,
sagemaker_client, test_file_dir):
download_dir = utils.mkdir(os.path.join(test_file_dir + "/generated"))
test_params = utils.load_params(
utils.replace_placeholders(
os.path.join(test_file_dir, "config.yaml"),
os.path.join(download_dir, "config.yaml"),
))
# 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:
workteam_name, workteam_arn = create_initial_workteam(
kfp_client,
experiment_id,
region,
sagemaker_client,
"resources/config/create-workteam",
download_dir,
)
test_params["Arguments"]["workteam_arn"] = workteam_arn
# Generate the ground_truth_train_job_name based on the workteam which will be used for labeling.
test_params["Arguments"][
"ground_truth_train_job_name"] = ground_truth_train_job_name = (
test_params["Arguments"]["ground_truth_train_job_name"] +
"-by-" + workteam_name)
run_id, _, _ = kfp_client_utils.compile_run_monitor_pipeline(
kfp_client,
experiment_id,
test_params["PipelineDefinition"],
test_params["Arguments"],
download_dir,
test_params["TestName"],
test_params["Timeout"],
test_params["StatusToCheck"],
)
response = sagemaker_utils.describe_labeling_job(
sagemaker_client, ground_truth_train_job_name)
assert response["LabelingJobStatus"] == "InProgress"
# Verify that the workteam has the specified labeling job
labeling_jobs = sagemaker_utils.list_labeling_jobs_for_workteam(
sagemaker_client, workteam_arn)
assert len(labeling_jobs["LabelingJobSummaryList"]) == 1
assert (labeling_jobs["LabelingJobSummaryList"][0]["LabelingJobName"]
== ground_truth_train_job_name)
# Test terminate functionality
print(
f"Terminating run: {run_id} where GT job_name: {ground_truth_train_job_name}"
)
kfp_client_utils.terminate_run(kfp_client, run_id)
response = sagemaker_utils.describe_labeling_job(
sagemaker_client, ground_truth_train_job_name)
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)
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)
