Enable ARM instances on Cortex (#2268)

Author: RobertLucian

.circleci/config.yml

@@ -121,7 +121,7 @@ jobs:
       - checkout
       - setup_remote_docker
       - install-go
-      - run: sudo apt update && sudo apt install parallel -y
+      - run: docker buildx create --name builder --platform linux/amd64,linux/arm64 --use && docker buildx inspect --bootstrap
       - run: sudo pip install awscli
       - run: make ci-build-cli
       - return-if-not-deployed-branch
@@ -174,6 +174,10 @@ jobs:
                 instance_type: inf1.xlarge
                 min_instances: 1
                 max_instances: 2
+              - name: arm
+                instance_type: a1.large
+                min_instances: 1
+                max_instances: 2
             EOF
       - run:
           name: Create/Update AWS User policy

Makefile

@@ -127,17 +127,24 @@ async-gateway-update:
 # docker images
 images-all:
 	@./dev/registry.sh update all
+images-all-multi-arch:
+	@./dev/registry.sh update all --include-arm64-arch
 images-all-skip-push:
 	@./dev/registry.sh update all --skip-push
 
 images-dev:
 	@./dev/registry.sh update dev
+images-dev-multi-arch:
+	@./dev/registry.sh update dev --include-arm64-arch
 images-dev-skip-push:
 	@./dev/registry.sh update dev --skip-push
 
 images-manager-skip-push:
 	@./dev/registry.sh update-single manager --skip-push
 
+images-clean-cache:
+	@./dev/registry.sh clean-cache
+
 registry-create:
 	@./dev/registry.sh create
 
@@ -151,7 +158,6 @@ tools:
 	@go get -u -v github.com/kyoh86/looppointer/cmd/looppointer
 	@go get -u -v github.com/VojtechVitek/rerun/cmd/rerun
 	@go get -u -v github.com/go-delve/delve/cmd/dlv
-	@if [[ "$$OSTYPE" == "darwin"* ]]; then brew install parallel; elif [[ "$$OSTYPE" == "linux"* ]]; then sudo apt-get install -y parallel; else echo "your operating system is not supported"; fi
 	@python3 -m pip install aiohttp black 'pydoc-markdown>=3.0.0,<4.0.0' boto3 pyyaml
 	@python3 -m pip install -e test/e2e
 

build/build-image.sh

@@ -22,8 +22,9 @@ ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")"/.. >/dev/null && pwd)"
 CORTEX_VERSION=master
 
 image=$1
+platforms=$2
 
 if [ "$image" == "inferentia" ]; then
   aws ecr get-login-password --region us-west-2 | docker login --username AWS --password-stdin 790709498068.dkr.ecr.us-west-2.amazonaws.com
 fi
-docker build "$ROOT" -f $ROOT/images/$image/Dockerfile -t quay.io/cortexlabs/${image}:${CORTEX_VERSION} -t cortexlabs/${image}:${CORTEX_VERSION}
+docker buildx build $ROOT -f $ROOT/images/$image/Dockerfile -t quay.io/cortexlabs/${image}:${CORTEX_VERSION} -t cortexlabs/${image}:${CORTEX_VERSION} --platform $platforms

build/build-images.sh

@@ -22,11 +22,10 @@ ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")"/.. >/dev/null && pwd)"
 source $ROOT/build/images.sh
 source $ROOT/dev/util.sh
 
-# if parallel utility is installed, the docker build commands will be parallelized
-if command -v parallel &> /dev/null && [ -n "${NUM_BUILD_PROCS+set}" ] && [ "$NUM_BUILD_PROCS" != "1" ]; then
-  ROOT=$ROOT SHELL=$(type -p /bin/bash) parallel --will-cite --halt now,fail=1 --eta --jobs $NUM_BUILD_PROCS $ROOT/build/build-image.sh {} ::: "${all_images[@]}"
-else
-  for image in "${all_images[@]}"; do
-    $ROOT/build/build-image.sh $image
-  done
-fi
+for image in "${all_images[@]}"; do
+  platforms="linux/amd64"
+  if in_array $image "multi_arch_images"; then
+    platforms+=",linux/arm64"
+  fi
+  $ROOT/build/build-image.sh $image $platforms
+done

