fix(al2023/nvidia): allow nvidia grid install bucket customization

Author: mselim00

doc/usage/al2023.md

@@ -34,6 +34,7 @@
 | `nodeadm_build_image` | Image to use as a build environment for nodeadm |
 | `nvidia_driver_major_version` | To be used only when ```enable_accelerator = nvidia```. Driver version to install, depends on what is available in NVIDIA repository. |
 | `nvidia_repository_url` | YUM/DNF Repository override for the NVIDIA driver packages |
+| `nvidia_grid_runfile_bucket_name` | Bucket name for sourcing the Nvidia GRID runfile |
 | `pause_container_image` | Image ref for the pause container image |
 | `remote_folder` | Directory path for shell provisioner scripts on the builder instance |
 | `runc_version` |  |

templates/al2023/provisioners/install-nvidia-driver.sh

@@ -29,15 +29,7 @@ echo "Installing NVIDIA ${NVIDIA_DRIVER_MAJOR_VERSION} drivers..."
 # of the AMI, we want to ensure that all three kernel modules (and also the userspace modules)
 # are on the same NVIDIA driver version. Currently, the script installs the NVIDIA GRID drivers
 # first and decides the full NVIDIA driver version that the AMI will adhere to
-if [[ "$AWS_REGION" == "us-isof-south-1" || "$AWS_REGION" == "eusc-de-east-1" || "$AWS_REGION" == "eu-isoe-west-1" ]]; then
-  EC2_GRID_DRIVER_S3_BUCKET="${BINARY_BUCKET_NAME}"
-  GRID_DRIVER_S3_SCAN_PATH="s3://${EC2_GRID_DRIVER_S3_BUCKET}/bin/nvidia-grid-drivers/"
-else
-  EC2_GRID_DRIVER_S3_BUCKET="ec2-linux-nvidia-drivers"
-  GRID_DRIVER_S3_SCAN_PATH="s3://${EC2_GRID_DRIVER_S3_BUCKET}/"
-fi
-
-NVIDIA_DRIVER_FULL_VERSION=$(aws s3 ls --recursive ${GRID_DRIVER_S3_SCAN_PATH} \
+NVIDIA_DRIVER_FULL_VERSION=$(aws s3 ls --recursive ${EC2_GRID_DRIVER_S3_BUCKET} \
   | grep -Eo "(NVIDIA-Linux-x86_64-)${NVIDIA_DRIVER_MAJOR_VERSION}\.[0-9]+\.[0-9]+(-grid-aws\.run)" \
   | cut -d'-' -f4 \
   | sort -V \
@@ -174,7 +166,7 @@ function archive-open-kmods() {
 
 function archive-grid-kmod() {
   local MACHINE
-  local NVIDIA_GRID_RUNFILE_NAME
+  local NVIDIA_GRID_RUNFILE_KEY
   local GRID_INSTALLATION_TEMP_DIR
   local EXTRACT_DIR
 
@@ -187,26 +179,28 @@ function archive-grid-kmod() {
   fi
 
   echo "Archiving NVIDIA GRID kernel modules for major version ${NVIDIA_DRIVER_MAJOR_VERSION}"
-  NVIDIA_GRID_RUNFILE_NAME=$(aws s3 ls --recursive ${GRID_DRIVER_S3_SCAN_PATH} \
+  NVIDIA_GRID_RUNFILE_KEY=$(aws s3 ls --recursive ${EC2_GRID_DRIVER_S3_BUCKET} \
     | grep "NVIDIA-Linux-x86_64-${NVIDIA_DRIVER_FULL_VERSION}" \
     | sort -k1,2 \
     | tail -1 \
     | awk '{print $4}')
 
-  if [[ -z "$NVIDIA_GRID_RUNFILE_NAME" ]]; then
+  if [[ -z "$NVIDIA_GRID_RUNFILE_KEY" ]]; then
     echo "ERROR: No GRID driver found for driver version ${NVIDIA_DRIVER_FULL_VERSION} in EC2 S3 bucket"
     exit 1
   fi
 
-  echo "Found GRID runfile: ${NVIDIA_GRID_RUNFILE_NAME}"
-  local GRID_RUNFILE_LOCAL_NAME
-  GRID_RUNFILE_LOCAL_NAME=$(basename "${NVIDIA_GRID_RUNFILE_NAME}")
+  echo "Found GRID runfile: ${NVIDIA_GRID_RUNFILE_KEY}"
+  local GRID_RUNFILE_NAME
+  GRID_RUNFILE_NAME=$(basename "${NVIDIA_GRID_RUNFILE_KEY}")
 
   echo "Downloading GRID driver runfile..."
-  aws s3 cp "s3://${EC2_GRID_DRIVER_S3_BUCKET}/${NVIDIA_GRID_RUNFILE_NAME}" "${GRID_INSTALLATION_TEMP_DIR}/${GRID_RUNFILE_LOCAL_NAME}"
-  chmod +x "${GRID_INSTALLATION_TEMP_DIR}/${GRID_RUNFILE_LOCAL_NAME}"
+  # This is the only command that requires the bucket name to actually just be the bucket (no prefix) b/c of how the
+  # s3 ls recursive output dumps the full object key regardless of the supplied prefix
+  aws s3 cp "s3://${EC2_GRID_DRIVER_S3_BUCKET%%/*}/${NVIDIA_GRID_RUNFILE_KEY}" "${GRID_INSTALLATION_TEMP_DIR}/${GRID_RUNFILE_NAME}"
+  chmod +x "${GRID_INSTALLATION_TEMP_DIR}/${GRID_RUNFILE_NAME}"
   echo "Extracting NVIDIA GRID driver runfile..."
-  sudo "${GRID_INSTALLATION_TEMP_DIR}/${GRID_RUNFILE_LOCAL_NAME}" --extract-only --target "${EXTRACT_DIR}"
+  sudo "${GRID_INSTALLATION_TEMP_DIR}/${GRID_RUNFILE_NAME}" --extract-only --target "${EXTRACT_DIR}"
 
   pushd "${EXTRACT_DIR}"
 

templates/al2023/template.json

@@ -30,6 +30,7 @@
     "nodeadm_build_image": null,
     "nvidia_driver_major_version": null,
     "nvidia_repository_url": null,
+    "nvidia_grid_runfile_bucket_name": null,
     "pause_container_image": null,
     "remote_folder": null,
     "runc_version": null,
@@ -267,6 +268,7 @@
         "BINARY_BUCKET_REGION={{user `binary_bucket_region`}}",
         "NVIDIA_DRIVER_MAJOR_VERSION={{user `nvidia_driver_major_version`}}",
         "NVIDIA_REPOSITORY={{user `nvidia_repository_url`}}",
+        "EC2_GRID_DRIVER_S3_BUCKET={{user `nvidia_grid_runfile_bucket_name`}}",
         "WORKING_DIR={{user `working_dir`}}"
       ]
     },
@@ -338,4 +340,4 @@
       }
     }
   ]
-}
+}

templates/al2023/variables-default.json

@@ -23,6 +23,7 @@
     "nodeadm_build_image": "public.ecr.aws/eks-distro-build-tooling/golang:1.25",
     "nvidia_driver_major_version": "580",
     "nvidia_repository_url": null,
+    "nvidia_grid_runfile_bucket_name": "ec2-linux-nvidia-drivers",
     "pause_container_image": "602401143452.dkr.ecr.us-west-2.amazonaws.com/eks/pause:3.10",
     "remote_folder": "/tmp",
     "runc_version": "*",