Hey! I know there's a lot of messages here about torch and pex and containers, but I haven't been able to figure out how to make gpu torch work with pex (and containers, but that's the next step). I have pants

2.19

and I try to make

torch==2.2.2

work and I don't have any additional indexes added in pants.yaml because I need a 2.2.2 version which is available in pip I have web server with loads torch models and

pants run web/server.py

(running just the script itslef) works just fine, torch is able to instantiate and grabs gpu, however, I also have a pex binary target defined like this:

pex_binary(
    name="server",
    dependencies=[
        ":code",
        "voices:voices", # resources
        "web:systems", # resources
        "//:reqs#gunicorn",
    ],
    entry_point="web.server:app",
    include_tools=True,
    layout="packed",
    execution_mode="venv",
)

After some digging in this channel I tried nvidia deps directly and also added some args, none of which worked. I ended up with the following target:

pex_binary(
    name="server",
    dependencies=[
        ":code",
        "voices:voices",
        "web:systems",
        # Gunicorn cannot be inferred from the entry point
        "//:reqs#gunicorn",
        "//:reqs#nvidia-nccl-cu12",
        "//:reqs#nvidia-cuda-runtime-cu12",
        "//:reqs#nvidia-cuda-cupti-cu12",
    ],
    args=[
        "--venv-site-packages-copies",
        "--pip-version latest",
        "--resolver-version pip-2020-resolver",
    ],
    entry_point="web.server:app",
    include_tools=True,
    layout="packed",
    execution_mode="venv",
)

The issue I'm facing is the following error:

File "/home/vitalybushaev/.pex/venvs/7e7646c854c2c7a7b98fddc406d5b4ebeeafe709/779eb2cc0ca9e2fdd204774cbc41848e4e7c5055/pex", line 282, in <module>
    module = importlib.import_module(module_name)
  File "/opt/conda/lib/python3.10/importlib/__init__.py", line 126, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
  File "<frozen importlib._bootstrap>", line 1050, in _gcd_import
  File "<frozen importlib._bootstrap>", line 1027, in _find_and_load
  File "<frozen importlib._bootstrap>", line 1006, in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 688, in _load_unlocked
  File "<frozen importlib._bootstrap_external>", line 883, in exec_module
  File "<frozen importlib._bootstrap>", line 241, in _call_with_frames_removed
  File "/home/vitalybushaev/.pex/venvs/7e7646c854c2c7a7b98fddc406d5b4ebeeafe709/779eb2cc0ca9e2fdd204774cbc41848e4e7c5055/lib/python3.10/site-packages/web/server.py", line 6, in <module>
    import torch
  File "/home/vitalybushaev/.pex/venvs/7e7646c854c2c7a7b98fddc406d5b4ebeeafe709/779eb2cc0ca9e2fdd204774cbc41848e4e7c5055/lib/python3.10/site-packages/torch/__init__.py", line 237, in <module>
    from torch._C import *  # noqa: F403
ImportError: /home/vitalybushaev/.pex/venvs/7e7646c854c2c7a7b98fddc406d5b4ebeeafe709/779eb2cc0ca9e2fdd204774cbc41848e4e7c5055/lib/python3.10/site-packages/torch/lib/libtorch_cuda.so: undefined symbol: ncclCommRegister

I've generated lockfile and it does contain all the nvidia dependencies it needs, especially the nccl one which the error seems to be about. I've also inspected PEX-INFO file in pex directory and confirmed that nvidia deps are listed there. I've also tried to add to my

pants.yaml

the following index, which should contain the cuda version I need, but the result didn't change

[python-repos]
indexes.add = ["<https://download.pytorch.org/whl/cu121>"]

Any suggestions where to go from here ?

I had a similar issue and solved it by installing https://developer.nvidia.com/nccl natively on the machine that's supposed to run the PyTorch application

Thanks! For future reference, I resolved the issue by passing the following parameter to the

pex_binary

target

venv_site_packages_copies=True