load-task: allow specifying caches on the command line

When debugging a task it can be useful to re-use e.g. the checkout cache
between load-task invocations, to not have to clone from scratch each
time.

Author: jcristau

src/taskgraph/docker.py

@@ -522,11 +522,6 @@ def load_task(
     # Add the task's environment variables.
     env.update(task_def["payload"].get("env", {}))  # type: ignore
 
-    # run-task expects the worker to mount a volume for each path defined in
-    # TASKCLUSTER_CACHES, delete them to avoid needing to do the same.
-    if "TASKCLUSTER_CACHES" in env:
-        del env["TASKCLUSTER_CACHES"]
-
     envfile = None
     initfile = None
     isatty = os.isatty(sys.stdin.fileno())
@@ -569,6 +564,14 @@ def load_task(
 
             command.extend(["-v", f"{initfile.name}:/builds/worker/.bashrc"])
 
+        # run-task expects the worker to mount a volume for each path defined in
+        # TASKCLUSTER_CACHES
+        if "TASKCLUSTER_CACHES" in env:
+            caches = env["TASKCLUSTER_CACHES"].split(";")
+            for cache in caches:
+                if not volumes or cache not in volumes:
+                    command.extend(["-v", cache])
+
         if volumes:
             for k, v in volumes.items():
                 command.extend(["-v", f"{k}:{v}"])

src/taskgraph/main.py

@@ -792,6 +792,14 @@ def image_digest(args):
     default="taskcluster",
     help="Relative path to the root of the Taskgraph definition.",
 )
+@argument(
+    "--volume",
+    "-v",
+    metavar="HOST_DIR:CONTAINER_DIR",
+    default=[],
+    action="append",
+    help="Mount local path into the container.",
+)
 def load_task(args):
     from taskgraph.config import load_graph_config  # noqa: PLC0415
     from taskgraph.docker import load_task  # noqa: PLC0415
@@ -806,6 +814,11 @@ def load_task(args):
         except ValueError:
             args["task"] = data  # assume it is a taskId
 
+    volumes = {}
+    for vol in args["volume"]:
+        k, v = vol.split(":", 1)
+        volumes[k] = v
+
     root = args["root"]
     graph_config = load_graph_config(root)
     return load_task(
@@ -815,6 +828,7 @@ def load_task(args):
         remove=args["remove"],
         user=args["user"],
         custom_image=args["image"],
+        volumes=volumes,
     )
 
 

test/test_docker.py

@@ -231,7 +231,7 @@ def test_load_task_env_init_and_remove(mocker, run_load_task):
     env_lines = written_env_content[0].split("\n")
 
     # Verify written env is expected
-    assert "TASKCLUSTER_CACHES=path" not in env_lines
+    assert "TASKCLUSTER_CACHES=path" in env_lines
     assert "FOO=BAR" in env_lines
     assert "BAZ=1" in env_lines
 
@@ -253,7 +253,7 @@ def test_load_task_env_init_and_remove(mocker, run_load_task):
     actual = mocks["subprocess_run"].call_args[0][0]
     assert actual[4] == "--rm"
     assert actual[5] == "--env-file=/tmp/test_envfile"
-    assert actual[6:8] == ["-v", "/tmp/test_initfile:/builds/worker/.bashrc"]
+    assert actual[6:10] == ["-v", "/tmp/test_initfile:/builds/worker/.bashrc", "-v", "path"]
 
 
 @pytest.mark.parametrize(

test/test_main.py

@@ -466,6 +466,7 @@ def test_load_task_command(run_load_task):
         remove=True,
         user=None,
         custom_image=None,
+        volumes={},
     )
 
     # Test with interactive flag
@@ -479,6 +480,7 @@ def test_load_task_command(run_load_task):
         remove=True,
         user=None,
         custom_image=None,
+        volumes={},
     )
 
 
@@ -503,6 +505,7 @@ def test_load_task_command_with_stdin(run_load_task):
         remove=True,
         user=None,
         custom_image=None,
+        volumes={},
     )
 
 
@@ -520,6 +523,7 @@ def test_load_task_command_with_task_id(run_load_task):
         remove=True,
         user=None,
         custom_image=None,
+        volumes={},
     )