Including psutil & setproctitle (#7031)

python/ray/__init__.py

@@ -25,6 +25,11 @@ pickle5_path = os.path.join(
     os.path.abspath(os.path.dirname(__file__)), "pickle5_files")
 sys.path.insert(0, pickle5_path)
 
+# Importing psutil & setproctitle. Must be before ray._raylet is initialized.
+thirdparty_files = os.path.join(
+    os.path.abspath(os.path.dirname(__file__)), "thirdparty_files")
+sys.path.insert(0, thirdparty_files)
+
 # Expose ray ABI symbols which may be dependent by other shared
 # libraries such as _streaming.so. See BUILD.bazel:_raylet
 so_path = os.path.join(dirname(__file__), "_raylet.so")

python/ray/exceptions.py

@@ -2,12 +2,8 @@ import os
 
 import colorama
 
-try:
-    import setproctitle
-except ImportError:
-    setproctitle = None
-
 import ray
+import setproctitle
 
 
 class RayError(Exception):

python/ray/memory_monitor.py

@@ -3,10 +3,9 @@ import os
 import sys
 import time
 
-try:
-    import psutil
-except ImportError:
-    psutil = None
+# Import ray before psutil will make sure we use psutil's bundled version
+import ray  # noqa F401
+import psutil  # noqa E402
 
 logger = logging.getLogger(__name__)
 
@@ -105,9 +104,6 @@ class MemoryMonitor:
         self.worker_name = worker_name
 
     def raise_if_low_memory(self):
-        if psutil is None:
-            return  # nothing we can do
-
         if time.time() - self.last_checked > self.check_interval:
             if "RAY_DEBUG_DISABLE_MEMORY_MONITOR" in os.environ:
                 return  # escape hatch, not intended for user use

python/ray/reporter.py

@@ -9,13 +9,8 @@ import grpc
 import subprocess
 from concurrent import futures
 
-try:
-    import psutil
-except ImportError:
-    print("The reporter requires psutil to run.")
-    import sys
-    sys.exit(1)
-
+import ray
+import psutil
 import ray.ray_constants as ray_constants
 import ray.services
 import ray.utils

python/ray/services.py

@@ -17,6 +17,7 @@ import pyarrow
 # Ray modules
 import ray
 import ray.ray_constants as ray_constants
+import psutil
 
 # True if processes are run in the valgrind profiler.
 RUN_RAYLET_PROFILER = False
@@ -91,11 +92,6 @@ def include_java_from_redis(redis_client):
 
 
 def find_redis_address_or_die():
-    try:
-        import psutil
-    except ImportError:
-        raise ImportError(
-            "Please install `psutil` to automatically detect the Ray cluster.")
     pids = psutil.pids()
     redis_addresses = set()
     for pid in pids:
@@ -998,13 +994,6 @@ def start_reporter(redis_address,
     if redis_password:
         command += ["--redis-password", redis_password]
 
-    try:
-        import psutil  # noqa: F401
-    except ImportError:
-        logger.warning("Failed to start the reporter. The reporter requires "
-                       "'pip install psutil'.")
-        return None
-
     process_info = start_ray_process(
         command,
         ray_constants.PROCESS_TYPE_REPORTER,
@@ -1066,8 +1055,6 @@ def start_dashboard(require_webui,
     webui_dependencies_present = True
     try:
         import aiohttp  # noqa: F401
-        import psutil  # noqa: F401
-        import setproctitle  # noqa: F401
         import grpc  # noqa: F401
     except ImportError:
         webui_dependencies_present = False

python/ray/tests/test_advanced_3.py

@@ -2,7 +2,6 @@
 import glob
 import logging
 import os
-import setproctitle
 import shutil
 import json
 import sys
@@ -20,6 +19,7 @@ from ray import signature
 import ray.ray_constants as ray_constants
 import ray.cluster_utils
 import ray.test_utils
+import setproctitle
 
 from ray.test_utils import RayTestTimeoutException
 

python/ray/tune/progress_reporter.py

@@ -215,6 +215,7 @@ class CLIReporter(TuneReporterBase):
 
 def memory_debug_str():
     try:
+        import ray  # noqa F401
         import psutil
         total_gb = psutil.virtual_memory().total / (1024**3)
         used_gb = total_gb - psutil.virtual_memory().available / (1024**3)

python/ray/tune/utils/util.py

@@ -7,14 +7,10 @@ from threading import Thread
 
 import numpy as np
 import ray
+import psutil
 
 logger = logging.getLogger(__name__)
 
-try:
-    import psutil
-except ImportError:
-    psutil = None
-
 try:
     import GPUtil
 except ImportError:

python/ray/utils.py

@@ -12,8 +12,10 @@ import threading
 import time
 import uuid
 
+import ray
 import ray.gcs_utils
 import ray.ray_constants as ray_constants
+import psutil
 
 logger = logging.getLogger(__name__)
 
@@ -438,12 +440,7 @@ def get_system_memory():
             docker_limit = int(f.read())
 
     # Use psutil if it is available.
-    psutil_memory_in_bytes = None
-    try:
-        import psutil
-        psutil_memory_in_bytes = psutil.virtual_memory().total
-    except ImportError:
-        pass
+    psutil_memory_in_bytes = psutil.virtual_memory().total
 
     if psutil_memory_in_bytes is not None:
         memory_in_bytes = psutil_memory_in_bytes
@@ -468,22 +465,7 @@ def estimate_available_memory():
         The total amount of available memory in bytes. It may be an
         overestimate if psutil is not installed.
     """
-
-    # Use psutil if it is available.
-    try:
-        import psutil
-        return psutil.virtual_memory().available
-    except ImportError:
-        pass
-
-    # Handle Linux.
-    if sys.platform == "linux" or sys.platform == "linux2":
-        bytes_in_kilobyte = 1024
-        return (
-            vmstat("total memory") - vmstat("used memory")) * bytes_in_kilobyte
-
-    # Give up
-    return get_system_memory()
+    return psutil.virtual_memory().available
 
 
 def get_shared_memory_bytes():

python/ray/worker.py

@@ -27,6 +27,8 @@ import ray.ray_constants as ray_constants
 import ray.remote_function
 import ray.serialization as serialization
 import ray.services as services
+import ray
+import setproctitle
 import ray.signature
 import ray.state
 
@@ -63,11 +65,6 @@ ERROR_KEY_PREFIX = b"Error:"
 # entry/init points.
 logger = logging.getLogger(__name__)
 
-try:
-    import setproctitle
-except ImportError:
-    setproctitle = None
-
 # Whether we should warn about slow put performance.
 if os.environ.get("OMP_NUM_THREADS") == "1":
     should_warn_of_slow_puts = True

python/setup.py

@@ -73,8 +73,8 @@ if "RAY_USE_NEW_GCS" in os.environ and os.environ["RAY_USE_NEW_GCS"] == "on":
     ]
 
 extras = {
-    "debug": ["psutil", "setproctitle", "py-spy >= 0.2.0"],
-    "dashboard": ["aiohttp", "google", "grpcio", "psutil", "setproctitle"],
+    "debug": [],
+    "dashboard": [],
     "serve": ["uvicorn", "pygments", "werkzeug", "flask", "pandas", "blist"],
     "tune": ["tabulate", "tensorboardX"],
 }
@@ -104,19 +104,16 @@ class build_ext(_build_ext.build_ext):
 
         # We also need to install pyarrow along with Ray, so make sure that the
         # relevant non-Python pyarrow files get copied.
-        pyarrow_files = []
-        for (root, dirs, filenames) in os.walk("./ray/pyarrow_files/pyarrow"):
-            for name in filenames:
-                pyarrow_files.append(os.path.join(root, name))
+        pyarrow_files = self.walk_directory("./ray/pyarrow_files/pyarrow")
 
         # We also need to install pickle5 along with Ray, so make sure that the
         # relevant non-Python pickle5 files get copied.
-        pickle5_files = []
-        for (root, dirs, filenames) in os.walk("./ray/pickle5_files/pickle5"):
-            for name in filenames:
-                pickle5_files.append(os.path.join(root, name))
+        pickle5_files = self.walk_directory("./ray/pickle5_files/pickle5")
 
-        files_to_include = ray_files + pyarrow_files + pickle5_files
+        thirdparty_files = self.walk_directory("./ray/thirdparty_files")
+
+        files_to_include = ray_files + pyarrow_files + pickle5_files + \
+            thirdparty_files
 
         # Copy over the autogenerated protobuf Python bindings.
         for directory in generated_python_directories:
@@ -135,6 +132,13 @@ class build_ext(_build_ext.build_ext):
                 print("Failed to copy optional file {}. This is ok."
                       .format(filename))
 
+    def walk_directory(self, directory):
+        file_list = []
+        for (root, dirs, filenames) in os.walk(directory):
+            for name in filenames:
+                file_list.append(os.path.join(root, name))
+        return file_list
+
     def move_file(self, filename):
         # TODO(rkn): This feels very brittle. It may not handle all cases. See
         # https://github.com/apache/arrow/blob/master/python/setup.py for an
@@ -183,6 +187,10 @@ requires = [
     "faulthandler;python_version<'3.3'",
     "protobuf >= 3.8.0",
     "cloudpickle",
+    "py-spy >= 0.2.0",
+    "aiohttp",
+    "google",
+    "grpcio"
 ]
 
 setup(