feat(ml): improve model loading (#31742)

* use recent mimalloc

* jit caching

* bump onnxruntime-gpu

* explicit memory release

* refactor in-memory cache

* configure hf org and revision

* refactor model load

* move orjson response

* disable direct submission for i915

* flush denormals

* update tests and lock

* disable suboptimal fusions

* respect resize_mode

* increase threads for clip text

machine-learning/Dockerfile

@@ -4,7 +4,7 @@ FROM python:3.13-slim-trixie@sha256:9d2e5553305c7c7b0097999bb17187c69b921ccd6bc9
 
 FROM builder-cpu AS builder-openvino
 
-FROM builder-cpu AS builder-cuda
+FROM python:3.13-slim-bookworm@sha256:2325bb286ec344af3e5898cc224b5844e2707ac6e26b1632516fd3edc84a5e26 AS builder-cuda
 
 FROM python:3.13-trixie@sha256:4a0820d8b1b1a1e012a23ff9a1296a3e28b48896a4d27a328a91ee197ce619d9 AS builder-armnn
 
@@ -108,17 +108,23 @@ ENV MACHINE_LEARNING_MODEL_ARENA=false
 
 ADD --chmod=644 --checksum=sha256:73993ed4b440460825f21611731564503cc1d5a0c123746477da6cd574f34885 "https://github.com/airockchip/rknn-toolkit2/raw/refs/tags/${RKNN_TOOLKIT_VERSION}/rknpu2/runtime/Linux/librknn_api/aarch64/librknnrt.so" /usr/lib/
 
+FROM python:3.13-slim-trixie@sha256:9d2e5553305c7c7b0097999bb17187c69b921ccd6bc9d40e4bb5ebe652c00285 AS mimalloc
+RUN echo "deb http://deb.debian.org/debian trixie-backports main" > /etc/apt/sources.list.d/backports.list && \
+    apt-get update && \
+    apt-get install -y --no-install-recommends libmimalloc3/trixie-backports && \
+    cp /usr/lib/*-linux-gnu/libmimalloc.so.3 /libmimalloc.so
+
 FROM prod-${DEVICE} AS prod
 
 ARG DEVICE
 
-RUN MIMALLOC=$(case "$DEVICE" in cuda) echo "libmimalloc2.0";; openvino|rocm);; *) echo "libmimalloc3";; esac); \
-    apt-get update && \
-    apt-get install -y --no-install-recommends tini ccache libgl1 libglib2.0-0 libgomp1 $MIMALLOC && \
+COPY --from=mimalloc /libmimalloc.so /usr/lib/libmimalloc.so
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends tini ccache libgomp1 && \
     apt-get autoremove -yqq && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/* && \
-    { [ -z "$MIMALLOC" ] || ls /usr/lib/*-linux-gnu/libmimalloc.so.? > /etc/ld.so.preload; }
+    case "$DEVICE" in openvino|rocm);; *) echo /usr/lib/libmimalloc.so > /etc/ld.so.preload;; esac
 
 WORKDIR /usr/src
 ENV TRANSFORMERS_CACHE=/cache \
@@ -157,6 +163,10 @@ ENV IMMICH_SOURCE_URL=https://github.com/immich-app/immich/commit/${BUILD_SOURCE
 # store huggingface in /cache subfolder to enable rootless deployment with a single volume
 ENV HF_HOME=/cache/hf-cache
 
+ENV CUDA_CACHE_PATH=/cache/jit/cuda \
+    AMD_COMGR_CACHE_DIR=/cache/jit/comgr \
+    NEO_CACHE_DIR=/cache/jit/neo
+
 ENTRYPOINT ["tini", "--"]
 CMD ["python", "-m", "immich_ml"]
 

machine-learning/conftest.py

@@ -9,6 +9,7 @@ from PIL import Image
 
 from immich_ml.config import log
 from immich_ml.main import app
+from immich_ml.models.base import InferenceModel
 
 
 @pytest.fixture
@@ -18,7 +19,7 @@ def pil_image() -> Image.Image:
 
 @pytest.fixture
 def mock_get_model() -> Iterator[mock.Mock]:
-    with mock.patch("immich_ml.models.cache.from_model_type", autospec=True) as mocked:
+    with mock.patch("immich_ml.models.cache.get_model_class") as mocked:
         yield mocked
 
 
@@ -187,5 +188,8 @@ def exception() -> Iterator[mock.Mock]:
 
 @pytest.fixture(scope="function")
 def snapshot_download() -> Iterator[mock.Mock]:
-    with mock.patch("immich_ml.models.base.snapshot_download") as mocked:
+    with (
+        mock.patch("immich_ml.models.base.snapshot_download") as mocked,
+        mock.patch.object(InferenceModel, "cached", new_callable=mock.PropertyMock, side_effect=lambda: mocked.called),
+    ):
         yield mocked

machine-learning/immich_ml/__main__.py

@@ -13,6 +13,11 @@ else:
 
 module_dir = Path(__file__).parent
 
+# on i915, Intel's runtime submits every allocation in the process with each run it gives OpenVINO's GPU plugin
+if any(driver.resolve().name == "i915" for driver in Path("/sys/class/drm").glob("renderD*/device/driver")):
+    os.environ.setdefault("NEOReadDebugKeys", "1")
+    os.environ.setdefault("EnableDirectSubmission", "0")
+
 
 def is_ipv6(host: str) -> bool:
     try:

machine-learning/immich_ml/allocator.py

@@ -0,0 +1,10 @@
+import ctypes
+
+_process = ctypes.CDLL(None)
+
+
+def release() -> None:
+    if hasattr(_process, "mi_collect"):
+        _process.mi_collect(True)
+    elif hasattr(_process, "malloc_trim"):
+        _process.malloc_trim(0)

machine-learning/immich_ml/config.py

@@ -13,7 +13,7 @@ from rich.logging import RichHandler
 from uvicorn import Server
 from uvicorn.workers import UvicornWorker
 
-from .schemas import ModelPrecision
+from .schemas import ModelOrganization, ModelPrecision
 
 
 class ClipSettings(BaseModel):
@@ -74,6 +74,12 @@ class Settings(BaseSettings):
     max_batch_size: MaxBatchSize | None = None
     openvino_precision: ModelPrecision = ModelPrecision.FP32
     rocm_precision: ModelPrecision = ModelPrecision.FP32
+    model_organization: ModelOrganization = ModelOrganization.APP
+    model_revision: str = "main"
+
+    @property
+    def legacy_models(self) -> bool:
+        return self.model_revision == "main"
 
     @property
     def device_id(self) -> str:

machine-learning/immich_ml/main.py

@@ -2,25 +2,26 @@ import asyncio
 import gc
 import os
 import signal
-import threading
 import time
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import asynccontextmanager
 from functools import partial
-from typing import Any, AsyncGenerator, Callable, Iterator
+from typing import Any, AsyncGenerator, Callable
 from zipfile import BadZipFile
 
 import orjson
 from fastapi import Depends, FastAPI, File, Form, HTTPException
-from fastapi.responses import PlainTextResponse
+from fastapi.responses import JSONResponse, PlainTextResponse
 from onnxruntime.capi.onnxruntime_pybind11_state import InvalidProtobuf, NoSuchFile
 from PIL.Image import Image
 from pydantic import ValidationError
 from starlette.formparsers import MultiPartParser
 
+from immich_ml import allocator
 from immich_ml.models import get_model_deps
 from immich_ml.models.base import InferenceModel
 from immich_ml.models.transforms import decode_pil
+from immich_ml.sessions.ort import flush_denormals
 
 from .config import PreloadModelData, log, settings
 from .models.cache import ModelCache
@@ -32,18 +33,24 @@ from .schemas import (
     ModelIdentity,
     ModelTask,
     ModelType,
-    ORJSONResponse,
     PipelineRequest,
     T,
 )
 
+
+class ORJSONResponse(JSONResponse):
+    def render(self, content: Any) -> bytes:
+        return orjson.dumps(content, option=orjson.OPT_SERIALIZE_NUMPY)
+
+
 MultiPartParser.spool_max_size = 2**26  # spools to disk if payload is 64 MiB or larger
 
-model_cache = ModelCache(revalidate=settings.model_ttl > 0)
+model_cache = ModelCache()
 thread_pool: ThreadPoolExecutor | None = None
-lock = threading.Lock()
+MODEL_FILE_ERRORS = (InvalidProtobuf, NoSuchFile)
 active_requests = 0
 last_called: float | None = None
+release: asyncio.TimerHandle | None = None
 
 
 @asynccontextmanager
@@ -57,19 +64,20 @@ async def lifespan(_: FastAPI) -> AsyncGenerator[None, None]:
     )
 
     try:
