helpy4/immich
1
0
Fork 0
mirror of https://github.com/immich-app/immich synced 2025-10-17 18:19:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

use arena by default in native installation

Browse source
This commit is contained in:
mertalev 2025-04-19 16:51:17 -04:00
parent 2018a7ff2d
commit b1b15d4ed6
No known key found for this signature in database
GPG key ID: DF6ABC77AAD98C95
4 changed files with 26 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

12
machine-learning/Dockerfile
View file

@ -70,7 +70,8 @@ RUN if [ "$DEVICE" = "rocm" ]; then \
FROM python:3.11-slim-bookworm@sha256:873f91540d53b36327ed4fb018c9669107a4e2a676719720edb4209c4b15d029 AS prod-cpu
ENV LD_PRELOAD=/usr/lib/libmimalloc.so.2
ENV LD_PRELOAD=/usr/lib/libmimalloc.so.2 \
MACHINE_LEARNING_MODEL_ARENA=false
FROM python:3.11-slim-bookworm@sha256:873f91540d53b36327ed4fb018c9669107a4e2a676719720edb4209c4b15d029 AS prod-openvino
@ -88,7 +89,8 @@ RUN apt-get update && \
FROM nvidia/cuda:12.2.2-runtime-ubuntu22.04@sha256:94c1577b2cd9dd6c0312dc04dff9cb2fdce2b268018abc3d7c2dbcacf1155000 AS prod-cuda
ENV LD_PRELOAD=/usr/lib/libmimalloc.so.2
ENV LD_PRELOAD=/usr/lib/libmimalloc.so.2 \
MACHINE_LEARNING_MODEL_ARENA=false
RUN apt-get update && \
apt-get install --no-install-recommends -yqq libcudnn9-cuda-12 && \
@ -104,7 +106,8 @@ FROM rocm/dev-ubuntu-22.04:6.4.3-complete@sha256:1f7e92ca7e3a3785680473329ed1091
FROM prod-cpu AS prod-armnn
ENV LD_LIBRARY_PATH=/opt/armnn \
LD_PRELOAD=/usr/lib/libmimalloc.so.2
LD_PRELOAD=/usr/lib/libmimalloc.so.2 \
MACHINE_LEARNING_MODEL_ARENA=false
RUN apt-get update && apt-get install -y --no-install-recommends ocl-icd-libopencl1 mesa-opencl-icd libgomp1 && \
rm -rf /var/lib/apt/lists/* && \
@ -127,7 +130,8 @@ FROM prod-cpu AS prod-rknn
# renovate: datasource=github-tags depName=airockchip/rknn-toolkit2
ARG RKNN_TOOLKIT_VERSION="v2.3.0"
ENV LD_PRELOAD=/usr/lib/libmimalloc.so.2
ENV LD_PRELOAD=/usr/lib/libmimalloc.so.2 \
MACHINE_LEARNING_MODEL_ARENA=false
ADD --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/

1
machine-learning/immich_ml/config.py
View file

@ -61,6 +61,7 @@ class Settings(BaseSettings):
request_threads: int = os.cpu_count() or 4
model_inter_op_threads: int = 0
model_intra_op_threads: int = 0
model_arena: bool = True
ann: bool = True
ann_fp16_turbo: bool = False
ann_tuning_level: int = 2

2
machine-learning/immich_ml/sessions/ort.py
View file

@ -123,7 +123,7 @@ class OrtSession:
@property
def _sess_options_default(self) -> ort.SessionOptions:
sess_options = ort.SessionOptions()
sess_options.enable_cpu_mem_arena = False
sess_options.enable_cpu_mem_arena = settings.model_arena
# avoid thread contention between models
if settings.model_inter_op_threads > 0:

16
machine-learning/test_main.py
View file

@ -309,6 +309,22 @@ class TestOrtSession:
assert session.sess_options.inter_op_num_threads == 2
assert session.sess_options.intra_op_num_threads == 4
def test_uses_arena_if_enabled(self, mocker: MockerFixture) -> None:
mock_settings = mocker.patch("immich_ml.sessions.ort.settings", autospec=True)
mock_settings.model_arena = True
session = OrtSession("ViT-B-32__openai", providers=["CPUExecutionProvider"])
assert session.sess_options.enable_cpu_mem_arena
def test_does_not_use_arena_if_disabled(self, mocker: MockerFixture) -> None:
mock_settings = mocker.patch("immich_ml.sessions.ort.settings", autospec=True)
mock_settings.model_arena = False
session = OrtSession("ViT-B-32__openai", providers=["CPUExecutionProvider"])
assert not session.sess_options.enable_cpu_mem_arena
def test_sets_sess_options_kwarg(self) -> None:
sess_options = ort.SessionOptions()
session = OrtSession(