feat(load-task): support specifying volumes

Author: ahal

src/taskgraph/docker.py

@@ -380,6 +380,7 @@ def load_task(
     user: Optional[str] = None,
     custom_image: Optional[str] = None,
     interactive: Optional[bool] = False,
+    volumes: Optional[Dict[str, str]] = None,
 ) -> int:
     """Load and run a task interactively in a Docker container.
 
@@ -498,6 +499,11 @@ def load_task(
 
         if task_command:
             command.extend(task_command)
+
+        if volumes:
+            for k, v in volumes.items():
+                command[2:2] = ["-v", f"{k}:{v}"]
+
         if remove:
             command.insert(2, "--rm")
 

test/test_docker.py

@@ -111,7 +111,12 @@ def mock_run(*popenargs, check=False, **kwargs):
 @pytest.fixture
 def run_load_task(mocker):
     def inner(
-        task, remove=False, custom_image=None, pass_task_def=False, interactive=True
+        task,
+        remove=False,
+        custom_image=None,
+        pass_task_def=False,
+        interactive=True,
+        volumes=None,
     ):
         proc = mocker.MagicMock()
         proc.returncode = 0
@@ -159,6 +164,7 @@ def inner(
             remove=remove,
             custom_image=custom_image,
             interactive=interactive,
+            volumes=volumes,
         )
         return ret, mocks
 
@@ -198,7 +204,9 @@ def test_load_task(run_load_task):
             "image": {"taskId": image_task_id, "type": "task-image"},
         },
     }
-    ret, mocks = run_load_task(task)
+    # Test with custom volumes
+    volumes = {"/host/path": "/container/path", "/another/host": "/another/container"}
+    ret, mocks = run_load_task(task, volumes=volumes)
     assert ret == 0
 
     if "get_task_definition" in mocks:
@@ -211,6 +219,10 @@ def test_load_task(run_load_task):
         "-v",
         re.compile(f"{tempfile.gettempdir()}/tmp.*:/builds/worker/.bashrc"),
         re.compile(f"--env-file={tempfile.gettempdir()}/tmp.*"),
+        "-v",
+        "/another/host:/another/container",
+        "-v",
+        "/host/path:/container/path",
         "-i",
         "-t",
         "image/tag",