Use psutil instead of our home-grown utilities

mdboom · mdboom · commit b59627b9a1a6 · 2015-11-02T11:14:53.000-05:00
diff --git a/lib/matplotlib/cbook.py b/lib/matplotlib/cbook.py
@@ -1438,42 +1438,6 @@ def restrict_dict(d, keys):
     return dict([(k, v) for (k, v) in six.iteritems(d) if k in keys])
 
 
-def report_memory(i=0):  # argument may go away
-    'return the memory, in bytes, consumed by process'
-    from matplotlib.compat.subprocess import Popen, PIPE
-    pid = os.getpid()
-    if sys.platform.startswith('linux'):
-        try:
-            a2 = Popen('ps -p %d -o rss,sz' % pid, shell=True,
-                       stdout=PIPE).stdout.readlines()
-        except OSError:
-            raise NotImplementedError(
-                "report_memory works on Linux only if "
-                "the 'ps' program is found")
-        mem = int(a2[1].split()[1]) * 1024
-    elif sys.platform.startswith('darwin'):
-        try:
-            a2 = Popen('ps -p %d -o rss,vsz' % pid, shell=True,
-                       stdout=PIPE).stdout.readlines()
-        except OSError:
-            raise NotImplementedError(
-                "report_memory works on Mac OS only if "
-                "the 'ps' program is found")
-        mem = int(a2[1].split()[0]) * 1024
-    elif sys.platform.startswith('win'):
-        try:
-            a2 = Popen(["tasklist", "/nh", "/fi", "pid eq %d" % pid],
-                       stdout=PIPE).stdout.read()
-        except OSError:
-            raise NotImplementedError(
-                "report_memory works on Windows only if "
-                "the 'tasklist' program is found")
-        mem = int(a2.strip().split()[-2].replace(',', '')) * 1024
-    else:
-        raise NotImplementedError(
-            "We don't have a memory monitor for %s" % sys.platform)
-    return mem
-
 _safezip_msg = 'In safezip, len(args[0])=%d but len(args[%d])=%d'
 
 
@@ -1505,58 +1469,6 @@ def safe_masked_invalid(x):
     return xm
 
 
-class MemoryMonitor(object):
-    def __init__(self, nmax=20000):
-        self._nmax = nmax
-        self._mem = np.zeros((self._nmax,), np.int32)
-        self.clear()
-
-    def clear(self):
-        self._n = 0
-        self._overflow = False
-
-    def __call__(self):
-        mem = report_memory()
-        if self._n < self._nmax:
-            self._mem[self._n] = mem
-            self._n += 1
-        else:
-            self._overflow = True
-        return mem
-
-    def report(self, segments=4):
-        n = self._n
-        segments = min(n, segments)
-        dn = int(n / segments)
-        ii = list(xrange(0, n, dn))
-        ii[-1] = n - 1
-        print()
-        print('memory report: i, mem, dmem, dmem/nloops')
-        print(0, self._mem[0])
-        for i in range(1, len(ii)):
-            di = ii[i] - ii[i - 1]
-            if di == 0:
-                continue
-            dm = self._mem[ii[i]] - self._mem[ii[i - 1]]
-            print('%5d %5d %3d %8.3f' % (ii[i], self._mem[ii[i]],
-                                         dm, dm / float(di)))
-        if self._overflow:
-            print("Warning: array size was too small for the number of calls.")
-
-    def xy(self, i0=0, isub=1):
-        x = np.arange(i0, self._n, isub)
-        return x, self._mem[i0:self._n:isub]
-
-    def plot(self, i0=0, isub=1, fig=None):
-        if fig is None:
-            from .pylab import figure
-            fig = figure()
-
-        ax = fig.add_subplot(111)
-        ax.plot(*self.xy(i0, isub))
-        fig.canvas.draw()
-
-
 def print_cycles(objects, outstream=sys.stdout, show_progress=False):
     """
     *objects*
diff --git a/unit/memleak.py b/unit/memleak.py
@@ -3,39 +3,48 @@
 from __future__ import print_function
 
 import gc
+
 try:
     import tracemalloc
 except ImportError:
     raise ImportError("This script requires Python 3.4 or later")
 
+try:
+    import psutil
+except ImportError:
+    raise ImportError("This script requires psutil")
+
 import numpy as np
 
 
 def run_memleak_test(bench, iterations, report):
-    from matplotlib.cbook import report_memory
-
     tracemalloc.start()
 
-    starti = min(10, iterations / 2)
+    starti = min(50, iterations / 2)
     endi = iterations
 
     malloc_arr = np.empty((endi,), dtype=np.int64)
     rss_arr = np.empty((endi,), dtype=np.int64)
     rss_peaks = np.empty((endi,), dtype=np.int64)
     nobjs_arr = np.empty((endi,), dtype=np.int64)
     garbage_arr = np.empty((endi,), dtype=np.int64)
+    open_files_arr = np.empty((endi,), dtype=np.int64)
     rss_peak = 0
 
+    p = psutil.Process()
+
     for i in range(endi):
         bench()
 
         gc.collect()
-        rss = report_memory()
+
+        rss = p.memory_info().rss
         malloc, peak = tracemalloc.get_traced_memory()
         nobjs = len(gc.get_objects())
         garbage = len(gc.garbage)
-        print("{0: 4d}: pymalloc {1: 10d}, rss {2: 10d}, nobjs {3: 10d}, garbage {4: 10d}".format(
-            i, malloc, rss, nobjs, garbage))
+        open_files = len(p.open_files())
+        print("{0: 4d}: pymalloc {1: 10d}, rss {2: 10d}, nobjs {3: 10d}, garbage {4: 4d}, files: {5: 4d}".format(
+            i, malloc, rss, nobjs, garbage, open_files))
 
         malloc_arr[i] = malloc
         rss_arr[i] = rss
@@ -44,23 +53,27 @@ def run_memleak_test(bench, iterations, report):
         rss_peaks[i] = rss_peak
         nobjs_arr[i] = nobjs
         garbage_arr[i] = garbage
+        open_files_arr[i] = open_files
 
     print('Average memory consumed per loop: %1.4f bytes\n' %
           (np.sum(rss_peaks[starti+1:] - rss_peaks[starti:-1]) / float(endi - starti)))
 
     from matplotlib import pyplot as plt
-    fig, (ax1, ax2) = plt.subplots(2)
-    ax3 = ax1.twinx()
-    ax1.plot(malloc_arr[5:], 'r')
-    ax3.plot(rss_arr[5:], 'b')
+    fig, (ax1, ax2, ax3) = plt.subplots(3)
+    ax1b = ax1.twinx()
+    ax1.plot(malloc_arr, 'r')
+    ax1b.plot(rss_arr, 'b')
     ax1.set_ylabel('pymalloc', color='r')
-    ax3.set_ylabel('rss', color='b')
+    ax1b.set_ylabel('rss', color='b')
 
-    ax4 = ax2.twinx()
-    ax2.plot(nobjs_arr[5:], 'r')
-    ax4.plot(garbage_arr[5:], 'b')
+    ax2b = ax2.twinx()
+    ax2.plot(nobjs_arr, 'r')
+    ax2b.plot(garbage_arr, 'b')
     ax2.set_ylabel('total objects', color='r')
-    ax4.set_ylabel('garbage objects', color='b')
+    ax2b.set_ylabel('garbage objects', color='b')
+
+    ax3.plot(open_files_arr)
+    ax3.set_ylabel('open file handles')
 
     if not report.endswith('.pdf'):
         report = report + '.pdf'
