Fix cuda pipeline (#467)

levkk · web-flow · commit d293bb7ca476 · 2022-10-27T16:38:43.000-07:00
diff --git a/pgml-extension/Dockerfile b/pgml-extension/Dockerfile
@@ -8,7 +8,7 @@ ARG PGVERSION=14
 ARG DEBIAN_FRONTEND=noninteractive
 ARG PACKAGE_VERSION=2.0.0
 ARG PACKAGE_PYTHON=true
-ARG PACKAGE_CUDA=true
+ARG PACKAGE_CUDA=false
 ENV TZ=Etc/UTC
 
 # Apt-fast
@@ -30,7 +30,6 @@ RUN echo "deb [signed-by=/usr/share/keyrings/kitware-archive-keyring.gpg] https:
 RUN apt-get update && apt-fast install postgresql-${PGVERSION} libopenblas-dev cmake postgresql-server-dev-${PGVERSION} pkg-config libssl-dev build-essential libclang-dev clang libpython3-dev python3-dev -y
 
 # CUDA
-# ARM64
 COPY docker/install_cuda.sh install_cuda.sh
 RUN bash install_cuda.sh
 
diff --git a/pgml-extension/docker/install_cuda.sh b/pgml-extension/docker/install_cuda.sh
@@ -10,13 +10,13 @@ else
 fi
 
 # ARM
-if [[ ${ARCH} == "aarch64" ]]; then
-    wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/cuda-keyring_1.0-1_all.deb
+if [[ ${ARCH} == "arm64" ]]; then
+    wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu${UBUNTU_VERSION//.}/sbsa/cuda-keyring_1.0-1_all.deb
 fi
 
 # Intel
 if [[ ${ARCH} == "amd64" ]]; then
-    wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.0-1_all.deb
+    wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu${UBUNTU_VERSION//.}/x86_64/cuda-keyring_1.0-1_all.deb
 fi
 
 dpkg -i cuda-keyring_1.0-1_all.deb
