fix(ml): stabilize MIGraphX inference (#28444)

* fix: stabilize ROCm MIGraphX inference

Serialize MIGraphX session runs so lazy compiles cannot overlap within a worker.

Use a fixed face-recognition batch size for MIGraphX to avoid compiling a new program for each detected face count.

* fix(ml): increase ROCm worker timeout

* fix(ml): narrow MIGraphX compile locking

* docs: format environment variables table

* docs: apply prettier to environment variables table

machine-learning/immich_ml/config.py

@@ -6,7 +6,7 @@ from pathlib import Path
 from socket import socket
 
 from gunicorn.arbiter import Arbiter
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 from pydantic_settings import BaseSettings, SettingsConfigDict
 from rich.console import Console
 from rich.logging import RichHandler
@@ -42,6 +42,10 @@ class MaxBatchSize(BaseModel):
     ocr: int | None = None
 
 
+def default_worker_timeout() -> int:
+    return 900 if os.environ.get("DEVICE") == "rocm" else 300
+
+
 class Settings(BaseSettings):
     model_config = SettingsConfigDict(
         env_prefix="MACHINE_LEARNING_",
@@ -54,7 +58,7 @@ class Settings(BaseSettings):
     model_ttl: int = 300
     model_ttl_poll_s: int = 10
     workers: int = 1
-    worker_timeout: int = 300
+    worker_timeout: int = Field(default_factory=default_worker_timeout)
     http_keepalive_timeout_s: int = 2
     test_full: bool = False
     request_threads: int = os.cpu_count() or 4

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

@@ -89,4 +89,10 @@ class FaceRecognizer(InferenceModel):
     @property
     def _batch_size_default(self) -> int | None:
         providers = ort.get_available_providers()
-        return None if self.model_format == ModelFormat.ONNX and "OpenVINOExecutionProvider" not in providers else 1
+        if (
+            self.model_format == ModelFormat.ONNX
+            and "MIGraphXExecutionProvider" not in providers
+            and "OpenVINOExecutionProvider" not in providers
+        ):
+            return None
+        return 1

machine-learning/immich_ml/sessions/ort.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from pathlib import Path
+from threading import Lock
 from typing import Any
 
 import numpy as np
@@ -12,6 +13,37 @@ from immich_ml.schemas import ModelPrecision, SessionNode
 
 from ..config import log, settings
 
+MigraphxInputSignature = tuple[tuple[str, str, tuple[int, ...]], ...]
+
+_migraphx_registry_lock = Lock()
+_migraphx_model_locks: dict[str, Lock] = {}
+_migraphx_compiled_inputs: set[tuple[str, MigraphxInputSignature]] = set()
+
+
+def _migraphx_get_model_lock(model_key: str) -> Lock:
+    with _migraphx_registry_lock:
+        lock = _migraphx_model_locks.get(model_key)
+        if lock is None:
+            lock = Lock()
+            _migraphx_model_locks[model_key] = lock
+        return lock
+
+
+def _migraphx_has_compiled_input(key: tuple[str, MigraphxInputSignature]) -> bool:
+    with _migraphx_registry_lock:
+        return key in _migraphx_compiled_inputs
+
+
+def _migraphx_mark_compiled_input(key: tuple[str, MigraphxInputSignature]) -> None:
+    with _migraphx_registry_lock:
+        _migraphx_compiled_inputs.add(key)
+
+
+def _migraphx_input_signature(
+    input_feed: dict[str, NDArray[np.float32]] | dict[str, NDArray[np.int32]],
+) -> MigraphxInputSignature:
+    return tuple((name, str(value.dtype), tuple(value.shape)) for name, value in sorted(input_feed.items()))
+
 
 class OrtSession:
     session: ort.InferenceSession
@@ -48,7 +80,21 @@ class OrtSession:
         input_feed: dict[str, NDArray[np.float32]] | dict[str, NDArray[np.int32]],
         run_options: Any = None,
     ) -> list[NDArray[np.float32]]:
-        outputs: list[NDArray[np.float32]] = self.session.run(output_names, input_feed, run_options)
+        if "MIGraphXExecutionProvider" in self.providers:
+            model_key = self.model_path.resolve().as_posix()
+            input_key = (model_key, _migraphx_input_signature(input_feed))
+            if not _migraphx_has_compiled_input(input_key):
+                model_lock = _migraphx_get_model_lock(model_key)
+                with model_lock:
+                    if not _migraphx_has_compiled_input(input_key):
+                        outputs: list[NDArray[np.float32]] = self.session.run(output_names, input_feed, run_options)
+                        _migraphx_mark_compiled_input(input_key)
+                        return outputs
+
+            outputs = self.session.run(output_names, input_feed, run_options)
+            return outputs
+
+        outputs = self.session.run(output_names, input_feed, run_options)
         return outputs
 
     @property