import custom class safely in controller by using runtime env

python/ray/_common/utils.py

@@ -9,27 +9,26 @@
 import sys
 import tempfile
 from inspect import signature
-from typing import Any, Coroutine, Dict, Optional
+from types import ModuleType
+from typing import Any, Coroutine, Dict, Optional, Tuple
 
 import psutil
 
 
-def import_attr(full_path: str, *, reload_module: bool = False):
-    """Given a full import path to a module attr, return the imported attr.
+def import_module_and_attr(
+    full_path: str, *, reload_module: bool = False
+) -> Tuple[ModuleType, Any]:
+    """Given a full import path to a module attr, return the imported module and attr.
 
     If `reload_module` is set, the module will be reloaded using `importlib.reload`.
 
-    For example, the following are equivalent:
-        MyClass = import_attr("module.submodule:MyClass")
-        MyClass = import_attr("module.submodule.MyClass")
-        from module.submodule import MyClass
+    Args:
+        full_path: The full import path to the module and attr.
+        reload_module: Whether to reload the module.
 
     Returns:
-        Imported attr
+        A tuple of the imported module and attr.
     """
-    if full_path is None:
-        raise TypeError("import path cannot be None")
-
     if ":" in full_path:
         if full_path.count(":") > 1:
             raise ValueError(
@@ -41,11 +40,26 @@ def import_attr(full_path: str, *, reload_module: bool = False):
         last_period_idx = full_path.rfind(".")
         module_name = full_path[:last_period_idx]
         attr_name = full_path[last_period_idx + 1 :]
-
     module = importlib.import_module(module_name)
     if reload_module:
         importlib.reload(module)
-    return getattr(module, attr_name)
+    return module, getattr(module, attr_name)
+
+
+def import_attr(full_path: str, *, reload_module: bool = False) -> Any:
+    """Given a full import path to a module attr, return the imported attr.
+
+    If `reload_module` is set, the module will be reloaded using `importlib.reload`.
+
+    For example, the following are equivalent:
+        MyClass = import_attr("module.submodule:MyClass")
+        MyClass = import_attr("module.submodule.MyClass")
+        from module.submodule import MyClass
+
+    Returns:
+        Imported attr
+    """
+    return import_module_and_attr(full_path, reload_module=reload_module)[1]
 
 
 def get_or_create_event_loop() -> asyncio.AbstractEventLoop: