Airflow workspace benchmark

.github/workflows/bench.yml

@@ -47,11 +47,21 @@ jobs:
           cargo run --bin uv -- pip compile test/requirements/jupyter.in --universal --exclude-newer 2024-08-08 --cache-dir .cache
           cargo run --bin uv -- pip compile test/requirements/airflow.in --universal --exclude-newer 2024-08-08 --cache-dir .cache
 
+      - name: "Clone airflow workspace"
+        run: |
+          git init airflow
+          git -C airflow remote add origin https://github.com/apache/airflow.git
+          git -C airflow fetch --depth 1 origin 0027d171d908908692f16755a77bc2e4dea42a25
+          git -C airflow checkout FETCH_HEAD
+
+      - name: "Copy airflow lock file"
+        run: cp crates/uv-bench/inputs/airflow.uv.lock airflow/uv.lock
+
       - name: "Build benchmarks"
         run: cargo codspeed build -m walltime --profile profiling -p uv-bench
 
       - name: "Create artifact archive"
-        run: tar -cvf benchmarks-walltime.tar target/codspeed target/debug/uv .cache
+        run: tar -cvf benchmarks-walltime.tar target/codspeed target/debug/uv .cache airflow
 
       - name: "Upload benchmark artifacts"
         uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2
@@ -125,6 +135,16 @@ jobs:
           cargo run --bin uv -- pip compile test/requirements/jupyter.in --universal --exclude-newer 2024-08-08 --cache-dir .cache
           cargo run --bin uv -- pip compile test/requirements/airflow.in --universal --exclude-newer 2024-08-08 --cache-dir .cache
 
+      - name: "Clone airflow workspace"
+        run: |
+          git init airflow
+          git -C airflow remote add origin https://github.com/apache/airflow.git
+          git -C airflow fetch --depth 1 origin 0027d171d908908692f16755a77bc2e4dea42a25
+          git -C airflow checkout FETCH_HEAD
+
+      - name: "Copy airflow lock file"
+        run: cp crates/uv-bench/inputs/airflow.uv.lock airflow/uv.lock
+
       - name: "Build benchmarks"
         run: cargo codspeed build --profile profiling -p uv-bench
 

crates/uv-bench/Cargo.toml

@@ -23,7 +23,9 @@ path = "benches/uv.rs"
 harness = false
 
 [dev-dependencies]
+uv = { path = "../uv" }
 uv-cache = { workspace = true }
+uv-cli = { workspace = true }
 uv-client = { workspace = true }
 uv-configuration = { workspace = true }
 uv-dispatch = { workspace = true }
@@ -42,9 +44,11 @@ uv-types = { workspace = true }
 uv-workspace = { workspace = true }
 
 anyhow = { workspace = true }
+clap = { workspace = true }
 criterion = { version = "4.0.3", default-features = false, package = "codspeed-criterion-compat", features = ["async_tokio"] }
 jiff = { workspace = true }
 tokio = { workspace = true }
+toml = { workspace = true }
 
 [package.metadata.cargo-shear]
 ignored = ["uv-extract"]

crates/uv-bench/benches/uv.rs

@@ -1,12 +1,14 @@
 use std::hint::black_box;
 use std::str::FromStr;
 
+use clap::Parser;
 use criterion::{Criterion, criterion_group, criterion_main, measurement::WallTime};
 use uv_cache::Cache;
+use uv_cli::Cli;
 use uv_client::{BaseClientBuilder, RegistryClientBuilder};
 use uv_distribution_types::Requirement;
 use uv_python::PythonEnvironment;
