Various taskgraph load-task fixes

src/taskgraph/docker.py

@@ -19,6 +19,7 @@
 
 from taskgraph.util import docker, json
 from taskgraph.util.taskcluster import (
+    find_task_id,
     get_artifact_url,
     get_root_url,
     get_session,
@@ -269,10 +270,8 @@ def load_task(task_id, remove=True, user=None):
     user = user or "worker"
     task_def = get_task_definition(task_id)
 
-    if (
-        impl := task_def.get("tags", {}).get("worker-implementation")
-    ) != "docker-worker":
-        print(f"Tasks with worker-implementation '{impl}' are not supported!")
+    if "payload" not in task_def or not (image := task_def["payload"].get("image")):
+        print("Tasks without a `payload.image` are not supported!")
         return 1
 
     command = task_def["payload"].get("command")
@@ -308,7 +307,14 @@ def load_task(task_id, remove=True, user=None):
         else:
             task_cwd = "$TASK_WORKDIR"
 
-    image_task_id = task_def["payload"]["image"]["taskId"]
+    if image["type"] == "task-image":
+        image_task_id = image["taskId"]
+    elif image["type"] == "indexed-image":
+        image_task_id = find_task_id(image["namespace"])
+    else:
+        print(f"Tasks with {image['type']} images are not supported!")
+        return 1
+
     image_tag = load_image_by_task_id(image_task_id)
 
     # Set some env vars the worker would normally set.
@@ -321,6 +327,11 @@ def load_task(task_id, remove=True, user=None):
     # Add the task's environment variables.
     env.update(task_def["payload"].get("env", {}))
 
+    # 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
     try:

test/test_docker.py

@@ -139,16 +139,19 @@ def test_load_task_invalid_task(run_load_task):
     task = {}
     assert run_load_task(task)[0] == 1
 
-    task["tags"] = {"worker-implementation": "generic-worker"}
+    task["payload"] = {}
     assert run_load_task(task)[0] == 1
 
-    task["tags"]["worker-implementation"] = "docker-worker"
-    task["payload"] = {"command": []}
+    task["payload"] = {"command": [], "image": {"type": "task-image"}}
     assert run_load_task(task)[0] == 1
 
     task["payload"]["command"] = ["echo", "foo"]
     assert run_load_task(task)[0] == 1
 
+    task["payload"]["image"]["type"] = "foobar"
+    task["payload"]["command"] = ["run-task", "--", "bash", "-c", "echo foo"]
+    assert run_load_task(task)[0] == 1
+
 
 def test_load_task(run_load_task):
     image_task_id = "def"
@@ -161,9 +164,8 @@ def test_load_task(run_load_task):
                 "--",
                 "echo foo",
             ],
-            "image": {"taskId": image_task_id},
+            "image": {"taskId": image_task_id, "type": "task-image"},
         },
-        "tags": {"worker-implementation": "docker-worker"},
     }
     ret, mocks = run_load_task(task)
     assert ret == 0
@@ -199,7 +201,33 @@ def test_load_task(run_load_task):
             assert exp == actual[i]
 
 
-def test_load_task_env_and_remove(run_load_task):
+def test_load_task_env_init_and_remove(mocker, run_load_task):
+    # Mock NamedTemporaryFile to capture what's written to it
+    mock_envfile = mocker.MagicMock()
+    mock_envfile.name = "/tmp/test_envfile"
+    mock_envfile.fileno.return_value = 123  # Mock file descriptor
+
+    written_env_content = []
+    mock_envfile.write = lambda content: written_env_content.append(content)
+    mock_envfile.close = mocker.MagicMock()
+
+    mock_initfile = mocker.MagicMock()
+    mock_initfile.name = "/tmp/test_initfile"
+    mock_initfile.fileno.return_value = 456  # Mock file descriptor
+    written_init_content = []
+    mock_initfile.write = lambda content: written_init_content.append(content)
+    mock_initfile.close = mocker.MagicMock()
+
+    # Return different mocks for each call to NamedTemporaryFile
+    mock_tempfile = mocker.patch.object(docker.tempfile, "NamedTemporaryFile")
+    mock_tempfile.side_effect = [mock_envfile, mock_initfile]
+
+    # Mock os.remove to prevent file deletion errors
+    mock_os_remove = mocker.patch.object(docker.os, "remove")
+
+    # Mock os.fchmod
+    mocker.patch.object(docker.os, "fchmod")
+
     image_task_id = "def"
     task = {
         "payload": {
@@ -210,15 +238,100 @@ def test_load_task_env_and_remove(run_load_task):
                 "--",
                 "echo foo",
             ],
-            "env": {"FOO": "BAR", "BAZ": 1},
-            "image": {"taskId": image_task_id},
+            "env": {"FOO": "BAR", "BAZ": "1", "TASKCLUSTER_CACHES": "path"},
+            "image": {"taskId": image_task_id, "type": "task-image"},
         },
