#!/bin/bash

# Do we have GPU support?
nvidia-smi > /dev/null 2>&1
if [ $? -eq 0 ]; then
  # check if we have nvidia-docker
  NVIDIA_DOCKER=`rpm -qa | grep -c nvidia-docker2`
  if [ $NVIDIA_DOCKER -eq 0 ]; then
    # Install nvidia-docker2
    DOCKER_VERSION=`yum list docker | tail -1 | awk '{print $2}' | head -c 2`

    if [ $DOCKER_VERSION -eq 17 ]; then
      DOCKER_PKG_VERSION='17.09.1ce-1.111.amzn1'
      NVIDIA_DOCKER_PKG_VERSION='2.0.3-1.docker17.09.1.ce.amzn1'
    else
      DOCKER_PKG_VERSION='18.06.1ce-3.17.amzn1'
      NVIDIA_DOCKER_PKG_VERSION='2.0.3-1.docker18.06.1.ce.amzn1'
    fi

    sudo yum -y remove docker
    sudo yum -y install docker-$DOCKER_PKG_VERSION

    sudo /etc/init.d/docker start

    curl -s -L https://nvidia.github.io/nvidia-docker/amzn1/nvidia-docker.repo | sudo tee /etc/yum.repos.d/nvidia-docker.repo
    sudo yum install -y nvidia-docker2-$NVIDIA_DOCKER_PKG_VERSION
    sudo cp daemon.json /etc/docker/daemon.json
    sudo pkill -SIGHUP dockerd
    echo "installed nvidia-docker2"
  else
    echo "nvidia-docker2 already installed. We are good to go!"
  fi
fi

# This is common for both GPU and CPU instances

# check if we have docker-compose
docker-compose version >/dev/null 2>&1
if [ $? -ne 0 ]; then
  # install docker compose
  pip install docker-compose
fi

# check if we need to configure our docker interface
SAGEMAKER_NETWORK=`docker network ls | grep -c sagemaker-local`
if [ $SAGEMAKER_NETWORK -eq 0 ]; then
  docker network create --driver bridge sagemaker-local
fi

# Notebook instance Docker networking fixes
RUNNING_ON_NOTEBOOK_INSTANCE=`sudo iptables -S OUTPUT -t nat | grep -c 169.254.0.2`

# Get the Docker Network CIDR and IP for the sagemaker-local docker interface.
SAGEMAKER_INTERFACE=br-`docker network ls | grep sagemaker-local | cut -d' ' -f1`
DOCKER_NET=`ip route | grep $SAGEMAKER_INTERFACE | cut -d" " -f1`
DOCKER_IP=`ip route | grep $SAGEMAKER_INTERFACE | cut -d" " -f12`

# check if both IPTables and the Route Table are OK.
IPTABLES_PATCHED=`sudo iptables -S PREROUTING -t nat | grep -c $SAGEMAKER_INTERFACE`
ROUTE_TABLE_PATCHED=`sudo ip route show table agent | grep -c $SAGEMAKER_INTERFACE`

if [ $RUNNING_ON_NOTEBOOK_INSTANCE -gt 0 ]; then

  if [ $ROUTE_TABLE_PATCHED -eq 0 ]; then
    # fix routing
    sudo ip route add $DOCKER_NET via $DOCKER_IP dev $SAGEMAKER_INTERFACE table agent
  else
    echo "SageMaker instance route table setup is ok. We are good to go."
  fi

  if [ $IPTABLES_PATCHED -eq 0 ]; then
    sudo iptables -t nat -A PREROUTING  -i $SAGEMAKER_INTERFACE -d 169.254.169.254/32 -p tcp -m tcp --dport 80 -j DNAT --to-destination 169.254.0.2:9081
    echo "iptables for Docker setup done"
  else
    echo "SageMaker instance routing for Docker is ok. We are good to go!"
  fi
fi


#!/bin/bash

sudo -n true
if [ $? -eq 0 ]; then
  echo "The user has root access."
else
  echo "The user does not have root access. Everything required to run the notebook is already installed and setup. We are good to go!"
  exit 0
fi

# Do we have GPU support?
nvidia-smi > /dev/null 2>&1
if [ $? -eq 0 ]; then
  # check if we have nvidia-docker
  NVIDIA_DOCKER=`rpm -qa | grep -c nvidia-docker2`
  if [ $NVIDIA_DOCKER -eq 0 ]; then
    # Install nvidia-docker2
    DOCKER_VERSION=`yum list docker | tail -1 | awk '{print $2}' | head -c 2`

    if [ $DOCKER_VERSION -eq 17 ]; then
      DOCKER_PKG_VERSION='17.09.1ce-1.111.amzn1'
      NVIDIA_DOCKER_PKG_VERSION='2.0.3-1.docker17.09.1.ce.amzn1'
    else
      DOCKER_PKG_VERSION='18.06.1ce-3.17.amzn1'
      NVIDIA_DOCKER_PKG_VERSION='2.0.3-1.docker18.06.1.ce.amzn1'
    fi

    sudo yum -y remove docker
    sudo yum -y install docker-$DOCKER_PKG_VERSION

    sudo /etc/init.d/docker start

    curl -s -L https://nvidia.github.io/nvidia-docker/amzn1/nvidia-docker.repo | sudo tee /etc/yum.repos.d/nvidia-docker.repo
    sudo yum install -y nvidia-docker2-$NVIDIA_DOCKER_PKG_VERSION
    sudo cp daemon.json /etc/docker/daemon.json
    sudo pkill -SIGHUP dockerd
    echo "installed nvidia-docker2"
  else
    echo "nvidia-docker2 already installed. We are good to go!"
  fi
fi

# This is common for both GPU and CPU instances

# check if we have docker-compose
docker-compose version >/dev/null 2>&1
if [ $? -ne 0 ]; then
  # install docker compose
  pip install docker-compose
fi

# check if we need to configure our docker interface
SAGEMAKER_NETWORK=`docker network ls | grep -c sagemaker-local`
if [ $SAGEMAKER_NETWORK -eq 0 ]; then
  docker network create --driver bridge sagemaker-local
fi

# Notebook instance Docker networking fixes
RUNNING_ON_NOTEBOOK_INSTANCE=`sudo iptables -S OUTPUT -t nat | grep -c 169.254.0.2`

# Get the Docker Network CIDR and IP for the sagemaker-local docker interface.
SAGEMAKER_INTERFACE=br-`docker network ls | grep sagemaker-local | cut -d' ' -f1`
DOCKER_NET=`ip route | grep $SAGEMAKER_INTERFACE | cut -d" " -f1`
DOCKER_IP=`ip route | grep $SAGEMAKER_INTERFACE | cut -d" " -f12`

# check if both IPTables and the Route Table are OK.
IPTABLES_PATCHED=`sudo iptables -S PREROUTING -t nat | grep -c $SAGEMAKER_INTERFACE`
ROUTE_TABLE_PATCHED=`sudo ip route show table agent | grep -c $SAGEMAKER_INTERFACE`

if [ $RUNNING_ON_NOTEBOOK_INSTANCE -gt 0 ]; then

  if [ $ROUTE_TABLE_PATCHED -eq 0 ]; then
    # fix routing
    sudo ip route add $DOCKER_NET via $DOCKER_IP dev $SAGEMAKER_INTERFACE table agent
  else
    echo "SageMaker instance route table setup is ok. We are good to go."
  fi

  if [ $IPTABLES_PATCHED -eq 0 ]; then
    sudo iptables -t nat -A PREROUTING  -i $SAGEMAKER_INTERFACE -d 169.254.169.254/32 -p tcp -m tcp --dport 80 -j DNAT --to-destination 169.254.0.2:9081
    echo "iptables for Docker setup done"
  else
    echo "SageMaker instance routing for Docker is ok. We are good to go!"
  fi
fi


#!/bin/bash
python -m pip install transformers==4.9.1
python pt_exporter.py


#!/bin/bash
python -m pip install transformers==4.9.1
python onnx_exporter.py
trtexec --onnx=model.onnx --saveEngine=model_bs16.plan --minShapes=token_ids:1x512,attn_mask:1x512 --optShapes=token_ids:16x512,attn_mask:16x512 --maxShapes=token_ids:128x512,attn_mask:128x512 --fp16 --verbose --workspace=14000 | tee conversion_bs16_dy.txt


#!/bin/bash

conda create -y -n sd_env python=3.8
source ~/anaconda3/etc/profile.d/conda.sh
source activate sd_env
export PYTHONNOUSERSITE=True
pip install torch
pip install transformers ftfy scipy accelerate
pip install diffusers==0.9.0
pip install transformers[onnxruntime]
pip install conda-pack
conda-pack


#!/bin/bash

pip install transformers ftfy scipy
pip install transformers[onnxruntime]
pip install diffusers
cd /mount
python export.py

# Accelerating VAE with TensorRT
trtexec --onnx=vae.onnx --saveEngine=vae.plan --minShapes=latent_sample:1x4x64x64 --optShapes=latent_sample:4x4x64x64 --maxShapes=latent_sample:8x4x64x64 --fp16


#!/bin/bash

conda create -y -n processing_env python=3.8
source ~/anaconda3/etc/profile.d/conda.sh
export PYTHONNOUSERSITE=True
conda activate processing_env && conda install -y transformers numpy conda-pack && conda install -y pytorch torchvision torchaudio -c pytorch
conda-pack -n processing_env


#!/bin/bash
MODEL_NAME=$1
python -m pip install transformers==4.9.1
python onnx_exporter.py --model $MODEL_NAME

trtexec \
    --onnx=model.onnx \
    --saveEngine=model.plan \
    --minShapes=token_ids:1x128,attn_mask:1x128 \
    --optShapes=token_ids:16x128,attn_mask:16x128 \
    --maxShapes=token_ids:32x128,attn_mask:32x128 \
    --verbose \
    --workspace=14000 \
| tee conversion.txt


#!/bin/bash

conda create -y -n processing_env python=3.8
source ~/anaconda3/etc/profile.d/conda.sh
source activate processing_env
export PYTHONNOUSERSITE=True
pip install torch
pip install transformers
pip install conda-pack
conda-pack


#!/bin/sh

# The name of our algorithm
repo=$1
registry_name=$2
image_tag=$3
docker_file=$4
platforms=$5

echo "[INFO]: registry_name=${registry_name}"
echo "[INFO]: image_tag=${image_tag}"
echo "[INFO]: docker_file=${docker_file}"
echo "[INFO]: platforms=${platforms}"

cd $repo

account=$(aws sts get-caller-identity --query Account --output text)

# Get the region defined in the current configuration (default to us-west-2 if none defined)
region=$(aws configure get region)

echo "[INFO]: Region ${region}"

fullname="${account}.dkr.ecr.${region}.amazonaws.com/${registry_name}:${image_tag}"

echo "[INFO]: Image name: ${fullname}"

# If the repository doesn't exist in ECR, create it.

aws ecr describe-repositories --repository-names "${registry_name}" > /dev/null 2>&1

aws ecr create-repository --repository-name "${registry_name}" > /dev/null

## If you are extending Amazon SageMaker Images, you need to login to the account
# Get the login command from ECR and execute it directly
password=$(aws ecr --region ${region} get-login-password)

docker login -u AWS -p ${password} "${account}.dkr.ecr.${region}.amazonaws.com"

if [ -z ${platforms} ]
then
  docker build -t ${fullname} -f ${docker_file} .
else
  echo "Provided platform = ${platforms}"
  docker build -t ${fullname} -f ${docker_file} . --platform=${platforms}
fi

docker push ${fullname}


#!/bin/sh

IMAGE_NAME=$1
REGISTRY_NAME=$2
SM_IMAGE_NAME=$3
ROLE_ARN=$4

if [ -z ${IMAGE_NAME} ]
then
  echo "[INFO]: IMAGE_NAME not passed"
  exit 1
fi

if [ -z ${REGISTRY_NAME} ]
then
  echo "[INFO]: REGISTRY_NAME not passed"
  exit 1
fi

if [ -z ${SM_IMAGE_NAME} ]
then
  echo "[INFO]: SM_IMAGE_NAME not passed"
  exit 1
fi

if [ -z ${ROLE_ARN} ]
then
  echo "[INFO]: ROLE_ARN not passed"
  exit 1
fi

echo "[INFO]: IMAGE_NAME=${IMAGE_NAME}"
echo "[INFO]: REGISTRY_NAME=${REGISTRY_NAME}"
echo "[INFO]: SM_IMAGE_NAME=${SM_IMAGE_NAME}"
echo "[INFO]: ROLE_ARN=${ROLE_ARN}"

aws sagemaker create-image \
    --image-name ${SM_IMAGE_NAME} \
    --role-arn ${ROLE_ARN} \
    || exit 1

account=$(aws sts get-caller-identity --query Account --output text)
region=$(aws configure get region)

aws sagemaker create-image-version \
    --image-name ${SM_IMAGE_NAME} \
    --base-image "${account}.dkr.ecr.${region}.amazonaws.com/${REGISTRY_NAME}:latest" \
    || exit 1

aws sagemaker delete-app-image-config --app-image-config-name ${SM_IMAGE_NAME}-config

aws sagemaker describe-image-version --image-name ${SM_IMAGE_NAME}

aws sagemaker create-app-image-config --cli-input-json file://${IMAGE_NAME}/app-image-config.json


#!/bin/sh

echo "studio-domain-config.json"

aws sagemaker update-domain --cli-input-json file://studio-domain-config.json


#!/bin/sh

IMAGE_NAME=$1
REGISTRY_NAME=$2
SM_IMAGE_NAME=$3
ROLE_ARN=$4

if [ -z ${IMAGE_NAME} ]
then
  echo "[INFO]: IMAGE_NAME not passed"
  exit 1
fi

if [ -z ${REGISTRY_NAME} ]
then
  echo "[INFO]: REGISTRY_NAME not passed"
  exit 1
fi

if [ -z ${SM_IMAGE_NAME} ]
then
  echo "[INFO]: SM_IMAGE_NAME not passed"
  exit 1
