depfile: buildcache support (#33315)

When installing some/all specs from a buildcache, build edges are pruned
from those specs. This can result in a much smaller effective DAG. Until
now, `spack env depfile` would always generate a full DAG.

Ths PR adds the `spack env depfile --use-buildcache` flag that was
introduced for `spack install` before. This way, not only can we drop
build edges, but also we can automatically set the right buildcache
related flags on the specific specs that are gonna get installed.

This way we get parallel installs of binary deps without redundancy,
which is useful for Gitlab CI.

2 files changed, 206 insertions, 35 deletions

diff --git a/lib/spack/spack/cmd/env.py b/lib/spack/spack/cmd/env.py
index e421e300e1..02d279b054 100644
--- a/lib/spack/spack/cmd/env.py
+++ b/lib/spack/spack/cmd/env.py
@@ -3,6 +3,7 @@
 #
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
 
+import argparse
 import os
 import shutil
 import sys
@@ -15,6 +16,8 @@ import llnl.util.tty as tty
 from llnl.util.tty.colify import colify
 from llnl.util.tty.color import colorize
 
+import spack.cmd
+import spack.cmd.common
 import spack.cmd.common.arguments
 import spack.cmd.common.arguments as arguments
 import spack.cmd.install
@@ -600,6 +603,15 @@ def env_depfile_setup_parser(subparser):
         help="disable POSIX jobserver support.",
     )
     subparser.add_argument(
+        "--use-buildcache",
+        dest="use_buildcache",
+        type=arguments.use_buildcache,
+        default="package:auto,dependencies:auto",
+        metavar="[{auto,only,never},][package:{auto,only,never},][dependencies:{auto,only,never}]",
+        help="When using `only`, redundant build dependencies are pruned from the DAG. "
+        "This flag is passed on to the generated spack install commands.",
+    )
+    subparser.add_argument(
         "-o",
         "--output",
         default=None,
@@ -613,6 +625,96 @@ def env_depfile_setup_parser(subparser):
         choices=("make",),
         help="specify the depfile type. Currently only make is supported.",
     )
+    subparser.add_argument(
+        metavar="specs",
+        dest="specs",
+        nargs=argparse.REMAINDER,
+        default=None,
+        help="generate a depfile only for matching specs in the environment",
+    )
+
+
+class SpecNode(object):
+    def __init__(self, spec, depth):
+        self.spec = spec
+        self.depth = depth
+
+    def key(self):
+        return self.spec.dag_hash()
+
+
+class UniqueNodesQueue(object):
+    def __init__(self, init=[]):
+        self.seen = set()
+        self.queue = []
+        for item in init:
+            self.push(item)
+
+    def push(self, item):
+        key = item.key()
+        if key in self.seen:
+            return
+        self.queue.append(item)
+        self.seen.add(key)
+
+    def empty(self):
+        return len(self.queue) == 0
+
+    def pop(self):
+        return self.queue.pop()
+
+
+def _deptypes(use_buildcache):
+    """What edges should we follow for a given node? If it's a cache-only
+    node, then we can drop build type deps."""
+    return ("link", "run") if use_buildcache == "only" else ("build", "link", "run")
+
+
+class MakeTargetVisitor(object):
+    """This visitor produces an adjacency list of a (reduced) DAG, which
+    is used to generate Makefile targets with their prerequisites."""
+
+    def __init__(self, target, pkg_buildcache, deps_buildcache):
+        """
+        Args:
+            target: function that maps dag_hash -> make target string
+            pkg_buildcache (str): "only", "never", "auto": when "only",
+                redundant build deps of roots are dropped
+            deps_buildcache (str): same as pkg_buildcache, but for non-root specs.
+        """
+        self.adjacency_list = []
+        self.target = target
+        self.pkg_buildcache = pkg_buildcache
+        self.deps_buildcache = deps_buildcache
+        self.deptypes_root = _deptypes(pkg_buildcache)
+        self.deptypes_deps = _deptypes(deps_buildcache)
+
+    def neighbors(self, node):
+        """Produce a list of spec to follow from node"""
+        deptypes = self.deptypes_root if node.depth == 0 else self.deptypes_deps
+        return node.spec.dependencies(deptype=deptypes)
+
+    def visit(self, node):
+        dag_hash = node.spec.dag_hash()
+        spec_str = node.spec.format("{name}{@version}{%compiler}{variants}{arch=architecture}")
+        buildcache = self.pkg_buildcache if node.depth == 0 else self.deps_buildcache
+        if buildcache == "only":
+            build_cache_flag = "--use-buildcache=only"
+        elif buildcache == "never":
+            build_cache_flag = "--use-buildcache=never"
+        else:
+            build_cache_flag = ""
+        prereqs = " ".join([self.target(dep.dag_hash()) for dep in self.neighbors(node)])
+        self.adjacency_list.append((dag_hash, spec_str, build_cache_flag, prereqs))
+
+
+def traverse_breadth_first(visitor, specs=[]):
+    queue = UniqueNodesQueue([SpecNode(s, 0) for s in specs])
+    while not queue.empty():
+        node = queue.pop()
+        visitor.visit(node)
+        for child in visitor.neighbors(node):
+            queue.push(SpecNode(child, node.depth + 1))
 
 
 def env_depfile(args):
