[AWS SageMaker] Integration tests automation

.gitignore

@@ -66,3 +66,5 @@ _artifacts
 # Generated Python SDK documentation
 docs/_build
 
+# Any environment variable files
+**/*/.env

components/aws/sagemaker/codebuild/integration-test.buildspec.yml

@@ -1,14 +1,21 @@
 version: 0.2      
+
+env:
+  variables:
+    CONTAINER_VARIABLES: "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI EKS_PRIVATE_SUBNETS EKS_PUBLIC_SUBNETS PYTEST_ADDOPTS S3_DATA_BUCKET EKS_EXISTING_CLUSTER SAGEMAKER_EXECUTION_ROLE_ARN REGION"
+
 phases:
   build:
     commands:
       - cd components/aws
-      - docker build . -f ./sagemaker/tests/integration_tests/Dockerfile -t amazon/integration-test-image --quiet
+      - docker build . -f ./sagemaker/tests/integration_tests/Dockerfile -t amazon/integration-test-image
+
+      - cd sagemaker/codebuild/scripts && export CONTAINER_VARIABLE_ARGUMENTS="$(./construct_environment_array.sh)"
 
       # Run the container and copy the results to /tmp
-      # Passes all host environment variables through to the container
-      - docker run --name integration-test-container $(env | cut -f1 -d= | sed 's/^/-e /') amazon/integration-test-image
-      - docker cp integration-test-container:/app/tests/integration_tests/integration_tests.log /tmp/results.xml
+      # Passes all listed host environment variables through to the container
+      - docker run --name integration-test-container $(echo $CONTAINER_VARIABLE_ARGUMENTS) amazon/integration-test-image
+      - docker cp integration-test-container:/tests/integration_tests/integration_tests.log /tmp/results.xml
       - docker rm -f integration-test-container
 
 reports:

components/aws/sagemaker/codebuild/scripts/construct_environment_array.sh

@@ -0,0 +1,10 @@
+#!/usr/bin/env bash
+
+# This script breaks up a string of environment variable names into a list of
+# parameters that `docker run` accepts. This needs to be made into a script
+# for CodeBuild because these commands do not run in dash - the default terminal
+# on the CodeBuild standard images.
+
+IFS=' ' read -a variable_array <<< $CONTAINER_VARIABLES
+printf -v CONTAINER_VARIABLE_ARGUMENTS -- "--env %s " "${variable_array[@]}"
+echo $CONTAINER_VARIABLE_ARGUMENTS

components/aws/sagemaker/tests/integration_tests/.env.example

@@ -0,0 +1,12 @@
+# If you would like to override the credentials for the container
+# AWS_ACCESS_KEY_ID=
+# AWS_SECRET_ACCESS_KEY=
+# AWS_SESSION_TOKEN=
+
+REGION=us-east-1
+
+SAGEMAKER_EXECUTION_ROLE_ARN=arn:aws:iam::123456789012:role/service-role/AmazonSageMaker-ExecutionRole-Example
+S3_DATA_BUCKET=my-data-bucket
+
+# If you hope to use an existing EKS cluster, rather than creating a new one.
+# EKS_EXISTING_CLUSTER=my-eks-cluster

components/aws/sagemaker/tests/integration_tests/Dockerfile

@@ -0,0 +1,45 @@
+FROM continuumio/miniconda:4.7.12
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    curl \
+    wget \
+    git \
+    vim \
+    sudo \
+    jq
+
+# Install eksctl
+RUN curl --location "https://github.com/weaveworks/eksctl/releases/download/0.19.0/eksctl_$(uname -s)_amd64.tar.gz" | tar xz -C /tmp \
+ && mv /tmp/eksctl /usr/local/bin
+
+# Install aws-iam-authenticator
+RUN curl -S -o /usr/local/bin/aws-iam-authenticator https://amazon-eks.s3.us-west-2.amazonaws.com/1.16.8/2020-04-16/bin/linux/amd64/aws-iam-authenticator \
+ && chmod +x /usr/local/bin/aws-iam-authenticator
+
+# Install Kubectl
+RUN curl -LO https://storage.googleapis.com/kubernetes-release/release/v1.18.0/bin/linux/amd64/kubectl \
+ && chmod +x ./kubectl \
+ && mv ./kubectl /usr/local/bin/kubectl
+
+# Install Argo CLI
+RUN curl -sSL -o /usr/local/bin/argo https://github.com/argoproj/argo/releases/download/v2.8.0/argo-linux-amd64 \
+ && chmod +x /usr/local/bin/argo
+
+# Copy conda environment early to avoid cache busting
+COPY ./sagemaker/tests/integration_tests/environment.yml environment.yml
+
+# Create conda environment for running tests and set as start-up environment
+RUN conda env create -f environment.yml
+RUN echo "source activate kfp_test_env" > ~/.bashrc
+ENV PATH "/opt/conda/envs/kfp_test_env/bin":$PATH
+
+# Environment variables to be used by tests
+ENV REGION="us-west-2"
+ENV SAGEMAKER_EXECUTION_ROLE_ARN="arn:aws:iam::1234567890:role/sagemaker-role"
+ENV S3_DATA_BUCKET="kfp-test-data"
+ENV MINIO_LOCAL_PORT=9000
+ENV KFP_NAMESPACE="kubeflow"
+
+COPY ./sagemaker/ .
+
+ENTRYPOINT [ "/bin/bash", "./tests/integration_tests/scripts/run_integration_tests" ]

components/aws/sagemaker/tests/integration_tests/resources/definition/create_endpoint_pipeline.py

@@ -34,7 +34,7 @@ def create_endpoint_pipeline(
         model_artifact_url=model_artifact_url,
         network_isolation=network_isolation,
         role=role,
-    ).apply(use_aws_secret("aws-secret", "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"))
+    )
 
     sagemaker_deploy_op(
         region=region,
@@ -46,7 +46,7 @@ def create_endpoint_pipeline(
         instance_type_1=instance_type_1,
         initial_instance_count_1=initial_instance_count_1,
         initial_variant_weight_1=initial_variant_weight_1,
-    ).apply(use_aws_secret("aws-secret", "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"))
+    )
 
 
 if __name__ == "__main__":

components/aws/sagemaker/tests/integration_tests/resources/definition/create_model_pipeline.py

@@ -26,7 +26,7 @@ def create_model_pipeline(
         model_artifact_url=model_artifact_url,
         network_isolation=network_isolation,
         role=role,
-    ).apply(use_aws_secret("aws-secret", "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"))
+    )
 
 
 if __name__ == "__main__":

components/aws/sagemaker/tests/integration_tests/resources/definition/hpo_pipeline.py

@@ -56,7 +56,7 @@ def hpo_pipeline(
         network_isolation=network_isolation,
         max_wait_time=max_wait_time,
         role=role,
-    ).apply(use_aws_secret("aws-secret", "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"))
+    )
 
 
 if __name__ == "__main__":

components/aws/sagemaker/tests/integration_tests/resources/definition/training_pipeline.py

@@ -46,7 +46,7 @@ def training_pipeline(
         max_wait_time=max_wait_time,
         checkpoint_config=checkpoint_config,
         role=role,
-    ).apply(use_aws_secret("aws-secret", "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"))
+    )
 
 
 if __name__ == "__main__":

components/aws/sagemaker/tests/integration_tests/resources/definition/transform_job_pipeline.py

@@ -40,7 +40,7 @@ def batch_transform_pipeline(
         model_artifact_url=model_artifact_url,
         network_isolation=network_isolation,
         role=role,
-    ).apply(use_aws_secret("aws-secret", "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"))
+    )
 
     sagemaker_batch_transform_op(
         region=region,
@@ -57,7 +57,7 @@ def batch_transform_pipeline(
         split_type=split_type,
         compression_type=compression_type,
         output_location=output_location,
-    ).apply(use_aws_secret("aws-secret", "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"))
+    )
 
 
 if __name__ == "__main__":

components/aws/sagemaker/tests/integration_tests/scripts/generate_iam_role

@@ -0,0 +1,68 @@
+#!/usr/bin/env bash
+
+# Helper script to generate an IAM Role needed to install operator using role-based authentication.
+#
+# Run as:
+# $ ./generate_iam_role ${cluster_arn/cluster_name} ${role_name} ${cluster_region} [optional: ${service_namespace} ${service_account}]
+#
+
+CLUSTER_ARN="${1}"
+ROLE_NAME="${2}"
+CLUSTER_REGION="${3:-us-east-1}"
+SERVICE_NAMESPACE="${4:-kubeflow}"
+SERVICE_ACCOUNT="${5:-pipeline-runner}"
+aws_account=$(aws sts get-caller-identity --query Account --output text)
+trustfile="trust.json"
+
+cwd=$(dirname $(realpath $0))
+
+# if using an existing cluster, use the cluster arn to get the region and cluster name
+# example, cluster_arn=arn:aws:eks:us-east-1:12345678910:cluster/test
+cluster_name=$(echo ${CLUSTER_ARN} | cut -d'/' -f2)
+
+# A function to get the OIDC_ID associated with an EKS cluster
+function get_oidc_id {
+  # TODO: Ideally this should be based on version compatibility instead of command failure
+  eksctl utils associate-iam-oidc-provider --cluster ${cluster_name} --region ${CLUSTER_REGION} --approve
+  if [[ $? -ge 1 ]]; then
+    eksctl utils associate-iam-oidc-provider --name ${cluster_name} --region ${CLUSTER_REGION} --approve
+  fi
+
+  local oidc=$(aws eks describe-cluster --name ${cluster_name} --region ${CLUSTER_REGION} --query cluster.identity.oidc.issuer --output text)
+  oidc_id=$(echo ${oidc} | rev | cut -d'/' -f1 | rev)
+}
+
+# A function that generates an IAM role for the given account, cluster, namespace, region
+# Parameter:
+#    $1: Name of the trust file to generate.
+function create_namespaced_iam_role {
+  local trustfile="${1}"
+  # Check if role already exists
+  aws iam get-role --role-name ${ROLE_NAME}
+  if [[ $? -eq 0 ]]; then
+    echo "A role for this cluster and namespace already exists in this account, assuming sagemaker access and proceeding."
+  else
+    echo "IAM Role does not exist, creating a new Role for the cluster"
+    aws iam create-role --role-name ${ROLE_NAME} --assume-role-policy-document file://${trustfile} --output=text --query "Role.Arn"
+    aws iam attach-role-policy --role-name ${ROLE_NAME}  --policy-arn arn:aws:iam::aws:policy/AmazonSageMakerFullAccess
+  fi
+}
+
+# Remove the generated trust file
+# Parameter:
+#    $1: Name of the trust file to delete.
+function delete_generated_file {
+  rm "${1}" 
+}
+
+echo "Get the OIDC ID for the cluster"
+get_oidc_id
+echo "Delete the trust json file if it already exists"
+delete_generated_file "${trustfile}"
+echo "Generate a trust json"
+"$cwd"/generate_trust_policy ${CLUSTER_REGION} ${aws_account} ${oidc_id} ${SERVICE_NAMESPACE} ${SERVICE_ACCOUNT} > "${trustfile}"
+echo "Create the IAM Role using these values"
+create_namespaced_iam_role "${trustfile}"
+echo "Cleanup for the next run"
+delete_generated_file "${trustfile}"
+

components/aws/sagemaker/tests/integration_tests/scripts/generate_trust_policy

@@ -0,0 +1,39 @@
+#!/usr/bin/env bash
+
+# Helper script to generate trust policy needed to install operator using role-based authentication.
+#
+# Run as:
+# $ ./generate_trust_policy ${EKS_CLUSTER_REGION} ${AWS_ACCOUNT_ID} ${OIDC_ID} ${SERVICE_NAMESPACE} ${SERVICE_ACCOUNT} > trust.json
+#
+# For example:
+# $ ./generate_trust_policy us-west-2 123456789012 D48675832CA65BD10A532F597OIDCID > trust.json
+# This will create a file `trust.json` containing a role policy that enables the operator in an EKS cluster to assume AWS roles.
+#
+# The SERVICE_NAMESPACE parameter is for when you want to run Kubeflow in a custom namespace other than "kubeflow".
+# The SERVICE_ACCOUNT parameter is for when you want to give permissions to a service account other than the default "pipeline-runner".
+
+cluster_region="$1"
+account_number="$2"
+oidc_id="$3"
+service_namespace="${4}"
+service_account="${5}"
+
+printf '{
+  "Version": "2012-10-17",
+  "Statement": [
+    {
+      "Effect": "Allow",
+      "Principal": {
+        "Federated": "arn:aws:iam::'"${account_number}"':oidc-provider/oidc.eks.'"${cluster_region}"'.amazonaws.com/id/'"${oidc_id}"'"
+      },
+      "Action": "sts:AssumeRoleWithWebIdentity",
+      "Condition": {
+        "StringEquals": {
+          "oidc.eks.'"${cluster_region}"'.amazonaws.com/id/'"${oidc_id}"':aud": "sts.amazonaws.com",
+          "oidc.eks.'"${cluster_region}"'.amazonaws.com/id/'"${oidc_id}"':sub": "system:serviceaccount:'"${service_namespace}"':'"${service_account}"'"
+        }
+      }
+    }
+  ]
+}
+'