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)