Flexibility for supporting package groups tied to specific sources (PyTorch issue)

[polarathene]

Basic single source for specific packages

This will install torch and other deps by prioritizing the packages from the configured source in pyproject.toml. include_packages is required due to respect-source-order = true to avoid undesired behaviour.

[project]
name = "example"
requires-python = ">=3.10"
dependencies = [
    # Implicitly resolves `torch` to `2.3.1+cu121` via configured PyTorch source below (likewise for the other deps)
    "torch",
    "torchvision",
    "torchaudio",
]

# Packages defined in the `include_packages` of a source will prioritize resolution to that source, otherwise PyPi will be used first
# (avoids attempts to resolve packages to the PyTorch source, which logs many failures via `pdm outdated` or similar)
[tool.pdm.resolution]
respect-source-order = true

[[tool.pdm.source]]
name = "pytorch-cuda-12.1"
url = "https://download.pytorch.org/whl/cu121"
# Also ensure the transitive deps of `nvidia-*` packages are installed from this source, not PyPi (CUDA version mismatch)
include_packages = ["torch", "torchvision", "torchaudio", "nvidia-*"]

Caveats when introducing multiple sources (cannot restrict packages to sources via groups)

Now you'd like the pyproject.toml for your app to support installing PyTorch with CUDA or CPU (there's also ROCm for AMD, but let's keep it simple):

[project]
name = "example"
requires-python = ">=3.10"
# Deps moved to optional groups
dependencies = []

# These will need to be explicitly requested via `-G` / `--group` during install:
[project.optional-dependencies]
torch_cpu = ["torch==2.3.1+cpu", "torchvision", "torchaudio"]
torch_cuda = ["torch==2.3.1+cu121", "torchvision", "torchaudio"]

[tool.pdm.resolution]
respect-source-order = true

# Now there are two PyTorch sources to consider:
[[tool.pdm.source]]
name = "pytorch-cuda-12.1"
url = "https://download.pytorch.org/whl/cu121"
include_packages = ["torch_cuda", "nvidia-*"]

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

You'd then choose to install either via:

pdm install --group torch_cpu
pdm install --group torch_cuda

But unlike the prior example, the torchvision and torchaudio deps will install via PyPi.

• The include_packages for the referenced group is likely doing nothing and incorrect here.
• If you reverted include_packages to include the explicit torch / torchvision / torchaudio deps you'd find the CUDA source is resolved as it is defined first, thus torch_cpu is not able to use the correct source.

With the groups I was more explicit about the local identifier being present (+cpu / +cu121), which only has support for == / != with a version (which is mandatory when paired with a local identifier unfortunately).

• It would seem that you would need to do this for the torchvision and torchaudio packages too, keeping track of their explicit version to pin.
• That would be a bit more problematic however for the nvidia-* pattern, if you want to support the other local identifiers for different versions of CUDA (additional sources).

There is one other caveat I've not mentioned thus far. PyTorch CPU source only has the +cpu local identifier for the x86_64 package, the aarch64 / ARM64 package does not have that assigned, which complicates support there (more details covered in a related uv discussion). Meanwhile the CUDA group has no aarch64 packages (despite nvidia having CUDA compatible products that run on aarch64, nvidia publishes it's own alternative to these PyTorch wheels sources for that).

---

Official advice for this scenario

Official advice appears to be to:

• Use groups like shown above, but with the local identifiers and then producing separate lockfiles to install from: Issues with cross-platform lockfiles #2066 (comment)
• Use a single PyTorch source named torch and configure pypi.torch.url (a bit vague, not enough context but presumably a CLI option?): Variable expansion at lock time for package indexes #2063 (comment)
• Only applicable for PyTorch specifically, use this plugin: https://github.com/EmbarkStudios/pdm-plugin-torch (presently this plugin has been broken for over 6 months despite a PR fix since March, it's also unclear if it really improves the DX due to how the plugin approaches solving the issue)

[polarathene]

During my investigation with how viable PDM would be for a project using PyTorch, it may also be worth noting that presently the cache feature PDM has does not work well with even the single PyTorch source (5GB is cached locally):

• Once you have the pdm.lock resolved (optimized for Dockerfile build by avoiding crossplatform and using static urls via pdm lock --strategy / -S) an install with no cache could take 1m40s, vs 1m15s with the cache.
• The reason for the poor performance is explained here (uv by comparison with cache is only 4 seconds), PDM may improve on that in future (uv integration), or offer some option to operate with those same assumptions that uv uses.

---

For caching, it seems some official advice is to cache the .venv (or equivalent) itself: #2472 (comment)

That reference also suggests for those using Github Actions to use the official action pdm-project/setup-pdm which handles that for you.