diff --git a/Cargo.lock b/Cargo.lock index 0f738bb09173..945773f7be27 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4573,6 +4573,7 @@ dependencies = [ "uv-shell", "uv-static", "uv-tool", + "uv-torch", "uv-trampoline-builder", "uv-types", "uv-version", @@ -4786,6 +4787,7 @@ dependencies = [ "uv-resolver", "uv-settings", "uv-static", + "uv-torch", "uv-version", "uv-warnings", ] @@ -4839,6 +4841,7 @@ dependencies = [ "uv-pypi-types", "uv-small-str", "uv-static", + "uv-torch", "uv-version", "uv-warnings", ] @@ -5341,6 +5344,7 @@ dependencies = [ "rustc-hash", "serde", "thiserror 2.0.11", + "uv-pep440", "uv-small-str", ] @@ -5636,6 +5640,7 @@ dependencies = [ "uv-python", "uv-resolver", "uv-static", + "uv-torch", "uv-warnings", ] @@ -5707,6 +5712,20 @@ dependencies = [ "uv-virtualenv", ] +[[package]] +name = "uv-torch" +version = "0.1.0" +dependencies = [ + "clap", + "either", + "schemars", + "serde", + "uv-distribution-types", + "uv-normalize", + "uv-pep440", + "uv-platform-tags", +] + [[package]] name = "uv-trampoline-builder" version = "0.0.1" diff --git a/Cargo.toml b/Cargo.toml index 8eb3b4eb4d55..e152c61db6b5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -63,6 +63,7 @@ uv-small-str = { path = "crates/uv-small-str" } uv-state = { path = "crates/uv-state" } uv-static = { path = "crates/uv-static" } uv-tool = { path = "crates/uv-tool" } +uv-torch = { path = "crates/uv-torch" } uv-trampoline-builder = { path = "crates/uv-trampoline-builder" } uv-types = { path = "crates/uv-types" } uv-version = { path = "crates/uv-version" } diff --git a/crates/uv-cli/Cargo.toml b/crates/uv-cli/Cargo.toml index 8dfc577690e9..debc21a74595 100644 --- a/crates/uv-cli/Cargo.toml +++ b/crates/uv-cli/Cargo.toml @@ -28,6 +28,7 @@ uv-python = { workspace = true, features = ["clap", "schemars"]} uv-resolver = { workspace = true, features = ["clap"] } uv-settings = { workspace = true, features = ["schemars"] } uv-static = { workspace = true } +uv-torch = { workspace = true, features = ["clap"] } uv-version = { workspace = true } uv-warnings = { workspace = true } diff --git a/crates/uv-cli/src/lib.rs b/crates/uv-cli/src/lib.rs index d168adf504a5..435aa2f4ce7a 100644 --- a/crates/uv-cli/src/lib.rs +++ b/crates/uv-cli/src/lib.rs @@ -21,6 +21,7 @@ use uv_pypi_types::VerbatimParsedUrl; use uv_python::{PythonDownloads, PythonPreference, PythonVersion}; use uv_resolver::{AnnotationStyle, ExcludeNewer, ForkStrategy, PrereleaseMode, ResolutionMode}; use uv_static::EnvVars; +use uv_torch::TorchMode; pub mod comma; pub mod compat; @@ -1258,6 +1259,19 @@ pub struct PipCompileArgs { #[arg(long, overrides_with("emit_index_annotation"), hide = true)] pub no_emit_index_annotation: bool, + /// The backend to use when fetching packages in the `PyTorch` ecosystem (e.g., `cu126` or `auto`) + /// + /// When set, uv will ignore the configured index URLs for packages in the `PyTorch` ecosystem, + /// and will instead use the defined backend. + /// + /// For example, when set to `cpu`, uv will use the CPU-only `PyTorch` index; when set to `cu126`, + /// uv will use the `PyTorch` index for CUDA 12.6. + /// + /// The `auto` mode will attempt to detect the appropriate `PyTorch` index based on the currently + /// installed CUDA drivers. + #[arg(long, value_enum)] + pub torch_backend: Option, + #[command(flatten)] pub compat_args: compat::PipCompileCompatArgs, } @@ -1499,6 +1513,19 @@ pub struct PipSyncArgs { #[arg(long)] pub dry_run: bool, + /// The backend to use when fetching packages in the `PyTorch` ecosystem (e.g., `cu126` or `auto`) + /// + /// When set, uv will ignore the configured index URLs for packages in the `PyTorch` ecosystem, + /// and will instead use the defined backend. + /// + /// For example, when set to `cpu`, uv will use the CPU-only `PyTorch` index; when set to `cu126`, + /// uv will use the `PyTorch` index for CUDA 12.6. + /// + /// The `auto` mode will attempt to detect the appropriate `PyTorch` index based on the currently + /// installed CUDA drivers. + #[arg(long, value_enum)] + pub torch_backend: Option, + #[command(flatten)] pub compat_args: compat::PipSyncCompatArgs, } @@ -1791,6 +1818,19 @@ pub struct PipInstallArgs { #[arg(long)] pub dry_run: bool, + /// The backend to use when fetching packages in the `PyTorch` ecosystem (e.g., `cu126` or `auto`) + /// + /// When set, uv will ignore the configured index URLs for packages in the `PyTorch` ecosystem, + /// and will instead use the defined backend. + /// + /// For example, when set to `cpu`, uv will use the CPU-only `PyTorch` index; when set to `cu126`, + /// uv will use the `PyTorch` index for CUDA 12.6. + /// + /// The `auto` mode will attempt to detect the appropriate `PyTorch` index based on the currently + /// installed CUDA drivers. + #[arg(long, value_enum)] + pub torch_backend: Option, + #[command(flatten)] pub compat_args: compat::PipInstallCompatArgs, } diff --git a/crates/uv-client/Cargo.toml b/crates/uv-client/Cargo.toml index c94687845555..5c5aa67ead10 100644 --- a/crates/uv-client/Cargo.toml +++ b/crates/uv-client/Cargo.toml @@ -25,6 +25,7 @@ uv-platform-tags = { workspace = true } uv-pypi-types = { workspace = true } uv-small-str = { workspace = true } uv-static = { workspace = true } +uv-torch = { workspace = true } uv-version = { workspace = true } uv-warnings = { workspace = true } diff --git a/crates/uv-client/src/registry_client.rs b/crates/uv-client/src/registry_client.rs index ec636775dd03..d9256008c44a 100644 --- a/crates/uv-client/src/registry_client.rs +++ b/crates/uv-client/src/registry_client.rs @@ -14,12 +14,6 @@ use tokio::sync::Semaphore; use tracing::{info_span, instrument, trace, warn, Instrument}; use url::Url; -use crate::base_client::{BaseClientBuilder, ExtraMiddleware}; -use crate::cached_client::CacheControl; -use crate::html::SimpleHtml; -use crate::remote_metadata::wheel_metadata_from_remote_zip; -use crate::rkyvutil::OwnedArchive; -use crate::{BaseClient, CachedClient, CachedClientError, Error, ErrorKind}; use uv_cache::{Cache, CacheBucket, CacheEntry, WheelCache}; use uv_configuration::KeyringProviderType; use uv_configuration::{IndexStrategy, TrustedHost}; @@ -34,12 +28,21 @@ use uv_pep508::MarkerEnvironment; use uv_platform_tags::Platform; use uv_pypi_types::{ResolutionMetadata, SimpleJson}; use uv_small_str::SmallString; +use uv_torch::TorchStrategy; + +use crate::base_client::{BaseClientBuilder, ExtraMiddleware}; +use crate::cached_client::CacheControl; +use crate::html::SimpleHtml; +use crate::remote_metadata::wheel_metadata_from_remote_zip; +use crate::rkyvutil::OwnedArchive; +use crate::{BaseClient, CachedClient, CachedClientError, Error, ErrorKind}; /// A builder for an [`RegistryClient`]. #[derive(Debug, Clone)] pub struct RegistryClientBuilder<'a> { index_urls: IndexUrls, index_strategy: IndexStrategy, + torch_backend: Option, cache: Cache, base_client_builder: BaseClientBuilder<'a>, } @@ -49,6 +52,7 @@ impl RegistryClientBuilder<'_> { Self { index_urls: IndexUrls::default(), index_strategy: IndexStrategy::default(), + torch_backend: None, cache, base_client_builder: BaseClientBuilder::new(), } @@ -68,6 +72,12 @@ impl<'a> RegistryClientBuilder<'a> { self } + #[must_use] + pub fn torch_backend(mut self, torch_backend: Option) -> Self { + self.torch_backend = torch_backend; + self + } + #[must_use] pub fn keyring(mut self, keyring_type: KeyringProviderType) -> Self { self.base_client_builder = self.base_client_builder.keyring(keyring_type); @@ -145,6 +155,7 @@ impl<'a> RegistryClientBuilder<'a> { RegistryClient { index_urls: self.index_urls, index_strategy: self.index_strategy, + torch_backend: self.torch_backend, cache: self.cache, connectivity, client, @@ -166,6 +177,7 @@ impl<'a> RegistryClientBuilder<'a> { RegistryClient { index_urls: self.index_urls, index_strategy: self.index_strategy, + torch_backend: self.torch_backend, cache: self.cache, connectivity, client, @@ -181,6 +193,7 @@ impl<'a> TryFrom> for RegistryClientBuilder<'a> { Ok(Self { index_urls: IndexUrls::default(), index_strategy: IndexStrategy::default(), + torch_backend: None, cache: Cache::temp()?, base_client_builder: value, }) @@ -194,6 +207,8 @@ pub struct RegistryClient { index_urls: IndexUrls, /// The strategy to use when fetching across multiple indexes. index_strategy: IndexStrategy, + /// The strategy to use when selecting a `PyTorch` backend, if any. + torch_backend: Option, /// The underlying HTTP client. client: CachedClient, /// Used for the remote wheel METADATA cache. @@ -230,6 +245,15 @@ impl RegistryClient { self.timeout } + /// Return the appropriate index URLs for the given [`PackageName`]. + fn index_urls_for(&self, package_name: &PackageName) -> impl Iterator { + self.torch_backend + .as_ref() + .and_then(|torch_backend| torch_backend.index_urls(package_name)) + .map(Either::Left) + .unwrap_or_else(|| Either::Right(self.index_urls.indexes().map(Index::url))) + } + /// Fetch a package from the `PyPI` simple API. /// /// "simple" here refers to [PEP 503 – Simple Repository API](https://peps.python.org/pep-0503/) @@ -243,23 +267,24 @@ impl RegistryClient { capabilities: &IndexCapabilities, download_concurrency: &Semaphore, ) -> Result)>, Error> { + // If `--no-index` is specified, avoid fetching regardless of whether the index is implicit, + // explicit, etc. + if self.index_urls.no_index() { + return Err(ErrorKind::NoIndex(package_name.to_string()).into()); + } + let indexes = if let Some(index) = index { Either::Left(std::iter::once(index)) } else { - Either::Right(self.index_urls.indexes().map(Index::url)) + Either::Right(self.index_urls_for(package_name)) }; - let mut it = indexes.peekable(); - if it.peek().is_none() { - return Err(ErrorKind::NoIndex(package_name.to_string()).into()); - } - let mut results = Vec::new(); match self.index_strategy { // If we're searching for the first index that contains the package, fetch serially. IndexStrategy::FirstIndex => { - for index in it { + for index in indexes { let _permit = download_concurrency.acquire().await; if let Some(metadata) = self .simple_single_index(package_name, index, capabilities) @@ -273,7 +298,7 @@ impl RegistryClient { // Otherwise, fetch concurrently. IndexStrategy::UnsafeBestMatch | IndexStrategy::UnsafeFirstMatch => { - results = futures::stream::iter(it) + results = futures::stream::iter(indexes) .map(|index| async move { let _permit = download_concurrency.acquire().await; let metadata = self diff --git a/crates/uv-distribution-types/src/index_url.rs b/crates/uv-distribution-types/src/index_url.rs index 69bce9664a9f..0744b9f9b6e3 100644 --- a/crates/uv-distribution-types/src/index_url.rs +++ b/crates/uv-distribution-types/src/index_url.rs @@ -492,6 +492,11 @@ impl<'a> IndexUrls { ) } } + + /// Return the `--no-index` flag. + pub fn no_index(&self) -> bool { + self.no_index + } } bitflags::bitflags! { diff --git a/crates/uv-platform-tags/Cargo.toml b/crates/uv-platform-tags/Cargo.toml index dcab932ad26c..2be566697278 100644 --- a/crates/uv-platform-tags/Cargo.toml +++ b/crates/uv-platform-tags/Cargo.toml @@ -16,6 +16,7 @@ doctest = false workspace = true [dependencies] +uv-pep440 = { workspace = true } uv-small-str = { workspace = true } memchr = { workspace = true } diff --git a/crates/uv-platform-tags/src/accelerator.rs b/crates/uv-platform-tags/src/accelerator.rs new file mode 100644 index 000000000000..014104c4687f --- /dev/null +++ b/crates/uv-platform-tags/src/accelerator.rs @@ -0,0 +1,17 @@ +use std::fmt; + +use uv_pep440::Version; + +#[derive(Debug, Clone, Eq, PartialEq, serde::Deserialize, serde::Serialize)] +#[serde(tag = "name", rename_all = "lowercase")] +pub enum Accelerator { + Cuda { driver_version: Version }, +} + +impl fmt::Display for Accelerator { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + match self { + Self::Cuda { driver_version } => write!(f, "CUDA {driver_version}"), + } + } +} diff --git a/crates/uv-platform-tags/src/lib.rs b/crates/uv-platform-tags/src/lib.rs index 1fe09c405555..9ffe68d004ac 100644 --- a/crates/uv-platform-tags/src/lib.rs +++ b/crates/uv-platform-tags/src/lib.rs @@ -1,10 +1,12 @@ pub use abi_tag::{AbiTag, ParseAbiTagError}; +pub use accelerator::Accelerator; pub use language_tag::{LanguageTag, ParseLanguageTagError}; pub use platform::{Arch, Os, Platform, PlatformError}; pub use platform_tag::{ParsePlatformTagError, PlatformTag}; pub use tags::{BinaryFormat, IncompatibleTag, TagCompatibility, TagPriority, Tags, TagsError}; mod abi_tag; +mod accelerator; mod language_tag; mod platform; mod platform_tag; diff --git a/crates/uv-python/python/get_interpreter_info.py b/crates/uv-python/python/get_interpreter_info.py index c134a0a6b24f..3f2a3d4ecbf7 100644 --- a/crates/uv-python/python/get_interpreter_info.py +++ b/crates/uv-python/python/get_interpreter_info.py @@ -395,6 +395,47 @@ def get_distutils_scheme(): return get_distutils_scheme() +def _detect_cuda_driver_version(): + """Detect the installed CUDA Driver version. + + Reads from the `UV_CUDA_DRIVER_VERSION` environment variable, if set; otherwise, + queries `nvidia-smi` for the driver version. + """ + driver_version = os.getenv("UV_CUDA_DRIVER_VERSION") + if driver_version is not None: + return driver_version + + import subprocess + + try: + result = subprocess.run( + [ + "nvidia-smi", + "--query-gpu=driver_version", + "--format=csv", + ], + check=True, + capture_output=True, + text=True, + ) + return result.stdout.splitlines()[-1] + except (FileNotFoundError, subprocess.CalledProcessError): + return None + + +def get_accelerator(): + cuda_driver_version = _detect_cuda_driver_version() + if cuda_driver_version is None: + accelerator = None + else: + accelerator = { + "name": "cuda", + "driver_version": cuda_driver_version, + } + + return accelerator + + def get_operating_system_and_architecture(): """Determine the Python interpreter architecture and operating system. @@ -523,6 +564,7 @@ def get_operating_system_and_architecture(): ) ) sys.exit(0) + return {"os": operating_system, "arch": architecture} @@ -541,21 +583,21 @@ def main() -> None: "sys_platform": sys.platform, } - os_and_arch = get_operating_system_and_architecture() + os_arch = get_operating_system_and_architecture() + accelerator = get_accelerator() manylinux_compatible = False - if os_and_arch["os"]["name"] == "manylinux": + if os_arch["os"]["name"] == "manylinux": # noinspection PyProtectedMember from .packaging._manylinux import _get_glibc_version, _is_compatible manylinux_compatible = _is_compatible( - arch=os_and_arch["arch"], version=_get_glibc_version() + arch=os_arch["arch"], version=_get_glibc_version() ) - elif os_and_arch["os"]["name"] == "musllinux": + elif os_arch["os"]["name"] == "musllinux": manylinux_compatible = True - # By default, pip uses sysconfig on Python 3.10+. # But Python distributors can override this decision by setting: # sysconfig._PIP_USE_SYSCONFIG = True / False @@ -613,10 +655,14 @@ def main() -> None: # Prior to the introduction of `sysconfig` patching, python-build-standalone installations would always use # "/install" as the prefix. With `sysconfig` patching, we rewrite the prefix to match the actual installation # location. So in newer versions, we also write a dedicated flag to indicate standalone builds. - "standalone": sysconfig.get_config_var("prefix") == "/install" or bool(sysconfig.get_config_var("PYTHON_BUILD_STANDALONE")), + "standalone": ( + sysconfig.get_config_var("prefix") == "/install" + or bool(sysconfig.get_config_var("PYTHON_BUILD_STANDALONE")) + ), "scheme": get_scheme(use_sysconfig_scheme), "virtualenv": get_virtualenv(), - "platform": os_and_arch, + "platform": os_arch, + "accelerator": accelerator, "manylinux_compatible": manylinux_compatible, # The `t` abiflag for freethreading Python. # https://peps.python.org/pep-0703/#build-configuration-changes diff --git a/crates/uv-python/src/interpreter.rs b/crates/uv-python/src/interpreter.rs index 228bad5054ec..bc4b9c028ee5 100644 --- a/crates/uv-python/src/interpreter.rs +++ b/crates/uv-python/src/interpreter.rs @@ -21,7 +21,7 @@ use uv_fs::{write_atomic_sync, PythonExt, Simplified}; use uv_install_wheel::Layout; use uv_pep440::Version; use uv_pep508::{MarkerEnvironment, StringVersion}; -use uv_platform_tags::Platform; +use uv_platform_tags::{Accelerator, Platform}; use uv_platform_tags::{Tags, TagsError}; use uv_pypi_types::{ResolverMarkerEnvironment, Scheme}; @@ -53,6 +53,7 @@ pub struct Interpreter { target: Option, prefix: Option, pointer_size: PointerSize, + accelerator: Option, gil_disabled: bool, } @@ -76,6 +77,7 @@ impl Interpreter { sys_prefix: info.sys_prefix, sys_base_exec_prefix: info.sys_base_exec_prefix, pointer_size: info.pointer_size, + accelerator: info.accelerator, gil_disabled: info.gil_disabled, sys_base_prefix: info.sys_base_prefix, sys_base_executable: info.sys_base_executable, @@ -172,12 +174,18 @@ impl Interpreter { Ok(base_python) } - /// Returns the path to the Python virtual environment. + /// Returns the [`Platform`] for this Python executable. #[inline] pub fn platform(&self) -> &Platform { &self.platform } + /// Returns the [`Accelerator`] for this Python executable. + #[inline] + pub fn accelerator(&self) -> Option<&Accelerator> { + self.accelerator.as_ref() + } + /// Returns the [`MarkerEnvironment`] for this Python executable. #[inline] pub const fn markers(&self) -> &MarkerEnvironment { @@ -729,6 +737,7 @@ struct InterpreterInfo { stdlib: PathBuf, standalone: bool, pointer_size: PointerSize, + accelerator: Option, gil_disabled: bool, } diff --git a/crates/uv-settings/Cargo.toml b/crates/uv-settings/Cargo.toml index 33729e8066c5..b06d90875cc6 100644 --- a/crates/uv-settings/Cargo.toml +++ b/crates/uv-settings/Cargo.toml @@ -17,18 +17,19 @@ workspace = true [dependencies] uv-cache-info = { workspace = true, features = ["schemars"] } -uv-configuration = { workspace = true, features = ["schemars", "clap"] } +uv-configuration = { workspace = true, features = ["schemars"] } uv-distribution-types = { workspace = true, features = ["schemars"] } uv-fs = { workspace = true } -uv-install-wheel = { workspace = true, features = ["schemars", "clap"] } +uv-install-wheel = { workspace = true, features = ["schemars"] } uv-macros = { workspace = true } uv-normalize = { workspace = true, features = ["schemars"] } uv-options-metadata = { workspace = true } uv-pep508 = { workspace = true } uv-pypi-types = { workspace = true } -uv-python = { workspace = true, features = ["schemars", "clap"] } -uv-resolver = { workspace = true, features = ["schemars", "clap"] } +uv-python = { workspace = true, features = ["schemars"] } +uv-resolver = { workspace = true, features = ["schemars"] } uv-static = { workspace = true } +uv-torch = { workspace = true, features = ["schemars"] } uv-warnings = { workspace = true } clap = { workspace = true } diff --git a/crates/uv-settings/src/combine.rs b/crates/uv-settings/src/combine.rs index 0acb6b6c564f..fa1a03d0710f 100644 --- a/crates/uv-settings/src/combine.rs +++ b/crates/uv-settings/src/combine.rs @@ -3,6 +3,7 @@ use std::path::PathBuf; use url::Url; +use crate::{FilesystemOptions, Options, PipOptions}; use uv_configuration::{ ConfigSettings, IndexStrategy, KeyringProviderType, RequiredVersion, TargetTriple, TrustedPublishing, @@ -12,8 +13,7 @@ use uv_install_wheel::LinkMode; use uv_pypi_types::{SchemaConflicts, SupportedEnvironments}; use uv_python::{PythonDownloads, PythonPreference, PythonVersion}; use uv_resolver::{AnnotationStyle, ExcludeNewer, ForkStrategy, PrereleaseMode, ResolutionMode}; - -use crate::{FilesystemOptions, Options, PipOptions}; +use uv_torch::TorchMode; pub trait Combine { /// Combine two values, preferring the values in `self`. @@ -95,6 +95,7 @@ impl_combine_or!(SchemaConflicts); impl_combine_or!(String); impl_combine_or!(SupportedEnvironments); impl_combine_or!(TargetTriple); +impl_combine_or!(TorchMode); impl_combine_or!(TrustedPublishing); impl_combine_or!(Url); impl_combine_or!(bool); diff --git a/crates/uv-settings/src/settings.rs b/crates/uv-settings/src/settings.rs index b47c5d6af066..d79bfb15e91b 100644 --- a/crates/uv-settings/src/settings.rs +++ b/crates/uv-settings/src/settings.rs @@ -19,6 +19,7 @@ use uv_pypi_types::{SupportedEnvironments, VerbatimParsedUrl}; use uv_python::{PythonDownloads, PythonPreference, PythonVersion}; use uv_resolver::{AnnotationStyle, ExcludeNewer, ForkStrategy, PrereleaseMode, ResolutionMode}; use uv_static::EnvVars; +use uv_torch::TorchMode; /// A `pyproject.toml` with an (optional) `[tool.uv]` section. #[allow(dead_code)] @@ -1533,6 +1534,25 @@ pub struct PipOptions { "# )] pub reinstall_package: Option>, + /// The backend to use when fetching packages in the `PyTorch` ecosystem. + /// + /// When set, uv will ignore the configured index URLs for packages in the `PyTorch` ecosystem, + /// and will instead use the defined backend. + /// + /// For example, when set to `cpu`, uv will use the CPU-only `PyTorch` index; when set to `cu126`, + /// uv will use the `PyTorch` index for CUDA 12.6. + /// + /// The `auto` mode will attempt to detect the appropriate `PyTorch` index based on the currently + /// installed CUDA drivers. + #[option( + default = "null", + value_type = "str", + example = r#" + torch-backend = "auto" + "#, + possible_values = true + )] + pub torch_backend: Option, } impl PipOptions { diff --git a/crates/uv-torch/Cargo.toml b/crates/uv-torch/Cargo.toml new file mode 100644 index 000000000000..869abee89a05 --- /dev/null +++ b/crates/uv-torch/Cargo.toml @@ -0,0 +1,24 @@ +[package] +name = "uv-torch" +version = "0.1.0" +edition.workspace = true +rust-version.workspace = true +homepage.workspace = true +documentation.workspace = true +repository.workspace = true +authors.workspace = true +license.workspace = true + +[dependencies] +uv-distribution-types = { workspace = true } +uv-normalize = { workspace = true } +uv-pep440 = { workspace = true } +uv-platform-tags = { workspace = true } + +clap = { workspace = true, optional = true } +either = { workspace = true } +schemars = { workspace = true, optional = true } +serde = { workspace = true } + +[lints] +workspace = true diff --git a/crates/uv-torch/src/lib.rs b/crates/uv-torch/src/lib.rs new file mode 100644 index 000000000000..c9963135f626 --- /dev/null +++ b/crates/uv-torch/src/lib.rs @@ -0,0 +1,413 @@ +//! `uv-torch` is a library for determining the appropriate `PyTorch` index based on the operating +//! system and CUDA driver version. +//! +//! This library is derived from `light-the-torch` by Philipp Meier, which is available under the +//! following BSD-3 Clause license: +//! +//! ```text +//! BSD 3-Clause License +//! +//! Copyright (c) 2020, Philip Meier +//! All rights reserved. +//! +//! Redistribution and use in source and binary forms, with or without +//! modification, are permitted provided that the following conditions are met: +//! +//! 1. Redistributions of source code must retain the above copyright notice, this +//! list of conditions and the following disclaimer. +//! +//! 2. Redistributions in binary form must reproduce the above copyright notice, +//! this list of conditions and the following disclaimer in the documentation +//! and/or other materials provided with the distribution. +//! +//! 3. Neither the name of the copyright holder nor the names of its +//! contributors may be used to endorse or promote products derived from +//! this software without specific prior written permission. +//! +//! THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +//! AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +//! IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +//! DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +//! FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +//! DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +//! SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +//! CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +//! OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +//! OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +//! ``` +//! +use std::str::FromStr; +use std::sync::LazyLock; + +use either::Either; + +use uv_distribution_types::IndexUrl; +use uv_normalize::PackageName; +use uv_pep440::Version; +use uv_platform_tags::{Accelerator, Os}; + +/// The strategy to use when determining the appropriate `PyTorch` index. +#[derive(Debug, Copy, Clone, Eq, PartialEq, serde::Deserialize, serde::Serialize)] +#[cfg_attr(feature = "clap", derive(clap::ValueEnum))] +#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] +pub enum TorchMode { + /// Select the appropriate `PyTorch` index based on the operating system and CUDA driver version. + Auto, + /// Use the CPU-only `PyTorch` index. + Cpu, + /// Use the `PyTorch` index for CUDA 12.6. + Cu126, + /// Use the `PyTorch` index for CUDA 12.5. + Cu125, + /// Use the `PyTorch` index for CUDA 12.4. + Cu124, + /// Use the `PyTorch` index for CUDA 12.3. + Cu123, + /// Use the `PyTorch` index for CUDA 12.2. + Cu122, + /// Use the `PyTorch` index for CUDA 12.1. + Cu121, + /// Use the `PyTorch` index for CUDA 12.0. + Cu120, + /// Use the `PyTorch` index for CUDA 11.8. + Cu118, + /// Use the `PyTorch` index for CUDA 11.7. + Cu117, + /// Use the `PyTorch` index for CUDA 11.6. + Cu116, + /// Use the `PyTorch` index for CUDA 11.5. + Cu115, + /// Use the `PyTorch` index for CUDA 11.4. + Cu114, + /// Use the `PyTorch` index for CUDA 11.3. + Cu113, + /// Use the `PyTorch` index for CUDA 11.2. + Cu112, + /// Use the `PyTorch` index for CUDA 11.1. + Cu111, + /// Use the `PyTorch` index for CUDA 11.0. + Cu110, + /// Use the `PyTorch` index for CUDA 10.2. + Cu102, + /// Use the `PyTorch` index for CUDA 10.1. + Cu101, + /// Use the `PyTorch` index for CUDA 10.0. + Cu100, + /// Use the `PyTorch` index for CUDA 9.2. + Cu92, + /// Use the `PyTorch` index for CUDA 9.1. + Cu91, + /// Use the `PyTorch` index for CUDA 9.0. + Cu90, + /// Use the `PyTorch` index for CUDA 8.0. + Cu80, +} + +/// The strategy to use when determining the appropriate `PyTorch` index. +#[derive(Debug, Clone, Eq, PartialEq)] +pub enum TorchStrategy { + /// Select the appropriate `PyTorch` index based on the operating system and CUDA driver version. + Auto { os: Os, driver_version: Version }, + /// Use the specified `PyTorch` index. + Backend(TorchBackend), +} + +impl TorchStrategy { + /// Determine the [`TorchStrategy`] from the given [`TorchMode`], [`Os`], and [`Accelerator`]. + pub fn from_mode(mode: TorchMode, os: &Os, accelerator: Option<&Accelerator>) -> Self { + match mode { + TorchMode::Auto => { + if let Some(Accelerator::Cuda { driver_version }) = accelerator { + Self::Auto { + os: os.clone(), + driver_version: driver_version.clone(), + } + } else { + Self::Backend(TorchBackend::Cpu) + } + } + TorchMode::Cpu => Self::Backend(TorchBackend::Cpu), + TorchMode::Cu126 => Self::Backend(TorchBackend::Cu126), + TorchMode::Cu125 => Self::Backend(TorchBackend::Cu125), + TorchMode::Cu124 => Self::Backend(TorchBackend::Cu124), + TorchMode::Cu123 => Self::Backend(TorchBackend::Cu123), + TorchMode::Cu122 => Self::Backend(TorchBackend::Cu122), + TorchMode::Cu121 => Self::Backend(TorchBackend::Cu121), + TorchMode::Cu120 => Self::Backend(TorchBackend::Cu120), + TorchMode::Cu118 => Self::Backend(TorchBackend::Cu118), + TorchMode::Cu117 => Self::Backend(TorchBackend::Cu117), + TorchMode::Cu116 => Self::Backend(TorchBackend::Cu116), + TorchMode::Cu115 => Self::Backend(TorchBackend::Cu115), + TorchMode::Cu114 => Self::Backend(TorchBackend::Cu114), + TorchMode::Cu113 => Self::Backend(TorchBackend::Cu113), + TorchMode::Cu112 => Self::Backend(TorchBackend::Cu112), + TorchMode::Cu111 => Self::Backend(TorchBackend::Cu111), + TorchMode::Cu110 => Self::Backend(TorchBackend::Cu110), + TorchMode::Cu102 => Self::Backend(TorchBackend::Cu102), + TorchMode::Cu101 => Self::Backend(TorchBackend::Cu101), + TorchMode::Cu100 => Self::Backend(TorchBackend::Cu100), + TorchMode::Cu92 => Self::Backend(TorchBackend::Cu92), + TorchMode::Cu91 => Self::Backend(TorchBackend::Cu91), + TorchMode::Cu90 => Self::Backend(TorchBackend::Cu90), + TorchMode::Cu80 => Self::Backend(TorchBackend::Cu80), + } + } + + /// Return the appropriate index URLs for the given [`TorchStrategy`] and [`PackageName`]. + pub fn index_urls( + &self, + package_name: &PackageName, + ) -> Option> { + if !matches!( + package_name.as_str(), + "torch" + | "torch-model-archiver" + | "torch-tb-profiler" + | "torcharrow" + | "torchaudio" + | "torchcsprng" + | "torchdata" + | "torchdistx" + | "torchserve" + | "torchtext" + | "torchvision" + | "pytorch-triton" + ) { + return None; + } + + match self { + TorchStrategy::Auto { os, driver_version } => { + // If this is a GPU-enabled package, and CUDA drivers are installed, use PyTorch's CUDA + // indexes. + // + // See: https://github.com/pmeier/light-the-torch/blob/33397cbe45d07b51ad8ee76b004571a4c236e37f/light_the_torch/_patch.py#L36-L49 + match os { + Os::Manylinux { .. } | Os::Musllinux { .. } => { + Some(Either::Left(Either::Left( + LINUX_DRIVERS + .iter() + .filter_map(move |(backend, version)| { + if driver_version >= version { + Some(backend.index_url()) + } else { + None + } + }) + .chain(std::iter::once(TorchBackend::Cpu.index_url())), + ))) + } + Os::Windows => Some(Either::Left(Either::Right( + WINDOWS_CUDA_VERSIONS + .iter() + .filter_map(move |(backend, version)| { + if driver_version >= version { + Some(backend.index_url()) + } else { + None + } + }) + .chain(std::iter::once(TorchBackend::Cpu.index_url())), + ))), + Os::Macos { .. } + | Os::FreeBsd { .. } + | Os::NetBsd { .. } + | Os::OpenBsd { .. } + | Os::Dragonfly { .. } + | Os::Illumos { .. } + | Os::Haiku { .. } + | Os::Android { .. } => Some(Either::Right(std::iter::once( + TorchBackend::Cpu.index_url(), + ))), + } + } + TorchStrategy::Backend(backend) => { + Some(Either::Right(std::iter::once(backend.index_url()))) + } + } + } +} + +/// The available backends for `PyTorch`. +#[derive(Debug, Clone, Eq, PartialEq)] +pub enum TorchBackend { + Cpu, + Cu126, + Cu125, + Cu124, + Cu123, + Cu122, + Cu121, + Cu120, + Cu118, + Cu117, + Cu116, + Cu115, + Cu114, + Cu113, + Cu112, + Cu111, + Cu110, + Cu102, + Cu101, + Cu100, + Cu92, + Cu91, + Cu90, + Cu80, +} + +impl TorchBackend { + /// Return the appropriate index URL for the given [`TorchBackend`]. + fn index_url(&self) -> &'static IndexUrl { + match self { + Self::Cpu => &CPU_INDEX_URL, + Self::Cu126 => &CU126_INDEX_URL, + Self::Cu125 => &CU125_INDEX_URL, + Self::Cu124 => &CU124_INDEX_URL, + Self::Cu123 => &CU123_INDEX_URL, + Self::Cu122 => &CU122_INDEX_URL, + Self::Cu121 => &CU121_INDEX_URL, + Self::Cu120 => &CU120_INDEX_URL, + Self::Cu118 => &CU118_INDEX_URL, + Self::Cu117 => &CU117_INDEX_URL, + Self::Cu116 => &CU116_INDEX_URL, + Self::Cu115 => &CU115_INDEX_URL, + Self::Cu114 => &CU114_INDEX_URL, + Self::Cu113 => &CU113_INDEX_URL, + Self::Cu112 => &CU112_INDEX_URL, + Self::Cu111 => &CU111_INDEX_URL, + Self::Cu110 => &CU110_INDEX_URL, + Self::Cu102 => &CU102_INDEX_URL, + Self::Cu101 => &CU101_INDEX_URL, + Self::Cu100 => &CU100_INDEX_URL, + Self::Cu92 => &CU92_INDEX_URL, + Self::Cu91 => &CU91_INDEX_URL, + Self::Cu90 => &CU90_INDEX_URL, + Self::Cu80 => &CU80_INDEX_URL, + } + } +} + +/// Linux CUDA driver versions and the corresponding CUDA versions. +/// +/// See: +static LINUX_DRIVERS: LazyLock<[(TorchBackend, Version); 23]> = LazyLock::new(|| { + [ + // Table 2 from + // https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html + (TorchBackend::Cu126, Version::new([525, 60, 13])), + (TorchBackend::Cu125, Version::new([525, 60, 13])), + (TorchBackend::Cu124, Version::new([525, 60, 13])), + (TorchBackend::Cu123, Version::new([525, 60, 13])), + (TorchBackend::Cu122, Version::new([525, 60, 13])), + (TorchBackend::Cu121, Version::new([525, 60, 13])), + (TorchBackend::Cu120, Version::new([525, 60, 13])), + // Table 2 from + // https://docs.nvidia.com/cuda/archive/11.8.0/cuda-toolkit-release-notes/index.html + (TorchBackend::Cu118, Version::new([450, 80, 2])), + (TorchBackend::Cu117, Version::new([450, 80, 2])), + (TorchBackend::Cu116, Version::new([450, 80, 2])), + (TorchBackend::Cu115, Version::new([450, 80, 2])), + (TorchBackend::Cu114, Version::new([450, 80, 2])), + (TorchBackend::Cu113, Version::new([450, 80, 2])), + (TorchBackend::Cu112, Version::new([450, 80, 2])), + (TorchBackend::Cu111, Version::new([450, 80, 2])), + (TorchBackend::Cu110, Version::new([450, 36, 6])), + // Table 1 from + // https://docs.nvidia.com/cuda/archive/10.2/cuda-toolkit-release-notes/index.html + (TorchBackend::Cu102, Version::new([440, 33])), + (TorchBackend::Cu101, Version::new([418, 39])), + (TorchBackend::Cu100, Version::new([410, 48])), + (TorchBackend::Cu92, Version::new([396, 26])), + (TorchBackend::Cu91, Version::new([390, 46])), + (TorchBackend::Cu90, Version::new([384, 81])), + (TorchBackend::Cu80, Version::new([375, 26])), + ] +}); + +/// Windows CUDA driver versions and the corresponding CUDA versions. +/// +/// See: +static WINDOWS_CUDA_VERSIONS: LazyLock<[(TorchBackend, Version); 23]> = LazyLock::new(|| { + [ + // Table 2 from + // https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html + (TorchBackend::Cu126, Version::new([528, 33])), + (TorchBackend::Cu125, Version::new([528, 33])), + (TorchBackend::Cu124, Version::new([528, 33])), + (TorchBackend::Cu123, Version::new([528, 33])), + (TorchBackend::Cu122, Version::new([528, 33])), + (TorchBackend::Cu121, Version::new([528, 33])), + (TorchBackend::Cu120, Version::new([528, 33])), + // Table 2 from + // https://docs.nvidia.com/cuda/archive/11.8.0/cuda-toolkit-release-notes/index.html + (TorchBackend::Cu118, Version::new([452, 39])), + (TorchBackend::Cu117, Version::new([452, 39])), + (TorchBackend::Cu116, Version::new([452, 39])), + (TorchBackend::Cu115, Version::new([452, 39])), + (TorchBackend::Cu114, Version::new([452, 39])), + (TorchBackend::Cu113, Version::new([452, 39])), + (TorchBackend::Cu112, Version::new([452, 39])), + (TorchBackend::Cu111, Version::new([452, 39])), + (TorchBackend::Cu110, Version::new([451, 22])), + // Table 1 from + // https://docs.nvidia.com/cuda/archive/10.2/cuda-toolkit-release-notes/index.html + (TorchBackend::Cu102, Version::new([441, 22])), + (TorchBackend::Cu101, Version::new([418, 96])), + (TorchBackend::Cu100, Version::new([411, 31])), + (TorchBackend::Cu92, Version::new([398, 26])), + (TorchBackend::Cu91, Version::new([391, 29])), + (TorchBackend::Cu90, Version::new([385, 54])), + (TorchBackend::Cu80, Version::new([376, 51])), + ] +}); + +static CPU_INDEX_URL: LazyLock = + LazyLock::new(|| IndexUrl::from_str("https://download.pytorch.org/whl/cpu").unwrap()); +static CU126_INDEX_URL: LazyLock = + LazyLock::new(|| IndexUrl::from_str("https://download.pytorch.org/whl/cu126").unwrap()); +static CU125_INDEX_URL: LazyLock = + LazyLock::new(|| IndexUrl::from_str("https://download.pytorch.org/whl/cu125").unwrap()); +static CU124_INDEX_URL: LazyLock = + LazyLock::new(|| IndexUrl::from_str("https://download.pytorch.org/whl/cu124").unwrap()); +static CU123_INDEX_URL: LazyLock = + LazyLock::new(|| IndexUrl::from_str("https://download.pytorch.org/whl/cu123").unwrap()); +static CU122_INDEX_URL: LazyLock = + LazyLock::new(|| IndexUrl::from_str("https://download.pytorch.org/whl/cu122").unwrap()); +static CU121_INDEX_URL: LazyLock = + LazyLock::new(|| IndexUrl::from_str("https://download.pytorch.org/whl/cu121").unwrap()); +static CU120_INDEX_URL: LazyLock = + LazyLock::new(|| IndexUrl::from_str("https://download.pytorch.org/whl/cu120").unwrap()); +static CU118_INDEX_URL: LazyLock = + LazyLock::new(|| IndexUrl::from_str("https://download.pytorch.org/whl/cu118").unwrap()); +static CU117_INDEX_URL: LazyLock = + LazyLock::new(|| IndexUrl::from_str("https://download.pytorch.org/whl/cu117").unwrap()); +static CU116_INDEX_URL: LazyLock = + LazyLock::new(|| IndexUrl::from_str("https://download.pytorch.org/whl/cu116").unwrap()); +static CU115_INDEX_URL: LazyLock = + LazyLock::new(|| IndexUrl::from_str("https://download.pytorch.org/whl/cu115").unwrap()); +static CU114_INDEX_URL: LazyLock = + LazyLock::new(|| IndexUrl::from_str("https://download.pytorch.org/whl/cu114").unwrap()); +static CU113_INDEX_URL: LazyLock = + LazyLock::new(|| IndexUrl::from_str("https://download.pytorch.org/whl/cu113").unwrap()); +static CU112_INDEX_URL: LazyLock = + LazyLock::new(|| IndexUrl::from_str("https://download.pytorch.org/whl/cu112").unwrap()); +static CU111_INDEX_URL: LazyLock = + LazyLock::new(|| IndexUrl::from_str("https://download.pytorch.org/whl/cu111").unwrap()); +static CU110_INDEX_URL: LazyLock = + LazyLock::new(|| IndexUrl::from_str("https://download.pytorch.org/whl/cu110").unwrap()); +static CU102_INDEX_URL: LazyLock = + LazyLock::new(|| IndexUrl::from_str("https://download.pytorch.org/whl/cu102").unwrap()); +static CU101_INDEX_URL: LazyLock = + LazyLock::new(|| IndexUrl::from_str("https://download.pytorch.org/whl/cu101").unwrap()); +static CU100_INDEX_URL: LazyLock = + LazyLock::new(|| IndexUrl::from_str("https://download.pytorch.org/whl/cu100").unwrap()); +static CU92_INDEX_URL: LazyLock = + LazyLock::new(|| IndexUrl::from_str("https://download.pytorch.org/whl/cu92").unwrap()); +static CU91_INDEX_URL: LazyLock = + LazyLock::new(|| IndexUrl::from_str("https://download.pytorch.org/whl/cu91").unwrap()); +static CU90_INDEX_URL: LazyLock = + LazyLock::new(|| IndexUrl::from_str("https://download.pytorch.org/whl/cu90").unwrap()); +static CU80_INDEX_URL: LazyLock = + LazyLock::new(|| IndexUrl::from_str("https://download.pytorch.org/whl/cu80").unwrap()); diff --git a/crates/uv/Cargo.toml b/crates/uv/Cargo.toml index 4cb1469dbfba..0a2d9fa503f7 100644 --- a/crates/uv/Cargo.toml +++ b/crates/uv/Cargo.toml @@ -49,10 +49,11 @@ uv-settings = { workspace = true, features = ["schemars"] } uv-shell = { workspace = true } uv-static = { workspace = true } uv-tool = { workspace = true } +uv-torch = { workspace = true } uv-trampoline-builder = { workspace = true } uv-types = { workspace = true } -uv-virtualenv = { workspace = true } uv-version = { workspace = true } +uv-virtualenv = { workspace = true } uv-warnings = { workspace = true } uv-workspace = { workspace = true } diff --git a/crates/uv/src/commands/pip/compile.rs b/crates/uv/src/commands/pip/compile.rs index 512769bfe642..f0884fc90cd7 100644 --- a/crates/uv/src/commands/pip/compile.rs +++ b/crates/uv/src/commands/pip/compile.rs @@ -38,6 +38,7 @@ use uv_resolver::{ InMemoryIndex, OptionsBuilder, PrereleaseMode, PythonRequirement, RequiresPython, ResolutionMode, ResolverEnvironment, }; +use uv_torch::{TorchMode, TorchStrategy}; use uv_types::{BuildIsolation, EmptyInstalledPackages, HashStrategy}; use uv_warnings::warn_user; @@ -80,6 +81,7 @@ pub(crate) async fn pip_compile( include_index_annotation: bool, index_locations: IndexLocations, index_strategy: IndexStrategy, + torch_backend: Option, dependency_metadata: DependencyMetadata, keyring_provider: KeyringProviderType, network_settings: &NetworkSettings, @@ -334,11 +336,29 @@ pub(crate) async fn pip_compile( } } + // Determine the PyTorch backend. + let torch_backend = torch_backend.map(|mode| { + if preview.is_disabled() { + warn_user!("The `--torch-backend` setting is experimental and may change without warning. Pass `--preview` to disable this warning."); + } + + TorchStrategy::from_mode( + mode, + python_platform + .map(TargetTriple::platform) + .as_ref() + .unwrap_or(interpreter.platform()) + .os(), + interpreter.accelerator(), + ) + }); + // Initialize the registry client. let client = RegistryClientBuilder::try_from(client_builder)? .cache(cache.clone()) .index_urls(index_locations.index_urls()) .index_strategy(index_strategy) + .torch_backend(torch_backend) .markers(interpreter.markers()) .platform(interpreter.platform()) .build(); diff --git a/crates/uv/src/commands/pip/install.rs b/crates/uv/src/commands/pip/install.rs index f9a599cae9fb..b386150f3330 100644 --- a/crates/uv/src/commands/pip/install.rs +++ b/crates/uv/src/commands/pip/install.rs @@ -6,6 +6,13 @@ use itertools::Itertools; use owo_colors::OwoColorize; use tracing::{debug, enabled, Level}; +use crate::commands::pip::loggers::{DefaultInstallLogger, DefaultResolveLogger, InstallLogger}; +use crate::commands::pip::operations::Modifications; +use crate::commands::pip::operations::{report_interpreter, report_target_environment}; +use crate::commands::pip::{operations, resolution_markers, resolution_tags}; +use crate::commands::{diagnostics, ExitStatus}; +use crate::printer::Printer; +use crate::settings::NetworkSettings; use uv_cache::Cache; use uv_client::{BaseClientBuilder, FlatIndexClient, RegistryClientBuilder}; use uv_configuration::{ @@ -33,15 +40,9 @@ use uv_resolver::{ DependencyMode, ExcludeNewer, FlatIndex, OptionsBuilder, PrereleaseMode, PythonRequirement, ResolutionMode, ResolverEnvironment, }; +use uv_torch::{TorchMode, TorchStrategy}; use uv_types::{BuildIsolation, HashStrategy}; - -use crate::commands::pip::loggers::{DefaultInstallLogger, DefaultResolveLogger, InstallLogger}; -use crate::commands::pip::operations::Modifications; -use crate::commands::pip::operations::{report_interpreter, report_target_environment}; -use crate::commands::pip::{operations, resolution_markers, resolution_tags}; -use crate::commands::{diagnostics, ExitStatus}; -use crate::printer::Printer; -use crate::settings::NetworkSettings; +use uv_warnings::warn_user; /// Install packages into the current environment. #[allow(clippy::fn_params_excessive_bools)] @@ -61,6 +62,7 @@ pub(crate) async fn pip_install( upgrade: Upgrade, index_locations: IndexLocations, index_strategy: IndexStrategy, + torch_backend: Option, dependency_metadata: DependencyMetadata, keyring_provider: KeyringProviderType, network_settings: &NetworkSettings, @@ -329,11 +331,29 @@ pub(crate) async fn pip_install( } } + // Determine the PyTorch backend. + let torch_backend = torch_backend.map(|mode| { + if preview.is_disabled() { + warn_user!("The `--torch-backend` setting is experimental and may change without warning. Pass `--preview` to disable this warning."); + } + + TorchStrategy::from_mode( + mode, + python_platform + .map(TargetTriple::platform) + .as_ref() + .unwrap_or(interpreter.platform()) + .os(), + interpreter.accelerator(), + ) + }); + // Initialize the registry client. let client = RegistryClientBuilder::try_from(client_builder)? .cache(cache.clone()) .index_urls(index_locations.index_urls()) .index_strategy(index_strategy) + .torch_backend(torch_backend) .markers(interpreter.markers()) .platform(interpreter.platform()) .build(); diff --git a/crates/uv/src/commands/pip/sync.rs b/crates/uv/src/commands/pip/sync.rs index 8c12252649ab..10c596f9926b 100644 --- a/crates/uv/src/commands/pip/sync.rs +++ b/crates/uv/src/commands/pip/sync.rs @@ -6,6 +6,13 @@ use anyhow::Result; use owo_colors::OwoColorize; use tracing::debug; +use crate::commands::pip::loggers::{DefaultInstallLogger, DefaultResolveLogger}; +use crate::commands::pip::operations::Modifications; +use crate::commands::pip::operations::{report_interpreter, report_target_environment}; +use crate::commands::pip::{operations, resolution_markers, resolution_tags}; +use crate::commands::{diagnostics, ExitStatus}; +use crate::printer::Printer; +use crate::settings::NetworkSettings; use uv_cache::Cache; use uv_client::{BaseClientBuilder, FlatIndexClient, RegistryClientBuilder}; use uv_configuration::{ @@ -30,15 +37,9 @@ use uv_resolver::{ DependencyMode, ExcludeNewer, FlatIndex, OptionsBuilder, PrereleaseMode, PythonRequirement, ResolutionMode, ResolverEnvironment, }; +use uv_torch::{TorchMode, TorchStrategy}; use uv_types::{BuildIsolation, HashStrategy}; - -use crate::commands::pip::loggers::{DefaultInstallLogger, DefaultResolveLogger}; -use crate::commands::pip::operations::Modifications; -use crate::commands::pip::operations::{report_interpreter, report_target_environment}; -use crate::commands::pip::{operations, resolution_markers, resolution_tags}; -use crate::commands::{diagnostics, ExitStatus}; -use crate::printer::Printer; -use crate::settings::NetworkSettings; +use uv_warnings::warn_user; /// Install a set of locked requirements into the current Python environment. #[allow(clippy::fn_params_excessive_bools)] @@ -52,6 +53,7 @@ pub(crate) async fn pip_sync( hash_checking: Option, index_locations: IndexLocations, index_strategy: IndexStrategy, + torch_backend: Option, dependency_metadata: DependencyMetadata, keyring_provider: KeyringProviderType, network_settings: &NetworkSettings, @@ -259,11 +261,29 @@ pub(crate) async fn pip_sync( } } + // Determine the PyTorch backend. + let torch_backend = torch_backend.map(|mode| { + if preview.is_disabled() { + warn_user!("The `--torch-backend` setting is experimental and may change without warning. Pass `--preview` to disable this warning."); + } + + TorchStrategy::from_mode( + mode, + python_platform + .map(TargetTriple::platform) + .as_ref() + .unwrap_or(interpreter.platform()) + .os(), + interpreter.accelerator(), + ) + }); + // Initialize the registry client. let client = RegistryClientBuilder::try_from(client_builder)? .cache(cache.clone()) .index_urls(index_locations.index_urls()) .index_strategy(index_strategy) + .torch_backend(torch_backend) .markers(interpreter.markers()) .platform(interpreter.platform()) .build(); diff --git a/crates/uv/src/lib.rs b/crates/uv/src/lib.rs index 363136f51afa..fbe89a1b3748 100644 --- a/crates/uv/src/lib.rs +++ b/crates/uv/src/lib.rs @@ -421,6 +421,7 @@ async fn run(mut cli: Cli) -> Result { args.settings.emit_index_annotation, args.settings.index_locations, args.settings.index_strategy, + args.settings.torch_backend, args.settings.dependency_metadata, args.settings.keyring_provider, &globals.network_settings, @@ -488,6 +489,7 @@ async fn run(mut cli: Cli) -> Result { args.settings.hash_checking, args.settings.index_locations, args.settings.index_strategy, + args.settings.torch_backend, args.settings.dependency_metadata, args.settings.keyring_provider, &globals.network_settings, @@ -576,6 +578,7 @@ async fn run(mut cli: Cli) -> Result { args.settings.upgrade, args.settings.index_locations, args.settings.index_strategy, + args.settings.torch_backend, args.settings.dependency_metadata, args.settings.keyring_provider, &globals.network_settings, diff --git a/crates/uv/src/settings.rs b/crates/uv/src/settings.rs index ae9c0e5f9142..725b6f8e8079 100644 --- a/crates/uv/src/settings.rs +++ b/crates/uv/src/settings.rs @@ -42,6 +42,7 @@ use uv_settings::{ ResolverInstallerOptions, ResolverOptions, }; use uv_static::EnvVars; +use uv_torch::TorchMode; use uv_warnings::warn_user_once; use uv_workspace::pyproject::DependencyType; @@ -1635,6 +1636,7 @@ impl PipCompileSettings { no_emit_marker_expression, emit_index_annotation, no_emit_index_annotation, + torch_backend, compat_args: _, } = args; @@ -1732,6 +1734,7 @@ impl PipCompileSettings { emit_marker_expression: flag(emit_marker_expression, no_emit_marker_expression), emit_index_annotation: flag(emit_index_annotation, no_emit_index_annotation), annotation_style, + torch_backend, ..PipOptions::from(resolver) }, filesystem, @@ -1783,6 +1786,7 @@ impl PipSyncSettings { strict, no_strict, dry_run, + torch_backend, compat_args: _, } = *args; @@ -1817,6 +1821,7 @@ impl PipSyncSettings { python_version, python_platform, strict: flag(strict, no_strict), + torch_backend, ..PipOptions::from(installer) }, filesystem, @@ -1883,6 +1888,7 @@ impl PipInstallSettings { strict, no_strict, dry_run, + torch_backend, compat_args: _, } = args; @@ -1972,6 +1978,7 @@ impl PipInstallSettings { python_platform, require_hashes: flag(require_hashes, no_require_hashes), verify_hashes: flag(verify_hashes, no_verify_hashes), + torch_backend, ..PipOptions::from(installer) }, filesystem, @@ -2681,6 +2688,7 @@ pub(crate) struct PipSettings { pub(crate) prefix: Option, pub(crate) index_strategy: IndexStrategy, pub(crate) keyring_provider: KeyringProviderType, + pub(crate) torch_backend: Option, pub(crate) no_build_isolation: bool, pub(crate) no_build_isolation_package: Vec, pub(crate) build_options: BuildOptions, @@ -2742,6 +2750,7 @@ impl PipSettings { no_index, find_links, index_strategy, + torch_backend, keyring_provider, no_build, no_binary, @@ -2939,6 +2948,7 @@ impl PipSettings { .config_settings .combine(config_settings) .unwrap_or_default(), + torch_backend: args.torch_backend.combine(torch_backend), python_version: args.python_version.combine(python_version), python_platform: args.python_platform.combine(python_platform), universal: args.universal.combine(universal).unwrap_or_default(), diff --git a/crates/uv/tests/it/lock.rs b/crates/uv/tests/it/lock.rs index 12af8d57e43b..1d9946caba2c 100644 --- a/crates/uv/tests/it/lock.rs +++ b/crates/uv/tests/it/lock.rs @@ -15223,9 +15223,7 @@ fn lock_explicit_default_index() -> Result<()> { DEBUG Searching for a compatible version of project @ file://[TEMP_DIR]/ (<0.1.0 | >0.1.0) DEBUG No compatible version found for: project × No solution found when resolving dependencies: - ╰─▶ Because anyio was not found in the provided package locations and your project depends on anyio, we can conclude that your project's requirements are unsatisfiable. - - hint: Packages were unavailable because index lookups were disabled and no additional package locations were provided (try: `--find-links `) + ╰─▶ Because anyio was not found in the package registry and your project depends on anyio, we can conclude that your project's requirements are unsatisfiable. "###); let lock = fs_err::read_to_string(context.temp_dir.join("uv.lock")).unwrap(); diff --git a/crates/uv/tests/it/show_settings.rs b/crates/uv/tests/it/show_settings.rs index 986dd60e0ecf..102a9ec433f8 100644 --- a/crates/uv/tests/it/show_settings.rs +++ b/crates/uv/tests/it/show_settings.rs @@ -51,7 +51,7 @@ fn resolve_uv_toml() -> anyhow::Result<()> { // Resolution should use the lowest direct version, and generate hashes. uv_snapshot!(context.filters(), add_shared_args(context.pip_compile(), context.temp_dir.path()) .arg("--show-settings") - .arg("requirements.in"), @r#" + .arg("requirements.in"), @r###" success: true exit_code: 0 ----- stdout ----- @@ -170,6 +170,7 @@ fn resolve_uv_toml() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -218,14 +219,14 @@ fn resolve_uv_toml() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); // Resolution should use the highest version, and generate hashes. uv_snapshot!(context.filters(), add_shared_args(context.pip_compile(), context.temp_dir.path()) .arg("--show-settings") .arg("requirements.in") - .arg("--resolution=highest"), @r#" + .arg("--resolution=highest"), @r###" success: true exit_code: 0 ----- stdout ----- @@ -344,6 +345,7 @@ fn resolve_uv_toml() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -392,7 +394,7 @@ fn resolve_uv_toml() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); // Resolution should use the highest version, and omit hashes. @@ -400,7 +402,7 @@ fn resolve_uv_toml() -> anyhow::Result<()> { .arg("--show-settings") .arg("requirements.in") .arg("--resolution=highest") - .arg("--no-generate-hashes"), @r#" + .arg("--no-generate-hashes"), @r###" success: true exit_code: 0 ----- stdout ----- @@ -519,6 +521,7 @@ fn resolve_uv_toml() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -567,7 +570,7 @@ fn resolve_uv_toml() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); Ok(()) @@ -607,7 +610,7 @@ fn resolve_pyproject_toml() -> anyhow::Result<()> { // Resolution should use the lowest direct version, and generate hashes. uv_snapshot!(context.filters(), add_shared_args(context.pip_compile(), context.temp_dir.path()) .arg("--show-settings") - .arg("requirements.in"), @r#" + .arg("requirements.in"), @r###" success: true exit_code: 0 ----- stdout ----- @@ -726,6 +729,7 @@ fn resolve_pyproject_toml() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -774,7 +778,7 @@ fn resolve_pyproject_toml() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); // Remove the `uv.toml` file. @@ -783,7 +787,7 @@ fn resolve_pyproject_toml() -> anyhow::Result<()> { // Resolution should use the highest version, and omit hashes. uv_snapshot!(context.filters(), add_shared_args(context.pip_compile(), context.temp_dir.path()) .arg("--show-settings") - .arg("requirements.in"), @r#" + .arg("requirements.in"), @r###" success: true exit_code: 0 ----- stdout ----- @@ -872,6 +876,7 @@ fn resolve_pyproject_toml() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -920,7 +925,7 @@ fn resolve_pyproject_toml() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); // Add configuration to the `pyproject.toml` file. @@ -938,7 +943,7 @@ fn resolve_pyproject_toml() -> anyhow::Result<()> { // Resolution should use the lowest direct version, and generate hashes. uv_snapshot!(context.filters(), add_shared_args(context.pip_compile(), context.temp_dir.path()) .arg("--show-settings") - .arg("requirements.in"), @r#" + .arg("requirements.in"), @r###" success: true exit_code: 0 ----- stdout ----- @@ -1057,6 +1062,7 @@ fn resolve_pyproject_toml() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -1105,7 +1111,7 @@ fn resolve_pyproject_toml() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); Ok(()) @@ -1137,7 +1143,7 @@ fn resolve_index_url() -> anyhow::Result<()> { uv_snapshot!(context.filters(), add_shared_args(context.pip_compile(), context.temp_dir.path()) .arg("--show-settings") - .arg("requirements.in"), @r#" + .arg("requirements.in"), @r###" success: true exit_code: 0 ----- stdout ----- @@ -1285,6 +1291,7 @@ fn resolve_index_url() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -1333,7 +1340,7 @@ fn resolve_index_url() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); // Providing an additional index URL on the command-line should be merged with the @@ -1342,7 +1349,7 @@ fn resolve_index_url() -> anyhow::Result<()> { .arg("--show-settings") .arg("requirements.in") .arg("--extra-index-url") - .arg("https://test.pypi.org/simple"), @r#" + .arg("https://test.pypi.org/simple"), @r###" success: true exit_code: 0 ----- stdout ----- @@ -1521,6 +1528,7 @@ fn resolve_index_url() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -1569,7 +1577,7 @@ fn resolve_index_url() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); Ok(()) @@ -1601,7 +1609,7 @@ fn resolve_find_links() -> anyhow::Result<()> { uv_snapshot!(context.filters(), add_shared_args(context.pip_compile(), context.temp_dir.path()) .arg("--show-settings") - .arg("requirements.in"), @r#" + .arg("requirements.in"), @r###" success: true exit_code: 0 ----- stdout ----- @@ -1720,6 +1728,7 @@ fn resolve_find_links() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -1768,7 +1777,7 @@ fn resolve_find_links() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); Ok(()) @@ -1799,7 +1808,7 @@ fn resolve_top_level() -> anyhow::Result<()> { uv_snapshot!(context.filters(), add_shared_args(context.pip_compile(), context.temp_dir.path()) .arg("--show-settings") - .arg("requirements.in"), @r#" + .arg("requirements.in"), @r###" success: true exit_code: 0 ----- stdout ----- @@ -1888,6 +1897,7 @@ fn resolve_top_level() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -1936,7 +1946,7 @@ fn resolve_top_level() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); // Write out to both the top-level (`tool.uv`) and the pip section (`tool.uv.pip`). The @@ -1960,7 +1970,7 @@ fn resolve_top_level() -> anyhow::Result<()> { uv_snapshot!(context.filters(), add_shared_args(context.pip_compile(), context.temp_dir.path()) .arg("--show-settings") - .arg("requirements.in"), @r#" + .arg("requirements.in"), @r###" success: true exit_code: 0 ----- stdout ----- @@ -2108,6 +2118,7 @@ fn resolve_top_level() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -2156,14 +2167,14 @@ fn resolve_top_level() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); // But the command-line should take precedence over both. uv_snapshot!(context.filters(), add_shared_args(context.pip_compile(), context.temp_dir.path()) .arg("--show-settings") .arg("requirements.in") - .arg("--resolution=lowest-direct"), @r#" + .arg("--resolution=lowest-direct"), @r###" success: true exit_code: 0 ----- stdout ----- @@ -2311,6 +2322,7 @@ fn resolve_top_level() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -2359,7 +2371,7 @@ fn resolve_top_level() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); Ok(()) @@ -2390,7 +2402,7 @@ fn resolve_user_configuration() -> anyhow::Result<()> { uv_snapshot!(context.filters(), add_shared_args(context.pip_compile(), context.temp_dir.path()) .arg("--show-settings") .arg("requirements.in") - .env(EnvVars::XDG_CONFIG_HOME, xdg.path()), @r#" + .env(EnvVars::XDG_CONFIG_HOME, xdg.path()), @r###" success: true exit_code: 0 ----- stdout ----- @@ -2479,6 +2491,7 @@ fn resolve_user_configuration() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -2527,7 +2540,7 @@ fn resolve_user_configuration() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); // Add a local configuration to generate hashes. @@ -2541,7 +2554,7 @@ fn resolve_user_configuration() -> anyhow::Result<()> { uv_snapshot!(context.filters(), add_shared_args(context.pip_compile(), context.temp_dir.path()) .arg("--show-settings") .arg("requirements.in") - .env(EnvVars::XDG_CONFIG_HOME, xdg.path()), @r#" + .env(EnvVars::XDG_CONFIG_HOME, xdg.path()), @r###" success: true exit_code: 0 ----- stdout ----- @@ -2630,6 +2643,7 @@ fn resolve_user_configuration() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -2678,7 +2692,7 @@ fn resolve_user_configuration() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); // Add a local configuration to override the user configuration. @@ -2692,7 +2706,7 @@ fn resolve_user_configuration() -> anyhow::Result<()> { uv_snapshot!(context.filters(), add_shared_args(context.pip_compile(), context.temp_dir.path()) .arg("--show-settings") .arg("requirements.in") - .env(EnvVars::XDG_CONFIG_HOME, xdg.path()), @r#" + .env(EnvVars::XDG_CONFIG_HOME, xdg.path()), @r###" success: true exit_code: 0 ----- stdout ----- @@ -2781,6 +2795,7 @@ fn resolve_user_configuration() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -2829,7 +2844,7 @@ fn resolve_user_configuration() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); // However, the user-level `tool.uv.pip` settings override the project-level `tool.uv` settings. @@ -2845,7 +2860,7 @@ fn resolve_user_configuration() -> anyhow::Result<()> { uv_snapshot!(context.filters(), add_shared_args(context.pip_compile(), context.temp_dir.path()) .arg("--show-settings") .arg("requirements.in") - .env(EnvVars::XDG_CONFIG_HOME, xdg.path()), @r#" + .env(EnvVars::XDG_CONFIG_HOME, xdg.path()), @r###" success: true exit_code: 0 ----- stdout ----- @@ -2934,6 +2949,7 @@ fn resolve_user_configuration() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -2982,7 +2998,7 @@ fn resolve_user_configuration() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); Ok(()) @@ -3178,7 +3194,7 @@ fn resolve_poetry_toml() -> anyhow::Result<()> { // Resolution should use the lowest direct version, and generate hashes. uv_snapshot!(context.filters(), add_shared_args(context.pip_compile(), context.temp_dir.path()) .arg("--show-settings") - .arg("requirements.in"), @r#" + .arg("requirements.in"), @r###" success: true exit_code: 0 ----- stdout ----- @@ -3267,6 +3283,7 @@ fn resolve_poetry_toml() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -3315,7 +3332,7 @@ fn resolve_poetry_toml() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); Ok(()) @@ -3357,7 +3374,7 @@ fn resolve_both() -> anyhow::Result<()> { // Resolution should succeed, but warn that the `pip` section in `pyproject.toml` is ignored. uv_snapshot!(context.filters(), add_shared_args(context.pip_compile(), context.temp_dir.path()) .arg("--show-settings") - .arg("requirements.in"), @r#" + .arg("requirements.in"), @r###" success: true exit_code: 0 ----- stdout ----- @@ -3476,6 +3493,7 @@ fn resolve_both() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -3525,7 +3543,7 @@ fn resolve_both() -> anyhow::Result<()> { ----- stderr ----- warning: Found both a `uv.toml` file and a `[tool.uv]` section in an adjacent `pyproject.toml`. The `[tool.uv]` section will be ignored in favor of the `uv.toml` file. - "# + "### ); Ok(()) @@ -3654,7 +3672,7 @@ fn resolve_config_file() -> anyhow::Result<()> { .arg("--show-settings") .arg("--config-file") .arg(config.path()) - .arg("requirements.in"), @r#" + .arg("requirements.in"), @r###" success: true exit_code: 0 ----- stdout ----- @@ -3773,6 +3791,7 @@ fn resolve_config_file() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -3821,7 +3840,7 @@ fn resolve_config_file() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); // Write in `pyproject.toml` schema. @@ -3929,7 +3948,7 @@ fn resolve_skip_empty() -> anyhow::Result<()> { uv_snapshot!(context.filters(), add_shared_args(context.pip_compile(), context.temp_dir.path()) .arg("--show-settings") .arg("requirements.in") - .current_dir(&child), @r#" + .current_dir(&child), @r###" success: true exit_code: 0 ----- stdout ----- @@ -4018,6 +4037,7 @@ fn resolve_skip_empty() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -4066,7 +4086,7 @@ fn resolve_skip_empty() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); // Adding a `tool.uv` section should cause us to ignore the `uv.toml`. @@ -4083,7 +4103,7 @@ fn resolve_skip_empty() -> anyhow::Result<()> { uv_snapshot!(context.filters(), add_shared_args(context.pip_compile(), context.temp_dir.path()) .arg("--show-settings") .arg("requirements.in") - .current_dir(&child), @r#" + .current_dir(&child), @r###" success: true exit_code: 0 ----- stdout ----- @@ -4172,6 +4192,7 @@ fn resolve_skip_empty() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -4220,7 +4241,7 @@ fn resolve_skip_empty() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); Ok(()) @@ -4245,7 +4266,7 @@ fn allow_insecure_host() -> anyhow::Result<()> { uv_snapshot!(context.filters(), add_shared_args(context.pip_compile(), context.temp_dir.path()) .arg("--show-settings") - .arg("requirements.in"), @r#" + .arg("requirements.in"), @r###" success: true exit_code: 0 ----- stdout ----- @@ -4345,6 +4366,7 @@ fn allow_insecure_host() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -4393,7 +4415,7 @@ fn allow_insecure_host() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); Ok(()) @@ -4421,7 +4443,7 @@ fn index_priority() -> anyhow::Result<()> { .arg("requirements.in") .arg("--show-settings") .arg("--index-url") - .arg("https://cli.pypi.org/simple"), @r#" + .arg("https://cli.pypi.org/simple"), @r###" success: true exit_code: 0 ----- stdout ----- @@ -4571,6 +4593,7 @@ fn index_priority() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -4619,14 +4642,14 @@ fn index_priority() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); uv_snapshot!(context.filters(), add_shared_args(context.pip_compile(), context.temp_dir.path()) .arg("requirements.in") .arg("--show-settings") .arg("--default-index") - .arg("https://cli.pypi.org/simple"), @r#" + .arg("https://cli.pypi.org/simple"), @r###" success: true exit_code: 0 ----- stdout ----- @@ -4776,6 +4799,7 @@ fn index_priority() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -4824,7 +4848,7 @@ fn index_priority() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); let config = context.temp_dir.child("uv.toml"); @@ -4837,7 +4861,7 @@ fn index_priority() -> anyhow::Result<()> { .arg("requirements.in") .arg("--show-settings") .arg("--default-index") - .arg("https://cli.pypi.org/simple"), @r#" + .arg("https://cli.pypi.org/simple"), @r###" success: true exit_code: 0 ----- stdout ----- @@ -4987,6 +5011,7 @@ fn index_priority() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -5035,7 +5060,7 @@ fn index_priority() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); // Prefer the `--index` from the CLI, but treat the index from the file as the default. @@ -5043,7 +5068,7 @@ fn index_priority() -> anyhow::Result<()> { .arg("requirements.in") .arg("--show-settings") .arg("--index") - .arg("https://cli.pypi.org/simple"), @r#" + .arg("https://cli.pypi.org/simple"), @r###" success: true exit_code: 0 ----- stdout ----- @@ -5193,6 +5218,7 @@ fn index_priority() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -5241,7 +5267,7 @@ fn index_priority() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); let config = context.temp_dir.child("uv.toml"); @@ -5256,7 +5282,7 @@ fn index_priority() -> anyhow::Result<()> { .arg("requirements.in") .arg("--show-settings") .arg("--index-url") - .arg("https://cli.pypi.org/simple"), @r#" + .arg("https://cli.pypi.org/simple"), @r###" success: true exit_code: 0 ----- stdout ----- @@ -5406,6 +5432,7 @@ fn index_priority() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -5454,7 +5481,7 @@ fn index_priority() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); // Prefer the `--extra-index-url` from the CLI, but not as the default. @@ -5462,7 +5489,7 @@ fn index_priority() -> anyhow::Result<()> { .arg("requirements.in") .arg("--show-settings") .arg("--extra-index-url") - .arg("https://cli.pypi.org/simple"), @r#" + .arg("https://cli.pypi.org/simple"), @r###" success: true exit_code: 0 ----- stdout ----- @@ -5612,6 +5639,7 @@ fn index_priority() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -5660,7 +5688,7 @@ fn index_priority() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); Ok(()) @@ -5681,7 +5709,7 @@ fn verify_hashes() -> anyhow::Result<()> { uv_snapshot!(context.filters(), add_shared_args(context.pip_install(), context.temp_dir.path()) .arg("-r") .arg("requirements.in") - .arg("--show-settings"), @r#" + .arg("--show-settings"), @r###" success: true exit_code: 0 ----- stdout ----- @@ -5771,6 +5799,7 @@ fn verify_hashes() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -5819,14 +5848,14 @@ fn verify_hashes() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); uv_snapshot!(context.filters(), add_shared_args(context.pip_install(), context.temp_dir.path()) .arg("-r") .arg("requirements.in") .arg("--no-verify-hashes") - .arg("--show-settings"), @r#" + .arg("--show-settings"), @r###" success: true exit_code: 0 ----- stdout ----- @@ -5916,6 +5945,7 @@ fn verify_hashes() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -5962,14 +5992,14 @@ fn verify_hashes() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); uv_snapshot!(context.filters(), add_shared_args(context.pip_install(), context.temp_dir.path()) .arg("-r") .arg("requirements.in") .arg("--require-hashes") - .arg("--show-settings"), @r#" + .arg("--show-settings"), @r###" success: true exit_code: 0 ----- stdout ----- @@ -6059,6 +6089,7 @@ fn verify_hashes() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -6107,14 +6138,14 @@ fn verify_hashes() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); uv_snapshot!(context.filters(), add_shared_args(context.pip_install(), context.temp_dir.path()) .arg("-r") .arg("requirements.in") .arg("--no-require-hashes") - .arg("--show-settings"), @r#" + .arg("--show-settings"), @r###" success: true exit_code: 0 ----- stdout ----- @@ -6204,6 +6235,7 @@ fn verify_hashes() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -6250,14 +6282,14 @@ fn verify_hashes() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); uv_snapshot!(context.filters(), add_shared_args(context.pip_install(), context.temp_dir.path()) .arg("-r") .arg("requirements.in") .env(EnvVars::UV_NO_VERIFY_HASHES, "1") - .arg("--show-settings"), @r#" + .arg("--show-settings"), @r###" success: true exit_code: 0 ----- stdout ----- @@ -6347,6 +6379,7 @@ fn verify_hashes() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -6393,7 +6426,7 @@ fn verify_hashes() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); uv_snapshot!(context.filters(), add_shared_args(context.pip_install(), context.temp_dir.path()) @@ -6401,7 +6434,7 @@ fn verify_hashes() -> anyhow::Result<()> { .arg("requirements.in") .arg("--verify-hashes") .arg("--no-require-hashes") - .arg("--show-settings"), @r#" + .arg("--show-settings"), @r###" success: true exit_code: 0 ----- stdout ----- @@ -6491,6 +6524,7 @@ fn verify_hashes() -> anyhow::Result<()> { prefix: None, index_strategy: FirstIndex, keyring_provider: Disabled, + torch_backend: None, no_build_isolation: false, no_build_isolation_package: [], build_options: BuildOptions { @@ -6539,7 +6573,7 @@ fn verify_hashes() -> anyhow::Result<()> { } ----- stderr ----- - "# + "### ); Ok(()) diff --git a/docs/reference/cli.md b/docs/reference/cli.md index 79c90ae36ea1..f620f5d6859e 100644 --- a/docs/reference/cli.md +++ b/docs/reference/cli.md @@ -6026,6 +6026,67 @@ uv pip compile [OPTIONS] ...

By default, uv uses the virtual environment in the current working directory or any parent directory, falling back to searching for a Python executable in PATH. The --system option instructs uv to avoid using a virtual environment Python and restrict its search to the system path.

May also be set with the UV_SYSTEM_PYTHON environment variable.

+
--torch-backend torch-backend

The backend to use when fetching packages in the PyTorch ecosystem (e.g., cu126 or auto)

+ +

When set, uv will ignore the configured index URLs for packages in the PyTorch ecosystem, and will instead use the defined backend.

+ +

For example, when set to cpu, uv will use the CPU-only PyTorch index; when set to cu126, uv will use the PyTorch index for CUDA 12.6.

+ +

The auto mode will attempt to detect the appropriate PyTorch index based on the currently installed CUDA drivers.

+ +

Possible values:

+ +
    +
  • auto: Select the appropriate PyTorch index based on the operating system and CUDA driver version
  • + +
  • cpu: Use the CPU-only PyTorch index
  • + +
  • cu126: Use the PyTorch index for CUDA 12.6
  • + +
  • cu125: Use the PyTorch index for CUDA 12.5
  • + +
  • cu124: Use the PyTorch index for CUDA 12.4
  • + +
  • cu123: Use the PyTorch index for CUDA 12.3
  • + +
  • cu122: Use the PyTorch index for CUDA 12.2
  • + +
  • cu121: Use the PyTorch index for CUDA 12.1
  • + +
  • cu120: Use the PyTorch index for CUDA 12.0
  • + +
  • cu118: Use the PyTorch index for CUDA 11.8
  • + +
  • cu117: Use the PyTorch index for CUDA 11.7
  • + +
  • cu116: Use the PyTorch index for CUDA 11.6
  • + +
  • cu115: Use the PyTorch index for CUDA 11.5
  • + +
  • cu114: Use the PyTorch index for CUDA 11.4
  • + +
  • cu113: Use the PyTorch index for CUDA 11.3
  • + +
  • cu112: Use the PyTorch index for CUDA 11.2
  • + +
  • cu111: Use the PyTorch index for CUDA 11.1
  • + +
  • cu110: Use the PyTorch index for CUDA 11.0
  • + +
  • cu102: Use the PyTorch index for CUDA 10.2
  • + +
  • cu101: Use the PyTorch index for CUDA 10.1
  • + +
  • cu100: Use the PyTorch index for CUDA 10.0
  • + +
  • cu92: Use the PyTorch index for CUDA 9.2
  • + +
  • cu91: Use the PyTorch index for CUDA 9.1
  • + +
  • cu90: Use the PyTorch index for CUDA 9.0
  • + +
  • cu80: Use the PyTorch index for CUDA 8.0
  • +
--universal

Perform a universal resolution, attempting to generate a single requirements.txt output file that is compatible with all operating systems, architectures, and Python implementations.

In universal mode, the current Python version (or user-provided --python-version) will be treated as a lower bound. For example, --universal --python-version 3.7 would produce a universal resolution for Python 3.7 and later.

@@ -6436,6 +6497,67 @@ uv pip sync [OPTIONS] ...

May also be set with the UV_SYSTEM_PYTHON environment variable.

--target target

Install packages into the specified directory, rather than into the virtual or system Python environment. The packages will be installed at the top-level of the directory

+
--torch-backend torch-backend

The backend to use when fetching packages in the PyTorch ecosystem (e.g., cu126 or auto)

+ +

When set, uv will ignore the configured index URLs for packages in the PyTorch ecosystem, and will instead use the defined backend.

+ +

For example, when set to cpu, uv will use the CPU-only PyTorch index; when set to cu126, uv will use the PyTorch index for CUDA 12.6.

+ +

The auto mode will attempt to detect the appropriate PyTorch index based on the currently installed CUDA drivers.

+ +

Possible values:

+ +
    +
  • auto: Select the appropriate PyTorch index based on the operating system and CUDA driver version
  • + +
  • cpu: Use the CPU-only PyTorch index
  • + +
  • cu126: Use the PyTorch index for CUDA 12.6
  • + +
  • cu125: Use the PyTorch index for CUDA 12.5
  • + +
  • cu124: Use the PyTorch index for CUDA 12.4
  • + +
  • cu123: Use the PyTorch index for CUDA 12.3
  • + +
  • cu122: Use the PyTorch index for CUDA 12.2
  • + +
  • cu121: Use the PyTorch index for CUDA 12.1
  • + +
  • cu120: Use the PyTorch index for CUDA 12.0
  • + +
  • cu118: Use the PyTorch index for CUDA 11.8
  • + +
  • cu117: Use the PyTorch index for CUDA 11.7
  • + +
  • cu116: Use the PyTorch index for CUDA 11.6
  • + +
  • cu115: Use the PyTorch index for CUDA 11.5
  • + +
  • cu114: Use the PyTorch index for CUDA 11.4
  • + +
  • cu113: Use the PyTorch index for CUDA 11.3
  • + +
  • cu112: Use the PyTorch index for CUDA 11.2
  • + +
  • cu111: Use the PyTorch index for CUDA 11.1
  • + +
  • cu110: Use the PyTorch index for CUDA 11.0
  • + +
  • cu102: Use the PyTorch index for CUDA 10.2
  • + +
  • cu101: Use the PyTorch index for CUDA 10.1
  • + +
  • cu100: Use the PyTorch index for CUDA 10.0
  • + +
  • cu92: Use the PyTorch index for CUDA 9.2
  • + +
  • cu91: Use the PyTorch index for CUDA 9.1
  • + +
  • cu90: Use the PyTorch index for CUDA 9.0
  • + +
  • cu80: Use the PyTorch index for CUDA 8.0
  • +
--verbose, -v

Use verbose output.

You can configure fine-grained logging using the RUST_LOG environment variable. (<https://docs.rs/tracing-subscriber/latest/tracing_subscriber/filter/struct.EnvFilter.html#directives>)

@@ -6916,6 +7038,67 @@ uv pip install [OPTIONS] |--editable May also be set with the UV_SYSTEM_PYTHON environment variable.

--target target

Install packages into the specified directory, rather than into the virtual or system Python environment. The packages will be installed at the top-level of the directory

+
--torch-backend torch-backend

The backend to use when fetching packages in the PyTorch ecosystem (e.g., cu126 or auto)

+ +

When set, uv will ignore the configured index URLs for packages in the PyTorch ecosystem, and will instead use the defined backend.

+ +

For example, when set to cpu, uv will use the CPU-only PyTorch index; when set to cu126, uv will use the PyTorch index for CUDA 12.6.

+ +

The auto mode will attempt to detect the appropriate PyTorch index based on the currently installed CUDA drivers.

+ +

Possible values:

+ +
    +
  • auto: Select the appropriate PyTorch index based on the operating system and CUDA driver version
  • + +
  • cpu: Use the CPU-only PyTorch index
  • + +
  • cu126: Use the PyTorch index for CUDA 12.6
  • + +
  • cu125: Use the PyTorch index for CUDA 12.5
  • + +
  • cu124: Use the PyTorch index for CUDA 12.4
  • + +
  • cu123: Use the PyTorch index for CUDA 12.3
  • + +
  • cu122: Use the PyTorch index for CUDA 12.2
  • + +
  • cu121: Use the PyTorch index for CUDA 12.1
  • + +
  • cu120: Use the PyTorch index for CUDA 12.0
  • + +
  • cu118: Use the PyTorch index for CUDA 11.8
  • + +
  • cu117: Use the PyTorch index for CUDA 11.7
  • + +
  • cu116: Use the PyTorch index for CUDA 11.6
  • + +
  • cu115: Use the PyTorch index for CUDA 11.5
  • + +
  • cu114: Use the PyTorch index for CUDA 11.4
  • + +
  • cu113: Use the PyTorch index for CUDA 11.3
  • + +
  • cu112: Use the PyTorch index for CUDA 11.2
  • + +
  • cu111: Use the PyTorch index for CUDA 11.1
  • + +
  • cu110: Use the PyTorch index for CUDA 11.0
  • + +
  • cu102: Use the PyTorch index for CUDA 10.2
  • + +
  • cu101: Use the PyTorch index for CUDA 10.1
  • + +
  • cu100: Use the PyTorch index for CUDA 10.0
  • + +
  • cu92: Use the PyTorch index for CUDA 9.2
  • + +
  • cu91: Use the PyTorch index for CUDA 9.1
  • + +
  • cu90: Use the PyTorch index for CUDA 9.0
  • + +
  • cu80: Use the PyTorch index for CUDA 8.0
  • +
--upgrade, -U

Allow package upgrades, ignoring pinned versions in any existing output file. Implies --refresh

--upgrade-package, -P upgrade-package

Allow upgrades for a specific package, ignoring pinned versions in any existing output file. Implies --refresh-package

diff --git a/docs/reference/settings.md b/docs/reference/settings.md index bf9ba62711e5..18d4e1c875cc 100644 --- a/docs/reference/settings.md +++ b/docs/reference/settings.md @@ -3246,6 +3246,67 @@ environment. The packages will be installed at the top-level of the directory. --- +#### [`torch-backend`](#pip_torch-backend) {: #pip_torch-backend } + + +The backend to use when fetching packages in the `PyTorch` ecosystem. + +When set, uv will ignore the configured index URLs for packages in the `PyTorch` ecosystem, +and will instead use the defined backend. + +For example, when set to `cpu`, uv will use the CPU-only `PyTorch` index; when set to `cu126`, +uv will use the `PyTorch` index for CUDA 12.6. + +The `auto` mode will attempt to detect the appropriate `PyTorch` index based on the currently +installed CUDA drivers. + +**Default value**: `null` + +**Possible values**: + +- `"auto"`: Select the appropriate PyTorch index based on the operating system and CUDA driver version +- `"cpu"`: Use the CPU-only PyTorch index +- `"cu126"`: Use the PyTorch index for CUDA 12.6 +- `"cu125"`: Use the PyTorch index for CUDA 12.5 +- `"cu124"`: Use the PyTorch index for CUDA 12.4 +- `"cu123"`: Use the PyTorch index for CUDA 12.3 +- `"cu122"`: Use the PyTorch index for CUDA 12.2 +- `"cu121"`: Use the PyTorch index for CUDA 12.1 +- `"cu120"`: Use the PyTorch index for CUDA 12.0 +- `"cu118"`: Use the PyTorch index for CUDA 11.8 +- `"cu117"`: Use the PyTorch index for CUDA 11.7 +- `"cu116"`: Use the PyTorch index for CUDA 11.6 +- `"cu115"`: Use the PyTorch index for CUDA 11.5 +- `"cu114"`: Use the PyTorch index for CUDA 11.4 +- `"cu113"`: Use the PyTorch index for CUDA 11.3 +- `"cu112"`: Use the PyTorch index for CUDA 11.2 +- `"cu111"`: Use the PyTorch index for CUDA 11.1 +- `"cu110"`: Use the PyTorch index for CUDA 11.0 +- `"cu102"`: Use the PyTorch index for CUDA 10.2 +- `"cu101"`: Use the PyTorch index for CUDA 10.1 +- `"cu100"`: Use the PyTorch index for CUDA 10.0 +- `"cu92"`: Use the PyTorch index for CUDA 9.2 +- `"cu91"`: Use the PyTorch index for CUDA 9.1 +- `"cu90"`: Use the PyTorch index for CUDA 9.0 +- `"cu80"`: Use the PyTorch index for CUDA 8.0 + +**Example usage**: + +=== "pyproject.toml" + + ```toml + [tool.uv.pip] + torch-backend = "auto" + ``` +=== "uv.toml" + + ```toml + [pip] + torch-backend = "auto" + ``` + +--- + #### [`universal`](#pip_universal) {: #pip_universal } diff --git a/uv.schema.json b/uv.schema.json index 19bf94be6c73..38c69637a8c4 100644 --- a/uv.schema.json +++ b/uv.schema.json @@ -1277,6 +1277,17 @@ "null" ] }, + "torch-backend": { + "description": "The backend to use when fetching packages in the `PyTorch` ecosystem.\n\nWhen set, uv will ignore the configured index URLs for packages in the `PyTorch` ecosystem, and will instead use the defined backend.\n\nFor example, when set to `cpu`, uv will use the CPU-only `PyTorch` index; when set to `cu126`, uv will use the `PyTorch` index for CUDA 12.6.\n\nThe `auto` mode will attempt to detect the appropriate `PyTorch` index based on the currently installed CUDA drivers.", + "anyOf": [ + { + "$ref": "#/definitions/TorchMode" + }, + { + "type": "null" + } + ] + }, "universal": { "description": "Perform a universal resolution, attempting to generate a single `requirements.txt` output file that is compatible with all operating systems, architectures, and Python implementations.\n\nIn universal mode, the current Python version (or user-provided `--python-version`) will be treated as a lower bound. For example, `--universal --python-version 3.7` would produce a universal resolution for Python 3.7 and later.", "type": [ @@ -2096,6 +2107,186 @@ }, "additionalProperties": false }, + "TorchMode": { + "description": "The strategy to use when determining the appropriate PyTorch index.", + "oneOf": [ + { + "description": "Select the appropriate PyTorch index based on the operating system and CUDA driver version.", + "type": "string", + "enum": [ + "Auto" + ] + }, + { + "description": "Use the CPU-only PyTorch index.", + "type": "string", + "enum": [ + "Cpu" + ] + }, + { + "description": "Use the PyTorch index for CUDA 12.6.", + "type": "string", + "enum": [ + "Cu126" + ] + }, + { + "description": "Use the PyTorch index for CUDA 12.5.", + "type": "string", + "enum": [ + "Cu125" + ] + }, + { + "description": "Use the PyTorch index for CUDA 12.4.", + "type": "string", + "enum": [ + "Cu124" + ] + }, + { + "description": "Use the PyTorch index for CUDA 12.3.", + "type": "string", + "enum": [ + "Cu123" + ] + }, + { + "description": "Use the PyTorch index for CUDA 12.2.", + "type": "string", + "enum": [ + "Cu122" + ] + }, + { + "description": "Use the PyTorch index for CUDA 12.1.", + "type": "string", + "enum": [ + "Cu121" + ] + }, + { + "description": "Use the PyTorch index for CUDA 12.0.", + "type": "string", + "enum": [ + "Cu120" + ] + }, + { + "description": "Use the PyTorch index for CUDA 11.8.", + "type": "string", + "enum": [ + "Cu118" + ] + }, + { + "description": "Use the PyTorch index for CUDA 11.7.", + "type": "string", + "enum": [ + "Cu117" + ] + }, + { + "description": "Use the PyTorch index for CUDA 11.6.", + "type": "string", + "enum": [ + "Cu116" + ] + }, + { + "description": "Use the PyTorch index for CUDA 11.5.", + "type": "string", + "enum": [ + "Cu115" + ] + }, + { + "description": "Use the PyTorch index for CUDA 11.4.", + "type": "string", + "enum": [ + "Cu114" + ] + }, + { + "description": "Use the PyTorch index for CUDA 11.3.", + "type": "string", + "enum": [ + "Cu113" + ] + }, + { + "description": "Use the PyTorch index for CUDA 11.2.", + "type": "string", + "enum": [ + "Cu112" + ] + }, + { + "description": "Use the PyTorch index for CUDA 11.1.", + "type": "string", + "enum": [ + "Cu111" + ] + }, + { + "description": "Use the PyTorch index for CUDA 11.0.", + "type": "string", + "enum": [ + "Cu110" + ] + }, + { + "description": "Use the PyTorch index for CUDA 10.2.", + "type": "string", + "enum": [ + "Cu102" + ] + }, + { + "description": "Use the PyTorch index for CUDA 10.1.", + "type": "string", + "enum": [ + "Cu101" + ] + }, + { + "description": "Use the PyTorch index for CUDA 10.0.", + "type": "string", + "enum": [ + "Cu100" + ] + }, + { + "description": "Use the PyTorch index for CUDA 9.2.", + "type": "string", + "enum": [ + "Cu92" + ] + }, + { + "description": "Use the PyTorch index for CUDA 9.1.", + "type": "string", + "enum": [ + "Cu91" + ] + }, + { + "description": "Use the PyTorch index for CUDA 9.0.", + "type": "string", + "enum": [ + "Cu90" + ] + }, + { + "description": "Use the PyTorch index for CUDA 8.0.", + "type": "string", + "enum": [ + "Cu80" + ] + } + ] + }, "TrustedHost": { "description": "A host or host-port pair.", "type": "string"