feat: add prover config argumetns to scripts

0xPolygonZero · Jul 24, 2024 · 542d5f3 · 542d5f3
1 parent 134cbe9
commit 542d5f3
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 2 deletions.
diff --git a/zero_bin/tools/prove_rpc.sh b/zero_bin/tools/prove_rpc.sh
@@ -64,6 +64,16 @@ OUTPUT_TO_TERMINAL="${OUTPUT_TO_TERMINAL:-false}"
 # Only generate proof by default
 RUN_VERIFICATION="${RUN_VERIFICATION:-false}"
 
+# 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}"
+if [[ $PROVER_SAVE_INPUTS_ON_ERROR=="true" ]]; then
+    PROVER_SAVE_INPUTS_ON_ERROR="--save-inputs-on-error"
+else
+    PROVER_SAVE_INPUTS_ON_ERROR=""
+fi
+
 mkdir -p $PROOF_OUTPUT_DIR
 
 if [ $IGNORE_PREVIOUS_PROOFS ]; then
@@ -122,7 +132,9 @@ 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 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 $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
@@ -91,6 +91,18 @@ else
 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}"
+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 we run ./prove_stdio.sh <witness file name> test_only, we'll generate a dummy
 # proof. This is useful for quickly testing decoding and all of the
 # other non-proving code.
@@ -108,7 +120,9 @@ fi
 cargo build --release --jobs "$num_procs"
 
 start_time=$(date +%s%N)
-"${TOOLS_DIR}/../../target/release/leader" --runtime in-memory --load-strategy on-demand stdio < $INPUT_FILE | tee $LEADER_OUT_PATH
+"${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
 end_time=$(date +%s%N)
 
 tail -n 1 $LEADER_OUT_PATH > $PROOFS_JSON_PATH