Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: faster & cleaner VK generation #2084

Merged
merged 10 commits into from
Jun 6, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions prover/Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

3 changes: 2 additions & 1 deletion prover/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ dialoguer = "0.11"
futures = "0.3"
hex = "0.4"
itertools = "0.10.5"
indicatif = "0.16"
jemallocator = "0.5"
local-ip-address = "0.5.0"
log = "0.4.20"
Expand All @@ -52,7 +53,7 @@ proptest = "1.2.0"
prover_dal = { path = "prover_dal" }
queues = "1.1.0"
rand = "0.8"
regex = "1.7.2"
regex = "1.10.4"
reqwest = "0.11"
serde = "1.0"
serde_derive = "1.0"
Expand Down
1 change: 1 addition & 0 deletions prover/vk_setup_data_generator_server_fri/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ toml_edit.workspace = true
md5.workspace = true
sha3.workspace = true
hex.workspace = true
indicatif.workspace = true

[dev-dependencies]
proptest.workspace = true
Expand Down
3 changes: 3 additions & 0 deletions prover/vk_setup_data_generator_server_fri/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,9 @@ circuit changes), first please make sure that you have a CRS file (used for SNAR
CRS_FILE=yyy ZKSYNC_HOME=xxx cargo run --release --bin key_generator generate-vk
```

You can also generate multiple keys in parallel (to speed things up), with `--jobs` flag, but you need at least 30 GB of
ram for each job.

### CRS FILE

The SNARK VK generation requires the `CRS_FILE` environment variable to be present and point to the correct file. The
Expand Down
55 changes: 47 additions & 8 deletions prover/vk_setup_data_generator_server_fri/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,13 @@ use std::collections::HashMap;
use anyhow::Context as _;
use clap::{Parser, Subcommand};
use commitment_generator::read_and_update_contract_toml;
use indicatif::{ProgressBar, ProgressStyle};
use tracing::level_filters::LevelFilter;
use zkevm_test_harness::{
compute_setups::{generate_base_layer_vks_and_proofs, generate_recursive_layer_vks_and_proofs},
compute_setups::{
basic_vk_count, generate_base_layer_vks, generate_recursive_layer_vks,
recursive_layer_vk_count,
},
data_source::{in_memory_data_source::InMemoryDataSource, SetupDataSource},
proof_wrapper_utils::{
check_trusted_setup_file_existace, get_wrapper_setup_and_vk_from_scheduler_vk,
Expand All @@ -30,20 +34,45 @@ mod commitment_generator;
#[cfg(test)]
mod tests;

fn generate_vks(keystore: &Keystore) -> anyhow::Result<()> {
/// Generates new verification keys, and stores them in `keystore`.
/// Jobs describe how many generators can run in parallel (each one is around 30 GB).
/// If quiet is true, it doesn't display any progress bar.
fn generate_vks(keystore: &Keystore, jobs: usize, quiet: bool) -> anyhow::Result<()> {
// Start by checking the trusted setup existence.
// This is used at the last step, but we want to fail early if user didn't configure everything
// correctly.
check_trusted_setup_file_existace();

let progress_bar = if quiet {
None
} else {
let count = basic_vk_count() + recursive_layer_vk_count() + 1;
let progress_bar = ProgressBar::new(count as u64);
progress_bar.set_style(ProgressStyle::default_bar()
.template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {pos:>7}/{len:7} ({eta})")
.progress_chars("#>-"));
Some(progress_bar)
};

let pb = std::sync::Arc::new(std::sync::Mutex::new(progress_bar));

let mut in_memory_source = InMemoryDataSource::new();
tracing::info!("Generating verification keys for Base layer.");
generate_base_layer_vks_and_proofs(&mut in_memory_source)
.map_err(|err| anyhow::anyhow!("Failed generating base vk's: {err}"))?;

generate_base_layer_vks(&mut in_memory_source, Some(jobs), || {
if let Some(p) = pb.lock().unwrap().as_ref() {
p.inc(1)
}
})
.map_err(|err| anyhow::anyhow!("Failed generating base vk's: {err}"))?;

tracing::info!("Generating verification keys for Recursive layer.");
generate_recursive_layer_vks_and_proofs(&mut in_memory_source)
.map_err(|err| anyhow::anyhow!("Failed generating recursive vk's: {err}"))?;
generate_recursive_layer_vks(&mut in_memory_source, Some(jobs), || {
if let Some(p) = pb.lock().unwrap().as_ref() {
p.inc(1)
}
})
.map_err(|err| anyhow::anyhow!("Failed generating recursive vk's: {err}"))?;

tracing::info!("Saving keys & hints");

Expand All @@ -63,6 +92,10 @@ fn generate_vks(keystore: &Keystore) -> anyhow::Result<()> {
.save_snark_verification_key(vk)
.context("save_snark_vk")?;

if let Some(p) = pb.lock().unwrap().as_ref() {
p.inc(1)
}

// Let's also update the commitments file.
keystore.save_commitments(&generate_commitments(keystore)?)
}
Expand Down Expand Up @@ -121,6 +154,12 @@ enum Command {
GenerateVerificationKeys {
#[arg(long)]
path: Option<String>,
/// Number of generators to run in parallel - each one consumes around 30 GB of RAM.
#[arg(short, long, default_value_t = 1)]
jobs: usize,
/// If true - disables progress bar.
#[arg(long)]
quiet: bool,
},
/// Generates setup keys (used by the CPU prover).
#[command(name = "generate-sk")]
Expand Down Expand Up @@ -208,13 +247,13 @@ fn main() -> anyhow::Result<()> {
let opt = Cli::parse();

match opt.command {
Command::GenerateVerificationKeys { path } => {
Command::GenerateVerificationKeys { path, jobs, quiet } => {
let keystore = keystore_from_optional_path(path, None);
tracing::info!(
"Generating verification keys and storing them inside {:?}",
keystore.get_base_path()
);
generate_vks(&keystore).context("generate_vks()")
generate_vks(&keystore, jobs, quiet).context("generate_vks()")
}
Command::UpdateCommitments { dryrun, path } => {
let keystore = keystore_from_optional_path(path, None);
Expand Down
Loading