sagemaker-presigned-url-v2.sh

sagemaker-presigned-url-v2.sh

#!/bin/bash

NOTEBOOK_INSTANCE_NAME="id00045V03"
INSTANCE_TYPE="ml.g5.xlarge"
INSTANCE_VOLUME="500"
AWS_REGION="us-east-1"
# Change
ROLE_ARN="arn:aws:iam::YOUR_AWS_ACCOUNT_ID:role/YOUR_SAGEMAKER_ROLE"

create_instance() {
  # For Code Server Setting
  curl -LO https://github.com/aws-samples/amazon-sagemaker-codeserver/releases/download/v0.2.0/amazon-sagemaker-codeserver-0.2.0.tar.gz
  tar -xvzf amazon-sagemaker-codeserver-0.2.0.tar.gz
  cd amazon-sagemaker-codeserver/install-scripts/notebook-instances

  aws sagemaker create-notebook-instance-lifecycle-config \
    --notebook-instance-lifecycle-config-name install-codeserver \
    --on-start Content="$((cat setup-codeserver.sh || echo '')| base64)" \
    --on-create Content="$((cat install-codeserver.sh || echo '')| base64)"
  aws sagemaker create-notebook-instance \
    --notebook-instance-name "$NOTEBOOK_INSTANCE_NAME" \
    --instance-type "$INSTANCE_TYPE" \
    --role-arn "$ROLE_ARN" \
    --direct-internet-access Enabled \
    --volume-size-in-gb 5 \
    --root-access Enabled \
    --platform-identifier notebook-al2-v2 \
    --region "$AWS_REGION" \
    --lifecycle-config-name install-codeserver

  return $?
}

check_instance_status() {
  local status=$(aws sagemaker describe-notebook-instance \
    --notebook-instance-name "$NOTEBOOK_INSTANCE_NAME" \
    --output text \
    --query 'NotebookInstanceStatus' \
    --region "$AWS_REGION" 2>/dev/null)

  if [ -z "$status" ]; then
    if [ "$1" == "--create" ]; then
      create_instance
      wait_for_instance_ready
      return $?
    else
      echo "Notebook Instance '$NOTEBOOK_INSTANCE_NAME' は存在しません。"
      return 1
    fi
  elif [ "$status" != "InService" ]; then
    echo "Notebook Instance '$NOTEBOOK_INSTANCE_NAME' は利用可能な状態ではありません。(ステータス: $status)"
    return 1
  fi

  return 0
}

wait_for_instance_ready() {
  local max_attempts=60
  local attempt=0

  while [ $attempt -lt $max_attempts ]; do
    local status=$(aws sagemaker describe-notebook-instance \
      --notebook-instance-name "$NOTEBOOK_INSTANCE_NAME" \
      --output text \
      --query 'NotebookInstanceStatus' \
      --region "$AWS_REGION" 2>/dev/null)

    if [ "$status" == "InService" ]; then
      return 0
    fi

    echo "Notebook Instance '$NOTEBOOK_INSTANCE_NAME' はまだ準備中です。(ステータス: $status) 10秒後に再試行します。"
    sleep 10
    ((attempt++))
  done

  echo "Notebook Instance '$NOTEBOOK_INSTANCE_NAME' が利用可能な状態になるのを待ちきれませんでした。"
  return 1
}

create_presigned_url() {
  PRESIGNED_URL=$(aws sagemaker create-presigned-notebook-instance-url --notebook-instance-name $NOTEBOOK_INSTANCE_NAME --query 'AuthorizedUrl' --output text --session-expiration-duration-in-seconds 3600)
  echo "URL:"
  echo "$PRESIGNED_URL" |tee .presignedURL
}

check_instance_status "$1"
if [ $? -eq 0 ]; then
  create_presigned_url
else
  exit 1
fi