fi

if [ -z ${ROLE_ARN} ]
then
  echo "[INFO]: ROLE_ARN not passed"
  exit 1
fi

echo "[INFO]: IMAGE_NAME=${IMAGE_NAME}"
echo "[INFO]: REGISTRY_NAME=${REGISTRY_NAME}"
echo "[INFO]: SM_IMAGE_NAME=${SM_IMAGE_NAME}"
echo "[INFO]: ROLE_ARN=${ROLE_ARN}"

account=$(aws sts get-caller-identity --query Account --output text)
region=$(aws configure get region)

aws sagemaker create-image-version \
    --image-name ${SM_IMAGE_NAME} \
    --base-image "${account}.dkr.ecr.${region}.amazonaws.com/${REGISTRY_NAME}:latest" \
    || exit 1

aws sagemaker delete-app-image-config --app-image-config-name ${SM_IMAGE_NAME}-config

aws sagemaker describe-image-version --image-name ${SM_IMAGE_NAME}

aws sagemaker create-app-image-config --cli-input-json file://${IMAGE_NAME}/app-image-config.json


#!/bin/bash

conda create -y -n preprocessing_env python=3.8
source ~/anaconda3/etc/profile.d/conda.sh
conda activate preprocessing_env
export PYTHONNOUSERSITE=True
conda install -y -c conda-forge pandas scikit-learn
pip install conda-pack
conda-pack


#!/bin/bash

set -e

sudo -u ec2-user -i <<'EOF'

mkdir -p rapids_kernel
cd rapids_kernel
wget -q https://rapidsai-data.s3.us-east-2.amazonaws.com/conda-pack/rapidsai/rapids21.06_cuda11.0_py3.8.tar.gz
echo "wget completed"
tar -xzf *.gz
echo "unzip completed"
source /home/ec2-user/rapids_kernel/bin/activate
conda-unpack
echo "unpack completed"
python -m ipykernel install --user --name rapids-2106
echo "kernel install completed"

EOF


#!/bin/bash
python -m pip install transformers==4.9.1
python onnx_exporter.py
python pt_exporter.py
trtexec --onnx=model.onnx --saveEngine=model_bs16.plan --minShapes=token_ids:1x128,attn_mask:1x128 --optShapes=token_ids:16x128,attn_mask:16x128 --maxShapes=token_ids:128x128,attn_mask:128x128 --fp16 --verbose --workspace=14000 | tee conversion_bs16_dy.txt


#!/bin/bash
python onnx_exporter.py
python pt_exporter.py
trtexec --onnx=model.onnx --saveEngine=model.plan --explicitBatch --minShapes=input:1x3x224x224 --optShapes=input:128x3x224x224 --maxShapes=input:128x3x224x224 --fp16 --verbose | tee conversion.txt


#!/bin/bash
set -e

if [[ "$1" = "serve" ]]; then
    shift 1
    printenv
    ls /opt
    torchserve --start --ts-config /home/model-server/config.properties --model-store /opt/ml/model --models all

else
    eval "$@"
fi

# prevent docker exit
tail -f /dev/null


#! /usr/bin/env bash
set -e

# Get the listen port from the SM env variable, otherwise default to 8080
LISTEN_PORT=${SAGEMAKER_BIND_TO_PORT:-8080}

# Set the number of gunicorn worker processes
GUNICORN_WORKER_COUNT=$(nproc)

# Start flask app
exec gunicorn -w $GUNICORN_WORKER_COUNT -b 0.0.0.0:$LISTEN_PORT main:app


#!/usr/bin/env bash
# This script shows how to build the Docker image and push it to ECR to be ready for use
# by SageMaker.
# The argument to this script is the image name. This will be used as the image on the local
# machine and combined with the account and region to form the repository name for ECR.
# set region

DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"

if [ "$#" -eq 4 ]; then
    dlc_account_id=$1
    region=$2
    image=$3
    tag=$4
else
    echo "usage: $0 <dlc-account-id> $1 <aws-region> $2 <image-repo> $3 <image-tag>"
    exit 1
fi

# Get the account number associated with the current IAM credentials
account=$(aws sts get-caller-identity --query Account --output text)

if [ $? -ne 0 ]
then
    exit 255
fi

fullname="${account}.dkr.ecr.${region}.amazonaws.com/${image}:${tag}"

# If the repository doesn't exist in ECR, create it.
aws ecr describe-repositories --region ${region} --repository-names "${image}" > /dev/null 2>&1
if [ $? -ne 0 ]; then
    echo "creating ECR repository : ${fullname} "
    aws ecr create-repository --region ${region} --repository-name "${image}" > /dev/null
fi

aws ecr get-login-password --region ${region} \
| docker login --username AWS --password-stdin ${account}.dkr.ecr.${region}.amazonaws.com

docker build . -t ${image} -f ${DIR}/Dockerfile  --build-arg dlc_account_id=${dlc_account_id} --build-arg region=${region}
docker tag ${image} ${fullname}
docker push ${fullname}

if [ $? -eq 0 ]; then
	echo "Amazon ECR URI: ${fullname}"
else
	echo "Error: Image build and push failed"
	exit 1
fi


export HDF5_USE_FILE_LOCKING=FALSE

python ./DeepLearningExamples/PyTorch/LanguageModeling/BERT/run_pretraining.py \
--input_dir=$SM_CHANNEL_TRAINING \
--output_dir=$SM_MODEL_DIR \
--config_file=./DeepLearningExamples/PyTorch/LanguageModeling/BERT/bert_config.json \
--bert_model=bert-large-uncased \
--train_batch_size=64 \
--max_seq_length=128 \
--max_predictions_per_seq=20 \
--max_steps=900864 \
--warmup_proportion=0.2843 \
--num_steps_per_checkpoint=900864 \
--learning_rate=1.12e-3 \
--seed=42 \
--fp16 \
--do_train \
--json-summary ./DeepLearningExamples/PyTorch/LanguageModeling/BERT/results/dllogger.json 2>&1


#!/usr/bin/env bash
# This script shows how to build the Docker image and push it to ECR to be ready for use
# by SageMaker.
# The argument to this script is the image name. This will be used as the image on the local
# machine and combined with the account and region to form the repository name for ECR.
# set region

DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"

# Change the active directory to the one that contains the training script
cd ${DIR}/SMDDP-Examples/pytorch/image_classification

if [ "$#" -eq 4 ]; then
    dlc_account_id=$1
    region=$2
    image=$3
    tag=$4
else
    echo "usage: $0 <dlc-account-id> $1 <aws-region> $2 <image-repo> $3 <image-tag>"
    exit 1
fi

# Get the account number associated with the current IAM credentials
account=$(aws sts get-caller-identity --query Account --output text)

if [ $? -ne 0 ]
then
    exit 255
fi

fullname="${account}.dkr.ecr.${region}.amazonaws.com/${image}:${tag}"

# If the repository doesn't exist in ECR, create it.
aws ecr describe-repositories --region ${region} --repository-names "${image}" > /dev/null 2>&1
if [ $? -ne 0 ]; then
    echo "creating ECR repository : ${fullname} "
    aws ecr create-repository --region ${region} --repository-name "${image}" > /dev/null
fi

aws ecr get-login-password --region ${region} \
| docker login --username AWS --password-stdin ${account}.dkr.ecr.${region}.amazonaws.com
docker build . -t ${image} -f ${DIR}/Dockerfile  --build-arg dlc_account_id=${dlc_account_id} --build-arg region=${region}
docker tag ${image} ${fullname}

aws ecr get-login-password --region ${region} \
| docker login --username AWS --password-stdin ${account}.dkr.ecr.${region}.amazonaws.com
docker push ${fullname}
if [ $? -eq 0 ]; then
	echo "Amazon ECR URI: ${fullname}"
else
	echo "Error: Image build and push failed"
	exit 1
fi


#!/usr/bin/env bash
# This script shows how to tag the Docker image and push it to ECR to be ready for use
# by SageMaker.
# The argument to this script is the image name. This will be used as the image on the local
# machine and combined with the account and region to form the repository name for ECR.
# set region

if [ "$#" -eq 3 ]; then
    region=$1
    image=$2
    tag=$3
else
    echo "usage: $0 <aws-region> $1 <image-repo> $2 <image-tag>"
    exit 1
fi

# Get the account number associated with the current IAM credentials
account=$(aws sts get-caller-identity --query Account --output text)

if [ $? -ne 0 ]
then
    exit 255
fi

fullname="${account}.dkr.ecr.${region}.amazonaws.com/${image}:${tag}"

# If the repository doesn't exist in ECR, create it.
aws ecr describe-repositories --region ${region} --repository-names "${image}" > /dev/null 2>&1
if [ $? -ne 0 ]; then
    echo "creating ECR repository : ${fullname} "
    aws ecr create-repository --region ${region} --repository-name "${image}" > /dev/null
fi

$(aws ecr get-login --no-include-email --region ${region}  --registry-ids 763104351884)
docker tag mlperf/rnn_speech_recognition_smddp:latest ${fullname}

# Get the login command from ECR and execute it directly
$(aws ecr get-login --region ${region} --no-include-email)
docker push ${fullname}
if [ $? -eq 0 ];
then
	echo "Amazon ECR URI: ${fullname}"
else
	echo "Error: Image build and push failed"
	exit 1
fi


#!/usr/bin/env bash
# This script shows how to build the Docker image and push it to ECR to be ready for use
# by SageMaker.
# The argument to this script is the image name. This will be used as the image on the local
# machine and combined with the account and region to form the repository name for ECR.
# set region

DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"

if [ "$#" -eq 3 ]; then
    region=$1
    image=$2
    tag=$3
else
    echo "usage: $0 <aws-region> $1 <image-repo> $2 <image-tag>"
    exit 1
fi

# Get the account number associated with the current IAM credentials
account=$(aws sts get-caller-identity --query Account --output text)

if [ $? -ne 0 ]
then
    exit 255
fi

fullname="${account}.dkr.ecr.${region}.amazonaws.com/${image}:${tag}"

# If the repository doesn't exist in ECR, create it.
aws ecr describe-repositories --region ${region} --repository-names "${image}" > /dev/null 2>&1
if [ $? -ne 0 ]; then
    echo "creating ECR repository : ${fullname} "
    aws ecr create-repository --region ${region} --repository-name "${image}" > /dev/null
fi

$(aws ecr get-login --no-include-email --region ${region}  --registry-ids 763104351884)
docker build ${DIR}/ -t ${image} -f ${DIR}/Dockerfile  --build-arg region=${region}
docker tag ${image} ${fullname}

# Get the login command from ECR and execute it directly
$(aws ecr get-login --region ${region} --no-include-email)
docker push ${fullname}
if [ $? -eq 0 ]; then
	echo "Amazon ECR URI: ${fullname}"
else
	echo "Error: Image build and push failed"
	exit 1
fi


#!/usr/bin/env bash
# This script shows how to build the Docker image and push it to ECR to be ready for use
# by SageMaker.
# The argument to this script is the image name. This will be used as the image on the local
# machine and combined with the account and region to form the repository name for ECR.
# set region

DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"

if [ "$#" -eq 3 ]; then
    region=$1
    image=$2
    tag=$3
else
    echo "usage: $0 <aws-region> $1 <image-repo> $2 <image-tag>"
    exit 1
fi

# Get the account number associated with the current IAM credentials
account=$(aws sts get-caller-identity --query Account --output text)

if [ $? -ne 0 ]
then
    exit 255
fi

fullname="${account}.dkr.ecr.${region}.amazonaws.com/${image}:${tag}"

# If the repository doesn't exist in ECR, create it.
aws ecr describe-repositories --region ${region} --repository-names "${image}" > /dev/null 2>&1
if [ $? -ne 0 ]; then
    echo "creating ECR repository : ${fullname} "
    aws ecr create-repository --region ${region} --repository-name "${image}" > /dev/null
fi

$(aws ecr get-login --no-include-email --region us-west-2  --registry-ids 763104351884)
docker build ${DIR}/ -t ${image} -f ${DIR}/Dockerfile  --build-arg region=${region}
docker tag ${image} ${fullname}

# Get the login command from ECR and execute it directly
$(aws ecr get-login --region ${region} --no-include-email)
docker push ${fullname}
if [ $? -eq 0 ]; then
	echo "Amazon ECR URI: ${fullname}"
else
	echo "Error: Image build and push failed"
	exit 1
fi


#!/usr/bin/env bash
# This script shows how to build the Docker image and push it to ECR to be ready for use
# by SageMaker.
# The argument to this script is the image name. This will be used as the image on the local
# machine and combined with the account and region to form the repository name for ECR.
# set region

DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"

if [ "$#" -eq 3 ]; then
    region=$1
    image=$2
    tag=$3
else
    echo "usage: $0 <aws-region> $1 <image-repo> $2 <image-tag>"
    exit 1
fi

# Get the account number associated with the current IAM credentials
account=$(aws sts get-caller-identity --query Account --output text)

if [ $? -ne 0 ]
then
    exit 255
fi

fullname="${account}.dkr.ecr.${region}.amazonaws.com/${image}:${tag}"

# If the repository doesn't exist in ECR, create it.
aws ecr describe-repositories --region ${region} --repository-names "${image}" > /dev/null 2>&1
if [ $? -ne 0 ]; then
    echo "creating ECR repository : ${fullname} "
    aws ecr create-repository --region ${region} --repository-name "${image}" > /dev/null
fi

aws ecr get-login-password --region ${region} \
| docker login --username AWS --password-stdin ${account}.dkr.ecr.${region}.amazonaws.com
docker build ${DIR}/ -t ${image} -f ${DIR}/Dockerfile  --build-arg region=${region}
docker tag ${image} ${fullname}

aws ecr get-login-password --region ${region} \
| docker login --username AWS --password-stdin ${account}.dkr.ecr.${region}.amazonaws.com
docker push ${fullname}
if [ $? -eq 0 ]; then
	echo "Amazon ECR URI: ${fullname}"