build/generate_ami_mapping.go

@@ -235,11 +235,15 @@ func main() {
 		fmt.Print(region)
 		sess := session.New(&aws.Config{Region: aws.String(region)})
 		svc := ec2.New(sess)
-		cpuAMI, err := FindImage(svc, EKSResourceAccountID(region), fmt.Sprintf("amazon-eks-node-%s-v*", k8sVersion))
+		cpuAmd64AMI, err := FindImage(svc, EKSResourceAccountID(region), fmt.Sprintf("amazon-eks-node-%s-v*", k8sVersion))
 		if err != nil {
 			log.Fatal(err.Error())
 		}
-		acceleratedAMI, err := FindImage(svc, EKSResourceAccountID(region), fmt.Sprintf("amazon-eks-gpu-node-%s-v*", k8sVersion))
+		cpuArm64AMI, err := FindImage(svc, EKSResourceAccountID(region), fmt.Sprintf("amazon-eks-arm64-node-%s-v*", k8sVersion))
+		if err != nil {
+			log.Fatal(err.Error())
+		}
+		acceleratedAmd64AMI, err := FindImage(svc, EKSResourceAccountID(region), fmt.Sprintf("amazon-eks-gpu-node-%s-v*", k8sVersion))
 		if err != nil {
 			log.Fatal(err.Error())
 		}
@@ -248,8 +252,9 @@ func main() {
 			k8sVersionMap[k8sVersion][region] = map[string]string{}
 		}
 		k8sVersionMap[k8sVersion][region] = map[string]string{
-			"cpu":         cpuAMI,
-			"accelerated": acceleratedAMI,
+			"cpu_amd64":         cpuAmd64AMI,
+			"cpu_arm64":         cpuArm64AMI,
+			"accelerated_amd64": acceleratedAmd64AMI,
 		}
 		fmt.Println(" ✓")
 	}

build/images.sh

@@ -50,6 +50,18 @@ non_dev_images=(
   "kubexit"
 )
 
+# for linux/amd64 and linux/arm64
+multi_arch_images=(
+  "proxy"
+  "async-gateway"
+  "enqueuer"
+  "dequeuer"
+  "fluent-bit"
+  "prometheus-node-exporter"
+  "kube-rbac-proxy"
+  "kubexit"
+)
+
 all_images=(
   "${dev_images[@]}"
   "${non_dev_images[@]}"

build/push-image.sh

@@ -17,10 +17,13 @@
 
 set -euo pipefail
 
+ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")"/.. >/dev/null && pwd)"
+
 CORTEX_VERSION=master
 
 host=$1
 image=$2
+platforms=$3
 
 echo "$DOCKER_PASSWORD" | docker login -u "$DOCKER_USERNAME" --password-stdin
-docker push $host/cortexlabs/${image}:${CORTEX_VERSION}
+docker buildx build $ROOT -f $ROOT/images/$image/Dockerfile $host/cortexlabs/${image}:${CORTEX_VERSION} --platform $platforms --push

build/push-images.sh

@@ -24,11 +24,10 @@ source $ROOT/dev/util.sh
 
 host=$1
 
