feat: ROCm 7.2 and MIGraphX support (#26178)

This commit is contained in:
Kishor Prins 2026-02-26 08:52:26 -08:00 committed by GitHub
parent 177d1c9a30
commit dd9046508d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 120 additions and 290 deletions

View file

@ -179,7 +179,7 @@ class TestOrtSession:
OV_EP = ["OpenVINOExecutionProvider", "CPUExecutionProvider"]
CUDA_EP_OUT_OF_ORDER = ["CPUExecutionProvider", "CUDAExecutionProvider"]
TRT_EP = ["TensorrtExecutionProvider", "CUDAExecutionProvider", "CPUExecutionProvider"]
ROCM_EP = ["ROCMExecutionProvider", "CPUExecutionProvider"]
ROCM_EP = ["MIGraphXExecutionProvider", "CPUExecutionProvider"]
COREML_EP = ["CoreMLExecutionProvider", "CPUExecutionProvider"]
@pytest.mark.providers(CPU_EP)
@ -289,12 +289,38 @@ class TestOrtSession:
assert session.provider_options == [{"arena_extend_strategy": "kSameAsRequested", "device_id": "1"}]
def test_sets_provider_options_for_rocm(self) -> None:
def test_sets_provider_options_for_rocm(self, mocker: MockerFixture) -> None:
model_path = "/cache/ViT-B-32__openai/textual/model.onnx"
os.environ["MACHINE_LEARNING_DEVICE_ID"] = "1"
mkdir = mocker.patch("immich_ml.sessions.ort.Path.mkdir")
session = OrtSession("ViT-B-32__openai", providers=["ROCMExecutionProvider"])
session = OrtSession(model_path, providers=["MIGraphXExecutionProvider"])
assert session.provider_options == [{"arena_extend_strategy": "kSameAsRequested", "device_id": "1"}]
assert session.provider_options == [
{
"device_id": "1",
"migraphx_model_cache_dir": "/cache/ViT-B-32__openai/textual/migraphx",
"migraphx_fp16_enable": "0",
}
]
mkdir.assert_called_once_with(parents=True, exist_ok=True)
def test_sets_rocm_to_fp16_if_enabled(self, path: mock.Mock, mocker: MockerFixture) -> None:
model_path = "/cache/ViT-B-32__openai/textual/model.onnx"
os.environ["MACHINE_LEARNING_DEVICE_ID"] = "1"
mocker.patch.object(settings, "rocm_precision", ModelPrecision.FP16)
mkdir = mocker.patch("immich_ml.sessions.ort.Path.mkdir")
session = OrtSession(model_path, providers=["MIGraphXExecutionProvider"])
assert session.provider_options == [
{
"device_id": "1",
"migraphx_model_cache_dir": "/cache/ViT-B-32__openai/textual/migraphx",
"migraphx_fp16_enable": "1",
}
]
mkdir.assert_called_once_with(parents=True, exist_ok=True)
def test_sets_provider_options_kwarg(self) -> None:
session = OrtSession(