Recommended way for multi-platform dependencies from different indexes?

[arthur-st]

I'm developing a PyTorch project on a Mac, and I would like to be able to handle via PDM the lockfile(s) needed to handle running the project on arm64 Mac, amd64 Linux, and amd64 Linux with CUDA. For all packages except torch this is fine basically as is, but the 3 target torch versions live in different package indexes.

My ideal solution would be for [[tool.pdm.source]] declarations to have a scope parameter that I can use to instruct pdm to only use this source for [applicable] packages declared within a specific [project.optional-dependencies] group. From what I can tell, this is not possible.

Continuing, I tried specifying different torch versions in different optional groups, e.g.:

[project.optional-dependencies]
cpu = ["torch==2.1.1+cpu", "optimum[onnxruntime]~=1.14"]
cuda = ["torch==2.1.1+cu118", "optimum[onnxruntime-gpu]~=1.14"]

But that is a conflict of dependencies, apparently, between these 2 optional torch declarations, and no torch in the main dependencies.

Then I found that there was a suggested way to pin a specific wheel. This is not very ergonomic, and I also am not sure how it would even work if I got a dependency conflict with the same thing in different optional groups.

Further in that discussion, however, #696 is referenced for find-links source. I'm not sure that I understand how that works. I tried this:

[[tool.pdm.source]]
url = "https://download.pytorch.org/whl/cpu"
verify_ssl = true
name = "torch"
include_packages = ["torch"]
type = "find_links"

[[tool.pdm.source]]
url = "https://download.pytorch.org/whl/cu118"
verify_ssl = true
name = "torch2"
include_packages = ["torch"]
type = "find_links"

Running pdm lock afterwards did fail with a CandidateNotFound error. However, going to both of these URLs I spotted that the torch package link resolves to the same URL. What I did then is replace the 2 sources above with this:

[[tool.pdm.source]]
url = "https://download.pytorch.org/whl/torch/"
verify_ssl = true
name = "torch"
include_packages = ["torch"]
type = "find_links"

And that kind-of worked, inasmuch as getting torch wheels goes:

[[package]]
name = "torch"
version = "2.1.1+rocm5.6"
requires_python = ">=3.8.0"
summary = "Tensors and Dynamic neural networks in Python with strong GPU acceleration"
dependencies = [
    "filelock",
    "fsspec",
    "jinja2",
    "networkx",
    "pytorch-triton-rocm==2.1.0",
    "sympy",
    "typing-extensions",
]
files = [
    {file = "torch-2.1.1+cpu-cp310-cp310-linux_x86_64.whl", hash = "sha256:16e6cdb1991ff1f15f469b12fffb5b20f00b1ecff8c1073c23c7760fba90fcbc"},
    {file = "torch-2.1.1+cpu-cp310-cp310-win_amd64.whl", hash = "sha256:8e4e47aa4a004d2f5edd516b296da6fba07981ae9d1fc0da862abf651dc07d96"},
    {file = "torch-2.1.1+cpu.cxx11.abi-cp310-cp310-linux_x86_64.whl", hash = "sha256:bd5275b330a05e5d5d42f89b6a6ebd9fb98aa9b0e104c2e34fb52c8efae3eb12"},
    {file = "torch-2.1.1+cu118-cp310-cp310-linux_x86_64.whl", hash = "sha256:8e2914484e74aeba08570a52c8057cc5d59c19b72a623a6ded29dc9b988151c0"},
    {file = "torch-2.1.1+cu118-cp310-cp310-win_amd64.whl", hash = "sha256:765e93911984c813ddf74427eecd70c1efc785af7c231777632954b1bd1429d3"},
    {file = "torch-2.1.1+cu121-cp310-cp310-linux_x86_64.whl", hash = "sha256:ec76d11350c8e887a34d36602cd37f50639bc9cda9faea4d43f2070e9792e4a4"},
    {file = "torch-2.1.1+cu121-cp310-cp310-win_amd64.whl", hash = "sha256:44d1f28af7bd2c51b633175e9b99465f88ced80038f0cad57f25794f994637d2"},
    {file = "torch-2.1.1+cu121.with.pypi.cudnn-cp310-cp310-linux_x86_64.whl", hash = "sha256:08ea6d0ab3de86010acd75f6774cd09d8927f4255fa2aee95443c5662b9981a4"},
    {file = "torch-2.1.1+rocm5.5-cp310-cp310-linux_x86_64.whl", hash = "sha256:6f27bea73d4dcd437edb71dc767fb9bede8adf18b522a2f7800bed0b25d991ce"},
    {file = "torch-2.1.1+rocm5.6-cp310-cp310-linux_x86_64.whl", hash = "sha256:3c6b4b79ee67c7a0c0b856703d132e9426125240f955b279639c722765dbaa88"},
    {file = "torch-2.1.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:84fefd63356416c0cd20578637ccdbb82164993400ed17b57c951dd6376dcee8"},
    {file = "torch-2.1.1-cp310-none-macosx_10_9_x86_64.whl", hash = "sha256:1e1e5faddd43a8f2c0e0e22beacd1e235a2e447794d807483c94a9e31b54a758"},
    {file = "torch-2.1.1-cp310-none-macosx_11_0_arm64.whl", hash = "sha256:e76bf3c5c354874f1da465c852a2fb60ee6cbce306e935337885760f080f9baa"},
]

