diff --git a/machine-learning/Dockerfile b/machine-learning/Dockerfile index fa24b0f87..b1ca31c88 100644 --- a/machine-learning/Dockerfile +++ b/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"] diff --git a/machine-learning/conftest.py b/machine-learning/conftest.py index a5acc059d..a968d6906 100644 --- a/machine-learning/conftest.py +++ b/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 diff --git a/machine-learning/immich_ml/__main__.py b/machine-learning/immich_ml/__main__.py index f305c7765..49e37aae9 100644 --- a/machine-learning/immich_ml/__main__.py +++ b/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: diff --git a/machine-learning/immich_ml/allocator.py b/machine-learning/immich_ml/allocator.py new file mode 100644 index 000000000..7cb3338bc --- /dev/null +++ b/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) diff --git a/machine-learning/immich_ml/config.py b/machine-learning/immich_ml/config.py index d1e44fd70..747d1a498 100644 --- a/machine-learning/immich_ml/config.py +++ b/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: diff --git a/machine-learning/immich_ml/main.py b/machine-learning/immich_ml/main.py index 4ac5b020a..fdad03eae 100644 --- a/machine-learning/immich_ml/main.py +++ b/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 diff --git a/machine-learning/immich_ml/models/__init__.py b/machine-learning/immich_ml/models/__init__.py index 6c3a74c96..f93939130 100644 --- a/machine-learning/immich_ml/models/__init__.py +++ b/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 diff --git a/machine-learning/immich_ml/models/base.py b/machine-learning/immich_ml/models/base.py index 5c40b21f3..a54a80485 100644 --- a/machine-learning/immich_ml/models/base.py +++ b/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 diff --git a/machine-learning/immich_ml/models/cache.py b/machine-learning/immich_ml/models/cache.py index d8f9ca81b..1a8dad2eb 100644 --- a/machine-learning/immich_ml/models/cache.py +++ b/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() diff --git a/machine-learning/immich_ml/models/clip/textual.py b/machine-learning/immich_ml/models/clip/textual.py index 421459fc0..003d30e56 100644 --- a/machine-learning/immich_ml/models/clip/textual.py +++ b/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) diff --git a/machine-learning/immich_ml/models/clip/visual.py b/machine-learning/immich_ml/models/clip/visual.py index 1d46d8fe1..37362fc53 100644 --- a/machine-learning/immich_ml/models/clip/visual.py +++ b/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}") diff --git a/machine-learning/immich_ml/models/facial_recognition/recognition.py b/machine-learning/immich_ml/models/facial_recognition/recognition.py index 2b0575ec3..09e48350a 100644 --- a/machine-learning/immich_ml/models/facial_recognition/recognition.py +++ b/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( diff --git a/machine-learning/immich_ml/schemas.py b/machine-learning/immich_ml/schemas.py index 0dcc540d0..a88faac9d 100644 --- a/machine-learning/immich_ml/schemas.py +++ b/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") diff --git a/machine-learning/immich_ml/sessions/ort.py b/machine-learning/immich_ml/sessions/ort.py index 95a40782c..1f887256b 100644 --- a/machine-learning/immich_ml/sessions/ort.py +++ b/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) diff --git a/machine-learning/pyproject.toml b/machine-learning/pyproject.toml index 6af41477b..84b1d2f4e 100644 --- a/machine-learning/pyproject.toml +++ b/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"] diff --git a/machine-learning/test_main.py b/machine-learning/test_main.py index a50bea00b..27191d67f 100644 --- a/machine-learning/test_main.py +++ b/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() diff --git a/machine-learning/uv.lock b/machine-learning/uv.lock index 6ca953035..5c84de4a4 100644 --- a/machine-learning/uv.lock +++ b/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" },