chore(ml): e2e gpu testing (#31998)

2026-10-17 05:12:35 +08:00 · 2026-10-03 17:15:09 -04:00
parent c5e06dcfd1
commit e89e034515
9 changed files with 227 additions and 27 deletions
@@ -578,7 +578,7 @@ jobs:

  success-check-e2e:
    name: End-to-End Tests Success
-    needs: [e2e-tests-server-cli, e2e-tests-web]
+    needs: [e2e-tests-server-cli, e2e-tests-web, e2e-tests-ml]
    permissions: {}
    runs-on: ubuntu-latest
    if: always()
@@ -626,8 +626,6 @@ jobs:
    needs: pre-job
    if: ${{ fromJSON(needs.pre-job.outputs.should_run).machine-learning == true }}
    runs-on: ubuntu-latest
-    permissions:
-      contents: read
    defaults:
      run:
        working-directory: ./machine-learning
@@ -648,10 +646,92 @@ jobs:
        uses: immich-app/devtools/actions/use-mise@06a9ef925332c91be647d6256642b86b398592c8 # use-mise-action-v3.2.1
        with:
          github_token: ${{ steps.token.outputs.token }}
+          working_directory: ./machine-learning

      - name: Run ci-unit
        run: mise run ci-unit

+  e2e-tests-ml:
+    name: End-to-End Tests (ML, ${{ matrix.runner }})
+    needs: pre-job
+    if: ${{ fromJSON(needs.pre-job.outputs.should_run).machine-learning == true && !github.event.pull_request.head.repo.fork }}
+    runs-on: ${{ matrix.runner }}
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - runner: pokedex-dgpu-nvidia-pascal
+            device: cuda
+            group: cuda-libs
+            provider: CUDAExecutionProvider
+          - runner: pokedex-dgpu-nvidia-ampere
+            device: cuda
+            group: cuda-libs
+            provider: nv_tensorrt_rtx
+          - runner: pokedex-igpu-intel-jasperlake
+            device: openvino
+            provider: OpenVINOExecutionProvider
+    defaults:
+      run:
+        working-directory: ./machine-learning
+    steps:
+      - id: token
+        uses: immich-app/devtools/actions/create-workflow-token@df46d635b905e618b88fc4c95baa920aeb30b309 # create-workflow-token-action-v3.0.1
+        with:
+          client-id: ${{ secrets.PUSH_O_MATIC_APP_CLIENT_ID }}
+          private-key: ${{ secrets.PUSH_O_MATIC_APP_KEY }}
+          permission-contents: read
+
+      - name: Checkout code
+        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+        with:
+          persist-credentials: false
+          submodules: 'recursive'
+          token: ${{ steps.token.outputs.token }}
+
+      - name: Setup Mise
+        uses: immich-app/devtools/actions/use-mise@06a9ef925332c91be647d6256642b86b398592c8 # use-mise-action-v3.2.1
+        with:
+          github_token: ${{ steps.token.outputs.token }}
+          working_directory: ./machine-learning
+
+      - uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 # zizmor: ignore[cache-poisoning] nothing here is published
+        with:
+          path: ~/.cache/uv
+          key: uv-ml-${{ runner.os }}-${{ runner.arch }}-${{ matrix.device }}-${{ hashFiles('machine-learning/uv.lock') }}
+
+      - if: ${{ matrix.device == 'openvino' }}
+        uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 # zizmor: ignore[cache-poisoning] nothing here is published
+        with:
+          path: ~/intel-runtime
+          key: intel-runtime-${{ hashFiles('machine-learning/scripts/install-intel-runtime.sh') }}
+
+      - name: Install the Intel GPU runtime
+        if: ${{ matrix.device == 'openvino' }}
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y --no-install-recommends ocl-icd-libopencl1
+          sudo scripts/install-intel-runtime.sh ~/intel-runtime
+
+      - name: Install
+        env:
+          DEVICE: ${{ matrix.device }}
+          GROUP: ${{ matrix.group }}
+        run: mise run install --extra "$DEVICE" ${GROUP:+--group "$GROUP"}
+
+      - name: Run the full tests
+        env:
+          DEVICE: ${{ matrix.device }}
+          MACHINE_LEARNING_TEST_PROVIDER: ${{ matrix.provider }}
+        run: |
+          if [ "$DEVICE" = cuda ]; then
+            nvidia-smi --query-gpu=name,driver_version,compute_cap --format=csv
+            # the CUDA libraries come from the venv here, where the image has them in its base
+            libs=(.venv/lib/python*/site-packages/nvidia/*/lib)
+            export LD_LIBRARY_PATH=$(IFS=:; echo "${libs[*]/#/$PWD/}")
+          fi
+          mise run test-full
+
  github-files-formatting:
    name: .github Files Formatting
    needs: pre-job
@@ -87,19 +87,11 @@ ENV MACHINE_LEARNING_MODEL_ARENA=false

 FROM python:3.13-slim-trixie@sha256:7c61056e61ac89e852de05f3dc6fa51a6dd2181797bceed46aa725dd7cb2cd3b AS prod-openvino

-RUN apt-get update && \
-    apt-get install --no-install-recommends -yqq ocl-icd-libopencl1 wget && \
-    wget -nv https://github.com/intel/intel-graphics-compiler/releases/download/v2.36.3/intel-igc-core-2_2.36.3+21719_amd64.deb && \
-    wget -nv https://github.com/intel/intel-graphics-compiler/releases/download/v2.36.3/intel-igc-opencl-2_2.36.3+21719_amd64.deb && \
-    wget -nv https://github.com/intel/compute-runtime/releases/download/26.22.38646.4/intel-opencl-icd_26.22.38646.4-0_amd64.deb &&  \
-    wget -nv https://github.com/intel/intel-graphics-compiler/releases/download/igc-1.0.17537.24/intel-igc-core_1.0.17537.24_amd64.deb && \
-    wget -nv https://github.com/intel/intel-graphics-compiler/releases/download/igc-1.0.17537.24/intel-igc-opencl_1.0.17537.24_amd64.deb && \
-    wget -nv https://github.com/intel/compute-runtime/releases/download/24.35.30872.36/intel-opencl-icd-legacy1_24.35.30872.36_amd64.deb && \
-    # TODO: Figure out how to get renovate to manage this differently versioned libigdgmm file
-    wget -nv https://github.com/intel/compute-runtime/releases/download/26.22.38646.4/libigdgmm12_22.10.0_amd64.deb && \
-    dpkg -i *.deb && \
-    rm *.deb && \
-    apt-get remove wget -yqq && \
+RUN --mount=type=bind,source=scripts/install-intel-runtime.sh,target=/tmp/install-intel-runtime.sh \
+    apt-get update && \
+    apt-get install --no-install-recommends -yqq ocl-icd-libopencl1 curl && \
+    /tmp/install-intel-runtime.sh && \
+    apt-get remove curl -yqq && \
    rm -rf /var/lib/apt/lists/*

 FROM nvidia/cuda:12.2.2-runtime-ubuntu22.04@sha256:94c1577b2cd9dd6c0312dc04dff9cb2fdce2b268018abc3d7c2dbcacf1155000 AS prod-cuda
@@ -62,6 +62,7 @@ class Settings(BaseSettings):
    worker_timeout: int = Field(default_factory=default_worker_timeout)
    http_keepalive_timeout_s: int = 2
    test_full: bool = False
+    test_provider: str | None = None
    request_threads: int = os.cpu_count() or 4
    model_inter_op_threads: int = 0
    model_intra_op_threads: int = 0
@@ -11,6 +11,10 @@ run = "uv run ruff check immich_ml scripts test_main.py conftest.py"
 [tasks.test]
 run = "uv run pytest --cov=immich_ml --cov-report term-missing"

+[tasks.test-full]
+env = { MACHINE_LEARNING_TEST_FULL = "true", MACHINE_LEARNING_MODEL_REVISION = "v2" }
+run = "uv run pytest test_main.py -k TestPredictionEndpoints"
+
 [tasks.format]
 run = "uv run ruff format immich_ml scripts test_main.py conftest.py"

@@ -47,6 +47,15 @@ lint = [
 dev = ["locust>=2.15.1", { include-group = "test" }, { include-group = "lint" }]
 # the image replaces onnxruntime-openvino with one built against and bundling this openvino
 onnxruntime-openvino-build = ["openvino>=2026.4.1,<2027", "setuptools", "wheel"]
+# the CUDA libraries the cuda image has in its nvidia/cuda base, for running outside it
+cuda-libs = [
+    "nvidia-cuda-runtime-cu12==12.2.140 ; sys_platform == 'linux' and platform_machine == 'x86_64'",
+    "nvidia-cuda-nvrtc-cu12==12.2.140 ; sys_platform == 'linux' and platform_machine == 'x86_64'",
+    "nvidia-cublas-cu12==12.2.5.6 ; sys_platform == 'linux' and platform_machine == 'x86_64'",
+    "nvidia-cufft-cu12==11.0.8.103 ; sys_platform == 'linux' and platform_machine == 'x86_64'",
+    "nvidia-curand-cu12==10.3.3.141 ; sys_platform == 'linux' and platform_machine == 'x86_64'",
+    "nvidia-cudnn-cu12==9.10.2.21 ; sys_platform == 'linux' and platform_machine == 'x86_64'",
+]

 [project.optional-dependencies]
 cpu = ["onnxruntime>=1.23.2,<2"]
@@ -0,0 +1,26 @@
+#!/usr/bin/env bash
+
+set -e
+
+if [ -n "$1" ]; then # the packages go to the directory given, which can be kept between runs, or to a temporary one
+  dir=$1
+else
+  dir=$(mktemp -d)
+  trap 'rm -rf "$dir"' EXIT
+fi
+mkdir -p "$dir"
+cd "$dir"
+urls=(
+  https://github.com/intel/intel-graphics-compiler/releases/download/v2.36.3/intel-igc-core-2_2.36.3+21719_amd64.deb
+  https://github.com/intel/intel-graphics-compiler/releases/download/v2.36.3/intel-igc-opencl-2_2.36.3+21719_amd64.deb
+  https://github.com/intel/compute-runtime/releases/download/26.22.38646.4/intel-opencl-icd_26.22.38646.4-0_amd64.deb
+  https://github.com/intel/intel-graphics-compiler/releases/download/igc-1.0.17537.24/intel-igc-core_1.0.17537.24_amd64.deb
+  https://github.com/intel/intel-graphics-compiler/releases/download/igc-1.0.17537.24/intel-igc-opencl_1.0.17537.24_amd64.deb
+  https://github.com/intel/compute-runtime/releases/download/24.35.30872.36/intel-opencl-icd-legacy1_24.35.30872.36_amd64.deb
+  # TODO: Figure out how to get renovate to manage the differently versioned libigdgmm file
+  https://github.com/intel/compute-runtime/releases/download/26.22.38646.4/libigdgmm12_22.10.0_amd64.deb
+)
+for url in "${urls[@]}"; do
+  [ -f "${url##*/}" ] || curl -fsSLO "$url"
+done
+dpkg -i "${urls[@]##*/}"
@@ -9,7 +9,7 @@ from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
 from random import randint
 from types import SimpleNamespace
-from typing import Any, Callable, TypeVar
+from typing import Any, Callable, Iterator, TypeVar
 from unittest import mock

 import numpy as np
@@ -1377,6 +1377,11 @@ def expected_box(cell_x: int, cell_y: int) -> list[float]:
    return [cx - 8, cy - 16, cx + 24, cy + 32]


+def cosine(a: list[float], b: list[float]) -> float:
+    x, y = np.asarray(a, dtype=np.float64), np.asarray(b, dtype=np.float64)
+    return float(x @ y / np.linalg.norm(x) / np.linalg.norm(y))
+
+
 M = TypeVar("M", bound=InferenceModel[Any])


@@ -2242,6 +2247,16 @@ def test_ping_endpoint(deployed_app: TestClient) -> None:
    reason="More time-consuming since it deploys the app and loads models.",
 )
 class TestPredictionEndpoints:
