[Core] Speed up single_client_tasks_sync (ray-project#41743)

--------- Signed-off-by: Jiajun Yao <jeromeyjj@gmail.com>
ShuN6211 · Dec 9, 2023 · ab36f94 · ab36f94
1 parent 2f29500
commit ab36f94
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 15 deletions.
diff --git a/python/ray/_private/accelerators/__init__.py b/python/ray/_private/accelerators/__init__.py
@@ -1,4 +1,4 @@
-from typing import Set
+from typing import Set, Optional
 
 from ray._private.accelerators.accelerator import AcceleratorManager
 from ray._private.accelerators.nvidia_gpu import NvidiaGPUAcceleratorManager
@@ -29,25 +29,34 @@ def get_all_accelerator_resource_names() -> Set[str]:
     }
 
 
-_resource_name_to_accelerator_manager = {
-    accelerator_manager.get_resource_name(): accelerator_manager
-    for accelerator_manager in get_all_accelerator_managers()
-}
-
-
-def get_accelerator_manager_for_resource(resource_name: str) -> AcceleratorManager:
+def get_accelerator_manager_for_resource(
+    resource_name: str,
+) -> Optional[AcceleratorManager]:
     """Get the corresponding accelerator manager for the given
     accelerator resource name
 
     E.g., TPUAcceleratorManager is returned if resource name is "TPU"
     """
-    if resource_name == "GPU":
+    try:
+        return get_accelerator_manager_for_resource._resource_name_to_accelerator_manager.get(  # noqa: E501
+            resource_name, None
+        )
+    except AttributeError:
+        # Lazy initialization.
+        resource_name_to_accelerator_manager = {
+            accelerator_manager.get_resource_name(): accelerator_manager
+            for accelerator_manager in get_all_accelerator_managers()
+        }
+        # Special handling for GPU resource name since multiple accelerator managers
+        # have the same GPU resource name.
         if IntelGPUAcceleratorManager.get_current_node_num_accelerators() > 0:
-            return IntelGPUAcceleratorManager
+            resource_name_to_accelerator_manager["GPU"] = IntelGPUAcceleratorManager
         else:
-            return NvidiaGPUAcceleratorManager
-    else:
-        return _resource_name_to_accelerator_manager.get(resource_name, None)
+            resource_name_to_accelerator_manager["GPU"] = NvidiaGPUAcceleratorManager
+        get_accelerator_manager_for_resource._resource_name_to_accelerator_manager = (
+            resource_name_to_accelerator_manager
+        )
+        return resource_name_to_accelerator_manager.get(resource_name, None)
 
 
 __all__ = [

diff --git a/python/ray/tests/accelerators/test_intel_gpu.py b/python/ray/tests/accelerators/test_intel_gpu.py
@@ -5,14 +5,18 @@
 
 import ray
 from ray._private.accelerators import IntelGPUAcceleratorManager as Accelerator
+from ray._private.accelerators import get_accelerator_manager_for_resource
 from ray.util.accelerators import INTEL_MAX_1550, INTEL_MAX_1100
 
 
 def test_visible_intel_gpu_ids(shutdown_only):
     with patch.object(Accelerator, "get_current_node_num_accelerators", return_value=4):
         os.environ["ONEAPI_DEVICE_SELECTOR"] = "level_zero:0,1,2"
+        # Delete the cache so it can be re-populated the next time
+        # we call get_accelerator_manager_for_resource
+        del get_accelerator_manager_for_resource._resource_name_to_accelerator_manager
         ray.init()
-        manager = ray._private.accelerators.get_accelerator_manager_for_resource("GPU")
+        manager = get_accelerator_manager_for_resource("GPU")
         assert manager.get_current_node_num_accelerators() == 4
         assert manager.__name__ == "IntelGPUAcceleratorManager"
         assert ray.available_resources()["GPU"] == 3
@@ -25,8 +29,9 @@ def test_visible_intel_gpu_type(shutdown_only):
         Accelerator, "get_current_node_accelerator_type", return_value=INTEL_MAX_1550
     ):
         os.environ["ONEAPI_DEVICE_SELECTOR"] = "level_zero:0,1,2"
+        del get_accelerator_manager_for_resource._resource_name_to_accelerator_manager
         ray.init()
-        manager = ray._private.accelerators.get_accelerator_manager_for_resource("GPU")
+        manager = get_accelerator_manager_for_resource("GPU")
         assert manager.get_current_node_accelerator_type() == INTEL_MAX_1550