plot function for mpmproc

Author: bbengfort

mpmprof

@@ -4,12 +4,22 @@ Multiprocessing version of memory profiling of Python programs.
 """
 
 import os
+import re
 import time
 import glob
 import argparse
 import subprocess
 import memory_profiler as mp
 
+from collections import defaultdict
+
+try:
+    import numpy as np
+    import matplotlib.pyplot as plt
+except ImportError:
+    plt = None
+    np  = None
+
 
 # Command Descriptions and Constants
 DESCRIPTION = "Multiprocessing memory profiling over time."
@@ -82,14 +92,130 @@ def run_action(args):
             # Sleep for the given interval
             time.sleep(args.interval)
 
+    # Return the results of the run action
     return "memory profile written to {}".format(args.output)
 
 
 def plot_action(args):
     """
     Use matplotlib to draw the memory usage of a mprofile .dat file.
     """
-    raise NotImplementedError("Not implemented yet.")
+    if plt is None:
+        raise ImportError(
+            "matplotlib is needed for plotting."
+        )
+
+    def read_mprofile_file(path):
+        """
+        Reads the specialized version of the mprofile for multiprocessing
+        """
+        # Regular expression line parsers for parsing data
+        cmdre = re.compile(r"^CMDLINE\s+(.+)$")
+        memre = re.compile(r"^MEM\s+([\d\.e]+)\s+([\d\.e]+)$")
+        cldre = re.compile(r"^CHLD(\d+)\s+([\d\.e]+)\s+([\d\.e]+)$")
+
+        # Data structure returned is a series of names (mem, ts) tuples.
+        series  = defaultdict(list)
+        command = None
+
+        with open(path, 'r') as f:
+            for line in f:
+
+                # Match children memory usage lines
+                match = cldre.match(line)
+                if match:
+                    idx, mem, ts = match.groups()
+                    series["child " + idx].append((float(mem), float(ts)))
+                    continue
+
+                # Match main process memory usage lines
+                match = memre.match(line)
+                if match:
+                    series['main'].append(tuple(map(float, match.groups())))
+                    continue
+
+                # Match command line(s)
+                # NOTE: mprofile files are openeded for appending, could be multiple
+                match = cmdre.match(line)
+                if match:
+                    command = match.groups()[0]
+
+        return command, series
+
+
+    def plot_mprofile_file(path, title=None):
+        """
+        Plots an mprofile file that contains specialized child process data.
+        """
+        # Parse the mprofile file to get the data
+        command, series = read_mprofile_file(path)
+        title = title or command
+
+        # Create and configure the figure
+        fig = plt.figure(figsize=(14, 6), dpi=90)
+        axe = fig.add_axes([0.1, 0.1, 0.6, 0.75])
+        axe.set_xlabel("time (in seconds)")
+        axe.set_ylabel("memory used (in MiB)")
+        axe.set_title(title)
+
+        # Find the start timestamp for the process and track the maximal memory point
+        # This currently assumes that the series were written in order
+        start  = series['main'][0][1]
+        mpoint = (0, 0)
+
+        # Plot all of the series, the main process and the child.
+        for proc, data in series.items():
+            # Create the numpy arrays from the series data
+            ts  = np.asarray([item[1] for item in data]) - start
+            mem = np.asarray([item[0] for item in data])
+
+            # Plot the line to the figure
+            plt.plot(ts, mem, "+-", label=proc)
+
+            # Detect the maximal memory point
+            max_mem = mem.max()
+            if max_mem > mpoint[1]:
+                mpoint = (mem.argmax(), max_mem)
+
+        # Add the marker lines for the maximal memory usage
+        plt.hlines(mpoint[1], plt.xlim()[0]+0.001, plt.xlim()[1] - 0.001, 'r', '--')
+        plt.vlines(ts[mpoint[0]], plt.ylim()[0]+0.001, plt.ylim()[1] - 0.001, 'r', '--')
+
+        # Add the legend
+        legend = axe.legend(loc='center left', bbox_to_anchor=(1, 0.5))
+        legend.get_frame().set_alpha(0.5)
+        axe.grid()
+
+    # Get the latest profile if no profile files were passed in.
+    if not args.profile:
+
+        # Glob profiles of our format and sort them.
+        profiles = glob.glob("mprofile_??????????????.dat")
+        profiles.sort()
+
+        if not profiles:
+            raise ValueError((
+                "No input file found.\nThis program looks for mprofile_*.dat "
+                "files generated by the `mpmprof run` command."
+            ))
+
+        # Assign the latest profile to visualize
+        args.profile = profiles[0:1]
+
+    # Filter out any files that do not exist
+    args.profile = list(filter(os.path.exists, args.profile))
+    if not args.profile:
+        raise ValueError("No input files found!")
+
+    # For each passed in file, create a figure from the mprofile.
+    for path in args.profile:
+        axe = plot_mprofile_file(path, args.title)
+        if args.output:
+            plt.savefig(args.output)
+        else:
+            plt.show()
+
+    return "{} memory profiles plotted.".format(len(args.profile))
 
 
 if __name__ == '__main__':