Support a --venv mode similar to --unzip mode. (#1153)

The new --venv execution mode builds a PEX file that includes pex.tools
and extracts itself into a venv under PEX_ROOT upon 1st execution or any
execution that might select a diffrent interpreter than the default.

In order to speed up the local build and execute case, --seed mode is
added to seed the PEX_ROOT caches that will be used at runtime. This is
important for --venv mode since venv seeding depends on the selected
interpreter and one is already selected during the PEX file build
process.

Fixes #962
Fixes #1097
Fixes #1115

pex/bin/pex.py

@@ -11,12 +11,12 @@
 import os
 import sys
 import tempfile
-import zipfile
 from argparse import Action, ArgumentDefaultsHelpFormatter, ArgumentParser, ArgumentTypeError
+from shlex import shlex
 from textwrap import TextWrapper
 
 from pex import pex_warnings
-from pex.common import die, safe_delete, safe_mkdtemp
+from pex.common import atomic_directory, die, open_zip, safe_mkdtemp
 from pex.inherit_path import InheritPath
 from pex.interpreter import PythonInterpreter
 from pex.interpreter_constraints import (
@@ -27,18 +27,19 @@
 from pex.network_configuration import NetworkConfiguration
 from pex.orderedset import OrderedSet
 from pex.pex import PEX
-from pex.pex_bootstrapper import iter_compatible_interpreters
+from pex.pex_bootstrapper import ensure_venv, iter_compatible_interpreters
 from pex.pex_builder import PEXBuilder
 from pex.pip import ResolverVersion
 from pex.platforms import Platform
 from pex.resolver import Unsatisfiable, parsed_platform, resolve_multi
 from pex.tracer import TRACER
 from pex.typing import TYPE_CHECKING
 from pex.variables import ENV, Variables
+from pex.venv_bin_path import BinPath
 from pex.version import __version__
 
 if TYPE_CHECKING:
-    from typing import List
+    from typing import List, Iterable
     from argparse import Namespace
 
 
@@ -88,6 +89,17 @@ def __call__(self, parser, namespace, value, option_str=None):
         setattr(namespace, self.dest, option_str == "--transitive")
 
 
+class HandleVenvAction(Action):
+    def __init__(self, *args, **kwargs):
+        kwargs["nargs"] = "?"
+        kwargs["choices"] = (BinPath.PREPEND.value, BinPath.APPEND.value)
+        super(HandleVenvAction, self).__init__(*args, **kwargs)
+
+    def __call__(self, parser, namespace, value, option_str=None):
+        bin_path = BinPath.FALSE if value is None else BinPath.for_value(value)
+        setattr(namespace, self.dest, bin_path)
+
+
 class PrintVariableHelpAction(Action):
     def __call__(self, parser, namespace, values, option_str=None):
         for variable_name, variable_type, variable_help in Variables.iter_help():
@@ -326,7 +338,8 @@ def configure_clp_pex_options(parser):
         "complete pex file, including dependencies, to be unzipped.",
     )
 
-    group.add_argument(
+    runtime_mode = group.add_mutually_exclusive_group()
+    runtime_mode.add_argument(
         "--unzip",
         "--no-unzip",
         dest="unzip",
@@ -336,6 +349,18 @@ def configure_clp_pex_options(parser):
         "be run multiple times under a stable runtime PEX_ROOT the unzipping will only be "
         "performed once and subsequent runs will enjoy lower startup latency.",
     )
+    runtime_mode.add_argument(
+        "--venv",
+        dest="venv",
+        metavar="{prepend,append}",
+        default=False,
+        action=HandleVenvAction,
+        help="Convert the pex file to a venv before executing it. If 'prepend' or 'append' is "
+        "specified, then all scripts and console scripts provided by distributions in the pex file "
+        "will be added to the PATH in the corresponding position. If the the pex file will be run "
+        "multiple times under a stable runtime PEX_ROOT, the venv creation will only be done once "
+        "and subsequent runs will enjoy lower startup latency.",
+    )
 
     group.add_argument(
         "--always-write-cache",
@@ -712,6 +737,16 @@ def configure_clp():
         help="Specify the temporary directory Pex and its subprocesses should use.",
     )
 
+    parser.add_argument(
+        "--seed",
+        "--no-seed",
+        dest="seed",
+        action=HandleBoolAction,
+        default=False,
+        help="Seed local Pex caches for the generated PEX and print out the command line to run "
+        "directly from the seed with.",
+    )
+
     parser.add_argument(
         "--help-variables",
         action=PrintVariableHelpAction,
@@ -824,7 +859,7 @@ def to_python_interpreter(full_path_or_basename):
         path=safe_mkdtemp(),
         interpreter=interpreter,
         preamble=preamble,
-        include_tools=options.include_tools,
+        include_tools=options.include_tools or options.venv,
     )
 
     if options.resources_directory:
@@ -844,6 +879,8 @@ def to_python_interpreter(full_path_or_basename):
     pex_info = pex_builder.info
     pex_info.zip_safe = options.zip_safe
     pex_info.unzip = options.unzip
+    pex_info.venv = bool(options.venv)
+    pex_info.venv_bin_path = options.venv
     pex_info.pex_path = options.pex_path
     pex_info.always_write_cache = options.always_write_cache
     pex_info.ignore_errors = options.ignore_errors
@@ -1005,14 +1042,14 @@ def warn_ignore_pex_root(set_via):
 
         if options.pex_name is not None:
             log("Saving PEX file to %s" % options.pex_name, V=options.verbosity)
-            tmp_name = options.pex_name + "~"
-            safe_delete(tmp_name)
             pex_builder.build(
-                tmp_name,
+                options.pex_name,
                 bytecode_compile=options.compile,
                 deterministic_timestamp=not options.use_system_time,
             )
-            os.rename(tmp_name, options.pex_name)
+            if options.seed:
+                execute_cached_args = seed_cache(options, pex)
+                print(" ".join(execute_cached_args))
         else:
             if not _compatible_with_current_platform(interpreter, options.platforms):
                 log("WARNING: attempting to run PEX with incompatible platforms!", V=1)
@@ -1030,5 +1067,36 @@ def warn_ignore_pex_root(set_via):
             sys.exit(pex.run(args=list(cmdline), env=patched_env))
 
 
+def seed_cache(
+    options,  # type: Namespace
+    pex,  # type: PEX
+):
+    # type: (...) -> Iterable[str]
+    pex_path = pex.path()
+    with TRACER.timed("Seeding local caches for {}".format(pex_path)):
+        if options.unzip:
+            unzip_dir = pex.pex_info().unzip_dir
+            if unzip_dir is None:
+                raise AssertionError(
+                    "Expected PEX-INFO for {} to have the components of an unzip directory".format(
+                        pex_path
+                    )
+                )
+            with atomic_directory(unzip_dir, exclusive=True) as chroot:
+                if chroot:
+                    with TRACER.timed("Extracting {}".format(pex_path)):
+                        with open_zip(options.pex_name) as pex_zip:
+                            pex_zip.extractall(chroot)
+            return [pex.interpreter.binary, unzip_dir]
+        elif options.venv:
+            with TRACER.timed("Creating venv from {}".format(pex_path)):
+                venv_pex = ensure_venv(pex)
+                return [venv_pex]
+        else:
+            with TRACER.timed("Extracting code and distributions for {}".format(pex_path)):
+                pex.activate()
+            return [os.path.abspath(options.pex_name)]
+
+
 if __name__ == "__main__":
     main()

pex/common.py

@@ -24,7 +24,7 @@
 from pex.typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Any, DefaultDict, Iterable, Iterator, NoReturn, Optional, Set
+    from typing import Any, DefaultDict, Iterable, Iterator, NoReturn, Optional, Set, Sized
 
 # We use the start of MS-DOS time, which is what zipfiles use (see section 4.4.6 of
 # https://pkware.cachefly.net/webdocs/casestudies/APPNOTE.TXT).
@@ -60,6 +60,22 @@ def die(msg, exit_code=1):
     sys.exit(exit_code)
 
 
+def pluralize(
+    subject,  # type: Sized
+    noun,  # type: str
+):
+    # type: (...) -> str
+    if noun == "":
+        return ""
+    count = len(subject)
+    if count == 1:
+        return noun
+    if noun[-1] in ("s", "x", "z") or noun[-2:] in ("sh", "ch"):
+        return noun + "es"
+    else:
+        return noun + "s"
+
+
 def safe_copy(source, dest, overwrite=False):
     # type: (str, str, bool) -> None
     def do_copy():

pex/environment.py

@@ -25,7 +25,7 @@
 from pex.util import CacheHelper, DistributionHelper
 
 if TYPE_CHECKING:
-    from typing import Container, Optional
+    from typing import Container, Iterator, Optional, Tuple, Iterable
 
 
 def _import_pkg_resources():
@@ -114,7 +114,7 @@ def explode_code(
         dest_dir,  # type: str
         exclude=(),  # type: Container[str]
     ):
-        # type: (...) -> None
+        # type: (...) -> Iterable[Tuple[str, str]]
         with TRACER.timed("Unzipping {}".format(pex_file)):
             with open_zip(pex_file) as pex_zip:
                 pex_files = (
@@ -125,6 +125,13 @@ def explode_code(
                     and name not in exclude
                 )
                 pex_zip.extractall(dest_dir, pex_files)
+                return [
+                    (
+                        "{pex_file}:{zip_path}".format(pex_file=pex_file, zip_path=f),
+                        os.path.join(dest_dir, f),
+                    )
+                    for f in pex_files
+                ]
 
     @classmethod
     def _force_local(cls, pex_file, pex_info):

pex/inherit_path.py

@@ -6,7 +6,7 @@
 from pex.typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Tuple, Union
+    from typing import Union
 
 
 class InheritPath(object):

pex/interpreter.py

@@ -607,7 +607,11 @@ def create_interpreter(stdout, check_binary=False):
         # would otherwise be unstable.
         #
         # See cls._REGEXEN for a related affordance.
-        path_id = binary.replace(os.sep, ".").lstrip(".")
+        #
+        # N.B.: The path for --venv mode interpreters can be quite long; so we just used a fixed
+        # length hash of the interpreter binary path to ensure uniqueness and not run afoul of file
+        # name length limits.
+        path_id = hashlib.sha1(binary.encode("utf-8")).hexdigest()
 
         cache_dir = os.path.join(os_cache_dir, interpreter_hash, path_id)
         cache_file = os.path.join(cache_dir, cls.INTERP_INFO_FILE)

pex/pex_bootstrapper.py

@@ -7,7 +7,7 @@
 import sys
 
 from pex import pex_warnings
-from pex.common import die
+from pex.common import atomic_directory, die
 from pex.inherit_path import InheritPath
 from pex.interpreter import PythonInterpreter
 from pex.interpreter_constraints import UnsatisfiableInterpreterConstraintsError
@@ -29,6 +29,8 @@
         Callable,
     )
 
+    from pex.pex import PEX
+
     InterpreterIdentificationError = Tuple[str, str]
     InterpreterOrError = Union[PythonInterpreter, InterpreterIdentificationError]
     PathFilter = Callable[[str], bool]
@@ -361,15 +363,51 @@ def _bootstrap(entry_point):
     return pex_info
 
 
+def ensure_venv(pex):
+    # type: (PEX) -> str
+    pex_info = pex.pex_info()
+    venv_dir = pex_info.venv_dir
+    if venv_dir is None:
+        raise AssertionError(
+            "Expected PEX-INFO for {} to have the components of a venv directory".format(pex.path())
+        )
+    with atomic_directory(venv_dir, exclusive=True) as venv:
+        if venv:
+            from .tools.commands.venv import populate_venv_with_pex
+            from .tools.commands.virtualenv import Virtualenv
+
+            virtualenv = Virtualenv.create(venv_dir=venv, interpreter=pex.interpreter)
+            populate_venv_with_pex(
+                virtualenv,
+                pex,
+                bin_path=pex_info.venv_bin_path,
+                python=os.path.join(venv_dir, "bin", os.path.basename(pex.interpreter.binary)),
+                collisions_ok=True,
+            )
+    return os.path.join(venv_dir, "pex")
+
+
 # NB: This helper is used by the PEX bootstrap __main__.py code.
 def bootstrap_pex(entry_point):
     # type: (str) -> None
     pex_info = _bootstrap(entry_point)
-    maybe_reexec_pex(pex_info.interpreter_constraints)
 
-    from . import pex
+    if not ENV.PEX_TOOLS and pex_info.venv:
+        try:
+            target = find_compatible_interpreter(
+                interpreter_constraints=pex_info.interpreter_constraints,
+            )
+        except UnsatisfiableInterpreterConstraintsError as e:
+            die(str(e))
+        from . import pex
+
+        venv_pex = ensure_venv(pex.PEX(entry_point, interpreter=target))
+        os.execv(venv_pex, [venv_pex] + sys.argv[1:])
+    else:
+        maybe_reexec_pex(pex_info.interpreter_constraints)
+        from . import pex
 
-    pex.PEX(entry_point).execute()
+        pex.PEX(entry_point).execute()
 
 
 # NB: This helper is used by third party libs - namely https://github.com/wickman/lambdex.