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-backendThe 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
targetInstall 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-backendThe 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
targetInstall 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-backendThe 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-packageAllow 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"