@@ -620,10 +722,6 @@ def env_depfile(args):
     spack.cmd.require_active_env(cmd_name="env depfile")
     env = ev.active_environment()
 
-    # Maps each hash in the environment to a string of install prereqs
-    hash_to_prereqs = {}
-    hash_to_spec = {}
-
     if args.make_target_prefix is None:
         target_prefix = os.path.join(env.env_subdir_path, "makedeps")
     else:
@@ -645,48 +743,44 @@ def env_depfile(args):
     def get_install_deps_target(name):
         return os.path.join(target_prefix, ".install-deps", name)
 
-    for _, spec in env.concretized_specs():
-        for s in spec.traverse(root=True):
-            hash_to_spec[s.dag_hash()] = s
-            hash_to_prereqs[s.dag_hash()] = [
-                get_install_target(dep.dag_hash()) for dep in s.dependencies()
-            ]
+    # What things do we build when running make? By default, we build the
+    # root specs. If specific specs are provided as input, we build those.
+    if args.specs:
+        abstract_specs = spack.cmd.parse_specs(args.specs)
+        roots = [env.matching_spec(s) for s in abstract_specs]
+    else:
+        roots = [s for _, s in env.concretized_specs()]
 
-    root_dags = [s.dag_hash() for _, s in env.concretized_specs()]
+    # Shallow means we will drop non-direct build deps from the DAG
+    pkg_buildcache, dep_buildcache = args.use_buildcache
+    visitor = MakeTargetVisitor(get_install_target, pkg_buildcache, dep_buildcache)
+    traverse_breadth_first(visitor, roots)
 
     # Root specs without deps are the prereqs for the environment target
-    root_install_targets = [get_install_target(h) for h in root_dags]
+    root_install_targets = [get_install_target(h.dag_hash()) for h in roots]
 
-    # All package install targets, not just roots.
-    all_install_targets = [get_install_target(h) for h in hash_to_spec.keys()]
-    all_install_deps_targets = [get_install_deps_target(h) for h, _ in hash_to_prereqs.items()]
+    # Cleanable targets...
+    cleanable_targets = [get_install_target(h) for h, _, _, _ in visitor.adjacency_list]
+    cleanable_targets.extend([get_install_deps_target(h) for h, _, _, _ in visitor.adjacency_list])
 
     buf = six.StringIO()
 
     template = spack.tengine.make_environment().get_template(os.path.join("depfile", "Makefile"))
 
-    fmt = "{name}{@version}{%compiler}{variants}{arch=architecture}"
-    hash_with_name = [(h, hash_to_spec[h].format(fmt)) for h in hash_to_prereqs.keys()]
-    targets_to_prereqs = [
-        (get_install_deps_target(h), " ".join(prereqs)) for h, prereqs in hash_to_prereqs.items()
-    ]
-
     rendered = template.render(
         {
             "all_target": get_target("all"),
             "env_target": get_target("env"),
             "clean_target": get_target("clean"),
-            "all_install_targets": " ".join(all_install_targets),
-            "all_install_deps_targets": " ".join(all_install_deps_targets),
+            "cleanable_targets": " ".join(cleanable_targets),
             "root_install_targets": " ".join(root_install_targets),
             "dirs_target": get_target("dirs"),
             "environment": env.path,
             "install_target": get_target(".install"),
             "install_deps_target": get_target(".install-deps"),
             "any_hash_target": get_target("%"),
-            "hash_with_name": hash_with_name,
             "jobserver_support": "+" if args.jobserver else "",
-            "targets_to_prereqs": targets_to_prereqs,
+            "adjacency_list": visitor.adjacency_list,
         }
     )
 
