Support running pexes from a working directory. (#12347)

Previously we only supported running them from the sandbox root.
Now we can run them from any subdir.

[ci skip-rust]

[ci skip-build-wheels]

Author: benjyw

src/python/pants/backend/python/util_rules/pex.py

@@ -6,6 +6,7 @@
 import dataclasses
 import json
 import logging
+import os
 import shlex
 from dataclasses import dataclass
 from pathlib import PurePath
@@ -523,11 +524,18 @@ def _create_venv_script(
             f"{name}={shlex.quote(value)}"
             for name, value in pex_environment.environment_dict(python_configured=True).items()
         )
+        # Ensure that the pex is executed from a path relative to the shim script, so
+        # that running from within a working directory works properly.
+        # NB: No method for determining the path of the current script is entirely foolproof,
+        #  but BASH_SOURCE works fine in our specific case (and has been available since bash-3.0,
+        #  released in 2004, so it seems safe to assume it's available).
+        local_pex_bin = '"${BASH_SOURCE%/*}/"' + shlex.quote(self.pex.name)
         target_venv_executable = shlex.quote(str(venv_executable))
         venv_dir = shlex.quote(str(self.venv_dir))
         execute_pex_args = " ".join(
-            shlex.quote(arg)
-            for arg in pex_environment.create_argv(self.pex.name, python=self.pex.python)
+            # Don't quote the BASH_SOURCE dereference in local_pex_bin.
+            arg if arg == local_pex_bin else shlex.quote(arg)
+            for arg in pex_environment.create_argv(local_pex_bin, python=self.pex.python)[1:]
         )
 
         script = dedent(
@@ -693,6 +701,7 @@ class PexProcess:
     description: str = dataclasses.field(compare=False)
     level: LogLevel
     input_digest: Digest | None
+    working_directory: str | None
     extra_env: FrozenDict[str, str] | None
     output_files: tuple[str, ...] | None
     output_directories: tuple[str, ...] | None
@@ -708,6 +717,7 @@ def __init__(
         argv: Iterable[str] = (),
         level: LogLevel = LogLevel.INFO,
         input_digest: Digest | None = None,
+        working_directory: str | None = None,
         extra_env: Mapping[str, str] | None = None,
         output_files: Iterable[str] | None = None,
         output_directories: Iterable[str] | None = None,
@@ -720,6 +730,7 @@ def __init__(
         self.description = description
         self.level = level
         self.input_digest = input_digest
+        self.working_directory = working_directory
         self.extra_env = FrozenDict(extra_env) if extra_env else None
         self.output_files = tuple(output_files) if output_files else None
         self.output_directories = tuple(output_directories) if output_directories else None
@@ -731,7 +742,12 @@ def __init__(
 @rule
 async def setup_pex_process(request: PexProcess, pex_environment: SandboxPexEnvironment) -> Process:
     pex = request.pex
-    argv = pex_environment.create_argv(pex.name, *request.argv, python=pex.python)
+    pex_bin = (
+        os.path.relpath(pex.name, request.working_directory)
+        if request.working_directory
+        else pex.name
+    )
+    argv = pex_environment.create_argv(pex_bin, *request.argv, python=pex.python)
     env = {
         **pex_environment.environment_dict(python_configured=pex.python is not None),
         **(request.extra_env or {}),
@@ -746,6 +762,7 @@ async def setup_pex_process(request: PexProcess, pex_environment: SandboxPexEnvi
         description=request.description,
         level=request.level,
         input_digest=input_digest,
+        working_directory=request.working_directory,
         env=env,
         output_files=request.output_files,
         output_directories=request.output_directories,
@@ -764,6 +781,7 @@ class VenvPexProcess:
     description: str = dataclasses.field(compare=False)
     level: LogLevel
     input_digest: Digest | None
+    working_directory: str | None
     extra_env: FrozenDict[str, str] | None
     output_files: tuple[str, ...] | None
     output_directories: tuple[str, ...] | None
@@ -779,6 +797,7 @@ def __init__(
         argv: Iterable[str] = (),
         level: LogLevel = LogLevel.INFO,
         input_digest: Digest | None = None,
+        working_directory: str | None = None,
         extra_env: Mapping[str, str] | None = None,
         output_files: Iterable[str] | None = None,
         output_directories: Iterable[str] | None = None,
@@ -791,6 +810,7 @@ def __init__(
         self.description = description
         self.level = level
         self.input_digest = input_digest
+        self.working_directory = working_directory
         self.extra_env = FrozenDict(extra_env) if extra_env else None
         self.output_files = tuple(output_files) if output_files else None
         self.output_directories = tuple(output_directories) if output_directories else None
@@ -804,7 +824,12 @@ async def setup_venv_pex_process(
     request: VenvPexProcess, pex_environment: SandboxPexEnvironment
 ) -> Process:
     venv_pex = request.venv_pex
-    argv = (venv_pex.pex.argv0, *request.argv)
+    pex_bin = (
+        os.path.relpath(venv_pex.pex.argv0, request.working_directory)
+        if request.working_directory
+        else venv_pex.pex.argv0
+    )
+    argv = (pex_bin, *request.argv)
     input_digest = (
         await Get(Digest, MergeDigests((venv_pex.digest, request.input_digest)))
         if request.input_digest
@@ -815,6 +840,7 @@ async def setup_venv_pex_process(
         description=request.description,
         level=request.level,
         input_digest=input_digest,
+        working_directory=request.working_directory,
         env=request.extra_env,
         output_files=request.output_files,
         output_directories=request.output_directories,

src/python/pants/backend/python/util_rules/pex_test.py

@@ -5,6 +5,7 @@
 
 import json
 import os.path
+import re
 import textwrap
 import zipfile
 from dataclasses import dataclass
@@ -30,7 +31,7 @@
 )
 from pants.backend.python.util_rules.pex import rules as pex_rules
 from pants.backend.python.util_rules.pex_cli import PexPEX
-from pants.engine.fs import CreateDigest, Digest, FileContent
+from pants.engine.fs import CreateDigest, Digest, Directory, FileContent
 from pants.engine.process import Process, ProcessResult
 from pants.testutil.rule_runner import QueryRule, RuleRunner
 
@@ -73,7 +74,7 @@ def rule_runner() -> RuleRunner:
             QueryRule(PexResolveInfo, (Pex,)),
             QueryRule(PexResolveInfo, (VenvPex,)),
             QueryRule(PexPEX, ()),
-        ]
+        ],
     )
 
 
@@ -268,6 +269,71 @@ def test_pex_environment(rule_runner: RuleRunner, pex_type: type[Pex | VenvPex])
     assert b"ftp_proxy=dummyproxy" in result.stdout
 
 
+@pytest.mark.parametrize("pex_type", [Pex, VenvPex])
+def test_pex_working_directory(rule_runner: RuleRunner, pex_type: type[Pex | VenvPex]) -> None:
+    sources = rule_runner.request(
+        Digest,
+        [
+            CreateDigest(
+                (
+                    FileContent(
+                        path="main.py",
+                        content=textwrap.dedent(
+                            """
+                            import os
+                            cwd = os.getcwd()
+                            print(f"CWD: {cwd}")
+                            for path, dirs, _ in os.walk(cwd):
+                                for name in dirs:
+                                    print(f"DIR: {os.path.relpath(os.path.join(path, name), cwd)}")
+                            """
+                        ).encode(),
+                    ),
+                )
+            ),
+        ],
+    )
+
+    pex_output = create_pex_and_get_all_data(
+        rule_runner,
+        pex_type=pex_type,
+        main=EntryPoint("main"),
+        sources=sources,
+        interpreter_constraints=InterpreterConstraints(["CPython>=3.6"]),
+    )
+
+    pex = pex_output["pex"]
+    pex_process_type = PexProcess if isinstance(pex, Pex) else VenvPexProcess
+
+    dirpath = "foo/bar/baz"
+    runtime_files = rule_runner.request(Digest, [CreateDigest([Directory(path=dirpath)])])
+
+    dirpath_parts = os.path.split(dirpath)
+    for i in range(0, len(dirpath_parts)):
+        working_dir = os.path.join(*dirpath_parts[:i]) if i > 0 else None
+        expected_subdir = os.path.join(*dirpath_parts[i:]) if i < len(dirpath_parts) else None
+        process = rule_runner.request(
+            Process,
+            [
+                pex_process_type(
+                    pex,
+                    description="Run the pex and check its cwd",
+                    working_directory=working_dir,
+                    input_digest=runtime_files,
+                )
+            ],
+        )
+        result = rule_runner.request(ProcessResult, [process])
+        output_str = result.stdout.decode()
+        mo = re.search(r"CWD: (.*)\n", output_str)
+        assert mo is not None
+        reported_cwd = mo.group(1)
+        if working_dir:
+            assert reported_cwd.endswith(working_dir)
+        if expected_subdir:
+            assert f"DIR: {expected_subdir}" in output_str
+
+
 def test_resolves_dependencies(rule_runner: RuleRunner) -> None:
     requirements = PexRequirements(["six==1.12.0", "jsonschema==2.6.0", "requests==2.23.0"])
     pex_info = create_pex_and_get_pex_info(rule_runner, requirements=requirements)