Disallow mixing requirements across PyTorch indexes

crates/uv-resolver/Cargo.toml

@@ -36,6 +36,7 @@ uv-python = { workspace = true }
 uv-requirements-txt = { workspace = true }
 uv-small-str = { workspace = true }
 uv-static = { workspace = true }
+uv-torch = { workspace = true }
 uv-types = { workspace = true }
 uv-warnings = { workspace = true }
 uv-workspace = { workspace = true }

crates/uv-resolver/src/options.rs

@@ -1,7 +1,9 @@
-use crate::fork_strategy::ForkStrategy;
-use crate::{DependencyMode, ExcludeNewer, PrereleaseMode, ResolutionMode};
 use uv_configuration::{BuildOptions, IndexStrategy};
 use uv_pypi_types::SupportedEnvironments;
+use uv_torch::TorchStrategy;
+
+use crate::fork_strategy::ForkStrategy;
+use crate::{DependencyMode, ExcludeNewer, PrereleaseMode, ResolutionMode};
 
 /// Options for resolving a manifest.
 #[derive(Debug, Default, Clone, PartialEq, Eq)]
@@ -15,6 +17,7 @@ pub struct Options {
     pub required_environments: SupportedEnvironments,
     pub flexibility: Flexibility,
     pub build_options: BuildOptions,
+    pub torch_backend: Option<TorchStrategy>,
 }
 
 /// Builder for [`Options`].
@@ -29,6 +32,7 @@ pub struct OptionsBuilder {
     required_environments: SupportedEnvironments,
     flexibility: Flexibility,
     build_options: BuildOptions,
+    torch_backend: Option<TorchStrategy>,
 }
 
 impl OptionsBuilder {
@@ -100,6 +104,13 @@ impl OptionsBuilder {
         self
     }
 
+    /// Sets the [`TorchStrategy`].
+    #[must_use]
+    pub fn torch_backend(mut self, torch_backend: Option<TorchStrategy>) -> Self {
+        self.torch_backend = torch_backend;
+        self
+    }
+
     /// Builds the options.
     pub fn build(self) -> Options {
         Options {
@@ -112,6 +123,7 @@ impl OptionsBuilder {
             required_environments: self.required_environments,
             flexibility: self.flexibility,
             build_options: self.build_options,
+            torch_backend: self.torch_backend,
         }
     }
 }

crates/uv-resolver/src/pubgrub/dependencies.rs

@@ -127,10 +127,11 @@ impl PubGrubDependency {
                         url,
                     }
                 }
-                PubGrubPackageInner::Root(_) => unreachable!("root package in dependencies"),
+                PubGrubPackageInner::Root(_) => unreachable!("Root package in dependencies"),
                 PubGrubPackageInner::Python(_) => {
-                    unreachable!("python package in dependencies")
+                    unreachable!("Python package in dependencies")
                 }
+                PubGrubPackageInner::System(_) => unreachable!("System package in dependencies"),
             }
         })
     }

crates/uv-resolver/src/pubgrub/package.rs

@@ -44,6 +44,8 @@ pub(crate) enum PubGrubPackageInner {
     Root(Option<PackageName>),
     /// A Python version.
     Python(PubGrubPython),
+    /// A system package, which is used to represent a non-Python package.
+    System(PackageName),
     /// A Python package.
     ///
     /// Note that it is guaranteed that `extra` and `dev` are never both
@@ -134,18 +136,21 @@ impl PubGrubPackage {
             // package is never returned by `get_dependencies`. So these cases never occur.
             PubGrubPackageInner::Root(None) | PubGrubPackageInner::Python(_) => None,
             PubGrubPackageInner::Root(Some(name))
+            | PubGrubPackageInner::System(name)
             | PubGrubPackageInner::Package { name, .. }
             | PubGrubPackageInner::Extra { name, .. }
             | PubGrubPackageInner::Dev { name, .. }
             | PubGrubPackageInner::Marker { name, .. } => Some(name),
         }
     }
 
