sdk/python/tests/performance/performance_profiler.py

   1 # Use the "profiled" decorator on a test to get profiling data.
   2 #
   3 # Usage:
   4 #   from performance_profiler import profiled
   5 #
   6 #   # See report in tmp/profile/foobar
   7 #   @profiled
   8 #   def foobar():
   9 #       baz = 1
  10 #
  11 #   See "test_a_sample.py" for a working example.
  12 #
  13 # Performance tests run as part of regular test suite.
  14 # You can also run only the performance tests using one of the following:
  15 #     python -m unittest discover tests.performance
  16 #     ./run-tests.sh WORKSPACE=~/arvados --only sdk/python sdk/python_test="--test-suite=tests.performance"
  17
  18 import functools
  19 import os
  20 import pstats
  21 import sys
  22 import unittest
  23 try:
  24     import cProfile as profile
  25 except ImportError:
  26     import profile
  27
  28 output_dir = os.path.abspath(os.path.join('tmp', 'profile'))
  29 if not os.path.exists(output_dir):
  30     os.makedirs(output_dir)
  31
  32 def profiled(function):
  33     @functools.wraps(function)
  34     def profiled_function(*args, **kwargs):
  35         outfile = open(os.path.join(output_dir, function.__name__), "w")
  36         caught = None
  37         pr = profile.Profile()
  38         pr.enable()
  39         try:
  40             return function(*args, **kwargs)
  41         finally:
  42             pr.disable()
  43             ps = pstats.Stats(pr, stream=outfile)
  44             ps.sort_stats('time').print_stats()
  45     return profiled_function