cuda for arm64

separate base image fix name make sure onnxruntime-gpu is installed
immich-app · Sep 8, 2024 · 0521147 · 0521147
1 parent d1ce9e4
commit 0521147
Show file tree

Hide file tree

Showing 4 changed files with 51 additions and 7 deletions.
diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml
@@ -111,6 +111,10 @@ jobs:
           - platforms: linux/amd64
             device: cuda
             suffix: -cuda
+
+          - platforms: linux/arm64
+            device: jetson
+            suffix: -jetson
 
           - platforms: linux/amd64
             device: openvino

diff --git a/machine-learning/Dockerfile b/machine-learning/Dockerfile
@@ -1,11 +1,14 @@
 ARG DEVICE=cpu
+ARG TARGETARCH
 
 FROM python:3.11-bookworm@sha256:20c1819af5af3acba0b2b66074a2615e398ceee6842adf03cd7ad5f8d0ee3daf AS builder-cpu
 
 FROM builder-cpu AS builder-openvino
 
 FROM builder-cpu AS builder-cuda
 
+FROM builder-cpu AS builder-jetson
+
 FROM builder-cpu AS builder-armnn
 
 ENV ARMNN_PATH=/opt/armnn
@@ -32,7 +35,11 @@ RUN poetry config installer.max-workers 10 && \
 RUN python3 -m venv /opt/venv
 
 COPY poetry.lock pyproject.toml ./
-RUN poetry install --sync --no-interaction --no-ansi --no-root --with ${DEVICE} --without dev
+RUN if [ "$DEVICE" = "jetson" ]; then \
+    # hack to work around poetry not setting the right filename for the wheel https://github.com/python-poetry/poetry/issues/4472
+    wget -q -O onnxruntime_gpu-1.18.0-cp311-cp311-manylinux_aarch64.whl https://nvidia.box.com/shared/static/fy55jvniujjbigr4gwkv8z1ma6ipgspg.whl; fi && \
+    poetry install --sync --no-interaction --no-ansi --no-root --with ${DEVICE} --without dev && \
+    if [ "$DEVICE" = "jetson" ]; then rm onnxruntime_gpu-1.18.0-cp311-cp311-manylinux_aarch64.whl; fi
 
 FROM python:3.11-slim-bookworm@sha256:ed4e985674f478c90ce879e9aa224fbb772c84e39b4aed5155b9e2280f131039 AS prod-cpu
 
@@ -55,6 +62,12 @@ 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
 COPY --from=builder-cuda /usr/local/lib/libpython3.11.so /usr/local/lib/libpython3.11.so
 
+FROM nvcr.io/nvidia/l4t-cuda:11.4.19-runtime@sha256:fb22ff080631990dda403fd768acb384dc3745a7e516f5ed1dc4c4944898da78 as prod-jetson
+
+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
+COPY --from=builder-cuda /usr/local/lib/libpython3.11.so /usr/local/lib/libpython3.11.so
+
 FROM prod-cpu AS prod-armnn
 
 ENV LD_LIBRARY_PATH=/opt/armnn
@@ -79,7 +92,7 @@ FROM prod-${DEVICE} AS prod
 ARG DEVICE
 
 RUN apt-get update && \
-    apt-get install -y --no-install-recommends tini $(if ! [ "$DEVICE" = "openvino" ]; then echo "libmimalloc2.0"; fi) && \
+    apt-get install -y --no-install-recommends tini $(if ! { [ "$DEVICE" = "openvino" ] || [ "$DEVICE" = "jetson" ]; }; then echo "libmimalloc2.0"; fi) && \
     apt-get autoremove -yqq && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*

diff --git a/machine-learning/poetry.lock b/machine-learning/poetry.lock
diff --git a/machine-learning/pyproject.toml b/machine-learning/pyproject.toml
@@ -4,15 +4,15 @@ version = "1.114.0"
 description = ""
 authors = ["Hau Tran <alex.tran1502@gmail.com>"]
 readme = "README.md"
-packages = [{include = "app"}]
+packages = [{ include = "app" }]
 
 [tool.poetry.dependencies]
 python = ">=3.10,<4.0"
 insightface = ">=0.7.3,<1.0"
 opencv-python-headless = ">=4.7.0.72,<5.0"
 pillow = ">=9.5.0,<11.0"
 fastapi-slim = ">=0.95.2,<1.0"
-uvicorn = {extras = ["standard"], version = ">=0.22.0,<1.0"}
+uvicorn = { extras = ["standard"], version = ">=0.22.0,<1.0" }
 pydantic = "^1.10.8"
 aiocache = ">=0.12.1,<1.0"
 rich = ">=13.4.2"
@@ -45,7 +45,13 @@ onnxruntime = "^1.15.0"
 optional = true
 
 [tool.poetry.group.cuda.dependencies]
-onnxruntime-gpu = {version = "^1.17.0", source = "cuda12"}
+onnxruntime-gpu = { version = "^1.17.0", source = "cuda12", markers = "platform_machine == 'x86_64'" }
+
+[tool.poetry.group.jetson]
+optional = true
+
+[tool.poetry.group.jetson.dependencies]
+onnxruntime-gpu = { python = "3.11", path = "onnxruntime_gpu-1.18.0-cp311-cp311-manylinux_aarch64.whl", markers = "platform_machine == 'aarch64'" }
 
 [tool.poetry.group.openvino]
 optional = true