Arvados-DCO-1.1-Signed-off-by: Radhika Chippada <radhika@curoverse.com>
[arvados.git] / 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