From 5723291db61f2677ded9d490461c165a08357efd Mon Sep 17 00:00:00 2001
From: Eric Liang <ekhliang@gmail.com>
Date: Thu, 15 Nov 2018 12:55:25 -0800
Subject: [PATCH] Raise exception if the node is nearly out of memory (#3323)

* wip

* add

* comment

* escape hatch

* update

* object store too

* .2
---
 python/ray/memory_monitor.py | 79 ++++++++++++++++++++++++++++++++++++
 python/ray/worker.py         |  3 ++
 2 files changed, 82 insertions(+)
 create mode 100644 python/ray/memory_monitor.py

diff --git a/python/ray/memory_monitor.py b/python/ray/memory_monitor.py
new file mode 100644
index 000000000..23d6c12b5
--- /dev/null
+++ b/python/ray/memory_monitor.py
@@ -0,0 +1,79 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import logging
+import os
+import time
+
+try:
+    import psutil
+except ImportError:
+    psutil = None
+
+logger = logging.getLogger(__name__)
+
+
+class RayOutOfMemoryError(Exception):
+    def __init__(self, msg):
+        Exception.__init__(self, msg)
+
+    @staticmethod
+    def get_message(used_gb, total_gb, threshold):
+        pids = psutil.pids()
+        proc_stats = []
+        for pid in pids:
+            proc = psutil.Process(pid)
+            proc_stats.append((proc.memory_info().rss, pid, proc.cmdline()))
+        proc_str = "PID\tMEM\tCOMMAND"
+        for rss, pid, cmdline in sorted(proc_stats, reverse=True)[:5]:
+            proc_str += "\n{}\t{}GB\t{}".format(
+                pid, round(rss / 1e9, 2), " ".join(cmdline)[:100].strip())
+        return ("More than {}% of the memory on ".format(int(
+            100 * threshold)) + "node {} is used ({} / {} GB). ".format(
+                os.uname()[1], round(used_gb, 2), round(total_gb, 2)) +
+                "The top 5 memory consumers are:\n\n{}".format(proc_str) +
+                "\n\nIn addition, ~{} GB of shared memory is ".format(
+                    round(psutil.virtual_memory().shared / 1e9, 2)) +
+                "currently being used by the Ray object store. You can set "
+                "the object store size with the `object_store_memory` "
+                "parameter when starting Ray.")
+
+
+class MemoryMonitor(object):
+    """Helper class for raising errors on low memory.
+
+    This presents a much cleaner error message to users than what would happen
+    if we actually ran out of memory.
+    """
+
+    def __init__(self, error_threshold=0.95, check_interval=1):
+        # Note: it takes ~50us to check the memory usage through psutil, so
+        # throttle this check at most once a second or so.
+        self.check_interval = check_interval
+        self.last_checked = time.time()
+        self.error_threshold = error_threshold
+        if not psutil:
+            logger.warning(
+                "WARNING: Not monitoring node memory since `psutil` is not "
+                "installed. Install this with `pip install psutil` to enable "
+                "debugging of memory-related crashes.")
+
+    def raise_if_low_memory(self):
+        if not psutil:
+            return  # nothing we can do
+
+        if "RAY_DEBUG_DISABLE_MEMORY_MONITOR" in os.environ:
+            return  # escape hatch, not intended for user use
+
+        if time.time() - self.last_checked > self.check_interval:
+            self.last_checked = time.time()
+            total_gb = psutil.virtual_memory().total / 1e9
+            used_gb = total_gb - psutil.virtual_memory().available / 1e9
+            if used_gb > total_gb * self.error_threshold:
+                raise RayOutOfMemoryError(
+                    RayOutOfMemoryError.get_message(used_gb, total_gb,
+                                                    self.error_threshold))
+            else:
+                logger.debug("Memory usage is {} / {}".format(
+                    used_gb, total_gb))
diff --git a/python/ray/worker.py b/python/ray/worker.py
index 5ed52f407..9c2358fce 100644
--- a/python/ray/worker.py
+++ b/python/ray/worker.py
@@ -25,6 +25,7 @@ import pyarrow.plasma as plasma
 import ray.cloudpickle as pickle
 import ray.experimental.state as state
 import ray.gcs_utils
+import ray.memory_monitor as memory_monitor
 import ray.remote_function
 import ray.serialization as serialization
 import ray.services as services
@@ -213,6 +214,7 @@ class Worker(object):
         # CUDA_VISIBLE_DEVICES environment variable.
         self.original_gpu_ids = ray.utils.get_cuda_visible_devices()
         self.profiler = profiling.Profiler(self)
+        self.memory_monitor = memory_monitor.MemoryMonitor()
         self.state_lock = threading.Lock()
         # A dictionary that maps from driver id to SerializationContext
         # TODO: clean up the SerializationContext once the job finished.
@@ -821,6 +823,7 @@ class Worker(object):
         try:
             if function_name != "__ray_terminate__":
                 self.reraise_actor_init_error()
+            self.memory_monitor.raise_if_low_memory()
             with profiling.profile("task:deserialize_arguments", worker=self):
                 arguments = self._get_arguments_for_execution(
                     function_name, args)