feat(zk_toolbox): Update prover support (#2533)

## What ❔

Update prover support:
* Witness vector generator thread selection
* Prover CloudType selection

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

zk_toolbox/README.md

@@ -148,6 +148,8 @@ You can specify the prover component to run by providing `--component <component
 If you are running `witness-generator` you can specify the round by providing `--round <round>` 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 <threads>`.
+
 ### Contract verifier
 
 Running the contract verifier:

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<InternalCloudType>,
 }
 
 #[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<InternalCloudType> 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<ProofStorageConfig>,
     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<InitBellmanCudaArgs> {
         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()
+    }
 }

zk_toolbox/crates/zk_inception/src/commands/prover/args/run.rs

@@ -1,15 +1,17 @@
 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 {
     #[clap(long)]
     pub component: Option<ProverComponent>,
     #[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<usize>,
+}
+
+impl WitnessVectorGeneratorArgs {
+    fn fill_values_with_prompt(&self, component: ProverComponent) -> anyhow::Result<Self> {
+        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<ProverRunArgs> {
         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,
         })
     }
 }

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

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<String> {
+    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,

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)
 }

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}")