From 486a33980d8977489a92c44c01303dfe098f9e25 Mon Sep 17 00:00:00 2001
From: Clark Zinzow <clarkzinzow@gmail.com>
Date: Mon, 8 Aug 2022 19:10:49 +0000
Subject: [PATCH] Make expensive subpackage imports dynamic.

---
 python/ray/__init__.py                        | 46 +++++++++++++------
 .../tests/test_runtime_env_working_dir_4.py   |  2 +-
 python/ray/tests/test_top_level_api.py        |  6 ++-
 3 files changed, 38 insertions(+), 16 deletions(-)

diff --git a/python/ray/__init__.py b/python/ray/__init__.py
index 387d7b9db1d0..269c5fc80325 100644
--- a/python/ray/__init__.py
+++ b/python/ray/__init__.py
@@ -1,6 +1,7 @@
 # isort: skip_file
 import logging
 import os
+import sys
 
 logger = logging.getLogger(__name__)
 
@@ -115,7 +116,7 @@ def _configure_system():
 
 import ray._raylet  # noqa: E402
 
-from ray._raylet import (  # noqa: E402
+from ray._raylet import (  # noqa: E402,F401
     ActorClassID,
     ActorID,
     NodeID,
@@ -134,7 +135,7 @@ def _configure_system():
 
 _config = _Config()
 
-from ray._private.state import (  # noqa: E402
+from ray._private.state import (  # noqa: E402,F401
     nodes,
     timeline,
     cluster_resources,
@@ -162,23 +163,20 @@ def _configure_system():
 # We import ray.actor because some code is run in actor.py which initializes
 # some functions in the worker.
 import ray.actor  # noqa: E402,F401
-from ray.actor import method  # noqa: E402
+from ray.actor import method  # noqa: E402,F401
 
 # TODO(qwang): We should remove this exporting in Ray2.0.
-from ray.cross_language import java_function, java_actor_class  # noqa: E402
-from ray.runtime_context import get_runtime_context  # noqa: E402
-from ray import autoscaler  # noqa:E402
-from ray import data  # noqa: E402,F401
+from ray.cross_language import java_function, java_actor_class  # noqa: E402,F401
+from ray.runtime_context import get_runtime_context  # noqa: E402,F401
 from ray import internal  # noqa: E402,F401
-from ray import util  # noqa: E402
+from ray import util  # noqa: E402,F401
 from ray import _private  # noqa: E402,F401
-from ray import workflow  # noqa: E402,F401
 
 # We import ClientBuilder so that modules can inherit from `ray.ClientBuilder`.
-from ray.client_builder import client, ClientBuilder  # noqa: E402
+from ray.client_builder import client, ClientBuilder  # noqa: E402,F401
 
 
-class _DeprecationWrapper(object):
+class _DeprecationWrapper:
     def __init__(self, name, real_worker):
         self._name = name
         self._real_worker = real_worker
@@ -201,18 +199,23 @@ def __getattr__(self, attr):
 serialization = _DeprecationWrapper("serialization", ray._private.serialization)
 state = _DeprecationWrapper("state", ray._private.state)
 
+
+_subpackages = [
+    "autoscaler",
+    "data",
+    "workflow",
+]
+
 __all__ = [
     "__version__",
     "_config",
     "get_runtime_context",
     "actor",
     "available_resources",
-    "autoscaler",
     "cancel",
     "client",
     "ClientBuilder",
     "cluster_resources",
-    "data",
     "get",
     "get_actor",
     "get_gpu_ids",
@@ -237,7 +240,7 @@ def __getattr__(self, attr):
     "LOCAL_MODE",
     "SCRIPT_MODE",
     "WORKER_MODE",
-]
+] + _subpackages
 
 # ID types
 __all__ += [
@@ -255,6 +258,21 @@ def __getattr__(self, attr):
     "PlacementGroupID",
 ]
 
+if sys.version_info < (3, 7):
+    # TODO(Clark): Remove this one we drop Python 3.6 support.
+    from ray import autoscaler  # noqa: F401
+    from ray import data  # noqa: F401
+    from ray import workflow  # noqa: F401
+else:
+    # Delay importing of expensive, isolated subpackages.
+    def __getattr__(name: str):
+        import importlib
+
+        if name in _subpackages:
+            return importlib.import_module("." + name, __name__)
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
 
 del os
 del logging
+del sys
diff --git a/python/ray/tests/test_runtime_env_working_dir_4.py b/python/ray/tests/test_runtime_env_working_dir_4.py
index f9a7fd742e75..590fdec539a6 100644
--- a/python/ray/tests/test_runtime_env_working_dir_4.py
+++ b/python/ray/tests/test_runtime_env_working_dir_4.py
@@ -150,7 +150,7 @@ def check(self):
         runtime_env = {"py_modules": [S3_PACKAGE_URI]}
 
     # Note: We should set a bigger timeout if downloads the s3 package slowly.
-    get_timeout = 10
+    get_timeout = 2
 
     # Start a task with runtime env
     if worker_register_timeout:
diff --git a/python/ray/tests/test_top_level_api.py b/python/ray/tests/test_top_level_api.py
index 5eb02a55ee8f..3a506e135485 100644
--- a/python/ray/tests/test_top_level_api.py
+++ b/python/ray/tests/test_top_level_api.py
@@ -36,9 +36,13 @@ def test_api_functions():
         "get_runtime_context",
     ]
 
+    IMPL_FUNCTIONS = ["__getattr__"]
+
     functions = getmembers(ray, isfunction)
     function_names = [f[0] for f in functions]
-    assert set(function_names) == set(PYTHON_API + OTHER_ALLOWED_FUNCTIONS)
+    assert set(function_names) == set(
+        PYTHON_API + OTHER_ALLOWED_FUNCTIONS + IMPL_FUNCTIONS
+    )
 
 
 def test_non_ray_modules():