feat(run-task): automatically interpolate $TASK_WORKDIR in environment

Author: ahal

src/taskgraph/run-task/run-task

@@ -1091,10 +1091,10 @@ def _display_python_version():
 
 
 def main(args):
-    os.environ["TASK_WORKDIR"] = os.getcwd()
+    task_workdir = os.environ["TASK_WORKDIR"] = os.getcwd()
     print_line(
         b"setup",
-        b"run-task started in %s\n" % os.environ["TASK_WORKDIR"].encode("utf-8"),
+        b"run-task started in %s\n" % task_workdir.encode("utf-8"),
     )
     print_line(
         b"setup",
@@ -1313,20 +1313,27 @@ def main(args):
     for repo in repositories:
         vcs_checkout_from_args(repo)
 
-    resource_process = None
-
-    try:
-        for k in ["MOZ_FETCHES_DIR", "UPLOAD_DIR"] + [
-            "{}_PATH".format(repository["project"].upper())
-            for repository in repositories
-        ]:
-            if k in os.environ:
-                os.environ[k] = os.path.abspath(os.environ[k])
+    # Interpolate environment variables with defined substitution patterns. For
+    # example, the variable `CACHE_DIR={task_workdir}/.cache` will be
+    # interpolated with the task's working directory.
+    env_subs = {
+        "task_workdir": task_workdir,
+    }
+    for k, v in os.environ.items():
+        for subname, subval in env_subs.items():
+            # We check for an exact match rather than using a format string to
+            # avoid accidentally trying to interpolate environment variables
+            # that happen to contain brackets for some other reason.
+            pattern = f"{{{subname}}}"
+            if pattern in v:
+                os.environ[k] = v.replace(pattern, subval)
                 print_line(
                     b"setup",
                     b"%s is %s\n" % (k.encode("utf-8"), os.environ[k].encode("utf-8")),
                 )
 
+    resource_process = None
+    try:
         if "MOZ_FETCHES" in os.environ:
             fetch_artifacts()
 

src/taskgraph/transforms/run/__init__.py

@@ -367,7 +367,7 @@ def cmp_artifacts(a):
             "task-reference": json.dumps(task_fetches, sort_keys=True)
         }
 
-        env.setdefault("MOZ_FETCHES_DIR", "fetches")
+        env.setdefault("MOZ_FETCHES_DIR", "{task_workdir}/fetches")
 
         yield task
 

src/taskgraph/transforms/run/common.py

@@ -10,6 +10,7 @@
 import hashlib
 import json
 
+from taskgraph.util import path
 from taskgraph.util.taskcluster import get_artifact_prefix
 
 
@@ -98,16 +99,16 @@ def support_vcs_checkout(config, task, taskdesc, repo_configs, sparse=False):
     assert is_mac or is_win or is_linux
 
     if is_win:
-        checkoutdir = "./build"
+        checkoutdir = "build"
         hgstore = "y:/hg-shared"
     elif is_docker:
         checkoutdir = "{workdir}/checkouts".format(**task["run"])
         hgstore = f"{checkoutdir}/hg-store"
     else:
-        checkoutdir = "./checkouts"
+        checkoutdir = "checkouts"
         hgstore = f"{checkoutdir}/hg-shared"
 
-    vcsdir = checkoutdir + "/" + get_vcsdir_name(worker["os"])
+    vcsdir = f"{checkoutdir}/{get_vcsdir_name(worker['os'])}"
     cache_name = "checkouts"
 
     # Robust checkout does not clean up subrepositories, so ensure  that tasks
@@ -138,7 +139,8 @@ def support_vcs_checkout(config, task, taskdesc, repo_configs, sparse=False):
             "REPOSITORIES": json.dumps(
                 {repo.prefix: repo.name for repo in repo_configs.values()}
             ),
-            "VCS_PATH": vcsdir,
+            # If vcsdir is already absolute this will return it unmodified.
+            "VCS_PATH": path.join("{task_workdir}", vcsdir),
         }
     )
     for repo_config in repo_configs.values():
