How tell pixi to use different pypi-options for different environments?

[asmith26]

Hi there,

I'm trying to install different pytorch dependencies depending on my computing environment (i.e. CPU or GPU). Following my answer at #1714 (comment), I'm trying:

[project]
name = "my_project"
requires-python = ">=3.12"

[tool.pixi.project]
channels = ["conda-forge"]
platforms = ["linux-64"]

[tool.pixi.dependencies]
python = ">=3.12"

[tool.pixi.feature.cpu]
pypi-options = { extra-index-urls = ["https://download.pytorch.org/whl/cpu"] }
pypi-dependencies = { "torch" = "*", "torchvision" = "*"}

[tool.pixi.feature.gpu]
system-requirements = {cuda = "12.4"}
pypi-options = { extra-index-urls = ["https://download.pytorch.org/whl/cu124"] }
pypi-dependencies = { "torch" = "*", "torchvision" = "*"}

[tool.pixi.environments]
default = {features = [], solve-group = "default"}
cpu = {features = ["cpu"], solve-group = "default"}
gpu = {features = ["gpu"], solve-group = "default"}

But when I look at pixi.lock I see:

  gpu:
    channels:
    - url: https://conda.anaconda.org/conda-forge/
    indexes:
    - https://pypi.org/simple
    - https://download.pytorch.org/whl/cpu
    - https://download.pytorch.org/whl/cu124

Any help great appreciated. Thanks!

[tdejager]

Oh hi! I missed this sorry, but the pypi-options are also combined over solve-groups there is not really a way around that currently. So to only use the indices for the environments, your only option now is to exclude the solve-groups.

[asmith26]

Many thanks for your help and info @tdejager. Can confirm this works for me:

# pyproject.toml  -  pypi

[project]
name = "my_project"
requires-python = ">=3.12"

[tool.pixi.project]
channels = ["conda-forge"]
platforms = ["linux-64"]

[tool.pixi.dependencies]
python = ">=3.12"

[tool.pixi.feature.cpu]
pypi-options = { extra-index-urls = ["https://download.pytorch.org/whl/cpu"] }
pypi-dependencies = { "torch" = "*", "torchvision" = "*"}

[tool.pixi.feature.gpu]
system-requirements = {cuda = "12.2"}
pypi-options = { extra-index-urls = ["https://download.pytorch.org/whl/cu124"] }
pypi-dependencies = { "torch" = "*", "torchvision" = "*"}

[tool.pixi.environments]
cpu = ["cpu"]
gpu = ["gpu"]

also (interestingly/strangely I did have to use requires-python = ">= 3.11" for this conda version):

# pyproject.toml  -  conda

[project]
name = "my_project"
version = "0.1.0"
description = "Add a short description here"
authors = [{name = "name", email = "email@email.com"}]
requires-python = ">= 3.11"

[tool.pixi]
project = {channels = ["conda-forge"], platforms = ["linux-64"]}
dependencies = {python = ">=3.11"}

[tool.pixi.feature.cpu]
channels = ["pytorch"]
dependencies = {"pytorch" = "*", "torchvision" = "*", "cpuonly" = "*"}

[tool.pixi.feature.gpu]
channels = ["pytorch", "nvidia"]
dependencies = {"pytorch" = "*", "torchvision" = "*", "pytorch-cuda" = "12.4"}

[tool.pixi.environments]
cpu = ["cpu"]
gpu = ["gpu"]

Then running the following works as expected:

$ pixi install && pixi install -e cpu && pixi install -e gpu

$ pixi run python -c "import torch; print(torch.cuda.is_available())"
...
ModuleNotFoundError: No module named 'torch'

$ pixi run -e cpu python -c "import torch; print(torch.cuda.is_available())"
False

$ pixi run -e gpu python -c "import torch; print(torch.cuda.is_available())"
True