Trying to use pex to build a ML runtime.

[jkbbwr]

I want to exploit --pex-path to have two pex files.

entry.pex containing my project and its entrypoint runtime.pex containing nearly 3-5gb of runtime ML libraries.

I tried this out in a small test case

pex torch -o torch.pex (this takes a while) pex --pex-path torch.pex

❯ pex --pex-path torch.pex
Python 3.11.8 (main, Feb 12 2024, 14:50:05) [GCC 13.2.1 20230801] on linux
Type "help", "copyright", "credits" or "license" for more information.
(InteractiveConsole)
>>> import torch
Traceback (most recent call last):
  File "<console>", line 1, in <module>
  File "/home/me/.pex/installed_wheels/d86664ec85902967d902e78272e97d1aff1d331f7619d398d3ffab1c9b8e9157/torch-2.2.1-cp311-cp311-manylinux1_x86_64.whl/torch/__init__.py", line 237, in <module>
    from torch._C import *  # noqa: F403
    ^^^^^^^^^^^^^^^^^^^^^^
ImportError: libcudnn.so.8: cannot open shared object file: No such file or directory

[jsirois]

pex torch -o torch.pex (this takes a while)

Try --no-pre-install-wheels:

:; time pex torch --no-pre-install-wheels -o torch-no-pre-install.pex
/home/jsirois/bin/pex.venv/lib/python3.12/site-packages/pex/pex_builder.py:113: PEXWarning: The PEX zip at torch-no-pre-install.pex~ is not a valid zipapp: Could not find the `__main__` module.
This is likely due to the zip requiring ZIP64 extensions due to size or the
number of file entries or both. You can work around this limitation in Python's
`zipimport` module by re-building the PEX with `--layout packed` or
`--layout loose`.
  pex_warnings.warn(message)

real    0m16.592s
user    0m8.318s
sys     0m5.997s

VS:

:; time pex torch -o torch.pex
/home/jsirois/bin/pex.venv/lib/python3.12/site-packages/pex/pex_builder.py:113: PEXWarning: The PEX zip at torch.pex~ is not a valid zipapp: Could not find the `__main__` module.
This is likely due to the zip requiring ZIP64 extensions due to size or the
number of file entries or both. You can work around this limitation in Python's
`zipimport` module by re-building the PEX with `--layout packed` or
`--layout loose`.
  pex_warnings.warn(message)

real    3m8.961s
user    3m19.135s
sys     0m10.558s

I tried this out in a small test case

If you need maximum compatibility with how a normal venv works (which is presumably how you usually interact with torch, just ask Pex to use one:

:; pex --pex-path torch.pex --venv
Python 3.12.2 (main, Feb 25 2024, 16:35:05) [GCC 11.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
(InteractiveConsole)
>>> import torch
/home/jsirois/.pex/venvs/6d511c42df4cb5031bed93bce1b322c725c201ee/d5ce79b0d5309f05fd3ed096297ca9d02fa4948e/lib/python3.12/site-packages/torch/nn/modules/transformer.py:20: UserWarning: Failed to initialize NumPy: No module named 'numpy' (Triggered internally at ../torch/csrc/utils/tensor_numpy.cpp:84.)
  device: torch.device = torch.device(torch._C._get_default_device()),  # torch.device('cpu'),
>>>
now exiting InteractiveConsole...

[jsirois]

@jkbbwr I'm going to close this as answered question. The summary is use --venv mode on the application PEX - that will ensure the PEX lays itself (and any PEX_PATH adjoints) out as a normal venv at runtime. The side-light is Torch is massive and optimizations like --no-pre-install-wheels can help. Unzipping and re-zipping ~4GB worth of wheels is slow even with native tools.

If you find you still have questions though, please feel free to ask.

[jkbbwr]

Looks good to me, thasnks for the answer and the side information as well.