But this does present a few separate questions, that I cannot answer on my own at the moment:

1. The version gets resolved to version = "2.1.1+rocm5.6" in some unclear way?
2. On Mac this should auto-resolve to 1 of 2 wheels based on the CPU, but how will this resolve on Linux?

I don't mind generating platform-specific lock files, if that is possible, but I would strongly prefer to be able to solve this out of a single pyproject.toml file. The problem, to clearly restate, is the ability to give pdm an instruction to install a specific subset of torch wheels (CPU/CUDA/ROCm flavour), out of different indexes, at any given occasion.

[arthur-st]

I've done some further testing, and I think I did something not too carefully in my earlier testing. Using this in my pyproject.toml:

[project.optional-dependencies]
mac = ["torch==2.1.1;sys_platform == 'darwin'", "optimum[onnxruntime]~=1.14"]
cpu = ["torch==2.1.1+cpu", "optimum[onnxruntime]~=1.14"]
gpu = ["torch==2.1.1+cu118", "optimum[onnxruntime-gpu]~=1.14"]

[[tool.pdm.source]]
url = "https://download.pytorch.org/whl/torch/"
verify_ssl = true
name = "torch"
include_packages = ["torch"]
type = "find_links"

pip lock -G cpu -L pdm.cpu.lock works as expected, ditto for gpu. mac optional dependency group, however, is not working as expected — ROCm local version gets locked as per above, and pdm install from a mac lock file fails due to platform distribution mismatch.

[zhukovgreen]

Submitted the ticket in pdm for this
#3394

Also tried this with uv (my first experience with this tool and sadly - it was shining in this task):
https://gist.github.com/zhukovgreen/923be42feb2f1d648da4057399b6148f

tried uv sync on linux and Mac - all works as expected

[arthur-st]

Yeah, I've switched all my projects to uv now, as the maintainers have specifically investigated and addressed common complaints and feature requests from the ML community. Between cross-platform compilation and universal lock files, it finally feels nice to manage the dependencies.

That said, I do hope that an official, PSF-approved lock file format does come our way eventually, for continuity purposes.

[zhukovgreen]

pdm actually supports this fine, just documentation is a bit confusing. Check my comment here:
#3394 (comment)

this is a solution on pdm:

❯ cat pyproject.toml
[project]
name = "cross-platform-torch"
version = "0.1.0"
description = "Default template for PDM package"
authors = [
    {name = "zhukovgreen", email = "[email protected]"},
]
dependencies = [
    "torch==1.13.1+cpu; sys_platform != 'darwin'",
    "torch==1.13.1; sys_platform == 'darwin'",
]
requires-python = "==3.10.*"
readme = "README.md"
license = {text = "MIT"}


[tool.pdm]
distribution = false

[project.optional-dependencies]

[[tool.pdm.source]]
url = "https://download.pytorch.org/whl/cpu"
verify_ssl = true
name = "torch_cpu"

And then:

❯ pdm lock --platform=macOS
❯ pdm lock --platform=linux --append

This gives multiplatform lock file