-use uv_resolver::Manifest;
+use uv_resolver::{Lock, Manifest};
 
 fn resolve_warm_jupyter(c: &mut Criterion<WallTime>) {
     let run = setup(Manifest::simple(vec![Requirement::from(
@@ -43,18 +45,79 @@ fn resolve_warm_airflow(c: &mut Criterion<WallTime>) {
 //     c.bench_function("resolve_warm_airflow_universal", |b| b.iter(|| run(true)));
 // }
 
+/// Benchmark `uv run python -V` in the airflow workspace with a satisfied lockfile.
+fn run_noop_airflow(c: &mut Criterion<WallTime>) {
+    let airflow_dir = std::path::absolute("../../airflow").unwrap();
+    if !airflow_dir.join("uv.lock").exists() {
+        return;
+    }
+    let cache_dir = std::path::absolute("../../.cache").unwrap();
+
+    let runtime = tokio::runtime::Builder::new_current_thread()
+        .enable_all()
+        .build()
+        .unwrap();
+
+    let airflow_dir = airflow_dir.to_string_lossy().to_string();
+    let cache_dir = cache_dir.to_string_lossy().to_string();
+
+    // Verify the setup works before benchmarking.
+    let cli = Cli::try_parse_from([
+        "uv",
+        "run",
+        "--directory",
+        &airflow_dir,
+        "--cache-dir",
+        &cache_dir,
+        "--no-sync",
+        "--quiet",
+        "python",
+        "-V",
+    ])
+    .unwrap();
+    runtime.block_on(uv::run(cli)).unwrap();
+
+    c.bench_function("run_noop_airflow", |b| {
+        b.iter(|| {
+            let cli = Cli::try_parse_from([
+                "uv",
+                "run",
+                "--directory",
+                black_box(&airflow_dir),
+                "--cache-dir",
+                black_box(&cache_dir),
+                "--no-sync",
+                "--quiet",
+                "python",
+                "-V",
+            ])
+            .unwrap();
+            runtime.block_on(uv::run(cli)).unwrap();
+        });
+    });
+}
+
+/// Benchmark lock file parsing for the airflow workspace (891 packages).
+fn parse_airflow_lockfile(c: &mut Criterion<WallTime>) {
+    let lockfile = include_str!("../inputs/airflow.uv.lock");
+
+    c.bench_function("parse_airflow_lockfile", |b| {
+        b.iter(|| toml::from_str::<Lock>(black_box(lockfile)).unwrap());
+    });
+}
+
 criterion_group!(
     uv,
     resolve_warm_jupyter,
     resolve_warm_jupyter_universal,
-    resolve_warm_airflow
+    resolve_warm_airflow,
+    run_noop_airflow,
+    parse_airflow_lockfile,
 );
 criterion_main!(uv);
 
 fn setup(manifest: Manifest) -> impl Fn(bool) {
     let runtime = tokio::runtime::Builder::new_current_thread()
-        // CodSpeed limits the total number of threads to 500
-        .max_blocking_threads(256)
         .enable_all()
         .build()
         .unwrap();

crates/uv/src/commands/mod.rs

@@ -96,7 +96,7 @@ mod venv;
 mod workspace;
 
 #[derive(Copy, Clone)]
-pub(crate) enum ExitStatus {
+pub enum ExitStatus {
     /// The command succeeded.
     Success,
 

crates/uv/src/lib.rs

@@ -60,7 +60,7 @@ use crate::settings::{
 };
 
 pub(crate) mod child;
-pub(crate) mod commands;
+pub mod commands;
 #[cfg(not(feature = "self-update"))]
 mod install_source;
 pub(crate) mod logging;
@@ -70,7 +70,7 @@ pub(crate) mod settings;
 mod windows_exception;
 
 #[instrument(skip_all)]
-async fn run(mut cli: Cli) -> Result<ExitStatus> {
+pub async fn run(mut cli: Cli) -> Result<ExitStatus> {
     // Enable flag to pick up warnings generated by workspace loading.
     if cli.top_level.global_args.quiet == 0 {
         uv_warnings::enable();
@@ -339,27 +339,48 @@ async fn run(mut cli: Cli) -> Result<ExitStatus> {
         &environment,
     );
 
-    // Set the global flags.
-    uv_flags::init(EnvironmentFlags::from(&environment))
-        .map_err(|()| anyhow::anyhow!("Flags are already initialized"))?;
-
-    // Configure the `tracing` crate, which controls internal logging.
-    #[cfg(feature = "tracing-durations-export")]
-    let (durations_layer, _duration_guard) =
-        logging::setup_durations(environment.tracing_durations_file.as_ref())?;
-    #[cfg(not(feature = "tracing-durations-export"))]
-    let durations_layer = None::<tracing_subscriber::layer::Identity>;
-    logging::setup_logging(
-        match globals.verbose {
-            0 => logging::Level::Off,
-            1 => logging::Level::DebugUv,
-            2 => logging::Level::TraceUv,
-            3.. => logging::Level::TraceAll,
-        },
-        durations_layer,
-        globals.color,
-        environment.log_context.unwrap_or_default(),
-    )?;
+    // Perform one-time global initialization. These use `OnceLock` or global subscribers
+    // internally, so they must only run once per process. The `AtomicBool` guard allows
+    // `run()` to be called multiple times (e.g., in benchmarks) without failing.
+    static INITIALIZED: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
+    if !INITIALIZED.swap(true, Ordering::SeqCst) {
+        // Set the global flags.
+        uv_flags::init(EnvironmentFlags::from(&environment))
+            .map_err(|()| anyhow::anyhow!("Flags are already initialized"))?;
+
+        // Configure the `tracing` crate, which controls internal logging.
+        #[cfg(feature = "tracing-durations-export")]
+        let (durations_layer, _duration_guard) =
+            logging::setup_durations(environment.tracing_durations_file.as_ref())?;
+        #[cfg(not(feature = "tracing-durations-export"))]
+        let durations_layer = None::<tracing_subscriber::layer::Identity>;
+        logging::setup_logging(
+            match globals.verbose {
+                0 => logging::Level::Off,
+                1 => logging::Level::DebugUv,
+                2 => logging::Level::TraceUv,
+                3.. => logging::Level::TraceAll,
+            },
+            durations_layer,
+            globals.color,
+            environment.log_context.unwrap_or_default(),
+        )?;
+
+        miette::set_hook(Box::new(|_| {
+            Box::new(
+                miette::MietteHandlerOpts::new()
+                    .break_words(false)
+                    .word_separator(textwrap::WordSeparator::AsciiSpace)
+                    .word_splitter(textwrap::WordSplitter::NoHyphenation)
+                    .wrap_lines(
+                        std::env::var(EnvVars::UV_NO_WRAP)
+                            .map(|_| false)
+                            .unwrap_or(true),
+                    )
+                    .build(),
+            )
+        }))?;
+    }
 
     debug!("uv {}", uv_cli::version::uv_self_version());
     if let Some(config_file) = cli.top_level.config_file.as_ref() {
@@ -449,21 +470,6 @@ async fn run(mut cli: Cli) -> Result<ExitStatus> {
 
     anstream::ColorChoice::write_global(globals.color.into());
 
-    miette::set_hook(Box::new(|_| {
-        Box::new(
-            miette::MietteHandlerOpts::new()
-                .break_words(false)
-                .word_separator(textwrap::WordSeparator::AsciiSpace)
-                .word_splitter(textwrap::WordSplitter::NoHyphenation)
-                .wrap_lines(
-                    std::env::var(EnvVars::UV_NO_WRAP)
-                        .map(|_| false)
-                        .unwrap_or(true),
-                )
-                .build(),
-        )
-    }))?;
-
     // Don't initialize the rayon threadpool yet, this is too costly when we're doing a noop sync.
     uv_configuration::RAYON_PARALLELISM.store(globals.concurrency.installs, Ordering::Relaxed);