Prepared flame graph mode

Author: Armand Foucault

mprof.py

@@ -456,6 +456,105 @@ def plot_file(filename, index=0, timestamps=True, children=True, options=None):
     return mprofile
 
 
+def flame_plotter(filename, index=0, timestamps=True, children=True, options=None):
+    try:
+        import pylab as pl
+    except ImportError as e:
+        print("matplotlib is needed for plotting.")
+        print(e)
+        sys.exit(1)
+    import numpy as np  # pylab requires numpy anyway
+    mprofile = read_mprofile_file(filename)
+
+    if len(mprofile['timestamp']) == 0:
+        print('** No memory usage values have been found in the profile '
+              'file.**\nFile path: {0}\n'
+              'File may be empty or invalid.\n'
+              'It can be deleted with "mprof rm {0}"'.format(
+            mprofile['filename']))
+        sys.exit(0)
+
+    # Merge function timestamps and memory usage together
+    ts = mprofile['func_timestamp']
+    t = mprofile['timestamp']
+    mem = mprofile['mem_usage']
+    chld = mprofile['children']
+
+    if len(ts) > 0:
+        for values in ts.values():
+            for v in values:
+                t.extend(v[:2])
+                mem.extend(v[2:4])
+
+    mem = np.asarray(mem)
+    t = np.asarray(t)
+    ind = t.argsort()
+    mem = mem[ind]
+    t = t[ind]
+
+    # Plot curves
+    global_start = float(t[0])
+    t = t - global_start
+
+    max_mem = mem.max()
+    max_mem_ind = mem.argmax()
+
+    all_colors = ("c", "y", "g", "r", "b")
+    mem_line_colors = ("k", "b", "r", "g", "c", "y", "m")
+    mem_line_label = time.strftime("%d / %m / %Y - start at %H:%M:%S",
+                                   time.localtime(global_start)) \
+                     + ".{0:03d}".format(int(round(math.modf(global_start)[0] * 1000)))
+
+    pl.plot(t, mem, "+-" + mem_line_colors[index % len(mem_line_colors)],
+            label=mem_line_label)
+
+    bottom, top = pl.ylim()
+    bottom += 0.001
+    top -= 0.001
+
+    # plot children, if any
+    if len(chld) > 0 and children:
+        cmpoint = (0,0) # maximal child memory
+
+        for idx, (proc, data) in enumerate(chld.items()):
+            # Create the numpy arrays from the series data
+            cts  = np.asarray([item[1] for item in data]) - global_start
+            cmem = np.asarray([item[0] for item in data])
+
+            # Plot the line to the figure
+            pl.plot(cts, cmem, "+-"  + mem_line_colors[(idx+1) % len(mem_line_colors)],
+                     label="child {}".format(proc))
+
+            # Detect the maximal child memory point
+            cmax_mem = cmem.max()
+            if cmax_mem > cmpoint[1]:
+                cmpoint = (cts[cmem.argmax()], cmax_mem)
+
+        # Add the marker lines for the maximal child memory usage
+        pl.vlines(cmpoint[0], pl.ylim()[0]+0.001, pl.ylim()[1] - 0.001, 'r', '--')
+        pl.hlines(cmpoint[1], pl.xlim()[0]+0.001, pl.xlim()[1] - 0.001, 'r', '--')
+
+    # plot timestamps, if any
+    if len(ts) > 0 and timestamps:
+        func_num = 0
+        f_labels = function_labels(ts.keys())
+        for f, exec_ts in ts.items():
+            for execution in exec_ts:
+                add_brackets(execution[:2], execution[2:], xshift=global_start,
+                             color=all_colors[func_num % len(all_colors)],
+                             label=f_labels[f]
+                                   + " %.3fs" % (execution[1] - execution[0]), options=options)
+            func_num += 1
+
+    if timestamps:
+        pl.hlines(max_mem,
+                  pl.xlim()[0] + 0.001, pl.xlim()[1] - 0.001,
+                  colors="r", linestyles="--")
+        pl.vlines(t[max_mem_ind], bottom, top,
+                  colors="r", linestyles="--")
+    return mprofile
+
+
 def function_labels(dotted_function_names):
     state = {}
 
@@ -507,6 +606,8 @@ def xlim_type(value):
                         help="Save plot to file instead of displaying it.")
     parser.add_argument("--window", "-w", dest="xlim", type=xlim_type,
                         help="Plot a time-subset of the data. E.g. to plot between 0 and 20.5 seconds: --window 0,20.5")
+    parser.add_argument("--flame", "-f", dest="flame_mode", action="store_true",
+                        help="Plot the timestamps as a flame-graph instead of the default brackets")
     parser.add_argument("--backend",
                       help="Specify the Matplotlib backend to use")
     parser.add_argument("profiles", nargs="*",
@@ -561,8 +662,11 @@ def xlim_type(value):
         timestamps = False
     else:
         timestamps = True
+    plotter = plot_file
+    if args.flame_mode:
+        plotter = flame_plotter
     for n, filename in enumerate(filenames):
-        mprofile = plot_file(filename, index=n, timestamps=timestamps, options=args)
+        mprofile = plotter(filename, index=n, timestamps=timestamps, options=args)
     pl.xlabel("time (in seconds)")
     pl.ylabel("memory used (in MiB)")