+        flush_denormals()  # for work that runs on the event loop's own thread
         if settings.request_threads > 0:
             # asyncio is a huge bottleneck for performance, so we use a thread pool to run blocking code
-            thread_pool = ThreadPoolExecutor(settings.request_threads) if settings.request_threads > 0 else None
+            thread_pool = ThreadPoolExecutor(settings.request_threads, initializer=flush_denormals)
             log.info(f"Initialized request thread pool with {settings.request_threads} threads.")
         if settings.model_ttl > 0 and settings.model_ttl_poll_s > 0:
             asyncio.ensure_future(idle_shutdown_task())
         if settings.preload is not None:
             await preload_models(settings.preload)
+            allocator.release()
         yield
     finally:
         log.handlers.clear()
-        for model in model_cache.cache._cache.values():
-            del model
+        model_cache.clear()
         if thread_pool is not None:
             thread_pool.shutdown()
         gc.collect()
@@ -78,10 +86,10 @@ async def lifespan(_: FastAPI) -> AsyncGenerator[None, None]:
 async def preload_models(preload: PreloadModelData) -> None:
     log.info(f"Preloading models: clip:{preload.clip} facial_recognition:{preload.facial_recognition}")
 
-    async def load_models(model_string: str, model_type: ModelType, model_task: ModelTask) -> None:
+    async def load_models(model_string: str, model_type: ModelType, model_task: ModelTask, **options: Any) -> None:
         for model_name in model_string.split(","):
             model_name = model_name.strip()
-            model = await model_cache.get(model_name, model_type, model_task)
+            model = model_cache.get(model_name, model_type, model_task, **options)
             await load(model)
 
     if preload.clip.textual is not None:
@@ -119,14 +127,19 @@ async def preload_models(preload: PreloadModelData) -> None:
         )
 
 
-def update_state() -> Iterator[None]:
-    global active_requests, last_called
+async def update_state() -> AsyncGenerator[None, None]:
+    global active_requests, last_called, release
     active_requests += 1
     last_called = time.time()
+    if release is not None:
+        release.cancel()
+        release = None
     try:
         yield
     finally:
         active_requests -= 1
+        if not active_requests:
+            release = asyncio.get_running_loop().call_later(5, allocator.release)
 
 
 def get_entries(entries: str = Form()) -> InferenceEntries:
@@ -187,9 +200,7 @@ async def run_inference(payload: Image | str, entries: InferenceEntries) -> Infe
     response: InferenceResponse = {}
 
     async def _run_inference(entry: InferenceEntry) -> None:
-        model = await model_cache.get(
-            entry["name"], entry["type"], entry["task"], ttl=settings.model_ttl, **entry["options"]
-        )
+        model = model_cache.get(entry["name"], entry["type"], entry["task"], ttl=settings.model_ttl, **entry["options"])
         inputs = [payload]
         for dep in model.depends:
             try:
@@ -198,14 +209,15 @@ async def run_inference(payload: Image | str, entries: InferenceEntries) -> Infe
                 message = f"Task {entry['task']} of type {entry['type']} depends on output of {dep}"
                 raise HTTPException(400, message)
         model = await load(model)
-        output = await run(model.predict, *inputs, **entry["options"])
+        output = await attempt(model, partial(model.predict, *inputs, **entry["options"]), MODEL_FILE_ERRORS)
         outputs[model.identity] = output
         response[entry["task"]] = output
 
-    without_deps, with_deps = entries
-    await asyncio.gather(*[_run_inference(entry) for entry in without_deps])
-    if with_deps:
-        await asyncio.gather(*[_run_inference(entry) for entry in with_deps])
+    for stage in entries:  # those that depend on another's output run after it
+        results = await asyncio.gather(*[_run_inference(entry) for entry in stage], return_exceptions=True)
+        for result in results:
+            if isinstance(result, BaseException):
+                raise result
     if isinstance(payload, Image):
         response["imageHeight"], response["imageWidth"] = payload.height, payload.width
 
@@ -220,42 +232,40 @@ async def run(func: Callable[..., T], *args: Any, **kwargs: Any) -> T:
 
 
 async def load(model: InferenceModel) -> InferenceModel:
-    if model.loaded:
-        return model
+    if not model.loaded:
+        await attempt(model, model.load, (OSError, BadZipFile, *MODEL_FILE_ERRORS))
+    return model
 
-    def _load(model: InferenceModel) -> InferenceModel:
-        if model.load_attempts > 1:
+
+async def attempt(model: InferenceModel, func: Callable[[], T], corrupt: tuple[type[Exception], ...]) -> T:
+    def _attempt() -> T:
+        if not model.loaded and model.load_attempts > 1:
             raise HTTPException(500, f"Failed to load model '{model.model_name}'")
-        with lock:
-            try:
-                model.load()
-            except FileNotFoundError as e:
-                if model.model_format == ModelFormat.ONNX:
-                    raise e
-                log.warning(
-                    f"{model.model_format.upper()} is available, but model '{model.model_name}' does not support it.",
-                    exc_info=e,
-                )
-                model.model_format = ModelFormat.ONNX
-                model.load()
-        return model
+        try:
+            return func()
+        except FileNotFoundError as e:
+            if model.model_format == ModelFormat.ONNX:
+                raise e
+            log.warning(
+                f"{model.model_format.upper()} is available, but model '{model.model_name}' does not support it.",
+                exc_info=e,
+            )
+            model.unload()
+            model.model_format = ModelFormat.ONNX
+            return func()
 
     try:
-        return await run(_load, model)
-    except (OSError, InvalidProtobuf, BadZipFile, NoSuchFile):
+        return await run(_attempt)
+    except corrupt:
         log.warning(f"Failed to load {model.model_type.replace('_', ' ')} model '{model.model_name}'. Clearing cache.")
+        model.unload()
         model.clear_cache()
-        return await run(_load, model)
+        return await run(_attempt)
 
 
 async def idle_shutdown_task() -> None:
     while True:
-        if (
-            last_called is not None
-            and not active_requests
-            and not lock.locked()
-            and time.time() - last_called > settings.model_ttl
-        ):
+        if last_called is not None and not active_requests and time.time() - last_called > settings.model_ttl:
             log.info("Shutting down due to inactivity.")
             os.kill(os.getpid(), signal.SIGINT)
             break

machine-learning/immich_ml/models/__init__.py

@@ -1,5 +1,3 @@
-from typing import Any
-
 from immich_ml.models.base import InferenceModel
 from immich_ml.models.clip.textual import MClipTextualEncoder, OpenClipTextualEncoder
 from immich_ml.models.clip.visual import OpenClipVisualEncoder
