From 167a46128ae18a5778cecd8e5452c9abe93160aa Mon Sep 17 00:00:00 2001
From: Ting Lu <tingl@nvidia.com>
Date: Sat, 18 May 2024 20:20:37 -0700
Subject: [PATCH 1/3] add nvidia driver for cuda arm wheel

---
 .github/actions/setup-nvidia/action.yml | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/.github/actions/setup-nvidia/action.yml b/.github/actions/setup-nvidia/action.yml
index 34f0308cb3..26cdc0ba48 100644
--- a/.github/actions/setup-nvidia/action.yml
+++ b/.github/actions/setup-nvidia/action.yml
@@ -8,6 +8,10 @@ inputs:
     required: false
     type: string
     default: "550.54.15"   # https://docs.nvidia.com/datacenter/tesla/tesla-release-notes-550-54-15/index.html
+  GPU_ARCH_TYPE:
+    required: true
+    type: string
+    description: GPU Arch type
 
 runs:
   using: composite
@@ -16,6 +20,11 @@ runs:
       uses: nick-fields/retry@3e91a01664abd3c5cd539100d10d33b9c5b68482
       env:
         DRIVER_VERSION: ${{ inputs.driver-version }}
+        GPU_ARCH_TYPE: ${{ inputs.GPU_ARCH_TYPE }}
+        ${{ if eq(inputs.GPU_ARCH_TYPE , 'cuda-aarch64') }}:
+            ARCH_TYPE: sbsa
+        ${{ else }}:
+            ARCH_TYPE: x86_64
       with:
         timeout_minutes: 10
         max_attempts: 3
@@ -25,7 +34,7 @@ runs:
           set -eou pipefail
 
           DISTRIBUTION=$(. /etc/os-release;echo $ID$VERSION_ID)
-          DRIVER_FN="NVIDIA-Linux-x86_64-${DRIVER_VERSION}.run"
+          DRIVER_FN="NVIDIA-Linux-${GPU_ARCH_TYPE}-${DRIVER_VERSION}.run"
           YUM_REPO_URL="https://nvidia.github.io/nvidia-docker/${DISTRIBUTION}/nvidia-docker.repo"
 
           install_nvidia_docker2_amzn2() {

From 55b882f07c33701aabd5dec2ed92f74de32213c9 Mon Sep 17 00:00:00 2001
From: Ting Lu <tingl@nvidia.com>
Date: Sat, 18 May 2024 22:02:18 -0700
Subject: [PATCH 2/3] update var name

---
 .github/actions/setup-nvidia/action.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/actions/setup-nvidia/action.yml b/.github/actions/setup-nvidia/action.yml
index 26cdc0ba48..01e955a54c 100644
--- a/.github/actions/setup-nvidia/action.yml
+++ b/.github/actions/setup-nvidia/action.yml
@@ -34,7 +34,7 @@ runs:
           set -eou pipefail
 
           DISTRIBUTION=$(. /etc/os-release;echo $ID$VERSION_ID)
-          DRIVER_FN="NVIDIA-Linux-${GPU_ARCH_TYPE}-${DRIVER_VERSION}.run"
+          DRIVER_FN="NVIDIA-Linux-${ARCH_TYPE}-${DRIVER_VERSION}.run"
           YUM_REPO_URL="https://nvidia.github.io/nvidia-docker/${DISTRIBUTION}/nvidia-docker.repo"
 
           install_nvidia_docker2_amzn2() {

From 63aa40bc725027bb29df08f5d9cb9ac915c0b12f Mon Sep 17 00:00:00 2001
From: Ting Lu <tingl@nvidia.com>
Date: Mon, 20 May 2024 07:57:05 -0700
Subject: [PATCH 3/3] Add var ARCH_TYPE and default to x86_64

---
 .github/actions/setup-nvidia/action.yml | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/.github/actions/setup-nvidia/action.yml b/.github/actions/setup-nvidia/action.yml
index 01e955a54c..b7f421355b 100644
--- a/.github/actions/setup-nvidia/action.yml
+++ b/.github/actions/setup-nvidia/action.yml
@@ -9,9 +9,15 @@ inputs:
     type: string
     default: "550.54.15"   # https://docs.nvidia.com/datacenter/tesla/tesla-release-notes-550-54-15/index.html
   GPU_ARCH_TYPE:
-    required: true
+    description: GPU Arch type #"cuda", "cuda-aarch64"
+    required: false
+    type: string
+  ARCH_TYPE:
+    description: CPU Arch Type
+    required: false
     type: string
-    description: GPU Arch type
+    default: "x86_64"   
+
 
 runs:
   using: composite
@@ -21,10 +27,9 @@ runs:
       env:
         DRIVER_VERSION: ${{ inputs.driver-version }}
         GPU_ARCH_TYPE: ${{ inputs.GPU_ARCH_TYPE }}
+        ARCH_TYPE: ${{ inputs.ARCH_TYPE }}
         ${{ if eq(inputs.GPU_ARCH_TYPE , 'cuda-aarch64') }}:
             ARCH_TYPE: sbsa
-        ${{ else }}:
-            ARCH_TYPE: x86_64
       with:
         timeout_minutes: 10
         max_attempts: 3