From cd80891ddb372b38d83ea5e46c778ad6c9de0a09 Mon Sep 17 00:00:00 2001 From: Robert Nishihara Date: Thu, 3 Jan 2019 23:07:24 -0800 Subject: [PATCH] Try to figure out the memory limit in a docker container. (#3605) * Try to figure out the memory limit in a docker container. * Update comment * Fix * Fix --- python/ray/utils.py | 26 ++++++++++++++++++++++---- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/python/ray/utils.py b/python/ray/utils.py index 8c61607a0..02cce4c01 100644 --- a/python/ray/utils.py +++ b/python/ray/utils.py @@ -327,20 +327,38 @@ def get_system_memory(): Returns: The total amount of system memory in bytes. """ + # Try to accurately figure out the memory limit if we are in a docker + # container. Note that this file is not specific to Docker and its value is + # often much larger than the actual amount of memory. + docker_limit = None + memory_limit_filename = "/sys/fs/cgroup/memory/memory.limit_in_bytes" + if os.path.exists(memory_limit_filename): + with open(memory_limit_filename, "r") as f: + docker_limit = int(f.read()) + # Use psutil if it is available. + psutil_memory_in_bytes = None try: import psutil - return psutil.virtual_memory().total + psutil_memory_in_bytes = psutil.virtual_memory().total except ImportError: pass - if sys.platform == "linux" or sys.platform == "linux2": + memory_in_bytes = None + if psutil_memory_in_bytes is not None: + memory_in_bytes = psutil_memory_in_bytes + elif sys.platform == "linux" or sys.platform == "linux2": # Handle Linux. bytes_in_kilobyte = 1024 - return vmstat("total memory") * bytes_in_kilobyte + memory_in_bytes = vmstat("total memory") * bytes_in_kilobyte else: # Handle MacOS. - return sysctl(["sysctl", "hw.memsize"]) + memory_in_bytes = sysctl(["sysctl", "hw.memsize"]) + + if docker_limit is not None: + return min(docker_limit, memory_in_bytes) + else: + return memory_in_bytes def get_shared_memory_bytes():