+    @pytest.fixture(autouse=True)
+    def providers(self, mocker: MockerFixture) -> Iterator[None]:
+        sessions = mocker.spy(ort, "InferenceSession")
+        yield
+        if settings.test_provider:
+            providers = [session.get_providers()[0] for session in sessions.spy_return_list]
+            assert providers == [settings.test_provider] * len(providers)
+            devices = [call.kwargs["provider_options"][0].get("device_type") for call in sessions.call_args_list]
+            assert all(device in (None, f"GPU.{settings.device_id}") for device in devices)
+
    def test_clip_image_endpoint(
        self, asset: Callable[[str], bytes], responses: dict[str, Any], deployed_app: TestClient
    ) -> None:
@@ -2252,7 +2267,7 @@ class TestPredictionEndpoints:
        )

        assert response.status_code == 200
-        assert np.allclose(orjson.loads(response.json()["clip"]), responses["clip"]["image"], atol=1e-3)
+        assert cosine(orjson.loads(response.json()["clip"]), responses["clip"]["image"]) >= 0.999

    def test_clip_text_endpoint(self, responses: dict[str, Any], deployed_app: TestClient) -> None:
        response = deployed_app.post(
@@ -2264,7 +2279,7 @@ class TestPredictionEndpoints:
        )

        assert response.status_code == 200