else
	echo "Error: Image build and push failed"
	exit 1
fi


#!/bin/bash

set -e

# This function builds, packages and deploys a region-specific operator to an ECR repo and output bucket.
# Parameter:
#    $1: The region of the ECR repo.
#    $2: The stage in the pipeline for the output account. (prod/beta/dev)
# e.g. build_canary us-east-1 prod
function build_canary()
{
  local account_region="$1"
  local stage="$2"

  # Match stage to pipeline type
  if [ "$stage" != "$PIPELINE_STAGE" ]; then
    return
  fi

  $(aws ecr get-login --no-include-email --region $CANARY_ECR_REGION --registry-ids $AWS_ACCOUNT_ID)

  # Download the operator for this canary
  pushd tests
  aws s3 cp s3://$ALPHA_TARBALL_BUCKET/${CODEBUILD_RESOLVED_SOURCE_VERSION}/sagemaker-k8s-operator-${account_region}.tar.gz sagemaker-k8s-operator.tar.gz

  # Build and push the newest canary tests with SHA1 and latest tags
  IMAGE=$AWS_ACCOUNT_ID.dkr.ecr.$CANARY_ECR_REGION.amazonaws.com/$CANARY_ECR_REPOSITORY
  SHA_REGION_TAG=${CODEBUILD_RESOLVED_SOURCE_VERSION}-${account_region}
  REGION_TAG=$account_region

  # Push a canary tagged with the SHA and the region
  IMG=$IMAGE:$SHA_REGION_TAG COMMIT_SHA=$CODEBUILD_RESOLVED_SOURCE_VERSION bash build_canary.sh
  docker push $IMAGE:$SHA_REGION_TAG

  # Tag it with just the region name for latest
  docker tag $IMAGE:$SHA_REGION_TAG $IMAGE:$REGION_TAG
  docker push $IMAGE:$REGION_TAG
  popd
}

# Replace JSON single quotes with double quotes for jq to understand
ACCOUNTS_ESCAPED=`echo $ACCOUNTS | sed "s/'/\"/g"`
for row in $(echo ${ACCOUNTS_ESCAPED} | jq -r '.[] | @base64'); do
  _jq() {
    echo ${row} | base64 --decode | jq -r ${1}
  }

  region="$(_jq '.region')"
  stage="$(_jq '.stage')"
  build_canary "$region" "$stage"
done


#!/bin/bash

# This script will build the integration test container. This container contains
# all the tools necessary for running the build and test steps for each of the
# CodeBuild projects. The script will also tag the container with the latest
# commit SHA, and with the "latest" tag, then push to an ECR repository.

set -x

# Build new integration test container
pushd tests
IMG=$INTEGRATION_CONTAINER_REPOSITORY bash build_integration.sh
popd

# Log into ECR
$(aws ecr get-login --no-include-email --region $REGION --registry-ids $AWS_ACCOUNT_ID)

# Tag the container with SHA and latest
docker tag $INTEGRATION_CONTAINER_REPOSITORY $AWS_ACCOUNT_ID.dkr.ecr.$REGION.amazonaws.com/$INTEGRATION_CONTAINER_REPOSITORY:$CODEBUILD_RESOLVED_SOURCE_VERSION
docker tag $INTEGRATION_CONTAINER_REPOSITORY $AWS_ACCOUNT_ID.dkr.ecr.$REGION.amazonaws.com/$INTEGRATION_CONTAINER_REPOSITORY:latest

# Push the newly tagged containers
docker push $AWS_ACCOUNT_ID.dkr.ecr.$REGION.amazonaws.com/$INTEGRATION_CONTAINER_REPOSITORY:$CODEBUILD_RESOLVED_SOURCE_VERSION
docker push $AWS_ACCOUNT_ID.dkr.ecr.$REGION.amazonaws.com/$INTEGRATION_CONTAINER_REPOSITORY:latest


RELEASE_BUCKET_NAME_FMT="%s-%s"

RELEASE_BINARY_PREFIX_FMT="s3://%s/kubectl-smlogs-plugin"
ALPHA_BINARY_PREFIX_FMT="s3://%s/%s"

ALPHA_LINUX_BINARY_PATH_FMT="%s/kubectl-smlogs-plugin.linux.amd64.tar.gz"
ALPHA_DARWIN_BINARY_PATH_FMT="%s/kubectl-smlogs-plugin.darwin.amd64.tar.gz"

RELEASE_LINUX_BINARY_PATH_FMT="%s/%s/linux.amd64.tar.gz"
RELEASE_DARWIN_BINARY_PATH_FMT="%s/%s/darwin.amd64.tar.gz"

PUBLIC_CP_ARGS="--acl public-read"


#!/bin/bash

BUNDLED_RELEASE_PATH="release/"
TMP_RELEASE_PATH="/tmp/release"

mkdir -p "${TMP_RELEASE_PATH}"

# Create installers from CRDs into temporary directory
INSTALLER_PATH="${TMP_RELEASE_PATH}" make create-installers

# Diff the two folders, excluding hidden files (such as .gitkeep)
diff -r --exclude=".*" "${TMP_RELEASE_PATH}" "${BUNDLED_RELEASE_PATH}"

if [ $? -ne 0 ]; then
  echo "Release files did not match. See diff above for more details"
  exit 1
fi


#!/bin/bash

source codebuild/scripts/package_operators.sh

set -x

# Login to alpha ECR
$(aws ecr get-login --no-include-email --region $ALPHA_REPOSITORY_REGION --registry-ids $ALPHA_ACCOUNT_ID)

# Build the image with a temporary tag
ALPHA_IMAGE=$ALPHA_ACCOUNT_ID.dkr.ecr.$ALPHA_REPOSITORY_REGION.amazonaws.com/$REPOSITORY_NAME
make docker-build docker-push IMG=$ALPHA_IMAGE:$CODEBUILD_RESOLVED_SOURCE_VERSION

# Ensure the smlogs are built and tarred into the alpha bucket
SHOULD_PUSH_SMLOGS_ALPHA="true"

# Release the operator into the private alpha repository
# Set as all to ensure it runs through the function
# Add the alpha prefix for integration testing
package_operator "$ALPHA_ACCOUNT_ID" "$ALPHA_REPOSITORY_REGION" "$REPOSITORY_NAME" "all" "-alpha"


#!/bin/bash

source codebuild/scripts/deployment_constants.sh

set -e

# Define alpha artifact locations
printf -v ALPHA_BUCKET_PREFIX $ALPHA_BINARY_PREFIX_FMT $ALPHA_TARBALL_BUCKET $CODEBUILD_RESOLVED_SOURCE_VERSION

printf -v ALPHA_LINUX_BINARY_PATH $ALPHA_LINUX_BINARY_PATH_FMT $ALPHA_BUCKET_PREFIX
printf -v ALPHA_DARWIN_BINARY_PATH $ALPHA_DARWIN_BINARY_PATH_FMT $ALPHA_BUCKET_PREFIX

# This function deploys a region-specific operator to an ECR prod repo from the existing
# image in the alpha repository. The function also copies across the smlogs binaries
# from the alpha tarball bucket into the production buckets.
# Parameter:
#    $1: The account ID for the ECR repo.
#    $2: The region of the ECR repo.
#    $3: The name of the ECR repository.
function deploy_from_alpha()
{
  local account_id="$1"
  local account_region="$2"
  local image_repository="$3"

  # Get the images from the alpha ECR repository
  local dest_ecr_image=${account_id}.dkr.ecr.${account_region}.amazonaws.com/${image_repository}
  local alpha_ecr_image=$ALPHA_ACCOUNT_ID.dkr.ecr.$ALPHA_REPOSITORY_REGION.amazonaws.com/${image_repository}

  # Login to the alpha repository
  $(aws ecr get-login --no-include-email --region $ALPHA_REPOSITORY_REGION --registry-ids $ALPHA_ACCOUNT_ID)
  docker pull $alpha_ecr_image:$CODEBUILD_RESOLVED_SOURCE_VERSION

  # Login to the prod repository
  $(aws ecr get-login --no-include-email --region ${account_region} --registry-ids ${account_id})

  # Clone the controller image to the repo
  docker tag $alpha_ecr_image:$CODEBUILD_RESOLVED_SOURCE_VERSION ${dest_ecr_image}:$CODEBUILD_RESOLVED_SOURCE_VERSION
  # TODO: Remove at the time of the next major version release
  # This moves the `latest` tag along with master branch, which will cause backwards incompatibility for any major
  # version updates.
  docker tag $alpha_ecr_image:$CODEBUILD_RESOLVED_SOURCE_VERSION ${dest_ecr_image}:latest

  # Push to the prod region
  docker push ${dest_ecr_image}:$CODEBUILD_RESOLVED_SOURCE_VERSION

  printf -v bucket_name $RELEASE_BUCKET_NAME_FMT $RELEASE_TARBALL_BUCKET_PREFIX $account_region
  printf -v binary_prefix $RELEASE_BINARY_PREFIX_FMT $bucket_name

  aws s3 cp "$ALPHA_LINUX_BINARY_PATH" "$(printf $RELEASE_LINUX_BINARY_PATH_FMT $binary_prefix $CODEBUILD_RESOLVED_SOURCE_VERSION)" $PUBLIC_CP_ARGS
  aws s3 cp "$ALPHA_DARWIN_BINARY_PATH" "$(printf $RELEASE_DARWIN_BINARY_PATH_FMT $binary_prefix $CODEBUILD_RESOLVED_SOURCE_VERSION)" $PUBLIC_CP_ARGS
}

# This function builds, packages and deploys a region-specific operator to an ECR repo and output bucket.
# Parameter:
#    $1: The account ID for the ECR repo.
#    $2: The region of the ECR repo.
#    $3: The name of the ECR repository.
#    $4: The stage in the pipeline for the output account. (prod/beta/dev)
#    $5: (Optional) A suffix for the operator install bundle tarball.
# e.g. package_operator 123456790 us-east-1 amazon-sagemaker-k8s-operator prod
function package_operator()
{
  local account_id="$1"
  local account_region="$2"
  local image_repository="$3"
  local stage="$4"
  local tarball_suffix="${5:-}"

  # Only build images that match the release pipeline stage
  if [ "$stage" != "$PIPELINE_STAGE" ] && [ "$stage" != "all" ]; then
    return 0
  fi

  # Only push to ECR repos if this is run on the prod pipeline
  if [ "$stage" == "prod" ] && [ "$PIPELINE_STAGE" == "prod" ]; then
    deploy_from_alpha "$account_id" "$account_region" "$image_repository"
  fi

  # Build, push and update the CRD with controller image and current git SHA, create the tarball and extract it to pack
  local ecr_image=${account_id}.dkr.ecr.${account_region}.amazonaws.com/${image_repository}
  make modify-base-kustomize-us IMG=${ecr_image}:$CODEBUILD_RESOLVED_SOURCE_VERSION
  make build-release-tarball
  pushd bin
    tar -xf sagemaker-k8s-operator-install-scripts.tar.gz
  popd

  # Create the smlog binary
  pushd smlogs-kubectl-plugin
    make build-release
  popd

  # Create a temporary dir and put all the necessary artifacts
  rm -rf /tmp/sagemaker-k8s-operator
  mkdir -p /tmp/sagemaker-k8s-operator
  mkdir -p /tmp/sagemaker-k8s-operator/smlogs-plugin/darwin.amd64
  mkdir -p /tmp/sagemaker-k8s-operator/smlogs-plugin/linux.amd64

  cp -r bin/sagemaker-k8s-operator-install-scripts /tmp/sagemaker-k8s-operator
  cp smlogs-kubectl-plugin/bin/kubectl-smlogs.linux.amd64 /tmp/sagemaker-k8s-operator/smlogs-plugin/linux.amd64/kubectl-smlogs
  cp smlogs-kubectl-plugin/bin/kubectl-smlogs.darwin.amd64 /tmp/sagemaker-k8s-operator/smlogs-plugin/darwin.amd64/kubectl-smlogs

  if [ "$SHOULD_PUSH_SMLOGS_ALPHA" == "true" ]; then
    # Create temp dirs per binary and put smlogs into it
    mkdir -p /tmp/kubectl-smlogs.linux.amd64
    mkdir -p /tmp/kubectl-smlogs.darwin.amd64

    cp smlogs-kubectl-plugin/bin/kubectl-smlogs.linux.amd64 /tmp/kubectl-smlogs.linux.amd64/kubectl-smlogs
    cp smlogs-kubectl-plugin/bin/kubectl-smlogs.darwin.amd64 /tmp/kubectl-smlogs.darwin.amd64/kubectl-smlogs

    pushd /tmp
      tar cvzf kubectl-smlogs-plugin.linux.amd64.tar.gz kubectl-smlogs.linux.amd64
      tar cvzf kubectl-smlogs-plugin.darwin.amd64.tar.gz kubectl-smlogs.darwin.amd64

      aws s3 cp kubectl-smlogs-plugin.linux.amd64.tar.gz "$ALPHA_LINUX_BINARY_PATH"
      aws s3 cp kubectl-smlogs-plugin.darwin.amd64.tar.gz "$ALPHA_DARWIN_BINARY_PATH"
    popd
  fi

  # Create a tar ball which has CRDs, smlog and sm spec generator binaries
  pushd /tmp
    tar cvzf sagemaker-k8s-operator.tar.gz sagemaker-k8s-operator

    # Upload the final tar ball to s3 with standard name and git SHA
    aws s3 cp sagemaker-k8s-operator.tar.gz "$ALPHA_BUCKET_PREFIX/sagemaker-k8s-operator-${account_region}${tarball_suffix}.tar.gz"
  popd
}


#!/bin/bash

source codebuild/scripts/package_operators.sh

