| # |
| # Copyright (C) 2017 Codethink Limited |
| # Copyright (C) 2019 Bloomberg Finance LP |
| # |
| # This program is free software; you can redistribute it and/or |
| # modify it under the terms of the GNU Lesser General Public |
| # License as published by the Free Software Foundation; either |
| # version 2 of the License, or (at your option) any later version. |
| # |
| # This library is distributed in the hope that it will be useful, |
| # but WITHOUT ANY WARRANTY; without even the implied warranty of |
| # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU |
| # Lesser General Public License for more details. |
| # |
| # You should have received a copy of the GNU Lesser General Public |
| # License along with this library. If not, see <http://www.gnu.org/licenses/>. |
| # |
| # Authors: |
| # Tristan Van Berkom <tristan.vanberkom@codethink.co.uk> |
| # James Ennis <james.ennis@codethink.co.uk> |
| # Benjamin Schubert <bschubert15@bloomberg.net> |
| |
| |
| import contextlib |
| import cProfile |
| import pstats |
| import os |
| import datetime |
| import time |
| |
| |
| # Use the topic values here to decide what to profile |
| # by setting them in the BST_PROFILE environment variable. |
| # |
| # Multiple topics can be set with the ':' separator. |
| # |
| # E.g.: |
| # |
| # BST_PROFILE=circ-dep-check:sort-deps bst <command> <args> |
| # |
| # The special 'all' value will enable all profiles. |
| class Topics(): |
| CIRCULAR_CHECK = 'circ-dep-check' |
| SORT_DEPENDENCIES = 'sort-deps' |
| LOAD_CONTEXT = 'load-context' |
| LOAD_PROJECT = 'load-project' |
| LOAD_PIPELINE = 'load-pipeline' |
| LOAD_SELECTION = 'load-selection' |
| SCHEDULER = 'scheduler' |
| ALL = 'all' |
| |
| |
| class _Profile: |
| def __init__(self, key, message): |
| self.profiler = cProfile.Profile() |
| self._additional_pstats_files = [] |
| |
| self.key = key |
| self.message = message |
| |
| self.start_time = time.time() |
| filename_template = os.path.join( |
| os.getcwd(), |
| "profile-{}-{}".format( |
| datetime.datetime.fromtimestamp(self.start_time).strftime("%Y%m%dT%H%M%S"), |
| self.key.replace("/", "-").replace(".", "-") |
| ) |
| ) |
| self.log_filename = "{}.log".format(filename_template) |
| self.cprofile_filename = "{}.cprofile".format(filename_template) |
| |
| def __enter__(self): |
| self.start() |
| |
| def __exit__(self, exc_type, exc_value, traceback): |
| self.stop() |
| self.save() |
| |
| def merge(self, profile): |
| self._additional_pstats_files.append(profile.cprofile_filename) |
| |
| def start(self): |
| self.profiler.enable() |
| |
| def stop(self): |
| self.profiler.disable() |
| |
| def save(self): |
| heading = "\n".join([ |
| "-" * 64, |
| "Profile for key: {}".format(self.key), |
| "Started at: {}".format(self.start_time), |
| "\n\t{}".format(self.message) if self.message else "", |
| "-" * 64, |
| "" # for a final new line |
| ]) |
| |
| with open(self.log_filename, "a") as fp: |
| stats = pstats.Stats(self.profiler, *self._additional_pstats_files, stream=fp) |
| |
| # Create the log file |
| fp.write(heading) |
| stats.sort_stats("cumulative") |
| stats.print_stats() |
| |
| # Dump the cprofile |
| stats.dump_stats(self.cprofile_filename) |
| |
| |
| class _Profiler: |
| def __init__(self, settings): |
| self.active_topics = set() |
| self.enabled_topics = set() |
| self._active_profilers = [] |
| |
| if settings: |
| self.enabled_topics = { |
| topic |
| for topic in settings.split(":") |
| } |
| |
| @contextlib.contextmanager |
| def profile(self, topic, key, message=None): |
| if not self._is_profile_enabled(topic): |
| yield |
| return |
| |
| if self._active_profilers: |
| # we are in a nested profiler, stop the parent |
| self._active_profilers[-1].stop() |
| |
| key = "{}-{}".format(topic, key) |
| |
| assert key not in self.active_topics |
| self.active_topics.add(key) |
| |
| profiler = _Profile(key, message) |
| self._active_profilers.append(profiler) |
| |
| with profiler: |
| yield |
| |
| self.active_topics.remove(key) |
| |
| # Remove the last profiler from the list |
| self._active_profilers.pop() |
| |
| if self._active_profilers: |
| # We were in a previous profiler, add the previous results to it |
| # and reenable it. |
| parent_profiler = self._active_profilers[-1] |
| parent_profiler.merge(profiler) |
| parent_profiler.start() |
| |
| def _is_profile_enabled(self, topic): |
| return topic in self.enabled_topics or Topics.ALL in self.enabled_topics |
| |
| |
| # Export a profiler to be used by BuildStream |
| PROFILER = _Profiler(os.getenv("BST_PROFILE")) |