Repository: pex-tool/pex

Issue:
How to know whether we are running from within pex
Hello there!
I recently bumped the pex version of our project and figured out that the hack (import _pex) that we were using to detect whether or not we are running from within a pex does not work anymore. Here is the initial discussion about it started in 2019 by one of my colleagues: https://github.com/pantsbuild/pex/issues/721
Is there a new way to detect that we are in a pex ?
Thanks in advance.


Discussion:
One solution could be:

```
import os

def running_from_pex() -> bool:
    pex_root = os.environ.get("PEX_ROOT", os.path.expanduser("~/.pex"))
    return any([pex_root in p for p in sys.path])
```

But this is not ideal.

Also, we were able to get the path to the pex from within the pex via
```
import _pex
os.path.abspath(os.path.dirname(os.path.dirname(os.path.dirname(_pex.__file__))))
```
This does not work anymore. Is there a (good) way to get that information ?
The first thing to examine is whether you really need to know this any longer. Pex no longer ever imports and runs code from a PEX zip file, it always either unzips the PEX file fully 1st (the default), or else installs the PEX in a virtualenv and runs from there (`--venv`). For PEX files built without `--venv` you can use this:
```console
$ python -mpex --layout loose -o pex.loose
$ python -mpex --layout packed -o pex.packed
$ python -mpex --layout zipapp -o pex.zip
$ python ./pex.loose -c "import pkgutil; print(pkgutil.get_data('__main__', 'PEX-INFO').decode('utf-8'))"
{"bootstrap_hash": "f3f9a6aeb1108b12fe1d284d8b0718e4394f0a94", "build_properties": {"pex_version": "2.1.51"}, "code_hash": "da39a3ee5e6b4b0d3255bfef95601890afd80709", "distributions": {}, "emit_warnings": true, "ignore_errors": false, "includes_tools": false, "inherit_path": "false", "interpreter_constraints": [], "pex_hash": "92de425e2d09b97c1648b5b96721025f77c59e69", "pex_path": null, "requirements": [], "strip_pex_env": true, "venv": false, "venv_bin_path": "false", "venv_copies": false}
$ python ./pex.packed -c "import pkgutil; print(pkgutil.get_data('__main__', 'PEX-INFO').decode('utf-8'))"
{"bootstrap_hash": "f3f9a6aeb1108b12fe1d284d8b0718e4394f0a94", "build_properties": {"pex_version": "2.1.51"}, "code_hash": "da39a3ee5e6b4b0d3255bfef95601890afd80709", "distributions": {}, "emit_warnings": true, "ignore_errors": false, "includes_tools": false, "inherit_path": "false", "interpreter_constraints": [], "pex_hash": "92de425e2d09b97c1648b5b96721025f77c59e69", "pex_path": null, "requirements": [], "strip_pex_env": true, "venv": false, "venv_bin_path": "false", "venv_copies": false}
$ python ./pex.zip -c "import pkgutil; print(pkgutil.get_data('__main__', 'PEX-INFO').decode('utf-8'))"
{"bootstrap_hash": "f3f9a6aeb1108b12fe1d284d8b0718e4394f0a94", "build_properties": {"pex_version": "2.1.51"}, "code_hash": "da39a3ee5e6b4b0d3255bfef95601890afd80709", "distributions": {}, "emit_warnings": true, "ignore_errors": false, "includes_tools": false, "inherit_path": "false", "interpreter_constraints": [], "pex_hash": "92de425e2d09b97c1648b5b96721025f77c59e69", "pex_path": null, "requirements": [], "strip_pex_env": true, "venv": false, "venv_bin_path": "false", "venv_copies": false}
```

For PEX files built with `--venv` there is no simple way to determine that Pex is involved at all, but that's also the point of having the PEX file install itself in a venv - to make the fact Pex is involved be completely irrelevant,
Hello @jsirois. Thank you for your answer. We've implemented a package, https://github.com/criteo/cluster-pack, that is used to automatically capture and "zip" the python user environment. This captured user environment can then be shipped on executors in a cluster to run user code.
We "zip" the python user environment only if active environment is a venv or conda env. If it is already a pex, we don't "zip" it, instead we use it directly. <- This is one reason why we would like to know whether we are running from within a pex or not.

Also, relative to the second question, we would like to retrieve the path of the pex when running within the pex to be able to call pip installed in the pex. pip don't clearly expose an API, so the best way to use pip from python is by starting a process (subprocess or other).
Thanks for the background - that helps.

