Skip to content
This repository
Browse code

New strategy to get memory consumption

  • Loading branch information...
commit 22f21ab0e1f951ab05f5ca575ebc696dcf40a444 1 parent a3fafb8
Fabian Pedregosa authored

Showing 1 changed file with 35 additions and 31 deletions. Show diff stats Hide diff stats

  1. 66  memory_profiler.py
66  memory_profiler.py
@@ -45,6 +45,30 @@ def _get_memory(pid):
45 45
         raise NotImplementedError('The psutil module is required for non-unix '
46 46
                                   'platforms')
47 47
 
  48
+try:
  49
+    from multiprocessing import Process, Event, Queue
  50
+except ImportError:
  51
+    from threading import Thread, Event, Queue
  52
+    Process = Thread
  53
+
  54
+class Timer(Process):
  55
+    """
  56
+    Fetch memory consumption from over a time interval
  57
+    """
  58
+
  59
+    def __init__(self, monitor_pid, interval, event, queue, *args, **kw):
  60
+        self.monitor_pid = monitor_pid
  61
+        self.q = queue
  62
+        self.interval = interval
  63
+        self.cont = True
  64
+        self.e = event
  65
+        super(Timer, self).__init__(*args, **kw)
  66
+
  67
+    def run(self):
  68
+        while not self.e.is_set():
  69
+            m = _get_memory(self.monitor_pid)
  70
+            self.q.put(m)
  71
+            time.sleep(self.interval)
48 72
 
49 73
 def memory_usage(proc=-1, interval=.1, timeout=None, run_in_place=False):
50 74
     """
@@ -87,13 +111,6 @@ def memory_usage(proc=-1, interval=.1, timeout=None, run_in_place=False):
87 111
         # for a Python function wait until it finishes
88 112
         max_iter = float('inf')
89 113
 
90  
-    if str(proc).endswith('.py'):
91  
-        filename = _find_script(proc)
92  
-        with open(filename) as f:
93  
-            proc = f.read()
94  
-        # XXX TODO
95  
-        raise NotImplementedError
96  
-
97 114
     if isinstance(proc, (list, tuple)):
98 115
 
99 116
         if len(proc) == 1:
@@ -114,30 +131,17 @@ def memory_usage(proc=-1, interval=.1, timeout=None, run_in_place=False):
114 131
             'Function expects %s value(s) but %s where given'
115 132
             % (n_args, len(args)))
116 133
 
117  
-        try:
118  
-            import multiprocessing
119  
-        except ImportError:
120  
-            print ('WARNING: cannot import module `multiprocessing`. Forcing to'
121  
-                   ' run inplace.')
122  
-            # force inplace
123  
-            run_in_place = True
124  
-        if run_in_place:
125  
-            import threading
126  
-            main_thread = threading.Thread(target=f, args=args, kwargs=kw)
127  
-        else:
128  
-            main_thread = multiprocessing.Process(target=f, args=args, kwargs=kw)
129  
-        i = 0
130  
-        main_thread.start()
131  
-        pid = getattr(main_thread, 'pid', os.getpid())
132  
-        while i < max_iter and main_thread.is_alive():
133  
-            m = _get_memory(pid)
134  
-            if m == -1:
135  
-                # this means the thread has finished
136  
-                break
137  
-            ret.append(m)
138  
-            time.sleep(interval)
139  
-            i += 1
140  
-        main_thread.join(1.)
  134
+        e = Event() # this will tell Timer when to stop measuring
  135
+        q = Queue() # this will store Timer's results
  136
+        p = Timer(os.getpid(), interval, e, q)
  137
+        p.start()
  138
+        f(*args, **kw)
  139
+        e.set()
  140
+        p.join(5 * interval)
  141
+        ret = []
  142
+        while not q.empty():
  143
+            ret.append(q.get_nowait())
  144
+
141 145
     else:
142 146
         # external process
143 147
         if proc == -1:

0 notes on commit 22f21ab

Please sign in to comment.
Something went wrong with that request. Please try again.