Backupppert/immich
1
1
Fork 0
mirror of https://github.com/immich-app/immich.git synced 2025-07-21 22:44:23 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 29

feat(ml): round-robin device assignment ()

Browse source 
* round-robin device assignment

* docs and tests

clarify doc
This commit is contained in:
Mert 2024-10-07 17:37:45 -04:00 • committed by GitHub
commit bd826b0b9b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 62 additions and 7 deletions
machine-learning/app
test_main.py

16
machine-learning/app/test_main.py
View file

@ -210,10 +210,24 @@ class TestOrtSession:
session = OrtSession(model_path, providers=["OpenVINOExecutionProvider", "CPUExecutionProvider"])
assert session.provider_options == [
{"device_type": "GPU", "precision": "FP32", "cache_dir": "/cache/ViT-B-32__openai/openvino"},
{"device_type": "GPU.0", "precision": "FP32", "cache_dir": "/cache/ViT-B-32__openai/openvino"},
{"arena_extend_strategy": "kSameAsRequested"},
]
def test_sets_device_id_for_openvino(self) -> None:
os.environ["MACHINE_LEARNING_DEVICE_ID"] = "1"
session = OrtSession("ViT-B-32__openai", providers=["OpenVINOExecutionProvider"])
assert session.provider_options[0]["device_type"] == "GPU.1"
def test_sets_device_id_for_cuda(self) -> None:
os.environ["MACHINE_LEARNING_DEVICE_ID"] = "1"
session = OrtSession("ViT-B-32__openai", providers=["CUDAExecutionProvider"])
assert session.provider_options[0]["device_id"] == "1"
def test_sets_provider_options_kwarg(self) -> None:
session = OrtSession(
"ViT-B-32__openai",