For the second bit with pip, you can run it like so - which is what Pex itself does: 
https://github.com/pantsbuild/pex/blob/da25bae605407b12461180fc8701cd0e78934b8e/pex/pip.py#L528-L588

```console
$ cat <<EOF > src/run_pip_via_runpy.py
import runpy

runpy.run_module(mod_name="pip", run_name="__main__", alter_sys=True)
EOF
$ pex pip cowsay -c cowsay -o contains-pip.pex -D src/
$ ./contains-pip.pex moo
  ___
| moo |
  ===
   \
    \
      ^__^
      (oo)\_______
      (__)\       )\/\
          ||----w |
          ||     ||
$ PEX_MODULE=run_pip_via_runpy ./contains-pip.pex --version
pip 21.3 from /home/jsirois/.pex/installed_wheels/487b7183d721c3e6138744c340b6d4f80b84168f/pip-21.3-py3-none-any.whl/pip (python 3.9)
```

Thinking a bit more on the 1st bit...
Hello @jsirois. Thanks for your answer.

> For the second bit with pip, you can run it like so - which is what Pex itself does:

I would need to provide the pex itself as interpreter. But for that, I need the path to the pex.

Another reason why we need to get the path to the pex from within the pex is because we want to be able to ship the pex on executors (in a cluster). Common use case: an app, running from a pex, submit a job on a cluster. The code to run the job on the cluster is in the pex.

We are unfortunately a bit stuck. Any idea to get the path to the pex ?
No good ideas with PEXes as-is. It should be possible to arrange sys.argv[0] to always be the path to the PEX, but that might need to be hidden behind a build time option.

Before going there though, I want to push on your use case harder. For example, does this solve your problem? It's not even Pex-specific - you could use it (redundantly) with a venv interpreter:
```console
$ cat src/main.py 
import os
import subprocess
import sys


def run_python_with_same_sys_path(*args: str) -> subprocess.Popen:
    return subprocess.Popen(
        [sys.executable, "-sS", *args],
        env={**os.environ, "PYTHONPATH": os.pathsep.join(sys.path)}
    )


sys.exit(run_python_with_same_sys_path("-m", "pip", "--version").wait())

$ pex -Dsrc pip -emain -opex-where.pex
$ ./pex-where.pex 
pip 21.3 from /home/jsirois/.pex/installed_wheels/487b7183d721c3e6138744c340b6d4f80b84168f/pip-21.3-py3-none-any.whl/pip (python 3.9)
```
Thanks for your quick answer.

This indeed helps me executing pip that have been embedded in the pex. Thanks!

But this does not tell me where is the pex. By knowing where is the pex, I wouldn't need to do what you suggest. I could directly do something like:
```
subprocess.Popen(
        ["/path/to/my/pex", "-m", "pip", "--version"],
)
```

What do you think ?
As I said, It should be possible to arrange sys.argv[0] to always be the path to the PEX, but that might need to be hidden behind a build time option. I wanted to unblock you though. It sounds like you are unblocked and don't actually need the path to the Pex to accomplish your goals. If I can circle back and add the sys.argv[0] functionality though, it looks like you could clean up your code slightly. Is that about right?
We are still blocked. I think I was not very clear when describing another use case, sorry for that.
We need to know the path to the pex for the following use case/example: I have a job A running from a pex that starts another job B on a remote cluster. The job B will also run from the pex. So before starting the job B, we need to ship the pex on the remote cluster.

To do so, what we were doing before bumping the pex version, was to detect that we are running from a pex (by trying to import _pex) and then by shipping the pex on the remote cluster (path of the pex could be found via os.path.abspath(os.path.dirname(os.path.dirname(os.path.dirname(_pex.__file__)))))
Aha - gotcha. The pip thing was distracting me since there are ~18 ways to do that without knowing the PEX location.
Let me experiment a bit with the sys.argv[0] idea. That would seem to be the most direct, since it applies to any program in any language - where am I at? `sys.argv[0]` of course.
OK, the `sys.argv[0]` trick works today for any PEX with no entrypoint defined, a `module:func` style entrypoint or a script or console script entrypoint. It only fails to work for a `module` style entrypoints. To get it working for `module` style entrypoints is easy, except for #1018, which it breaks. You either get sys.argv[0] == the PEX file or you get working pickle, but not both.

1st: @nateagr does sys.argv[0] unblock you currently then? I.E.: Do you or must you use `module` entry points for your PEXes as opposed to `module:func` entry points or script / console script entry points?

