feat: support setting repo wide default values for cache selection

Author: ahal

packages/pytest-taskgraph/src/pytest_taskgraph/fixtures/gen.py

@@ -204,7 +204,11 @@ def inner(
         if extra_params:
             parameters.update(extra_params)
         if extra_graph_config:
-            graph_config._config.update(extra_graph_config)
+            # We need this intermediate variable because `GraphConfig` is
+            # frozen and we can't set attributes on it.
+            new_graph_config = merge(graph_config._config, extra_graph_config)
+            graph_config._config.update(new_graph_config)
+
         return TransformConfig(
             "test",
             str(here),
@@ -220,12 +224,12 @@ def inner(
 
 @pytest.fixture
 def run_transform(make_transform_config):
-    def inner(func, tasks, config=None):
+    def inner(func, tasks, config=None, **kwargs):
         if not isinstance(tasks, list):
             tasks = [tasks]
 
         if not config:
-            config = make_transform_config()
+            config = make_transform_config(**kwargs)
         return list(func(config, tasks))
 
     return inner

src/taskgraph/config.py

@@ -12,6 +12,7 @@
 from voluptuous import All, Any, Extra, Length, Optional, Required
 
 from .util import path
+from .util.caches import CACHES
 from .util.python_path import find_object
 from .util.schema import Schema, optionally_keyed_by, validate_schema
 from .util.yaml import load_yaml
@@ -74,6 +75,16 @@
                 "index-path-regexes",
                 description="Regular expressions matching index paths to be summarized.",
             ): [str],
+            Optional(
+                "run",
+                description="Configuration related to the 'run' transforms.",
+            ): {
+                Optional(
+                    "use-caches",
+                    description="List of caches to enable, or a boolean to "
+                    "enable/disable all of them.",
+                ): Any(bool, list(CACHES.keys())),
+            },
             Required("repositories"): All(
                 {
                     str: {
@@ -106,8 +117,8 @@ def __getitem__(self, name):
     def __contains__(self, name):
         return name in self._config
 
-    def get(self, name):
-        return self._config.get(name)
+    def get(self, name, default=None):
+        return self._config.get(name, default)
 
     def register(self):
         """

src/taskgraph/transforms/run/common.py

@@ -7,12 +7,12 @@
 consistency.
 """
 
-import hashlib
 import json
 from typing import Any, Dict, List, Union
 
 from taskgraph.transforms.base import TransformConfig
 from taskgraph.util import path
+from taskgraph.util.caches import CACHES, get_checkout_dir
 from taskgraph.util.taskcluster import get_artifact_prefix
 
 
@@ -23,54 +23,6 @@ def get_vcsdir_name(os):
         return "vcs"
 
 
-def get_checkout_dir(task: Dict[str, Any]) -> str:
-    worker = task["worker"]
-    if worker["os"] == "windows":
-        return "build"
-    elif worker["implementation"] == "docker-worker":
-        return f"{task['run']['workdir']}/checkouts"
-    else:
-        return "checkouts"
-
-
-def get_checkout_cache_name(config: TransformConfig, task: Dict[str, Any]) -> str:
-    repo_configs = config.repo_configs
-    cache_name = "checkouts"
-
-    # Robust checkout does not clean up subrepositories, so ensure  that tasks
-    # that checkout different sets of paths have separate caches.
-    # See https://bugzilla.mozilla.org/show_bug.cgi?id=1631610
-    if len(repo_configs) > 1:
-        checkout_paths = {
-            "\t".join([repo_config.path, repo_config.prefix])
-            for repo_config in sorted(
-                repo_configs.values(), key=lambda repo_config: repo_config.path
-            )
-        }
-        checkout_paths_str = "\n".join(checkout_paths).encode("utf-8")
-        digest = hashlib.sha256(checkout_paths_str).hexdigest()
-        cache_name += f"-repos-{digest}"
-
-    # Sparse checkouts need their own cache because they can interfere
-    # with clients that aren't sparse aware.
-    if task["run"]["sparse-profile"]:
-        cache_name += "-sparse"
-
-    return cache_name
-
-
-CACHES = {
-    "cargo": {"env": "CARGO_HOME"},
-    "checkout": {
-        "cache_dir": get_checkout_dir,
-        "cache_name": get_checkout_cache_name,
-    },
-    "npm": {"env": "npm_config_cache"},
-    "pip": {"env": "PIP_CACHE_DIR"},
-    "uv": {"env": "UV_CACHE_DIR"},
-}
-
-
 def add_cache(task, taskdesc, name, mount_point, skip_untrusted=False):
     """Adds a cache based on the worker's implementation.
 
@@ -216,8 +168,15 @@ def support_caches(
         workdir = workdir or "/builds/worker"
         base_cache_dir = path.join(workdir, base_cache_dir)
 
-    # Default to checkout cache if no selection is specified.
-    use_caches = run.get("use-caches", ["checkout"])
+    use_caches = run.get("use-caches")
+    if use_caches is None:
+        # Use project default values for filtering caches, default to
+        # checkout cache if no selection is specified.
+        use_caches = (
+            config.graph_config.get("taskgraph", {})
+            .get("run", {})
+            .get("use-caches", ["checkout"])
+        )
 
     for name, cache_cfg in CACHES.items():
         if not should_use_cache(name, use_caches, run["checkout"]):

src/taskgraph/transforms/run/run_task.py

@@ -12,12 +12,12 @@
 
 from taskgraph.transforms.run import run_task_using
 from taskgraph.transforms.run.common import (
-    CACHES,
     support_caches,
     support_vcs_checkout,
 )
 from taskgraph.transforms.task import taskref_or_string
 from taskgraph.util import path, taskcluster
+from taskgraph.util.caches import CACHES
 from taskgraph.util.schema import Schema
 
 EXEC_COMMANDS = {

src/taskgraph/util/caches.py

@@ -0,0 +1,57 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+import hashlib
+from typing import TYPE_CHECKING, Any, Dict
+
+if TYPE_CHECKING:
+    from taskgraph.transforms.base import TransformConfig
+
+
+def get_checkout_dir(task: Dict[str, Any]) -> str:
+    worker = task["worker"]
+    if worker["os"] == "windows":
+        return "build"
+    elif worker["implementation"] == "docker-worker":
+        return f"{task['run']['workdir']}/checkouts"
+    else:
+        return "checkouts"
+
+
+def get_checkout_cache_name(config: "TransformConfig", task: Dict[str, Any]) -> str:
+    repo_configs = config.repo_configs
+    cache_name = "checkouts"
+
+    # Robust checkout does not clean up subrepositories, so ensure  that tasks
+    # that checkout different sets of paths have separate caches.
+    # See https://bugzilla.mozilla.org/show_bug.cgi?id=1631610
+    if len(repo_configs) > 1:
+        checkout_paths = {
+            "\t".join([repo_config.path, repo_config.prefix])
+            for repo_config in sorted(
+                repo_configs.values(), key=lambda repo_config: repo_config.path
+            )
+        }
+        checkout_paths_str = "\n".join(checkout_paths).encode("utf-8")
+        digest = hashlib.sha256(checkout_paths_str).hexdigest()
+        cache_name += f"-repos-{digest}"
+
+    # Sparse checkouts need their own cache because they can interfere
+    # with clients that aren't sparse aware.
+    if task["run"]["sparse-profile"]:
+        cache_name += "-sparse"
+
+    return cache_name
+
+
+CACHES = {
+    "cargo": {"env": "CARGO_HOME"},
+    "checkout": {
+        "cache_dir": get_checkout_dir,
+        "cache_name": get_checkout_cache_name,
+    },
+    "npm": {"env": "npm_config_cache"},
+    "pip": {"env": "PIP_CACHE_DIR"},
+    "uv": {"env": "UV_CACHE_DIR"},
+}

taskcluster/docker/python/Dockerfile

@@ -6,7 +6,7 @@ FROM debian:bookworm-slim
 LABEL maintainer="Release Engineering <release@mozilla.com>"
 
 VOLUME /builds/worker/checkouts
-VOLUME /builds/worker/.cache
+VOLUME /builds/worker/.task-cache/uv
 
 # Add worker user
 RUN mkdir -p /builds && \

test/test_transforms_run.py

@@ -43,13 +43,13 @@ def transform(monkeypatch, run_transform):
     # Needed by 'generic_worker_run_task'
     monkeypatch.setenv("TASK_ID", "fakeid")
 
-    def inner(task_input):
+    def inner(task_input, **kwargs):
         defaults = deepcopy(TASK_DEFAULTS)
         task = merge(defaults, task_input)
 
         with patch("taskgraph.transforms.run.configure_taskdesc_for_run") as m:
             # This forces the generator to be evaluated
-            run_transform(run.transforms, task)
+            run_transform(run.transforms, task, **kwargs)
             return m.call_args[0]
 
     return inner

test/test_transforms_run_run_task.py

@@ -8,8 +8,8 @@
 import pytest
 
 from taskgraph.transforms.run import make_task_description
-from taskgraph.transforms.run.common import CACHES
 from taskgraph.transforms.task import payload_builders
+from taskgraph.util.caches import CACHES
 from taskgraph.util.schema import Schema, validate_schema
 from taskgraph.util.templates import merge
 
@@ -36,9 +36,9 @@ def run_task_using(mocker, run_transform):
         "taskcluster/scripts/toolchain/run.ps1",
     ]
 
-    def inner(task):
+    def inner(task, **kwargs):
         task = merge(TASK_DEFAULTS, task)
-        return run_transform(make_task_description, task)[0]
+        return run_transform(make_task_description, task, **kwargs)[0]
 
     return inner
 
@@ -225,12 +225,12 @@ def test_run_task(monkeypatch, request, run_task_using, task):
 
 @pytest.fixture
 def run_caches(run_task_using):
-    def inner(task):
+    def inner(task, **kwargs):
         task.setdefault("run", {}).setdefault("using", "run-task")
         impl = task.setdefault("worker", {}).setdefault(
             "implementation", "generic-worker"
         )
-        result = run_task_using(task)
+        result = run_task_using(task, **kwargs)
 
         key = "mounts" if impl == "generic-worker" else "caches"
 
@@ -281,3 +281,13 @@ def test_caches_explicit(run_caches):
     assert run_caches(task) == [
         {"cache-name": "cargo", "directory": "/builds/worker/.task-cache/cargo"}
     ]
+
+
+def test_caches_project_explicit(run_caches):
+    caches = run_caches(
+        {},
+        extra_graph_config={"taskgraph": {"run": {"use-caches": ["cargo"]}}},
+    )
+    assert caches == [
+        {"cache-name": "cargo", "directory": "/builds/worker/.task-cache/cargo"}
+    ]