From b89d2fca87282ee47d469e9b42f9448cf80252fc Mon Sep 17 00:00:00 2001
From: stark <stark@yunify.com>
Date: Mon, 11 Mar 2024 17:16:25 +0800
Subject: [PATCH] Add ubuntu2204 cuda for iso repo

---
 hack/gen-repository-iso/dockerfile.ubuntu2204 | 14 ++++++++++++++
 hack/gen-repository-iso/packages.yaml         |  3 +++
 2 files changed, 17 insertions(+)

diff --git a/hack/gen-repository-iso/dockerfile.ubuntu2204 b/hack/gen-repository-iso/dockerfile.ubuntu2204
index 7a92912a..42ebb4d9 100644
--- a/hack/gen-repository-iso/dockerfile.ubuntu2204
+++ b/hack/gen-repository-iso/dockerfile.ubuntu2204
@@ -1,5 +1,6 @@
 FROM ubuntu:22.04 as ubuntu2204
 ARG TARGETARCH
+ARG DISTRO=ubuntu2204
 ARG OS_RELEASE=jammy
 ARG DIR=ubuntu-22.04-${TARGETARCH}-debs
 ARG PKGS=.common[],.debs[],.ubuntu[],.ubuntu2204[]
@@ -15,6 +16,19 @@ RUN apt update -qq \
     && echo "deb [arch=$TARGETARCH] https://download.docker.com/linux/ubuntu ${OS_RELEASE} stable" > /etc/apt/sources.list.d/docker.list\
     && apt update -qq
 
+# install NVIDIA CUDA
+RUN if [ "${TARGETARCH}" = "amd64" ]; then \
+        ARCH=x86_64; \
+    else \
+        ARCH=${TARGETARCH}; \
+    fi \
+    && wget https://developer.download.nvidia.com/compute/cuda/repos/${DISTRO}/${ARCH}/cuda-archive-keyring.gpg \
+    && mv cuda-archive-keyring.gpg /usr/share/keyrings/cuda-archive-keyring.gpg \
+    && echo "deb [signed-by=/usr/share/keyrings/cuda-archive-keyring.gpg] https://developer.download.nvidia.com/compute/cuda/repos/${DISTRO}/${ARCH}/ /" | tee /etc/apt/sources.list.d/cuda-${DISTRO}-${ARCH}.list \
+    && wget https://developer.download.nvidia.com/compute/cuda/repos/${DISTRO}/${ARCH}/cuda-${DISTRO}.pin \
+    && mv cuda-${DISTRO}.pin /etc/apt/preferences.d/cuda-repository-pin-600 \
+    && apt-get update
+
 WORKDIR /package
 COPY packages.yaml .
 
diff --git a/hack/gen-repository-iso/packages.yaml b/hack/gen-repository-iso/packages.yaml
index 65d89a3e..4708bd56 100644
--- a/hack/gen-repository-iso/packages.yaml
+++ b/hack/gen-repository-iso/packages.yaml
@@ -39,6 +39,9 @@ debs:
   - openssh-server
   - software-properties-common
   - sudo
+  - cuda-toolkit-12-4
+  - nvidia-driver-550-open
+  - cuda-drivers-550
 
 centos:
   - containerd.io