diff --git a/prover/setup.sh b/prover/setup.sh index 0f2b2d3301f6..e755f5ae433a 100755 --- a/prover/setup.sh +++ b/prover/setup.sh @@ -23,32 +23,5 @@ rm ../etc/env/base/fri_proof_compressor.toml.backup zk config compile dev - -SHOULD_GENERATE_KEYS=false - -for i in {1..13} -do - if ! [ -f vk_setup_data_generator_server_fri/data/setup_basic_${i}_data.bin ]; then - SHOULD_GENERATE_KEYS=true - fi -done - -if ! [ -f vk_setup_data_generator_server_fri/data/setup_scheduler_data.bin ]; then - SHOULD_GENERATE_KEYS=true -fi - -if ! [ -f vk_setup_data_generator_server_fri/data/setup_node_data.bin ]; then - SHOULD_GENERATE_KEYS=true -fi - -for i in {3..15} -do - if ! [ -f vk_setup_data_generator_server_fri/data/setup_leaf_${i}_data.bin ]; then - SHOULD_GENERATE_KEYS=true - fi -done - - -if [ "$SHOULD_GENERATE_KEYS" = "true" ]; then - zk f cargo run $GPU_FLAG --release --bin key_generator -- $GENERATE_SK_COMMAND all -fi \ No newline at end of file +# Update setup keys (only if they are not present) +zk f cargo run $GPU_FLAG --release --bin key_generator -- $GENERATE_SK_COMMAND all --recompute-if-missing diff --git a/prover/vk_setup_data_generator_server_fri/src/keystore.rs b/prover/vk_setup_data_generator_server_fri/src/keystore.rs index dcbb7af96038..ccb729a1cf1c 100644 --- a/prover/vk_setup_data_generator_server_fri/src/keystore.rs +++ b/prover/vk_setup_data_generator_server_fri/src/keystore.rs @@ -1,4 +1,8 @@ -use std::{fs, fs::File, io::Read}; +use std::{ + fs::{self, File}, + io::Read, + path::Path, +}; use anyhow::Context as _; use circuit_definitions::{ @@ -293,6 +297,10 @@ impl Keystore { }) } + pub fn is_setup_data_present(&self, key: &ProverServiceDataKey) -> bool { + Path::new(&self.get_file_path(key.clone(), ProverServiceDataType::SetupData)).exists() + } + pub fn save_setup_data_for_circuit_type( &self, key: ProverServiceDataKey, diff --git a/prover/vk_setup_data_generator_server_fri/src/main.rs b/prover/vk_setup_data_generator_server_fri/src/main.rs index dadbd4b72f9e..bf5f4a86bb63 100644 --- a/prover/vk_setup_data_generator_server_fri/src/main.rs +++ b/prover/vk_setup_data_generator_server_fri/src/main.rs @@ -104,6 +104,11 @@ struct GeneratorOptions { #[arg(long, default_value = "false")] dry_run: bool, + /// If true, then generate the setup key, only if it is missing. + // Warning: this doesn't check the correctness of the existing file. + #[arg(long, default_value = "false")] + recompute_if_missing: bool, + #[arg(long)] path: Option, @@ -170,7 +175,7 @@ fn generate_setup_keys( ) -> anyhow::Result<()> { let circuit_type = match options.circuits_type { CircuitSelector::All => { - let digests = generator.generate_all(options.dry_run)?; + let digests = generator.generate_all(options.dry_run, options.recompute_if_missing)?; tracing::info!("Setup keys md5(s):"); print_stats(digests)?; return Ok(()); @@ -189,7 +194,7 @@ fn generate_setup_keys( }; let digest = generator - .generate_and_write_setup_data(circuit_type, options.dry_run) + .generate_and_write_setup_data(circuit_type, options.dry_run, options.recompute_if_missing) .context("generate_setup_data()")?; tracing::info!("digest: {:?}", digest); Ok(()) diff --git a/prover/vk_setup_data_generator_server_fri/src/setup_data_generator.rs b/prover/vk_setup_data_generator_server_fri/src/setup_data_generator.rs index 371bfbe0327b..3b2724251846 100644 --- a/prover/vk_setup_data_generator_server_fri/src/setup_data_generator.rs +++ b/prover/vk_setup_data_generator_server_fri/src/setup_data_generator.rs @@ -85,7 +85,15 @@ pub trait SetupDataGenerator { &self, circuit: ProverServiceDataKey, dry_run: bool, + recompute_if_missing: bool, ) -> anyhow::Result { + if recompute_if_missing && self.keystore().is_setup_data_present(&circuit) { + tracing::info!( + "Skipping setup computation for {:?} as file is already present.", + circuit.name(), + ); + return Ok("Skipped".to_string()); + } let serialized = self.generate_setup_data(circuit.clone())?; let digest = md5::compute(&serialized); @@ -100,12 +108,16 @@ pub trait SetupDataGenerator { } /// Generate all setup keys for boojum circuits. - fn generate_all(&self, dry_run: bool) -> anyhow::Result> { + fn generate_all( + &self, + dry_run: bool, + recompute_if_missing: bool, + ) -> anyhow::Result> { Ok(ProverServiceDataKey::all_boojum() .iter() .map(|circuit| { let digest = self - .generate_and_write_setup_data(circuit.clone(), dry_run) + .generate_and_write_setup_data(circuit.clone(), dry_run, recompute_if_missing) .context(circuit.name()) .unwrap(); (circuit.name(), digest)