/
profile.py
29 lines (25 loc) · 1.2 KB
/
profile.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
#***************************************************************************************************
# Copyright 2015, 2019 National Technology & Engineering Solutions of Sandia, LLC (NTESS).
# Under the terms of Contract DE-NA0003525 with NTESS, the U.S. Government retains certain rights
# in this software.
# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
# in compliance with the License. You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0 or in the LICENSE file in the root pyGSTi directory.
#***************************************************************************************************
import cProfile
from mpi4py import MPI
def profile(filename=None, comm=MPI.COMM_WORLD):
def prof_decorator(f):
def wrap_f(*args, **kwargs):
pr = cProfile.Profile()
pr.enable()
result = f(*args, **kwargs)
pr.disable()
if filename is None:
pr.print_stats()
else:
filename_r = filename # + ".{}".format(comm.rank)
pr.dump_stats(filename_r)
return result
return wrap_f
return prof_decorator