# Replace JSON single quotes with double quotes for jq to understand
ACCOUNTS_ESCAPED=`echo $ACCOUNTS | sed "s/'/\"/g"`
for row in $(echo ${ACCOUNTS_ESCAPED} | jq -r '.[] | @base64'); do
  _jq() {
    echo ${row} | base64 --decode | jq -r ${1}
  }

  repository_account="$(_jq '.repositoryAccount')"
  region="$(_jq '.region')"
  image_repository="${REPOSITORY_NAME}"
  stage="$(_jq '.stage')"

  package_operator "$repository_account" "$region" "$image_repository" "$stage"
done


#!/bin/bash

source codebuild/scripts/deployment_constants.sh

set -e

# Define alpha artifact locations
printf -v ALPHA_BUCKET_PREFIX $ALPHA_BINARY_PREFIX_FMT $ALPHA_TARBALL_BUCKET $CODEBUILD_RESOLVED_SOURCE_VERSION

printf -v ALPHA_LINUX_BINARY_PATH $ALPHA_LINUX_BINARY_PATH_FMT $ALPHA_BUCKET_PREFIX
printf -v ALPHA_DARWIN_BINARY_PATH $ALPHA_DARWIN_BINARY_PATH_FMT $ALPHA_BUCKET_PREFIX

# This function will pull an existing image + tag and push it with a new tag.
# Parameter:
#    $1: The repository and image to pull from.
#    $2: The previous tag for that image.
#    $3: The new tag to push for that image.
function retag_image()
{
  local image="$1"
  local old_tag="$2"
  local new_tag="$3"

  docker pull $image:$old_tag
  docker tag $image:$old_tag $image:$new_tag
  docker push $image:$new_tag
}

# This function will push artifacts to their own folder with a given tag.
# Parameter:
#    $1: The new tag to push for the artifacts.
#    $2: The region of the new artifacts.
function retag_binaries()
{
  local new_tag="$1"
  local region="$2"

  printf -v release_bucket $RELEASE_BUCKET_NAME_FMT $RELEASE_TARBALL_BUCKET_PREFIX $region
  printf -v binary_prefix $RELEASE_BINARY_PREFIX_FMT $release_bucket

  aws s3 cp "$ALPHA_LINUX_BINARY_PATH" "$(printf $RELEASE_LINUX_BINARY_PATH_FMT $binary_prefix $new_tag)" $PUBLIC_CP_ARGS
  aws s3 cp "$ALPHA_DARWIN_BINARY_PATH" "$(printf $RELEASE_DARWIN_BINARY_PATH_FMT $binary_prefix $new_tag)" $PUBLIC_CP_ARGS
}

GIT_TAG="$(git describe --tags --exact-match 2>/dev/null)"

# Only run the release process for tagged commits
if [ "$GIT_TAG" == "" ]; then
  exit 0
fi

# Replace JSON single quotes with double quotes for jq to understand
ACCOUNTS_ESCAPED=`echo $ACCOUNTS | sed "s/'/\"/g"`
for row in $(echo ${ACCOUNTS_ESCAPED} | jq -r '.[] | @base64'); do
  _jq() {
    echo ${row} | base64 --decode | jq -r ${1}
  }

  repository_account="$(_jq '.repositoryAccount')"
  region="$(_jq '.region')"
  image_repository="${REPOSITORY_NAME}"
  stage="$(_jq '.stage')"

  if [ "$stage" != "$PIPELINE_STAGE" ] && [ "$stage" != "all" ]; then
    continue
  fi

  $(aws ecr get-login --no-include-email --region ${region} --registry-ids ${repository_account})

  image=${repository_account}.dkr.ecr.${region}.amazonaws.com/${image_repository}
  old_tag="${CODEBUILD_RESOLVED_SOURCE_VERSION}"
  full_tag="${GIT_TAG}"

  # Get minor and major version tags
  [[ $GIT_TAG =~ ^v[0-9]+\.[0-9]+ ]] && minor_tag="${BASH_REMATCH[0]}"
  [[ $GIT_TAG =~ ^v[0-9]+ ]] && major_tag="${BASH_REMATCH[0]}"

  echo "Tagging $region with $full_tag"

  retag_image "$image" "$old_tag" "$full_tag"
  retag_image "$image" "$old_tag" "$minor_tag"
  retag_image "$image" "$old_tag" "$major_tag"

  retag_binaries "$full_tag" "$region"
  retag_binaries "$minor_tag" "$region"
  retag_binaries "$major_tag" "$region"

  echo "Finished tagging $region with $full_tag"
done


#!/bin/bash

docker build -f images/Dockerfile.canary . -t ${IMG:-canary-test-container} --build-arg DATA_BUCKET --build-arg COMMIT_SHA --build-arg RESULT_BUCKET


#!/bin/bash

docker build -f images/Dockerfile.integration -t ${IMG:-integration-test-container} .


#!/bin/bash

default_operator_namespace="sagemaker-k8s-operator-system"

# Create a resource (run a test) given a specification within a given namespace.
# Parameter:
#    $1: Target namespace
#    $2: Filename of the test
function run_test()
{
  local target_namespace="$1"
  local file_name="$2"

  kubectl apply -n "$target_namespace" -f "$file_name"
}

# This function gets the sagemaker model name from k8s model
# Parameter:
#    $1: Namespace of the k8s model
#    $2: Name of k8s model
# e.g. get_sagemaker_model_from_k8s_model default k8s_model
function get_sagemaker_model_from_k8s_model()
{
  local model_namespace="$1"
  local k8s_model_name="$2"
  # Get the second line and 3rd column, since valid output will be following
  # NAME                    STATUS    SAGE-MAKER-MODEL-NAME
  # xgboost-model           Created   model-5c06b18921e411ea91230292f5024981
  kubectl get -n "${model_namespace}" model "${k8s_model_name}" | sed -n '2 p' |  awk '{print $3}'
}

# This function waits until a CRD has reached a particular state, or times out
# Parameter:
#    $1: Namespace of CRD
#    $2: Kind of CRD
#    $3: Instance of CRD
#    $4: Timeout to complete the test
#    $5: The status that verifies the job has succeeded.
function wait_for_crd_status()
{
  local crd_namespace="$1"
  local crd_type="$2"
  local crd_instance="$3"
  local timeout="$4"
  local desired_status="$5"

  timeout "${timeout}" bash -c \
    'until [ "$(kubectl get -n "$0" "$1" "$2" -o=custom-columns=STATUS:.status | grep -i "$3" | wc -l)" -eq "1" ]; do \
      sleep 5; \
    done' "$crd_namespace" "$crd_type" "$crd_instance" "$desired_status"

  if [ $? -ne 0 ]; then
    return 1
  fi
}

# Force deletes all trainingJobs which might have been left dangling.
# Parameter:
#    $1: Namespace of CRD
function force_delete_training_jobs()
{
  local crd_namespace="$1"
  training_jobs=$(kubectl get trainingjobs -n "$crd_namespace" -ojson | jq -r '.items | .[] | .metadata.name')

  for job in $training_jobs
  do
      echo "Removing finalizer for ${job}"
      kubectl patch -n "$crd_namespace" trainingjob $job -p '{"metadata":{"finalizers":null}}' --type=merge
  done

  kubectl delete -n "$crd_namespace" trainingjob --all
}

# Cleans up all resources created during tests.
# Parameter:
#    $1: Namespace of CRD
function delete_all_resources()
{
  local crd_namespace="$1"
  kubectl delete -n "$crd_namespace" hyperparametertuningjob --all
  kubectl delete -n "$crd_namespace" processingjob --all
  kubectl delete -n "$crd_namespace" batchtransformjob --all
  # HAP must be deleted before hostingdeployment
  kubectl delete -n "$crd_namespace" hostingautoscalingpolicies --all
  kubectl delete -n "$crd_namespace" endpointconfig --all
  kubectl delete -n "$crd_namespace" hostingdeployment --all
  kubectl delete -n "$crd_namespace" model --all

  kubectl delete -n "$crd_namespace" trainingjob --all --timeout=3m
  if [ $? -ne 0 ]; then
    echo "Delete failed, will need to force delete"
  fi

  force_delete_training_jobs "$crd_namespace"
}

# A helper function to generate an IAM Role name for the current cluster and specified namespace
# Parameter:
#    $1: Namespace of CRD
function generate_iam_role_name {
    local crd_namespace="$1"
    local cluster=$(echo "${cluster_name}" | cut -d'/' -f2)

    role_name="${cluster}-${crd_namespace}"
    # IAM Role name must have length less than 64
    role_name=`echo $role_name|cut -c1-64`
}

# A function that cleans up an IAM Role after detaching the sagemaker access policy.
# Parameter:
#    $1: Name of the Role to be deleted
function delete_generated_role {
    local role_to_delete="${1}"
    # Delete the role associated with the cluster thats being deleted
    aws iam detach-role-policy --role-name "${role_to_delete}" --policy-arn arn:aws:iam::aws:policy/AmazonSageMakerFullAccess
    aws iam delete-role --role-name "${role_to_delete}"
}

# A function to print the operator logs for a given namespace
# Parameter:
#    $1: Namespace of the operator
function get_manager_logs {
    local crd_namespace="${1:-$default_operator_namespace}"
    if [ "${PRINT_DEBUG}" != "false" ]; then
        echo "Controller manager logs in the ${crd_namespace} namespace"
        kubectl -n "${crd_namespace}" logs "$(kubectl get pods -n "${crd_namespace}" | grep sagemaker-k8s-operator-controller-manager | awk '{print $1}')" manager
    fi
}

# A function that cleans up Jobs, CRDs and Operator in the default namespace
function cleanup_default_namespace {
    set +e
    get_manager_logs
    delete_all_resources "default"
    rolebased_operator_install_or_delete "${default_operator_namespace}" "config/installers/rolebasedcreds" "${default_role_name}" "delete"
}

# A helper function that generates the namespace-scoped operator installer yaml file with updated namespace and role.
# It also applies/deletes the generated installer as specified.
# Parameter:
#    $1: Namespace to which the op pod is installed. For cluster scope, this value is not used.
#    $2: Path to the kustomize source
#    $3: Name of the IAM Role to use
#    $4: kubectl action - apply or delete the resources defined in the installer
# TODO:  Investigate if it is possible to overlay values when we build the Kustomize targets instead.
function rolebased_operator_install_or_delete {
    local crd_namespace="$1"
    local kustomize_file_path="$2"
    local role="$3"
    local kubectl_action="$4"
    local aws_account=$(aws sts get-caller-identity --query Account --output text)

    kustomize build "${kustomize_file_path}" > temp_file.yaml
    sed -i "s/PLACEHOLDER-NAMESPACE/$crd_namespace/g" temp_file.yaml
    sed -i "s/123456789012/$aws_account/g" temp_file.yaml
    sed -i "s/DELETE_ME/$role/g" temp_file.yaml

    kubectl "${kubectl_action}" -f temp_file.yaml
    rm temp_file.yaml
}

# A function that builds the kustomize target, then deploys the CRDs and operator in cluster scope.
# Parameter:
#    $1: Name of the IAM Role to use
function operator_clusterscope_deploy {
    local role="$1"

    # Allow for overriding the installation of the CRDs/controller image from the
    # build scripts if we want to use our own installation

        pushd sagemaker-k8s-operator/sagemaker-k8s-operator-install-scripts
            echo "Deploying the operator to the default namespace"
            rolebased_operator_install_or_delete "${default_operator_namespace}" "config/installers/rolebasedcreds" "${role}" "apply"
        popd
        echo "Waiting for controller pod to be Ready"
        # Wait to increase chance that pod is ready
        # TODO: Should upgrade kubectl to version that supports `kubectl wait pods --all`
        sleep 60
        kubectl get pods --all-namespaces | grep sagemaker
}

# A function that builds the kustomize target, then deploys the CRDs (cluster scope) and operator (namespace scope).
# Parameter:
#    $1: Namespace of CRD
#    $2: Name of the IAM Role to use
function operator_namespace_deploy {
    local crd_namespace="$1"
    local role="$2"

    # Goto directory that holds the CRD
    pushd sagemaker-k8s-operator/sagemaker-k8s-operator-install-scripts
        kustomize build config/crd | kubectl apply -f -
        rolebased_operator_install_or_delete ${crd_namespace} "config/installers/rolebasedcreds/namespaced" "${role}" "apply"
    popd
    echo "Waiting for controller pod to be Ready"
    # Wait to increase chance that pod is ready
    # TODO: Should upgrade kubectl to version that supports `kubectl wait pods --all`
    sleep 60
    echo "Print manager pod status"
    kubectl get pods --all-namespaces | grep sagemaker
}


#!/bin/bash

source common.sh
source inject_tests.sh

# Applies each of the resources needed for the canary tests.
# Parameter:
#    $1: Namespace of the CRD
function run_canary_tests
{
  local crd_namespace="$1"

  echo "Injecting variables into tests"
  inject_all_variables

  echo "Starting Canary Tests"
  run_test "${crd_namespace}" testfiles/xgboost-mnist-trainingjob.yaml
  run_test "${crd_namespace}" testfiles/kmeans-mnist-processingjob.yaml
  run_test "${crd_namespace}" testfiles/xgboost-mnist-hpo.yaml
  # Special code for batch transform till we fix issue-59
  run_test "${crd_namespace}" testfiles/xgboost-model.yaml
  # We need to get sagemaker model before running batch transform
  verify_test "${crd_namespace}" Model xgboost-model 5m Created
  yq w -i testfiles/xgboost-mnist-batchtransform.yaml "spec.modelName" "$(get_sagemaker_model_from_k8s_model "${crd_namespace}" xgboost-model)"
  run_test "${crd_namespace}" testfiles/xgboost-mnist-batchtransform.yaml
  run_test "${crd_namespace}" testfiles/xgboost-hosting-deployment.yaml
}

