diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 672d6183e..707853e98 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -1,4 +1,4 @@ ---- # Workflow with multiple jobs to test different parts of the project +--- # Workflow with multiple jobs to test different parts of the project name: Continuous Integration @@ -171,7 +171,7 @@ jobs: uses: actions/checkout@v4 - name: Run the script - run: ./scripts/prove_stdio.sh artifacts/witness_b19807080.json + run: cargo xtask prove-stdio verify artifacts/witness_b19807080.json simple_proof_witness_only: name: Execute bash script to generate the proof witness for a small block. @@ -182,7 +182,7 @@ jobs: uses: actions/checkout@v4 - name: Run the script - run: ./scripts/prove_stdio.sh artifacts/witness_b19807080.json test_only + run: cargo xtask prove-stdio test artifacts/witness_b19807080.json multi_blocks_proof_regular: name: Execute bash script to generate and verify a proof for multiple blocks using parallel proving. @@ -193,4 +193,4 @@ jobs: uses: actions/checkout@v4 - name: Run the script - run: ./scripts/prove_stdio.sh artifacts/witness_b3_b6.json "" use_test_config + run: cargo xtask prove-stdio verify artifacts/witness_b3_b6.json --use-test-config diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 36ffd9682..f8d5ba41e 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -1,4 +1,4 @@ ---- # Rust lint related checks +--- # Rust lint related checks name: lint diff --git a/Cargo.lock b/Cargo.lock index 970ea3082..df690b963 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2890,9 +2890,9 @@ checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" [[package]] name = "libc" -version = "0.2.158" +version = "0.2.166" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d8adc4bb1803a324070e64a98ae98f38934d91957a99cfb3a43dcbc01bc56439" +checksum = "c2ccc108bbc0b1331bd061864e7cd823c0cab660bbe6970e66e2c0614decde36" [[package]] name = "libgit2-sys" @@ -4856,9 +4856,9 @@ dependencies = [ [[package]] name = "sysinfo" -version = "0.32.0" +version = "0.32.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e3b5ae3f4f7d64646c46c4cae4e3f01d1c5d255c7406fdd7c7f999a94e488791" +checksum = "4c33cd241af0f2e9e3b5c32163b873b29956890b5342e6745b917ce9d490f4af" dependencies = [ "core-foundation-sys", "libc", @@ -5954,6 +5954,7 @@ dependencies = [ "alloy", "anyhow", "clap", + "num_cpus", "serde", "serde_json", "sysinfo", diff --git a/scripts/Cargo.toml b/scripts/Cargo.toml index 75727776a..bfc481b56 100644 --- a/scripts/Cargo.toml +++ b/scripts/Cargo.toml @@ -13,6 +13,7 @@ publish = false alloy.workspace = true anyhow.workspace = true clap = { workspace = true, features = ["derive"] } +num_cpus = "1.0" serde = { workspace = true, features = ["derive"] } serde_json.workspace = true sysinfo = "0.32.0" diff --git a/scripts/prove_stdio.rs b/scripts/prove_stdio.rs new file mode 100644 index 000000000..7cc0b2f43 --- /dev/null +++ b/scripts/prove_stdio.rs @@ -0,0 +1,165 @@ +use std::{fs::File, path::PathBuf, process::Command}; + +use anyhow::ensure; +use clap::{arg, Args, ValueEnum, ValueHint}; + +#[derive(ValueEnum, Copy, Clone)] +enum RunMode { + /// Dummy proof is generated. Useful for quickly testing decoding and + /// all other non-proving logic. + Test, + /// The proof is generated and verified. + Verify, +} + +#[derive(Args)] +pub struct ProveStdioArgs { + /// Whether to generate a proof and verify it or not. + mode: RunMode, + /// JSON file containing the witness data. + #[arg(value_hint = ValueHint::DirPath)] + input_witness_file: PathBuf, + /// The end of the block range to prove. If None, start_block-1 is used. + #[arg(long, default_value_t = false)] + use_test_config: bool, + /// The batch size for block fetching. + #[arg(long, default_value_t = 8)] + block_batch_size: u32, + /// The directory to output the proof files. If it does not exist, it will + /// recursively be created. + #[arg(short = 'o', long, value_hint = ValueHint::DirPath, default_value = "./proofs")] + output_dir: PathBuf, +} + +/// Runs prover against the provided witness file and outputs the proof to the +/// specified directory. +pub fn prove_via_stdio(args: ProveStdioArgs) -> anyhow::Result<()> { + // Get number of cores of the system. + let num_cpus = num_cpus::get().to_string(); + let mut envs = vec![ + ("RUST_MIN_STACK", "33554432"), + ("RUST_BACKTRACE", "full"), + ("RUST_LOG", "info"), + ("RUSTFLAGS", "-C target-cpu=native -Zlinker-features=-lld"), + ("RAYON_NUM_THREADS", num_cpus.as_str()), + ("TOKIO_WORKER_THREADS", num_cpus.as_str()), + ]; + + match args.mode { + RunMode::Test => { + let mut cmd = prove_command(args, envs)?; + let status = cmd.spawn()?.wait()?; + ensure!(status.success(), "command failed with {}", status); + Ok(()) + } + RunMode::Verify => { + // Build the targets before timing. + let status = Command::new("cargo") + .envs(envs.clone()) + .args(["build", "--release", "--jobs", num_cpus.as_str()]) + .spawn()? + .wait()?; + ensure!(status.success(), "command failed with {}", status); + + // Construct the command to run. + add_verify_envs(&args, &mut envs)?; + let mut cmd = prove_command(args, envs)?; + + // Time the proving. + let start = std::time::Instant::now(); + let status = cmd.spawn()?.wait()?; + ensure!(status.success(), "command failed with {}", status); + let elapsed = start.elapsed(); + println!("Proving duration: {elapsed:?}"); + Ok(()) + } + } +} + +/// Adds environment variables to the command for verifying the proof based on +/// the input file. +fn add_verify_envs(args: &ProveStdioArgs, envs: &mut Vec<(&str, &str)>) -> anyhow::Result<()> { + // Get the witness filename. + let witness_filename = args + .input_witness_file + .to_str() + .ok_or(anyhow::anyhow!("Invalid witness file path"))?; + + // Set envs based on filename. + if witness_filename.contains("witness_b19807080") { + envs.extend([ + ("ARITHMETIC_CIRCUIT_SIZE", "16..18"), + ("BYTE_PACKING_CIRCUIT_SIZE", "8..15"), + ("CPU_CIRCUIT_SIZE", "9..20"), + ("KECCAK_CIRCUIT_SIZE", "7..18"), + ("KECCAK_SPONGE_CIRCUIT_SIZE", "8..14"), + ("LOGIC_CIRCUIT_SIZE", "5..17"), + ("MEMORY_CIRCUIT_SIZE", "17..22"), + ("MEMORY_BEFORE_CIRCUIT_SIZE", "16..20"), + ("MEMORY_AFTER_CIRCUIT_SIZE", "7..20"), + ("POSEIDON_CIRCUIT_SIZE", "4..8"), + ]); + } else if witness_filename.contains("witness_b3_b6") { + envs.extend([ + ("ARITHMETIC_CIRCUIT_SIZE", "16..18"), + ("BYTE_PACKING_CIRCUIT_SIZE", "8..15"), + ("CPU_CIRCUIT_SIZE", "10..20"), + ("KECCAK_CIRCUIT_SIZE", "4..13"), + ("KECCAK_SPONGE_CIRCUIT_SIZE", "8..9"), + ("LOGIC_CIRCUIT_SIZE", "4..14"), + ("MEMORY_CIRCUIT_SIZE", "17..22"), + ("MEMORY_BEFORE_CIRCUIT_SIZE", "16..18"), + ("MEMORY_AFTER_CIRCUIT_SIZE", "7..8"), + ("POSEIDON_CIRCUIT_SIZE", "4..8"), + ]); + } else { + envs.extend([ + ("ARITHMETIC_CIRCUIT_SIZE", "16..18"), + ("BYTE_PACKING_CIRCUIT_SIZE", "8..15"), + ("CPU_CIRCUIT_SIZE", "9..20"), + ("KECCAK_CIRCUIT_SIZE", "7..18"), + ("KECCAK_SPONGE_CIRCUIT_SIZE", "8..14"), + ("LOGIC_CIRCUIT_SIZE", "5..17"), + ("MEMORY_CIRCUIT_SIZE", "17..22"), + ("MEMORY_BEFORE_CIRCUIT_SIZE", "16..20"), + ("MEMORY_AFTER_CIRCUIT_SIZE", "7..20"), + // TODO(Robin): update Poseidon ranges here and below once Kernel ASM supports + ("POSEIDON_CIRCUIT_SIZE", "4..8"), + ]); + } + Ok(()) +} + +/// Constructs the command to run for proving via stdin. +fn prove_command(args: ProveStdioArgs, envs: Vec<(&str, &str)>) -> anyhow::Result { + let witness_file = File::open(&args.input_witness_file)?; + let mut cmd = Command::new("cargo"); + cmd.envs(envs).stdin(witness_file); + cmd.args([ + "run", + "--release", + "--package", + "zero", + "--bin", + "leader", + "--", + "--runtime", + "in-memory", + "--load-strategy", + "on-demand", + "--block-batch-size", + args.block_batch_size.to_string().as_str(), + "--proof-output-dir", + args.output_dir + .to_str() + .ok_or(anyhow::anyhow!("Invalid output dir path"))?, + ]); + if let RunMode::Test = args.mode { + cmd.arg("--test-only"); + } + if args.use_test_config { + cmd.arg("--use-test-config"); + } + cmd.arg("stdio"); + Ok(cmd) +} diff --git a/scripts/prove_stdio.sh b/scripts/prove_stdio.sh deleted file mode 100755 index 04bbbbe6b..000000000 --- a/scripts/prove_stdio.sh +++ /dev/null @@ -1,167 +0,0 @@ -#!/bin/bash -# ------------------------------------------------------------------------------ -set -exo pipefail - -# Run prover with the parsed input from the standard terminal. -# To generate the json input file, use the `rpc` tool, for example: -# `cargo run --bin rpc -- fetch --rpc-url http://127.0.0.1:8546 --start-block 2 --end-block 5 > witness.json` - -# Args: -# 1 --> Input witness json file -# 2 --> Test run only flag `test_only` (optional) - -# We're going to set the parallelism in line with the total cpu count -if [[ "$OSTYPE" == "darwin"* ]]; then - num_procs=$(sysctl -n hw.physicalcpu) -else - num_procs=$(nproc) -fi - -# Force the working directory to always be the `tools/` directory. -REPO_ROOT=$(git rev-parse --show-toplevel) -PROOF_OUTPUT_DIR="${REPO_ROOT}/proofs" - -BLOCK_BATCH_SIZE="${BLOCK_BATCH_SIZE:-8}" -echo "Block batch size: $BLOCK_BATCH_SIZE" - -OUTPUT_LOG="${REPO_ROOT}/output.log" -PROOFS_FILE_LIST="${PROOF_OUTPUT_DIR}/proof_files.json" -TEST_OUT_PATH="${REPO_ROOT}/test.out" - -# Configured Rayon and Tokio with rough defaults -export RAYON_NUM_THREADS=$num_procs -export TOKIO_WORKER_THREADS=$num_procs - -export RUST_MIN_STACK=33554432 -export RUST_BACKTRACE=full -export RUST_LOG=info -# Script users are running locally, and might benefit from extra perf. -# See also .cargo/config.toml. -export RUSTFLAGS='-C target-cpu=native -Zlinker-features=-lld' - -INPUT_FILE=$1 -TEST_ONLY=$2 -USE_TEST_CONFIG=$3 - -if [[ $INPUT_FILE == "" ]]; then - echo "Please provide witness json input file, e.g. artifacts/witness_b19240705.json" - exit 1 -fi - -# 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" - export ARITHMETIC_CIRCUIT_SIZE="16..18" - export BYTE_PACKING_CIRCUIT_SIZE="8..15" - export CPU_CIRCUIT_SIZE="9..20" - export KECCAK_CIRCUIT_SIZE="7..18" - export KECCAK_SPONGE_CIRCUIT_SIZE="8..14" - export LOGIC_CIRCUIT_SIZE="5..17" - export MEMORY_CIRCUIT_SIZE="17..22" - export MEMORY_BEFORE_CIRCUIT_SIZE="16..20" - export MEMORY_AFTER_CIRCUIT_SIZE="7..20" - # TODO(Robin): update Poseidon ranges here and below once Kernel ASM supports Poseidon ops - export POSEIDON_CIRCUIT_SIZE="4..8" - elif [[ $INPUT_FILE == *"witness_b3_b6"* ]]; then - # These sizes are configured specifically for custom blocks 3 to 6. Don't use this in other scenarios - echo "Using specific circuit sizes for witness_b3_b6.json" - export ARITHMETIC_CIRCUIT_SIZE="16..18" - export BYTE_PACKING_CIRCUIT_SIZE="8..15" - export CPU_CIRCUIT_SIZE="10..20" - export KECCAK_CIRCUIT_SIZE="4..13" - export KECCAK_SPONGE_CIRCUIT_SIZE="8..9" - export LOGIC_CIRCUIT_SIZE="4..14" - export MEMORY_CIRCUIT_SIZE="17..22" - export MEMORY_BEFORE_CIRCUIT_SIZE="16..18" - export MEMORY_AFTER_CIRCUIT_SIZE="7..8" - export POSEIDON_CIRCUIT_SIZE="4..8" - else - export ARITHMETIC_CIRCUIT_SIZE="16..21" - export BYTE_PACKING_CIRCUIT_SIZE="8..21" - export CPU_CIRCUIT_SIZE="8..21" - export KECCAK_CIRCUIT_SIZE="4..20" - export KECCAK_SPONGE_CIRCUIT_SIZE="8..17" - export LOGIC_CIRCUIT_SIZE="4..21" - export MEMORY_CIRCUIT_SIZE="17..24" - export MEMORY_BEFORE_CIRCUIT_SIZE="16..23" - export MEMORY_AFTER_CIRCUIT_SIZE="7..23" - export POSEIDON_CIRCUIT_SIZE="4..8" - fi -fi - - -# If we run ./prove_stdio.sh test_only, we'll generate a dummy -# proof. This is useful for quickly testing decoding and all of the -# other non-proving code. -if [[ $TEST_ONLY == "test_only" ]]; then - cargo run --quiet --release --package zero --bin leader -- \ - --test-only \ - --runtime in-memory \ - --load-strategy on-demand \ - --block-batch-size "$BLOCK_BATCH_SIZE" \ - --proof-output-dir "$PROOF_OUTPUT_DIR" \ - 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" - exit - else - # Some error occurred, display the logs and exit. - cat "$TEST_OUT_PATH" - echo "Failed to create proof witnesses. See $TEST_OUT_PATH for more details." - exit 1 - fi -fi - -cargo build --release --jobs "$num_procs" - -start_time=$(date +%s%N) - -cmd=("${REPO_ROOT}/target/release/leader" --runtime in-memory \ - --load-strategy on-demand \ - --block-batch-size "$BLOCK_BATCH_SIZE") - -if [[ "$USE_TEST_CONFIG" == "use_test_config" ]]; then - cmd+=("--use-test-config") -fi - -"${cmd[@]}" --proof-output-dir "$PROOF_OUTPUT_DIR" stdio < "$INPUT_FILE" &> "$OUTPUT_LOG" -end_time=$(date +%s%N) - -grep "Successfully wrote to disk proof file " "$OUTPUT_LOG" | awk '{print $NF}' | tee "$PROOFS_FILE_LIST" -if [ ! -s "$PROOFS_FILE_LIST" ]; then - # Some error occurred, display the logs and exit. - cat "$OUTPUT_LOG" - echo "Proof list not generated, some error happened. For more details check the log file $OUTPUT_LOG" - exit 1 -fi - -while read -r proof_file; -do - echo "Verifying proof file $proof_file" - verify_file=$PROOF_OUTPUT_DIR/verify_$(basename "$proof_file").out - "${REPO_ROOT}/target/release/verifier" -f "$proof_file" | tee "$verify_file" - if grep -q 'All proofs verified successfully!' "$verify_file"; then - echo "Proof verification for file $proof_file successful"; - rm "$verify_file" # we keep the generated proof for potential reuse - else - # Some error occurred with verification, display the logs and exit. - cat "$verify_file" - echo "There was an issue with proof verification. See $verify_file for more details."; - exit 1 - fi -done < "$PROOFS_FILE_LIST" - -duration_ns=$((end_time - start_time)) -duration_sec=$(echo "$duration_ns / 1000000000" | bc -l) - -echo "Success!" -echo "Proving duration: $duration_sec seconds" -echo "Note, this duration is inclusive of circuit handling and overall process initialization"; - -# Clean up in case of success -rm "$OUTPUT_LOG" - diff --git a/scripts/xtask.rs b/scripts/xtask.rs index 0a28d5ef0..616343183 100644 --- a/scripts/xtask.rs +++ b/scripts/xtask.rs @@ -2,11 +2,13 @@ mod outdated; mod prove_rpc; +mod prove_stdio; use anyhow::Result; use clap::Parser; use outdated::list_outdated_deps; use prove_rpc::{prove_via_rpc, ProveRpcArgs}; +use prove_stdio::{prove_via_stdio, ProveStdioArgs}; #[derive(Parser)] enum Args { @@ -20,11 +22,14 @@ enum Args { Outdated, /// Execute proving via RPC endpoint. ProveRpc(Box), + /// Execute proving via stdin. + ProveStdio(ProveStdioArgs), } fn main() -> Result<()> { match Args::parse() { Args::Outdated => list_outdated_deps(), Args::ProveRpc(args) => prove_via_rpc(*args), + Args::ProveStdio(args) => prove_via_stdio(args), } }