Choosing specific Pytorch version when configuring accelerators with optional dependencies

[eginhard]

On Ubuntu with uv 0.5.5 with the following pyproject.toml, adapted from https://docs.astral.sh/uv/guides/integration/pytorch/#configuring-accelerators-with-optional-dependencies (torchvision removed for simplicity and minimum Pytorch version lowered):

[project]
name = "project"
version = "0.1.0"
requires-python = ">=3.12.0"
dependencies = []

[project.optional-dependencies]
cpu = [
  "torch>=2.0",
]
cu124 = [
  "torch>=2.0",
]

[tool.uv]
conflicts = [
  [
    { extra = "cpu" },
    { extra = "cu124" },
  ],
]

[tool.uv.sources]
torch = [
  { index = "pytorch-cpu", extra = "cpu", marker = "platform_system != 'Darwin'" },
  { index = "pytorch-cu124", extra = "cu124" },
]

[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true

[[tool.uv.index]]
name = "pytorch-cu124"
url = "https://download.pytorch.org/whl/cu124"
explicit = true

My goal is to test specific CPU Pytorch versions in CI and I'm not sure which of the below commands should work. The following works, but counter-intuitively installs the GPU version:

$ uv run --extra cpu --with torch==2.2.2 python -c "import torch; print(torch.__version__)"
2.2.2+cu121

This doesn't work:

$ uv run --extra cpu --with torch==2.2.2+cpu python -c "import torch; print(torch.__version__)"
  × No solution found when resolving `--with` dependencies:
  ╰─▶ Because there is no version of torch==2.2.2+cpu and you require torch==2.2.2+cpu, we
      can conclude that your requirements are unsatisfiable.

This fails when there is no lockfile present yet:

$ uv run --with torch==2.2.2 --index https://download.pytorch.org/whl/cpu python -c "import torch; print(torch.__version__)"
error: Found duplicate package `torch==2.5.1+cpu @ registry+https://download.pytorch.org/whl/cpu`

But it works if I sync first:

$ uv sync
$ uv run --with torch==2.2.2 --index https://download.pytorch.org/whl/cpu python -c "import torch; print(torch.__version__)"
2.2.2+cpu

[charliermarsh]

Hmm... Yeah I think --with isn't quite what you want here... --with is typically for overlaying additional dependencies, rather than applying further constraints to the project.

I'm tempted to say that the "right" way to do this is to add a tool.uv.constraint-dependencies to your pyproject.toml within the CI job itself, then run uv run...

[eginhard]

Oh, I see, that makes sense. There doesn't seem to be a way to set this directly from the command line? --constraint/UV_CONSTRAINT only seem to be available for uv pip compile/sync.

Otherwise the following works as well to not have to edit the pyproject.toml:

echo "torch==2.2.2" > constraints.txt
uv pip compile pyproject.toml --constraints constraints.txt --extra cpu > requirements.txt
uv pip sync requirements.txt --index https://download.pytorch.org/whl/cpu --index-strategy unsafe-best-match

[charliermarsh]

Yeah, that looks right. We have an open issue to allow passing these directly on the command-line which would help here! #9614