# Applies each of the resources needed for the canary tests.
# Parameter:
#    $1: Namespace of the CRD
function run_canary_tests_china
{
  local crd_namespace="$1"

  echo "Injecting variables into tests"
  inject_all_variables

  echo "Starting Canary Tests China"
  run_test "${crd_namespace}" testfiles/xgboost-mnist-trainingjob-china.yaml
  run_test "${crd_namespace}" testfiles/xgboost-mnist-hpo-china.yaml
  # Special code for batch transform till we fix issue-59
  run_test "${crd_namespace}" testfiles/xgboost-model-china.yaml
  # We need to get sagemaker model before running batch transform
  verify_test "${crd_namespace}" Model xgboost-model-china 5m Created
  yq w -i testfiles/xgboost-mnist-batchtransform-china.yaml "spec.modelName" "$(get_sagemaker_model_from_k8s_model "${crd_namespace}" xgboost-model-china)"
  run_test "${crd_namespace}" testfiles/xgboost-mnist-batchtransform-china.yaml
}

# Applies each of the resources needed for the integration tests.
# Parameter:
#    $1: Namespace of the CRD
function run_integration_tests
{
  local crd_namespace="$1"
  run_canary_tests "${crd_namespace}"

  # TODO: Automate creation/testing of EFS file systems for relevant jobs
  # Build prerequisite resources
  if [ "$FSX_ID" == "" ]; then
    echo "Skipping build_fsx_from_s3 as fsx tests are disabled"
    #build_fsx_from_s3
  fi

  echo "Starting integration tests"
  run_test "${crd_namespace}" testfiles/spot-xgboost-mnist-trainingjob.yaml
  run_test "${crd_namespace}" testfiles/xgboost-mnist-custom-endpoint.yaml
  # run_test "${crd_namespace}" testfiles/efs-xgboost-mnist-trainingjob.yaml
  # run_test "${crd_namespace}" testfiles/fsx-kmeans-mnist-trainingjob.yaml
  run_test "${crd_namespace}" testfiles/spot-xgboost-mnist-hpo.yaml
  run_test "${crd_namespace}" testfiles/xgboost-mnist-hpo-custom-endpoint.yaml
  run_test "${crd_namespace}" testfiles/xgboost-mnist-trainingjob-debugger.yaml
  run_test "${crd_namespace}" testfiles/xgboost-hosting-deployment-multi-container.yaml
  run_hap_test "${crd_namespace}" named-xgboost-hosting testfiles/xgboost-hostingautoscaling.yaml testfiles/xgboost-hostingautoscaling-custom.yaml
  run_test "${crd_namespace}" testfiles/hd-retain-varient-properties.yaml
}

# Verifies that all resources were created and are running/completed for the canary tests.
# Parameter:
#    $1: Namespace of the CRD
function verify_canary_tests
{
  local crd_namespace="$1"
  echo "Verifying canary tests"
  verify_test "${crd_namespace}" TrainingJob xgboost-mnist 20m Completed
  verify_test "${crd_namespace}" ProcessingJob kmeans-mnist 20m Completed
  verify_test "${crd_namespace}" HyperparameterTuningJob xgboost-mnist-hpo 20m Completed
  verify_test "${crd_namespace}" BatchTransformJob xgboost-batch 20m Completed
  verify_test "${crd_namespace}" HostingDeployment xgboost-hosting 90m InService
}

# Verifies that all resources were created and are running/completed for the canary tests.
# Parameter:
#    $1: Namespace of the CRD
function verify_canary_tests_china
{
  local crd_namespace="$1"
  echo "Verifying canary tests"
  verify_test "${crd_namespace}" TrainingJob xgboost-mnist-china 20m Completed
  verify_test "${crd_namespace}" HyperparameterTuningJob xgboost-mnist-hpo-china 20m Completed
  verify_test "${crd_namespace}" BatchTransformJob xgboost-batch-china 20m Completed
}

# Verifies that all resources were created and are running/completed for the integration tests.
# Parameter:
#    $1: Namespace of the CRD
function verify_integration_tests
{
  local crd_namespace="$1"
  echo "Verifying integration tests"
  verify_canary_tests "${crd_namespace}"

  verify_test "${crd_namespace}" TrainingJob spot-xgboost-mnist 20m Completed
  verify_test "${crd_namespace}" TrainingJob xgboost-mnist-custom-endpoint 20m Completed
  # verify_test "${crd_namespace}" TrainingJob efs-xgboost-mnist 20m Completed
  # verify_test "${crd_namespace}" TrainingJob fsx-kmeans-mnist 20m Completed
  verify_test "${crd_namespace}" HyperparameterTuningJob spot-xgboost-mnist-hpo 20m Completed
  verify_test "${crd_namespace}" HyperparameterTuningJob xgboost-mnist-hpo-custom-endpoint 20m Completed
  verify_test "${crd_namespace}" HostingDeployment xgboost-hosting-multi-container 90m InService
  verify_test "${crd_namespace}" TrainingJob xgboost-mnist-debugger 20m Completed
  # Verify that debug job has status
  verify_debug_test "${crd_namespace}" TrainingJob xgboost-mnist-debugger 20m NoIssuesFound
  verify_test "${crd_namespace}" HostingAutoscalingPolicy hap-predefined 5m Created
  verify_test "${crd_namespace}" HostingAutoscalingPolicy hap-custom-metric 5m Created
  # verify_hap_test "3"
  verify_retain_varient_properties "${crd_namespace}" testfiles/hd-retain-varient-properties.yaml
}

function verify_retain_varient_properties(){
  local crd_namespace="$1"
  local hostingdeployment_yaml_filepath="$2"
  local autoscaling_yaml_filepath="testfiles/hd-autoscaling-retain-varient-properties.yaml"

  echo "Varifying retain varient properties"

  # endpoint is already created with instance count 1 and instance weight 2
  enpoint_name=$(yq r $hostingdeployment_yaml_filepath "spec.endpointName")
  endpoint_region=$(yq r $hostingdeployment_yaml_filepath "spec.region")
  hostingdeployment_name=$(yq r $hostingdeployment_yaml_filepath "metadata.name")
  wait_for_crd_status $crd_namespace HostingDeployment $hostingdeployment_name 90m InService

  # verify that the already created endpoint has instance count 1 and weight 2
  instance_count=$(aws sagemaker describe-endpoint --endpoint-name $enpoint_name --region $endpoint_region --query ProductionVariants[0].CurrentInstanceCount)
  weight=$(aws sagemaker describe-endpoint --endpoint-name $enpoint_name --region $endpoint_region --query ProductionVariants[0].CurrentWeight | awk '{printf "%.0f\n", $1}')
  if [ "${instance_count}" == "1" ] && [ "${weight}" == "2" ]; then
    echo "Initial verification that instance_count is 1 and weight is 2"
  else
    echo "[FAILED] Ininitial hosting deployment has wrong properties. Maybe the default namespace test didn't reset the values"
    echo "instance_count: $instance_count, weight: $weight"
    exit 1
  fi

  # autoscalling increases instance count to 2 (while keeping the previous weight 2)
  yq w -i $autoscaling_yaml_filepath "spec.resourceId[0].endpointName" $enpoint_name
  kubectl apply -n $crd_namespace -f $autoscaling_yaml_filepath
  wait_for_crd_status $crd_namespace HostingDeployment $hostingdeployment_name 90m Updating
  wait_for_crd_status $crd_namespace HostingDeployment $hostingdeployment_name 90m InService

  # verify instance count is 2
  instance_count=$(aws sagemaker describe-endpoint --endpoint-name $enpoint_name --region $endpoint_region --query ProductionVariants[0].CurrentInstanceCount)
  if [ "${instance_count}" == "2" ]; then
    echo "Autoscalling applied to increase instance count to 2"
  else
    echo "[FAILED] Failed to apply autoscaling policy in retain varient properties test"
    echo "instance_count: $instance_count"
    exit 1
  fi

  # re-apply hosting deployment with instance count 1 and instance weight 1
  # retainAllVariantProperties "true" and excludeRetainedVariantProperties "DesiredWeight"
  #
  # This should retain the previous instance count of 2 and
  # shouldn't retain the previous instance weight of 2 and change it to 3
  yq w -i $hostingdeployment_yaml_filepath "spec.productionVariants[0].initialVariantWeight" 3
  kubectl apply -n $crd_namespace -f $hostingdeployment_yaml_filepath
  wait_for_crd_status $crd_namespace HostingDeployment $hostingdeployment_name 90m Updating
  wait_for_crd_status $crd_namespace HostingDeployment $hostingdeployment_name 90m InService
  # Check if it retained the previous instance count and did not retain previous weight
  # sleep 5 && wait_for_crd_status "${crd_namespace}" HostingDeployment $hostingdeployment_name 40m InService && sleep 5
  instance_count=$(aws sagemaker describe-endpoint --endpoint-name $enpoint_name --region $endpoint_region --query ProductionVariants[0].CurrentInstanceCount)
  weight=$(aws sagemaker describe-endpoint --endpoint-name $enpoint_name --region $endpoint_region --query ProductionVariants[0].CurrentWeight | awk '{printf "%.0f\n", $1}')
  if [ "${instance_count}" == "2" ] && [ "${weight}" == "3" ]; then
    echo "[PASSED] Successfully tested retainAllVariantProperties and excludeRetainedVariantProperties"
  else
    echo "[FAILED] Failed to retain the variant properties"
    echo "instance_count: $instance_count, weight: $weight"
    exit 1
  fi

  # changing weight back to 2 for namaspace based tests
  yq w -i $hostingdeployment_yaml_filepath "spec.productionVariants[0].initialVariantWeight" 2
}

# Replaces the names of the endpoint generated in the previous test into the hap spec file and runs the test.
# Parameter:
#    $1: Target namespace
#    $2: K8s Name and endpointName of the hostingdeployment to apply autoscaling
#    $3: Filename of the hap test to run
#    $4: Filename of the custom metric hap test to run
function run_hap_test()
{
  local target_namespace="$1"
  local hosting_deployment_1="$2"
  local file_name="$3"
  local file_name_custom="$4"
  local hosting_deployment_2="${hosting_deployment_1}-2"
  local hosting_deployment_3="${hosting_deployment_1}-3"
  local hostingdeployment_type="hostingdeployment"

  # Create the 3 Endpoints
  update_hostingdeployment_input_metadata "${crd_namespace}" "${hosting_deployment_1}" "true"
  update_hostingdeployment_input_metadata "${crd_namespace}" "${hosting_deployment_2}" "true"
  update_hostingdeployment_input_metadata "${crd_namespace}" "${hosting_deployment_3}" "true"

  # This can be removed, but will make debugging easier with the same runtime.
  verify_test "${crd_namespace}" HostingDeployment "${hosting_deployment_1}" 90m InService
  verify_test "${crd_namespace}" HostingDeployment "${hosting_deployment_2}" 90m InService
  verify_test "${crd_namespace}" HostingDeployment "${hosting_deployment_3}" 90m InService

  # HAP Test 1: Using the Predefined Metric
  yq w -i "$file_name" "spec.resourceId[0].endpointName" "${hosting_deployment_1}"
  yq w -i "$file_name" "spec.resourceId[1].endpointName" "${hosting_deployment_2}"
  run_test "$target_namespace" "$file_name"

  # HAP Test 2: Using the Custom Metric
  yq w -i "$file_name_custom" "spec.resourceId[0].endpointName" "${hosting_deployment_3}"
  yq w -i "$file_name_custom" "spec.targetTrackingScalingPolicyConfiguration.customizedMetricSpecification.dimensions[0].value" "${hosting_deployment_3}"
  run_test "$target_namespace" "$file_name_custom"

  update_hostingdeployment_input_metadata "${crd_namespace}" "${hosting_deployment_1}" "false"
}

# This function verifies that the number of scaling policies applied is as expected
# Parameter:
#    $1: The expected number of policies
function verify_hap_test()
{
  local expected_number_of_policies="$1"
  scaling_policies="$(aws application-autoscaling describe-scaling-policies --region us-west-2 --service-namespace sagemaker | jq .ScalingPolicies)"
  number_of_policies_applied="$(echo $scaling_policies | jq length)"

  if [ "${number_of_policies_applied}" == "${expected_number_of_policies}" ]; then
    echo "[PASSED] All Scaling Policies were successfully applied to the specified number of endpoints"
  else
    echo "[FAILED] ${number_of_policies_applied} scaling policy/policies were applied. Expected ${expected_number_of_policies}."
    exit 1
  fi
}

# This function verifies that a given debug job has specific status
# Parameter:
#    $1: CRD namespace
#    $1: CRD type
#    $2: Instance of CRD
#    $3: Single debug job status
function verify_debug_test
{
  local crd_namespace="$1"
  local crd_type="$2"
  local crd_instance="$3"
  local timeout="$4"
  local expected_debug_job_status="$5"
  # First verify that trainingjob has been completed
  verify_test "${crd_namespace}" TrainingJob xgboost-mnist-debugger $timeout Completed

  # TODO extend this for multiple debug job with debug job statuses parameter

  echo "Waiting for debug job to finish"
  timeout "${timeout}" bash -c \
    'until [ "$(kubectl get -n "$0" "$1" "$2" -o json | jq -r .status.debugRuleEvaluationStatuses[0].ruleEvaluationStatus)" == "$3" ]; do \
      echo "Debug job has not completed yet"; \
      sleep 1; \
    done' "${crd_namespace}" "${crd_type}" "${crd_instance}" "${expected_debug_job_status}"

  if [ $? -ne 0 ]; then
    echo "[FAILED] Debug job ${crd_type} ${crd_instance} with expected status ${expected_debug_job_status} has timed out"
    exit 1
  fi

  echo "Debug job has completed"
}

