Merge pull request #285 from cr-dani-koretsky/profile-existing-process

mprof attach option - profile existing process

Author: fabianp

mprof.py

@@ -8,18 +8,20 @@
 import math
 import logging
 import itertools
+from ast import literal_eval
 
 from collections import defaultdict
 from argparse import ArgumentParser, ArgumentError, REMAINDER, RawTextHelpFormatter
 
 import importlib
 import memory_profiler as mp
 
-ALL_ACTIONS = ("run", "rm", "clean", "list", "plot")
+ALL_ACTIONS = ("run", "rm", "clean", "list", "plot", "attach")
 help_msg = """
 Available commands:
 
     run      run a given command or python file
+    attach   alias for 'run --attach': attach to an existing process by pid or name
     rm       remove a given file generated by mprof
     clean    clean the current directory from files created by mprof
     list     display existing profiles, with indices
@@ -176,6 +178,16 @@ def get_cmd_line(args):
     args = [s if blanks.isdisjoint(s) else "'" + s + "'" for s in args]
     return ' '.join(args)
 
+def find_first_process(name):
+    for i in mp.psutil.process_iter():
+        if name in i.name():
+            return i
+    return None
+
+def attach_action():
+    argv = sys.argv
+    sys.argv = argv[:1] + ['--attach'] + argv[1:]
+    run_action()
 
 def run_action():
     import time, subprocess
@@ -192,6 +204,10 @@ def run_action():
     parser.add_argument("--multiprocess", "-M", dest="multiprocess", action="store_true",
                         help="""Monitors forked processes creating individual plots for each child (disables --python features)""")
     parser.add_argument("--exit-code", "-E", dest="exit_code", action="store_true", help="""Propagate the exit code""")
+    attach_arg = parser.add_argument("--attach", "-a", dest="attach_existing", action="store_true",
+                        help="Attach to an existing process, by process name or by pid")
+    parser.add_argument("--timeout", "-t", dest="timeout", action="store", type=int,
+                        help="timeout in seconds for the profiling, default new process has no timeout, attach existing is 1 hour")
     parser.add_argument("--output", "-o", dest="filename",
                         default="mprofile_%s.dat" % time.strftime("%Y%m%d%H%M%S", time.localtime()),
                         help="""File to store results in, defaults to 'mprofile_<YYYYMMDDhhmmss>.dat' in the current directory,
@@ -214,34 +230,53 @@ def run_action():
 
     mprofile_output = args.filename
 
-    # .. TODO: more than one script as argument ? ..
     program = args.program
-    if program[0].endswith('.py') and not args.nopython:
-        if args.multiprocess:
-            # in multiprocessing mode you want to spawn a separate
-            # python process
+    if args.attach_existing:
+        print('attaching to existing process, using hint: {}'.format(program[0]))
+        if program[0].isdigit():
+            p = literal_eval(program[0])
+            cmd_line = get_cmd_line(program)
+        else:
+            proc = find_first_process(program[0])
+            if proc is None:
+                raise ArgumentError(attach_arg, '\nWhen attaching, program should be process name or pid.\nFailed to find a process using hint: {}'.format(program[0]))
+            
+            p = proc.pid
+            try:
+                cmd_line = proc.cmdline()
+            except:
+                cmd_line = get_cmd_line(program)
+        if args.timeout is None:
+            args.timeout = 3600
+    else:
+        print('running new process')
+        # .. TODO: more than one script as argument ? ..
+        if program[0].endswith('.py') and not args.nopython:
+            if args.multiprocess:
+                # in multiprocessing mode you want to spawn a separate
+                # python process
+                if not program[0].startswith("python"):
+                    program.insert(0, sys.executable)
+                args.python = False
+            else:
+                args.python = True
+        if args.python:
+            print("running as a Python program...")
             if not program[0].startswith("python"):
                 program.insert(0, sys.executable)
-            args.python = False
+            cmd_line = get_cmd_line(program)
+            extra_args = ["-m", "memory_profiler", "--timestamp", "-o", mprofile_output]
+            if args.include_children:
+                extra_args.append("--include-children")
+            program[1:1] = extra_args
+            p = subprocess.Popen(program)
         else:
-            args.python = True
-    if args.python:
-        print("running as a Python program...")
-        if not program[0].startswith("python"):
-            program.insert(0, sys.executable)
-        cmd_line = get_cmd_line(program)
-        extra_args = ["-m", "memory_profiler", "--timestamp", "-o", mprofile_output]
-        if args.include_children:
-            extra_args.append("--include-children")
-        program[1:1] = extra_args
-        p = subprocess.Popen(program)
-    else:
-        cmd_line = get_cmd_line(program)
-        p = subprocess.Popen(program)
+            cmd_line = get_cmd_line(program)
+            p = subprocess.Popen(program)
 
     with open(mprofile_output, "a") as f:
         f.write("CMDLINE {0}\n".format(cmd_line))
-        mp.memory_usage(proc=p, interval=args.interval, timestamps=True,
+        mp.memory_usage(proc=p, interval=args.interval, timeout=args.timeout, timestamps=True,
                         include_children=args.include_children,
                         multiprocess=args.multiprocess, stream=f)
 
@@ -837,6 +872,7 @@ def main():
                "clean": clean_action,
                "list": list_action,
                "run": run_action,
+               "attach": attach_action,
                "plot": plot_action}
     actions[get_action()]()