If not, I think I will need to put this behavior behind a flag so you can make the choice to break pickle in top level entry point modules in favor of having sys.argv[0] reflect the PEX file. Alternatively I'll need to expose a non-scrubbed `PEX_FILE` env var, also maybe behind a flag.
We mostly run modules with PEXes (https://github.com/criteo/cluster-pack/blob/master/cluster_pack/skein/skein_config_builder.py#L122).

> Alternatively I'll need to expose a non-scrubbed PEX_FILE env var, also maybe behind a flag.

I'd vote for that one but you are the decider :)
Hello @jsirois , did you have time to think about it ? I can contribute to a PR if needed
I think a PEX_FILE-style env var is the way to go. I was away for a few weeks on vacation / moving into a new house, but I'll take a crack and add you to the review.

Resolving patch:
diff --git a/pex/pex_builder.py b/pex/pex_builder.py
--- a/pex/pex_builder.py
+++ b/pex/pex_builder.py
@@ -59,12 +59,8 @@ class Value(Enum.Value):
 __INSTALLED_FROM__ = '__PEX_EXE__'
 
 
-def __re_exec__(pex, args):
-  # N.B.: This is read upon re-exec below to point sys.argv[0] back to the original pex before
-  # unconditionally scrubbing the env var and handing off to user code.
-  os.environ[__INSTALLED_FROM__] = pex
-
-  os.execv(args[0], args)
+def __re_exec__(argv0, *extra_launch_args):
+  os.execv(argv0, [argv0] + list(extra_launch_args) + sys.argv[1:])
 
 
 def __maybe_install_pex__(pex, pex_root, pex_hash):
@@ -75,8 +71,12 @@ def __maybe_install_pex__(pex, pex_root, pex_hash):
   if not installed_location:
     return
 
+  # N.B.: This is read upon re-exec below to point sys.argv[0] back to the original pex before
+  # unconditionally scrubbing the env var and handing off to user code.
+  os.environ[__INSTALLED_FROM__] = pex
+
   TRACER.log('Executing installed PEX for {{}} at {{}}'.format(pex, installed_location))
-  __re_exec__(pex, [sys.executable, installed_location] + sys.argv[1:])
+  __re_exec__(sys.executable, installed_location)
 
 
 def __maybe_run_venv__(pex, pex_root, pex_path):
@@ -97,7 +97,7 @@ def __maybe_run_venv__(pex, pex_root, pex_path):
     return
 
   TRACER.log('Executing venv PEX for {{}} at {{}}'.format(pex, venv_pex))
-  __re_exec__(pex, [venv_pex] + sys.argv[1:])
+  __re_exec__(venv_pex)
 
 
 __entry_point__ = None
@@ -121,6 +121,7 @@ def __maybe_run_venv__(pex, pex_root, pex_path):
 sys.path.insert(0, os.path.abspath(os.path.join(__entry_point__, {bootstrap_dir!r})))
 
 if not __installed_from__:
+    os.environ['PEX'] = os.path.realpath(__entry_point__)
     from pex.variables import ENV, Variables
     __pex_root__ = Variables.PEX_ROOT.value_or(ENV, {pex_root!r})
     if not ENV.PEX_TOOLS and Variables.PEX_VENV.value_or(ENV, {is_venv!r}):
@@ -130,6 +131,8 @@ def __maybe_run_venv__(pex, pex_root, pex_path):
         pex_path=Variables.PEX_PATH.value_or(ENV, {pex_path!r}),
       )
     __maybe_install_pex__(__entry_point__, pex_root=__pex_root__, pex_hash={pex_hash!r})
+else:
+    os.environ['PEX'] = os.path.realpath(__installed_from__)
 
 from pex.pex_bootstrapper import bootstrap_pex
 bootstrap_pex(__entry_point__)
diff --git a/pex/tools/commands/venv.py b/pex/tools/commands/venv.py
--- a/pex/tools/commands/venv.py
+++ b/pex/tools/commands/venv.py
@@ -190,8 +190,6 @@ def iter_valid_venv_pythons():
                     "PEX_PYTHON_PATH",
                     "PEX_VERBOSE",
                     "PEX_EMIT_WARNINGS",
-                    # This is used in re-exec.
-                    "__PEX_EXE__",
                     # This is used by the vendoring system.
                     "__PEX_UNVENDORED__",
                     # This is _not_ used (it is ignored), but it's present under CI and simplest to


