Immich Machine Learning 模块解析：uv 环境配置、ONNX Runtime 硬件加速与 Locust 推理压测实践

Immich 的机器学习模块（machine-learning/ 目录）是整个自托管照片/视频管理系统背后的"推理引擎"：它提供 CLIP 视觉与文本向量编码、人脸检测/识别等能力，供服务端任务队列消费。本篇以 machine-learning/README.md 为主体，完整覆盖其 uv 环境搭建、依赖管理、Locust 压测方法与 InsightFace 模型许可边界，并结合 pyproject.toml、immich_ml/main.py、immich_ml/sessions/ort.py 等源码，把 README 中的每一步操作落到可验证的实现细节上，帮助读者独立完成本地开发环境搭建、推理端点压测与硬件加速选型。

模块职责与技术栈概览

README 开篇即声明了该模块的两大核心能力：

• CLIP embeddings：将图像与文本编码到同一向量空间，支撑 Immich 的语义搜索；
• Facial recognition：人脸检测与特征向量提取，支撑人脸聚类功能。

从源码结构看，这是一套以 FastAPI + Gunicorn/Uvicorn 为服务框架、以 ONNX Runtime 为推理运行时、以 aiocache 内存缓存 管理模型生命周期的独立 HTTP 服务：

• 入口为 immich_ml/main.py，它通过 subprocess 拉起 gunicorn immich_ml.main:app，工作进程类为自定义的 CustomUvicornWorker，并读取 gunicorn_conf.py 日志配置；
• 应用主体 immich_ml/main.py 暴露 GET /、GET /ping 与核心的 POST /predict 端点；
• 核心依赖在 pyproject.toml 中声明：fastapi、onnx、onnxruntime-*（按设备条件引入）、opencv-python-headless、huggingface-hub、tokenizers、rapidocr 等，要求 Python >=3.11,<4.0。

服务默认监听地址为 [::]:3003（见 config.py 中 NonPrefixedSettings 的 immich_host/immich_port 默认值），这与部署时 docker/docker-compose.yml 中 immich-machine-learning 服务的定位一致。

环境搭建：基于 uv 的隔离虚拟环境

README 的 Setup 一节指明：项目使用 uv 管理依赖，需先安装 uv，再执行：

uv sync --extra cpu

这会在隔离的虚拟环境中安装 CPU 推理所需的全部依赖。若要使用硬件加速 API，将 --extra cpu 替换为对应设备 extras：

uv sync --extra cuda      # NVIDIA GPU
uv sync --extra rocm      # AMD GPU（MIGraphX）
uv sync --extra openvino  # Intel 设备

README 特别提醒：CUDA 路径要求 GPU 的 compute capability ≥ 5.2。

对照 pyproject.toml 中 [project.optional-dependencies] 的定义，每个 extra 实际拉入的 onnxruntime 发行版如下，这是"一套代码、多后端运行时"的实现基础：

extra	引入的包	说明
cpu	onnxruntime>=1.23.2,<2	纯 CPU 推理
cuda	onnxruntime-gpu>=1.23.2,<2	NVIDIA GPU
rocm	onnxruntime-migraphx>=1.23.2,<2	AMD GPU（MIGraphX EP）
openvino	onnxruntime-openvino>=1.24.1,<2	Intel 设备
armnn	onnxruntime	ARM Mali（配合仓库内置的 ann/ C++ 扩展）
rknn	onnxruntime + rknn-toolkit-lite2>=2.3.0,<3	Rockchip NPU

依赖的日常维护命令 README 也给出了：

uv add $PACKAGE_NAME    # 添加依赖
uv remove $PACKAGE_NAME # 移除依赖
uv lock                 # 重新生成锁定文件

并要求将 uv.lock 与 pyproject.toml 一并提交——仓库根目录下确实存在 uv.lock，说明该约定在项目中是强制执行的，以保证 CI 与 Docker 构建（Dockerfile 中以 uv sync --frozen --extra ${DEVICE} 消费锁文件）获得可复现的环境。

服务启动与关键运行参数

python -m immich_ml 启动后（即 Dockerfile 的 CMD），Gunicorn 按 main.py 传参运行：-w 对应 settings.workers，-t 对应 settings.worker_timeout，--keep-alive 对应 http_keepalive_timeout_s，--graceful-timeout 10。

所有这些参数均可通过环境变量覆盖。config.py 中的 Settings 使用 pydantic-settings，env_prefix="MACHINE_LEARNING_"、嵌套分隔符 __，常用项与默认值如下（均可作为部署调优的入口）：

