CELERY_BENCH: Show RSS memory usag at start and end of benchmark

celery/utils/compat.py

@@ -153,3 +153,18 @@ else:
                         stat = os.stat(self.baseFilename)
                     self.dev, self.ino = stat[ST_DEV], stat[ST_INO]
                 logging.FileHandler.emit(self, record)
+
+
+############## format(int, ',d') ##########################
+
+if sys.version_info >= (2, 7):
+    def format_d(i):
+        return format(i, ',d')
+else:
+    def format_d(i):
+        s = '%d' % i
+        groups = []
+        while s and s[-1].isdigit():
+            groups.append(s[-3:])
+            s = s[:-3]
+        return s + ','.join(reversed(groups))

celery/worker/state.py

@@ -72,15 +72,34 @@ if os.environ.get("CELERY_BENCH"):  # pragma: no cover
 
     from time import time
     from billiard import current_process
+    from celery.utils.compat import format_d
 
     all_count = 0
     bench_first = None
+    bench_mem_first = None
     bench_start = None
     bench_last = None
     bench_every = int(os.environ.get("CELERY_BENCH_EVERY", 1000))
     bench_sample = []
     __reserved = task_reserved
     __ready = task_ready
+    _process = None
+
+    def ps():
+        global _process
+        if _process is None:
+            try:
+                from psutil import Process
+            except ImportError:
+                return None
+            _process = Process(os.getpid())
+        return _process
+
+    def mem_rss():
+        p = ps()
+        if p is None:
+            return "(psutil not installed)"
+        return "%s MB" % (format_d(p.get_memory_info().rss // 1024), )
 
     if current_process()._name == 'MainProcess':
         @atexit.register
@@ -89,15 +108,19 @@ if os.environ.get("CELERY_BENCH"):  # pragma: no cover
                 print("\n- Time spent in benchmark: %r" % (
                     bench_last - bench_first))
                 print("- Avg: %s" % (sum(bench_sample) / len(bench_sample)))
+                print("- RSS: %s --> %s" % (bench_mem_first, mem_rss()))
 
     def task_reserved(request):  # noqa
         global bench_start
         global bench_first
+        global bench_mem_first
         now = None
         if bench_start is None:
             bench_start = now = time()
         if bench_first is None:
             bench_first = now
+        if bench_mem_first is None:
+            bench_mem_first = mem_rss()
 
         return __reserved(request)