Including psutil & setproctitle (#7031)

.gitignore

@@ -3,6 +3,7 @@
 /python/ray/core
 /python/ray/pyarrow_files/
 /python/ray/pickle5_files/
+/python/ray/thirdparty_files/
 /python/ray/jars/
 /python/build
 /python/dist

.travis.yml

@@ -82,7 +82,7 @@ matrix:
         - sphinx-build -W -b html -d _build/doctrees source _build/html
         - cd ..
         # Run Python linting, ignore dict vs {} (C408), others are defaults
-        - flake8 --inline-quotes '"' --no-avoid-escape --exclude=python/ray/core/generated/,streaming/python/generated,doc/source/conf.py,python/ray/cloudpickle/ --ignore=C408,E121,E123,E126,E226,E24,E704,W503,W504,W605
+        - flake8 --inline-quotes '"' --no-avoid-escape --exclude=python/ray/core/generated/,streaming/python/generated,doc/source/conf.py,python/ray/cloudpickle/,python/ray/thirdparty_files --ignore=C408,E121,E123,E126,E226,E24,E704,W503,W504,W605
         - ./ci/travis/format.sh --all
         # Make sure that the README is formatted properly.
         - cd python

build.sh

@@ -115,6 +115,10 @@ if [[ "$PYTHON_VERSION" == "3.5" || "$PYTHON_VERSION" == "3.6" || "$PYTHON_VERSI
   popd
 fi
 
+
+"$PYTHON_EXECUTABLE" -m pip install -q psutil setproctitle \
+        --target="$ROOT_DIR/python/ray/thirdparty_files"
+
 export PYTHON3_BIN_PATH="$PYTHON_EXECUTABLE"
 export PYTHON2_BIN_PATH="$PYTHON_EXECUTABLE"
 

ci/travis/check_import_order.py

@@ -0,0 +1,114 @@
+import argparse
+import re
+import sys
+import tempfile
+from pathlib import Path
+
+exit_with_error = False
+
+
+def check_import(file):
+    check_to_lines = {
+        "import ray": -1,
+        "import psutil": -1,
+        "import setproctitle": -1
+    }
+
+    with open(file) as f:
+        for i, line in enumerate(f):
+            for check in check_to_lines.keys():
+                if re.match(r"\s+" + check + r"(\s|$)", line):
+                    check_to_lines[check] = i
+
+    for import_lib in ["import psutil", "import setproctitle"]:
+        if check_to_lines[import_lib] != -1:
+            import_psutil_line = check_to_lines[import_lib]
+            import_ray_line = check_to_lines["import ray"]
+            if import_ray_line == -1 or import_ray_line > import_psutil_line:
+                print(
+                    "{}:{}".format(str(file), import_psutil_line + 1),
+                    "{} without explicit import ray before it.".format(
+                        import_lib))
+                global exit_with_error
+                exit_with_error = True
+
+
+# Run the test with pytest file_name
+def test_check_import():
+    global exit_with_error
+    _, path = tempfile.mkstemp()
+
+    with open(path, "w") as f:
+        f.write("""
+        import psutil
+        import ray
+        """)
+    check_import(path)
+    assert exit_with_error
+    exit_with_error = False
+
+    with open(path, "w") as f:
+        f.write("""
+        import psutil
+        """)
+    check_import(path)
+    assert exit_with_error
+    exit_with_error = False
+
+    with open(path, "w") as f:
+        f.write("""
+        import random_lib
+        """)
+    check_import(path)
+    assert not exit_with_error
+    exit_with_error = False
+
+    with open(path, "w") as f:
+        f.write("""
+        import setproctitle
+        import ray
+        """)
+    check_import(path)
+    assert exit_with_error
+    exit_with_error = False
+
+    with open(path, "w") as f:
+        f.write("""
+        import ray
+        import psutil
+        """)
+    check_import(path)
+    assert not exit_with_error
+    exit_with_error = False
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("path", help="File path to check. e.g. '.' or './src'")
+    parser.add_argument(
+        "-s", "--skip", action="append", help="Skip certian directory")
+    args = parser.parse_args()
+
+    file_path = Path(args.path)
+    if file_path.is_dir():
+        all_py_files = file_path.rglob("*.py")
+    else:
+        all_py_files = [file_path]
+
+    if args.skip is not None:
+        filtered_py_files = []
+        for py_file in all_py_files:
+            should_skip = False
+            for skip_dir in args.skip:
+                if str(py_file).startswith(skip_dir):
+                    should_skip = True
+            if not should_skip:
+                filtered_py_files.append(py_file)
+        all_py_files = filtered_py_files
+
+    for py_file in all_py_files:
+        check_import(py_file)
+
+    if exit_with_error:
+        print("check import ordering failed")
+        sys.exit(1)

ci/travis/format.sh

@@ -81,6 +81,7 @@ YAPF_EXCLUDES=(
     '--exclude' 'python/build/*'
     '--exclude' 'python/ray/pyarrow_files/*'
     '--exclude' 'python/ray/core/src/ray/gcs/*'
+    '--exclude' 'python/ray/thirdparty_files/*'
 )
 
 # Format specified files
@@ -104,14 +105,14 @@ format_changed() {
              yapf --in-place "${YAPF_EXCLUDES[@]}" "${YAPF_FLAGS[@]}"
         if which flake8 >/dev/null; then
             git diff --name-only --diff-filter=ACRM "$MERGEBASE" -- '*.py' | xargs -P 5 \
-                 flake8 --inline-quotes '"' --no-avoid-escape --exclude=python/ray/core/generated/,streaming/python/generated,doc/source/conf.py,python/ray/cloudpickle/ --ignore=C408,E121,E123,E126,E226,E24,E704,W503,W504,W605
+                 flake8 --inline-quotes '"' --no-avoid-escape --exclude=python/ray/core/generated/,streaming/python/generated,doc/source/conf.py,python/ray/cloudpickle/,python/ray/thirdparty_files/ --ignore=C408,E121,E123,E126,E226,E24,E704,W503,W504,W605
         fi
     fi
 
     if ! git diff --diff-filter=ACRM --quiet --exit-code "$MERGEBASE" -- '*.pyx' '*.pxd' '*.pxi' &>/dev/null; then
         if which flake8 >/dev/null; then
             git diff --name-only --diff-filter=ACRM "$MERGEBASE" -- '*.pyx' '*.pxd' '*.pxi' | xargs -P 5 \
-                 flake8 --inline-quotes '"' --no-avoid-escape --exclude=python/ray/core/generated/,streaming/python/generated,doc/source/conf.py,python/ray/cloudpickle/ --ignore=C408,E121,E123,E126,E211,E225,E226,E227,E24,E704,E999,W503,W504,W605
+                 flake8 --inline-quotes '"' --no-avoid-escape --exclude=python/ray/core/generated/,streaming/python/generated,doc/source/conf.py,python/ray/cloudpickle/,python/ray/thirdparty_files/ --ignore=C408,E121,E123,E126,E211,E225,E226,E227,E24,E704,E999,W503,W504,W605
         fi
     fi
 
@@ -141,6 +142,11 @@ else
     format_changed
 fi
 
+# Ensure import ordering
+# Make sure that for every import psutil; import setpproctitle
+# There's a import ray above it.
+python ci/travis/check_import_order.py . -s ci -s python/ray/pyarrow_files
+
 if ! git diff --quiet &>/dev/null; then
     echo 'Reformatted changed files. Please review and stage the changes.'
     echo 'Files updated:'

ci/travis/install-dependencies.sh

@@ -61,3 +61,6 @@ if [[ "$PYTHON" == "3.5" ]] || [[ "$MAC_WHEELS" == "1" ]]; then
   source $HOME/.nvm/nvm.sh
   nvm install node
 fi
+
+pip install -q psutil setproctitle \
+        --target="$ROOT_DIR/../../python/ray/thirdparty_files"

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(