Как настроить uv: маркеры окружения для CUDA в PyTorch на x86_64 и aarch64

Почему uv тянет torch+cu124 на aarch64 и как это исправить: используем маркеры окружения в pyproject.toml. Резолвинг CUDA 12.4/12.8 для x86_64 и aarch64.

Поддержка нескольких вариантов CUDA для разных архитектур кажется простой — пока резолвер не решит иначе. Типичный сценарий: Linux x86_64 должен оставаться на CUDA 12.4, тогда как aarch64 — на CUDA 12.8. PyTorch не выпускает колёса CUDA 12.4 для arm64, поэтому резолвер надо изначально отвести от попыток собрать такую связку.

Как воспроизвести проблему

Ниже приведённая конфигурация опирается на extras и required-environments, чтобы направлять uv к корректной сборке CUDA в зависимости от платформы. Несмотря на это, uv всё равно пытается подобрать колесо CUDA 12.4 для aarch64 и терпит неудачу.

[project]
name = "uv_extras_case"
version = "0.1.1"
description = ""
readme = "README.md"
requires-python = ">=3.13"
dependencies = []

[tool.uv]
conflicts = [
    [{extra="cu128"}, {extra="cu124"}]
]
required-environments = [
    "sys_platform == 'linux' and platform_machine == 'x86_64'  and extra == 'cu124'",
    "sys_platform == 'linux' and platform_machine == 'aarch64' and extra == 'cu128'",
]
index-strategy = "unsafe-best-match"

[project.optional-dependencies]
cu124 = [
    "torch==2.6.0+cu124",
]

cu128 = [
    "torch==2.7.0+cu128",
]

[[tool.uv.index]]
name = "pt-cu128"
url = "https://download.pytorch.org/whl/cu128"

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

Что именно идёт не так

Резолвер ведёт себя так, будто extra в required-environments не мешает ему исследовать несовместимые комбинации. Итог — несогласованный набор ограничений для aarch64, когда uv пытается получить torch==2.6.0+cu124, у которого просто нет колёс под эту платформу.

Решение не найдено … Поскольку для torch==2.6.0+cu124 нет … совместимых колёс для aarch64 … можно заключить, что … требования невыполнимы.

Иными словами, ограничивать extras через required-environments недостаточно, чтобы не позволить uv рассматривать связку torch+cu124 на arm.

Решение

Включите платформенное ограничение прямо в зависимость с помощью маркера окружения. Так uv недвусмысленно понимает, что torch+cu124 ставится только на Linux x86_64 и не должен пытаться устанавливаться на aarch64.

[project.optional-dependencies]
cu124 = [
    "torch==2.6.0+cu124; sys_platform == 'linux' and platform_machine == 'x86_64'",
]

cu128 = [
    "torch==2.7.0+cu128",
]

С этим правилом uv перестанет подбирать torch+cu124 под arm64, а для aarch64 резолвинг пойдёт по нужному пути — с билдом под CUDA 12.8.

Почему это важно

Мультиархитектурные окружения часто сочетают разные графы зависимостей и индексы. Если пакет не выпускается для конкретной пары платформа-акселератор, позволять резолверу доходить до этого перебором — дорого и ненадёжно. Задав платформенное ограничение прямо в строке зависимости, вы избегаете невыполнимых резолюций и делаете определение окружения самоочевидным.

Выводы

Если пакет зависит от платформы, фиксируйте это ограничение прямо в зависимости с помощью маркеров окружения. Опираться только на extras и общие правила выбора окружения может оставить пространство для неверных комбинаций — как в случае с torch+cu124 на aarch64. Небольшой, но явный маркер убирает двусмысленность и приводит резолвер к корректному решению.