@@ -40,9 +38,5 @@ def get_model_class(model_name: str, model_type: ModelType, model_task: ModelTas
             raise ValueError(f"Unknown model combination: {source}, {model_type}, {model_task}")
 
 
-def from_model_type(model_name: str, model_type: ModelType, model_task: ModelTask, **kwargs: Any) -> InferenceModel:
-    return get_model_class(model_name, model_type, model_task)(model_name, **kwargs)
-
-
 def get_model_deps(model_name: str, model_type: ModelType, model_task: ModelTask) -> list[tuple[ModelType, ModelTask]]:
     return get_model_class(model_name, model_type, model_task).depends

machine-learning/immich_ml/models/base.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from abc import ABC, abstractmethod
 from pathlib import Path
 from shutil import rmtree
+from threading import Lock
 from typing import Any, ClassVar
 
 from huggingface_hub import snapshot_download
@@ -15,10 +16,17 @@ from ..config import clean_name, log, settings
 from ..schemas import ModelFormat, ModelIdentity, ModelSession, ModelTask, ModelType
 from ..sessions.ann import AnnSession
 
+_IGNORED_PATTERNS: dict[ModelFormat, list[str]] = {
+    ModelFormat.ONNX: ["*.armnn", "*.rknn"],
+    ModelFormat.ARMNN: ["*.rknn"],
+    ModelFormat.RKNN: ["*.armnn"],
+}
+
 
 class InferenceModel(ABC):
     depends: ClassVar[list[ModelIdentity]]
     identity: ClassVar[ModelIdentity]
+    threads: ClassVar[int] = 2  # how many threads one run takes for CPU inference
 
     def __init__(
         self,
@@ -30,6 +38,7 @@ class InferenceModel(ABC):
     ) -> None:
         self.loaded = session is not None
         self.load_attempts = 0
+        self._load_lock = Lock()
         self.model_name = clean_name(model_name)
         self.cache_dir = Path(cache_dir) if cache_dir is not None else self._cache_dir_default
         self.model_format = model_format if model_format is not None else self._model_format_default
@@ -37,50 +46,47 @@ class InferenceModel(ABC):
             self.session = session
 
     def download(self) -> None:
-        if not self.cached:
-            model_type = self.model_type.replace("-", " ")
-            log.info(f"Downloading {model_type} model '{self.model_name}' to {self.model_path}. This may take a while.")
-            self._download()
+        if self.cached:
+            return
+        model_type = self.model_type.replace("-", " ")
+        log.info(f"Downloading {model_type} model '{self.model_name}' to {self.model_dir}. This may take a while.")
+        snapshot_download(
+            f"{settings.model_organization}/{self.model_name}",
+            revision=settings.model_revision,
+            cache_dir=self.cache_dir,
+            local_dir=self.cache_dir,
+            ignore_patterns=_IGNORED_PATTERNS.get(self.model_format, []),
+        )
+        if not self.cached:  # the repository has nothing in this format, which is what sends a model back to ONNX
+            raise FileNotFoundError(f"Model file not found: {self.model_path}")
 
     def load(self) -> None:
-        if self.loaded:
-            return
-        self.load_attempts += 1
+        with self._load_lock:  # two requests may arrive for a model neither has loaded
+            if self.loaded:
+                return
+            self.load_attempts += 1
 
-        self.download()
-        attempt = f"Attempt #{self.load_attempts} to load" if self.load_attempts > 1 else "Loading"
-        log.info(f"{attempt} {self.model_type.replace('-', ' ')} model '{self.model_name}' to memory")
-        self.session = self._load()
-        self.loaded = True
+            self.download()
+            attempt = f"Attempt #{self.load_attempts} to load" if self.load_attempts > 1 else "Loading"
+            log.info(f"{attempt} {self.model_type.replace('-', ' ')} model '{self.model_name}' to memory")
+            self.session = self._load()
+            self.loaded = True
+
+    def unload(self) -> None:
+        with self._load_lock:
+            if self.loaded:
+                del self.session
+                self.loaded = False
 
     def predict(self, *inputs: Any, **model_kwargs: Any) -> Any:
         self.load()
-        if model_kwargs:
-            self.configure(**model_kwargs)
         return self._predict(*inputs, **model_kwargs)
 
     @abstractmethod
     def _predict(self, *inputs: Any, **model_kwargs: Any) -> Any: ...
 
-    def configure(self, **kwargs: Any) -> None:
-        pass
-
-    def _download(self) -> None:
-        ignored_patterns: dict[ModelFormat, list[str]] = {
-            ModelFormat.ONNX: ["*.armnn", "*.rknn"],
-            ModelFormat.ARMNN: ["*.rknn"],
-            ModelFormat.RKNN: ["*.armnn"],
-        }
-
-        snapshot_download(
-            f"immich-app/{clean_name(self.model_name)}",
-            cache_dir=self.cache_dir,
-            local_dir=self.cache_dir,
-            ignore_patterns=ignored_patterns.get(self.model_format, []),
-        )
-
     def _load(self) -> ModelSession:
-        return self._make_session(self.model_path)
+        return self._make_session()
 
     def clear_cache(self) -> None:
         if not self.cache_dir.exists():
@@ -104,20 +110,14 @@ class InferenceModel(ABC):
             self.cache_dir.unlink()
         self.cache_dir.mkdir(parents=True, exist_ok=True)
 
-    def _make_session(self, model_path: Path) -> ModelSession:
-        if not model_path.is_file():
-            raise FileNotFoundError(f"Model file not found: {model_path}")
-
-        match model_path.suffix:
-            case ".armnn":
-                session: ModelSession = AnnSession(model_path)
-            case ".onnx":
-                session = OrtSession(model_path)
-            case ".rknn":
-                session = rknn.RknnSession(model_path)
-            case _:
-                raise ValueError(f"Unsupported model file type: {model_path.suffix}")
-        return session
+    def _make_session(self) -> ModelSession:
+        match self.model_format:
+            case ModelFormat.ARMNN:
+                return AnnSession(self.model_path)
+            case ModelFormat.ONNX:
+                return OrtSession(self.model_path, threads=self.threads)
+            case ModelFormat.RKNN:
+                return rknn.RknnSession(self.model_path)
 
     def model_path_for_format(self, model_format: ModelFormat) -> Path:
         model_path_prefix = rknn.model_prefix if model_format == ModelFormat.RKNN else None

machine-learning/immich_ml/models/cache.py

@@ -1,60 +1,43 @@
+import asyncio
 from typing import Any
 
-from aiocache.backends.memory import SimpleMemoryCache
-from aiocache.lock import OptimisticLock
-from aiocache.plugins import TimingPlugin
-
-from immich_ml.models import from_model_type
+from immich_ml import allocator
+from immich_ml.models import get_model_class
 from immich_ml.models.base import InferenceModel
 
-from ..schemas import ModelTask, ModelType, has_profiling
+from ..schemas import ModelTask, ModelType
 
 
 class ModelCache:
-    """Fetches a model from an in-memory cache, instantiating it if it's missing."""
+    def __init__(self) -> None:
+        self._models: dict[tuple[Any, ...], InferenceModel] = {}
+        self._expiries: dict[tuple[Any, ...], asyncio.TimerHandle] = {}
 
-    def __init__(
-        self,
-        revalidate: bool = False,
-        timeout: int | None = None,
-        profiling: bool = False,
-    ) -> None:
-        """
-        Args:
-            revalidate: Resets TTL on cache hit. Useful to keep models in memory while active. Defaults to False.
-            timeout: Maximum allowed time for model to load. Disabled if None. Defaults to None.
-            profiling: Collects metrics for cache operations, adding slight overhead. Defaults to False.
-        """
-
-        plugins = []
-
-        if profiling:
-            plugins.append(TimingPlugin())
-
-        self.should_revalidate = revalidate
-
-        self.cache = SimpleMemoryCache(timeout=timeout, plugins=plugins, namespace=None)
-
-    async def get(
-        self, model_name: str, model_type: ModelType, model_task: ModelTask, **model_kwargs: Any
+    def get(
+        self, model_name: str, model_type: ModelType, model_task: ModelTask, ttl: int | None = None, **options: Any
     ) -> InferenceModel:
-        key = f"{model_name}{model_type}{model_task}"
+        model_cls = get_model_class(model_name, model_type, model_task)
+        key = (model_name, model_type, model_task)
+        if key not in self._models:
+            self._models[key] = model_cls(model_name, **options)
+        elif key not in self._expiries:
+            ttl = None  # a model that was preloaded stays
+        if ttl:
+            if key in self._expiries:
+                self._expiries[key].cancel()
+            self._expiries[key] = asyncio.get_running_loop().call_later(ttl, self._evict, key)
+        return self._models[key]
 
-        async with OptimisticLock(self.cache, key) as lock:
-            model: InferenceModel | None = await self.cache.get(key)
-            if model is None:
-                model = from_model_type(model_name, model_type, model_task, **model_kwargs)
-                await lock.cas(model, ttl=model_kwargs.get("ttl", None))
-            elif self.should_revalidate:
-                await self.revalidate(key, model_kwargs.get("ttl", None))
-        return model
+    def clear(self) -> None:
+        if not self._models:
+            return
+        self._models.clear()
+        for key in list(self._expiries.keys()):
+            self._expiries.pop(key).cancel()
+        allocator.release()
 
-    async def get_profiling(self) -> dict[str, float] | None:
-        if not has_profiling(self.cache):
-            return None
-
-        return self.cache.profiling
-
-    async def revalidate(self, key: str, ttl: int | None) -> None:
-        if ttl is not None and key in self.cache._handlers:
-            await self.cache.expire(key, ttl)
+    def _evict(self, key: tuple[Any, ...]) -> None:
+        if key not in self._models:
+            return
+        del self._models[key], self._expiries[key]
+        allocator.release()

machine-learning/immich_ml/models/clip/textual.py

@@ -18,6 +18,7 @@ from immich_ml.schemas import ModelSession, ModelTask, ModelType
 class BaseCLIPTextualEncoder(InferenceModel):
     depends = []
     identity = (ModelType.TEXTUAL, ModelTask.SEARCH)
+    threads = 4  # lower search latency
 
     def _predict(self, inputs: str, language: str | None = None) -> str:
         tokens = self.tokenize(inputs, language=language)

machine-learning/immich_ml/models/clip/visual.py

@@ -70,8 +70,16 @@ class OpenClipVisualEncoder(BaseCLIPVisualEncoder):
         return super()._load()
 
     def transform(self, image: Image.Image) -> dict[str, NDArray[np.float32]]:
-        image = resize_pil(image, self.size)
-        image = crop_pil(image, self.size)
+        image = self._resize(image)
         image_np = to_numpy(image)
         image_np = normalize(image_np, self.mean, self.std)
         return {"image": np.expand_dims(image_np.transpose(2, 0, 1), 0)}
+
+    def _resize(self, image: Image.Image) -> Image.Image:
+        match self.preprocess_cfg.get("resize_mode", "shortest"):  # open_clip's default for older configs
+            case "squash":
+                return image.resize((self.size, self.size), resample=self.resampling)
+            case "shortest":
+                return crop_pil(resize_pil(image, self.size, self.resampling), self.size)
+            case mode:
+                raise ValueError(f"Unsupported resize_mode {mode!r} in {self.preprocess_cfg_path}")

machine-learning/immich_ml/models/facial_recognition/recognition.py

@@ -33,10 +33,10 @@ class FaceRecognizer(InferenceModel):
         self.batch_size = max_batch_size if max_batch_size else self._batch_size_default
 
     def _load(self) -> ModelSession:
-        session = self._make_session(self.model_path)
+        session = self._make_session()
         if (not self.batch_size or self.batch_size > 1) and str(session.get_inputs()[0].shape[0]) != "batch":
             self._add_batch_axis(self.model_path)
-            session = self._make_session(self.model_path)
+            session = self._make_session()
         return session
 
     def _predict(

machine-learning/immich_ml/schemas.py

@@ -1,19 +1,12 @@
 from collections.abc import Sequence
 from enum import Enum
-from typing import Any, Literal, Protocol, TypeGuard, TypeVar
+from typing import Any, Literal, Protocol, TypeVar
 
 import numpy as np
 import numpy.typing as npt
-import orjson
-from fastapi.responses import JSONResponse
 from typing_extensions import TypedDict
 
 
-class ORJSONResponse(JSONResponse):
-    def render(self, content: Any) -> bytes:
-        return orjson.dumps(content, option=orjson.OPT_SERIALIZE_NUMPY)
-
-
 class StrEnum(str, Enum):
     __str__ = str.__str__
 
@@ -56,6 +49,11 @@ class ModelPrecision(StrEnum):
     FP32 = "FP32"
 
 
+class ModelOrganization(StrEnum):
+    APP = "immich-app"
+    TESTING = "immich-testing"
+
+
 ModelIdentity = tuple[ModelType, ModelTask]
 
 
@@ -84,10 +82,6 @@ class ModelSession(Protocol):
     def get_metadata(self) -> dict[str, str]: ...
 
 
-class HasProfiling(Protocol):
-    profiling: dict[str, float]
-
-
 class FaceDetectionOutput(TypedDict):
     boxes: npt.NDArray[np.float32]
     scores: npt.NDArray[np.float32]
@@ -124,8 +118,4 @@ InferenceEntries = tuple[list[InferenceEntry], list[InferenceEntry]]
 InferenceResponse = dict[ModelTask | Literal["imageHeight"] | Literal["imageWidth"], Any]
 
 
-def has_profiling(obj: Any) -> TypeGuard[HasProfiling]:
-    return hasattr(obj, "profiling") and isinstance(obj.profiling, dict)
-
-
 T = TypeVar("T")

machine-learning/immich_ml/sessions/ort.py

@@ -1,5 +1,8 @@
 from __future__ import annotations
 
+import ctypes
+import platform
+import sys
 from collections.abc import Sequence
 from pathlib import Path
 from threading import Lock
@@ -53,9 +56,13 @@ class OrtSession:
         providers: list[str] | None = None,
         provider_options: list[dict[str, Any]] | None = None,
         sess_options: ort.SessionOptions | None = None,
+        threads: int = 2,
     ):
         self.model_path = Path(model_path)
+        self.threads = threads
         self.providers = providers if providers is not None else self._providers_default
+        self.disabled_optimizers = _disabled_optimizers_default(self.providers)
+        log.debug(f"Setting disabled_optimizers to {self.disabled_optimizers}")
         self.provider_options = provider_options if provider_options is not None else self._provider_options_default
         self.sess_options = sess_options if sess_options is not None else self._sess_options_default
         self.session = ort.InferenceSession(
@@ -63,6 +70,7 @@ class OrtSession:
             providers=self.providers,
             provider_options=self.provider_options,
             sess_options=self.sess_options,
+            disabled_optimizers=self.disabled_optimizers,
         )
 
     def get_inputs(self) -> Sequence[SessionNode]:
@@ -184,6 +192,8 @@ class OrtSession:
     @property
     def _sess_options_default(self) -> ort.SessionOptions:
         sess_options = ort.SessionOptions()
+        # some CPUs slow down many times over on subnormal operands, and ORT clears the flush in its threads otherwise
+        sess_options.add_session_config_entry("session.set_denormal_as_zero", "1")
         sess_options.enable_cpu_mem_arena = settings.model_arena
 
         # avoid thread contention between models
@@ -198,9 +208,32 @@ class OrtSession:
         if settings.model_intra_op_threads > 0:
             sess_options.intra_op_num_threads = settings.model_intra_op_threads
         elif settings.model_intra_op_threads == 0 and self.providers == ["CPUExecutionProvider"]:
-            sess_options.intra_op_num_threads = 2
+            sess_options.intra_op_num_threads = self.threads
 
         if sess_options.inter_op_num_threads > 1:
             sess_options.execution_mode = ort.ExecutionMode.ORT_PARALLEL
 
         return sess_options
+
+
+def _disabled_optimizers_default(providers: list[str]) -> list[str]:
+    disabled_optimizers: list[str] = []
+    if platform.machine() in ("arm64", "aarch64"):  # as macOS and Linux name the same architecture
+        disabled_optimizers.append("ConvAddActivationFusion")
+
+    # the Gemm it makes runs slower than the pair it replaces there, and on CUDA it also keeps BiasGelu from fusing
+    if "CoreMLExecutionProvider" in providers or "CUDAExecutionProvider" in providers:
+        disabled_optimizers.append("MatMulAddFusion")
+
+    return disabled_optimizers
+
+
+def flush_denormals() -> None:
+    """Reads subnormal floats as zero on the calling thread, which ORT does only on its own threads."""
+    if sys.platform != "linux" or platform.machine() != "x86_64":
+        return
+    libm = ctypes.CDLL("libm.so.6")
+    env = (ctypes.c_uint32 * 8)()  # glibc's fenv_t here: the x87 environment, then MXCSR
+    libm.fegetenv(env)
+    env[7] |= 0x8040  # denormals are zero, flush to zero
+    libm.fesetenv(env)

machine-learning/pyproject.toml

@@ -6,7 +6,6 @@ authors = [{ name = "Hau Tran", email = "alex.tran1502@gmail.com" }]
 requires-python = ">=3.12,<4.0"
 readme = "README.md"
 dependencies = [
-    "aiocache>=0.12.1,<1.0",
     "fastapi>=0.95.2,<1.0",
     "gunicorn>=21.1.0",
     "huggingface-hub>=1.0,<2.0",
@@ -48,7 +47,7 @@ dev = ["locust>=2.15.1", { include-group = "test" }, { include-group = "lint" }]
 
 [project.optional-dependencies]
 cpu = ["onnxruntime>=1.23.2,<2"]
-cuda = ["onnxruntime-gpu>=1.23.2,<2"]
+cuda = ["onnxruntime-gpu>=1.27.1,<1.28"] # TODO: last ORT release supporting CUDA 12.2, also need to handle CUDA 13
 openvino = ["onnxruntime-openvino>=1.24.1,<2"]
 armnn = ["onnxruntime>=1.23.2,<2"]
 rknn = ["onnxruntime>=1.23.2,<2", "rknn-toolkit-lite2>=2.3.0,<3"]
@@ -56,6 +55,16 @@ rocm = ["onnxruntime-migraphx>=1.23.2,<2"]
 
 [tool.uv]
 compile-bytecode = true
+override-dependencies = ["opencv-python ; sys_platform == 'never'"] # only install the headless version
+
+[tool.uv.sources]
+onnxruntime-gpu = { index = "onnxruntime-cuda-12" }
+
+[[tool.uv.index]]
+name = "onnxruntime-cuda-12"
+url = "https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/"
+explicit = true
+
 
 [tool.hatch.build.targets.sdist]
 include = ["immich_ml"]

machine-learning/test_main.py

@@ -1,5 +1,10 @@
 import json
 import os
+import platform
+import sys
+import threading
+import time
+from concurrent.futures import ThreadPoolExecutor
 from io import BytesIO
 from pathlib import Path
 from random import randint
@@ -17,8 +22,10 @@ from PIL import Image
 from pytest import MonkeyPatch
 from pytest_mock import MockerFixture
 
-from immich_ml.config import MaxBatchSize, Settings, settings
-from immich_ml.main import load, preload_models
+from immich_ml import allocator
+from immich_ml.config import MaxBatchSize, PreloadModelData, Settings, settings
+from immich_ml.main import app, lifespan, load, preload_models, update_state
+from immich_ml.main import run_inference as run_request
 from immich_ml.models.base import InferenceModel
 from immich_ml.models.cache import ModelCache
 from immich_ml.models.clip.textual import MClipTextualEncoder, OpenClipTextualEncoder
@@ -29,7 +36,7 @@ from immich_ml.models.ocr.detection import TextDetector
 from immich_ml.models.ocr.recognition import TextRecognizer
 from immich_ml.schemas import ModelFormat, ModelPrecision, ModelTask, ModelType
 from immich_ml.sessions.ann import AnnSession
-from immich_ml.sessions.ort import OrtSession
+from immich_ml.sessions.ort import OrtSession, flush_denormals
 from immich_ml.sessions.rknn import RknnSession, run_inference
 
 
@@ -162,6 +169,7 @@ class TestBase:
 
         snapshot_download.assert_called_once_with(
             "immich-app/ViT-B-32__openai",
+            revision=settings.model_revision,
             cache_dir=encoder.cache_dir,
             local_dir=encoder.cache_dir,
             ignore_patterns=["*.armnn", "*.rknn"],
@@ -171,27 +179,18 @@ class TestBase:
         encoder = OpenClipTextualEncoder("ViT-B-32__openai", model_format=ModelFormat.ARMNN)
         encoder.download()
 
-        snapshot_download.assert_called_once_with(
-            "immich-app/ViT-B-32__openai",
-            cache_dir=encoder.cache_dir,
-            local_dir=encoder.cache_dir,
-            ignore_patterns=["*.rknn"],
-        )
+        assert snapshot_download.call_args.kwargs["ignore_patterns"] == ["*.rknn"]
 
     def test_download_downloads_rknn_if_preferred_format(self, snapshot_download: mock.Mock) -> None:
         encoder = OpenClipTextualEncoder("ViT-B-32__openai", model_format=ModelFormat.RKNN)
         encoder.download()
 
-        snapshot_download.assert_called_once_with(
-            "immich-app/ViT-B-32__openai",
-            cache_dir=encoder.cache_dir,
-            local_dir=encoder.cache_dir,
-            ignore_patterns=["*.armnn"],
-        )
+        assert snapshot_download.call_args.kwargs["ignore_patterns"] == ["*.armnn"]
 
     def test_throws_exception_if_model_path_does_not_exist(
-        self, snapshot_download: mock.Mock, ort_session: mock.Mock, path: mock.Mock
+        self, ort_session: mock.Mock, path: mock.Mock, mocker: MockerFixture
     ) -> None:
+        snapshot_download = mocker.patch("immich_ml.models.base.snapshot_download")  # which brings no such file
         path.return_value.__truediv__.return_value.__truediv__.return_value.is_file.return_value = False
 
         encoder = OpenClipTextualEncoder("ViT-B-32__openai", cache_dir=path)
@@ -202,6 +201,18 @@ class TestBase:
         snapshot_download.assert_called_once()
         ort_session.assert_not_called()
 
+    @pytest.mark.parametrize(
+        ("symbols", "called"), [(["mi_collect", "malloc_trim"], "mi_collect"), (["malloc_trim"], "malloc_trim")]
+    )
+    def test_returns_memory_through_the_allocator_the_process_runs_on(
+        self, mocker: MockerFixture, symbols: list[str], called: str
+    ) -> None:
+        process = mocker.patch("immich_ml.allocator._process", mock.Mock(spec=symbols))
+
+        allocator.release()
+
+        getattr(process, called).assert_called_once()
+
 
 @pytest.mark.usefixtures("ort_session")
 class TestOrtSession:
@@ -370,12 +381,63 @@ class TestOrtSession:
 
         assert session.provider_options == []
 
+    @pytest.mark.parametrize(
+        ("machine", "disabled"),
+        [("arm64", ["ConvAddActivationFusion"]), ("aarch64", ["ConvAddActivationFusion"]), ("x86_64", [])],
+    )
+    def test_disables_the_fusion_that_is_slower_on_arm_whatever_the_os_calls_it(
+        self, ort_session: mock.Mock, mocker: MockerFixture, machine: str, disabled: list[str]
+    ) -> None:
+        mocker.patch("immich_ml.sessions.ort.platform.machine", return_value=machine)
+
+        OrtSession("ViT-B-32__openai", providers=["CPUExecutionProvider"])
+
+        assert ort_session.call_args.kwargs["disabled_optimizers"] == disabled
+
+    @pytest.mark.parametrize(
+        ("providers", "disabled"),
+        [
+            (["CUDAExecutionProvider", "CPUExecutionProvider"], ["MatMulAddFusion"]),
+            (["CoreMLExecutionProvider", "CPUExecutionProvider"], ["MatMulAddFusion"]),
+            (["CPUExecutionProvider"], []),
+        ],
+    )
+    def test_disables_the_fusion_that_is_slower_on_cuda_and_coreml(
+        self, ort_session: mock.Mock, mocker: MockerFixture, providers: list[str], disabled: list[str]
+    ) -> None:
+        mocker.patch("immich_ml.sessions.ort.platform.machine", return_value="x86_64")
+
+        OrtSession("ViT-B-32__openai", providers=providers)
+
+        assert ort_session.call_args.kwargs["disabled_optimizers"] == disabled
+
     def test_sets_default_sess_options_if_cpu(self) -> None:
         session = OrtSession("ViT-B-32__openai", providers=["CPUExecutionProvider"])
 
         assert session.sess_options.execution_mode == ort.ExecutionMode.ORT_SEQUENTIAL
         assert session.sess_options.inter_op_num_threads == 1
         assert session.sess_options.intra_op_num_threads == 2
+        assert session.sess_options.get_session_config_entry("session.set_denormal_as_zero") == "1"
+
+    @pytest.mark.skipif(sys.platform != "linux" or platform.machine() != "x86_64", reason="an x86 register")
+    def test_flushes_denormals_on_the_calling_thread(self) -> None:
+        results: list[float] = []
+
+        def probe() -> None:  # a fresh thread, as a request thread is
+            flush_denormals()
+            results.append(float((np.array([1e-39], np.float32) * np.float32(1))[0]))
+
+        thread = threading.Thread(target=probe)
+        thread.start()
+        thread.join()
+
+        assert results == [0.0]
+
+    def test_gives_a_model_the_cpu_threads_it_asks_for(self) -> None:
+        session = OrtSession("ViT-B-32__openai", providers=["CPUExecutionProvider"], threads=4)
+
+        assert session.sess_options.intra_op_num_threads == 4
+        assert OpenClipTextualEncoder.threads == 4 and OpenClipVisualEncoder.threads == 2
 
     @pytest.mark.ov_device_ids(["CPU"])
     def test_sets_default_sess_options_if_openvino_cpu(self, ov_device_ids: list[str]) -> None:
@@ -595,6 +657,22 @@ class TestCLIP:
         assert len(embedding) == clip_model_cfg["embed_dim"]
         mocked.run.assert_called_once()
 
+    def test_visual_squashes_when_the_tower_was_calibrated_that_way(
+        self, mocker: MockerFixture, clip_model_cfg: dict[str, Any], clip_preprocess_cfg: dict[str, Any]
+    ) -> None:
+        mocker.patch.object(OpenClipVisualEncoder, "download")
+        mocker.patch.object(OpenClipVisualEncoder, "model_cfg", clip_model_cfg)
+        mocker.patch.object(OpenClipVisualEncoder, "preprocess_cfg", clip_preprocess_cfg | {"resize_mode": "squash"})
+        session = mocker.patch.object(InferenceModel, "_make_session", autospec=True).return_value
+        session.run.return_value = [[self.embedding]]
+
+        # a stripe far enough left that a shortest-side crop of this frame would discard it
+        image = Image.new("RGB", (600, 200), "black")
+        image.paste(Image.new("RGB", (20, 200), "white"), (0, 0))
+        OpenClipVisualEncoder("ViT-B-32__openai", cache_dir="test_cache").predict(image)
+
+        assert (session.run.call_args.args[1]["image"][0, :, 0, 0] > 0).all()  # white, where black normalizes below 0
+
     def test_basic_text(
         self,
         mocker: MockerFixture,
@@ -1179,56 +1257,76 @@ class TestOcr:
 class TestCache:
     async def test_caches(self, mock_get_model: mock.Mock) -> None:
         model_cache = ModelCache()
-        await model_cache.get("test_model_name", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION)
-        await model_cache.get("test_model_name", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION)
-        assert len(model_cache.cache._cache) == 1
-        mock_get_model.assert_called_once()
+        model_cache.get("test_model_name", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION)
+        model_cache.get("test_model_name", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION)
+        assert len(model_cache._models) == 1
+        mock_get_model.return_value.assert_called_once()
 
     async def test_kwargs_used(self, mock_get_model: mock.Mock) -> None:
         model_cache = ModelCache()
-        await model_cache.get(
-            "test_model_name", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION, cache_dir="test_cache"
-        )
-        mock_get_model.assert_called_once_with(
-            "test_model_name", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION, cache_dir="test_cache"
-        )
+        model_cache.get("test_model_name", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION, cache_dir="test_cache")
+        mock_get_model.return_value.assert_called_once_with("test_model_name", cache_dir="test_cache")
 
     async def test_different_clip(self, mock_get_model: mock.Mock) -> None:
         model_cache = ModelCache()
-        await model_cache.get("test_model_name", ModelType.VISUAL, ModelTask.SEARCH)
-        await model_cache.get("test_model_name", ModelType.TEXTUAL, ModelTask.SEARCH)
-        mock_get_model.assert_has_calls(
-            [
-                mock.call("test_model_name", ModelType.VISUAL, ModelTask.SEARCH),
-                mock.call("test_model_name", ModelType.TEXTUAL, ModelTask.SEARCH),
-            ]
-        )
-        assert len(model_cache.cache._cache) == 2
+        model_cache.get("test_model_name", ModelType.VISUAL, ModelTask.SEARCH)
+        model_cache.get("test_model_name", ModelType.TEXTUAL, ModelTask.SEARCH)
+        assert mock_get_model.call_args_list == [
+            mock.call("test_model_name", ModelType.VISUAL, ModelTask.SEARCH),
+            mock.call("test_model_name", ModelType.TEXTUAL, ModelTask.SEARCH),
+        ]
+        assert len(model_cache._models) == 2
 
-    @mock.patch("immich_ml.models.cache.OptimisticLock", autospec=True)
-    async def test_model_ttl(self, mock_lock_cls: mock.Mock, mock_get_model: mock.Mock) -> None:
+    async def test_lets_go_of_a_model_unused_for_its_ttl_and_returns_its_memory(
+        self, mock_get_model: mock.Mock, mocker: MockerFixture
+    ) -> None:
+        events: list[str] = []
+
+        class Model:
+            def __init__(self, *args: Any, **kwargs: Any) -> None:
+                pass
+
+            def __del__(self) -> None:
+                events.append("destroyed")
+
+        mock_get_model.return_value = Model
+        mocker.patch("immich_ml.models.cache.allocator.release", side_effect=lambda: events.append("released"))
+        loop = mocker.patch("immich_ml.models.cache.asyncio.get_running_loop").return_value
         model_cache = ModelCache()
-        await model_cache.get("test_model_name", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION, ttl=100)
-        mock_lock_cls.return_value.__aenter__.return_value.cas.assert_called_with(mock.ANY, ttl=100)
 
-    @mock.patch("immich_ml.models.cache.SimpleMemoryCache.expire")
-    async def test_revalidate_get(self, mock_cache_expire: mock.Mock, mock_get_model: mock.Mock) -> None:
-        model_cache = ModelCache(revalidate=True)
-        await model_cache.get("test_model_name", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION, ttl=100)
-        await model_cache.get("test_model_name", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION, ttl=100)
-        mock_cache_expire.assert_called_once_with(mock.ANY, 100)
+        model_cache.get("test_model_name", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION, ttl=100)
+        delay, evict, key = loop.call_later.call_args.args
+        evict(key)
 
-    async def test_profiling(self, mock_get_model: mock.Mock) -> None:
-        model_cache = ModelCache(profiling=True)
-        await model_cache.get("test_model_name", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION, ttl=100)
-        profiling = await model_cache.get_profiling()
-        assert isinstance(profiling, dict)
-        assert profiling == model_cache.cache.profiling
+        assert delay == 100
+        assert events == ["destroyed", "released"]
+        assert not model_cache._models
+
+    async def test_a_use_starts_the_ttl_over(self, mock_get_model: mock.Mock, mocker: MockerFixture) -> None:
+        loop = mocker.patch("immich_ml.models.cache.asyncio.get_running_loop").return_value
+        model_cache = ModelCache()
+
+        model_cache.get("test_model_name", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION, ttl=100)
+        model_cache.get("test_model_name", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION, ttl=100)
+
+        loop.call_later.return_value.cancel.assert_called_once()
+        assert loop.call_later.call_count == 2
+
+    async def test_keeps_a_preloaded_model_whatever_ttl_a_request_names(
+        self, mock_get_model: mock.Mock, mocker: MockerFixture
+    ) -> None:
+        loop = mocker.patch("immich_ml.models.cache.asyncio.get_running_loop").return_value
+        model_cache = ModelCache()
+
+        model_cache.get("test_model_name", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION)
+        model_cache.get("test_model_name", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION, ttl=100)
+
+        loop.call_later.assert_not_called()
 
     async def test_loads_mclip(self) -> None:
         model_cache = ModelCache()
 
-        model = await model_cache.get("XLM-Roberta-Large-Vit-B-32", ModelType.TEXTUAL, ModelTask.SEARCH)
+        model = model_cache.get("XLM-Roberta-Large-Vit-B-32", ModelType.TEXTUAL, ModelTask.SEARCH)
 
         assert isinstance(model, MClipTextualEncoder)
         assert model.model_name == "XLM-Roberta-Large-Vit-B-32"
@@ -1238,13 +1336,13 @@ class TestCache:
         model_cache = ModelCache()
 
         with pytest.raises(ValueError):
-            await model_cache.get("XLM-Roberta-Large-Vit-B-32", ModelType.TEXTUAL, invalid)
+            model_cache.get("XLM-Roberta-Large-Vit-B-32", ModelType.TEXTUAL, invalid)
 
     async def test_raises_exception_if_unknown_model_name(self) -> None:
         model_cache = ModelCache()
 
         with pytest.raises(ValueError):
-            await model_cache.get("test_model_name", ModelType.TEXTUAL, ModelTask.SEARCH)
+            model_cache.get("test_model_name", ModelType.TEXTUAL, ModelTask.SEARCH)
 
     async def test_preloads_clip_models(self, monkeypatch: MonkeyPatch, mock_get_model: mock.Mock) -> None:
         os.environ["MACHINE_LEARNING_PRELOAD__CLIP__TEXTUAL"] = "ViT-B-32__openai"
@@ -1528,3 +1626,64 @@ class TestPredictionEndpoints:
             parsed_embedding = orjson.loads(embedding)
             assert np.allclose(expected_face["embedding"], parsed_embedding)
             assert np.allclose(expected_face["score"], actual_face["score"])
+
+
+@pytest.mark.asyncio
+async def test_waits_for_every_entry_before_failing_the_request(mocker: MockerFixture) -> None:
+    finished = threading.Event()
+
+    def model(task: ModelTask, load: Callable[[], None]) -> mock.Mock:
+        stub = mock.Mock(
+            spec=InferenceModel, depends=[], loaded=False, load_attempts=0, identity=(ModelType.VISUAL, task)
+        )
+        stub.load.side_effect = load
+        return stub
+
+    def slow() -> None:
+        time.sleep(0.2)
+        finished.set()
+
+    models = {
+        ModelTask.SEARCH: model(ModelTask.SEARCH, slow),
+        ModelTask.OCR: model(ModelTask.OCR, mock.Mock(side_effect=RuntimeError("fails at once"))),
+    }
+    mocker.patch("immich_ml.main.model_cache.get", side_effect=lambda name, type, task, **options: models[task])
+    entries: Any = [
+        {"name": "a", "task": ModelTask.SEARCH, "type": ModelType.VISUAL, "options": {}},
+        {"name": "b", "task": ModelTask.OCR, "type": ModelType.DETECTION, "options": {}},
+    ]
+
+    with ThreadPoolExecutor(2) as pool, pytest.raises(RuntimeError, match="fails at once"):
+        mocker.patch("immich_ml.main.thread_pool", pool)
+        await run_request("text", (entries, []))
+
+    assert finished.is_set()  # or it would still be loading with no request counted as active
+
+
+@pytest.mark.asyncio
+async def test_returns_what_a_preload_freed(mocker: MockerFixture) -> None:
+    events: list[str] = []
+    mocker.patch.object(settings, "preload", PreloadModelData())
+    mocker.patch("immich_ml.main.preload_models", side_effect=lambda _: events.append("preloaded"))
+    mocker.patch("immich_ml.main.allocator.release", side_effect=lambda: events.append("released"))
+
+    async with lifespan(app):
+        assert events == ["preloaded", "released"]
+
+
+@pytest.mark.asyncio
+async def test_returns_memory_once_requests_stop_unless_another_arrives(mocker: MockerFixture) -> None:
+    mocker.patch("immich_ml.main.release", None)
+    loop = mocker.patch("immich_ml.main.asyncio.get_running_loop").return_value
+    first, second, third = update_state(), update_state(), update_state()
+    await anext(first)
+    await anext(second)
+
+    await first.aclose()
+    loop.call_later.assert_not_called()
+    await second.aclose()
+    loop.call_later.assert_called_once_with(5, allocator.release)
+
+    await anext(third)
+    loop.call_later.return_value.cancel.assert_called_once()
+    await third.aclose()

machine-learning/uv.lock

@@ -13,14 +13,8 @@ resolution-markers = [
     "(python_full_version < '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version < '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')",
 ]
 
-[[package]]
-name = "aiocache"
-version = "0.12.3"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/7a/64/b945b8025a9d1e6e2138845f4022165d3b337f55f50984fbc6a4c0a1e355/aiocache-0.12.3.tar.gz", hash = "sha256:f528b27bf4d436b497a1d0d1a8f59a542c153ab1e37c3621713cb376d44c4713", size = 132196, upload-time = "2024-09-25T13:20:23.823Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/37/d7/15d67e05b235d1ed8c3ce61688fe4d84130e72af1657acadfaac3479f4cf/aiocache-0.12.3-py2.py3-none-any.whl", hash = "sha256:889086fc24710f431937b87ad3720a289f7fc31c4fd8b68e9f918b9bacd8270d", size = 28199, upload-time = "2024-09-25T13:20:22.688Z" },
-]
+[manifest]
+overrides = [{ name = "opencv-python", marker = "sys_platform == 'never'" }]
 
 [[package]]
 name = "annotated-doc"
@@ -685,7 +679,6 @@ name = "immich-ml"
 version = "3.2.0"
 source = { editable = "." }
 dependencies = [
-    { name = "aiocache" },
     { name = "fastapi" },
     { name = "gunicorn" },
     { name = "huggingface-hub" },
@@ -766,7 +759,6 @@ types = [
 
 [package.metadata]
 requires-dist = [
-    { name = "aiocache", specifier = ">=0.12.1,<1.0" },
     { name = "fastapi", specifier = ">=0.95.2,<1.0" },
     { name = "gunicorn", specifier = ">=21.1.0" },
     { name = "huggingface-hub", specifier = ">=1.0,<2.0" },
@@ -775,7 +767,7 @@ requires-dist = [
     { name = "onnxruntime", marker = "extra == 'armnn'", specifier = ">=1.23.2,<2" },
     { name = "onnxruntime", marker = "extra == 'cpu'", specifier = ">=1.23.2,<2" },
     { name = "onnxruntime", marker = "extra == 'rknn'", specifier = ">=1.23.2,<2" },
-    { name = "onnxruntime-gpu", marker = "extra == 'cuda'", specifier = ">=1.23.2,<2" },
+    { name = "onnxruntime-gpu", marker = "extra == 'cuda'", specifier = ">=1.27.1,<1.28", index = "https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/" },
     { name = "onnxruntime-migraphx", marker = "extra == 'rocm'", specifier = ">=1.23.2,<2" },
     { name = "onnxruntime-openvino", marker = "extra == 'openvino'", specifier = ">=1.24.1,<2" },
     { name = "opencv-python-headless", specifier = ">=4.7.0.72,<5.0" },
@@ -1326,8 +1318,8 @@ wheels = [
 
 [[package]]
 name = "onnxruntime-gpu"
-version = "1.26.0"
-source = { registry = "https://pypi.org/simple" }
+version = "1.27.1"
+source = { registry = "https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/" }
 dependencies = [
     { name = "flatbuffers" },
     { name = "numpy" },
@@ -1335,14 +1327,14 @@ dependencies = [
     { name = "protobuf" },
 ]
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/94/fd/59bee7cffaa435da44fefdeb63e29c61de4dbfa4b279852f59cd02c042ae/onnxruntime_gpu-1.26.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3c01119ed4d9449d60367fa8ccffcd02bd3fe736754284e4b198d131f54edad6", size = 276971796, upload-time = "2026-05-08T19:15:46.192Z" },
-    { url = "https://files.pythonhosted.org/packages/a4/e4/9b378a5466ea0bed65e5beb8e09254973c580a6522810a38afbcc45e5105/onnxruntime_gpu-1.26.0-cp312-cp312-win_amd64.whl", hash = "sha256:5f49c44689894650990e4c8a857d2edafc276fbd79bba57ceb224bd18d25d491", size = 226548963, upload-time = "2026-05-08T19:09:34.925Z" },
-    { url = "https://files.pythonhosted.org/packages/dd/97/fe8979f44b9275654b42f7bb556e30789b71a1b22998c83b540df2b1b774/onnxruntime_gpu-1.26.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:cfda2fad535595bfc3e570eb588092717711dcb2957656d814695e0c9ceb1508", size = 276974871, upload-time = "2026-05-08T19:15:58.052Z" },
-    { url = "https://files.pythonhosted.org/packages/67/3f/59f1777a394625ecc9a85636de57dc47c25dbb5f888da050f1463955a0ce/onnxruntime_gpu-1.26.0-cp313-cp313-win_amd64.whl", hash = "sha256:6ab9f9c741d2e239b2e321ab0d389c04329d4ab7f11e3b92dd3aa7db1c59dee4", size = 226548083, upload-time = "2026-05-08T19:09:44.408Z" },
-    { url = "https://files.pythonhosted.org/packages/89/96/360328e3c463f7ea08e853c4239c397e83363dd0204de71a710dc1a544bd/onnxruntime_gpu-1.26.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bcf6f347cad9f88a9a625c2b352cf9de927528aedb627ebbc089a201f1990b94", size = 276992052, upload-time = "2026-05-08T19:16:09.892Z" },
-    { url = "https://files.pythonhosted.org/packages/fd/c8/aa2dc0e79bba577f37d5448bcb32fea79977e07506684d8138c19a0f1077/onnxruntime_gpu-1.26.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9e6e4fb1ec9ae1cf456534d9115f106ab2a1ae96fa513b4ed0f4795302b4a2c6", size = 276978254, upload-time = "2026-05-08T19:16:22.096Z" },
-    { url = "https://files.pythonhosted.org/packages/41/e7/923298431e669567d7ccc2a4c898b6534a47641a051569fd97165fe6d9b8/onnxruntime_gpu-1.26.0-cp314-cp314-win_amd64.whl", hash = "sha256:3e592439b0183d303c2374517b5b392599a3d50b2dc9de949b9b15731ac921c9", size = 229142768, upload-time = "2026-05-08T19:09:54.589Z" },
-    { url = "https://files.pythonhosted.org/packages/97/91/93ffe5431d154989f5e04864a25a97eea480997d771232bcbbc538188241/onnxruntime_gpu-1.26.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:56dc7b73954ff4bdc71f5b8ab306b6f61be5d007881b6ef423a609e2b9cd088b", size = 276991545, upload-time = "2026-05-08T19:16:33.347Z" },
+    { url = "https://aiinfra.pkgs.visualstudio.com/2692857e-05ef-43b4-ba9c-ccf1c22c437c/_packaging/9387c3aa-d9ad-4513-968c-383f6f7f53b8/pypi/download/onnxruntime-gpu/1.27.1/onnxruntime_gpu-1.27.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b258525fb253b97a5c215aeb9bc40f32fd2f562b7cf64c0992bc09793af7e4ba", size = 304202950, upload-time = "2026-07-09T19:59:11Z" },
+    { url = "https://aiinfra.pkgs.visualstudio.com/2692857e-05ef-43b4-ba9c-ccf1c22c437c/_packaging/9387c3aa-d9ad-4513-968c-383f6f7f53b8/pypi/download/onnxruntime-gpu/1.27.1/onnxruntime_gpu-1.27.1-cp312-cp312-win_amd64.whl", hash = "sha256:1f4a0b56de1510db3d418bf00760224a00dbbbbb52114ff6a85bb3a5e2e279d4", size = 240898739, upload-time = "2026-07-09T19:56:11Z" },
+    { url = "https://aiinfra.pkgs.visualstudio.com/2692857e-05ef-43b4-ba9c-ccf1c22c437c/_packaging/9387c3aa-d9ad-4513-968c-383f6f7f53b8/pypi/download/onnxruntime-gpu/1.27.1/onnxruntime_gpu-1.27.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:cccabc89bc1c146849e09e482ee5435709c5d7f2e6bc743f4ad68c1733c23c15", size = 304194132, upload-time = "2026-07-09T19:57:24Z" },
+    { url = "https://aiinfra.pkgs.visualstudio.com/2692857e-05ef-43b4-ba9c-ccf1c22c437c/_packaging/9387c3aa-d9ad-4513-968c-383f6f7f53b8/pypi/download/onnxruntime-gpu/1.27.1/onnxruntime_gpu-1.27.1-cp313-cp313-win_amd64.whl", hash = "sha256:da603fef4bc9ffbcf057714f086757c54a295c8bb13976be37077b8bff1810b7", size = 240893326, upload-time = "2026-07-09T19:56:30Z" },
+    { url = "https://aiinfra.pkgs.visualstudio.com/2692857e-05ef-43b4-ba9c-ccf1c22c437c/_packaging/9387c3aa-d9ad-4513-968c-383f6f7f53b8/pypi/download/onnxruntime-gpu/1.27.1/onnxruntime_gpu-1.27.1-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:684cc8b3114da29927db3db7d9b007b579ec3cd283af4e9bc213ee00984d4be2", size = 304224252, upload-time = "2026-07-09T19:58:48Z" },
+    { url = "https://aiinfra.pkgs.visualstudio.com/2692857e-05ef-43b4-ba9c-ccf1c22c437c/_packaging/9387c3aa-d9ad-4513-968c-383f6f7f53b8/pypi/download/onnxruntime-gpu/1.27.1/onnxruntime_gpu-1.27.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:72a0d1490c8c6ca8458178dd125968045bbed1b43f6bb617cab54ed74e147a2e", size = 304202506, upload-time = "2026-07-09T19:58:00Z" },
+    { url = "https://aiinfra.pkgs.visualstudio.com/2692857e-05ef-43b4-ba9c-ccf1c22c437c/_packaging/9387c3aa-d9ad-4513-968c-383f6f7f53b8/pypi/download/onnxruntime-gpu/1.27.1/onnxruntime_gpu-1.27.1-cp314-cp314-win_amd64.whl", hash = "sha256:e879c96037fe282157fce15cedbc0e998dfdac17f75ee6cd455d08fad6b49dd3", size = 243537318, upload-time = "2026-07-09T19:55:48Z" },
+    { url = "https://aiinfra.pkgs.visualstudio.com/2692857e-05ef-43b4-ba9c-ccf1c22c437c/_packaging/9387c3aa-d9ad-4513-968c-383f6f7f53b8/pypi/download/onnxruntime-gpu/1.27.1/onnxruntime_gpu-1.27.1-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0c32338c64cab70fa6fddc906564f26c7a6ebd8b5ffd7704c6be396e31965b09", size = 304220218, upload-time = "2026-07-09T19:58:25Z" },
 ]
 
 [[package]]
@@ -1384,17 +1376,9 @@ name = "opencv-python"
 version = "4.11.0.86"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "numpy" },
+    { name = "numpy", marker = "(platform_machine != 'aarch64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux')" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/17/06/68c27a523103dad5837dc5b87e71285280c4f098c60e4fe8a8db6486ab09/opencv-python-4.11.0.86.tar.gz", hash = "sha256:03d60ccae62304860d232272e4a4fda93c39d595780cb40b161b310244b736a4", size = 95171956, upload-time = "2025-01-16T13:52:24.737Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/05/4d/53b30a2a3ac1f75f65a59eb29cf2ee7207ce64867db47036ad61743d5a23/opencv_python-4.11.0.86-cp37-abi3-macosx_13_0_arm64.whl", hash = "sha256:432f67c223f1dc2824f5e73cdfcd9db0efc8710647d4e813012195dc9122a52a", size = 37326322, upload-time = "2025-01-16T13:52:25.887Z" },
-    { url = "https://files.pythonhosted.org/packages/3b/84/0a67490741867eacdfa37bc18df96e08a9d579583b419010d7f3da8ff503/opencv_python-4.11.0.86-cp37-abi3-macosx_13_0_x86_64.whl", hash = "sha256:9d05ef13d23fe97f575153558653e2d6e87103995d54e6a35db3f282fe1f9c66", size = 56723197, upload-time = "2025-01-16T13:55:21.222Z" },
-    { url = "https://files.pythonhosted.org/packages/f3/bd/29c126788da65c1fb2b5fb621b7fed0ed5f9122aa22a0868c5e2c15c6d23/opencv_python-4.11.0.86-cp37-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1b92ae2c8852208817e6776ba1ea0d6b1e0a1b5431e971a2a0ddd2a8cc398202", size = 42230439, upload-time = "2025-01-16T13:51:35.822Z" },
-    { url = "https://files.pythonhosted.org/packages/2c/8b/90eb44a40476fa0e71e05a0283947cfd74a5d36121a11d926ad6f3193cc4/opencv_python-4.11.0.86-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6b02611523803495003bd87362db3e1d2a0454a6a63025dc6658a9830570aa0d", size = 62986597, upload-time = "2025-01-16T13:52:08.836Z" },
-    { url = "https://files.pythonhosted.org/packages/fb/d7/1d5941a9dde095468b288d989ff6539dd69cd429dbf1b9e839013d21b6f0/opencv_python-4.11.0.86-cp37-abi3-win32.whl", hash = "sha256:810549cb2a4aedaa84ad9a1c92fbfdfc14090e2749cedf2c1589ad8359aa169b", size = 29384337, upload-time = "2025-01-16T13:52:13.549Z" },
-    { url = "https://files.pythonhosted.org/packages/a4/7d/f1c30a92854540bf789e9cd5dde7ef49bbe63f855b85a2e6b3db8135c591/opencv_python-4.11.0.86-cp37-abi3-win_amd64.whl", hash = "sha256:085ad9b77c18853ea66283e98affefe2de8cc4c1f43eda4c100cf9b2721142ec", size = 39488044, upload-time = "2025-01-16T13:52:21.928Z" },
-]
 
 [[package]]
 name = "opencv-python-headless"
@@ -1933,7 +1917,7 @@ dependencies = [
     { name = "colorlog" },
     { name = "numpy" },
     { name = "omegaconf" },
-    { name = "opencv-python" },
+    { name = "opencv-python", marker = "sys_platform == 'never'" },
     { name = "pillow" },
     { name = "pyclipper" },
     { name = "pyyaml" },