feat(machine-learning): support cuda 12 (#7569)

* feat: support cuda12 * fix: group optional * move to cuda 12 * pr feedback
immich-app · Mar 3, 2024 · 8ce18b3 · 8ce18b3
1 parent fd3503e
commit 8ce18b3
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 16 deletions.
diff --git a/machine-learning/Dockerfile b/machine-learning/Dockerfile
@@ -39,7 +39,7 @@ FROM python:3.11-slim-bookworm@sha256:ce81dc539f0aedc9114cae640f8352fad83d37461c
 FROM openvino/ubuntu22_runtime:2023.1.0@sha256:002842a9005ba01543b7169ff6f14ecbec82287f09c4d1dd37717f0a8e8754a7 as prod-openvino
 USER root
 
-FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04@sha256:85fb7ac694079fff1061a0140fd5b5a641997880e12112d92589c3bbb1e8b7ca as prod-cuda
+FROM nvidia/cuda:12.2.2-cudnn8-runtime-ubuntu22.04@sha256:8b51b1fe922964d73c482a267b5b519e990d90bf744ec7a40419923737caff6d as prod-cuda
 
 COPY --from=builder-cuda /usr/local/bin/python3 /usr/local/bin/python3
 COPY --from=builder-cuda /usr/local/lib/python3.11 /usr/local/lib/python3.11

diff --git a/machine-learning/poetry.lock b/machine-learning/poetry.lock
diff --git a/machine-learning/pyproject.toml b/machine-learning/pyproject.toml
@@ -45,7 +45,7 @@ onnxruntime = "^1.15.0"
 optional = true
 
 [tool.poetry.group.cuda.dependencies]
-onnxruntime-gpu = "^1.15.0"
+onnxruntime-gpu = {version = "^1.17.0", source = "cuda12"}
 
 [tool.poetry.group.openvino]
 optional = true
@@ -59,6 +59,11 @@ optional = true
 [tool.poetry.group.armnn.dependencies]
 onnxruntime = "^1.15.0"
 
+[[tool.poetry.source]]
+name = "cuda12"
+url = "https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/"
+priority = "explicit"
+
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"