From bc2ff6e8cdd19ae7968e4e9f1ac478c840ba3baa Mon Sep 17 00:00:00 2001 From: Marko Atanasievski Date: Tue, 4 Jun 2024 13:23:23 +0200 Subject: [PATCH] fix: proving logic --- common/src/block_interval.rs | 16 +++++- leader/src/cli.rs | 4 +- leader/src/jerigon.rs | 15 +++-- leader/src/main.rs | 4 +- prover/src/lib.rs | 62 +++++++++----------- tools/prove_blocks.sh | 106 ++++++++++++++++++++++++++--------- 6 files changed, 132 insertions(+), 75 deletions(-) diff --git a/common/src/block_interval.rs b/common/src/block_interval.rs index 427db47e..5493208d 100644 --- a/common/src/block_interval.rs +++ b/common/src/block_interval.rs @@ -105,7 +105,7 @@ impl BlockInterval { /// Convert the block interval into a stream of block numbers. pub fn into_bounded_stream(self) -> Result, BlockIntervalError> { match self { - BlockInterval::Single(num) => Ok(tokio_stream::iter(num..num)), + BlockInterval::Single(num) => Ok(tokio_stream::iter(num..num + 1)), BlockInterval::Range(range) => Ok(tokio_stream::iter(range)), _ => Err(BlockIntervalError::IntoBoundedStreamError), } @@ -238,4 +238,18 @@ mod test { ))) ) } + + #[tokio::test] + async fn can_into_bounded_stream() { + use tokio_stream::StreamExt; + let mut result = Vec::new(); + let mut stream = BlockInterval::new("1..10") + .unwrap() + .into_bounded_stream() + .unwrap(); + while let Some(val) = stream.next().await { + result.push(val); + } + assert_eq!(result, Vec::from_iter(1u64..10u64)); + } } diff --git a/leader/src/cli.rs b/leader/src/cli.rs index e5726552..e67aa207 100644 --- a/leader/src/cli.rs +++ b/leader/src/cli.rs @@ -46,10 +46,10 @@ pub(crate) enum Command { /// The previous proof output. #[arg(long, short = 'f', value_hint = ValueHint::FilePath)] previous_proof: Option, - /// If provided, write the generated proof to this file instead of + /// If provided, write the generated proofs to this directory instead of /// stdout. #[arg(long, short = 'o', value_hint = ValueHint::FilePath)] - proof_output_path: Option, + proof_output_dir: Option, /// If true, save the public inputs to disk on error. #[arg(short, long, default_value_t = false)] save_inputs_on_error: bool, diff --git a/leader/src/jerigon.rs b/leader/src/jerigon.rs index 23080e55..802bd218 100644 --- a/leader/src/jerigon.rs +++ b/leader/src/jerigon.rs @@ -18,7 +18,7 @@ pub(crate) async fn jerigon_main( block_interval: BlockInterval, checkpoint_block_number: u64, previous_proof: Option, - proof_output_path_opt: Option, + proof_output_dir_opt: Option, save_inputs_on_error: bool, ) -> Result<()> { let prover_input = rpc::prover_input( @@ -39,12 +39,11 @@ pub(crate) async fn jerigon_main( runtime.close().await?; for block_proof in block_proofs { - //todo fix proof dump to disk (for every block there it may be different file - // in one directory) + let blokck_proof_str = serde_json::to_vec(&block_proof.intern)?; write_proof( - serde_json::to_vec(&block_proof.intern)?, - proof_output_path_opt.clone().map(|mut path| { - path.push(format!("_block_{}", block_proof.b_height)); + blokck_proof_str, + proof_output_dir_opt.clone().map(|mut path| { + path.push(format!("b{}.zkproof", block_proof.b_height)); path }), )?; @@ -52,8 +51,8 @@ pub(crate) async fn jerigon_main( Ok(()) } -fn write_proof(proof: Vec, proof_output_path_opt: Option) -> Result<()> { - match proof_output_path_opt { +fn write_proof(proof: Vec, proof_output_dir_opt: Option) -> Result<()> { + match proof_output_dir_opt { Some(p) => { if let Some(parent) = p.parent() { create_dir_all(parent)?; diff --git a/leader/src/main.rs b/leader/src/main.rs index cd7d5911..2f6ba365 100644 --- a/leader/src/main.rs +++ b/leader/src/main.rs @@ -95,7 +95,7 @@ async fn main() -> Result<()> { block_interval, checkpoint_block_number, previous_proof, - proof_output_path, + proof_output_dir, save_inputs_on_error, } => { let previous_proof = get_previous_proof(previous_proof)?; @@ -118,7 +118,7 @@ async fn main() -> Result<()> { block_interval, checkpoint_block_number, previous_proof, - proof_output_path, + proof_output_dir, save_inputs_on_error, ) .await?; diff --git a/prover/src/lib.rs b/prover/src/lib.rs index c3fe1fc6..c27132d3 100644 --- a/prover/src/lib.rs +++ b/prover/src/lib.rs @@ -130,40 +130,43 @@ pub struct ProverInput { } impl ProverInput { - #[cfg(not(feature = "test_only"))] pub async fn prove( self, runtime: &Runtime, checkpoint_proof: Option, save_inputs_on_error: bool, ) -> Result> { - let results: Arc>> = + let results: Arc>>> = Arc::new(Mutex::new(HashMap::new())); if let Some(checkpoint_proof) = checkpoint_proof { results .lock() .await - .insert(checkpoint_proof.b_height, checkpoint_proof.intern); + .insert(checkpoint_proof.b_height, Some(checkpoint_proof.intern)); }; for block in self.blocks { - //todo this will be further reorganized with the new BlockProofFuture - + //todo this will be further reorganized with the new BlockProofFuture, + // running multiple block proof generation in parallel, and awaiting for the + // previous block proof when needed to prove the next block. + // For now prove blocks one by one sequentially and assume previous proof is + // already available let results = results.clone(); - // for now prove blocks one by one async move { - //todo handle genesis block case for previous block number - let previous_block_number = block - .get_block_number() + let block_number = block.get_block_number(); + info!("Proving block {block_number}"); + // For genesis block we don't have a previous proof, so + // previous_block_number would be None + let previous_block_number: Option = block_number .checked_sub(U256::from(1)) - .ok_or(anyhow::Error::msg("block number smaller than 1"))? - .to_u64(); - let previous_proof = if let Some(previous_block_number) = previous_block_number { - //todo optimize previous block proof clone - results.lock().await.get(&previous_block_number).cloned() + .and_then(|n| n.to_u64()); + let previous_proof = if let Some(number) = previous_block_number { + //todo could we optimize this to avoid this large proof cloning? + results.lock().await.get(&number).cloned().flatten() } else { None }; + // Prove the block let block_proof = block .prove(runtime, previous_proof, save_inputs_on_error) .await?; @@ -172,34 +175,25 @@ impl ProverInput { .b_height .to_u64() .ok_or(anyhow::Error::msg("block number u64 overflow"))?, - block_proof.intern, + Some(block_proof.intern), ); + info!("Proving block {block_number} finished!"); Result::<(), anyhow::Error>::Ok(()) } .await?; } let mut results = results.lock().await; - Ok(results + results .drain() - .map(|(block_number, intern)| GeneratedBlockProof { - b_height: block_number, - intern, + .map(|(block_number, intern)| { + Ok::(GeneratedBlockProof { + b_height: block_number, + intern: intern.ok_or_else(|| { + anyhow::Error::msg("missing proof for block {block_number}") + })?, + }) }) - .collect()) - } - - #[cfg(feature = "test_only")] - pub async fn prove( - self, - runtime: &Runtime, - _previous: Option, - save_inputs_on_error: bool, - ) -> Result> { - //todo test and update this - self.blocks - .into_iter() - .map(|block| block.prove(runtime, None, save_inputs_on_error))? - .collect() + .collect::, _>>() } } diff --git a/tools/prove_blocks.sh b/tools/prove_blocks.sh index c27317ef..21193acc 100755 --- a/tools/prove_blocks.sh +++ b/tools/prove_blocks.sh @@ -5,48 +5,94 @@ # 2 --> End block index (inclusive) # 3 --> Rpc endpoint:port (eg. http://35.246.1.96:8545) # 4 --> Ignore previous proofs (boolean) +# 5 --> Test only flag +export RUST_MIN_STACK=33554432 export RUST_BACKTRACE=1 -export RUST_LOG=mpt_trie=info,trace_decoder=info,plonky2=info,evm_arithmetization=trace,leader=info -export RUSTFLAGS='-Ctarget-cpu=native' +export RUST_LOG=info +# Disable the lld linker for now, as it's causing issues with the linkme package. +# https://github.com/rust-lang/rust/pull/124129 +# https://github.com/dtolnay/linkme/pull/88 +export RUSTFLAGS='-C target-cpu=native -Zlinker-features=-lld' -export ARITHMETIC_CIRCUIT_SIZE="16..23" -export BYTE_PACKING_CIRCUIT_SIZE="9..21" -export CPU_CIRCUIT_SIZE="12..25" -export KECCAK_CIRCUIT_SIZE="14..20" -export KECCAK_SPONGE_CIRCUIT_SIZE="9..15" -export LOGIC_CIRCUIT_SIZE="12..18" -export MEMORY_CIRCUIT_SIZE="17..28" +if [[ $5 == "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="14..15" + export KECCAK_SPONGE_CIRCUIT_SIZE="9..10" + export LOGIC_CIRCUIT_SIZE="12..13" + export MEMORY_CIRCUIT_SIZE="17..18" +else + export ARITHMETIC_CIRCUIT_SIZE="16..23" + export BYTE_PACKING_CIRCUIT_SIZE="9..21" + export CPU_CIRCUIT_SIZE="12..25" + export KECCAK_CIRCUIT_SIZE="14..20" + export KECCAK_SPONGE_CIRCUIT_SIZE="9..15" + export LOGIC_CIRCUIT_SIZE="12..18" + export MEMORY_CIRCUIT_SIZE="17..28" +fi PROOF_OUTPUT_DIR="proofs" +OUT_LOG_PATH="${PROOF_OUTPUT_DIR}/b${i}.log" ALWAYS_WRITE_LOGS=0 # Change this to `1` if you always want logs to be written. - TOT_BLOCKS=$(($2-$1+1)) + +START_BLOCK=$1 +END_BLOCK=$2 +NODE_RPC_URL=$3 IGNORE_PREVIOUS_PROOFS=$4 -echo "Proving blocks ${1}..=${2}... (Total: ${TOT_BLOCKS})" + mkdir -p $PROOF_OUTPUT_DIR -for ((i=$1; i<=$2; i++)) -do - echo "Proving block ${i}..." - OUT_PROOF_PATH="${PROOF_OUTPUT_DIR}/b${i}.zkproof" - OUT_LOG_PATH="${PROOF_OUTPUT_DIR}/b${i}.log" +if [ $IGNORE_PREVIOUS_PROOFS ]; then + # Set checkpoint height to previous block number for the first block in range + prev_proof_num=$(($1-1)) + PREV_PROOF_EXTRA_ARG="--checkpoint-block-number ${prev_proof_num}" +else + if [ $1 -gt 1 ]; then + prev_proof_num=$(($1-1)) + PREV_PROOF_EXTRA_ARG="-f ${PROOF_OUTPUT_DIR}/b${prev_proof_num}.zkproof" + fi +fi + +# Convert hex do decimal parameters +if [[ $START_BLOCK == 0x* ]]; then + START_BLOCK=$((16#${START_BLOCK#"0x"})) +fi +if [[ $END_BLOCK == 0x* ]]; then + END_BLOCK=$((16#${END_BLOCK#"0x"})) +fi + +if [ $START_BLOCK == $END_BLOCK ]; then + BLOCK_INTERVAL=$((16#${START_BLOCK#"0x"})) +else + BLOCK_INTERVAL=$START_BLOCK..=$END_BLOCK +fi - if [ $IGNORE_PREVIOUS_PROOFS ]; then - # Set checkpoint height to previous block number - prev_proof_num=$((i-1)) - PREV_PROOF_EXTRA_ARG="--checkpoint-block-number ${prev_proof_num}" + +# If we set test_only flag, we'll generate a dummy +# proof. This is useful for quickly testing decoding and all of the +# other non-proving code. +if [[ $5 == "test_only" ]]; then + # test only run + echo "Proving blocks ${BLOCK_INTERVAL} in a test_only mode now... (Total: ${TOT_BLOCKS})" + cargo r --release --features test_only --bin leader -- --runtime in-memory --load-strategy on-demand jerigon --rpc-url "$NODE_RPC_URL" --block-interval $BLOCK_INTERVAL --proof-output-dir $PROOF_OUTPUT_DIR $PREV_PROOF_EXTRA_ARG | tee test-jerigon.out + if grep 'Successfully generated witness for block' test-jerigon.out; then + echo "Success - Note this was just a test, not a proof" + exit else - if [ $i -gt 1 ]; then - prev_proof_num=$((i-1)) - PREV_PROOF_EXTRA_ARG="-f ${PROOF_OUTPUT_DIR}/b${prev_proof_num}.zkproof" - fi + echo "Failed to create a witness" + exit 1 fi +else + # normal run + echo "Proving blocks ${BLOCK_INTERVAL} now... (Total: ${TOT_BLOCKS})" + cargo r --release --bin leader -- --runtime in-memory --load-strategy on-demand jerigon --rpc-url "$3" --block-interval $BLOCK_INTERVAL --proof-output-dir $PROOF_OUTPUT_DIR $PREV_PROOF_EXTRA_ARG > $OUT_LOG_PATH 2>&1 - cargo r --release --bin leader -- --runtime in-memory jerigon --rpc-url "$3" --block-number $i --proof-output-path $OUT_PROOF_PATH $PREV_PROOF_EXTRA_ARG > $OUT_LOG_PATH 2>&1 - retVal=$? if [ $retVal -ne 0 ]; then # Some error occured. @@ -58,6 +104,10 @@ do rm $OUT_LOG_PATH fi fi -done -echo "Successfully generated ${TOT_BLOCKS} proofs!" \ No newline at end of file + echo "Successfully generated ${TOT_BLOCKS} proofs!" +fi + + + +