From 5089dfc54376f771cf5703f416babc2cc073400e Mon Sep 17 00:00:00 2001 From: Andrew Halberstadt Date: Fri, 5 Sep 2025 13:33:44 -0400 Subject: [PATCH 1/3] fix(load-task): use feature detection rather than worker-implementation Previously we were checking whether the worker-implementation was docker-worker by looking at the task tags. But many projects are missing this tag, and especially Decision tasks don't tend to have it. Instead, check if the `payload.image` key is defined. --- src/taskgraph/docker.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/src/taskgraph/docker.py b/src/taskgraph/docker.py index e1f67fba2..1e41354a5 100644 --- a/src/taskgraph/docker.py +++ b/src/taskgraph/docker.py @@ -269,10 +269,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 +306,7 @@ def load_task(task_id, remove=True, user=None): else: task_cwd = "$TASK_WORKDIR" - image_task_id = task_def["payload"]["image"]["taskId"] + image_task_id = image["taskId"] image_tag = load_image_by_task_id(image_task_id) # Set some env vars the worker would normally set. From 56ea1a8f3188dab8fdbf10b49bc951e01229a185 Mon Sep 17 00:00:00 2001 From: Andrew Halberstadt Date: Fri, 5 Sep 2025 14:44:17 -0400 Subject: [PATCH 2/3] fix(load-task): support tasks with indexed-images --- src/taskgraph/docker.py | 10 +++++- test/test_docker.py | 75 +++++++++++++++++++++++++++++++++++++---- 2 files changed, 78 insertions(+), 7 deletions(-) diff --git a/src/taskgraph/docker.py b/src/taskgraph/docker.py index 1e41354a5..12d4d8dcf 100644 --- a/src/taskgraph/docker.py +++ b/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, @@ -306,7 +307,14 @@ def load_task(task_id, remove=True, user=None): else: task_cwd = "$TASK_WORKDIR" - image_task_id = 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. diff --git a/test/test_docker.py b/test/test_docker.py index 4f8d79d4d..a4abf02da 100644 --- a/test/test_docker.py +++ b/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,7 +164,7 @@ 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"}, } @@ -211,9 +214,8 @@ def test_load_task_env_and_remove(run_load_task): "echo foo", ], "env": {"FOO": "BAR", "BAZ": 1}, - "image": {"taskId": image_task_id}, + "image": {"taskId": image_task_id, "type": "task-image"}, }, - "tags": {"worker-implementation": "docker-worker"}, } ret, mocks = run_load_task(task, remove=True) assert ret == 0 @@ -222,3 +224,64 @@ def test_load_task_env_and_remove(run_load_task): actual = mocks["subprocess_run"].call_args[0][0] assert re.match(r"--env-file=/tmp/tmp.*", actual[4]) 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, + }, + "tags": {"worker-implementation": "docker-worker"}, + } + + 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"}, + }, + "tags": {"worker-implementation": "docker-worker"}, + } + + ret, mocks = run_load_task(task) + assert ret == 1 + + out, _ = capsys.readouterr() + assert "Tasks with unsupported-type images are not supported!" in out From ffb3ed98951cc75d1c24516e80b288bfc888ffa7 Mon Sep 17 00:00:00 2001 From: Andrew Halberstadt Date: Fri, 5 Sep 2025 14:48:34 -0400 Subject: [PATCH 3/3] fix(load-task): remove TASKCLUSTER_CACHES env run-task expects the worker to mount a volume at each path passed in here. Let's avoid needing the user of `taskgraph load-task` needing to do the same. --- src/taskgraph/docker.py | 5 ++++ test/test_docker.py | 62 +++++++++++++++++++++++++++++++++++++---- 2 files changed, 61 insertions(+), 6 deletions(-) diff --git a/src/taskgraph/docker.py b/src/taskgraph/docker.py index 12d4d8dcf..7d237ab6f 100644 --- a/src/taskgraph/docker.py +++ b/src/taskgraph/docker.py @@ -327,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: diff --git a/test/test_docker.py b/test/test_docker.py index a4abf02da..a1bf4628c 100644 --- a/test/test_docker.py +++ b/test/test_docker.py @@ -166,7 +166,6 @@ def test_load_task(run_load_task): ], "image": {"taskId": image_task_id, "type": "task-image"}, }, - "tags": {"worker-implementation": "docker-worker"}, } ret, mocks = run_load_task(task) assert ret == 0 @@ -202,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": { @@ -213,16 +238,43 @@ def test_load_task_env_and_remove(run_load_task): "--", "echo foo", ], - "env": {"FOO": "BAR", "BAZ": 1}, + "env": {"FOO": "BAR", "BAZ": "1", "TASKCLUSTER_CACHES": "path"}, "image": {"taskId": image_task_id, "type": "task-image"}, }, } 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" @@ -254,7 +306,6 @@ def test_load_task_with_different_image_types( ], "image": image, }, - "tags": {"worker-implementation": "docker-worker"}, } mocker.patch.object(docker, "find_task_id", return_value=image_task_id) @@ -277,7 +328,6 @@ def test_load_task_with_unsupported_image_type(capsys, run_load_task): ], "image": {"type": "unsupported-type", "path": "/some/path"}, }, - "tags": {"worker-implementation": "docker-worker"}, } ret, mocks = run_load_task(task)