From 75e599490c1d335aa61529e243f9e7857db70dd7 Mon Sep 17 00:00:00 2001 From: Todd Gamblin Date: Tue, 17 Sep 2024 15:32:08 -0700 Subject: [PATCH] `spack solve`: add a `--profile` option We often want to know what the solver is thinking about in order to optimize it. We can generate a profile of which atoms are most frequently propagated and/or undone by the solver using a custom propagator class. This adds a `spack.solver.profiler` module with a simple propagator that counts `undo` and `propagate` operations and associates them with atoms in the solve. To simplify things, rather than enumerating *every* atom in the logic program, we coarsen to show only function names for everything but `attr()`, and we show the first argument of any `attr()` atom (e.g. `attr("node")`). `spack solve --profile` will print the top 40 most propagated atoms in the program, along with how many times they are undone (these are usually but not always similar). The list can give people working with the concretizer a rough idea of where the solver is spending its time, and hopefully point to areas of the program to optimize. This is based on some code that Tom Scogland wrote a year or so ago. I did the work of integrating it as an option so that I do not have to keep looking it up and running this from the CLI :). Co-authored-by: Tom Scogland Signed-off-by: Todd Gamblin --- lib/spack/spack/cmd/solve.py | 9 ++ lib/spack/spack/solver/asp.py | 32 ++++- lib/spack/spack/solver/profiler.py | 111 ++++++++++++++++++ lib/spack/spack/test/concretization/core.py | 12 ++ .../tutorial/.spack-env/transaction_lock | 0 share/spack/spack-completion.bash | 2 +- share/spack/spack-completion.fish | 5 +- 7 files changed, 166 insertions(+), 5 deletions(-) create mode 100644 lib/spack/spack/solver/profiler.py create mode 100755 share/spack/gitlab/cloud_pipelines/stacks/tutorial/.spack-env/transaction_lock diff --git a/lib/spack/spack/cmd/solve.py b/lib/spack/spack/cmd/solve.py index 8adc06fdfbe..957dd039b95 100644 --- a/lib/spack/spack/cmd/solve.py +++ b/lib/spack/spack/cmd/solve.py @@ -3,6 +3,7 @@ # # SPDX-License-Identifier: (Apache-2.0 OR MIT) +import argparse import re import sys @@ -48,6 +49,13 @@ def setup_parser(subparser): subparser.add_argument( "--stats", action="store_true", default=False, help="print out statistics from clingo" ) + subparser.add_argument( + "--profile", + action="store_true", + default=False, + help="profile the solve phase and print out statistics on atoms", + ) + subparser.add_argument("specs", nargs=argparse.REMAINDER, help="specs of packages") spack.cmd.spec.setup_parser(subparser) @@ -147,6 +155,7 @@ def solve(parser, args): stats=args.stats, setup_only=setup_only, allow_deprecated=allow_deprecated, + profile=args.profile, ) if not setup_only: _process_result(result, show, required_format, kwargs) diff --git a/lib/spack/spack/solver/asp.py b/lib/spack/spack/solver/asp.py index 5310c4e3a84..21b4429cbd6 100644 --- a/lib/spack/spack/solver/asp.py +++ b/lib/spack/spack/solver/asp.py @@ -64,6 +64,7 @@ parse_term, ) from .counter import FullDuplicatesCounter, MinimalDuplicatesCounter, NoDuplicatesCounter +from .profiler import ProfilePropagator from .version_order import concretization_version_order GitOrStandardVersion = Union[spack.version.GitVersion, spack.version.StandardVersion] @@ -800,7 +801,16 @@ def __init__(self, cores=True): # This attribute will be reset at each call to solve self.control = None - def solve(self, setup, specs, reuse=None, output=None, control=None, allow_deprecated=False): + def solve( + self, + setup, + specs, + reuse=None, + output=None, + control=None, + allow_deprecated=False, + profile=False, + ): """Set up the input and solve for dependencies of ``specs``. Arguments: @@ -826,6 +836,11 @@ def solve(self, setup, specs, reuse=None, output=None, control=None, allow_depre # Initialize the control object for the solver self.control = control or default_clingo_control() + # if profiling is enabled, register a profiling propagator + if profile: + propagator = ProfilePropagator() + self.control.register_propagator(propagator) + # ensure core deps are present on Windows # needs to modify active config scope, so cannot be run within # bootstrap config scope @@ -925,13 +940,18 @@ def on_model(model): result.cores.extend(cores) if output.timers: + tty.msg("Timers:") timer.write_tty() print() if output.stats: - print("Statistics:") + tty.msg("Statistics:") pprint.pprint(self.control.statistics) + if profile: + tty.msg("Profile:") + propagator.print_profile(40) + result.raise_if_unsat() if result.satisfiable and result.unsolved_specs and setup.concretize_everything: @@ -4361,6 +4381,7 @@ def solve_with_stats( tests=False, setup_only=False, allow_deprecated=False, + profile=False, ): """ Concretize a set of specs and track the timing and statistics for the solve @@ -4382,7 +4403,12 @@ def solve_with_stats( setup = SpackSolverSetup(tests=tests) output = OutputConfiguration(timers=timers, stats=stats, out=out, setup_only=setup_only) return self.driver.solve( - setup, specs, reuse=reusable_specs, output=output, allow_deprecated=allow_deprecated + setup, + specs, + reuse=reusable_specs, + output=output, + allow_deprecated=allow_deprecated, + profile=profile, ) def solve(self, specs, **kwargs): diff --git a/lib/spack/spack/solver/profiler.py b/lib/spack/spack/solver/profiler.py new file mode 100644 index 00000000000..4a50f135ee6 --- /dev/null +++ b/lib/spack/spack/solver/profiler.py @@ -0,0 +1,111 @@ +# Copyright 2013-2024 Lawrence Livermore National Security, LLC and other +# Spack Project Developers. See the top-level COPYRIGHT file for details. +# +# SPDX-License-Identifier: (Apache-2.0 OR MIT) +"""Profiling propagator for clingo solves.""" + +import re +from typing import Any, Dict, List + +import llnl.util.tty as tty + + +class Data: + """Counters for propagations and undos of an atom.""" + + __slots__ = ["atom", "literal", "prop", "undo"] + + # currently we use Any for clingo types because clingo has a bunch of import + # wrappers around it that make typing difficult (see spack.solver.core for details) + def __init__(self, atom: Any, literal: int, prop: int, undo: int): + self.atom = atom + self.literal = literal + self.prop = prop + self.undo = undo + + +class AggregatedData: + """Aggregated data for a profile, constructed from ``Data``. + + We coarsen from atom granularity to string keys when aggregating. + """ + + __slots__ = ["name", "prop", "undo"] + + def __init__(self, name: str, prop: int, undo: int): + self.name = name + self.prop = prop + self.undo = undo + + +class ProfilePropagator: + """Profiling propagator for `spack solve --profile`. + + Register this with the ``clingo.Control`` object to profile a solve. + """ + + _literal_to_atom: Dict + _profile: Dict[int, Data] + + def init(self, init) -> None: + self._literal_to_atom = {} + self._profile = {} + for atom in init.symbolic_atoms: + solver_literal = init.solver_literal(atom.literal) + self._profile[solver_literal] = Data(atom, solver_literal, 0, 0) + init.add_watch(solver_literal) + + def propagate(self, ctl, changes: List[int]) -> bool: + """Record a propagation in the solve.""" + for literal in changes: + data = self._profile[literal] + data.prop += 1 + return True + + def undo(self, solver_id: int, assign, undo: List[int]) -> None: + """Record an undo in the solve.""" + for literal in undo: + data = self._profile[literal] + data.undo += 1 + + def color_sym(self, string: str) -> str: + """Colorize a symbol for profile output""" + string = re.sub(r"^(\w+)", r"@C{\1}", string) + string = re.sub(r'("[^"]*")', r"@G{\1}", string) + string = re.sub(r"([\(\)])", r"@b{\1}", string) + return tty.color.colorize(string) + + def key(self, atom) -> str: + """Convert an atom into an aggregate key for our profile. + + Currently this compresses most things to their function name, and expands + ``attr("name", ...)`` to ``attr("name")`` so we can see which attributes affect + the solve most. + + """ + sym = atom.symbol + return f"attr({sym.arguments[0]})" if sym.name == "attr" else sym.name + + def print_profile(self, n_atoms: int) -> None: + """Aggregate and print nicely formatted profile data.""" + aggregated = {} + for data in self._profile.values(): + name = self.key(data.atom) + if name not in aggregated: + aggregated[name] = AggregatedData(name, data.prop, data.undo) + else: + agg = aggregated[name] + agg.prop += data.prop + agg.undo += data.undo + + values = sorted( + (x for x in aggregated.values() if x.prop), key=lambda x: x.prop, reverse=True + ) + + # format the output nicely + w = 10 # width for number fields + print(tty.color.colorize(f" @*{{{'Prop':<{w}}{'Undo':<{w}}{'Symbol'}}}")) + for a in values[:n_atoms]: + print(f" {a.prop:<{w}}{a.undo:<{w}}{self.color_sym(a.name)}") + if len(values) > n_atoms: + print(" ...") diff --git a/lib/spack/spack/test/concretization/core.py b/lib/spack/spack/test/concretization/core.py index e33f9761dad..be2a632916b 100644 --- a/lib/spack/spack/test/concretization/core.py +++ b/lib/spack/spack/test/concretization/core.py @@ -2871,6 +2871,18 @@ def test_specifying_different_versions_build_deps(self): assert any(x.satisfies(hdf5_str) for x in result.specs) assert any(x.satisfies(pinned_str) for x in result.specs) + def test_solve_with_profile(self, capsys): + """For now, just ensure that the profiler runs.""" + solver = spack.solver.asp.Solver() + solver.solve([Spec("hdf5")], profile=True) + + out, _ = capsys.readouterr() + assert "Profile:" in out + assert "Symbol" in out + assert "Prop" in out + assert "Undo" in out + assert "internal_error" in out # symbol is always in small solves + @pytest.mark.parametrize( "v_str,v_opts,checksummed", diff --git a/share/spack/gitlab/cloud_pipelines/stacks/tutorial/.spack-env/transaction_lock b/share/spack/gitlab/cloud_pipelines/stacks/tutorial/.spack-env/transaction_lock new file mode 100755 index 00000000000..e69de29bb2d diff --git a/share/spack/spack-completion.bash b/share/spack/spack-completion.bash index 706fb7e2d59..3ea04e0dd61 100644 --- a/share/spack/spack-completion.bash +++ b/share/spack/spack-completion.bash @@ -1849,7 +1849,7 @@ _spack_restage() { _spack_solve() { if $list_options then - SPACK_COMPREPLY="-h --help --show --timers --stats -l --long -L --very-long -N --namespaces -I --install-status --no-install-status -y --yaml -j --json --format -c --cover -t --types -U --fresh --reuse --fresh-roots --reuse-deps --deprecated" + SPACK_COMPREPLY="-h --help --show --timers --stats --profile -l --long -L --very-long -N --namespaces -I --install-status --no-install-status -y --yaml -j --json --format -c --cover -t --types -U --fresh --reuse --fresh-roots --reuse-deps --deprecated" else _all_packages fi diff --git a/share/spack/spack-completion.fish b/share/spack/spack-completion.fish index 52153f01c0e..2950e27dcba 100644 --- a/share/spack/spack-completion.fish +++ b/share/spack/spack-completion.fish @@ -2799,8 +2799,9 @@ complete -c spack -n '__fish_spack_using_command restage' -s h -l help -f -a hel complete -c spack -n '__fish_spack_using_command restage' -s h -l help -d 'show this help message and exit' # spack solve -set -g __fish_spack_optspecs_spack_solve h/help show= timers stats l/long L/very-long N/namespaces I/install-status no-install-status y/yaml j/json format= c/cover= t/types U/fresh reuse fresh-roots deprecated +set -g __fish_spack_optspecs_spack_solve h/help show= timers stats profile l/long L/very-long N/namespaces I/install-status no-install-status y/yaml j/json format= c/cover= t/types U/fresh reuse fresh-roots deprecated complete -c spack -n '__fish_spack_using_command_pos_remainder 0 solve' -f -k -a '(__fish_spack_specs_or_id)' +complete -c spack -n '__fish_spack_using_command_pos_remainder 1 solve' -f -k -a '(__fish_spack_specs_or_id)' complete -c spack -n '__fish_spack_using_command solve' -s h -l help -f -a help complete -c spack -n '__fish_spack_using_command solve' -s h -l help -d 'show this help message and exit' complete -c spack -n '__fish_spack_using_command solve' -l show -r -f -a show @@ -2809,6 +2810,8 @@ complete -c spack -n '__fish_spack_using_command solve' -l timers -f -a timers complete -c spack -n '__fish_spack_using_command solve' -l timers -d 'print out timers for different solve phases' complete -c spack -n '__fish_spack_using_command solve' -l stats -f -a stats complete -c spack -n '__fish_spack_using_command solve' -l stats -d 'print out statistics from clingo' +complete -c spack -n '__fish_spack_using_command solve' -l profile -f -a profile +complete -c spack -n '__fish_spack_using_command solve' -l profile -d 'profile the solve phase and print out statistics on atoms' complete -c spack -n '__fish_spack_using_command solve' -s l -l long -f -a long complete -c spack -n '__fish_spack_using_command solve' -s l -l long -d 'show dependency hashes as well as versions' complete -c spack -n '__fish_spack_using_command solve' -s L -l very-long -f -a very_long