Create --pex-path argument for pex cli and load pex path into pex-inf…

…o metadata (#417)

Problem:
Pex environments rely on the PEX_PATH environment variable to resolve modules from other pexes when composing them into a single pex. Python scripts that re-execute (like a server listening for fs changes, a Jupyter notebook) throw an ImportError upon re-exec due to environment scrubbing that removes this PEX_PATH information.

Solution:
Add a --pex-path argument to the pex client and plumb the data through to the pex-info metadata. Resolve pexes to compose into the current pex being built by reading from the new pex_path property of the pex-info metadata. This will ensure a self-contained environment that does not lose context on a python script re-exec.

Relates to: pantsbuild/pants#4682

pex/bin/pex.py

@@ -149,6 +149,13 @@ def configure_clp_pex_resolution(parser, builder):
       callback_args=(builder,),
       help='Whether to use pypi to resolve dependencies; Default: use pypi')
 
+  group.add_option(
+    '--pex-path',
+    dest='pex_path',
+    type=str,
+    default=None,
+    help='A colon separated list of other pex files to merge into the runtime environment.')
+
   group.add_option(
       '-f', '--find-links', '--repo',
       metavar='PATH/URL',
@@ -533,6 +540,7 @@ def build_pex(args, options, resolver_option_builder):
 
   pex_info = pex_builder.info
   pex_info.zip_safe = options.zip_safe
+  pex_info.pex_path = options.pex_path
   pex_info.always_write_cache = options.always_write_cache
   pex_info.ignore_errors = options.ignore_errors
   pex_info.inherit_path = options.inherit_path

pex/pex.py

@@ -71,11 +71,12 @@ def _activate(self):
       pex_info.update(self._pex_info_overrides)
       self._envs.append(PEXEnvironment(self._pex, pex_info))
 
-      # set up other environments as specified in PEX_PATH
-      for pex_path in filter(None, self._vars.PEX_PATH.split(os.pathsep)):
-        pex_info = PexInfo.from_pex(pex_path)
-        pex_info.update(self._pex_info_overrides)
-        self._envs.append(PEXEnvironment(pex_path, pex_info))
+      if pex_info.pex_path:
+        # set up other environments as specified in PEX_PATH
+        for pex_path in filter(None, pex_info.pex_path.split(os.pathsep)):
+          pex_info = PexInfo.from_pex(pex_path)
+          pex_info.update(self._pex_info_overrides)
+          self._envs.append(PEXEnvironment(pex_path, pex_info))
 
       # activate all of them
       for env in self._envs:

pex/pex_info.py

@@ -98,6 +98,7 @@ def from_env(cls, env=ENV):
       'inherit_path': supplied_env.PEX_INHERIT_PATH,
       'ignore_errors': supplied_env.PEX_IGNORE_ERRORS,
       'always_write_cache': supplied_env.PEX_ALWAYS_CACHE,
+      'pex_path': supplied_env.PEX_PATH,
     }
     # Filter out empty entries not explicitly set in the environment.
     return cls(info=dict((k, v) for (k, v) in pex_info.items() if v is not None))
@@ -165,6 +166,19 @@ def zip_safe(self):
   def zip_safe(self, value):
     self._pex_info['zip_safe'] = bool(value)
 
+  @property
+  def pex_path(self):
+    """A colon separated list of other pex files to merge into the runtime environment.
+
+    This pex info property is used to persist the PEX_PATH environment variable into the pex info
+    metadata for reuse within a built pex.
+    """
+    return self._pex_info.get('pex_path')
+
+  @pex_path.setter
+  def pex_path(self, value):
+    self._pex_info['pex_path'] = value
+
   @property
   def inherit_path(self):
     """Whether or not this PEX should be allowed to inherit system dependencies.

tests/test_integration.py

@@ -186,3 +186,95 @@ def test_pex_multi_resolve():
     assert len(included_dists) == 4
     for dist_substr in ('-cp27-', '-cp36-', '-manylinux1_x86_64', '-macosx_'):
       assert any(dist_substr in f for f in included_dists)
+
+
+@pytest.mark.xfail(reason='See https://github.com/pantsbuild/pants/issues/4682')
+def test_pex_re_exec_failure():
+  with temporary_dir() as output_dir:
+
+    # create 2 pex files for PEX_PATH
+    pex1_path = os.path.join(output_dir, 'pex1.pex')
+    res1 = run_pex_command(['--disable-cache', 'requests', '-o', pex1_path])
+    res1.assert_success()
+    pex2_path = os.path.join(output_dir, 'pex2.pex')
+    res2 = run_pex_command(['--disable-cache', 'flask', '-o', pex2_path])
+    res2.assert_success()
+    pex_path = ':'.join(os.path.join(output_dir, name) for name in ('pex1.pex', 'pex2.pex'))
+
+    # create test file test.py that attmepts to import modules from pex1/pex2
+    test_file_path = os.path.join(output_dir, 'test.py')
+    with open(test_file_path, 'w') as fh:
+      fh.write(dedent('''
+        import requests
+        import flask
+        import sys
+        import os
+        import subprocess
+        if 'RAN_ONCE' in os.environ::
+          print('Hello world')
+        else:
+          env = os.environ.copy()
+          env['RAN_ONCE'] = '1'
+          subprocess.call([sys.executable] + sys.argv, env=env)
+          sys.exit()
+        '''))
+
+    # set up env for pex build with PEX_PATH in the environment
+    env = os.environ.copy()
+    env['PEX_PATH'] = pex_path
+
+    # build composite pex of pex1/pex1
+    pex_out_path = os.path.join(output_dir, 'out.pex')
+    run_pex_command(['--disable-cache',
+      'wheel',
+      '-o', pex_out_path])
+
+    # run test.py with composite env
+    stdout, rc = run_simple_pex(pex_out_path, [test_file_path], env=env)
+
+    assert rc == 0
+    assert stdout == b'Hello world\n'
+
+
+def test_pex_path_arg():
+  with temporary_dir() as output_dir:
+
+    # create 2 pex files for PEX_PATH
+    pex1_path = os.path.join(output_dir, 'pex1.pex')
+    res1 = run_pex_command(['--disable-cache', 'requests', '-o', pex1_path])
+    res1.assert_success()
+    pex2_path = os.path.join(output_dir, 'pex2.pex')
+    res2 = run_pex_command(['--disable-cache', 'flask', '-o', pex2_path])
+    res2.assert_success()
+    pex_path = ':'.join(os.path.join(output_dir, name) for name in ('pex1.pex', 'pex2.pex'))
+
+    # parameterize the pex arg for test.py
+    pex_out_path = os.path.join(output_dir, 'out.pex')
+    # create test file test.py that attempts to import modules from pex1/pex2
+    test_file_path = os.path.join(output_dir, 'test.py')
+    with open(test_file_path, 'w') as fh:
+      fh.write(dedent('''
+        import requests
+        import flask
+        import sys
+        import os
+        import subprocess
+        if 'RAN_ONCE' in os.environ:
+          print('Success!')
+        else:
+          env = os.environ.copy()
+          env['RAN_ONCE'] = '1'
+          subprocess.call([sys.executable] + ['%s'] + sys.argv, env=env)
+          sys.exit()
+        ''' % pex_out_path))
+
+    # build out.pex composed from pex1/pex1
+    run_pex_command(['--disable-cache',
+      '--pex-path={}'.format(pex_path),
+      'wheel',
+      '-o', pex_out_path])
+
+    # run test.py with composite env
+    stdout, rc = run_simple_pex(pex_out_path, [test_file_path])
+    assert rc == 0
+    assert stdout == b'Success!\n'