Test patch:
diff --git a/tests/test_pex_builder.py b/tests/test_pex_builder.py
--- a/tests/test_pex_builder.py
+++ b/tests/test_pex_builder.py
@@ -5,6 +5,7 @@
 import os
 import stat
 import subprocess
+import sys
 import zipfile
 
 import pytest
@@ -15,7 +16,7 @@
 from pex.layout import Layout
 from pex.pex import PEX
 from pex.pex_builder import CopyMode, PEXBuilder
-from pex.testing import built_wheel, make_bdist
+from pex.testing import built_wheel, make_bdist, make_env
 from pex.testing import write_simple_pex as write_pex
 from pex.typing import TYPE_CHECKING
 from pex.variables import ENV
@@ -404,3 +405,86 @@ def test_pex_builder_exclude_bootstrap_testing(
     assert not [
         f for f in bootstrap_files if f.endswith(("testing.py", "testing.pyc"))
     ], "Expected testing support files to be stripped from the Pex `.bootstrap`."
+
+
+@pytest.mark.parametrize(
+    "strip_pex_env", [pytest.param(True, id="StripPexEnv"), pytest.param(False, id="NoStripPexEnv")]
+)
+@pytest.mark.parametrize(
+    "layout", [pytest.param(layout, id=layout.value) for layout in Layout.values()]
+)
+@pytest.mark.parametrize("venv", [pytest.param(True, id="VENV"), pytest.param(False, id="UNZIP")])
+def test_pex_env_var_issues_1485(
+    tmpdir,  # type: Any
+    strip_pex_env,  # type: bool
+    venv,  # type: bool
+    layout,  # type: Layout.Value
+):
+    # type: (...) -> None
+
+    pex_path = os.path.join(str(tmpdir), "empty.pex")
+    pb = PEXBuilder()
+    pb.info.strip_pex_env = strip_pex_env
+    pb.info.venv = venv
+    pb.build(pex_path, layout=layout)
+
+    launch_args = [pex_path] if layout == Layout.ZIPAPP else [sys.executable, pex_path]
+    pex_root = os.path.join(str(tmpdir), "pex_root")
+
+    def assert_pex_env_var(
+        script,  # type: str
+        expected_pex_env_var,  # type: str
+        expect_empty_pex_root,  # type: bool
+        expected_layout=layout,  # type: Layout.Value
+    ):
+        # type: (...) -> None
+        if expect_empty_pex_root:
+            assert not os.path.exists(pex_root)
+        else:
+            assert len(os.listdir(pex_root)) > 0
+        output = subprocess.check_output(
+            launch_args + ["-c", script], env=make_env(PEX_ROOT=pex_root)
+        )
+        actual = output.decode("utf-8").strip()
+        assert os.path.realpath(expected_pex_env_var) == actual
+        if Layout.ZIPAPP == expected_layout:
+            assert zipfile.is_zipfile(actual)
+        else:
+            assert os.path.isdir(actual)
+
+    print_pex_env_var_script = "import os; print(os.environ['PEX'])"
+
+    assert_pex_env_var(
+        script=print_pex_env_var_script,
+        expected_pex_env_var=pex_path,
+        expect_empty_pex_root=True,
+    )
+    assert_pex_env_var(
+        script=print_pex_env_var_script,
+        expected_pex_env_var=pex_path,
+        expect_empty_pex_root=False,
+    )
+
+    other_pex_path = os.path.join(str(tmpdir), "other.pex")
+    other_pb = PEXBuilder()
+    other_pb.info.includes_tools = True
+    other_pex_main = os.path.join(str(tmpdir), "main.py")
+    with open(other_pex_main, mode="w") as fp:
+        fp.write(print_pex_env_var_script)
+    other_pb.set_executable(other_pex_main)
+    other_pb.build(other_pex_path, layout=Layout.ZIPAPP)
+
+    def assert_pex_env_var_nested(**env):
+        # type: (**Any) -> None
+        assert_pex_env_var(
+            script="import subprocess; subprocess.check_call([{other_pex!r}], env={env!r})".format(
+                other_pex=other_pex_path, env=make_env(**env)
+            ),
+            expected_pex_env_var=other_pex_path,
+            expect_empty_pex_root=False,
+            expected_layout=Layout.ZIPAPP,
+        )
+
+    assert_pex_env_var_nested()
+    assert_pex_env_var_nested(PEX_VENV=False)
+    assert_pex_env_var_nested(PEX_VENV=True)