fix(scripts): check for ORT modules

ssube · Mar 18, 2023 · c465b61 · c465b61
1 parent 84718e5
commit c465b61
Showing 1 changed file with 17 additions and 1 deletion.
diff --git a/api/scripts/check-env.py b/api/scripts/check-env.py
@@ -3,7 +3,7 @@
 
 REQUIRED_MODULES = ["onnx", "diffusers", "safetensors", "torch"]
 
-REQUIRED_RUNTIME = [
+RUNTIME_MODULES = [
     "onnxruntime",
     "onnxruntime_gpu",
     "onnxruntime_rocm",
@@ -26,6 +26,21 @@ def check_modules() -> List[str]:
     return results
 
 
+def check_runtimes() -> List[str]:
+    results = []
+    for name in RUNTIME_MODULES:
+        try:
+            __import__(name)
+            module_version = version(name)
+            results.append(
+                f"runtime module {name} is present at version {module_version}"
+            )
+        except ImportError as e:
+            results.append(f"unable to import runtime module {name}: {e}")
+
+    return results
+
+
 def check_providers() -> List[str]:
     results = []
     try:
@@ -46,6 +61,7 @@ def check_providers() -> List[str]:
 
 ALL_CHECKS = [
     check_modules,
+    check_runtimes,
     check_providers,
 ]