From eaaf614eaa9fc0572b7898d5731c19039b83c57f Mon Sep 17 00:00:00 2001 From: Marko Atanasievski Date: Thu, 25 Jul 2024 15:14:26 +0200 Subject: [PATCH] feat: fix aggregation add segmentation ci test --- .github/workflows/ci.yml | 21 +++++++++++++++--- proof_gen/src/proof_gen.rs | 8 +++---- proof_gen/src/proof_types.rs | 32 +++++++++++++------------- zero_bin/ops/src/lib.rs | 20 ++++++++--------- zero_bin/prover/src/lib.rs | 42 +++++++++++++++++------------------ zero_bin/tools/prove_rpc.sh | 8 ++++++- zero_bin/tools/prove_stdio.sh | 26 +++++++++++++--------- 7 files changed, 92 insertions(+), 65 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index da2a3d615..f9bd28607 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -176,6 +176,20 @@ jobs: CARGO_INCREMENTAL: 1 RUST_BACKTRACE: 1 + simple_proof_witness_only: + name: Execute bash script to generate the proof witness for a small block. + runs-on: zero-ci + + steps: + - name: Checkout code + uses: actions/checkout@v3 + + - name: Run the script + run: | + pushd zero_bin/tools + ./prove_stdio.sh artifacts/witness_b19240705.json test_only + + simple_proof_regular: name: Execute bash script to generate and verify a proof for a small block. runs-on: zero-ci @@ -189,8 +203,8 @@ jobs: pushd zero_bin/tools ./prove_stdio.sh artifacts/witness_b19240705.json - simple_proof_witness_only: - name: Execute bash script to generate the proof witness for a small block. + simple_proof_using_continuations: + name: Execute bash script to generate and verify a proof for a small block utilizing custom batch and segment chunk size runs-on: zero-ci steps: @@ -200,7 +214,8 @@ jobs: - name: Run the script run: | pushd zero_bin/tools - ./prove_stdio.sh artifacts/witness_b19240705.json test_only + time PROVER_BATCH_SIZE=5 PROVER_SEGMENT_CHUNK_SIZE=4 PROVER_MAX_CPU_LEN_LOG=17 ./prove_stdio.sh ./artifacts/witness_b19240705.json + multi_blocks_proof_regular: name: Execute bash script to generate and verify a proof for multiple blocks using parallel proving. diff --git a/proof_gen/src/proof_gen.rs b/proof_gen/src/proof_gen.rs index cd6d48a52..c387460d6 100644 --- a/proof_gen/src/proof_gen.rs +++ b/proof_gen/src/proof_gen.rs @@ -16,8 +16,8 @@ use plonky2::{ use crate::{ proof_types::{ - GeneratedBlockProof, GeneratedSegmentAggProof, GeneratedSegmentProof, GeneratedTxnAggProof, - SegmentAggregatableProof, TxnAggregatableProof, + BatchAggregatableProof, GeneratedBlockProof, GeneratedSegmentAggProof, + GeneratedSegmentProof, GeneratedTxnAggProof, SegmentAggregatableProof, }, prover_state::ProverState, types::{Field, PlonkyProofIntern, EXTENSION_DEGREE}, @@ -121,8 +121,8 @@ pub fn generate_segment_agg_proof( /// Note that the child proofs may be either transaction or aggregation proofs. pub fn generate_transaction_agg_proof( p_state: &ProverState, - lhs_child: &TxnAggregatableProof, - rhs_child: &TxnAggregatableProof, + lhs_child: &BatchAggregatableProof, + rhs_child: &BatchAggregatableProof, ) -> ProofGenResult { let (b_proof_intern, p_vals) = p_state .state diff --git a/proof_gen/src/proof_types.rs b/proof_gen/src/proof_types.rs index fea8f845f..9807f0b25 100644 --- a/proof_gen/src/proof_types.rs +++ b/proof_gen/src/proof_types.rs @@ -67,7 +67,7 @@ pub enum SegmentAggregatableProof { /// we can combine it into an agg proof. For these cases, we want to abstract /// away whether or not the proof was a txn or agg proof. #[derive(Clone, Debug, Deserialize, Serialize)] -pub enum TxnAggregatableProof { +pub enum BatchAggregatableProof { /// The underlying proof is a segment proof. It first needs to be aggregated /// with another segment proof, or a dummy one. Segment(GeneratedSegmentProof), @@ -100,28 +100,28 @@ impl SegmentAggregatableProof { } } -impl TxnAggregatableProof { +impl BatchAggregatableProof { pub(crate) fn public_values(&self) -> PublicValues { match self { - TxnAggregatableProof::Segment(info) => info.p_vals.clone(), - TxnAggregatableProof::Txn(info) => info.p_vals.clone(), - TxnAggregatableProof::Agg(info) => info.p_vals.clone(), + BatchAggregatableProof::Segment(info) => info.p_vals.clone(), + BatchAggregatableProof::Txn(info) => info.p_vals.clone(), + BatchAggregatableProof::Agg(info) => info.p_vals.clone(), } } pub(crate) fn is_agg(&self) -> bool { match self { - TxnAggregatableProof::Segment(_) => false, - TxnAggregatableProof::Txn(_) => false, - TxnAggregatableProof::Agg(_) => true, + BatchAggregatableProof::Segment(_) => false, + BatchAggregatableProof::Txn(_) => false, + BatchAggregatableProof::Agg(_) => true, } } pub(crate) fn intern(&self) -> &PlonkyProofIntern { match self { - TxnAggregatableProof::Segment(info) => &info.intern, - TxnAggregatableProof::Txn(info) => &info.intern, - TxnAggregatableProof::Agg(info) => &info.intern, + BatchAggregatableProof::Segment(info) => &info.intern, + BatchAggregatableProof::Txn(info) => &info.intern, + BatchAggregatableProof::Agg(info) => &info.intern, } } } @@ -138,23 +138,23 @@ impl From for SegmentAggregatableProof { } } -impl From for TxnAggregatableProof { +impl From for BatchAggregatableProof { fn from(v: GeneratedSegmentAggProof) -> Self { Self::Txn(v) } } -impl From for TxnAggregatableProof { +impl From for BatchAggregatableProof { fn from(v: GeneratedTxnAggProof) -> Self { Self::Agg(v) } } -impl From for TxnAggregatableProof { +impl From for BatchAggregatableProof { fn from(v: SegmentAggregatableProof) -> Self { match v { - SegmentAggregatableProof::Agg(agg) => TxnAggregatableProof::Txn(agg), - SegmentAggregatableProof::Seg(seg) => TxnAggregatableProof::Segment(seg), + SegmentAggregatableProof::Agg(agg) => BatchAggregatableProof::Txn(agg), + SegmentAggregatableProof::Seg(seg) => BatchAggregatableProof::Segment(seg), } } } diff --git a/zero_bin/ops/src/lib.rs b/zero_bin/ops/src/lib.rs index f2c025913..4ff93a93f 100644 --- a/zero_bin/ops/src/lib.rs +++ b/zero_bin/ops/src/lib.rs @@ -9,7 +9,7 @@ use paladin::{ use proof_gen::{ proof_gen::{generate_block_proof, generate_segment_agg_proof, generate_transaction_agg_proof}, proof_types::{ - GeneratedBlockProof, GeneratedTxnAggProof, SegmentAggregatableProof, TxnAggregatableProof, + BatchAggregatableProof, GeneratedBlockProof, GeneratedTxnAggProof, SegmentAggregatableProof, }, }; use serde::{Deserialize, Serialize}; @@ -234,23 +234,23 @@ impl Monoid for SegmentAggProof { } #[derive(Deserialize, Serialize, RemoteExecute)] -pub struct TxnAggProof { +pub struct BatchAggProof { pub save_inputs_on_error: bool, } -fn get_agg_proof_public_values(elem: TxnAggregatableProof) -> PublicValues { +fn get_agg_proof_public_values(elem: BatchAggregatableProof) -> PublicValues { match elem { - TxnAggregatableProof::Segment(info) => info.p_vals, - TxnAggregatableProof::Txn(info) => info.p_vals, - TxnAggregatableProof::Agg(info) => info.p_vals, + BatchAggregatableProof::Segment(info) => info.p_vals, + BatchAggregatableProof::Txn(info) => info.p_vals, + BatchAggregatableProof::Agg(info) => info.p_vals, } } -impl Monoid for TxnAggProof { - type Elem = TxnAggregatableProof; +impl Monoid for BatchAggProof { + type Elem = BatchAggregatableProof; fn combine(&self, a: Self::Elem, b: Self::Elem) -> Result { let lhs = match a { - TxnAggregatableProof::Segment(segment) => TxnAggregatableProof::from( + BatchAggregatableProof::Segment(segment) => BatchAggregatableProof::from( generate_segment_agg_proof( p_state(), &SegmentAggregatableProof::from(segment.clone()), @@ -263,7 +263,7 @@ impl Monoid for TxnAggProof { }; let rhs = match b { - TxnAggregatableProof::Segment(segment) => TxnAggregatableProof::from( + BatchAggregatableProof::Segment(segment) => BatchAggregatableProof::from( generate_segment_agg_proof( p_state(), &SegmentAggregatableProof::from(segment.clone()), diff --git a/zero_bin/prover/src/lib.rs b/zero_bin/prover/src/lib.rs index e40c4ff2a..bc109ec35 100644 --- a/zero_bin/prover/src/lib.rs +++ b/zero_bin/prover/src/lib.rs @@ -77,21 +77,21 @@ impl BlockProverInput { ); // Create segment proof - let seg_ops = ops::SegmentProof { + let seg_prove_ops = ops::SegmentProof { save_inputs_on_error, }; // Generate segment data. - let agg_ops = ops::SegmentAggProof { + let seg_agg_ops = ops::SegmentAggProof { save_inputs_on_error, }; - // Aggregate transaction proofs - let txn_agg_proof = ops::TxnAggProof { + // Aggregate batch proofs to a single proof + let batch_agg_ops = ops::BatchAggProof { save_inputs_on_error, }; - let mut all_block_txn_aggregatable_proofs = Vec::new(); + let mut all_block_batch_proofs = Vec::new(); // Loop for all generation inputs in the block for generation_inputs in block_generation_inputs { let mut segment_data_iter = SegmentDataIterator { @@ -102,42 +102,42 @@ impl BlockProverInput { let mut chunk_segment_iter = SegmentDataChunkIterator::new(&mut segment_data_iter, segment_chunk_size); - let mut chunk_txn_aggregatable_proofs = Vec::new(); + let mut chunk_seg_aggregatable_proofs = Vec::new(); // We take one chunk of segments, perform proving and - // aggregate it to `TxnAggregatableProof` + // aggregate it into a single chunk proof while let Some(chunk) = chunk_segment_iter.next() { - chunk_txn_aggregatable_proofs.push( - Directive::map(IndexedStream::from(chunk.into_iter()), &seg_ops) - .fold(&agg_ops) + chunk_seg_aggregatable_proofs.push( + Directive::map(IndexedStream::from(chunk.into_iter()), &seg_prove_ops) + .fold(&seg_agg_ops) .run(runtime) - .map(move |e| e.map(proof_gen::proof_types::TxnAggregatableProof::from)) .await, ); } - // Fold all the generation input transaction proofs - // into a single transaction proof - let generation_input_txn_proof = Directive::fold( - IndexedStream::from(chunk_txn_aggregatable_proofs.into_iter().collect::, anyhow::Error, >>( )?), - &txn_agg_proof, + &seg_agg_ops, ) .run(runtime) + .map(move |e| e.map(proof_gen::proof_types::BatchAggregatableProof::from)) .await?; - all_block_txn_aggregatable_proofs.push(generation_input_txn_proof); + all_block_batch_proofs.push(batch_aggregated_proof); } - // Fold all the agg transaction proofs into a single transaction proof + // Fold all the aggregated batch proofs into a single batch proof let final_txn_proof = Directive::fold( - IndexedStream::from(all_block_txn_aggregatable_proofs.into_iter()), - &txn_agg_proof, + IndexedStream::from(all_block_batch_proofs.into_iter()), + &batch_agg_ops, ) .run(runtime) .await?; - if let proof_gen::proof_types::TxnAggregatableProof::Agg(proof) = final_txn_proof { + if let proof_gen::proof_types::BatchAggregatableProof::Agg(proof) = final_txn_proof { let block_number = block_number .to_u64() .context("block number overflows u64")?; diff --git a/zero_bin/tools/prove_rpc.sh b/zero_bin/tools/prove_rpc.sh index 477698301..40aebc4f4 100755 --- a/zero_bin/tools/prove_rpc.sh +++ b/zero_bin/tools/prove_rpc.sh @@ -74,6 +74,12 @@ else PROVER_SAVE_INPUTS_ON_ERROR="" fi +if [ -n "$NUM_WORKERS" ]; then + SET_NUM_WORKERS="--num-workers $NUM_WORKERS" +else + SET_NUM_WORKERS="" +fi + mkdir -p $PROOF_OUTPUT_DIR if [ $IGNORE_PREVIOUS_PROOFS ]; then @@ -132,7 +138,7 @@ if [[ $8 == "test_only" ]]; then else # normal run echo "Proving blocks ${BLOCK_INTERVAL} now... (Total: ${TOT_BLOCKS})" - command='cargo r --release --bin leader -- --runtime in-memory --load-strategy on-demand --batch-size $PROVER_BATCH_SIZE --max-cpu-len-log $PROVER_MAC_CPU_LEN_LOG --segment-chunk-size $PROVER_SEGMENT_CHUNK_SIZE $PROVER_SAVE_INPUTS_ON_ERROR rpc --rpc-type "$NODE_RPC_TYPE" --rpc-url "$3" --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 -- --runtime in-memory --load-strategy on-demand --batch-size $PROVER_BATCH_SIZE --max-cpu-len-log $PROVER_MAC_CPU_LEN_LOG --segment-chunk-size $PROVER_SEGMENT_CHUNK_SIZE $SET_NUM_WORKERS $PROVER_SAVE_INPUTS_ON_ERROR rpc --rpc-type "$NODE_RPC_TYPE" --rpc-url "$3" --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 echo -e "Proof generation finished with result: $?" diff --git a/zero_bin/tools/prove_stdio.sh b/zero_bin/tools/prove_stdio.sh index 86b4d8bf8..c157aa4a9 100755 --- a/zero_bin/tools/prove_stdio.sh +++ b/zero_bin/tools/prove_stdio.sh @@ -57,14 +57,14 @@ else # These sizes are configured specifically for block 19240705. Don't use this in other scenarios. echo "Using specific circuit sizes for witness_b19240705.json" export ARITHMETIC_CIRCUIT_SIZE="16..19" - export BYTE_PACKING_CIRCUIT_SIZE="11..15" + export BYTE_PACKING_CIRCUIT_SIZE="10..15" export CPU_CIRCUIT_SIZE="18..21" - export KECCAK_CIRCUIT_SIZE="14..18" - export KECCAK_SPONGE_CIRCUIT_SIZE="9..13" - export LOGIC_CIRCUIT_SIZE="12..17" - export MEMORY_CIRCUIT_SIZE="20..23" - export MEMORY_BEFORE_CIRCUIT_SIZE="16..17" - export MEMORY_AFTER_CIRCUIT_SIZE="7..8" + export KECCAK_CIRCUIT_SIZE="12..18" + export KECCAK_SPONGE_CIRCUIT_SIZE="7..13" + export LOGIC_CIRCUIT_SIZE="10..17" + export MEMORY_CIRCUIT_SIZE="19..23" + export MEMORY_BEFORE_CIRCUIT_SIZE="15..19" + export MEMORY_AFTER_CIRCUIT_SIZE="7..19" elif [[ $INPUT_FILE == *"witness_b2_b7"* ]]; then # These sizes are configured specifically for custom small blocks. Don't use this in other scenarios. echo "Using specific circuit sizes for witness_b2_b7.json" @@ -94,13 +94,19 @@ fi # Prover config. Override the defaults if needed by setting the env variables. PROVER_BATCH_SIZE="${PROVER_BATCH_SIZE:-1}" PROVER_SEGMENT_CHUNK_SIZE="${PROVER_SEGMENT_CHUNK_SIZE:-64}" -PROVER_MAC_CPU_LEN_LOG="${PROVER_MAC_CPU_LEN_LOG:-20}" +PROVER_MAX_CPU_LEN_LOG="${PROVER_MAX_CPU_LEN_LOG:-20}" if [[ $PROVER_SAVE_INPUTS_ON_ERROR == "true" ]]; then PROVER_SAVE_INPUTS_ON_ERROR="--save-inputs-on-error" else PROVER_SAVE_INPUTS_ON_ERROR="" fi +if [ -n "$NUM_WORKERS" ]; then + SET_NUM_WORKERS="--num-workers $NUM_WORKERS" +else + SET_NUM_WORKERS="" +fi + # If we run ./prove_stdio.sh test_only, we'll generate a dummy @@ -121,8 +127,8 @@ cargo build --release --jobs "$num_procs" start_time=$(date +%s%N) "${TOOLS_DIR}/../../target/release/leader" --runtime in-memory --load-strategy on-demand --batch-size $PROVER_BATCH_SIZE \ - --max-cpu-len-log $PROVER_MAC_CPU_LEN_LOG --segment-chunk-size $PROVER_SEGMENT_CHUNK_SIZE \ - $PROVER_SAVE_INPUTS_ON_ERROR stdio < $INPUT_FILE | tee $LEADER_OUT_PATH + --max-cpu-len-log $PROVER_MAX_CPU_LEN_LOG --segment-chunk-size $PROVER_SEGMENT_CHUNK_SIZE \ + $SET_NUM_WORKERS $PROVER_SAVE_INPUTS_ON_ERROR stdio < $INPUT_FILE | tee $LEADER_OUT_PATH end_time=$(date +%s%N) tail -n 1 $LEADER_OUT_PATH > $PROOFS_JSON_PATH