# This function verifies that job has started and not failed
# Parameter:
#    $1: Namespace of the CRD
#    $2: Kind of CRD
#    $3: Instance of CRD
#    $4: Timeout to complete the test
#    $5: The status that verifies the job has succeeded.
# e.g. verify_test default trainingjobs xgboost-mnist
function verify_test()
{
  local crd_namespace="$1"
  local crd_type="$2"
  local crd_instance="$3"
  local timeout="$4"
  local desired_status="$5"

  # Check if job exist
  kubectl get -n "${crd_namespace}" "${crd_type}" "${crd_instance}"
  if [ $? -ne 0 ]; then
    echo "[FAILED] ${crd_type} ${crd_namespace}:${crd_instance} job does not exist"
    exit 1
  fi

  # Wait until Training Job has started, there will be two rows, header(STATUS) and entry.
  # Entry will be none if the job has not started yet. When job starts none will disappear
  # and real status will be present.
  echo "Waiting for ${crd_type} ${crd_namespace}:${crd_instance} job to start"
  timeout 1m bash -c \
    'until [ "$(kubectl get -n "$0" "$1" "$2" -o=custom-columns=STATUS:.status | grep -i none | wc -l)" -eq "0" ]; do \
      echo "${crd_type} ${crd_namespace}:${crd_instance} job has not started yet"; \
      sleep 1; \
    done' "${crd_namespace}" "${crd_type}" "${crd_instance}"

  # Job has started, check whether it has failed or not
  kubectl get -n "${crd_namespace}" "${crd_type}" "${crd_instance}" -o=custom-columns=NAME:.status | grep -i fail
  if [ $? -eq 0 ]; then
    echo "[FAILED] ${crd_type} ${crd_namespace}:${crd_instance} job has failed"
    exit 1
  fi

  echo "Waiting for ${crd_type} ${crd_namespace}:${crd_instance} job to complete"
  if ! wait_for_crd_status "$crd_namespace" "$crd_type" "$crd_instance" "$timeout" "$desired_status"; then
      echo "[FAILED] Waiting for status Completed failed"
      exit 1
  fi

  # Check weather job has completed or not
  echo "[PASSED] Verified ${crd_type} ${crd_namespace}:${crd_instance} has completed"
}

# Build a new FSX file system from an S3 data source to be used by FSX integration tests.
function build_fsx_from_s3()
{
  echo "Building FSX file system from S3 data source"

  NEW_FS="$(aws fsx create-file-system \
  --file-system-type LUSTRE \
  --lustre-configuration ImportPath=s3://${DATA_BUCKET}/kmeans_mnist_example \
  --storage-capacity 1200 \
  --subnet-ids subnet-187e9960 \
  --tags Key="Name",Value="$(date '+%Y-%m-%d-%H-%M-%S')" \
  --region us-west-2)"

  echo $NEW_FS

  FSX_ID="$(echo "$NEW_FS" | jq -r ".FileSystem.FileSystemId")"
  FS_AVAILABLE=CREATING
  until [[ "${FS_AVAILABLE}" != "CREATING" ]]; do
    FS_AVAILABLE="$(aws fsx --region us-west-2 describe-file-systems --file-system-id "${FSX_ID}" | jq -r ".FileSystems[0].Lifecycle")"
    sleep 30
  done
  aws fsx --region us-west-2 describe-file-systems --file-system-id ${FSX_ID}

  if [[ "${FS_AVAILABLE}" != "AVAILABLE" ]]; then
    exit 1
  fi

  export FSX_ID=$FSX_ID
}

# Replaces the names of the endpoint and optionally  creates it
# Parameter:
#    $1: Target namespace
#    $2: K8s Name and endpointName of the hostingdeployment to apply autoscaling
#    $3: Boolean to indicate if the modified HostingDeployment should be applied
function update_hostingdeployment_input_metadata()
{
  local crd_namespace="$1"
  local hosting_deployment="$2"
  local is_create="$3"
  local file_name_hd="testfiles/xgboost-hosting-deployment-with-name.yaml"

  yq w -i "${file_name_hd}" "metadata.name" $hosting_deployment
  yq w -i "${file_name_hd}" "spec.endpointName" $hosting_deployment

  if [[ "${is_create}" == "true" ]]; then
    run_test "${crd_namespace}" "${file_name_hd}"
  fi
}


#!/bin/bash

source common.sh
source create_tests.sh # TODO: Remove temporary import for "verify_test" method

# Verify that k8s and SageMaker resources are deleted when the user uses the delete verb.
# Parameter:
#    $1: CRD namespace
function run_delete_canary_tests
{
  # Run verbose output for delete tests for simpler debugging
  set -x
  local crd_namespace="$1"
  echo "Running delete canary tests"
  verify_delete "${crd_namespace}" TrainingJob testfiles/xgboost-mnist-trainingjob-debugger.yaml
  verify_delete "${crd_namespace}" ProcessingJob testfiles/kmeans-mnist-processingjob.yaml
  verify_delete "${crd_namespace}" HyperparameterTuningJob testfiles/xgboost-mnist-hpo.yaml

  # Create model before running batch delete test
  run_test "${crd_namespace}" testfiles/xgboost-model.yaml
  verify_test "${crd_namespace}" Model xgboost-model 5m Created
  yq w -i testfiles/xgboost-mnist-batchtransform.yaml "spec.modelName" "$(get_sagemaker_model_from_k8s_model "${crd_namespace}" xgboost-model)"

  verify_delete "${crd_namespace}" BatchTransformJob testfiles/xgboost-mnist-batchtransform.yaml
  set +x
}

# Parameter:
#    $1: CRD namespace
function run_delete_integration_tests
{
  echo "Running delete integration tests"
  local crd_namespace="$1"
  run_delete_canary_tests "${crd_namespace}"
}

# Applies a k8s resource, waits for it to start and then immediately deletes
# and ensure the state of the resource in k8s is stopping.
# Parameter:
#    $1: CRD namespace
#    $2: CRD type
#    $3: Filename of test
#    $4: (Optional) Timeout
# e.g. verify_delete default TrainingJob xgboost-mnist.yaml 60s
function verify_delete()
{
  local crd_namespace="$1"
  local crd_type="$(echo "$2" | awk '{print tolower($0)}')"
  local file_name="$3"
  local timeout=${4:-"10m"}

  # Apply file and get name
  local k8s_resource_name="$(kubectl apply -n "${crd_namespace}" -f "${file_name}" -o json | jq -r ".metadata.name")"

  # Wait until job has started.
  if ! wait_for_crd_status "$crd_namespace" "$crd_type" "$k8s_resource_name" "$timeout" "InProgress"; then
      echo "[FAILED] Waiting for ${crd_type} ${crd_namespace}:${k8s_resource_name} job status InProgress failed"
      exit 1
  fi

  local jobNamePath
  case $crd_type in
    trainingjob)
      jobNamePath=".status.sageMakerTrainingJobName"
      ;;
    processingjob)
      jobNamePath=".status.sageMakerProcessingJobName"
      ;;
    hyperparametertuningjob)
      jobNamePath=".status.sageMakerHyperParameterTuningJobName"
      ;;
    batchtransformjob)
      jobNamePath=".status.sageMakerTransformJobName"
      ;;
    hostingdeployment)
      jobNamePath=".status.endpointName"
      ;;
    *)
      echo "[FAILED] Delete did not recognise CRD type ${crd_type}"
      exit 1
      ;;
  esac

  local job_name="$(kubectl get -n "${crd_namespace}" "${crd_type}" "${k8s_resource_name}" -o json | jq -r "${jobNamePath}")"

  # Actually delete the resource
  kubectl delete -n "${crd_namespace}" -f "${file_name}" &

  # Wait for the status to be Stopping. If it does not happen, then check if the job still exists. If it doesn't, then the delete occurred. If it still exists, then the delete failed.
  if ! wait_for_crd_status "$crd_namespace" "${crd_type}" "${k8s_resource_name}" "${timeout}" "Stopping" &&
  kubectl get -n "${crd_namespace}" "${crd_type}" "${k8s_resource_name}"; then
      echo "[FAILED] Waiting for ${crd_type} ${crd_namespace}:${k8s_resource_name} job status Stopping failed"
      exit 1
  fi
  # Check that it has been stopped in SageMaker
  verify_sm_resource_stopping_else_fail "${crd_type}" "${job_name}"

  echo "[PASSED] Verified ${crd_type} ${crd_namespace}:${k8s_resource_name} deleted from k8s successfully"
}

# Queries AWS SageMaker to ensure that the job created by a k8s resource has
# been marked as stopped or stopping.
# Parameter:
#    $1: CRD type
#    $2: SageMaker job name
function verify_sm_resource_stopping_else_fail()
{
  local crd_type="$(echo "$1" | awk '{print tolower($0)}')"
  local job_name="$2"

  local jobStatusPath sageMakerResourceType
  case $crd_type in
    trainingjob)
      jobStatusPath=".TrainingJobStatus"
      sageMakerResourceType="training-job"
      ;;

    processingjob)
      jobStatusPath=".ProcessingJobStatus"
      sageMakerResourceType="processing-job"
      ;;

    hyperparametertuningjob)
      jobStatusPath=".HyperParameterTuningJobStatus"
      sageMakerResourceType="hyper-parameter-tuning-job"
      ;;

    batchtransformjob)
      jobStatusPath=".TransformJobStatus"
      sageMakerResourceType="transform-job"
      ;;

    hostingdeployment)
      jobStatusPath=".EndpointStatus"
      sageMakerResourceType="endpoint"
      ;;

    *)
      echo "[FAILED] Delete did not recognise CRD type ${crd_type}"
      exit 1
      ;;
  esac

  job_status="$(aws sagemaker "describe-${sageMakerResourceType}" "--${sageMakerResourceType}-name" "${job_name}" --region us-west-2 | jq -r "${jobStatusPath}")"
  # Job can go from Stopping to completed State if job is completing already like uploading model
  if [ "${job_status}" != "Stopping" ] && [ "${job_status}" != "Stopped" ] && ["${job_status}" != "Completed"]; then
    echo "[FAILED] AWS SageMaker resource \"${job_name}\" did not stop (status \"${job_status}\")"
    exit 1
  else
    echo "[PASSED] AWS SageMaker resource \"${job_name}\" deleted from SageMaker"
  fi
}


#!/bin/bash

source common.sh
source inject_tests.sh

# Applies each of the resources needed for the canary tests.
function run_feature_canary_tests
{
  echo "Running feature canary tests"
  inject_all_variables
}

# Applies each of the resources needed for the integration tests.
# Parameter:
#    $1: CRD namespace
function run_feature_integration_tests
{
  echo "Running feature integration tests"
  local crd_namespace="$1"
  run_feature_canary_tests

  run_test "${crd_namespace}" testfiles/failing-xgboost-mnist-hpo.yaml
  run_test "${crd_namespace}" testfiles/failing-xgboost-mnist-trainingjob.yaml
}

# Creates a training job in the specified namespace
# For this test, the job is created in a namespace that does not have the operator.
# Parameter:
#    $1: CRD namespace
function run_namespace_deployment_tests
{
  echo "Running feature namespaced test"
  local crd_namespace="$1"
  # Create the noop-namespace if it does not exist
  kubectl create namespace "${crd_namespace}"
  run_feature_canary_tests
  run_test "${crd_namespace}" testfiles/xgboost-mnist-trainingjob-namespaced.yaml
}

# Verifies that the job created in an incorrect namespace does not gain a status or sagemaker name.
# Parameter:
#    $1: CRD namespace
function verify_job_fails_outside_operator_namespace
{
  echo "Verifying namespace deployment test"
  local crd_namespace="$1"
  sleep 5

  if verify_job_has_no_status "$crd_namespace" TrainingJob "xgboost-mnist-namespaced"; then
    echo "[FAILED] TrainingJob deployed to $crd_namespace namespace was created and has a job status"
    exit 1
  fi
  if verify_trainingjob_has_no_sagemaker_name "$crd_namespace" "xgboost-mnist-namespaced"; then
    echo "[FAILED] TrainingJob deployed to $crd_namespace namespace was created and has a sagemaker job name"
    exit 1
  fi

}

# Verify that each canary feature test has completed successfully.
function verify_feature_canary_tests
{
  echo "Verifying feature canary tests"
}

# Verifies that each integration feature test has completed successfully.
# Parameter:
#    $1: CRD namespace
function verify_feature_integration_tests
{
  echo "Verifying feature integration tests"
  local crd_namespace="$1"
  verify_feature_canary_tests

  if ! wait_for_crd_status "$crd_namespace" TrainingJob failing-xgboost-mnist 5m Failed; then
    echo "[FAILED] Failing training job never reached status Failed"
    exit 1
  fi
  if ! verify_resource_has_additional "$crd_namespace" TrainingJob failing-xgboost-mnist; then
    echo "[FAILED] Failing training job does not have any additional status set"
    exit 1
  fi
  echo "[SUCCESS] TrainingJob with Failed status has additional set"

  if ! wait_for_crd_status "$crd_namespace" HyperParameterTuningJob failing-xgboost-mnist-hpo 10m Failed; then
    echo "[FAILED] Failing hyperparameter tuning job never reached status Failed"
    exit 1
  fi
  if ! verify_resource_has_additional "$crd_namespace" HyperParameterTuningJob failing-xgboost-mnist-hpo; then
    echo "[FAILED] Failing hyperparameter tuning job does not have any additional status set"
    exit 1
  fi
  if ! verify_failed_trainingjobs_from_hpo_have_additional "$crd_namespace" failing-xgboost-mnist-hpo; then
    echo "[FAILED] Not all failed training jobs in failing HPO job contained the additional in their statuses"
    exit 1
  fi
  echo "[SUCCESS] HyperParameterTuningJob with Failed status has its additional status set and set for all TrainingJobs"
}

# This function verifies that a given training job has a failure reason in the
# additional part of the status.
# Parameter:
#    $1: CRD namespace
#    $2: CRD type
#    $3: Instance of CRD
function verify_resource_has_additional
{
  local crd_namespace="$1"
  local crd_type="$2"
  local crd_instance="$3"
  local additional="$(kubectl get -n "$crd_namespace" "$crd_type" "$crd_instance" -o json | jq -r '.status.additional')"

  if [ -z "$additional" ]; then
    return 1
  fi
}

