astral-sh / uv

An extremely fast Python package and project manager, written in Rust.
https://docs.astral.sh/uv
Apache License 2.0
27.17k stars 785 forks source link

Cannot add torch dependent lib with Pytorch env following official guide. #9353

Closed sebastian-sz closed 5 hours ago

sebastian-sz commented 6 hours ago

Hi everyone. Thank you for the amazing work that has been done with uv and Pytorch integration.

I am following the official documentation to setup a Pytorch project. All works until I try to add a dependency that depends on Pytorch, e.g. uv add torchmetrics, which fails with:

error: Distribution `torch==2.5.1 @ registry+https://download.pytorch.org/whl/cpu` can't be installed because it doesn't have a source distribution or wheel for the current platform

but running uv pip install torchmetrics works with no issues.

To reproduce:

pyproject.toml

[project]
name = "uv-torchmetrics"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.10"
dependencies = [
]

[project.optional-dependencies]
cpu = [
  "torch>=2.5.1",
  "torchvision>=0.20.1",
]
cu124 = [
  "torch>=2.5.1",
  "torchvision>=0.20.1",
]

[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" },
]
torchvision = [
  { 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

And the command: uv add torchmetrics

Platform: Ubuntu 22.04 uv version: 0.5.3

If there is anything I could do to improve my workflow I'd appreciate the help!

charliermarsh commented 5 hours ago

I think what you ultimately want is this (you need to edit yourself, not via uv add):

[project]
name = "uv-torchmetrics"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.10"
dependencies = [
]

[project.optional-dependencies]
cpu = [
  "torch>=2.5.1",
  "torchvision>=0.20.1",
  "torchmetrics",
]
cu124 = [
  "torch>=2.5.1",
  "torchvision>=0.20.1",
  "torchmetrics",
]

[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" },
]
torchvision = [
  { index = "pytorch-cpu", extra = "cpu", marker = "platform_system != 'Darwin'" },
  { index = "pytorch-cu124", extra = "cu124" },
]
torchmetrics = [
  { 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

However, I think this is blocked on #9289. I'll come back to this once that issue is closed out.

charliermarsh commented 5 hours ago

Actually, sorry, I think what I posted above should work now?

sebastian-sz commented 5 hours ago

@charliermarsh You are absolutely right. The above works. Thank you for the swift response!