feat(sdk): async app prover (#2154)

Adds `AsyncAppProver` which executes metered serially and then
re-executes while spawning prover tasks for each continuation segment.
The prover tasks use a semaphore to limit the max concurrency.

Note this PR includes the fix
#2155

---------

Co-authored-by: Ayush Shukla <[email protected]>

Author: jonathanpwang

Cargo.lock

@@ -113,11 +113,11 @@ lto = "thin"
 
 [workspace.dependencies]
 # Stark Backend
-openvm-stark-backend = { git = "https://github.com/openvm-org/stark-backend.git", branch="fix/async-errors", default-features = false }
-openvm-stark-sdk = { git = "https://github.com/openvm-org/stark-backend.git", branch="fix/async-errors", default-features = false }
-openvm-cuda-backend = { git = "https://github.com/openvm-org/stark-backend.git", branch="fix/async-errors", default-features = false }
-openvm-cuda-builder = { git = "https://github.com/openvm-org/stark-backend.git", branch="fix/async-errors", default-features = false }
-openvm-cuda-common = { git = "https://github.com/openvm-org/stark-backend.git", branch="fix/async-errors", default-features = false }
+openvm-stark-backend = { git = "https://github.com/openvm-org/stark-backend.git", branch="main", default-features = false }
+openvm-stark-sdk = { git = "https://github.com/openvm-org/stark-backend.git", branch="main", default-features = false }
+openvm-cuda-backend = { git = "https://github.com/openvm-org/stark-backend.git", branch="main", default-features = false }
+openvm-cuda-builder = { git = "https://github.com/openvm-org/stark-backend.git", branch="main", default-features = false }
+openvm-cuda-common = { git = "https://github.com/openvm-org/stark-backend.git", branch="main", default-features = false }
 
 # OpenVM
 openvm-sdk = { path = "crates/sdk", default-features = false }
@@ -233,6 +233,7 @@ dashmap = "6.1.0"
 memmap2 = "0.9.5"
 libc = "0.2.175"
 tracing-subscriber = { version = "0.3.20", features = ["std", "env-filter"] }
+tokio = "1" # >=1.0.0 to allow downstream flexibility
 
 # default-features = false for no_std for use in guest programs
 itertools = { version = "0.14.0", default-features = false }

Cargo.toml

@@ -19,9 +19,9 @@ openvm-native-circuit.workspace = true
 openvm-native-compiler.workspace = true
 openvm-native-recursion = { workspace = true, features = ["test-utils"] }
 
-clap = { version = "4.5.9", features = ["derive", "env"] }
+clap = { workspace = true, features = ["derive", "env"] }
 eyre.workspace = true
-tokio = { version = "1.43.1", features = ["rt", "rt-multi-thread", "macros", "time"] }
+tokio = { workspace = true, features = ["rt", "rt-multi-thread", "macros"] }
 rand_chacha = { version = "0.3", default-features = false }
 k256 = { workspace = true, features = ["ecdsa"] }
 tiny-keccak.workspace = true
@@ -33,11 +33,12 @@ metrics.workspace = true
 [dev-dependencies]
 
 [features]
-default = ["parallel", "jemalloc", "metrics"]
+default = ["parallel", "jemalloc", "metrics", "async"]
 metrics = ["openvm-sdk/metrics"]
 tco = ["openvm-sdk/tco"]
 perf-metrics = ["openvm-sdk/perf-metrics", "metrics"]
 stark-debug = ["openvm-sdk/stark-debug"]
+async = ["openvm-sdk/async"]
 # runs leaf aggregation benchmarks:
 aggregation = []
 evm = ["openvm-sdk/evm-verify"]
@@ -63,3 +64,8 @@ path = "src/bin/fib_e2e.rs"
 [[bin]]
 name = "kitchen_sink"
 path = "src/bin/kitchen_sink.rs"
+
+[[bin]]
+name = "async_regex"
+path = "src/bin/async_regex.rs"
+required-features = ["async"]

benchmarks/prove/Cargo.toml

@@ -1,22 +1,14 @@
-use std::{
-    env::var,
-    sync::{
-        atomic::{AtomicUsize, Ordering},
-        Arc,
-    },
-    time::Duration,
-};
+use std::env::var;
 
 use clap::Parser;
 use openvm_benchmarks_prove::util::BenchmarkCli;
 use openvm_benchmarks_utils::get_programs_dir;
 use openvm_sdk::{
     config::{SdkVmBuilder, SdkVmConfig},
-    prover::AppProver,
+    prover::AsyncAppProver,
     DefaultStarkEngine, Sdk, StdIn, F,
 };
 use openvm_stark_sdk::config::setup_tracing;
-use tokio::{spawn, task::spawn_blocking, time::sleep};
 
 #[tokio::main]
 async fn main() -> eyre::Result<()> {
@@ -47,46 +39,15 @@ async fn main() -> eyre::Result<()> {
     let (app_pk, _app_vk) = sdk.app_keygen();
 
     let max_par_jobs: usize = var("MAX_PAR_JOBS").map(|m| m.parse()).unwrap_or(Ok(1))?;
-    let num_jobs: usize = var("NUM_JOBS").map(|m| m.parse()).unwrap_or(Ok(20))?;
-    let cur_num_jobs = Arc::new(AtomicUsize::new(0));
 
-    let mut tasks = Vec::with_capacity(num_jobs);
-    for idx in 0..num_jobs {
-        let cur_num_jobs = cur_num_jobs.clone();
-        let app_exe = app_exe.clone();
-        let app_pk = app_pk.clone();
-        let input = input.clone();
-        let task = spawn(async move {
-            loop {
-                let c = cur_num_jobs.fetch_add(1, Ordering::SeqCst);
-                if c < max_par_jobs {
-                    tracing::info!("Acquired job {}, cur num jobs {}", idx, c + 1);
-                    break;
-                }
-                cur_num_jobs.fetch_sub(1, Ordering::SeqCst);
-                sleep(Duration::from_millis(100)).await;
-            }
-            let res = spawn_blocking(move || -> eyre::Result<()> {
-                let mut prover = AppProver::<DefaultStarkEngine, _>::new(
-                    SdkVmBuilder,
-                    &app_pk.app_vm_pk,
-                    app_exe,
-                    app_pk.leaf_verifier_program_commit(),
-                )?;
-                let _proof = prover.prove(input)?;
-                Ok(())
-            })
-            .await?;
-            let prev_num = cur_num_jobs.fetch_sub(1, Ordering::SeqCst);
-            tracing::info!("Decrement cur_num_jobs {} to {}", prev_num, prev_num - 1);
-            res
-        });
-        tasks.push(task);
-        sleep(Duration::from_millis(1000)).await;
-    }
-    for task in tasks {
-        task.await??;
-    }
+    let prover = AsyncAppProver::<DefaultStarkEngine, _>::new(
+        SdkVmBuilder,
+        app_pk.app_vm_pk.clone(),
+        app_exe,
+        app_pk.leaf_verifier_program_commit(),
+        max_par_jobs,
+    )?;
+    let _proof = prover.prove(input).await?;
 
     Ok(())
 }

benchmarks/prove/src/bin/async_regex.rs

@@ -61,9 +61,10 @@ forge-fmt = { workspace = true, optional = true }
 rrs-lib.workspace = true
 num-bigint.workspace = true
 cfg-if.workspace = true
+tokio = { workspace = true, features = ["rt", "sync"], optional = true }
 
 [features]
-default = ["parallel", "jemalloc"]
+default = ["parallel", "jemalloc", "async"]
 evm-prove = [
     "openvm-continuations/static-verifier",
     "openvm-native-recursion/evm-prove",
@@ -101,6 +102,7 @@ perf-metrics = [
 # turns on stark-backend debugger in all proofs
 stark-debug = ["openvm-circuit/stark-debug"]
 test-utils = ["openvm-circuit/test-utils"]
+async = ["tokio"]
 # performance features:
 # (rayon is always imported because of halo2, so "parallel" feature is redundant)
 parallel = ["openvm-circuit/parallel"]