Merge pull request #8017 from hauntsaninja/patch-1

Speed up cupy environment duplicate detection
cupy · Dec 20, 2023 · 891869e · 891869e
2 parents 90a566c + ce68149
commit 891869e
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/cupy/_environment.py b/cupy/_environment.py
@@ -451,7 +451,7 @@ def _preload_warning(lib, exc):
 
 def _detect_duplicate_installation():
     # List of all CuPy packages, including out-dated ones.
-    known = [
+    known = {
         'cupy',
         'cupy-cuda80',
         'cupy-cuda90',
@@ -476,10 +476,10 @@ def _detect_duplicate_installation():
         'cupy-rocm-4-2',
         'cupy-rocm-4-3',
         'cupy-rocm-5-0',
-    ]
-    cupy_installed = [
-        name for name in known
-        if list(importlib.metadata.distributions(name=name))]
+    }
+    installed_names = {d.metadata["Name"]
+                       for d in importlib.metadata.distributions()}
+    cupy_installed = known & installed_names
     if 1 < len(cupy_installed):
         cupy_packages_list = ', '.join(sorted(cupy_installed))
         warnings.warn(f'''