-    /// Returns the name of this PubGrub package, if it is not the root package or a Python version
-    /// constraint.
+    /// Returns the name of this PubGrub package, if it is not the root package, a Python version
+    /// constraint, or a system package.
     pub(crate) fn name_no_root(&self) -> Option<&PackageName> {
         match &**self {
-            PubGrubPackageInner::Root(_) | PubGrubPackageInner::Python(_) => None,
+            PubGrubPackageInner::Root(_)
+            | PubGrubPackageInner::Python(_)
+            | PubGrubPackageInner::System(_) => None,
             PubGrubPackageInner::Package { name, .. }
             | PubGrubPackageInner::Extra { name, .. }
             | PubGrubPackageInner::Dev { name, .. }
@@ -159,7 +164,9 @@ impl PubGrubPackage {
         match &**self {
             // A root can never be a dependency of another package, and a `Python` pubgrub
             // package is never returned by `get_dependencies`. So these cases never occur.
-            PubGrubPackageInner::Root(_) | PubGrubPackageInner::Python(_) => MarkerTree::TRUE,
+            PubGrubPackageInner::Root(_)
+            | PubGrubPackageInner::Python(_)
+            | PubGrubPackageInner::System(_) => MarkerTree::TRUE,
             PubGrubPackageInner::Package { marker, .. }
             | PubGrubPackageInner::Extra { marker, .. }
             | PubGrubPackageInner::Dev { marker, .. } => *marker,
@@ -177,6 +184,7 @@ impl PubGrubPackage {
             // package is never returned by `get_dependencies`. So these cases never occur.
             PubGrubPackageInner::Root(_)
             | PubGrubPackageInner::Python(_)
+            | PubGrubPackageInner::System(_)
             | PubGrubPackageInner::Package { extra: None, .. }
             | PubGrubPackageInner::Dev { .. }
             | PubGrubPackageInner::Marker { .. } => None,
@@ -198,6 +206,7 @@ impl PubGrubPackage {
             // package is never returned by `get_dependencies`. So these cases never occur.
             PubGrubPackageInner::Root(_)
             | PubGrubPackageInner::Python(_)
+            | PubGrubPackageInner::System(_)
             | PubGrubPackageInner::Package { dev: None, .. }
             | PubGrubPackageInner::Extra { .. }
             | PubGrubPackageInner::Marker { .. } => None,
@@ -256,7 +265,9 @@ impl PubGrubPackage {
     /// reporting where this routine is used.
     pub(crate) fn simplify_markers(&mut self, python_requirement: &PythonRequirement) {
         match *Arc::make_mut(&mut self.0) {
-            PubGrubPackageInner::Root(_) | PubGrubPackageInner::Python(_) => {}
+            PubGrubPackageInner::Root(_)
+            | PubGrubPackageInner::Python(_)
+            | PubGrubPackageInner::System(_) => {}
             PubGrubPackageInner::Package { ref mut marker, .. }
             | PubGrubPackageInner::Extra { ref mut marker, .. }
             | PubGrubPackageInner::Dev { ref mut marker, .. }
@@ -272,6 +283,7 @@ impl PubGrubPackage {
         match &**self {
             PubGrubPackageInner::Root(_) => "root",
             PubGrubPackageInner::Python(_) => "python",
+            PubGrubPackageInner::System(_) => "system",
             PubGrubPackageInner::Package { .. } => "package",
             PubGrubPackageInner::Extra { .. } => "extra",
             PubGrubPackageInner::Dev { .. } => "dev",
@@ -304,6 +316,7 @@ impl std::fmt::Display for PubGrubPackageInner {
                 }
             }
             Self::Python(_) => write!(f, "Python"),
+            Self::System(name) => write!(f, "system:{name}"),
             Self::Package {
                 name,
                 extra: None,

crates/uv-resolver/src/pubgrub/priority.rs

@@ -129,6 +129,7 @@ impl PubGrubPriorities {
             PubGrubPackageInner::Python(PubGrubPython::Target) => {
                 (PubGrubPriority::Root, PubGrubTiebreaker::from(2))
             }
+            PubGrubPackageInner::System(_) => (PubGrubPriority::Root, PubGrubTiebreaker::from(3)),
             PubGrubPackageInner::Marker { name, .. }
             | PubGrubPackageInner::Extra { name, .. }
             | PubGrubPackageInner::Dev { name, .. }

crates/uv-resolver/src/resolver/mod.rs

@@ -63,19 +63,20 @@ use crate::resolver::environment::{
     fork_version_by_marker, fork_version_by_python_requirement, ForkingPossibility,
 };
 pub(crate) use crate::resolver::fork_map::{ForkMap, ForkSet};
-pub(crate) use crate::resolver::urls::Urls;
-use crate::universal_marker::{ConflictMarker, UniversalMarker};
-pub(crate) use provider::MetadataUnavailable;
-
 pub use crate::resolver::index::InMemoryIndex;
 use crate::resolver::indexes::Indexes;
 pub use crate::resolver::provider::{
     DefaultResolverProvider, MetadataResponse, PackageVersionsResult, ResolverProvider,
     VersionsResponse, WheelMetadataResult,
 };
 pub use crate::resolver::reporter::{BuildId, Reporter};
+use crate::resolver::system::SystemDependency;
+pub(crate) use crate::resolver::urls::Urls;
+use crate::universal_marker::{ConflictMarker, UniversalMarker};
 use crate::yanks::AllowedYanks;
 use crate::{marker, DependencyMode, Exclusions, FlatIndex, Options, ResolutionMode, VersionMap};
+pub(crate) use provider::MetadataUnavailable;
+use uv_torch::TorchStrategy;
 
 mod availability;
 mod batch_prefetch;
@@ -86,6 +87,7 @@ mod index;
 mod indexes;
 mod provider;
 mod reporter;
+mod system;
 mod urls;
 
 /// The number of conflicts a package may accumulate before we re-prioritize and backtrack.
@@ -598,6 +600,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
                     next_id,
                     next_package,
                     &version,
+                    &state.pins,
                     &state.fork_urls,
                     &state.env,
                     &state.python_requirement,
@@ -1055,6 +1058,15 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
                 Ok(None)
             }
 
+            PubGrubPackageInner::System(_) => {
+                // We don't care what the actual version is here, just that it's consistent across
+                // the dependency graph.
+                let Some(version) = range.as_singleton() else {
+                    return Ok(None);
+                };
+                Ok(Some(ResolverVersion::Unforked(version.clone())))
+            }
+
             PubGrubPackageInner::Marker { name, .. }
             | PubGrubPackageInner::Extra { name, .. }
             | PubGrubPackageInner::Dev { name, .. }
@@ -1641,6 +1653,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
         id: Id<PubGrubPackage>,
         package: &PubGrubPackage,
         version: &Version,
+        pins: &FilePins,
         fork_urls: &ForkUrls,
         env: &ResolverEnvironment,
         python_requirement: &PythonRequirement,
@@ -1650,6 +1663,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
             id,
             package,
             version,
+            pins,
             fork_urls,
             env,
             python_requirement,
@@ -1674,6 +1688,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
         id: Id<PubGrubPackage>,
         package: &PubGrubPackage,
         version: &Version,
+        pins: &FilePins,
         fork_urls: &ForkUrls,
         env: &ResolverEnvironment,
         python_requirement: &PythonRequirement,
@@ -1781,6 +1796,24 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
                     }
                 }
 
+                // Identify any system dependencies based on the index URL.
+                let system_dependencies = self
+                    .options
+                    .torch_backend
+                    .as_ref()
+                    .filter(|torch_backend| matches!(torch_backend, TorchStrategy::Auto { .. }))
+                    .and_then(|_| pins.get(name, version).and_then(ResolvedDist::index))
+                    .map(IndexUrl::url)
+                    .and_then(SystemDependency::from_index)
+                    .into_iter()
+                    .inspect(|system_dependency| {
+                        debug!(
+                            "Adding system dependency `{}` for `{package}@{version}`",
+                            system_dependency
+                        );
+                    })
+                    .map(PubGrubDependency::from);
+
                 let requirements = self.flatten_requirements(
                     &metadata.requires_dist,
                     &metadata.dependency_groups,
@@ -1800,11 +1833,14 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
                             Some(name),
                         )
                     })
+                    .chain(system_dependencies)
                     .collect()
             }
 
             PubGrubPackageInner::Python(_) => return Ok(Dependencies::Unforkable(Vec::default())),
 
+            PubGrubPackageInner::System(_) => return Ok(Dependencies::Unforkable(Vec::default())),
+
             // Add a dependency on both the marker and base package.
             PubGrubPackageInner::Marker { name, marker } => {
                 return Ok(Dependencies::Unforkable(
@@ -2562,6 +2598,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
             match &**package {
                 PubGrubPackageInner::Root(_) => {}
                 PubGrubPackageInner::Python(_) => {}
+                PubGrubPackageInner::System(_) => {}
                 PubGrubPackageInner::Marker { .. } => {}
                 PubGrubPackageInner::Extra { .. } => {}
                 PubGrubPackageInner::Dev { .. } => {}

crates/uv-resolver/src/resolver/system.rs

@@ -0,0 +1,84 @@
+use std::str::FromStr;
+
+use pubgrub::Ranges;
+use url::Url;
+
+use uv_normalize::PackageName;
+use uv_pep440::Version;
+use uv_torch::TorchBackend;
+
+use crate::pubgrub::{PubGrubDependency, PubGrubPackage, PubGrubPackageInner};
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub(super) struct SystemDependency {
+    /// The name of the system dependency (e.g., `cuda`).
+    name: PackageName,
+    /// The version of the system dependency (e.g., `12.4`).
+    version: Version,
+}
+
+impl SystemDependency {
+    /// Extract a [`SystemDependency`] from an index URL.
+    ///
+    /// For example, given `https://download.pytorch.org/whl/cu124`, returns CUDA 12.4.
+    pub(super) fn from_index(index: &Url) -> Option<Self> {
+        let backend = TorchBackend::from_index(index)?;
+        let cuda_version = backend.cuda_version()?;
+        Some(Self {
+            name: PackageName::from_str("cuda").unwrap(),
+            version: cuda_version,
+        })
+    }
+}
+
+impl std::fmt::Display for SystemDependency {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}@{}", self.name, self.version)
+    }
+}
+
+impl From<SystemDependency> for PubGrubDependency {
+    fn from(value: SystemDependency) -> Self {
+        PubGrubDependency {
+            package: PubGrubPackage::from(PubGrubPackageInner::System(value.name)),
+            version: Ranges::singleton(value.version),
+            url: None,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::str::FromStr;
+
+    use url::Url;
+
+    use uv_normalize::PackageName;
+    use uv_pep440::Version;
+
+    use crate::resolver::system::SystemDependency;
+
+    #[test]
+    fn pypi() {
+        let url = Url::parse("https://pypi.org/simple").unwrap();
+        assert_eq!(SystemDependency::from_index(&url), None);
+    }
+
+    #[test]
+    fn pytorch_cuda_12_4() {
+        let url = Url::parse("https://download.pytorch.org/whl/cu124").unwrap();
+        assert_eq!(
+            SystemDependency::from_index(&url),
+            Some(SystemDependency {
+                name: PackageName::from_str("cuda").unwrap(),
+                version: Version::new([12, 4]),
+            })
+        );
+    }
+
+    #[test]
+    fn pytorch_cpu() {
+        let url = Url::parse("https://download.pytorch.org/whl/cpu").unwrap();
+        assert_eq!(SystemDependency::from_index(&url), None);
+    }
+}