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

@@ -523,9 +523,18 @@ def load_task(
     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.
+    # TASKCLUSTER_CACHES; delete them to avoid needing to do the same, unless
+    # they're passed in as volumes.
     if "TASKCLUSTER_CACHES" in env:
-        del env["TASKCLUSTER_CACHES"]
+        if volumes:
+            caches = env["TASKCLUSTER_CACHES"].split(";")
+            caches = [cache for cache in caches if cache in volumes.values()]
+        else:
+            caches = []
+        if caches:
+            env["TASKCLUSTER_CACHES"] = ";".join(caches)
+        else:
+            del env["TASKCLUSTER_CACHES"]
 
     envfile = None
     initfile = None

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

@@ -216,11 +216,11 @@ def test_load_task_env_init_and_remove(mocker, run_load_task):
                 "--",
                 "echo foo",
             ],
-            "env": {"FOO": "BAR", "BAZ": "1", "TASKCLUSTER_CACHES": "path"},
+            "env": {"FOO": "BAR", "BAZ": "1", "TASKCLUSTER_CACHES": "/path;/cache"},
             "image": {"taskId": image_task_id, "type": "task-image"},
         },
     }
-    ret, mocks = run_load_task(task, remove=True)
+    ret, mocks = run_load_task(task, remove=True, volumes={"/host/path": "/cache"})
     assert ret == 0
 
     # NamedTemporaryFile was called twice (once for env, once for init)
@@ -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=/cache" in env_lines
     assert "FOO=BAR" in env_lines
     assert "BAZ=1" in env_lines
 

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={},
     )