-        assert np.allclose(orjson.loads(response.json()["clip"]), responses["clip"]["text"], atol=1e-3)
+        assert cosine(orjson.loads(response.json()["clip"]), responses["clip"]["text"]) >= 0.999

    def test_face_endpoint(
        self, asset: Callable[[str], bytes], responses: dict[str, Any], deployed_app: TestClient
@@ -2293,9 +2308,9 @@ class TestPredictionEndpoints:
        assert len(actual["facial-recognition"]) == len(expected["faces"])

        for expected_face, actual_face in zip(expected["faces"], actual["facial-recognition"]):
-            assert actual_face["boundingBox"] == expected_face["boundingBox"]
-            assert actual_face["score"] == pytest.approx(expected_face["score"], abs=1e-3)
-            assert np.allclose(orjson.loads(actual_face["embedding"]), expected_face["embedding"], atol=1e-3)
+            assert actual_face["boundingBox"] == pytest.approx(expected_face["boundingBox"], abs=1)
+            assert actual_face["score"] == pytest.approx(expected_face["score"], abs=0.01)
+            assert cosine(orjson.loads(actual_face["embedding"]), expected_face["embedding"]) >= 0.999

    def test_ocr_endpoint(
        self, asset: Callable[[str], bytes], responses: dict[str, Any], deployed_app: TestClient
@@ -2322,7 +2337,13 @@ class TestPredictionEndpoints:

        actual = response.json()["ocr"]
        assert response.status_code == 200
-        assert actual["text"] == expected["text"]
-        assert np.allclose(actual["box"], expected["box"], atol=1e-3)
-        assert np.allclose(actual["boxScore"], expected["boxScore"], atol=1e-3)
-        assert np.allclose(actual["textScore"], expected["textScore"], atol=1e-3)
+        # lines are matched by text, since one read differently can move in reading order
+        found = {text: line for line, text in enumerate(actual["text"])}
+        shared = [(found[text], line) for line, text in enumerate(expected["text"]) if text in found]
+        assert len(shared) >= 0.9 * len(expected["text"])
+        assert len(shared) >= 0.9 * len(actual["text"])
+        for found_line, expected_line in shared:
+            box = slice(8 * found_line, 8 * found_line + 8)
+            assert actual["box"][box] == pytest.approx(
+                expected["box"][8 * expected_line : 8 * expected_line + 8], abs=0.01
+            )
@@ -720,6 +720,14 @@ rocm = [
 ]

 [package.dev-dependencies]
+cuda-libs = [
+    { name = "nvidia-cublas-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
+    { name = "nvidia-cuda-nvrtc-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
+    { name = "nvidia-cuda-runtime-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
+    { name = "nvidia-cudnn-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
+    { name = "nvidia-cufft-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
+    { name = "nvidia-curand-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
+]
 dev = [
    { name = "httpx" },
    { name = "locust" },
@@ -794,6 +802,14 @@ requires-dist = [
 provides-extras = ["cpu", "cuda", "openvino", "armnn", "rknn", "rocm"]

 [package.metadata.requires-dev]
+cuda-libs = [
+    { name = "nvidia-cublas-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'", specifier = "==12.2.5.6" },
+    { name = "nvidia-cuda-nvrtc-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'", specifier = "==12.2.140" },
+    { name = "nvidia-cuda-runtime-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'", specifier = "==12.2.140" },
+    { name = "nvidia-cudnn-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'", specifier = "==9.10.2.21" },
+    { name = "nvidia-cufft-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'", specifier = "==11.0.8.103" },
+    { name = "nvidia-curand-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'", specifier = "==10.3.3.141" },
+]
 dev = [
    { name = "httpx", specifier = ">=0.24.1" },
    { name = "locust", specifier = ">=2.15.1" },
@@ -1271,6 +1287,57 @@ wheels = [
    { url = "https://files.pythonhosted.org/packages/fb/0b/b12a2df5d1b774bd9007a6fdff9381145b6223d37f11afc9c37ab0efd9a1/numpy-2.5.3-cp315-cp315t-win_arm64.whl", hash = "sha256:befa1ae5bd6030b3f512b43ff3fa5290bbed6b84411a44244b14adf835f5b89d", size = 10850807, upload-time = "2026-09-06T16:27:43.868Z" },
 ]

+[[package]]
+name = "nvidia-cublas-cu12"
+version = "12.2.5.6"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/b6/6a/e8cca34f85b18a0280e3a19faca1923f6a04e7d587e9d8e33bc295a52b6d/nvidia_cublas_cu12-12.2.5.6-py3-none-manylinux1_x86_64.whl", hash = "sha256:f448d8db463f9a2c7dd9a1e607c0961d3e95f825b86839c5882f7d073d406df0", size = 417825305, upload-time = "2023-08-29T16:03:20.821Z" },
+]
+
+[[package]]
+name = "nvidia-cuda-nvrtc-cu12"
+version = "12.2.140"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/cb/84/bb975ca63e3d4017824ad3d62e141db8492bc53cec6dc84705fb0650c22d/nvidia_cuda_nvrtc_cu12-12.2.140-py3-none-manylinux1_x86_64.whl", hash = "sha256:b1cdcae6443cb7cdcfb5d1bf51896a148c9f82f9e65b27db9ab095eaed689aa7", size = 23430053, upload-time = "2023-08-29T16:01:36.949Z" },
+]
+
+[[package]]
+name = "nvidia-cuda-runtime-cu12"
+version = "12.2.140"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/95/46/6361d45c7a6fe3b3bb8d5fa35eb43c1dcd12d14799a0dc6faef3d76eaf41/nvidia_cuda_runtime_cu12-12.2.140-py3-none-manylinux1_x86_64.whl", hash = "sha256:acb9cb7e44594a3512533497ea790d9f0b0e399abcc888584c157d1ffa080e41", size = 845802, upload-time = "2023-08-29T16:00:12.485Z" },
+]
+
+[[package]]
+name = "nvidia-cudnn-cu12"
+version = "9.10.2.21"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "nvidia-cublas-cu12", marker = "(platform_machine != 'aarch64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux')" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/ba/51/e123d997aa098c61d029f76663dedbfb9bc8dcf8c60cbd6adbe42f76d049/nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:949452be657fa16687d0930933f032835951ef0892b37d2d53824d1a84dc97a8", size = 706758467, upload-time = "2025-06-06T21:54:08.597Z" },
+]
+
+[[package]]
+name = "nvidia-cufft-cu12"
+version = "11.0.8.103"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/8f/86/311c32926aa321ebc169d6c910a1ecd30339ae198014e0d4a7cf448f6ce9/nvidia_cufft_cu12-11.0.8.103-py3-none-manylinux1_x86_64.whl", hash = "sha256:72e961e578f9be9e316ccb4db8652c421dd0b57442624faa71061d452f449a67", size = 98570970, upload-time = "2023-08-29T16:03:46.694Z" },
+]
+
+[[package]]
+name = "nvidia-curand-cu12"
+version = "10.3.3.141"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/c1/2c/9677fd666335801f565505bf831891b6ac6f645fe9e7689acc7d09f9a7fb/nvidia_curand_cu12-10.3.3.141-py3-none-manylinux1_x86_64.whl", hash = "sha256:7b6bad6817b5df2a6f5ec221206fb895d9195fdd7d2443f3e86554108c141a0a", size = 56463566, upload-time = "2023-08-29T16:04:11.71Z" },
+]
+
 [[package]]
 name = "omegaconf"
 version = "2.3.0"