Source code for tests.system.providers.amazon.aws.example_sagemaker
# Licensed to the Apache Software Foundation (ASF) under one# or more contributor license agreements. See the NOTICE file# distributed with this work for additional information# regarding copyright ownership. The ASF licenses this file# to you under the Apache License, Version 2.0 (the# "License"); you may not use this file except in compliance# with the License. You may obtain a copy of the License at## http://www.apache.org/licenses/LICENSE-2.0## Unless required by applicable law or agreed to in writing,# software distributed under the License is distributed on an# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY# KIND, either express or implied. See the License for the# specific language governing permissions and limitations# under the License.from__future__importannotationsimportjsonimportloggingimportsubprocessfromdatetimeimportdatetimefromtempfileimportNamedTemporaryFileimportboto3fromairflowimportDAGfromairflow.decoratorsimporttaskfromairflow.models.baseoperatorimportchainfromairflow.operators.pythonimportget_current_contextfromairflow.providers.amazon.aws.hooks.ecrimportEcrHookfromairflow.providers.amazon.aws.operators.s3import(S3CreateBucketOperator,S3CreateObjectOperator,S3DeleteBucketOperator,)fromairflow.providers.amazon.aws.operators.sagemakerimport(SageMakerAutoMLOperator,SageMakerCreateExperimentOperator,SageMakerDeleteModelOperator,SageMakerModelOperator,SageMakerProcessingOperator,SageMakerRegisterModelVersionOperator,SageMakerStartPipelineOperator,SageMakerStopPipelineOperator,SageMakerTrainingOperator,SageMakerTransformOperator,SageMakerTuningOperator,)fromairflow.providers.amazon.aws.sensors.sagemakerimport(SageMakerAutoMLSensor,SageMakerPipelineSensor,SageMakerTrainingSensor,SageMakerTransformSensor,SageMakerTuningSensor,)fromairflow.utils.trigger_ruleimportTriggerRulefromtests.system.providers.amazon.aws.utilsimportENV_ID_KEY,SystemTestContextBuilder,prune_logs
# The URI of a Docker image for handling KNN model training.# To find the URI of a free Amazon-provided image that can be used, substitute your# desired region in the following link and find the URI under "Registry Path".# https://docs.aws.amazon.com/sagemaker/latest/dg/ecr-us-east-1.html#knn-us-east-1.title# This URI should be in the format of {12-digits}.dkr.ecr.{region}.amazonaws.com/knn
# This script will be the entrypoint for the docker image which will handle preprocessing the raw data# NOTE: The following string must remain dedented as it is being written to a file.
[docs]PREPROCESS_SCRIPT_TEMPLATE="""import boto3import numpy as npimport pandas as pddef main(): # Load the dataset from {input_path}/input.csv, split it into train/test # subsets, and write them to {output_path}/ for the Processing Operator. data = pd.read_csv('{input_path}/input.csv') # Split into test and train data data_train, data_test = np.split( data.sample(frac=1, random_state=np.random.RandomState()), [int(0.7 * len(data))] ) # Remove the "answers" from the test set data_test.drop(['class'], axis=1, inplace=True) # Write the splits to disk data_train.to_csv('{output_path}/train.csv', index=False, header=False) data_test.to_csv('{output_path}/test.csv', index=False, header=False) print('Preprocessing Done.')if __name__ == "__main__": main()"""
def_create_ecr_repository(repo_name):execution_role_arn=boto3.client("sts").get_caller_identity()["Arn"]access_policy={"Version":"2012-10-17","Statement":[{"Sid":"Allow access to the system test execution role","Effect":"Allow","Principal":{"AWS":execution_role_arn},"Action":"ecr:*",}],}client=boto3.client("ecr")repo=client.create_repository(repositoryName=repo_name)["repository"]client.set_repository_policy(repositoryName=repo["repositoryName"],policyText=json.dumps(access_policy))returnrepo["repositoryUri"]def_build_and_upload_docker_image(preprocess_script,repository_uri):""" We need a Docker image with the following requirements: - Has numpy, pandas, requests, and boto3 installed - Has our data preprocessing script mounted and set as the entry point """ecr_region=repository_uri.split(".")[3]# Fetch ECR Token to be used for dockercreds=EcrHook(region_name=ecr_region).get_temporary_credentials()[0]withNamedTemporaryFile(mode="w+t")aspreprocessing_script,NamedTemporaryFile(mode="w+t")asdockerfile:preprocessing_script.write(preprocess_script)preprocessing_script.flush()dockerfile.write(f""" FROM public.ecr.aws/amazonlinux/amazonlinux COPY {preprocessing_script.name.split('/')[2]} /preprocessing.py ADD credentials /credentials ENV AWS_SHARED_CREDENTIALS_FILE=/credentials RUN yum install python3 pip -y RUN pip3 install boto3 pandas requests CMD [ "python3", "/preprocessing.py"] """)dockerfile.flush()docker_build_and_push_commands=f""" cp /root/.aws/credentials /tmp/credentials && # login to public ecr repo containing amazonlinux image docker login --username {creds.username} --password {creds.password} public.ecr.aws docker build --platform=linux/amd64 -f {dockerfile.name} -t {repository_uri} /tmp && rm /tmp/credentials && # login again, this time to the private repo we created to hold that specific image aws ecr get-login-password --region {ecr_region} | docker login --username {creds.username} --password {creds.password}{repository_uri} && docker push {repository_uri} """docker_build=subprocess.Popen(docker_build_and_push_commands,shell=True,stdout=subprocess.PIPE,stderr=subprocess.PIPE,)_,stderr=docker_build.communicate()ifdocker_build.returncode!=0:raiseRuntimeError("Failed to prepare docker image for the preprocessing job.\n"f"The following error happened while executing the sequence of bash commands:\n{stderr}")
[docs]defgenerate_data()->str:"""generates a very simple csv dataset with headers"""content="class,x,y\n"# headersforiinrange(SAMPLE_SIZE):content+=f"{i%100},{i},{SAMPLE_SIZE-i}\n"returncontent
[docs]defdelete_ecr_repository(repository_name):client=boto3.client("ecr")# All images must be removed from the repo before it can be deleted.image_ids=client.list_images(repositoryName=repository_name)["imageIds"]client.batch_delete_image(repositoryName=repository_name,imageIds=[{"imageDigest":image["imageDigest"]forimageinimage_ids}],)client.delete_repository(repositoryName=repository_name)
@task(trigger_rule=TriggerRule.ALL_DONE)
[docs]defdelete_model_group(group_name,model_version_arn):sgmk_client=boto3.client("sagemaker")# need to destroy model registered in group firstsgmk_client.delete_model_package(ModelPackageName=model_version_arn)sgmk_client.delete_model_package_group(ModelPackageGroupName=group_name)
[docs]defdelete_docker_image(image_name):docker_build=subprocess.Popen(f"docker rmi {image_name}",shell=True,stdout=subprocess.PIPE,stderr=subprocess.PIPE,)_,stderr=docker_build.communicate()ifdocker_build.returncode!=0:logging.error("Failed to delete local docker image. ""Run 'docker images' to see if you need to clean it yourself.\n"f"error message: {stderr}")
env_id=test_context[ENV_ID_KEY]test_setup=set_up(env_id=env_id,role_arn=test_context[ROLE_ARN_KEY],)create_bucket=S3CreateBucketOperator(task_id="create_bucket",bucket_name=test_setup["bucket_name"],)upload_dataset=S3CreateObjectOperator(task_id="upload_dataset",s3_bucket=test_setup["bucket_name"],s3_key=test_setup["raw_data_s3_key"],data=generate_data(),replace=True,)# [START howto_operator_sagemaker_auto_ml]automl=SageMakerAutoMLOperator(task_id="auto_ML",job_name=test_setup["auto_ml_job_name"],s3_input=test_setup["input_data_uri"],target_attribute="class",s3_output=test_setup["output_data_uri"],role_arn=test_context[ROLE_ARN_KEY],time_limit=30,# will stop the job before it can do anything, but it's not the point here)# [END howto_operator_sagemaker_auto_ml]automl.wait_for_completion=False# just to be able to test the sensor next# [START howto_sensor_sagemaker_auto_ml]await_automl=SageMakerAutoMLSensor(job_name=test_setup["auto_ml_job_name"],task_id="await_auto_ML")# [END howto_sensor_sagemaker_auto_ml]await_automl.poke_interval=10# [START howto_operator_sagemaker_start_pipeline]start_pipeline1=SageMakerStartPipelineOperator(task_id="start_pipeline1",pipeline_name=test_setup["pipeline_name"],)# [END howto_operator_sagemaker_start_pipeline]# [START howto_operator_sagemaker_stop_pipeline]stop_pipeline1=SageMakerStopPipelineOperator(task_id="stop_pipeline1",pipeline_exec_arn=start_pipeline1.output,)# [END howto_operator_sagemaker_stop_pipeline]start_pipeline2=SageMakerStartPipelineOperator(task_id="start_pipeline2",pipeline_name=test_setup["pipeline_name"],)# [START howto_sensor_sagemaker_pipeline]await_pipeline2=SageMakerPipelineSensor(task_id="await_pipeline2",pipeline_exec_arn=start_pipeline2.output,)# [END howto_sensor_sagemaker_pipeline]await_pipeline2.poke_interval=10# [START howto_operator_sagemaker_experiment]create_experiment=SageMakerCreateExperimentOperator(task_id="create_experiment",name=test_setup["experiment_name"])# [END howto_operator_sagemaker_experiment]# [START howto_operator_sagemaker_processing]preprocess_raw_data=SageMakerProcessingOperator(task_id="preprocess_raw_data",config=test_setup["processing_config"],)# [END howto_operator_sagemaker_processing]# [START howto_operator_sagemaker_training]train_model=SageMakerTrainingOperator(task_id="train_model",config=test_setup["training_config"],)# [END howto_operator_sagemaker_training]# SageMakerTrainingOperator waits by default, setting as False to test the Sensor below.train_model.wait_for_completion=False# [START howto_sensor_sagemaker_training]await_training=SageMakerTrainingSensor(task_id="await_training",job_name=test_setup["training_job_name"],)# [END howto_sensor_sagemaker_training]# [START howto_operator_sagemaker_model]create_model=SageMakerModelOperator(task_id="create_model",config=test_setup["model_config"],)# [END howto_operator_sagemaker_model]# [START howto_operator_sagemaker_register]register_model=SageMakerRegisterModelVersionOperator(task_id="register_model",image_uri=test_setup["inference_code_image"],model_url=test_setup["model_trained_weights"],package_group_name=test_setup["model_package_group_name"],)# [END howto_operator_sagemaker_register]# [START howto_operator_sagemaker_tuning]tune_model=SageMakerTuningOperator(task_id="tune_model",config=test_setup["tuning_config"],)# [END howto_operator_sagemaker_tuning]# SageMakerTuningOperator waits by default, setting as False to test the Sensor below.tune_model.wait_for_completion=False# [START howto_sensor_sagemaker_tuning]await_tuning=SageMakerTuningSensor(task_id="await_tuning",job_name=test_setup["tuning_job_name"],)# [END howto_sensor_sagemaker_tuning]# [START howto_operator_sagemaker_transform]test_model=SageMakerTransformOperator(task_id="test_model",config=test_setup["transform_config"],)# [END howto_operator_sagemaker_transform]# SageMakerTransformOperator waits by default, setting as False to test the Sensor below.test_model.wait_for_completion=False# [START howto_sensor_sagemaker_transform]await_transform=SageMakerTransformSensor(task_id="await_transform",job_name=test_setup["transform_job_name"],)# [END howto_sensor_sagemaker_transform]# [START howto_operator_sagemaker_delete_model]delete_model=SageMakerDeleteModelOperator(task_id="delete_model",config={"ModelName":test_setup["model_name"]},)# [END howto_operator_sagemaker_delete_model]delete_model.trigger_rule=TriggerRule.ALL_DONEdelete_bucket=S3DeleteBucketOperator(task_id="delete_bucket",trigger_rule=TriggerRule.ALL_DONE,bucket_name=test_setup["bucket_name"],force_delete=True,)log_cleanup=prune_logs([# Format: ('log group name', 'log stream prefix')("/aws/sagemaker/ProcessingJobs",env_id),("/aws/sagemaker/TrainingJobs",env_id),("/aws/sagemaker/TransformJobs",env_id),])chain(# TEST SETUPtest_context,test_setup,create_bucket,upload_dataset,# TEST BODYautoml,await_automl,start_pipeline1,start_pipeline2,stop_pipeline1,await_pipeline2,create_experiment,preprocess_raw_data,train_model,await_training,create_model,register_model,tune_model,await_tuning,test_model,await_transform,# TEST TEARDOWNdelete_ecr_repository(test_setup["ecr_repository_name"]),delete_model_group(test_setup["model_package_group_name"],register_model.output),delete_model,delete_bucket,delete_experiment(test_setup["experiment_name"]),delete_pipeline(test_setup["pipeline_name"]),delete_docker_image(test_setup["docker_image"]),log_cleanup,)fromtests.system.utils.watcherimportwatcher# This test needs watcher in order to properly mark success/failure# when "tearDown" task with trigger rule is part of the DAGlist(dag.tasks)>>watcher()fromtests.system.utilsimportget_test_run# noqa: E402# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest)