-        "tags": {"worker-implementation": "docker-worker"},
     }
     ret, mocks = run_load_task(task, remove=True)
     assert ret == 0
 
+    # NamedTemporaryFile was called twice (once for env, once for init)
+    assert mock_tempfile.call_count == 2
+
+    # Verify the environment content written to the file
+    assert len(written_env_content) == 1
+    env_lines = written_env_content[0].split("\n")
+
+    # Verify written env is expected
+    assert "TASKCLUSTER_CACHES=path" not in env_lines
+    assert "FOO=BAR" in env_lines
+    assert "BAZ=1" in env_lines
+
+    # Check that the default env vars were included
+    assert any("RUN_ID=0" in line for line in env_lines)
+    assert any("TASK_ID=abc" in line for line in env_lines)
+    assert any("TASK_GROUP_ID=" in line for line in env_lines)
+    assert any("TASKCLUSTER_ROOT_URL=" in line for line in env_lines)
+
+    # Both files were closed and removed
+    mock_envfile.close.assert_called_once()
+    mock_initfile.close.assert_called_once()
+    assert mock_os_remove.call_count == 2
+    assert mock_os_remove.call_args_list[0] == mocker.call("/tmp/test_envfile")
+    assert mock_os_remove.call_args_list[1] == mocker.call("/tmp/test_initfile")
+
+    # Verify subprocess was called with the correct env file and init file
     mocks["subprocess_run"].assert_called_once()
     actual = mocks["subprocess_run"].call_args[0][0]
-    assert re.match(r"--env-file=/tmp/tmp.*", actual[4])
+    assert actual[3] == "/tmp/test_initfile:/builds/worker/.bashrc"
+    assert actual[4] == "--env-file=/tmp/test_envfile"
     assert actual[5] == "--rm"
+
+
+@pytest.mark.parametrize(
+    "image",
+    [
+        pytest.param({"type": "task-image", "taskId": "xyz"}, id="task_image"),
+        pytest.param(
+            {"type": "indexed-image", "namespace": "project.some-namespace.latest"},
+            id="indexed_image",
+        ),
+    ],
+)
+def test_load_task_with_different_image_types(
+    mocker,
+    run_load_task,
+    image,
+):
+    task_id = "abc"
+    image_task_id = "xyz"
+    task = {
+        "payload": {
+            "command": [
+                "/usr/bin/run-task",
+                "--task-cwd=/builds/worker",
+                "--",
+                "echo",
+                "test",
+            ],
+            "image": image,
+        },
+    }
+
+    mocker.patch.object(docker, "find_task_id", return_value=image_task_id)
+
+    ret, mocks = run_load_task(task)
+    assert ret == 0
+
+    mocks["get_task_definition"].assert_called_once_with(task_id)
+    mocks["load_image_by_task_id"].assert_called_once_with(image_task_id)
+
+
+def test_load_task_with_unsupported_image_type(capsys, run_load_task):
+    task = {
+        "payload": {
+            "command": [
+                "/usr/bin/run-task",
+                "--task-cwd=/builds/worker",
+                "--",
+                "echo foo",
+            ],
+            "image": {"type": "unsupported-type", "path": "/some/path"},
+        },
+    }
+
+    ret, mocks = run_load_task(task)
+    assert ret == 1
+
+    out, _ = capsys.readouterr()
+    assert "Tasks with unsupported-type images are not supported!" in out