sdk/python/tests/performance/performance_profiler.py

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