Repository: pex-tool/pex

Issue:
Control of the TMPDIR used by Pex and its subprocesses (notably Pip).
The Pex CLI can consume alot of disk space in the process of downloading and installing wheels in PEX files. Although Pex does most of its work in the --pex-root, Pip subprocesses still do some of their heavy lifting in TMPDIR. This can be controlled by exporting TMPDIR but its not necessarily obvious to do this. It probably makes sense for Pex to add an option to control this or else default to using a TMPDIR in PEX_ROOT.


Discussion:
Motivating issue: https://github.com/pantsbuild/pants/issues/10936

Resolving patch:
diff --git a/pex/bin/pex.py b/pex/bin/pex.py
--- a/pex/bin/pex.py
+++ b/pex/bin/pex.py
@@ -10,6 +10,7 @@
 
 import os
 import sys
+import tempfile
 from optparse import OptionGroup, OptionParser, OptionValueError
 from textwrap import TextWrapper
 
@@ -704,6 +705,14 @@ def configure_clp():
         "[Default: {}]".format(ENV.PEX_ROOT),
     )
 
+    parser.add_option(
+        "--tmpdir",
+        dest="tmpdir",
+        default=tempfile.gettempdir(),
+        help="Specify the temporary directory Pex and its subprocesses should use. "
+        "[Default: %default]",
+    )
+
     parser.add_option(
         "--help-variables",
         action="callback",
@@ -946,6 +955,15 @@ def main(args=None):
 
     options, reqs = parser.parse_args(args=args)
 
+    # Ensure the TMPDIR is an absolute path (So subprocesses that change CWD can find it) and
+    # that it exists.
+    tmpdir = os.path.realpath(options.tmpdir)
+    if not os.path.exists(tmpdir):
+        die("The specified --tmpdir does not exist: {}".format(tmpdir))
+    if not os.path.isdir(tmpdir):
+        die("The specified --tmpdir is not a directory: {}".format(tmpdir))
+    tempfile.tempdir = os.environ["TMPDIR"] = tmpdir
+
     if options.cache_dir:
         pex_warnings.warn("The --cache-dir option is deprecated, use --pex-root instead.")
         if options.pex_root and options.cache_dir != options.pex_root:
@@ -976,7 +994,9 @@ def warn_ignore_pex_root(set_via):
     if options.python and options.interpreter_constraint:
         die('The "--python" and "--interpreter-constraint" options cannot be used together.')
 
-    with ENV.patch(PEX_VERBOSE=str(options.verbosity), PEX_ROOT=pex_root) as patched_env:
+    with ENV.patch(
+        PEX_VERBOSE=str(options.verbosity), PEX_ROOT=pex_root, TMPDIR=tmpdir
+    ) as patched_env:
         with TRACER.timed("Building pex"):
             pex_builder = build_pex(reqs, options, cache=ENV.PEX_ROOT)
 


Test patch:
diff --git a/tests/test_integration.py b/tests/test_integration.py
--- a/tests/test_integration.py
+++ b/tests/test_integration.py
@@ -20,7 +20,15 @@
 import pytest
 
 from pex import pex_builder
-from pex.common import safe_copy, safe_mkdir, safe_open, safe_rmtree, safe_sleep, temporary_dir
+from pex.common import (
+    safe_copy,
+    safe_mkdir,
+    safe_open,
+    safe_rmtree,
+    safe_sleep,
+    temporary_dir,
+    touch,
+)
 from pex.compatibility import WINDOWS, nested, to_bytes
 from pex.executor import Executor
 from pex.interpreter import PythonInterpreter
@@ -2342,3 +2350,53 @@ def create_platform_pex(args):
         )
         assert 0 == returncode
         assert int(output.strip()) >= multiprocessing.cpu_count()
+
+
+@pytest.fixture
+def tmp_workdir():
+    # type: () -> Iterator[str]
+    with temporary_dir() as tmpdir:
+        os.chdir(tmpdir)
+        yield os.path.realpath(tmpdir)
+
+
+def test_tmpdir_absolute(tmp_workdir):
+    # type: (str) -> None
+    result = run_pex_command(
+        args=[
+            "--tmpdir",
+            ".",
+            "--",
+            "-c",
+            dedent(
+                """\
+                import os
+                import tempfile
+                
+                print(os.environ["TMPDIR"])
+                print(tempfile.gettempdir())
+                """
+            ),
+        ]
+    )
+    result.assert_success()
+    assert [tmp_workdir, tmp_workdir] == result.output.strip().splitlines()
+
+
+def test_tmpdir_dne(tmp_workdir):
+    # type: (str) -> None
+    tmpdir_dne = os.path.join(tmp_workdir, ".tmp")
+    result = run_pex_command(args=["--tmpdir", ".tmp", "--", "-c", ""])
+    result.assert_failure()
+    assert tmpdir_dne in result.error
+    assert "does not exist" in result.error
+
+
+def test_tmpdir_file(tmp_workdir):
+    # type: (str) -> None
+    tmpdir_file = os.path.join(tmp_workdir, ".tmp")
+    touch(tmpdir_file)
+    result = run_pex_command(args=["--tmpdir", ".tmp", "--", "-c", ""])
+    result.assert_failure()
+    assert tmpdir_file in result.error
+    assert "is not a directory" in result.error