环境变量	默认值	作用
MACHINE_LEARNING_WORKERS	1	Gunicorn worker 数
MACHINE_LEARNING_WORKER_TIMEOUT	300（ROCm 设备下 900，见 default_worker_timeout）	worker 超时（秒）
MACHINE_LEARNING_MODEL_TTL	300	模型空闲卸载时间（秒），0 表示禁用卸载
MACHINE_LEARNING_MODEL_TTL_POLL_S	10	空闲检测轮询间隔
MACHINE_LEARNING_REQUEST_THREADS	CPU 核数	推理线程池大小（绕过 asyncio 阻塞瓶颈）
MACHINE_LEARNING_MODEL_INTRA_OP_THREADS / INTER_OP	0（自动）	ONNX Runtime 线程数覆盖
MACHINE_LEARNING_MODEL_ARENA	true	是否启用 CPU 内存 arena（CPU 镜像在 Dockerfile 中显式关闭）
MACHINE_LEARNING_DEVICE_ID	0	GPU/设备编号
MACHINE_LEARNING_CACHE_FOLDER	~/.cache/immich_ml	模型下载缓存目录（容器内设为 /cache）
MACHINE_LEARNING_PRELOAD__*	未设置	启动时预热模型，如 MACHINE_LEARNING_PRELOAD__CLIP__VISUAL 等
MACHINE_LEARNING_MAX_BATCH_SIZE__*	未设置	限制人脸/OCR 批量推理的 batch 上限

两个值得注意的源码级机制：

1. 请求线程池：main.py 在 lifespan 中创建 ThreadPoolExecutor(settings.request_threads)，所有阻塞推理通过 run_in_executor 调度，注释明确指出 "asyncio is a huge bottleneck for performance"；
2. 空闲自杀：idle_shutdown_task 轮询检查——当无活跃请求、无模型加载锁、且距上次请求超过 model_ttl 时，进程向自己发送 SIGINT 优雅退出。这对按需拉起的部署形态（例如无定时任务时 ML 容器可以自行释放内存）非常关键。preload 配置则可以在启动时提前加载指定模型，避免首次请求的冷启动延迟，加载逻辑见 preload_models。

推理接口：POST /predict 的请求模型

压测文件与实现共同揭示了统一推理端点的协议形态。/predict 接受 multipart/form-data：

• entries：一段 JSON，描述"任务 → 类型 → 模型与选项"的推理计划；
• image：待推理图片（视觉任务必需）；
• text：待编码文本（CLIP 文本任务使用）。

run_inference 将 entries 拆为 without_deps 与 with_deps 两组并行执行，依赖项（例如人脸识别依赖人脸检测的输出）排在依赖组之后，通过模型输出字典串联——这解释了人脸请求为何必须同时下发 detection 与 recognition 两个条目。模型实例由 ModelCache 以乐观锁方式从内存缓存获取/创建，缓存键为 model_name + type + task，并按 model_ttl 过期，与上文"空闲自杀"共同构成模型的完整生命周期管理。

人脸识别的实现类 FaceRecognizer 声明了 depends = [(ModelType.DETECTION, ModelTask.FACIAL_RECOGNITION)]，并对超过 batch_size 的人脸裁剪自动分批推理；当模型输入不含动态 batch 维时，它会用 onnx.tools.update_model_dims 就地改写 ONNX 模型添加 batch 轴。

压测：Locust 推理吞吐与延迟测量

README 的 Load Testing 一节完整给出了方法：使用 Locust 加载现成的 locustfile.py，因为 Locust 直接查询模型端点并聚合统计，被测服务必须先部署好。启动命令：

locust --web-host 127.0.0.1

然后浏览器打开 localhost:8089 访问控制台 UI，可在界面上调整并发用户数、启动速率等。

locustfile.py 通过命令行参数暴露了可调的实验变量，默认值如下：

参数	默认值	含义
--clip-model	ViT-B-32::openai	用于编码的 CLIP 模型名
--face-model	buffalo_l	人脸检测/识别模型名
--face-min-score	0.034	人脸检测最低置信度；README 注释指出该默认值大约"每请求返回 1 张人脸"，设为 0 会把返回的人脸数放大到成千上万
--image-size	1000	测试图片边长（像素），测试图是程序生成的纯色 JPEG

压测流量模型由三个用户类构成，全部 POST 到 http://127.0.0.1:3003/predict（host 定义）：

• CLIPTextFormDataLoadTest：发送 clip.textual 条目 + text 表单字段，模拟语义搜索的文本编码；
• CLIPVisionFormDataLoadTest：发送 clip.visual 条目 + image 文件，模拟图片向量编码；
• RecognitionFormDataLoadTest：一次请求同时携带 facial-recognition.detection 与 facial-recognition.recognition 两个条目，且 detection 带上 minScore 选项，完整模拟生产环境的人脸推理管线。

README 特别解释了 Locust 的并发术语换算，这也是实际压测中最容易搞错的一点：Locust 的 users 是全局并发用户数，每个用户一次只执行一个任务。要得到"每个端点固定 N 个并发请求"，应令 users = N × 端点数。按 locustfile 中的三个用户类计算，若希望每个端点同时有 8 个请求在途，用户数应设为 8 × 3 = 24。

硬件加速执行提供商与镜像形态

README 声明支持 CUDA、ROCm 与 OpenVINO 三类加速 API。在源码中，提供商优先级集中定义在 constants.py 的 SUPPORTED_PROVIDERS：