@@ -162,3 +164,5 @@ def support_vcs_checkout(config, task, taskdesc, repo_configs, sparse=False):
     # only some worker platforms have taskcluster-proxy enabled
     if task["worker"]["implementation"] in ("docker-worker",):
         taskdesc["worker"]["taskcluster-proxy"] = True
+
+    return vcsdir

src/taskgraph/transforms/run/run_task.py

@@ -70,15 +70,14 @@ def common_setup(config, task, taskdesc, command):
                 for (repo, config) in run["checkout"].items()
             }
 
-        support_vcs_checkout(
+        vcs_path = support_vcs_checkout(
             config,
             task,
             taskdesc,
             repo_configs=repo_configs,
             sparse=bool(run["sparse-profile"]),
         )
 
-        vcs_path = taskdesc["worker"]["env"]["VCS_PATH"]
         for repo_config in repo_configs.values():
             checkout_path = path.join(vcs_path, repo_config.path)
             command.append(f"--{repo_config.prefix}-checkout={checkout_path}")

test/test_scripts_run_task.py

@@ -438,7 +438,7 @@ def test_display_python_version_should_output_python_versions(run_task_mod, caps
 
 
 @pytest.fixture
-def run_main(tmp_path, capsys, mocker, mock_stdin, run_task_mod):
+def run_main(tmp_path, mocker, mock_stdin, run_task_mod):
     base_args = [
         f"--task-cwd={str(tmp_path)}",
     ]
@@ -463,21 +463,26 @@ def inner(extra_args=None, env=None):
         args.extend(base_command_args)
 
         result = run_task_mod.main(args)
-        out, err = capsys.readouterr()
-        return result, out, err, env
+        return result, env
 
     return inner
 
 
 def test_main_interpolate_environment(run_main):
-    result, out, err, env = run_main(
-        env={"MOZ_FETCHES_DIR": "file", "UPLOAD_DIR": "file", "FOO": "file"}
+    result, env = run_main(
+        env={
+            "MOZ_FETCHES_DIR": "{task_workdir}/file",
+            "UPLOAD_DIR": "$TASK_WORKDIR/file",
+            "FOO": "{foo}/file",
+            "BAR": "file",
+        }
     )
     assert result == 0
 
     assert env == {
         "MOZ_FETCHES_DIR": "/builds/worker/file",
-        "UPLOAD_DIR": "/builds/worker/file",
-        "FOO": "file",
+        "UPLOAD_DIR": "$TASK_WORKDIR/file",
+        "FOO": "{foo}/file",
+        "BAR": "file",
         "TASK_WORKDIR": "/builds/worker",
     }

test/test_transforms_run_run_task.py

@@ -99,7 +99,7 @@ def assert_generic_worker(task):
         "worker": {
             "command": [
                 "C:/mozilla-build/python3/python3.exe run-task "
-                '--ci-checkout=./build/src/ -- bash -cx "echo hello '
+                '--ci-checkout=build/src/ -- bash -cx "echo hello '
                 'world"'
             ],
             "env": {
@@ -111,11 +111,11 @@ def assert_generic_worker(task):
                 "HG_STORE_PATH": "y:/hg-shared",
                 "MOZ_SCM_LEVEL": "1",
                 "REPOSITORIES": '{"ci": "Taskgraph"}',
-                "VCS_PATH": "./build/src",
+                "VCS_PATH": "{task_workdir}/build/src",
             },
             "implementation": "generic-worker",
             "mounts": [
-                {"cache-name": "checkouts", "directory": "./build"},
+                {"cache-name": "checkouts", "directory": "build"},
                 {
                     "content": {
                         "url": "https://tc-tests.localhost/api/queue/v1/task/<TASK_ID>/artifacts/public/run-task"  # noqa
@@ -159,7 +159,7 @@ def assert_run_task_command_generic_worker(task):
         [
             "/foo/bar/python3",
             "run-task",
-            "--ci-checkout=./checkouts/vcs/",
+            "--ci-checkout=checkouts/vcs/",
             "--",
             "bash",
             "-cx",