diff --git a/crates/uv-bench/benches/uv.rs b/crates/uv-bench/benches/uv.rs index d8b1def9c..d6ac82c3d 100644 --- a/crates/uv-bench/benches/uv.rs +++ b/crates/uv-bench/benches/uv.rs @@ -3,44 +3,48 @@ use std::str::FromStr; use criterion::{Criterion, criterion_group, criterion_main, measurement::WallTime}; use uv_cache::Cache; -use uv_client::{BaseClientBuilder, RegistryClientBuilder}; +use uv_client::{BaseClientBuilder, Connectivity, RegistryClientBuilder}; use uv_distribution_types::Requirement; use uv_python::PythonEnvironment; use uv_resolver::Manifest; fn resolve_warm_jupyter(c: &mut Criterion) { - let run = setup(Manifest::simple(vec![Requirement::from( + let manifest = Manifest::simple(vec![Requirement::from( uv_pep508::Requirement::from_str("jupyter==1.0.0").unwrap(), - )])); - c.bench_function("resolve_warm_jupyter", |b| b.iter(|| run(false))); + )]); + let run = setup(manifest, false); + c.bench_function("resolve_warm_jupyter", |b| b.iter(&run)); } fn resolve_warm_jupyter_universal(c: &mut Criterion) { - let run = setup(Manifest::simple(vec![Requirement::from( + let manifest = Manifest::simple(vec![Requirement::from( uv_pep508::Requirement::from_str("jupyter==1.0.0").unwrap(), - )])); - c.bench_function("resolve_warm_jupyter_universal", |b| b.iter(|| run(true))); + )]); + let run = setup(manifest, true); + c.bench_function("resolve_warm_jupyter_universal", |b| b.iter(&run)); } fn resolve_warm_airflow(c: &mut Criterion) { - let run = setup(Manifest::simple(vec![ + let manifest = Manifest::simple(vec![ Requirement::from(uv_pep508::Requirement::from_str("apache-airflow[all]==2.9.3").unwrap()), Requirement::from( uv_pep508::Requirement::from_str("apache-airflow-providers-apache-beam>3.0.0").unwrap(), ), - ])); - c.bench_function("resolve_warm_airflow", |b| b.iter(|| run(false))); + ]); + let run = setup(manifest, false); + c.bench_function("resolve_warm_airflow", |b| b.iter(&run)); } // This takes >5m to run in CodSpeed. // fn resolve_warm_airflow_universal(c: &mut Criterion) { -// let run = setup(Manifest::simple(vec![ +// let manifest = Manifest::simple(vec![ // Requirement::from(uv_pep508::Requirement::from_str("apache-airflow[all]").unwrap()), // Requirement::from( // uv_pep508::Requirement::from_str("apache-airflow-providers-apache-beam>3.0.0").unwrap(), // ), -// ])); -// c.bench_function("resolve_warm_airflow_universal", |b| b.iter(|| run(true))); +// ]); +// let run = setup(manifest, true); +// c.bench_function("resolve_warm_airflow_universal", |b| b.iter(&run)); // } criterion_group!( @@ -51,7 +55,7 @@ criterion_group!( ); criterion_main!(uv); -fn setup(manifest: Manifest) -> impl Fn(bool) { +fn setup(manifest: Manifest, universal: bool) -> impl Fn() { let runtime = tokio::runtime::Builder::new_current_thread() // CodSpeed limits the total number of threads to 500 .max_blocking_threads(256) @@ -68,7 +72,29 @@ fn setup(manifest: Manifest) -> impl Fn(bool) { .into_interpreter(); let client = RegistryClientBuilder::new(BaseClientBuilder::default(), cache.clone()).build(); - move |universal| { + // Prime the cache: First run for performance the network operation, the second run primes + // reading from the cache from the first run. If they are already primed, we only lose ~1s for + // the large airflow benchmark. + for _ in 0..2 { + runtime + .block_on(resolver::resolve( + black_box(manifest.clone()), + black_box(cache.clone()), + black_box(&client), + &interpreter, + universal, + )) + .unwrap(); + } + + // No matter how long the benchmarks run, never do fresh network requests + let client = RegistryClientBuilder::new( + BaseClientBuilder::default().connectivity(Connectivity::Offline), + cache.clone(), + ) + .build(); + + move || { runtime .block_on(resolver::resolve( black_box(manifest.clone()),