handle libcudnn9

neggles · Jul 15, 2024 · 26ac80a · 26ac80a
1 parent 09d60ed
commit 26ac80a
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/docker-bake.hcl b/docker-bake.hcl
@@ -191,7 +191,7 @@ target "base" {
 }
 
 target xformers-wheel {
-  inherits = ["base-cu121-torch240"]
+  inherits = ["base-cu124-torch240"]
   target   = "xformers-wheel"
   tags = [
     repoImage("xformers", "v0.0.27", cudaName("12.4.1"), torchName("2.4.0"))

diff --git a/docker/base/Dockerfile b/docker/base/Dockerfile
@@ -91,12 +91,18 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     python3-venv\
   && apt-get clean
 
-# Install CUDNN dev package to match existing binary package
+# Install CUDNN dev package to match existing binary package (if installed)
 RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     --mount=type=cache,target=/var/lib/apt,sharing=locked \
     apt-get update \
-  && apt-get -y install --no-install-recommends \
-    libcudnn8-dev=$(dpkg-query --showformat='${Version}' --show libcudnn8) \
+  && if $(dpkg-query --show libcudnn8 2>&1 >/dev/null); then  \
+    apt-get -y install --no-install-recommends \
+        libcudnn8-dev=$(dpkg-query --showformat='${Version}' --show libcudnn8) \
+    ; fi \
+  && if $(dpkg-query --show libcudnn9 2>&1 >/dev/null); then  \
+    apt-get -y install --no-install-recommends \
+        libcudnn9-dev=$(dpkg-query --showformat='${Version}' --show libcudnn9) \
+    ; fi \
   && apt-get clean
 
 # Install TensorRT libraries