diff --git a/Cargo.lock b/Cargo.lock index 3ba8345f7..da480ead5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6802,9 +6802,12 @@ dependencies = [ "rustc-hash", "serde", "serde_json", + "thiserror 2.0.16", "tracing", "uv-distribution-filename", + "uv-normalize", "uv-once-map", + "uv-pep440", "uv-pep508", "uv-pypi-types", ] diff --git a/crates/uv-distribution/src/distribution_database.rs b/crates/uv-distribution/src/distribution_database.rs index 526558457..4bf6ab9d8 100644 --- a/crates/uv-distribution/src/distribution_database.rs +++ b/crates/uv-distribution/src/distribution_database.rs @@ -1,12 +1,14 @@ -use std::future::Future; -use std::io; -use std::path::Path; -use std::pin::Pin; -use std::sync::Arc; -use std::task::{Context, Poll}; - use futures::{FutureExt, StreamExt, TryStreamExt}; use rustc_hash::FxHashMap; +use std::ffi::OsString; +use std::future::Future; +use std::path::Path; +use std::path::PathBuf; +use std::pin::Pin; +use std::str::FromStr; +use std::sync::Arc; +use std::task::{Context, Poll}; +use std::{env, io}; use tempfile::TempDir; use tokio::io::{AsyncRead, AsyncSeekExt, ReadBuf}; use tokio::sync::Semaphore; @@ -27,13 +29,15 @@ use uv_distribution_types::{ }; use uv_extract::hash::Hasher; use uv_fs::write_atomic; -use uv_pep508::{MarkerEnvironment, MarkerVariantsUniversal, VariantNamespace}; +use uv_pep440::VersionSpecifiers; +use uv_pep508::{MarkerEnvironment, MarkerVariantsUniversal, VariantNamespace, VersionOrUrl}; use uv_platform_tags::Tags; use uv_pypi_types::{HashDigest, HashDigests, PyProjectToml}; use uv_redacted::DisplaySafeUrl; use uv_types::{BuildContext, BuildStack, VariantsTrait}; use uv_variants::VariantProviderOutput; use uv_variants::resolved_variants::ResolvedVariants; +use uv_variants::variant_lock::{VariantLock, VariantLockProvider, VariantLockResolved}; use uv_variants::variants_json::{Provider, VariantsJsonContent}; use crate::archive::Archive; @@ -594,7 +598,33 @@ impl<'a, Context: BuildContext> DistributionDatabase<'a, Context> { marker_env: &MarkerEnvironment, debug_filename: &VariantsJson, ) -> Result { + // TODO: parse_boolish_environment_variable + let locked_and_inferred = + env::var_os("UV_VARIANT_LOCK_INCOMPLETE").is_some_and(|var| var == "1"); + // TODO(konsti): Integrate this properly, and add this to the CLI. + let variant_lock = if let Some(variant_lock_path) = env::var_os("UV_VARIANT_LOCK") { + let variant_lock: VariantLock = toml::from_slice( + &fs_err::read(&variant_lock_path).map_err(Error::VariantLockRead)?, + ) + .map_err(|err| Error::VariantLockParse(PathBuf::from(&variant_lock_path), err))?; + // TODO(konsti): If parsing fails, check the version + if !VersionSpecifiers::from_str(">=0.1,<0.2") + .unwrap() + .contains(&variant_lock.metadata.version) + { + return Err(Error::VariantLockVersion( + PathBuf::from(variant_lock_path), + variant_lock.metadata.version, + )); + } + + Some((variant_lock_path, variant_lock)) + } else { + None + }; + // Query the install time provider. + // TODO(konsti): Don't use threads if we're fully static. let mut provider_outputs: FxHashMap> = futures::stream::iter(variants_json.providers.iter().filter(|(_, provider)| { provider.plugin_use.unwrap_or_default().run_on_install() @@ -602,10 +632,11 @@ impl<'a, Context: BuildContext> DistributionDatabase<'a, Context> { .enable_if .evaluate(marker_env, MarkerVariantsUniversal, &[]) })) - .map(|(name, provider)| self.query_variant_provider(name, provider)) + .map(|(name, provider)| { + self.resolve_provider(locked_and_inferred, variant_lock.as_ref(), name, provider) + }) // TODO(konsti): Buffer size .buffered(8) - .map_ok(|config| (config.namespace.clone(), config)) .try_collect() .await?; @@ -638,6 +669,42 @@ impl<'a, Context: BuildContext> DistributionDatabase<'a, Context> { }) } + async fn resolve_provider( + &self, + locked_and_inferred: bool, + variant_lock: Option<&(OsString, VariantLock)>, + name: &VariantNamespace, + provider: &Provider, + ) -> Result<(VariantNamespace, Arc), Error> { + if let Some((variant_lock_path, variant_lock)) = &variant_lock { + if let Some(static_provider) = variant_lock + .provider + .iter() + .find(|static_provider| satisfies_provider_requires(provider, static_provider)) + { + Ok(( + static_provider.namespace.clone(), + Arc::new(VariantProviderOutput { + namespace: static_provider.namespace.clone(), + features: static_provider.properties.clone().into_iter().collect(), + }), + )) + } else if locked_and_inferred { + let config = self.query_variant_provider(name, provider).await?; + Ok((config.namespace.clone(), config)) + } else { + Err(Error::VariantLockMissing { + variant_lock: PathBuf::from(variant_lock_path), + requires: provider.requires.clone(), + plugin_api: provider.plugin_api.clone().unwrap_or_default().clone(), + }) + } + } else { + let config = self.query_variant_provider(name, provider).await?; + Ok((config.namespace.clone(), config)) + } + } + async fn query_variant_provider( &self, name: &VariantNamespace, @@ -1465,6 +1532,43 @@ fn add_tar_zst_extension(mut url: DisplaySafeUrl) -> DisplaySafeUrl { url } +fn satisfies_provider_requires( + requested_provider: &Provider, + static_provider: &VariantLockProvider, +) -> bool { + // TODO(konsti): Correct plugin_api inference. + if static_provider.plugin_api.clone().or(static_provider + .resolved + .first() + .map(|resolved| resolved.name().to_string())) + != requested_provider.plugin_api.clone().or(static_provider + .resolved + .first() + .map(|resolved| resolved.name().to_string())) + { + return false; + } + requested_provider.requires.iter().all(|requested| { + static_provider.resolved.iter().any(|resolved| { + // We ignore extras for simplicity. + &requested.name == resolved.name() + && match (&requested.version_or_url, resolved) { + (None, _) => true, + ( + Some(VersionOrUrl::VersionSpecifier(requested)), + VariantLockResolved::Version(_name, resolved), + ) => requested.contains(resolved), + ( + Some(VersionOrUrl::Url(resolved_url)), + VariantLockResolved::Url(_name, url), + ) => resolved_url == &**url, + // We don't support URL<->version matching + _ => false, + } + }) + }) +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/uv-distribution/src/error.rs b/crates/uv-distribution/src/error.rs index 539a8a027..7b5280f0e 100644 --- a/crates/uv-distribution/src/error.rs +++ b/crates/uv-distribution/src/error.rs @@ -1,5 +1,6 @@ use std::path::PathBuf; +use itertools::Itertools; use owo_colors::OwoColorize; use tokio::task::JoinError; use zip::result::ZipError; @@ -11,8 +12,8 @@ use uv_distribution_types::{InstalledDist, InstalledDistError, IsBuildBackendErr use uv_fs::Simplified; use uv_normalize::PackageName; use uv_pep440::{Version, VersionSpecifiers}; -use uv_pep508::VariantNamespace; -use uv_pypi_types::{HashAlgorithm, HashDigest}; +use uv_pep508::{Requirement, VariantNamespace}; +use uv_pypi_types::{HashAlgorithm, HashDigest, VerbatimParsedUrl}; use uv_redacted::DisplaySafeUrl; use uv_types::AnyErrorBuild; @@ -84,6 +85,23 @@ pub enum Error { declared: VariantNamespace, actual: VariantNamespace, }, + #[error("Failed to read variant lock")] + VariantLockRead(#[source] std::io::Error), + #[error("Variant lock has an unsupported format: {}", _0.user_display())] + VariantLockParse(PathBuf, #[source] toml::de::Error), + #[error("Variant lock has an unsupported version {}, only version 0.1 is supported: {}", _1, _0.user_display())] + VariantLockVersion(PathBuf, Version), + #[error( + "Variant lock is missing a matching provider and `UV_VARIANT_LOCK_INCOMPLETE` is not set\n variant lock: {}\n requires: `{}`\n plugin-api: {}", + variant_lock.user_display(), + requires.iter().join("`, `"), + plugin_api + )] + VariantLockMissing { + variant_lock: PathBuf, + requires: Vec>, + plugin_api: String, + }, #[error("Failed to parse metadata from built wheel")] Metadata(#[from] uv_pypi_types::MetadataError), #[error("Failed to read metadata: `{}`", _0.user_display())] diff --git a/crates/uv-variants/Cargo.toml b/crates/uv-variants/Cargo.toml index 328796d1f..84d9b12d8 100644 --- a/crates/uv-variants/Cargo.toml +++ b/crates/uv-variants/Cargo.toml @@ -11,7 +11,9 @@ license.workspace = true [dependencies] uv-distribution-filename = { workspace = true } +uv-normalize = { workspace = true } uv-once-map = { workspace = true } +uv-pep440 = { workspace = true, features = ["serde"] } uv-pep508 = { workspace = true, features = ["serde"] } uv-pypi-types = { workspace = true } @@ -19,6 +21,7 @@ indexmap = { workspace = true } indoc = { workspace = true } rustc-hash = { workspace = true } serde = { workspace = true } +thiserror = { workspace = true } tracing = { workspace = true } [dev-dependencies] diff --git a/crates/uv-variants/src/lib.rs b/crates/uv-variants/src/lib.rs index 19b9daf32..b6346e629 100644 --- a/crates/uv-variants/src/lib.rs +++ b/crates/uv-variants/src/lib.rs @@ -4,6 +4,7 @@ use uv_pep508::{VariantFeature, VariantNamespace, VariantValue}; pub mod cache; pub mod resolved_variants; +pub mod variant_lock; pub mod variants_json; /// Wire format between with the Python shim for provider plugins. diff --git a/crates/uv-variants/src/variant_lock.rs b/crates/uv-variants/src/variant_lock.rs new file mode 100644 index 000000000..2c74b4539 --- /dev/null +++ b/crates/uv-variants/src/variant_lock.rs @@ -0,0 +1,94 @@ +use std::str::FromStr; + +use indexmap::IndexMap; +use serde::{Deserialize, Deserializer}; +use thiserror::Error; + +use uv_normalize::PackageName; +use uv_pep440::Version; +use uv_pep508::{UnnamedRequirementUrl, VariantFeature, VariantNamespace, VariantValue}; +use uv_pypi_types::VerbatimParsedUrl; + +#[derive(Debug, Error)] +pub enum VariantLockError { + #[error("Invalid resolved requirement format: {0}")] + InvalidResolvedFormat(String), +} + +#[derive(Debug, Clone, Deserialize)] +#[serde(rename_all = "kebab-case")] +pub struct VariantLock { + pub metadata: VariantLockMetadata, + pub provider: Vec, +} + +#[derive(Debug, Clone, Deserialize)] +#[serde(rename_all = "kebab-case")] +pub struct VariantLockMetadata { + pub created_by: String, + pub version: Version, +} + +#[derive(Debug, Clone, Deserialize)] +#[serde(rename_all = "kebab-case")] +pub struct VariantLockProvider { + pub resolved: Vec, + pub plugin_api: Option, + pub namespace: VariantNamespace, + pub properties: IndexMap>, +} + +/// A resolved requirement in the form `==` or ` @ ` +#[derive(Debug, Clone)] +pub enum VariantLockResolved { + Version(PackageName, Version), + Url(PackageName, Box), +} + +impl VariantLockResolved { + pub fn name(&self) -> &PackageName { + match self { + Self::Version(name, _) => name, + Self::Url(name, _) => name, + } + } +} + +impl FromStr for VariantLockResolved { + type Err = VariantLockError; + + fn from_str(resolved: &str) -> Result { + if let Some((name, version)) = resolved.split_once("==") { + Ok(Self::Version( + PackageName::from_str(name.trim()) + .map_err(|_| VariantLockError::InvalidResolvedFormat(resolved.to_string()))?, + Version::from_str(version.trim()) + .map_err(|_| VariantLockError::InvalidResolvedFormat(resolved.to_string()))?, + )) + } else if let Some((name, url)) = resolved.split_once(" @ ") { + Ok(Self::Url( + PackageName::from_str(name.trim()) + .map_err(|_| VariantLockError::InvalidResolvedFormat(resolved.to_string()))?, + Box::new( + VerbatimParsedUrl::parse_unnamed_url(url.trim()).map_err(|_| { + VariantLockError::InvalidResolvedFormat(resolved.to_string()) + })?, + ), + )) + } else { + Err(VariantLockError::InvalidResolvedFormat( + resolved.to_string(), + )) + } + } +} + +impl<'de> Deserialize<'de> for VariantLockResolved { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let s = String::deserialize(deserializer)?; + Self::from_str(&s).map_err(serde::de::Error::custom) + } +} diff --git a/files/built_by_uv-0.1.0-variants.json b/files/built_by_uv-0.1.0-variants.json index b2cc5ac41..eeba3ae68 100644 --- a/files/built_by_uv-0.1.0-variants.json +++ b/files/built_by_uv-0.1.0-variants.json @@ -24,7 +24,7 @@ "enable-if": "platform_machine == 'x86_64' or platform_machine == 'AMD64'", "plugin-api": "cpu_level_provider", "requires": [ - "cpu_level_provider" + "cpu_level_provider >= 0.1, <0.2" ] } }, diff --git a/files/variant-lock-incomplete.toml b/files/variant-lock-incomplete.toml new file mode 100644 index 000000000..c08e8d89a --- /dev/null +++ b/files/variant-lock-incomplete.toml @@ -0,0 +1,12 @@ +[metadata] +created-by = "tool-that-created-the-file" # Freeform string +version = "0.1" # PEP 440 version + +# We don't have static information for the requested provider + +[[provider]] +resolved = ["foo==0.0.1"] +plugin-api = "foo" +namespace = "foo" + +[provider.properties] diff --git a/files/variant-lock-none.toml b/files/variant-lock-none.toml new file mode 100644 index 000000000..8ceade36d --- /dev/null +++ b/files/variant-lock-none.toml @@ -0,0 +1,11 @@ +[metadata] +created-by = "tool-that-created-the-file" # Freeform string +version = "0.1" # PEP 440 version + +[[provider]] +resolved = ["cpu_level_provider==0.1"] +plugin-api = "cpu_level_provider" +namespace = "cpu_level" + +[provider.properties] +x86_64_level = [] diff --git a/files/variant-lock.toml b/files/variant-lock.toml new file mode 100644 index 000000000..68fc32261 --- /dev/null +++ b/files/variant-lock.toml @@ -0,0 +1,20 @@ +[metadata] +created-by = "tool-that-created-the-file" # Freeform string +version = "0.1" # PEP 440 version + +[[provider]] +resolved = ["cpu_level_provider==0.1"] +plugin-api = "cpu_level_provider" +namespace = "cpu_level" + +[provider.properties] +x86_64_level = ["v1", "v2"] + +[[provider]] +resolved = ["nvidia-variant-provider==0.0.1"] +plugin-api = "nvidia_variant_provider.plugin:NvidiaVariantPlugin" +namespace = "nvidia" + +[provider.properties] +cuda_version_lower_bound = ["12.8"] +sm_arch = ["100_real", "120_real", "70_real", "75_real", "80_real", "86_real", "90_real"] diff --git a/test_variants.sh b/test_variants.sh index 9702baa90..3617efed2 100755 --- a/test_variants.sh +++ b/test_variants.sh @@ -18,6 +18,15 @@ uv venv -c -q && UV_CPU_LEVEL_OVERRIDE=2 ${uv} pip install built-by-uv --no-inde echo "# Matching cpu2 variant wheel, to be preferred over the non-variant wheel and the sdist" uv venv -c -q && UV_CPU_LEVEL_OVERRIDE=2 RUST_LOG=uv_distribution_types=debug ${uv} pip install built-by-uv --no-index --no-cache --no-progress --find-links ./files --find-links ./files_wheel --find-links ./files_sdist +echo "Valid variants lock" +uv venv -c -q && UV_VARIANT_LOCK=files/variant-lock.toml RUST_LOG=uv_distribution_types=debug ${uv} pip install built-by-uv --no-index --no-cache --no-progress --find-links ./files +echo "No matching variants wheel with variants lock" +uv venv -c -q && ( ( UV_VARIANT_LOCK=files/variant-lock-none.toml RUST_LOG=uv_distribution_types=debug ${uv} pip install built-by-uv --no-index --no-cache --no-progress --find-links ./files && exit 1 ) || exit 0 ) +echo "No matching variant provider in variants lock" +uv venv -c -q && ( ( UV_VARIANT_LOCK=files/variant-lock-incomplete.toml RUST_LOG=uv_distribution_types=debug ${uv} pip install built-by-uv --no-index --no-cache --no-progress --find-links ./files && exit 1 ) || exit 0 ) +echo "No matching variants in variants lock, but UV_VARIANT_LOCK_INCOMPLETE set" +uv venv -c -q && UV_VARIANT_LOCK_INCOMPLETE=1 UV_VARIANT_LOCK=files/variant-lock-incomplete.toml RUST_LOG=uv_distribution_types=debug ${uv} pip install built-by-uv --no-index --no-cache --no-progress --find-links ./files + echo "# sync without a compatible variant wheel (fresh)" ( cd scripts/packages/cpu_user && rm -f uv.lock && ( ( uv venv -c -q && UV_CPU_LEVEL_OVERRIDE=0 ${uv} sync && exit 1 ) || exit 0 ) ) echo "# sync without a compatible variant wheel (existing lockfile)"