-# if parallel utility is installed, the docker push commands will be parallelized
-if command -v parallel &> /dev/null && [ -n "${NUM_BUILD_PROCS+set}" ] && [ "$NUM_BUILD_PROCS" != "1" ]; then
-  ROOT=$ROOT DOCKER_USERNAME=$DOCKER_USERNAME DOCKER_PASSWORD=$DOCKER_PASSWORD SHELL=$(type -p /bin/bash) parallel --will-cite --halt now,fail=1 --eta --jobs $NUM_BUILD_PROCS $ROOT/build/push-image.sh $host {} ::: "${all_images[@]}"
-else
-  for image in "${all_images[@]}"; do
-    $ROOT/build/push-image.sh $host $image
-  done
-fi
+for image in "${all_images[@]}"; do
+  platforms="linux/amd64"
+  if in_array $image "multi_arch_images"; then
+    platforms+=",linux/arm64"
+  fi
+  $ROOT/build/push-image.sh $host $image $platforms
+done

dev/registry.sh

@@ -23,7 +23,7 @@ ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")"/.. >/dev/null && pwd)"
 source $ROOT/build/images.sh
 source $ROOT/dev/util.sh
 
-images_with_builders="operator proxy async-gateway enqueuer dequeuer controller-manager"
+images_that_can_run_locally="operator manager"
 
 if [ -f "$ROOT/dev/config/env.sh" ]; then
   source $ROOT/dev/config/env.sh
@@ -33,6 +33,7 @@ AWS_ACCOUNT_ID=${AWS_ACCOUNT_ID:-}
 AWS_REGION=${AWS_REGION:-}
 
 skip_push="false"
