feat: add neuron and efa device plugin updater (#643)

* feat: add neuron and efa device plugin updater

* Update hack/update-image-tags.sh

Co-authored-by: Nick Baker <[email protected]>

---------

Co-authored-by: Nick Baker <[email protected]>

Author: mselim00

Makefile

@@ -3,3 +3,9 @@ include ${BGO_MAKEFILE}
 pre-release::
 	go test -c -tags=e2e ./test/... -o $(GOBIN)
 	go install sigs.k8s.io/kubetest2/...@latest
+
+update-deps:
+	for SCRIPT in ./hack/update-*.sh; do \
+    	"$$SCRIPT" ; \
+	done
+

hack/update-image-tags.sh

@@ -4,14 +4,61 @@ set -o nounset
 set -o errexit
 set -o pipefail
 
+ECR_PUBLIC_REGISTRY="public.ecr.aws"
+EKS_CONTAINER_REGISTRY="602401143452.dkr.ecr.us-west-2.amazonaws.com"
+
+# get_ecr_image_tags <REGISTRY> <REPOSITORY>
+# e.g. get_ecr_image_tags $ECR_PUBLIC_REGISTRY amazonlinux/amazonlinux
+get_ecr_image_tags() {
+    set -e
+    local REGISTRY=$1 
+    local REPOSITORY=$2
+    local TOKEN
+
+    # Get ECR public token if image is from a public registry, otherwise use a private token
+    # An authorization token is required for every ECR HTTP request
+    if [ "$REGISTRY" = "$ECR_PUBLIC_REGISTRY" ]; then
+        TOKEN=$(aws ecr-public get-authorization-token --region us-east-1 --output=text --query 'authorizationData.authorizationToken')
+        local AUTHORIZATION_TYPE="Bearer"
+    else 
+        TOKEN=$(aws ecr get-authorization-token --output text --query 'authorizationData[].authorizationToken')
+        local AUTHORIZATION_TYPE="Basic"
+    fi
+
+    curl -s -H "Authorization: ${AUTHORIZATION_TYPE} $TOKEN" "https://$REGISTRY/v2/$REPOSITORY/tags/list" | jq '.tags'
+}
+
+# update_image_uris REPOSITORY IMAGE_TAG
+update_image_uris() {
+    local REPOSITORY=$1
+    local NEW_TAG=$2
+    PREFIX="image: ${REPOSITORY}"
+    find ./test/manifests -type f -exec sed -i "s#$PREFIX:.*#$PREFIX:$NEW_TAG#g" {} +
+}
+
 # update the nvidia k8s device plugin
+echo "Updating Nvidia device plugin image"
+NVIDIA_DEVICE_PLUGIN_TAG=$(curl -s 'https://catalog.ngc.nvidia.com/api/containers/images?orgName=nvidia&name=k8s-device-plugin&isPublic=true' | jq -r '.images | sort_by(.updatedDate) | reverse | map(select(.tag | test("^v[0-9]+.[0-9]+.[0-9]+$"))) | first | .tag')
+update_image_uris nvcr.io/nvidia/k8s-device-plugin $NVIDIA_DEVICE_PLUGIN_TAG
 
-NVIDIA_DEVICE_PLUGIN=$(curl -s 'https://catalog.ngc.nvidia.com/api/containers/images?orgName=nvidia&name=k8s-device-plugin&isPublic=true' | jq -r '.images | sort_by(.updatedDate) | reverse | map(select(.tag | test("^v[0-9]+.[0-9]+.[0-9]+$"))) | first | .tag')
-PREFIX="image: nvcr.io/nvidia/k8s-device-plugin"
-find ./test/manifests -type f -exec sed -i "s#$PREFIX:.*#$PREFIX:$NVIDIA_DEVICE_PLUGIN#g" {} +
+# below updates require authentication and should not exit early with a failure.
+# TODO: remove this once the aws credentials are setup and the paths are expected to succeed.
+set +e
 
 # update the neuron k8s device plugin
-# TODO
+echo "Updating Neuron device plugin image"
+NEURON_DEVICE_PLUGIN_REPOSITORY_NAME="neuron/neuron-device-plugin"
+NEURON_DEVICE_PLUGIN_TAGS=$(get_ecr_image_tags $ECR_PUBLIC_REGISTRY $NEURON_DEVICE_PLUGIN_REPOSITORY_NAME)
+if [ $? -eq 0 ]; then
+    LATEST_NEURON_DEVICE_PLUGIN_TAG=$(echo $NEURON_DEVICE_PLUGIN_TAGS | jq -r 'max_by(split(".") | map(tonumber))')
+    update_image_uris "${ECR_PUBLIC_REGISTRY}/${NEURON_DEVICE_PLUGIN_REPOSITORY_NAME}" $LATEST_NEURON_DEVICE_PLUGIN_TAG
+fi
 
 # update the efa k8s device plugin
-# TODO
+echo "Updating EFA device plugin image"
+EFA_DEVICE_PLUGIN_REPOSITORY_NAME="eks/aws-efa-k8s-device-plugin"
+EFA_DEVICE_PLUGIN_TAGS=$(get_ecr_image_tags $EKS_CONTAINER_REGISTRY $EFA_DEVICE_PLUGIN_REPOSITORY_NAME)
+if [ $? -eq 0 ]; then
+    LATEST_EFA_DEVICE_PLUGIN_TAG=$(echo $EFA_DEVICE_PLUGIN_TAGS | jq -r 'map(split("-") | .[0]) | max_by(sub("^v"; "") | split(".") | map(tonumber))')
+    update_image_uris "${EKS_CONTAINER_REGISTRY}/${EFA_DEVICE_PLUGIN_REPOSITORY_NAME}" $LATEST_EFA_DEVICE_PLUGIN_TAG
+fi

test/manifests/assets/efa-device-plugin.yaml

@@ -29,7 +29,7 @@ spec:
       priorityClassName: "system-node-critical"
       hostNetwork: true
       containers:
-        - image: 602401143452.dkr.ecr.us-west-2.amazonaws.com/eks/aws-efa-k8s-device-plugin:v0.5.6
+        - image: 602401143452.dkr.ecr.us-west-2.amazonaws.com/eks/aws-efa-k8s-device-plugin:v0.5.8
           name: aws-efa-k8s-device-plugin
           securityContext:
             allowPrivilegeEscalation: false

test/manifests/assets/k8s-neuron-device-plugin.yml

@@ -53,7 +53,7 @@ spec:
                       - trn2u.48xlarge
       containers:
         # Find all neuron-device-plugin images at https://gallery.ecr.aws/neuron/neuron-device-plugin
-      - image: public.ecr.aws/neuron/neuron-device-plugin:2.25.24.0
+      - image: public.ecr.aws/neuron/neuron-device-plugin:2.26.26.0
         imagePullPolicy: Always
         name: neuron-device-plugin
         env: