Automatically infer the PyTorch index via --torch-backend=auto

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" }

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 }
 

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<TorchMode>,
+
     #[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<TorchMode>,
+
     #[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<TorchMode>,
+
     #[command(flatten)]
     pub compat_args: compat::PipInstallCompatArgs,
 }

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 }
 

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<TorchStrategy>,
     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<TorchStrategy>) -> 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<BaseClientBuilder<'a>> 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<TorchStrategy>,
     /// 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<Item = &IndexUrl> {
+        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<Vec<(&'index IndexUrl, OwnedArchive<SimpleMetadata>)>, 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

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! {

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 }

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}"),
+        }
+    }
+}

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;