diff --git a/lib/spack/spack/test/cmd/env.py b/lib/spack/spack/test/cmd/env.py
index 5cf4b09cae..6257431ee9 100644
--- a/lib/spack/spack/test/cmd/env.py
+++ b/lib/spack/spack/test/cmd/env.py
@@ -3030,29 +3030,106 @@ def test_read_legacy_lockfile_and_reconcretize(mock_stage, mock_fetch, install_m
     assert current_versions == expected_versions
 
 
-def test_environment_depfile_makefile(tmpdir, mock_packages):
+@pytest.mark.parametrize(
+    "depfile_flags,expected_installs",
+    [
+        # This installs the full environment
+        (
+            ["--use-buildcache=never"],
+            [
+                "dtbuild1",
+                "dtbuild2",
+                "dtbuild3",
+                "dtlink1",
+                "dtlink2",
+                "dtlink3",
+                "dtlink4",
+                "dtlink5",
+                "dtrun1",
+                "dtrun2",
+                "dtrun3",
+                "dttop",
+            ],
+        ),
+        # This prunes build deps at depth > 0
+        (
+            ["--use-buildcache=package:never,dependencies:only"],
+            [
+                "dtbuild1",
+                "dtlink1",
+                "dtlink2",
+                "dtlink3",
+                "dtlink4",
+                "dtlink5",
+                "dtrun1",
+                "dtrun2",
+                "dtrun3",
+                "dttop",
+            ],
+        ),
+        # This prunes all build deps
+        (
+            ["--use-buildcache=only"],
+            [
+                "dtlink1",
+                "dtlink3",
+                "dtlink4",
+                "dtlink5",
+                "dtrun1",
+                "dtrun3",
+                "dttop",
+            ],
+        ),
+        # Test whether pruning of build deps is correct if we explicitly include one
+        # that is also a dependency of a root.
+        (
+            ["--use-buildcache=only", "dttop", "dtbuild1"],
+            [
+                "dtbuild1",
+                "dtlink1",
+                "dtlink2",
+                "dtlink3",
+                "dtlink4",
+                "dtlink5",
+                "dtrun1",
+                "dtrun2",
+                "dtrun3",
+                "dttop",
+            ],
+        ),
+    ],
+)
+def test_environment_depfile_makefile(depfile_flags, expected_installs, tmpdir, mock_packages):
     env("create", "test")
     make = Executable("make")
     makefile = str(tmpdir.join("Makefile"))
     with ev.read("test"):
-        add("libdwarf")
+        add("dttop")
         concretize()
 
     # Disable jobserver so we can do a dry run.
     with ev.read("test"):
         env(
-            "depfile", "-o", makefile, "--make-disable-jobserver", "--make-target-prefix", "prefix"
+            "depfile",
+            "-o",
+            makefile,
+            "--make-disable-jobserver",
+            "--make-target-prefix=prefix",
+            *depfile_flags
         )
 
     # Do make dry run.
-    all_out = make("-n", "-f", makefile, output=str)
+    out = make("-n", "-f", makefile, output=str)
 
-    # Check whether `make` installs everything
-    with ev.read("test") as e:
-        for _, root in e.concretized_specs():
-            for spec in root.traverse(root=True):
-                tgt = os.path.join("prefix", ".install", spec.dag_hash())
-                assert "touch {}".format(tgt) in all_out
+    # Spack install commands are of the form "spack install ... # <spec>",
+    # so we just parse the spec again, for simplicity.
+    specs_that_make_would_install = [
+        Spec(line.split("# ")[1]).name for line in out.splitlines() if line.startswith("spack")
+    ]
+
+    # Check that all specs are there (without duplicates)
+    assert set(specs_that_make_would_install) == set(expected_installs)
+    assert len(specs_that_make_would_install) == len(expected_installs)
 
 
 def test_environment_depfile_out(tmpdir, mock_packages):