From a0dcf0e3cb5fd843f4c5e9a8a4ba78dd0fff03bb Mon Sep 17 00:00:00 2001 From: Hamy Ratoanina Date: Wed, 21 Aug 2024 14:43:25 -0400 Subject: [PATCH] Make test-only a CLI argument (#521) * Make test-only a CLI argument * Delete unused struct * Update cargo files * Apply offline comments * Send dummy proof --- zero_bin/leader/Cargo.toml | 1 - zero_bin/leader/src/client.rs | 72 ++++++++++++++++++----------------- zero_bin/leader/src/http.rs | 30 ++++++++++----- zero_bin/leader/src/main.rs | 26 +++++++------ zero_bin/leader/src/stdio.rs | 2 +- zero_bin/ops/Cargo.toml | 1 - zero_bin/ops/src/lib.rs | 17 +++------ zero_bin/prover/Cargo.toml | 1 - zero_bin/prover/src/cli.rs | 5 +++ zero_bin/prover/src/lib.rs | 72 ++++++++++++++++++++++------------- zero_bin/tools/prove_rpc.sh | 16 ++------ zero_bin/tools/prove_stdio.sh | 16 ++------ 12 files changed, 136 insertions(+), 123 deletions(-) diff --git a/zero_bin/leader/Cargo.toml b/zero_bin/leader/Cargo.toml index 5d0881da1..7f3655961 100644 --- a/zero_bin/leader/Cargo.toml +++ b/zero_bin/leader/Cargo.toml @@ -34,7 +34,6 @@ zero_bin_common = { workspace = true } [features] default = [] -test_only = ["ops/test_only", "prover/test_only"] [build-dependencies] cargo_metadata = { workspace = true } diff --git a/zero_bin/leader/src/client.rs b/zero_bin/leader/src/client.rs index 74910f621..9510b1d75 100644 --- a/zero_bin/leader/src/client.rs +++ b/zero_bin/leader/src/client.rs @@ -63,47 +63,49 @@ pub(crate) async fn client_main( runtime.close().await?; let proved_blocks = proved_blocks?; - if cfg!(feature = "test_only") { + if params.prover_config.test_only { info!("All proof witnesses have been generated successfully."); } else { info!("All proofs have been generated successfully."); } - if params.keep_intermediate_proofs { - if params.proof_output_dir.is_some() { - // All proof files (including intermediary) are written to disk and kept - warn!("Skipping cleanup, intermediate proof files are kept"); + if !params.prover_config.test_only { + if params.keep_intermediate_proofs { + if params.proof_output_dir.is_some() { + // All proof files (including intermediary) are written to disk and kept + warn!("Skipping cleanup, intermediate proof files are kept"); + } else { + // Output all proofs to stdout + std::io::stdout().write_all(&serde_json::to_vec( + &proved_blocks + .into_iter() + .filter_map(|(_, block)| block) + .collect::>(), + )?)?; + } + } else if let Some(proof_output_dir) = params.proof_output_dir.as_ref() { + // Remove intermediary proof files + proved_blocks + .into_iter() + .rev() + .skip(1) + .map(|(block_number, _)| { + generate_block_proof_file_name(&proof_output_dir.to_str(), block_number) + }) + .for_each(|path| { + if let Err(e) = std::fs::remove_file(path) { + error!("Failed to remove intermediate proof file: {e}"); + } + }); } else { - // Output all proofs to stdout - std::io::stdout().write_all(&serde_json::to_vec( - &proved_blocks - .into_iter() - .filter_map(|(_, block)| block) - .collect::>(), - )?)?; - } - } else if let Some(proof_output_dir) = params.proof_output_dir.as_ref() { - // Remove intermediary proof files - proved_blocks - .into_iter() - .rev() - .skip(1) - .map(|(block_number, _)| { - generate_block_proof_file_name(&proof_output_dir.to_str(), block_number) - }) - .for_each(|path| { - if let Err(e) = std::fs::remove_file(path) { - error!("Failed to remove intermediate proof file: {e}"); - } - }); - } else { - // Output only last proof to stdout - if let Some(last_block) = proved_blocks - .into_iter() - .filter_map(|(_, block)| block) - .last() - { - std::io::stdout().write_all(&serde_json::to_vec(&last_block)?)?; + // Output only last proof to stdout + if let Some(last_block) = proved_blocks + .into_iter() + .filter_map(|(_, block)| block) + .last() + { + std::io::stdout().write_all(&serde_json::to_vec(&last_block)?)?; + } } } diff --git a/zero_bin/leader/src/http.rs b/zero_bin/leader/src/http.rs index 9137622be..39c7333e1 100644 --- a/zero_bin/leader/src/http.rs +++ b/zero_bin/leader/src/http.rs @@ -71,15 +71,27 @@ async fn prove( let block_number = payload.prover_input.get_block_number(); - match payload - .prover_input - .prove( - &runtime, - payload.previous.map(futures::future::ok), - prover_config, - ) - .await - { + let proof_res = if prover_config.test_only { + payload + .prover_input + .prove_test( + &runtime, + payload.previous.map(futures::future::ok), + prover_config, + ) + .await + } else { + payload + .prover_input + .prove( + &runtime, + payload.previous.map(futures::future::ok), + prover_config, + ) + .await + }; + + match proof_res { Ok(b_proof) => match write_to_file(output_dir, block_number, &b_proof) { Ok(file) => { info!("Successfully wrote proof to {}", file.display()); diff --git a/zero_bin/leader/src/main.rs b/zero_bin/leader/src/main.rs index 7f28cc31c..32c9baa20 100644 --- a/zero_bin/leader/src/main.rs +++ b/zero_bin/leader/src/main.rs @@ -9,6 +9,7 @@ use dotenvy::dotenv; use ops::register; use paladin::runtime::Runtime; use proof_gen::proof_types::GeneratedBlockProof; +use prover::ProverConfig; use tracing::{info, warn}; use zero_bin_common::{ block_interval::BlockInterval, prover_state::persistence::set_circuit_cache_dir_env_if_not_set, @@ -53,22 +54,25 @@ async fn main() -> Result<()> { } let args = cli::Cli::parse(); - if let paladin::config::Runtime::InMemory = args.paladin.runtime { - // If running in emulation mode, we'll need to initialize the prover - // state here. - args.prover_state_config - .into_prover_state_manager() - .initialize()?; - } let runtime = Runtime::from_config(&args.paladin, register()).await?; - let cli_prover_config = args.prover_config; + let prover_config: ProverConfig = args.prover_config.into(); + + // If not in test_only mode and running in emulation mode, we'll need to + // initialize the prover state here. + if !prover_config.test_only { + if let paladin::config::Runtime::InMemory = args.paladin.runtime { + args.prover_state_config + .into_prover_state_manager() + .initialize()?; + } + } match args.command { Command::Stdio { previous_proof } => { let previous_proof = get_previous_proof(previous_proof)?; - stdio::stdio_main(runtime, previous_proof, cli_prover_config.into()).await?; + stdio::stdio_main(runtime, previous_proof, prover_config).await?; } Command::Http { port, output_dir } => { // check if output_dir exists, is a directory, and is writable @@ -80,7 +84,7 @@ async fn main() -> Result<()> { panic!("output-dir is not a writable directory"); } - http::http_main(runtime, port, output_dir, cli_prover_config.into()).await?; + http::http_main(runtime, port, output_dir, prover_config).await?; } Command::Rpc { rpc_url, @@ -120,7 +124,7 @@ async fn main() -> Result<()> { checkpoint_block_number, previous_proof, proof_output_dir, - prover_config: cli_prover_config.into(), + prover_config, keep_intermediate_proofs, }, ) diff --git a/zero_bin/leader/src/stdio.rs b/zero_bin/leader/src/stdio.rs index d74f4dce6..0acaf88ad 100644 --- a/zero_bin/leader/src/stdio.rs +++ b/zero_bin/leader/src/stdio.rs @@ -26,7 +26,7 @@ pub(crate) async fn stdio_main( runtime.close().await?; let proved_blocks = proved_blocks?; - if cfg!(feature = "test_only") { + if prover_config.test_only { info!("All proof witnesses have been generated successfully."); } else { info!("All proofs have been generated successfully."); diff --git a/zero_bin/ops/Cargo.toml b/zero_bin/ops/Cargo.toml index 3d2742f3c..8975cc8d5 100644 --- a/zero_bin/ops/Cargo.toml +++ b/zero_bin/ops/Cargo.toml @@ -21,4 +21,3 @@ zero_bin_common = { path = "../common" } [features] default = [] -test_only = [] diff --git a/zero_bin/ops/src/lib.rs b/zero_bin/ops/src/lib.rs index 286ae4bca..4de5e0c30 100644 --- a/zero_bin/ops/src/lib.rs +++ b/zero_bin/ops/src/lib.rs @@ -1,16 +1,12 @@ -#[cfg(not(feature = "test_only"))] use std::time::Instant; -#[cfg(not(feature = "test_only"))] use evm_arithmetization::generation::TrimmedGenerationInputs; use evm_arithmetization::proof::PublicValues; -#[cfg(feature = "test_only")] use evm_arithmetization::{prover::testing::simulate_execution_all_segments, GenerationInputs}; use paladin::{ operation::{FatalError, FatalStrategy, Monoid, Operation, Result}, registry, RemoteExecute, }; -#[cfg(feature = "test_only")] use proof_gen::types::Field; use proof_gen::{ proof_gen::{generate_block_proof, generate_segment_agg_proof, generate_transaction_agg_proof}, @@ -20,7 +16,6 @@ use proof_gen::{ }; use serde::{Deserialize, Serialize}; use tracing::error; -#[cfg(not(feature = "test_only"))] use tracing::{event, info_span, Level}; use zero_bin_common::{debug_utils::save_inputs_to_disk, prover_state::p_state}; @@ -31,7 +26,6 @@ pub struct SegmentProof { pub save_inputs_on_error: bool, } -#[cfg(not(feature = "test_only"))] impl Operation for SegmentProof { type Input = evm_arithmetization::AllData; type Output = proof_gen::proof_types::SegmentAggregatableProof; @@ -72,8 +66,12 @@ impl Operation for SegmentProof { } } -#[cfg(feature = "test_only")] -impl Operation for SegmentProof { +#[derive(Deserialize, Serialize, RemoteExecute)] +pub struct SegmentProofTestOnly { + pub save_inputs_on_error: bool, +} + +impl Operation for SegmentProofTestOnly { type Input = (GenerationInputs, usize); type Output = (); @@ -107,14 +105,12 @@ impl Operation for SegmentProof { /// /// - When created, it starts a span with the transaction proof id. /// - When dropped, it logs the time taken by the transaction proof. -#[cfg(not(feature = "test_only"))] struct SegmentProofSpan { _span: tracing::span::EnteredSpan, start: Instant, descriptor: String, } -#[cfg(not(feature = "test_only"))] impl SegmentProofSpan { /// Get a unique id for the transaction proof. fn get_id(ir: &TrimmedGenerationInputs, segment_index: usize) -> String { @@ -174,7 +170,6 @@ impl SegmentProofSpan { } } -#[cfg(not(feature = "test_only"))] impl Drop for SegmentProofSpan { fn drop(&mut self) { event!( diff --git a/zero_bin/prover/Cargo.toml b/zero_bin/prover/Cargo.toml index 0266421d5..9f2050746 100644 --- a/zero_bin/prover/Cargo.toml +++ b/zero_bin/prover/Cargo.toml @@ -30,4 +30,3 @@ clap = {workspace = true} [features] default = [] -test_only = ["ops/test_only"] diff --git a/zero_bin/prover/src/cli.rs b/zero_bin/prover/src/cli.rs index 885d7456d..a49eb2d1d 100644 --- a/zero_bin/prover/src/cli.rs +++ b/zero_bin/prover/src/cli.rs @@ -14,6 +14,10 @@ pub struct CliProverConfig { /// If true, save the public inputs to disk on error. #[arg(short='i', long, help_heading = HELP_HEADING, default_value_t = false)] save_inputs_on_error: bool, + /// If true, only test the trace decoder and witness generation without + /// generating a proof. + #[arg(long, help_heading = HELP_HEADING, default_value_t = false)] + test_only: bool, } impl From for crate::ProverConfig { @@ -22,6 +26,7 @@ impl From for crate::ProverConfig { batch_size: cli.batch_size, max_cpu_len_log: cli.max_cpu_len_log, save_inputs_on_error: cli.save_inputs_on_error, + test_only: cli.test_only, } } } diff --git a/zero_bin/prover/src/lib.rs b/zero_bin/prover/src/lib.rs index f96039edf..d5d62f430 100644 --- a/zero_bin/prover/src/lib.rs +++ b/zero_bin/prover/src/lib.rs @@ -21,6 +21,7 @@ pub struct ProverConfig { pub batch_size: usize, pub max_cpu_len_log: usize, pub save_inputs_on_error: bool, + pub test_only: bool, } #[derive(Debug, Deserialize, Serialize)] @@ -34,7 +35,6 @@ impl BlockProverInput { self.other_data.b_data.b_meta.block_number.into() } - #[cfg(not(feature = "test_only"))] pub async fn prove( self, runtime: &Runtime, @@ -50,6 +50,7 @@ impl BlockProverInput { max_cpu_len_log, batch_size, save_inputs_on_error, + test_only: _, } = prover_config; let block_number = self.get_block_number(); @@ -123,8 +124,7 @@ impl BlockProverInput { } } - #[cfg(feature = "test_only")] - pub async fn prove( + pub async fn prove_test( self, runtime: &Runtime, previous: Option>>, @@ -139,6 +139,7 @@ impl BlockProverInput { max_cpu_len_log, batch_size, save_inputs_on_error, + test_only: _, } = prover_config; let block_number = self.get_block_number(); @@ -147,7 +148,7 @@ impl BlockProverInput { let block_generation_inputs = trace_decoder::entrypoint(self.block_trace, self.other_data, batch_size)?; - let batch_ops = ops::SegmentProof { + let seg_ops = ops::SegmentProofTestOnly { save_inputs_on_error, }; @@ -157,7 +158,7 @@ impl BlockProverInput { .into_iter() .zip(repeat(max_cpu_len_log)), ), - &batch_ops, + &seg_ops, ); simulation @@ -211,28 +212,45 @@ impl ProverInput { // Prove the block let proof_output_dir = proof_output_dir.clone(); - let fut = block - .prove(runtime, prev.take(), prover_config) - .then(move |proof| async move { - let proof = proof?; - let block_number = proof.b_height; - - // Write latest generated proof to disk if proof_output_dir is provided - let return_proof: Option = - if proof_output_dir.is_some() { - ProverInput::write_proof(proof_output_dir, &proof).await?; - None - } else { - Some(proof.clone()) - }; - - if tx.send(proof).is_err() { - anyhow::bail!("Failed to send proof"); - } - - Ok((block_number, return_proof)) - }) - .boxed(); + let fut = if prover_config.test_only { + block + .prove_test(runtime, prev.take(), prover_config) + .then(move |proof| async move { + let proof = proof?; + let block_number = proof.b_height; + + if tx.send(proof).is_err() { + anyhow::bail!("Failed to send proof"); + } + + // We ignore the returned dummy proof in test-only mode. + Ok((block_number, None)) + }) + .boxed() + } else { + block + .prove(runtime, prev.take(), prover_config) + .then(move |proof| async move { + let proof = proof?; + let block_number = proof.b_height; + + // Write latest generated proof to disk if proof_output_dir is provided. + let return_proof: Option = + if proof_output_dir.is_some() { + ProverInput::write_proof(proof_output_dir, &proof).await?; + None + } else { + Some(proof.clone()) + }; + + if tx.send(proof).is_err() { + anyhow::bail!("Failed to send proof"); + } + + Ok((block_number, return_proof)) + }) + .boxed() + }; prev = Some(Box::pin(rx.map_err(anyhow::Error::new))); diff --git a/zero_bin/tools/prove_rpc.sh b/zero_bin/tools/prove_rpc.sh index 539bb71f8..9c0f4bad9 100755 --- a/zero_bin/tools/prove_rpc.sh +++ b/zero_bin/tools/prove_rpc.sh @@ -17,18 +17,8 @@ export RUST_LOG=info # See also .cargo/config.toml. export RUSTFLAGS='-C target-cpu=native -Zlinker-features=-lld' -if [[ $8 == "test_only" ]]; then - # Circuit sizes don't matter in test_only mode, so we keep them minimal. - export ARITHMETIC_CIRCUIT_SIZE="16..17" - export BYTE_PACKING_CIRCUIT_SIZE="9..10" - export CPU_CIRCUIT_SIZE="12..13" - export KECCAK_CIRCUIT_SIZE="4..5" - export KECCAK_SPONGE_CIRCUIT_SIZE="9..10" - export LOGIC_CIRCUIT_SIZE="12..13" - export MEMORY_CIRCUIT_SIZE="17..18" - export MEMORY_BEFORE_CIRCUIT_SIZE="7..8" - export MEMORY_AFTER_CIRCUIT_SIZE="7..8" -else +# Circuit sizes only matter in non test_only mode. +if ! [[ $8 == "test_only" ]]; then export ARITHMETIC_CIRCUIT_SIZE="16..21" export BYTE_PACKING_CIRCUIT_SIZE="8..21" export CPU_CIRCUIT_SIZE="8..21" @@ -112,7 +102,7 @@ fi if [[ $8 == "test_only" ]]; then # test only run echo "Proving blocks ${BLOCK_INTERVAL} in a test_only mode now... (Total: ${TOT_BLOCKS})" - command='cargo r --release --features test_only --bin leader -- --runtime in-memory --load-strategy on-demand rpc --rpc-type "$NODE_RPC_TYPE" --rpc-url "$NODE_RPC_URL" --block-interval $BLOCK_INTERVAL --proof-output-dir $PROOF_OUTPUT_DIR $PREV_PROOF_EXTRA_ARG --backoff "$BACKOFF" --max-retries "$RETRIES" ' + command='cargo r --release --bin leader -- --test-only --runtime in-memory --load-strategy on-demand rpc --rpc-type "$NODE_RPC_TYPE" --rpc-url "$NODE_RPC_URL" --block-interval $BLOCK_INTERVAL --proof-output-dir $PROOF_OUTPUT_DIR $PREV_PROOF_EXTRA_ARG --backoff "$BACKOFF" --max-retries "$RETRIES" ' if [ "$OUTPUT_TO_TERMINAL" = true ]; then eval $command retVal=$? diff --git a/zero_bin/tools/prove_stdio.sh b/zero_bin/tools/prove_stdio.sh index 80ef76fbd..64c140023 100755 --- a/zero_bin/tools/prove_stdio.sh +++ b/zero_bin/tools/prove_stdio.sh @@ -44,18 +44,8 @@ if [[ $INPUT_FILE == "" ]]; then exit 1 fi -if [[ $TEST_ONLY == "test_only" ]]; then - # Circuit sizes don't matter in test_only mode, so we keep them minimal. - export ARITHMETIC_CIRCUIT_SIZE="16..17" - export BYTE_PACKING_CIRCUIT_SIZE="9..10" - export CPU_CIRCUIT_SIZE="12..13" - export KECCAK_CIRCUIT_SIZE="4..5" - export KECCAK_SPONGE_CIRCUIT_SIZE="9..10" - export LOGIC_CIRCUIT_SIZE="12..13" - export MEMORY_CIRCUIT_SIZE="17..18" - export MEMORY_BEFORE_CIRCUIT_SIZE="7..8" - export MEMORY_AFTER_CIRCUIT_SIZE="7..8" -else +# Circuit sizes only matter in non test_only mode. +if ! [[ $TEST_ONLY == "test_only" ]]; then if [[ $INPUT_FILE == *"witness_b19807080"* ]]; then # These sizes are configured specifically for block 19807080. Don't use this in other scenarios echo "Using specific circuit sizes for witness_b19807080.json" @@ -98,7 +88,7 @@ fi # proof. This is useful for quickly testing decoding and all of the # other non-proving code. if [[ $TEST_ONLY == "test_only" ]]; then - cargo run --release --features test_only --bin leader -- --runtime in-memory --load-strategy on-demand stdio < $INPUT_FILE &> $TEST_OUT_PATH + cargo run --release --bin leader -- --test-only --runtime in-memory --load-strategy on-demand stdio < $INPUT_FILE &> $TEST_OUT_PATH if grep -q 'All proof witnesses have been generated successfully.' $TEST_OUT_PATH; then echo -e "\n\nSuccess - Note this was just a test, not a proof" rm $TEST_OUT_PATH