# This function verifies that each Failed trainingjob for a given HPO job have
# the additional field in their status set.
# Parameter:
#    $1: Job namespace
#    $2: Instance of HyperParameterTuningJob
function verify_failed_trainingjobs_from_hpo_have_additional
{
  local crd_namespace="$1"
  local crd_instance="$2"

  local sagemaker_hpo_job_name="$(kubectl get -n "$crd_namespace" hyperparametertuningjob "$crd_instance" -o json | jq -r '.status.sageMakerHyperParameterTuningJobName')"

  # Loop over every failed training job and get the k8s resource name
  for training_job in $(kubectl get -n "$crd_namespace" trainingjob | grep Failed | grep "$sagemaker_hpo_job_name" | awk '{print $1}'); do
    if ! verify_resource_has_additional "$crd_namespace" TrainingJob "$training_job"; then
      return 1
    fi
  done
}

# This function verifies that a given job does not have a status.
# Returns 1 if job creation fails.
# Parameter:
#    $1: CRD namespace
#    $2: CRD type
#    $3: Instance of CRD
function verify_job_has_no_status
{
  local crd_namespace="$1"
  local crd_type="$2"
  local crd_instance="$3"
  local sagemaker_job_status="$(kubectl get -n "$crd_namespace" "$crd_type" "$crd_instance" -o json | jq -r '.status')"

  if [ "$sagemaker_job_status" = null ]; then
    echo "[SUCCESS] $crd_type deployed to $crd_namespace namespace does not have a status"
    return 1
  fi
}

# This function verifies that a given training job does not have a sagemaker job name.
# Returns 1 if job creation fails.
# Parameter:
#    $1: CRD namespace
#    $2: CRD type
#    $3: Instance of CRD
function verify_trainingjob_has_no_sagemaker_name
{
  local crd_namespace="$1"
  local crd_instance="$2"
  local crd_type="TrainingJob"
  local sagemaker_training_job_name="$(kubectl get -n "$crd_namespace" "$crd_type" "$crd_instance" -o json | jq -r '.status.sageMakerTrainingJobName')"

  if [ "$sagemaker_training_job_name" = null ]; then
    echo "[SUCCESS] $crd_type deployed to $crd_namespace namespace does not have a sagemaker job name"
    return 1
  fi
}


#!/usr/bin/env bash

# Helper script to generate an IAM Role needed to install operator using role-based authentication.
# https://sagemaker.readthedocs.io/en/stable/amazon_sagemaker_operators_for_kubernetes.html#create-an-iam-role
#
# Run as:
# $ ./generate_iam_role.sh ${cluster_arn/cluster_name} ${operator_namespace} ${role_name} ${cluster_region}
#

CLUSTER_ARN="${1}"
OPERATOR_NAMESPACE="${2}"
ROLE_NAME="${3}"
CLUSTER_REGION="${4:-us-east-1}"
aws_account=$(aws sts get-caller-identity --query Account --output text)
trustfile="trust.json"

# 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
        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"
./generate_trust_policy.sh ${CLUSTER_REGION} ${aws_account} ${oidc_id} ${OPERATOR_NAMESPACE} > "${trustfile}"
echo "Create the IAM Role using these values"
create_namespaced_iam_role "${trustfile}"
echo "Cleanup for the next run!"
delete_generated_file "${trustfile}"


#!/usr/bin/env bash

# Helper script to generate trust policy needed to install operator using role-based authentication.
# https://sagemaker.readthedocs.io/en/stable/amazon_sagemaker_operators_for_kubernetes.html#create-an-iam-role
#
# Run as:
# $ ./generate_trust_policy ${EKS_CLUSTER_REGION} ${AWS_ACCOUNT_ID} ${OIDC_ID} [optional: ${OPERATOR_NAMESPACE}] > 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 OPERATOR_NAMESPACE parameter is for when you want to run the operator in a custom namespace other than "sagemaker-k8s-operator-system".

cluster_region="$1"
account_number="$2"
oidc_id="$3"
operator_namespace="${4:-sagemaker-k8s-operator-system}"

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:'"${operator_namespace}"':sagemaker-k8s-operator-default"
        }
      }
    }
  ]
}
'


#!/bin/bash

# Inject environment variables into the job YAMLs
function inject_variables()
{
  variables=("ROLE_ARN" "DATA_BUCKET" "FSX_ID")

  local file_name="$1"
  for i in "${variables[@]}"
  do
    local curr_var=${!i}
    sed -i "s|{$i}|${curr_var}|g" "${file_name}"
  done
}

# Injects necessary environment variables into resource yaml specs. This allows
# for generic integration tests to be created, substituting values that are
# specific to the account they are run within.
function inject_all_variables
{
  inject_variables testfiles/xgboost-mnist-trainingjob.yaml
  inject_variables testfiles/kmeans-mnist-processingjob.yaml
  inject_variables testfiles/spot-xgboost-mnist-trainingjob.yaml
  inject_variables testfiles/xgboost-mnist-custom-endpoint.yaml
  # inject_variables testfiles/efs-xgboost-mnist-trainingjob.yaml
  # inject_variables testfiles/fsx-kmeans-mnist-trainingjob.yaml
  inject_variables testfiles/xgboost-mnist-hpo.yaml
  inject_variables testfiles/spot-xgboost-mnist-hpo.yaml
  inject_variables testfiles/xgboost-mnist-hpo-custom-endpoint.yaml
  inject_variables testfiles/xgboost-model.yaml
  inject_variables testfiles/xgboost-mnist-batchtransform.yaml
  inject_variables testfiles/xgboost-hosting-deployment.yaml
  inject_variables testfiles/xgboost-hosting-deployment-with-name.yaml
  inject_variables testfiles/xgboost-hosting-deployment-multi-container.yaml
  inject_variables testfiles/failing-xgboost-mnist-trainingjob.yaml
  inject_variables testfiles/failing-xgboost-mnist-hpo.yaml
  inject_variables testfiles/xgboost-mnist-trainingjob-debugger.yaml
  inject_variables testfiles/xgboost-mnist-trainingjob-namespaced.yaml
  inject_variables testfiles/xgboost-mnist-trainingjob-china.yaml
  inject_variables testfiles/xgboost-mnist-hpo-china.yaml
  inject_variables testfiles/xgboost-model-china.yaml
  inject_variables testfiles/xgboost-mnist-batchtransform-china.yaml
  inject_variables testfiles/hd-retain-varient-properties.yaml
}


#!/bin/bash

source create_tests.sh
source feature_tests.sh
source delete_tests.sh
source inject_tests.sh
source smlogs_tests.sh

run_canary_tests "default"
run_feature_canary_tests "default"
verify_canary_tests "default"
run_update_canary_tests "default"
verify_update_canary_tests "default"
verify_feature_canary_tests "default"
run_smlogs_canary_tests "default"
delete_all_resources "default" # Delete all existing resources to re-use metadata names
run_delete_canary_tests "default"


#!/bin/bash

source create_tests.sh
source feature_tests.sh
source delete_tests.sh
source inject_tests.sh
source smlogs_tests.sh

run_canary_tests_china "default"
verify_canary_tests_china "default"
run_update_canary_tests "default"
verify_update_canary_tests "default"
delete_all_resources "default" # Delete all existing resources to re-use metadata names


#!/bin/bash

source create_tests.sh
source delete_tests.sh
source inject_tests.sh
source feature_tests.sh

crd_namespace="${1}"
noop_namespace="noop-namespace"

run_canary_tests "${crd_namespace}"
verify_canary_tests "${crd_namespace}"
run_namespace_deployment_tests "${noop_namespace}" # Deploy job to namespace without operator.
verify_job_fails_outside_operator_namespace "${noop_namespace}"
delete_all_resources "${crd_namespace}" # Delete all existing resources to re-use metadata names
delete_all_resources "${noop_namespace}"
run_delete_canary_tests "${crd_namespace}"


#!/bin/bash

source create_tests.sh
source update_tests.sh
source feature_tests.sh
source delete_tests.sh
source inject_tests.sh
source smlogs_tests.sh

crd_namespace=${1}

run_integration_tests ${crd_namespace}
run_feature_integration_tests ${crd_namespace}
verify_integration_tests ${crd_namespace}
# run_update_integration_tests ${crd_namespace}
# verify_update_integration_tests ${crd_namespace}
verify_feature_integration_tests ${crd_namespace}
run_smlogs_integration_tests ${crd_namespace}
delete_all_resources ${crd_namespace} # Delete all existing resources to re-use metadata names
run_delete_integration_tests ${crd_namespace}


#!/bin/bash

# Verbose trace of commands, helpful since test iteration takes a long time.
set -x

# A function to delete cluster, if cluster was not launched this will fail, so test will fail ultimately too
function cleanup {
    # We want to run every command in this function, even if some fail.
    set +e

    # Managable to print logs for ephemeral cluster
    get_manager_logs

    delete_all_resources "default"

    if [ -z "${USE_EXISTING_CLUSTER}" ]
    then
        # Delete the role associated with the cluster thats being deleted
        delete_generated_role "${role_name}"
        # Tear down the cluster if we set it up.
        echo "USE_EXISTING_CLUSTER is false, tearing down cluster we created."
        eksctl delete cluster --name "${cluster_name}" --region "${CLUSTER_REGION}"
    fi
}

# Set the trap to clean up resources
# In case of error or normal exit delete the cluster
trap cleanup EXIT

# If any command fails, exit the script with an error code.
set -e

source ./common.sh

# TODOs
# 1. Add validation for each steps and abort the test if steps fails
# Build environment `Docker image` has all prerequisite setup and credentials are being passed using AWS system manager

CLUSTER_REGION=${CLUSTER_REGION:-us-east-1}
CLUSTER_VERSION=1.21

# Define the list of optional subnets for the EKS test cluster
CLUSTER_PUBLIC_SUBNETS=${CLUSTER_PUBLIC_SUBNETS:-}
CLUSTER_PRIVATE_SUBNETS=${CLUSTER_PRIVATE_SUBNETS:-}

# Output the commit SHA for logging sake
echo "Launching canary test for ${COMMIT_SHA}"

# Launch EKS cluster if we need to and define cluster_name,CLUSTER_REGION.
echo "Launching the cluster"

cluster_name="sagemaker-k8s-canary-"$(date '+%Y-%m-%d-%H-%M-%S')""

if [ -z "${USE_EXISTING_CLUSTER}" ]
then
   eksctl_args=( --nodes 1 --node-type=c5.xlarge --region "${CLUSTER_REGION}" --version "${CLUSTER_VERSION}" )
   [ "${CLUSTER_PUBLIC_SUBNETS}" != "" ] && eksctl_args+=( --vpc-public-subnets="${CLUSTER_PUBLIC_SUBNETS}" )
   [ "${CLUSTER_PRIVATE_SUBNETS}" != "" ] && eksctl_args+=( --vpc-private-subnets="${CLUSTER_PRIVATE_SUBNETS}" )

   eksctl create cluster "${cluster_name}" "${eksctl_args[@]}" --dry-run > generated-cluster.yaml
   yq w -i generated-cluster.yaml managedNodeGroups[0].disableIMDSv1 true

   eksctl create cluster -f generated-cluster.yaml --auto-kubeconfig --timeout=40m

   echo "Setting kubeconfig"
   export KUBECONFIG="/root/.kube/eksctl/clusters/${cluster_name}"
else
   cluster_name="non-ephemeral-cluster"
   aws eks update-kubeconfig --name "${cluster_name}" --region "${CLUSTER_REGION}"
fi

# Download the CRD
tar -xf sagemaker-k8s-operator.tar.gz

# jump to the root dir of operator
pushd sagemaker-k8s-operator

    # Setup the PATH for smlogs
    mv smlogs-plugin/linux.amd64/kubectl-smlogs /usr/bin/kubectl-smlogs
popd

generate_iam_role_name "${default_operator_namespace}"
./generate_iam_role.sh "${cluster_name}" "${default_operator_namespace}" "${role_name}" "${CLUSTER_REGION}"

#smlogs depends on clusterscope
operator_clusterscope_deploy "${role_name}"

# Run the integration test file
./run_all_sample_canary_tests.sh

# Send results back to results bucket
FILE_NAME=`TZ=UTC date +%Y-%m-%d-%H-%M-%S`
touch /tmp/$FILE_NAME
aws s3 cp /tmp/$FILE_NAME s3://${RESULT_BUCKET}/${CLUSTER_REGION}/$FILE_NAME


#!/usr/bin/env bash

# Env variables that need to be set before running this script
# ROLE_ARN : SageMaker executor role arn
# DATA_BUCKET : S3 bucket name where input data is stored

if [[ -z "${ROLE_ARN}" ]]; then
  echo "ROLE_ARN environment variable not found"
  exit 1
fi
if [[ -z "${DATA_BUCKET}" ]]; then
  echo "DATA_BUCKET environment variable not found"
  exit 1
fi

# Optional Env variables

if [ "${EKS_PRIVATE_SUBNET_1}" -a "${EKS_PRIVATE_SUBNET_2}" -a "${EKS_PUBLIC_SUBNET_1}" -a "${EKS_PUBLIC_SUBNET_2}" ];then
  echo "Will use existing subnets to create the EKS cluster"
  USE_EXISTING_SUBNET=True
else
  echo "Existing subnets not provided. Will Create new VPC during EKS cluster creation"
fi

# local variables
DEPLOYMENT_NAME="ephemeral-operator-canary-"$(date '+%Y-%m-%d-%H-%M-%S')""
CLUSTER_NAME=${DEPLOYMENT_NAME}-cluster
CLUSTER_REGION=${CLUSTER_REGION:-cn-northwest-1}
OIDC_ROLE_NAME=pod-role-$DEPLOYMENT_NAME
AWS_ACC_NUM=$(aws sts get-caller-identity --region $CLUSTER_REGION   --query Account --output text)
CLUSTER_VERSION=1.21

function download_installer_china(){
  if [ -f ./installer_china.yaml ]; then
    # codebuild spec copies this file from release/rolebased/china/installer_china.yaml to here
    echo "installer_china.yaml found in current path"
    return 0
  fi

  echo "installer_china.yaml not found in current path"
  echo "Trying to download it from github"

  n=0
  until [ "$n" -ge 3 ]
  do
    wget --retry-connrefused --waitretry=30 --read-timeout=20 --timeout=15 -t 3 \
      -O installer_china.yaml https://raw.githubusercontent.com/aws/amazon-sagemaker-operator-for-k8s/master/release/rolebased/china/installer_china.yaml \
      && break
    n=$((n+1))
    if [[ "$n" -ge 3 ]]; then
      echo "Failed to download installer_china.yaml"
      exit 1
    fi
    echo "Sleeping for" $((60*2*n)) s
    sleep $((60*2*n))
  done
}

