#18965 Improved workflows for Torch (and Tensorflow?) Issue created by tgolsson This thread is intended to focus on actionable solutions for #18293, potentially multi-platform Pex locks, maybe with other components as well. Is your feature request related to a problem? Please describe. The big neural network libraries Tensorflow and Pytorch have significant hurdles in usage based on how they're versioned. In general, the basic wheels published to PyPi aren't usable by default, and any important usage requires custom indexes and compute-API-tagged builds. In the case of

tensorflow

these are differently named packages, and in the case of

torch

they are differentiated by local versions. These specific packages also differ in what platforms they support, and the local versions may not exist for all platforms. Currently, the universal lock provided by Pex specifies a maximum set of

{Linux, Mac}

but makes no guarantee either platform will be supported. This compounds the complexity of version specification, as local versions shall be preferred. In practice, one has to jump through hoops to ensure a basic resolve works. For example:

# pants.toml
[python-repos]
indexes = [
    "<https://pypi.org/simple/>",
    "<https://download.pytorch.org/whl/cpu/>",
]

# BUILD
python_requirement(
    name="torch",
    requirements=["torch==1.11.0"],
    resolve="generic",
)

python_requirement(
    name="torch_cpu",
    requirements=["torch==1.11.0+cpu"],
    resolve="cpu-only",  # only works on Linux since there's no +cpu for Mac.
)

This seems like we'd end up with one lock that works on both Mac (CPU) and Linux (CPU + random CUDA), and one for Linux with only CPU support. However; due to the PEP440 requirements both locks end up picking +CPU, and having no install candidates for Mac. The proper requirement is to specify it like this:

# BUILD
python_requirement(
    name="torch",
    requirements=["torch==1.11.0,!=1.11.0+cpu"],
    resolve="generic",
)

This means that the more different compute variants you need the more specific the generic constraints have to be. If we allowed Pex to separate the versions between Mac and Linux (and Windows?) we would have an easier time ensuring something reasonable gets picked. However; even so it'd be required to exclude all available versions in each resolve as we'd not want a massive

+cuda

result in the generic for Linux. Describe the solution you'd like The thread on Slack discusses multi-platform locks as a potential solution, but I think that still runs afoul of some footguns with the local-version specifiers as one still needs to exclude them. It might therefore be necessary to combine them with other solutions such as per-resolve extra-indexes. Additional context https://pantsbuild.slack.com/archives/C046T6T9U/p1683187396031649 pantsbuild/pants