Repository: mars-project/mars

Issue:
[ENH] More accurate memory stats using memory.stat in cgroup
Given kernel doc https://www.kernel.org/doc/Documentation/cgroup-v1/memory.txt, we may use ``/sys/fs/cgroup/memory/memory.stat`` to provide more accurate memory size estimations in Docker. Note that parsing this file is not supported in ``psutil`` now, thus we may parse it ourselves.


Discussion:

Resolving patch:
diff --git a/mars/resource.py b/mars/resource.py
--- a/mars/resource.py
+++ b/mars/resource.py
@@ -23,22 +23,25 @@
 
 from .lib import nvutils
 
+CGROUP_MEM_STAT_FILE = '/sys/fs/cgroup/memory/memory.stat'
+
 _proc = psutil.Process()
 _timer = getattr(time, 'monotonic', time.time)
 
 _cpu_use_process_stat = bool(int(os.environ.get('MARS_CPU_USE_PROCESS_STAT', '0').strip('"')))
 _mem_use_process_stat = bool(int(os.environ.get('MARS_MEM_USE_PROCESS_STAT', '0').strip('"')))
+_mem_use_cgroup_stat = bool(int(os.environ.get('MARS_MEM_USE_CGROUP_STAT', '0').strip('"')))
 
 if 'MARS_USE_PROCESS_STAT' in os.environ:
     _cpu_use_process_stat = _mem_use_process_stat = \
         bool(int(os.environ['MARS_USE_PROCESS_STAT'].strip('"')))
 
-if 'MARS_CPU_TOTAL' in os.environ:
+if _cpu_use_process_stat and 'MARS_CPU_TOTAL' in os.environ:
     _cpu_total = int(os.environ['MARS_CPU_TOTAL'].strip('"'))
 else:
     _cpu_total = psutil.cpu_count(logical=True)
 
-if 'MARS_MEMORY_TOTAL' in os.environ:
+if _mem_use_process_stat and 'MARS_MEMORY_TOTAL' in os.environ:
     _mem_total = int(os.environ['MARS_MEMORY_TOTAL'].strip('"'))
 else:
     _mem_total = None
@@ -53,9 +56,28 @@
     _shm_path = _shm_path[0]
 
 
+def _read_cgroup_stat_file():
+    with open(CGROUP_MEM_STAT_FILE, 'r') as cg_file:
+        contents = cg_file.read()
+    kvs = dict()
+    for l in contents.splitlines():
+        parts = l.split(' ')
+        if len(parts) == 2:
+            kvs[parts[0]] = int(parts[1])
+    return kvs
+
+
 def virtual_memory():
     sys_mem = psutil.virtual_memory()
-    if not _mem_use_process_stat:
+    if _mem_use_cgroup_stat:
+        # see section 5.5 in https://www.kernel.org/doc/Documentation/cgroup-v1/memory.txt
+        cgroup_mem_info = _read_cgroup_stat_file()
+        total = cgroup_mem_info['hierarchical_memory_limit']
+        used = cgroup_mem_info['cache'] + cgroup_mem_info['rss'] + cgroup_mem_info.get('swap', 0)
+        available = free = total - used
+        percent = 100.0 * (total - available) / total
+        return _virt_memory_stat(total, available, percent, used, free)
+    elif not _mem_use_process_stat:
         total = sys_mem.total
         used = sys_mem.used + getattr(sys_mem, 'shared', 0)
         available = sys_mem.available


Test patch:
diff --git a/mars/tests/test_resource.py b/mars/tests/test_resource.py
--- a/mars/tests/test_resource.py
+++ b/mars/tests/test_resource.py
@@ -13,11 +13,25 @@
 # limitations under the License.
 
 import os
+import tempfile
 import time
 import unittest
 
 from mars.compat import reload_module
 
+# just a fragment of real memory.stat
+_memory_stat_content = """
+cache 489275392
+rss 218181632
+mapped_file 486768640
+swap 0
+inactive_anon 486744064
+active_anon 218103808
+inactive_file 2457600
+active_file 73728
+hierarchical_memory_limit 1073741824
+"""
+
 
 class Test(unittest.TestCase):
     def testStats(self):
@@ -75,3 +89,24 @@ def testUseProcessStats(self):
             del os.environ['MARS_CPU_TOTAL']
             del os.environ['MARS_MEMORY_TOTAL']
             reload_module(resource)
+
+    def testUseCGroupStats(self):
+        from mars import resource
+        fd, mem_stat_path = tempfile.mkstemp(prefix='test-mars-res-')
+        with os.fdopen(fd, 'w') as f:
+            f.write(_memory_stat_content)
+
+        old_stat_file = resource.CGROUP_MEM_STAT_FILE
+        try:
+            os.environ['MARS_MEM_USE_CGROUP_STAT'] = '1'
+            resource = reload_module(resource)
+            resource.CGROUP_MEM_STAT_FILE = mem_stat_path
+
+            mem_stats = resource.virtual_memory()
+            self.assertEqual(mem_stats.total, 1073741824)
+            self.assertEqual(mem_stats.used, 707457024)
+        finally:
+            resource.CGROUP_MEM_STAT_FILE = old_stat_file
+            del os.environ['MARS_MEM_USE_CGROUP_STAT']
+            os.unlink(mem_stat_path)
+            reload_module(resource)