function create_eks_cluster() {
  eksctl_args=( --managed --nodes 1 --node-type=c5.xlarge --region "$CLUSTER_REGION" --version "${CLUSTER_VERSION}" )
  [ ! -z "${USE_EXISTING_SUBNET}" ] && eksctl_args+=( --vpc-public-subnets="${EKS_PUBLIC_SUBNET_1},${EKS_PUBLIC_SUBNET_2}" )
  [ ! -z "${USE_EXISTING_SUBNET}" ] && eksctl_args+=( --vpc-private-subnets="${EKS_PRIVATE_SUBNET_1},${EKS_PRIVATE_SUBNET_2}" )

  eksctl create cluster "${CLUSTER_NAME}" "${eksctl_args[@]}" --dry-run > generated-cluster.yaml
  yq w -i generated-cluster.yaml managedNodeGroups[0].disableIMDSv1 true

  eksctl create cluster -f generated-cluster.yaml --auto-kubeconfig --timeout=40m
}

function install_k8s_operators() {
  echo "Installing SageMaker operators"
  aws --region $CLUSTER_REGION eks update-kubeconfig --name $CLUSTER_NAME
  eksctl utils associate-iam-oidc-provider --cluster $CLUSTER_NAME --region $CLUSTER_REGION --approve

  OIDC_URL=$(aws eks describe-cluster --region $CLUSTER_REGION --name $CLUSTER_NAME --query "cluster.identity.oidc.issuer" --output text | cut -c9-)

  printf '{
    "Version": "2012-10-17",
    "Statement": [
      {
        "Effect": "Allow",
        "Principal": {
          "Federated": "arn:aws-cn:iam::'$AWS_ACC_NUM':oidc-provider/'$OIDC_URL'"
        },
        "Action": "sts:AssumeRoleWithWebIdentity",
        "Condition": {
          "StringEquals": {
            "'$OIDC_URL':aud": "sts.amazonaws.com",
            "'$OIDC_URL':sub": "system:serviceaccount:sagemaker-k8s-operator-system:sagemaker-k8s-operator-default"
          }
        }
      }
    ]
  }
  ' > ./trust.json

  aws --region $CLUSTER_REGION iam create-role --role-name $OIDC_ROLE_NAME --assume-role-policy-document file://trust.json
  aws --region $CLUSTER_REGION iam attach-role-policy --role-name $OIDC_ROLE_NAME --policy-arn arn:aws-cn:iam::aws:policy/AmazonSageMakerFullAccess
  OIDC_ROLE_ARN=$(aws --region $CLUSTER_REGION iam get-role --role-name $OIDC_ROLE_NAME --output text --query 'Role.Arn')

  download_installer_china

  FIND_STR=$(yq r -d'*' installer_china.yaml 'metadata.annotations."eks.amazonaws.com/role-arn"')
  sed -i "s#$FIND_STR#$OIDC_ROLE_ARN#g" installer_china.yaml
  kubectl apply -f installer_china.yaml

  echo "Waiting for controller pod to be Ready"
  # Wait to increase chance that pod is ready
  # TODO: Should upgrade kubectl to version that supports `kubectl wait pods --all`
  sleep 60

}

function delete_generated_oidc_role {
    echo "Deleting generated OIDC role"
    aws iam detach-role-policy --region $CLUSTER_REGION  --role-name ${OIDC_ROLE_NAME} --policy-arn arn:aws-cn:iam::aws:policy/AmazonSageMakerFullAccess
    aws iam delete-role --region $CLUSTER_REGION --role-name ${OIDC_ROLE_NAME}
}

function delete_generated_cluster {
  echo "Deleting cluster"
  eksctl delete cluster --region $CLUSTER_REGION $CLUSTER_NAME
}

function cleanup() {
  set +e

  delete_generated_oidc_role
  delete_generated_cluster
}

trap cleanup EXIT

create_eks_cluster
install_k8s_operators

set -e
./run_all_sample_canary_tests_china.sh


#!/bin/bash

# TODOs
# 1. Add validation for each steps and abort the test if steps fails
# Build environment `Docker image` has all prerequisite setup and credentials are being passed using AWS system manager
# crd_namespace is currently hardcoded here for cleanup. Should be an env variable set in the pipeline and .env file

source tests/codebuild/common.sh
crd_namespace=${NAMESPACE_OVERRIDE:-"test-namespace"}

# Verbose trace of commands, helpful since test iteration takes a long time.
set -x

# A function to cleanup resources before exit. If cluster was not launched by the script, only the CRDs, jobs and operator will be deleted.
function cleanup {
    # We want to run every command in this function, even if some fail.
    set +e
    echo "[SECTION] Running final cluster cleanup, some commands may fail if resources do not exist."

    get_manager_logs "${crd_namespace}"
    delete_all_resources "${crd_namespace}"
    cleanup_default_namespace

    # Tear down the cluster if we set it up.
    if [ "${need_setup_cluster}" == "true" ]; then
        echo "need_setup_cluster is true, tearing down cluster we created."
        eksctl delete cluster --name "${cluster_name}" --region "${cluster_region}"
        # Delete the role associated with the cluster thats being deleted
        delete_generated_role "${role_name}"
        delete_generated_role "${default_role_name}"
    else
        echo "need_setup_cluster is not true, will remove operator without deleting cluster"
        kustomize build "$path_to_installer/config/default" | kubectl delete -f -
        # Delete the namespaced operator. TODO: This can be cleaner if parameterized
        kustomize build "$path_to_installer/config/crd" | kubectl delete -f -
        rolebased_operator_install_or_delete ${crd_namespace} "config/installers/rolebasedcreds/namespaced" "${role_name}" "delete"
    fi

    if [ "${existing_fsx}" == "false" ] && [ "$FSX_ID" != "" ]; then
        aws fsx --region us-west-2 delete-file-system --file-system-id "$FSX_ID"
    fi
}

# Set the trap to clean up resources
# In case of error or normal exit delete the cluster
trap cleanup EXIT

# When we run integration tests locally, the default kubeconfig file will be present.
# If we are not running locally, we will set up an EKS cluster for testing.
if [ ! -f ~/.kube/config ]; then
    echo "No kubeconfig found, going to set up own cluster."
    need_setup_cluster="true"
else
    echo "kubeconfig present, no need to set up cluster."
    need_setup_cluster="false"
fi
readonly need_setup_cluster

if [ "${FSX_ID}" != "" ]; then
    existing_fsx="true"
else
    existing_fsx="false"
fi
readonly existing_fsx

# If any command fails, exit the script with an error code.
set -e

# Launch EKS cluster if we need to and define cluster_name,cluster_region.
echo "[SECTION] Lauch EKS Cluster if needed"
if [ "${need_setup_cluster}" == "true" ]; then
    echo "Launching the cluster"
    readonly cluster_name="sagemaker-k8s-pipeline-"$(date '+%Y-%m-%d-%H-%M-%S')""
    readonly cluster_region="us-east-1"

    # By default eksctl picks random AZ, which time to time leads to capacity issue.
    eksctl create cluster "${cluster_name}" --timeout=40m --region "${cluster_region}" --zones us-east-1a,us-east-1b,us-east-1c --auto-kubeconfig --version=1.21 --fargate
    eksctl create fargateprofile --namespace "${crd_namespace}" --cluster "${cluster_name}" --name namespace-profile --region "${cluster_region}"
    eksctl create fargateprofile --namespace "${default_operator_namespace}" --cluster "${cluster_name}" --name operator-profile --region "${cluster_region}"

    echo "Setting kubeconfig"
    export KUBECONFIG="/root/.kube/eksctl/clusters/${cluster_name}"
else
    readonly cluster_info="$(kubectl config get-contexts | awk '$1 == "*" {print $3}' | sed 's/\./ /g')"
    readonly cluster_name="$(echo "${cluster_info}" | awk '{print $1}')"
    readonly cluster_region="$(echo "${cluster_info}" | awk '{print $2}')"
fi

echo "[SECTION] Setup for integration tests"
# Download the CRD from the tarball artifact bucket
aws s3 cp s3://$ALPHA_TARBALL_BUCKET/${CODEBUILD_RESOLVED_SOURCE_VERSION}/sagemaker-k8s-operator-us-west-2-alpha.tar.gz sagemaker-k8s-operator.tar.gz
tar -xf sagemaker-k8s-operator.tar.gz
# Jump to the root dir of the operator
pushd sagemaker-k8s-operator
# Setup the PATH for smlogs
    mv smlogs-plugin/linux.amd64/kubectl-smlogs /usr/bin/kubectl-smlogs
    path_to_installer=$(pwd)/sagemaker-k8s-operator-install-scripts
popd

echo "[SECTION] Run integration tests for the cluster scoped operator deployment"

generate_iam_role_name "${default_operator_namespace}"
default_role_name="${role_name}"  # role_name will get overwritten, save for cleanup.
pushd tests/codebuild
    ./generate_iam_role.sh "${cluster_name}" "${default_operator_namespace}" "${default_role_name}" "${cluster_region}"
popd

if [ "${SKIP_INSTALLATION}" == "true" ]; then
    echo "Skipping installation of CRDs and operator"
else
    operator_clusterscope_deploy "${default_role_name}"
fi

echo "[SECTION] Running integration tests in default namespace"
pushd tests/codebuild
    ./run_all_sample_test.sh "default"
popd

echo "Skipping private link test"
#pushd private-link-test && ./run_private_link_integration_test "${cluster_name}" "us-west-2" && popd

echo "[SECTION] Run integration tests for the namespaced operator deployment"

# Cleanup
echo "[SECTION] Cleaning the default namespace to test namespace deployment"
cleanup_default_namespace

# If any command fails, exit the script with an error code.
# Cleanup unsets this
set -e

#Create the IAM Role for the given namespace
generate_iam_role_name "${crd_namespace}"
pushd tests/codebuild
    ./generate_iam_role.sh "${cluster_name}" "${crd_namespace}" "${role_name}" "${cluster_region}"
popd

# Allow for overriding the installation of the CRDs/controller image from the
# build scripts if we want to use our own installation
if [ "${SKIP_INSTALLATION}" == "true" ]; then
    echo "Skipping installation of CRDs and operator"
else
    operator_namespace_deploy "${crd_namespace}" "${role_name}"
fi

# Run the integration test file
echo "[SECTION] Running integration tests for namespace deployment"
pushd tests/codebuild/
    ./run_all_sample_namespace_tests.sh "${crd_namespace}"
popd


#!/bin/bash

# Ensure that the smlogs kubectl plugin is able to detect and properly process the output of each job type.
function run_smlogs_canary_tests
{
  echo "Running smlogs canary tests"
  # Verify smlogs worked.
  if [ "$(kubectl smlogs trainingjob xgboost-mnist | wc -l)" -lt "1" ]; then
    echo "smlogs trainingjob did not produce any output."
    exit 1
  fi
  if [ "$(kubectl smlogs batchtransformjob xgboost-batch | wc -l)" -lt "1" ]; then
    echo "smlogs batchtransformjob did not produce any output."
    exit 1
  fi
}

function run_smlogs_integration_tests
{
  echo "Running smlogs integration tests"
  run_smlogs_canary_tests
}


#!/bin/bash

source common.sh
source inject_tests.sh

# Updates the spec and re-applies to ensure updates work as expected.
# Parameter:
#    $1: Namespace of the CRD
function run_update_canary_tests
{
  local crd_namespace="$1"
}

# Parameter:
#    $1: CRD namespace
function run_update_integration_tests
{
  echo "Running update integration tests"
  local crd_namespace="$1"
  run_update_canary_tests "${crd_namespace}"

  echo "Injecting variables into tests"
  inject_all_variables

  echo "Starting Update Tests"
  update_hap_test "${crd_namespace}" named-xgboost-hosting testfiles/xgboost-hostingautoscaling.yaml
}

# Verifies that all resources were created and are running/completed for the canary tests.
# Parameter:
#    $1: Namespace of the CRD
function verify_update_canary_tests
{
  local crd_namespace="$1"
  echo "Verifying update tests"
}

# Verifies that each integration update test has completed successfully.
# Parameter:
#    $1: CRD namespace
function verify_update_integration_tests
{
  echo "Verifying Update integration tests"
  local crd_namespace="$1"
  verify_update_canary_tests

  # At this point there are two variants in total(1 predefined and 1 custom) that have HAP applied.
  verify_test "${crd_namespace}" HostingAutoscalingPolicy hap-predefined 5m Created
  # verify_hap_test "2"
}

# Updates the ResourceID List and the MaxCapacity in the spec to check for updates
# Parameter:
#    $1: Target namespace
#    $2: K8s Name of the hostingdeployment to apply autoscaling
#    $3: Filename of the hap test to update
function update_hap_test()
{
  local target_namespace="$1"
  local hosting_deployment_1="$2"
  local file_name="$3"
  local hostingdeployment_type="hostingdeployment"
  local updated_filename="${file_name}-updated-${target_namespace}"

  # Copy and update the test file
  cp "${file_name}" "${updated_filename}"
  # Update the Resource ID list to remove one endpoint/variant
  yq d -i "${updated_filename}" "spec.resourceId[1]"
  yq w -i "$updated_filename" "spec.resourceId[0].endpointName" ${hosting_deployment_1}
  yq w -i "$updated_filename" "spec.maxCapacity" 3

  # HAP Test 1: Using the Predefined Metric
  run_test "$target_namespace" "$updated_filename"

  yq w -i testfiles/xgboost-hosting-deployment.yaml "metadata.name" ${hosting_deployment_1}
  rm "${updated_filename}"
}