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
10 changed files with 120 additions and 290 deletions
+30 -4
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(