SUPPORTED_PROVIDERS = [
    "CUDAExecutionProvider",
    "MIGraphXExecutionProvider",
    "OpenVINOExecutionProvider",
    "CoreMLExecutionProvider",
    "CPUExecutionProvider",
]

OrtSession 构造时按该顺序过滤当前 onnxruntime 中实际可用的提供商，因此同一份模型代码在 CPU/CUDA/ROCm/OpenVINO 环境下无差别运行。各提供商的会话选项由 provider_options 分支生成，几个关键细节：

• CUDA：透传 device_id（来自 MACHINE_LEARNING_DEVICE_ID），内存 arena 策略设为 kSameAsRequested；
• MIGraphX（ROCm）：自动创建模型目录下的 migraphx 缓存目录（否则运行时会崩溃），并按 MACHINE_LEARNING_ROCM_PRECISION 决定 FP16 开关；ROCm 场景下 OrtSession.run 还会用模型级锁串行化首次推理（MIGraphX 首次运行需在线编译，见 run 方法）；
• OpenVINO：优先选择 GPU.{device_id}，无 GPU 时回退 CPU，精度由 MACHINE_LEARNING_OPENVINO_PRECISION 控制，缓存目录为模型旁的 openvino 目录；
• CPU 线程策略：_sess_options_default 仅在纯 CPU 提供商下默认 inter_op=1、intra_op=2（注释说明 ORT 默认线程数在 GPU 场景会造成瓶颈），显式配置 0 表示不干预。

部署侧的对应关系见 docker/docker-compose.yml：镜像基础标签 ghcr.io/immich-app/immich-machine-learning:${IMMICH_VERSION:-release}，需要加速时"向 image tag 追加 -cuda/-rocm/-openvino/-armnn/-rknn 之一"，并可通过 hwaccel.ml.yml 扩展服务定义（compose 注释中给出的做法）；更完整的硬件加速配置说明可参阅仓库内文档 ml-hardware-acceleration。Dockerfile 以构建参数 ARG DEVICE=cpu 驱动整套多阶段构建：builder 阶段用对应 extra 执行 uv sync --frozen，prod-${DEVICE} 阶段分别安装 CUDA 12.2 runtime + cuDNN 9.10（注释注明 9.11 起弃用 Pascal 支持）、ROCm 的 migraphx/MIOpen、Intel OpenCL/IGC 组件等；所有 prod 镜像统一将模型与 Hugging Face 缓存指向 /cache（MACHINE_LEARNING_CACHE_FOLDER=/cache、HF_HOME=/cache/hf-cache），与 compose 中的 model-cache 卷一一对应，并以内置 healthcheck.py 做健康检查。

人脸识别模型与许可边界

README 的 Facial Recognition 一节说明模型来自 InsightFace 项目的 model zoo，明确列出了四个可用模型组：antelopev2、buffalo_l、buffalo_m、buffalo_s。这与源码 constants.py 中白名单 _INSIGHTFACE_MODELS 完全一致——服务端只接受这四个名字（经 clean_name 归一化后）作为 InsightFace 来源的模型，再由 get_model_source 分派到对应下载源；locustfile 中压测使用的默认 buffalo_l 也在该集合内。

许可条款方面，README 给出了明确的法律事实：项目方于 2023 年 3 月 18 日 通过邮件获得 Jia Guo（[e-mail]）对 InsightFace 人脸模型在本项目内使用的授权；但该授权不延伸至第三方对这些模型的分发或商业使用，使用者应自行遵守 InsightFace 仓库给出的许可条款。任何将 Immich ML 集成到商业产品、或自行二次分发这些模型文件的场景，都应以该条款为约束边界。

从能力链路看，detection 输出（框、关键点、置信度）会被 recognition 阶段消费：FaceRecognizer._predict 对每张检测到的人脸执行对齐裁剪（align_face）、归一化后批量提取 embedding，最终 postprocess 输出 boundingBox + embedding + score 结构，供上层数据库存储与聚类。

小结与关键路径

本篇围绕 machine-learning/README.md 的三个核心章节展开，并落到仓库内的可验证实现：

• 环境：uv sync --extra {cpu|cuda|rocm|openvino}，extras 与 onnxruntime 发行版的映射定义在 pyproject.toml，锁文件提交是硬约定；
• 服务与调优：MACHINE_LEARNING_* 环境变量族覆盖 worker、TTL、线程池、预加载与批大小，实现在 config.py，空闲自杀与线程池调度在 main.py；
• 压测：locust --web-host 127.0.0.1 + locustfile.py，记住 users = 每端点并发 × 端点数 的换算；
• 加速：执行提供商优先级在 constants.py，会话选项与线程策略在 ort.py，容器形态由 Dockerfile 的 DEVICE 构建参数与 docker-compose.yml 的镜像 tag 后缀共同决定；
• 合规：InsightFace 四个模型组的授权范围仅限本项目内部使用，禁止第三方分发与商用。