Best way to handle architecture dependent environments

Question

Hi,

I have a project dependent on libraries like torch and jax which needs to install different packages based on the accelerator they use. While this is not a problem in development, when I submit a job to a hybrid GPU cluster, the environment needs to have the correct package based on which hardware the code is running on (e.g. torch-cuda or torch-rocm). Therefore, I prepared the following pyproject.toml file and defined the rocm and cuda packages to be mutually exclusive:

[project]
name = "new-project"
version = "0.1.0"
description = "desc"
requires-python = ">=3.14"
dependencies = [
    "jaxtyping>=0.2",
    "scipy>=1.14",
    "numpy>=2.1",
    "jaxlib",
]

# hardware targets as extras
[project.optional-dependencies]
rocm = [
    "torch>=2.13.0,<2.14",
    "torchvision",
    "triton-rocm",
    "jax[rocm7-local]",
]
cuda = [
    "torch>=2.13.0,<2.14",
    "torchvision",
    "triton",
    "jax[cuda13]>=0.4",
]

# rocm and cuda are mutually exclusive
[tool.uv]
conflicts = [
    [
        { extra = "rocm" },
        { extra = "cuda" },
    ],
]

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

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

[tool.uv.sources]
torch = [
  { index = "pytorch-cuda", extra = "cuda" },
  { index = "pytorch-rocm", extra = "rocm" },
]
torchvision = [
  { index = "pytorch-cuda", extra = "cuda" },
  { index = "pytorch-rocm", extra = "rocm" },
]
triton-rocm = { index = "pytorch-rocm", extra = "rocm" }

And when submitting an sbatch file to the cluster I use:

if lspci | grep -qi "NVIDIA"; then module load CUDA; else module load rocm; fi
if nvidia-smi > /dev/null 2>&1; then
  uv sync --extra cuda
elif rocminfo > /dev/null 2>&1; then
  uv sync --extra rocm
else
  uv sync
fi

srun ...

so that the architecture the code is running on is detected and proper packages are installed before the run. However, this causes problems when I submit multiple jobs to the queue and some of them go to AMD GPUs and some go to NVIDIA GPUs as the uv sync command changes the files inside .venv.

So, I would like to know if there is a way to use multiple .venv environments in one project or having a virtual environment where some packages point to different files based on hardware the code is running on?

Kind regards, Abdullah

Platform

Linux 5.14.0-611.34.1.el9_7.x86_64 x86_64 GNU/Linux

Version

uv 0.12.12 (x86_64-unknown-linux-gnu)