From 63c92b6205fb156f4b50dee581674b814f44f874 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mat=C3=ADas=20Ignacio=20Gonz=C3=A1lez?= Date: Mon, 29 Jul 2024 18:44:25 +0200 Subject: [PATCH] feat(zk_toolbox): Update prover support (#2533) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## What ❔ Update prover support: * Witness vector generator thread selection * Prover CloudType selection --- etc/env/file_based/general.yaml | 1 + zk_toolbox/README.md | 2 + .../src/commands/prover/args/init.rs | 38 ++++++++++++++++++- .../src/commands/prover/args/run.rs | 33 +++++++++++++++- .../zk_inception/src/commands/prover/init.rs | 1 + .../src/commands/prover/init_bellman_cuda.rs | 5 +++ .../zk_inception/src/commands/prover/run.rs | 16 ++++++-- .../crates/zk_inception/src/messages.rs | 2 + 8 files changed, 90 insertions(+), 8 deletions(-) diff --git a/etc/env/file_based/general.yaml b/etc/env/file_based/general.yaml index 059b993d326b..a36e3db2b6cd 100644 --- a/etc/env/file_based/general.yaml +++ b/etc/env/file_based/general.yaml @@ -177,6 +177,7 @@ prover: availability_check_interval_in_secs: 10000 zone_read_url: http://metadata.google.internal/computeMetadata/v1/instance/zone shall_save_to_public_bucket: true + cloud_type: LOCAL witness_generator: generation_timeout_in_secs: 900 max_attempts: 10 diff --git a/zk_toolbox/README.md b/zk_toolbox/README.md index d97d05f1459a..71eb9af00ec5 100644 --- a/zk_toolbox/README.md +++ b/zk_toolbox/README.md @@ -148,6 +148,8 @@ You can specify the prover component to run by providing `--component ` argument. Possible rounds are: `all-rounds, basic-circuits, leaf-aggregation, node-aggregation, recursion-tip, scheduler` +For `witness-vector-generator`, specify the number of WVG jobs with `--threads `. + ### Contract verifier Running the contract verifier: diff --git a/zk_toolbox/crates/zk_inception/src/commands/prover/args/init.rs b/zk_toolbox/crates/zk_inception/src/commands/prover/args/init.rs index 4943c596a1d6..3eca80f56220 100644 --- a/zk_toolbox/crates/zk_inception/src/commands/prover/args/init.rs +++ b/zk_toolbox/crates/zk_inception/src/commands/prover/args/init.rs @@ -3,14 +3,15 @@ use common::{logger, Prompt, PromptConfirm, PromptSelect}; use serde::{Deserialize, Serialize}; use strum::{EnumIter, IntoEnumIterator}; use xshell::Shell; +use zksync_config::configs::fri_prover::CloudType; use super::init_bellman_cuda::InitBellmanCudaArgs; use crate::{ commands::prover::gcs::get_project_ids, consts::{DEFAULT_CREDENTIALS_FILE, DEFAULT_PROOF_STORE_DIR}, messages::{ - MSG_CREATE_GCS_BUCKET_LOCATION_PROMPT, MSG_CREATE_GCS_BUCKET_NAME_PROMTP, - MSG_CREATE_GCS_BUCKET_PROJECT_ID_NO_PROJECTS_PROMPT, + MSG_CLOUD_TYPE_PROMPT, MSG_CREATE_GCS_BUCKET_LOCATION_PROMPT, + MSG_CREATE_GCS_BUCKET_NAME_PROMTP, MSG_CREATE_GCS_BUCKET_PROJECT_ID_NO_PROJECTS_PROMPT, MSG_CREATE_GCS_BUCKET_PROJECT_ID_PROMPT, MSG_CREATE_GCS_BUCKET_PROMPT, MSG_DOWNLOAD_SETUP_KEY_PROMPT, MSG_GETTING_PROOF_STORE_CONFIG, MSG_GETTING_PUBLIC_STORE_CONFIG, MSG_PROOF_STORE_CONFIG_PROMPT, MSG_PROOF_STORE_DIR_PROMPT, @@ -52,6 +53,9 @@ pub struct ProverInitArgs { #[clap(flatten)] #[serde(flatten)] pub setup_key_config: SetupKeyConfigTmp, + + #[clap(long)] + cloud_type: Option, } #[derive(Debug, Clone, ValueEnum, EnumIter, strum::Display, PartialEq, Eq)] @@ -61,6 +65,24 @@ enum ProofStoreConfig { GCS, } +#[derive( + Debug, Clone, ValueEnum, EnumIter, strum::Display, PartialEq, Eq, Deserialize, Serialize, +)] +#[allow(clippy::upper_case_acronyms)] +enum InternalCloudType { + GCP, + Local, +} + +impl From for CloudType { + fn from(cloud_type: InternalCloudType) -> Self { + match cloud_type { + InternalCloudType::GCP => CloudType::GCP, + InternalCloudType::Local => CloudType::Local, + } + } +} + #[derive(Clone, Debug, Serialize, Deserialize, Parser, Default)] pub struct ProofStorageGCSTmp { #[clap(long)] @@ -144,6 +166,7 @@ pub struct ProverInitArgsFinal { pub public_store: Option, pub setup_key_config: SetupKeyConfig, pub bellman_cuda_config: InitBellmanCudaArgs, + pub cloud_type: CloudType, } impl ProverInitArgs { @@ -156,11 +179,14 @@ impl ProverInitArgs { let public_store = self.fill_public_storage_values_with_prompt(shell)?; let setup_key_config = self.fill_setup_key_values_with_prompt(setup_key_path); let bellman_cuda_config = self.fill_bellman_cuda_values_with_prompt()?; + let cloud_type = self.get_cloud_type_with_prompt(); + Ok(ProverInitArgsFinal { proof_store, public_store, setup_key_config, bellman_cuda_config, + cloud_type, }) } @@ -406,4 +432,12 @@ impl ProverInitArgs { fn fill_bellman_cuda_values_with_prompt(&self) -> anyhow::Result { self.bellman_cuda_config.clone().fill_values_with_prompt() } + + fn get_cloud_type_with_prompt(&self) -> CloudType { + let cloud_type = self.cloud_type.clone().unwrap_or_else(|| { + PromptSelect::new(MSG_CLOUD_TYPE_PROMPT, InternalCloudType::iter()).ask() + }); + + cloud_type.into() + } } diff --git a/zk_toolbox/crates/zk_inception/src/commands/prover/args/run.rs b/zk_toolbox/crates/zk_inception/src/commands/prover/args/run.rs index 4b485099cc80..c2d5cef26ad4 100644 --- a/zk_toolbox/crates/zk_inception/src/commands/prover/args/run.rs +++ b/zk_toolbox/crates/zk_inception/src/commands/prover/args/run.rs @@ -1,8 +1,8 @@ use clap::{Parser, ValueEnum}; -use common::PromptSelect; +use common::{Prompt, PromptSelect}; use strum::{EnumIter, IntoEnumIterator}; -use crate::messages::{MSG_ROUND_SELECT_PROMPT, MSG_RUN_COMPONENT_PROMPT}; +use crate::messages::{MSG_ROUND_SELECT_PROMPT, MSG_RUN_COMPONENT_PROMPT, MSG_THREADS_PROMPT}; #[derive(Debug, Clone, Parser, Default)] pub struct ProverRunArgs { @@ -10,6 +10,8 @@ pub struct ProverRunArgs { pub component: Option, #[clap(flatten)] pub witness_generator_args: WitnessGeneratorArgs, + #[clap(flatten)] + pub witness_vector_generator_args: WitnessVectorGeneratorArgs, } #[derive( @@ -50,6 +52,28 @@ pub enum WitnessGeneratorRound { Scheduler, } +#[derive(Debug, Clone, Parser, Default)] +pub struct WitnessVectorGeneratorArgs { + #[clap(long)] + pub threads: Option, +} + +impl WitnessVectorGeneratorArgs { + fn fill_values_with_prompt(&self, component: ProverComponent) -> anyhow::Result { + if component != ProverComponent::WitnessVectorGenerator { + return Ok(Self::default()); + } + + let threads = self + .threads + .unwrap_or_else(|| Prompt::new(MSG_THREADS_PROMPT).default("1").ask()); + + Ok(Self { + threads: Some(threads), + }) + } +} + impl ProverRunArgs { pub fn fill_values_with_prompt(&self) -> anyhow::Result { let component = self.component.unwrap_or_else(|| { @@ -60,9 +84,14 @@ impl ProverRunArgs { .witness_generator_args .fill_values_with_prompt(component)?; + let witness_vector_generator_args = self + .witness_vector_generator_args + .fill_values_with_prompt(component)?; + Ok(ProverRunArgs { component: Some(component), witness_generator_args, + witness_vector_generator_args, }) } } diff --git a/zk_toolbox/crates/zk_inception/src/commands/prover/init.rs b/zk_toolbox/crates/zk_inception/src/commands/prover/init.rs index 54b53b8576db..a27e5f1b0bec 100644 --- a/zk_toolbox/crates/zk_inception/src/commands/prover/init.rs +++ b/zk_toolbox/crates/zk_inception/src/commands/prover/init.rs @@ -59,6 +59,7 @@ pub(crate) async fn run(args: ProverInitArgs, shell: &Shell) -> anyhow::Result<( } else { prover_config.shall_save_to_public_bucket = false; } + prover_config.cloud_type = args.cloud_type; general_config.prover_config = Some(prover_config); let mut proof_compressor_config = general_config diff --git a/zk_toolbox/crates/zk_inception/src/commands/prover/init_bellman_cuda.rs b/zk_toolbox/crates/zk_inception/src/commands/prover/init_bellman_cuda.rs index c6c5d3ef23d9..75535587c42c 100644 --- a/zk_toolbox/crates/zk_inception/src/commands/prover/init_bellman_cuda.rs +++ b/zk_toolbox/crates/zk_inception/src/commands/prover/init_bellman_cuda.rs @@ -37,6 +37,11 @@ pub(crate) async fn run(shell: &Shell, args: InitBellmanCudaArgs) -> anyhow::Res } fn clone_bellman_cuda(shell: &Shell) -> anyhow::Result { + let path = shell.current_dir().join(BELLMAN_CUDA_DIR); + if shell.path_exists(path.clone()) { + return Ok(path.to_str().context(MSG_BELLMAN_CUDA_DIR_ERR)?.to_string()); + } + let spinner = Spinner::new(MSG_CLONING_BELLMAN_CUDA_SPINNER); let path = git::clone( shell, diff --git a/zk_toolbox/crates/zk_inception/src/commands/prover/run.rs b/zk_toolbox/crates/zk_inception/src/commands/prover/run.rs index 898cf0e45d66..5497db8a21e0 100644 --- a/zk_toolbox/crates/zk_inception/src/commands/prover/run.rs +++ b/zk_toolbox/crates/zk_inception/src/commands/prover/run.rs @@ -4,7 +4,10 @@ use config::{ChainConfig, EcosystemConfig}; use xshell::{cmd, Shell}; use super::{ - args::run::{ProverComponent, ProverRunArgs, WitnessGeneratorArgs, WitnessGeneratorRound}, + args::run::{ + ProverComponent, ProverRunArgs, WitnessGeneratorArgs, WitnessGeneratorRound, + WitnessVectorGeneratorArgs, + }, utils::get_link_to_prover, }; use crate::messages::{ @@ -32,7 +35,7 @@ pub(crate) async fn run(args: ProverRunArgs, shell: &Shell) -> anyhow::Result<() run_witness_generator(shell, &chain, args.witness_generator_args)? } Some(ProverComponent::WitnessVectorGenerator) => { - run_witness_vector_generator(shell, &chain)? + run_witness_vector_generator(shell, &chain, args.witness_vector_generator_args)? } Some(ProverComponent::Prover) => run_prover(shell, &chain)?, Some(ProverComponent::Compressor) => run_compressor(shell, &chain, &ecosystem_config)?, @@ -76,12 +79,17 @@ fn run_witness_generator( cmd.run().context(MSG_RUNNING_WITNESS_GENERATOR_ERR) } -fn run_witness_vector_generator(shell: &Shell, chain: &ChainConfig) -> anyhow::Result<()> { +fn run_witness_vector_generator( + shell: &Shell, + chain: &ChainConfig, + args: WitnessVectorGeneratorArgs, +) -> anyhow::Result<()> { logger::info(MSG_RUNNING_WITNESS_VECTOR_GENERATOR); let config_path = chain.path_to_general_config(); let secrets_path = chain.path_to_secrets_config(); - let mut cmd = Cmd::new(cmd!(shell, "cargo run --release --bin zksync_witness_vector_generator -- --config-path={config_path} --secrets-path={secrets_path}")); + let threads = args.threads.unwrap_or(1).to_string(); + let mut cmd = Cmd::new(cmd!(shell, "cargo run --release --bin zksync_witness_vector_generator -- --config-path={config_path} --secrets-path={secrets_path} --threads={threads}")); cmd = cmd.with_force_run(); cmd.run().context(MSG_RUNNING_WITNESS_VECTOR_GENERATOR_ERR) } diff --git a/zk_toolbox/crates/zk_inception/src/messages.rs b/zk_toolbox/crates/zk_inception/src/messages.rs index 4e1ad9074389..710290158cc6 100644 --- a/zk_toolbox/crates/zk_inception/src/messages.rs +++ b/zk_toolbox/crates/zk_inception/src/messages.rs @@ -283,6 +283,8 @@ pub(super) const MSG_BELLMAN_CUDA_ORIGIN_SELECT: &str = "Select the origin of bellman-cuda repository"; pub(super) const MSG_BELLMAN_CUDA_SELECTION_CLONE: &str = "Clone for me (recommended)"; pub(super) const MSG_BELLMAN_CUDA_SELECTION_PATH: &str = "I have the code already"; +pub(super) const MSG_CLOUD_TYPE_PROMPT: &str = "Select the cloud type:"; +pub(super) const MSG_THREADS_PROMPT: &str = "Provide the number of threads:"; pub(super) fn msg_bucket_created(bucket_name: &str) -> String { format!("Bucket created successfully with url: gs://{bucket_name}")