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