+include_arm64_arch="false"
 positional_args=()
 while [[ $# -gt 0 ]]; do
   key="$1"
@@ -41,6 +42,10 @@ while [[ $# -gt 0 ]]; do
     skip_push="true"
     shift
     ;;
+    --include-arm64-arch)
+    include_arm64_arch="true"
+    shift
+    ;;
     *)
     positional_args+=("$1")
     shift
@@ -105,59 +110,55 @@ function create_ecr_repository() {
 
 ### HELPERS ###
 
-function build() {
-  local image=$1
-  local tag=$2
-  local dir="${ROOT}/images/${image}"
-
-  tag_args=""
-  if [ -n "$AWS_ACCOUNT_ID" ] && [ -n "$AWS_REGION" ]; then
-    tag_args+=" -t $AWS_ACCOUNT_ID.dkr.ecr.$AWS_REGION.amazonaws.com/cortexlabs/$image:$tag"
-  fi
-
-  blue_echo "Building $image:$tag..."
-  docker build $ROOT -f $dir/Dockerfile -t cortexlabs/$image:$tag $tag_args
-  green_echo "Built $image:$tag\n"
-}
-
-function cache_builder() {
+function build_and_push() {
   local image=$1
+  local include_arm64_arch=$2
   local dir="${ROOT}/images/${image}"
 
-  blue_echo "Building $image-builder..."
-  docker build $ROOT -f $dir/Dockerfile -t cortexlabs/$image-builder:$CORTEX_VERSION --target builder
-  green_echo "Built $image-builder\n"
-}
+  set -euo pipefail
 
-function push() {
-  if [ "$skip_push" = "true" ]; then
-    return
+  if ! in_array $image "multi_arch_images"; then
+    include_arm64_arch="false"
   fi
 
+  if [ ! -n "$AWS_ACCOUNT_ID" ] || [ ! -n "$AWS_REGION" ]; then
+    echo "AWS_ACCOUNT_ID or AWS_REGION env vars not found"
+    exit 1
+  fi
   registry_login
 
-  local image=$1
-  local tag=$2
+  tag=$CORTEX_VERSION
+  if [ "$include_arm64_arch" = "true" ]; then
+    blue_echo "Building and pushing $image:$tag (amd64 and arm64)..."
+  else
+    blue_echo "Building and pushing $image:$tag (amd64)..."
+  fi
 
-  blue_echo "Pushing $image:$tag..."
-  docker push $registry_push_url/cortexlabs/$image:$tag
-  green_echo "Pushed $image:$tag\n"
-}
+  platforms="linux/amd64"
+  if [ "$include_arm64_arch" = "true" ]; then
+    platforms+=",linux/arm64"
+  fi
 
-function build_and_push() {
-  local image=$1
+  docker buildx build $ROOT -f $dir/Dockerfile -t $AWS_ACCOUNT_ID.dkr.ecr.$AWS_REGION.amazonaws.com/cortexlabs/$image:$tag --platform $platforms --push
 
-  set -euo pipefail  # necessary since this is called in a new shell by parallel
+  if [ "$include_arm64_arch" = "true" ]; then
+    green_echo "Built and pushed $image:$tag (amd64 and arm64)..."
+  else
+    green_echo "Built and pushed $image:$tag (amd64)..."
+  fi
 
-  tag=$CORTEX_VERSION
-  build $image $tag
-  push $image $tag
+  if [[ " $images_that_can_run_locally " =~ " $image " ]] && [[ "$include_arm64_arch" == "false" ]]; then
+    blue_echo "Exporting $image:$tag to local docker..."
+    docker buildx build $ROOT -f $dir/Dockerfile -t cortexlabs/$image:$tag -t $AWS_ACCOUNT_ID.dkr.ecr.$AWS_REGION.amazonaws.com/cortexlabs/$image:$tag --platform $platforms --load
+    green_echo "Exported $image:$tag to local docker..."
+  fi
 }
 
 function cleanup_local() {
   echo "cleaning local repositories..."
   docker container prune -f
   docker image prune -f
+  docker buildx prune -f
 }
 
 function cleanup_ecr() {
@@ -195,14 +196,6 @@ function validate_env() {
   fi
 }
 
-# export functions for parallel command
-export -f build_and_push
-export -f push
-export -f build
-export -f blue_echo
-export -f green_echo
-export -f registry_login
-
 # validate environment is correctly set on env.sh
 validate_env
 
@@ -218,13 +211,9 @@ elif [ "$cmd" = "create" ]; then
 # usage: registry.sh update-single IMAGE
 elif [ "$cmd" = "update-single" ]; then
   image=$sub_cmd
-  if [[ " $images_with_builders " =~ " $image " ]]; then
-    cache_builder $image
-  fi
-  build_and_push $image
+  build_and_push $image $include_arm64_arch
 
 # usage: registry.sh update all|dev|api
-# if parallel utility is installed, the docker build commands will be parallelized
 elif [ "$cmd" = "update" ]; then
   images_to_build=()
 
@@ -236,20 +225,12 @@ elif [ "$cmd" = "update" ]; then
     images_to_build+=( "${dev_images[@]}" )
   fi
 
-  for image in $images_with_builders; do
-    if [[ " ${images_to_build[@]} " =~ " $image " ]]; then
-      cache_builder $image
-    fi
+  for image in "${images_to_build[@]}"; do
+    build_and_push $image $include_arm64_arch
   done
 
-  if command -v parallel &> /dev/null && [ -n "${NUM_BUILD_PROCS+set}" ] && [ "$NUM_BUILD_PROCS" != "1" ]; then
-    is_registry_logged_in=$is_registry_logged_in ROOT=$ROOT registry_push_url=$registry_push_url SHELL=$(type -p /bin/bash) parallel --will-cite --halt now,fail=1 --eta --jobs $NUM_BUILD_PROCS build_and_push "{}" ::: "${images_to_build[@]}"
-  else
-    for image in "${images_to_build[@]}"; do
-      build_and_push $image
-    done
-  fi
-
+# usage: registry.sh clean-cache
+elif [ "$cmd" = "clean-cache" ]; then
   cleanup_local
 
 else

dev/util.sh

@@ -45,3 +45,13 @@ function join_by() {
   done
   echo
 }
+
+# Check if array contains item [$1: item, $2: array name]
+function in_array() {
+    local needle="$1" item
+    local -n arrref="$2"
+    for item in "${arrref[@]}"; do
+        [[ "${item}" == "${needle}" ]] && return 0
+    done
+    return 1
+}