Add a profiling context manager that can be easily enabled

Change-Id: Icd27abe032b7e4daf78dc9a9e80f5daeee8f078e

1 files changed, 78 insertions, 9 deletions

diff --git a/tools/speed_test.py b/tools/speed_test.py
index 9038594..45bca78 100644
--- a/tools/speed_test.py
+++ b/tools/speed_test.py
@@ -12,6 +12,15 @@
 #    License for the specific language governing permissions and limitations
 #    under the License.
 
+"""
+Profile a simple engine build/load/compile/prepare/validate/run.
+"""
+
+import argparse
+import cProfile as profiler
+import pstats
+
+import six
 from six.moves import range as compat_range
 
 from taskflow import engines
@@ -20,10 +29,48 @@ from taskflow import task
 from taskflow.types import timing
 
 
+def print_header(name):
+    if name:
+        header_footer = "-" * len(name)
+        print(header_footer)
+        print(name)
+        print(header_footer)
+
+
+class ProfileIt(object):
+    stats_ordering = ('cumulative', 'calls',)
+
+    def __init__(self, name, args):
+        self.name = name
+        self.profile = profiler.Profile()
+        self.args = args
+
+    def __enter__(self):
+        self.profile.enable()
+
+    def __exit__(self, exc_tp, exc_v, exc_tb):
+        self.profile.disable()
+        buf = six.StringIO()
+        ps = pstats.Stats(self.profile, stream=buf)
+        ps = ps.sort_stats(*self.stats_ordering)
+        percent_limit = max(0.0, max(1.0, self.args.limit / 100.0))
+        ps.print_stats(percent_limit)
+        print_header(self.name)
+        needs_newline = False
+        for line in buf.getvalue().splitlines():
+            line = line.lstrip()
+            if line:
+                print(line)
+                needs_newline = True
+        if needs_newline:
+            print("")
+
+
 class TimeIt(object):
-    def __init__(self, name):
+    def __init__(self, name, args):
         self.watch = timing.StopWatch()
         self.name = name
+        self.args = args
 
     def __enter__(self):
         self.watch.restart()
@@ -31,7 +78,8 @@ class TimeIt(object):
     def __exit__(self, exc_tp, exc_v, exc_tb):
         self.watch.stop()
         duration = self.watch.elapsed()
-        print("Took %0.3f seconds to run '%s'" % (duration, self.name))
+        print_header(self.name)
+        print("- Took %0.3f seconds to run" % (duration))
 
 
 class DummyTask(task.Task):
@@ -40,20 +88,41 @@ class DummyTask(task.Task):
 
 
 def main():
-    dummy_am = 100
-    with TimeIt("building"):
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('--profile', "-p",
+                        dest='profile', action='store_true',
+                        default=False,
+                        help='profile instead of gather timing'
+                             ' (default: False)')
+    parser.add_argument('--dummies', "-d",
+                        dest='dummies', action='store', type=int,
+                        default=100, metavar="<number>",
+                        help='how many dummy/no-op tasks to inject'
+                             ' (default: 100)')
+    parser.add_argument('--limit', '-l',
+                        dest='limit', action='store', type=float,
+                        default=100.0, metavar="<number>",
+                        help='percentage of profiling output to show'
+                             ' (default: 100%%)')
+    args = parser.parse_args()
+    if args.profile:
+        ctx_manager = ProfileIt
+    else:
+        ctx_manager = TimeIt
+    dummy_am = max(0, args.dummies)
+    with ctx_manager("Building linear flow with %s tasks" % dummy_am, args):
         f = lf.Flow("root")
         for i in compat_range(0, dummy_am):
             f.add(DummyTask(name="dummy_%s" % i))
-    with TimeIt("loading"):
+    with ctx_manager("Loading", args):
         e = engines.load(f)
-    with TimeIt("compiling"):
+    with ctx_manager("Compiling", args):
         e.compile()
-    with TimeIt("preparing"):
+    with ctx_manager("Preparing", args):
         e.prepare()
-    with TimeIt("validating"):
+    with ctx_manager("Validating", args):
         e.validate()
-    with TimeIt("running"):
+    with ctx_manager("Running", args):
         e.run()