[AWS SageMaker] Fix small bugs (#4161)

* fix small bugs * add SKIP_OIDC_SETUP config * address comments * address comments and add KFP_VERSION to .env * typo
kubeflow · Sep 2, 2020 · 05398cf · 05398cf
1 parent 7840d30
commit 05398cf
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 11 deletions.
diff --git a/components/aws/sagemaker/tests/integration_tests/.env.example b/components/aws/sagemaker/tests/integration_tests/.env.example
@@ -10,9 +10,13 @@ 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
+# SKIP_KFP_OIDC_SETUP=true
 
 # If you would like to skip the FSx set-up and tests
 # SKIP_FSX_TESTS=true
 
+# Version of the KFP to install on the cluster
+# KFP_VERSION=0.5.1
+
 # If you have an IAM role that the EKS cluster should assume for the "assume role" tests
 # ASSUMED_ROLE_NAME=my-assumed-role
diff --git a/components/aws/sagemaker/tests/integration_tests/scripts/run_integration_tests b/components/aws/sagemaker/tests/integration_tests/scripts/run_integration_tests
@@ -21,6 +21,8 @@ EKS_CLUSTER_VERSION=${EKS_CLUSTER_VERSION:-"1.15"} # EKS cluster K8s version
 EKS_NODE_COUNT=${EKS_NODE_COUNT:-"1"} # The initial node count of the EKS cluster
 EKS_PUBLIC_SUBNETS=${EKS_PUBLIC_SUBNETS:-""}
 EKS_PRIVATE_SUBNETS=${EKS_PRIVATE_SUBNETS:-""}
+SKIP_KFP_OIDC_SETUP=${SKIP_KFP_OIDC_SETUP:-"false"}
+KFP_VERSION=${KFP_VERSION:-"0.5.1"}
 
 ### Testing parameters
 MINIO_LOCAL_PORT=${MINIO_LOCAL_PORT:-9000}
@@ -81,7 +83,8 @@ function cleanup() {
 
   cleanup_kfp
   delete_assumed_role
-  delete_oidc_role
+
+  [ "${SKIP_KFP_OIDC_SETUP}" == "false" ] && delete_oidc_role
 
   if [[ "${SKIP_FSX_TESTS}" == "false" ]]; then
     delete_fsx_instance
@@ -120,10 +123,9 @@ function delete_eks() {
 function install_kfp() {
   echo "[Installing KFP] Applying KFP manifests"
 
-  PIPELINE_VERSION=0.5.1
-  kubectl apply -k github.com/kubeflow/pipelines/manifests/kustomize/cluster-scoped-resources?ref=$PIPELINE_VERSION
+  kubectl apply -k github.com/kubeflow/pipelines/manifests/kustomize/cluster-scoped-resources?ref=$KFP_VERSION
   kubectl wait --for condition=established --timeout=60s crd/applications.app.k8s.io
-  kubectl apply -k github.com/kubeflow/pipelines/manifests/kustomize/env/dev?ref=$PIPELINE_VERSION
+  kubectl apply -k github.com/kubeflow/pipelines/manifests/kustomize/env/dev?ref=$KFP_VERSION
 
   echo "[Installing KFP] Port-forwarding Minio"
 
@@ -229,11 +231,13 @@ else
   wait
 fi
 
-generate_oidc_role_name
-"$cwd"/generate_iam_role ${EKS_CLUSTER_NAME} ${OIDC_ROLE_NAME} ${REGION} ${KFP_NAMESPACE} ${KFP_SERVICE_ACCOUNT}
-install_kfp
-install_oidc_role
+if [[ "${SKIP_KFP_OIDC_SETUP}" == "false" ]]; then
+  generate_oidc_role_name
+  "$cwd"/generate_iam_role ${EKS_CLUSTER_NAME} ${OIDC_ROLE_NAME} ${REGION} ${KFP_NAMESPACE} ${KFP_SERVICE_ACCOUNT}
+fi
 
+install_kfp
+[ "${SKIP_KFP_OIDC_SETUP}" == "false" ] && install_oidc_role
 generate_assumed_role
 
 pytest_args=( --region "${REGION}" --role-arn "${SAGEMAKER_EXECUTION_ROLE_ARN}" \
@@ -252,4 +256,4 @@ fi
 
 DIR_THIS_FILE="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
 
-cd $DIR_THIS_FILE/../ && python -m pytest "${pytest_args[@]}" --junitxml ./integration_tests.log -n $(nproc)
+cd $DIR_THIS_FILE/../ && python -m pytest "${pytest_args[@]}" --junitxml ./integration_tests.log -n $(nproc)
diff --git a/samples/contrib/aws-samples/simple_train_pipeline/training-pipeline.py b/samples/contrib/aws-samples/simple_train_pipeline/training-pipeline.py
@@ -34,7 +34,7 @@
 }
 
 channelObj['ChannelName'] = 'train'
-channelObj['DataSource']['S3DataSource']['S3Uri'] = 's3://kubeflow-pipeline-data/mnist_kmeans_example/data'
+channelObj['DataSource']['S3DataSource']['S3Uri'] = 's3://kubeflow-pipeline-data/mnist_kmeans_example/train_data'
 channelObjList.append(copy.deepcopy(channelObj))
 
 
@@ -53,7 +53,7 @@ def training(
         instance_count=1,
         volume_size=50,
         max_run_time=3600,
-        model_artifact_path='s3://kubeflow-pipeline-data/mnist_kmeans_example/data',
+        model_artifact_path='s3://kubeflow-pipeline-data/mnist_kmeans_example/output',
         output_encryption_key='',
         network_isolation=True,
         traffic_encryption=False,