From 2e0f8be23bd392e7c3388c2fb42c3c753698c242 Mon Sep 17 00:00:00 2001
From: Sergiy <sergiy.popovich@gmail.com>
Date: Fri, 7 Feb 2025 16:52:45 +0000
Subject: [PATCH] fix: install libnvrtc

---
 docker/Dockerfile.all     |  2 ++
 docker/Dockerfile.all_pcg |  2 ++
 install_zutils.py         | 15 +++++++++++++--
 3 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/docker/Dockerfile.all b/docker/Dockerfile.all
index 4ccdea336..033c0c909 100644
--- a/docker/Dockerfile.all
+++ b/docker/Dockerfile.all
@@ -24,6 +24,8 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     zstd parallel coreutils libjemalloc2 \
     # waterz
     libboost-dev g++ \
+    # torch JIT
+    cuda-nvrtc-12-0 \
   && find /usr/local/lib /opt/conda -type f -executable -exec ldd '{}' ';' \
       | awk '/=>/ { print $(NF-1) }' \
       | sort -u \
diff --git a/docker/Dockerfile.all_pcg b/docker/Dockerfile.all_pcg
index 8c6b27716..fbd85740b 100644
--- a/docker/Dockerfile.all_pcg
+++ b/docker/Dockerfile.all_pcg
@@ -24,6 +24,8 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     zstd parallel coreutils libjemalloc2 \
     # waterz
     libboost-dev g++ \
+     # torch JIT
+    cuda-nvrtc-12-0 \
   && find /usr/local/lib /opt/conda -type f -executable -exec ldd '{}' ';' \
       | awk '/=>/ { print $(NF-1) }' \
       | sort -u \
diff --git a/install_zutils.py b/install_zutils.py
index adffc0749..c0ab24b46 100755
--- a/install_zutils.py
+++ b/install_zutils.py
@@ -625,10 +625,21 @@ def main():
         sudo_prefix = ""
 
     if not args.skip_apt:
+        run_command(f"{sudo_prefix}apt-get update", "Updating package lists")
+        run_command(
+            command=f"{sudo_prefix}apt-get install -y wget gnupg software-properties-common && "
+            "wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-ubuntu2004.pin && "
+            f"{sudo_prefix}mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600 && "
+            f"{sudo_prefix}apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/7fa2af80.pub && "
+            f"{sudo_prefix}apt-key adv --keyserver keyserver.ubuntu.com --recv-keys A4B469963BF863CC && "
+            f'{sudo_prefix}add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/ /"',
+            description="Adding NVIDIA apt repository",
+        )
+
+        run_command(f"{sudo_prefix}apt-get update", "Updating package lists")
         apt_packages = [
             "git",
             "build-essential",
-            "wget",
             "ffmpeg",  # opencv
             "libsm6",  # opencv
             "libxext6",  # opencv
@@ -645,8 +656,8 @@ def main():
             "libtbb-dev",  # abiss
             "libboost-dev",  # waterz
             "unixodbc-dev",  # ???
+            "cuda-nvrtc-12-0",  # torch jit
         ]
-        run_command(f"{sudo_prefix}apt-get update", "Updating package lists")
         apt_install_flags = '-y --no-install-recommends -o Dpkg::Options::="--force-confdef" -o Dpkg::Options::="--force-confold"'
         run_command(
             f"{sudo_prefix} apt-get install {apt_install_flags} {' '.join(apt_packages)}",