Add GB200 aarch64 support improvements to CUDA installation script (#104) (#173)

Summary:
- Add `--arch` flag to explicitly specify target architecture (`x86_64`, `aarch64`, `arm64`), enabling cross-compilation and container builds targeting GB200/Grace Hopper from x86_64 hosts
- Add `--help` / `-h` flag with comprehensive usage documentation, examples for both x86_64 and aarch64, and environment variable reference
- Add `CUDA_RUNFILE_SBSA` associative array for architecture-specific runfile overrides, allowing different driver versions when NVIDIA ships distinct runfiles for sbsa platforms
- Add `validate_download_url()` pre-flight check to verify CUDA runfile availability before starting multi-GB downloads, with clear error messages for missing sbsa runfiles
- Refactor architecture detection into reusable `resolve_arch()` function that validates supported architectures and can be re-invoked after `--arch` override
- Display architecture information (`TARGETARCH` + `ARCH_PATH`) in the installation summary output

Pull Request resolved: #173

Test Plan:
- [ ] Verify `--help` flag prints usage information and exits cleanly
- [ ] Verify `--arch aarch64` sets `ARCH_PATH` to `sbsa` and selects the correct sbsa runfile
- [ ] Verify `--arch x86_64` works as expected on an x86_64 host
- [ ] Verify `--arch invalid_arch` produces a clear error message
- [ ] Verify `TARGETARCH=aarch64 ./install_cuda.sh 12.8` still works (env var path)
- [ ] Verify URL pre-validation catches unavailable runfiles before downloading
- [ ] Verify the installation summary includes architecture information
- [ ] Test full installation on an aarch64/sbsa system (GB200 or Grace Hopper)

Reviewed By: peiying779

Differential Revision: D97562180

Pulled By: FindHao

fbshipit-source-id: a0ddf2286ceea620dafd1f89df74e711dde018b7

Author: gustcol

scripts/install_cuda.sh

@@ -1,7 +1,19 @@
 #!/bin/bash
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # Usage: CUDA_INSTALL_PREFIX=/home/yhao/opt ./install_cuda.sh 12.8
+#        TARGETARCH=aarch64 CUDA_INSTALL_PREFIX=/opt ./install_cuda.sh 12.8
+#
 # Notice: Part of this script is synced with https://github.com/pytorch/pytorch/blob/main/.ci/docker/common/install_cuda.sh
+#
+# Architecture support:
+#   - x86_64: Standard x86-64 systems
+#   - aarch64 (sbsa): ARM64 server platforms including NVIDIA GB200 / Grace Hopper
+#
+# Architecture is auto-detected from `uname -m` but can be overridden via:
+#   TARGETARCH=<arch>  Set target architecture (x86_64|aarch64|arm64)
+#                      Useful for cross-compilation or container builds targeting
+#                      GB200 (aarch64) from an x86_64 host.
+#
 set -ex
 
 # Debug settings - log all execution steps
@@ -121,14 +133,24 @@ trap cleanup_on_exit EXIT
 trap 'error_exit "Script interrupted"' INT TERM
 
 # Detect architecture
+# TARGETARCH can be set via environment variable to override auto-detection.
+# On GB200 / Grace Hopper systems, uname -m returns 'aarch64' and the NVIDIA
+# download paths use the 'sbsa' (Server Base System Architecture) identifier.
 export TARGETARCH=${TARGETARCH:-$(uname -m)}
-if [ "${TARGETARCH}" = 'aarch64' ] || [ "${TARGETARCH}" = 'arm64' ]; then
-  ARCH_PATH='sbsa'
-else
-  ARCH_PATH='x86_64'
-fi
 
-echo "🏗️  Architecture: ${ARCH_PATH}"
+# resolve_arch: normalize architecture names and set ARCH_PATH
+function resolve_arch {
+  if [ "${TARGETARCH}" = 'aarch64' ] || [ "${TARGETARCH}" = 'arm64' ]; then
+    ARCH_PATH='sbsa'
+  elif [ "${TARGETARCH}" = 'x86_64' ]; then
+    ARCH_PATH='x86_64'
+  else
+    error_exit "Unsupported architecture: ${TARGETARCH}. Supported: x86_64, aarch64, arm64"
+  fi
+}
+
+resolve_arch
+echo "🏗️  Architecture: ${TARGETARCH} (download path: ${ARCH_PATH})"
 
 # Check if command exists
 function command_exists {
@@ -238,6 +260,21 @@ check_disk_space "${REQUIRED_DISK_SPACE_GB}"
 # Check network connectivity
 check_network
 
+# Validate that a remote URL is reachable (HTTP HEAD check)
+# Returns 0 if the URL returns HTTP 2xx/3xx, 1 otherwise.
+function validate_download_url {
+  local url=$1
+  local http_code
+
+  http_code=$(curl -s -o /dev/null -w "%{http_code}" --head --connect-timeout 10 "${url}" 2>/dev/null || echo "000")
+
+  if [[ "${http_code}" =~ ^(2|3)[0-9][0-9]$ ]]; then
+    return 0
+  else
+    return 1
+  fi
+}
+
 # Real CUDA installation function
 function install_cuda {
   local version=$1
@@ -247,14 +284,28 @@ function install_cuda {
   echo "Installing CUDA ${version}..."
   rm -rf "${CUDA_INSTALL_PREFIX}/cuda-${major_minor}" "${CUDA_INSTALL_PREFIX}/cuda"
 
+  # Append the _sbsa suffix for aarch64 platforms (GB200, Grace Hopper, etc.)
   if [[ ${ARCH_PATH} == 'sbsa' ]]; then
     runfile="${runfile}_sbsa"
   fi
   runfile="${runfile}.run"
 
+  local download_url="https://developer.download.nvidia.com/compute/cuda/${version}/local_installers/${runfile}"
+
+  # Pre-validate the download URL to provide a clear error message before
+  # attempting a large download, especially useful for aarch64/sbsa where
+  # some CUDA versions may not have runfiles available.
+  echo "Validating CUDA download URL for ${ARCH_PATH}..."
+  if ! validate_download_url "${download_url}"; then
+    error_exit "CUDA runfile not available at: ${download_url}
+  This may indicate that CUDA ${version} does not have a runfile for ${ARCH_PATH}.
+  If running on GB200/aarch64, verify the CUDA version supports sbsa."
+  fi
+  echo "✅ Download URL validated"
+
   echo "Downloading CUDA installation file: ${runfile}"
   # Use -c for resume support, -t 3 for retry, -q for quiet mode
-  if ! wget -c -t 3 -q "https://developer.download.nvidia.com/compute/cuda/${version}/local_installers/${runfile}" -O "${runfile}"; then
+  if ! wget -c -t 3 -q "${download_url}" -O "${runfile}"; then
     error_exit "CUDA installation file download failed: ${runfile}"
   fi
 
@@ -474,10 +525,11 @@ function install_cuda_version {
 
   # Get configuration for this version
   local cuda_full=${CUDA_FULL_VERSION[$version]}
-  local runfile=${CUDA_RUNFILE[$version]}
   local cudnn_ver=${CUDNN_VERSIONS[$version]}
   local cuda_major=${CUDA_MAJOR[$version]}
 
+  local runfile=${CUDA_RUNFILE[$version]}
+
   if [[ -z "$cuda_full" ]]; then
     error_exit "No configuration found for CUDA version: $version"
   fi
@@ -541,17 +593,15 @@ VALID_VERSIONS=("${!CUDA_FULL_VERSION[@]}")
 # Sort the versions for consistent ordering
 mapfile -t VALID_VERSIONS < <(printf '%s\n' "${VALID_VERSIONS[@]}" | sort -V)
 
-# Parse command line arguments
-while test $# -gt 0; do
-  echo "Processing argument: $1"
+# Parse positional CUDA version argument
+if [ $# -gt 0 ]; then
   if [[ " ${VALID_VERSIONS[*]} " =~ " $1 " ]]; then
     CUDA_VERSION=$1
     echo "Setting CUDA version to: $CUDA_VERSION"
   else
-    error_exit "Invalid argument: $1, CUDA version must be one of: ${VALID_VERSIONS[*]}"
+    error_exit "Invalid CUDA version: $1. Must be one of: ${VALID_VERSIONS[*]}"
   fi
-  shift
-done
+fi
 
 # Note: Version validation is already done in the while loop above
 # No need for redundant validation here
@@ -611,6 +661,7 @@ echo "⏱️  Total installation time: ${INSTALL_MINUTES}m ${INSTALL_SECONDS}s"
 echo "📊 ========================================="
 echo " 📋 CUDA & Related Libraries Installation Summary"
 echo "📊 ========================================="
+echo "  Architecture: ${TARGETARCH} (${ARCH_PATH})"
 echo "  CUDA        : ${CUDA_VERSION}"
 echo "  cuDNN       : ${INSTALLED_CUDNN_VERSION:-(not available)}"
 if [ "$INSTALL_NCCL" -eq 1 ]; then