cerebrium-deploy

#!/usr/bin/env bash

# To download this script run:
# curl -L -o cerebrium-deploy "https://gist.github.com/jonoirwinrsa/dae834757905a63cfeaf0abce3bd7b5c/raw/cerebrium_deploy" && chmod +x cerebrium-deploy && sudo mv cerebrium-deploy /usr/local/bin

cerebrium_version="0.1.3"
echo "Cerebrium Builder v$cerebrium_version"

ENV=${CEREBRIUM_ENV:prod}

if [ "$ENV" = "dev" ]; then
  dashboard_url="https://dev-dashboard.cerebrium.ai"
  api_url="https://dev-rest-api.cerebrium.ai"
  echo Environment: "$ENV"
else
  dashboard_url="https://dashboard.cerebrium.ai"
  api_url="https://rest-api.cerebrium.ai"
fi

key=""
name=""
hardware=""
cpu=""
memory=""

while getopts ":k:n:h:c:m:" opt; do
  case $opt in
  k)
    key="$OPTARG"
    ;;
  n)
    name="$OPTARG"
    ;;
  h)
    hardware="$OPTARG"
    ;;
  c)
    cpu="$OPTARG"
    ;;
  m)
    memory="$OPTARG"
    ;;
  \?)
    echo "Invalid option: -$OPTARG" >&2
    exit 1
    ;;
  :)
    echo "Option -$OPTARG requires an argument." >&2
    exit 1
    ;;
  esac
done

if [[ -z "$key" || -z "$name" ]]; then
  echo "Usage: $0 -k <api_key> -n <name>"
  exit 1
fi

api_key="$key"
name="$name"

echo "Deploying $name to Cerebrium..."

echo "Authorising..."

# Call API endpoint and check response status code
response=$(curl -s -w "\n%{http_code}" -H "Authorization: $api_key" "$api_url/builder-deploy?name=$name&hardware=$hardware&cpu=$cpu&memory=$memory&cerebrium_version=$cerebrium_version")
response_status_code=$(echo "$response" | sed -n '$p')
response_body=$(echo "$response" | sed '$d')

if [ "$response_status_code" -ne 200 ]; then
  echo "Error: API request failed with status code $response_status_code"
  echo "$response_body"
  exit 1
fi

# Parse JSON response to get signed S3 URL and project ID
signed_s3_url=$(echo "$response_body" | jq -r '.uploadUrl')
project_id=$(echo "$response_body" | jq -r '.projectId')
zip_file_name=$(echo "$response_body" | jq -r '.keyName')
endpoint=$(echo "$response_body" | jq -r '.internalEndpoint')
buildId=$(echo "$response_body" | jq -r '.buildId')

echo "You can view your project at $dashboard_url/projects/$project_id"

# Check if jq is installed
if [ -z "$signed_s3_url" ] || [ -z "$project_id" ] || [ -z "$zip_file_name" ]; then
  echo "Please make sure jq is installed: https://stedolan.github.io/jq/download/"
  exit 1
fi

echo "Zipping files..."

zip -r "$zip_file_name" . -x "*/.*/*"

echo "Uploading files..."

# Upload zip file to S3 and check response status code
upload_response=$(curl -s -w "\n%{http_code}" -X PUT -T "${zip_file_name}" -H "Content-Type: application/zip" "${signed_s3_url}")
upload_status_code=$(echo "$upload_response" | sed -n '$p')
upload_response_body=$(echo "$$upload_response" | sed '$d')

if [ "$upload_status_code" -ne 200 ]; then
  echo "Error: Upload to S3 failed with status code $upload_status_code"
  echo "Response body: $upload_response_body"
  exit 1
fi

# Clean up
rm "$zip_file_name"

# Function to poll an endpoint for a maximum of 2 minutes until the response body.status is "deployed"
poll_endpoint() {
  local url=$1
  local timeout=600 # 10 minutes in seconds
  local elapsed_time=0
  local poll_interval=5 # Adjust the sleep time (in seconds) between polling attempts as needed
  local status_code=0
  local status=""
  printed_created_timestamps=()

  while [ "$status" != "deployed" ] && [ $elapsed_time -lt $timeout ]; do
    response=$(curl -s -w "%{http_code}" -H "Authorization: $api_key" -H "Content-Type: application/json" "$url")
    status_code=${response: -3}
    body=${response%???}

    if [ "$status_code" -ne 200 ]; then
      echo "Error: Received a non-200 status code ($status_code)"
      return 1
    fi

    status=$(echo "$body" | jq -r '.status')
    logs=$(echo "$body" | jq -r '.logs')

    created_and_messages=$(echo "$logs" | jq -r '.[] | "\(.created)|\(.message)"')

    # Log each message if its created timestamp hasn't been printed before
    while read -r created_and_message; do
      created="${created_and_message%%|*}"
      message="${created_and_message#*|}"
      if ! printf '%s\n' "${printed_created_timestamps[@]}" | grep -q -F -x "$created"; then
        echo "$message"
        printed_created_timestamps+=("$created")
      fi
    done <<<"$created_and_messages"

    if [ "$status" != "deployed" ] || [ "$status" != "deployed" ]; then
      sleep $poll_interval
      elapsed_time=$((elapsed_time + poll_interval))
    fi
  done

  if [ "$status" == "deployed" ]; then
    echo "Endpoint status is 'deployed'."
    echo "Please note that it can take a few minutes to download model weights depending on the size of your model."
    echo "-------------------------------------------------------"
    echo "View stats: $dashboard_url/projects/$project_id/models/$project_id-$name?tab=overview"
    echo "View builds: $dashboard_url/projects/$project_id/models/$project_id-$name?tab=builds"
    echo "View runs: $dashboard_url/projects/$project_id/models/$project_id-$name?tab=runs"
    echo "-------------------------------------------------------"
    echo -e "Model callable using the following curl command:\n\
    curl --request POST --location '$endpoint' \\\\\n\
    --header 'Authorization: $api_key' \\\\\n\
    --header 'Content-Type: application/json' \\\\\n\
    --data '{}'"
  elif [ "$status" == "failed" ]; then
    echo "Deployment failed'."
  else
    echo "Polling timed out after 10 minutes."
  fi
}

poll_endpoint $api_url/streamBuildLogs?buildId="$buildId"