feat(zk_toolbox): Add prover run

ZkStack.yaml

@@ -1,6 +1,7 @@
 name: zk
 l1_network: Localhost
 link_to_code: .
+bellman_cuda_dir: .
 chains: ./chains
 config: ./configs/
 default_chain: era

zk_toolbox/crates/config/src/ecosystem.rs

@@ -25,6 +25,7 @@ struct EcosystemConfigInternal {
     pub name: String,
     pub l1_network: L1Network,
     pub link_to_code: PathBuf,
+    pub bellman_cuda_dir: PathBuf,
     pub chains: PathBuf,
     pub config: PathBuf,
     pub default_chain: String,
@@ -40,6 +41,7 @@ pub struct EcosystemConfig {
     pub name: String,
     pub l1_network: L1Network,
     pub link_to_code: PathBuf,
+    pub bellman_cuda_dir: PathBuf,
     pub chains: PathBuf,
     pub config: PathBuf,
     pub default_chain: String,
@@ -72,6 +74,11 @@ impl<'de> Deserialize<'de> for EcosystemConfig {
                 .absolutize()
                 .expect("Failed to parse zksync-era path")
                 .to_path_buf(),
+            bellman_cuda_dir: config
+                .bellman_cuda_dir
+                .absolutize()
+                .expect("Failed to parse bellman-cuda path")
+                .to_path_buf(),
             chains: config.chains.clone(),
             config: config.config.clone(),
             default_chain: config.default_chain.clone(),
@@ -202,6 +209,11 @@ impl EcosystemConfig {
                 .absolutize()
                 .expect("Failed to parse zksync-era path")
                 .into(),
+            bellman_cuda_dir: self
+                .bellman_cuda_dir
+                .absolutize()
+                .expect("Failed to parse bellman-cuda path")
+                .into(),
             chains: self.chains.clone(),
             config: self.config.clone(),
             default_chain: self.default_chain.clone(),

zk_toolbox/crates/zk_inception/src/commands/ecosystem/create.rs

@@ -21,6 +21,7 @@ use crate::{
             create_configs::{create_erc20_deployment_config, create_initial_deployments_config},
         },
     },
+    consts::BELLMAN_CUDA_DIR,
     messages::{
         MSG_CLONING_ERA_REPO_SPINNER, MSG_CREATED_ECOSYSTEM, MSG_CREATING_DEFAULT_CHAIN_SPINNER,
         MSG_CREATING_ECOSYSTEM, MSG_CREATING_INITIAL_CONFIGURATIONS_SPINNER,
@@ -76,6 +77,7 @@ fn create(args: EcosystemCreateArgs, shell: &Shell) -> anyhow::Result<()> {
         name: ecosystem_name.clone(),
         l1_network: args.l1_network,
         link_to_code: link_to_code.clone(),
+        bellman_cuda_dir: shell.current_dir().join(BELLMAN_CUDA_DIR),
         chains: chains_path.clone(),
         config: configs_path,
         era_chain_id: get_default_era_chain_id(),

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

@@ -0,0 +1,29 @@
+use std::path::PathBuf;
+
+use clap::Parser;
+use common::Prompt;
+
+use crate::messages::MSG_BELLMAN_CUDA_DIR_PROMPT;
+
+#[derive(Debug, Clone, Parser, Default)]
+pub struct InitBellmanCudaArgs {
+    #[clap(long)]
+    pub bellman_cuda_dir: Option<String>,
+}
+
+impl InitBellmanCudaArgs {
+    pub fn fill_values_with_prompt(
+        self,
+        default_bellman_cuda_dir: &PathBuf,
+    ) -> anyhow::Result<InitBellmanCudaArgs> {
+        let bellman_cuda_dir = self.bellman_cuda_dir.unwrap_or_else(|| {
+            Prompt::new(MSG_BELLMAN_CUDA_DIR_PROMPT)
+                .default(default_bellman_cuda_dir.to_str().unwrap_or_default())
+                .ask()
+        });
+
+        Ok(InitBellmanCudaArgs {
+            bellman_cuda_dir: Some(bellman_cuda_dir),
+        })
+    }
+}

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

@@ -1,2 +1,3 @@
 pub mod init;
+pub mod init_bellman_cuda;
 pub mod run;

zk_toolbox/crates/zk_inception/src/commands/prover/generate_sk.rs

@@ -1,12 +1,14 @@
 use anyhow::Ok;
-use common::{cmd::Cmd, logger, spinner::Spinner};
+use common::{check_prover_prequisites, cmd::Cmd, logger, spinner::Spinner};
 use config::EcosystemConfig;
 use xshell::{cmd, Shell};
 
 use super::utils::get_link_to_prover;
 use crate::messages::{MSG_GENERATING_SK_SPINNER, MSG_SK_GENERATED};
 
 pub(crate) async fn run(shell: &Shell) -> anyhow::Result<()> {
+    check_prover_prequisites(shell);
+
     let ecosystem_config = EcosystemConfig::from_file(shell)?;
     let link_to_prover = get_link_to_prover(&ecosystem_config);
     shell.change_dir(&link_to_prover);

zk_toolbox/crates/zk_inception/src/commands/prover/init.rs

@@ -12,10 +12,9 @@ use super::{
     utils::get_link_to_prover,
 };
 use crate::{
-    consts::{BELLMAN_CUDA_DIR, PROVER_STORE_MAX_RETRIES},
+    consts::PROVER_STORE_MAX_RETRIES,
     messages::{
-        MSG_BUILDING_BELLMAN_CUDA_SPINNER, MSG_CHAIN_NOT_FOUND_ERR,
-        MSG_CLONING_BELLMAN_CUDA_SPINNER, MSG_DOWNLOADING_SETUP_KEY_SPINNER,
+        MSG_CHAIN_NOT_FOUND_ERR, MSG_DOWNLOADING_SETUP_KEY_SPINNER,
         MSG_GENERAL_CONFIG_NOT_FOUND_ERR, MSG_PROOF_COMPRESSOR_CONFIG_NOT_FOUND_ERR,
         MSG_PROVER_CONFIG_NOT_FOUND_ERR, MSG_PROVER_INITIALIZED, MSG_SETUP_KEY_PATH_ERROR,
     },
@@ -68,8 +67,6 @@ pub(crate) async fn run(args: ProverInitArgs, shell: &Shell) -> anyhow::Result<(
 
     chain_config.save_zksync_general_config(&general_config)?;
 
-    init_bellman_cuda(shell)?;
-
     logger::outro(MSG_PROVER_INITIALIZED);
     Ok(())
 }
@@ -136,27 +133,3 @@ fn get_object_store_config(
 
     Ok(object_store)
 }
-
-fn init_bellman_cuda(shell: &Shell) -> anyhow::Result<()> {
-    if shell.path_exists(BELLMAN_CUDA_DIR) {
-        return Ok(());
-    }
-
-    let spinner = Spinner::new(MSG_CLONING_BELLMAN_CUDA_SPINNER);
-    Cmd::new(cmd!(
-        shell,
-        "git clone https://github.com/matter-labs/era-bellman-cuda"
-    ))
-    .run()?;
-    spinner.finish();
-
-    let spinner = Spinner::new(MSG_BUILDING_BELLMAN_CUDA_SPINNER);
-    Cmd::new(cmd!(
-        shell,
-        "cmake -Bera-bellman-cuda/build -Sera-bellman-cuda/ -DCMAKE_BUILD_TYPE=Release"
-    ))
-    .run()?;
-    Cmd::new(cmd!(shell, "cmake --build era-bellman-cuda/build")).run()?;
-    spinner.finish();
-    Ok(())
-}

zk_toolbox/crates/zk_inception/src/commands/prover/init_bellman_cuda.rs

@@ -0,0 +1,50 @@
+use common::{check_prover_prequisites, cmd::Cmd, logger, spinner::Spinner};
+use config::{traits::SaveConfigWithBasePath, EcosystemConfig};
+use xshell::{cmd, Shell};
+
+use crate::messages::{
+    MSG_BELLMAN_CUDA_DIR_ERR, MSG_BELLMAN_CUDA_INITIALIZED_MSG, MSG_BUILDING_BELLMAN_CUDA_SPINNER,
+    MSG_CLONING_BELLMAN_CUDA_SPINNER,
+};
+
+use super::args::init_bellman_cuda::InitBellmanCudaArgs;
+
+pub(crate) async fn run(shell: &Shell, args: InitBellmanCudaArgs) -> anyhow::Result<()> {
+    check_prover_prequisites(shell);
+
+    let mut ecosystem_config = EcosystemConfig::from_file(shell)?;
+
+    let args = args.fill_values_with_prompt(&ecosystem_config.bellman_cuda_dir)?;
+
+    let bellman_cuda_dir = args.bellman_cuda_dir.expect(MSG_BELLMAN_CUDA_DIR_ERR);
+    ecosystem_config.bellman_cuda_dir = bellman_cuda_dir.clone().into();
+
+    if shell.path_exists(&bellman_cuda_dir) {
+        return save_bellman_cuda_dir(shell, &ecosystem_config);
+    }
+
+    let spinner = Spinner::new(MSG_CLONING_BELLMAN_CUDA_SPINNER);
+    Cmd::new(cmd!(
+        shell,
+        "git clone https://github.com/matter-labs/era-bellman-cuda {bellman_cuda_dir}"
+    ))
+    .run()?;
+    spinner.finish();
+
+    let spinner = Spinner::new(MSG_BUILDING_BELLMAN_CUDA_SPINNER);
+    Cmd::new(cmd!(
+        shell,
+        "cmake -B{bellman_cuda_dir}/build -S{bellman_cuda_dir}/ -DCMAKE_BUILD_TYPE=Release"
+    ))
+    .run()?;
+    Cmd::new(cmd!(shell, "cmake --build {bellman_cuda_dir}/build")).run()?;
+    spinner.finish();
+
+    save_bellman_cuda_dir(shell, &ecosystem_config)
+}
+
+fn save_bellman_cuda_dir(shell: &Shell, ecosystem: &EcosystemConfig) -> anyhow::Result<()> {
+    ecosystem.save_with_base_path(shell, ".")?;
+    logger::outro(MSG_BELLMAN_CUDA_INITIALIZED_MSG);
+    Ok(())
+}

zk_toolbox/crates/zk_inception/src/commands/prover/mod.rs

@@ -1,11 +1,12 @@
-use args::{init::ProverInitArgs, run::ProverRunArgs};
+use args::{init::ProverInitArgs, init_bellman_cuda::InitBellmanCudaArgs, run::ProverRunArgs};
 use clap::Subcommand;
 use xshell::Shell;
 
 mod args;
 mod gcs;
 mod generate_sk;
 mod init;
+mod init_bellman_cuda;
 mod run;
 mod utils;
 
@@ -17,12 +18,15 @@ pub enum ProverCommands {
     GenerateSK,
     /// Run prover
     Run(ProverRunArgs),
+    /// Initialize bellman-cuda
+    InitBellmanCuda(Box<InitBellmanCudaArgs>),
 }
 
 pub(crate) async fn run(shell: &Shell, args: ProverCommands) -> anyhow::Result<()> {
     match args {
         ProverCommands::Init(args) => init::run(*args, shell).await,
         ProverCommands::GenerateSK => generate_sk::run(shell).await,
         ProverCommands::Run(args) => run::run(args, shell).await,
+        ProverCommands::InitBellmanCuda(args) => init_bellman_cuda::run(shell, *args).await,
     }
 }

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

@@ -1,25 +1,18 @@
-use std::path::PathBuf;
-
 use anyhow::Context;
-use common::{cmd::Cmd, logger};
+use common::{check_prover_prequisites, cmd::Cmd, logger};
 use config::{ChainConfig, EcosystemConfig};
-use path_absolutize::Absolutize;
 use xshell::{cmd, Shell};
 
 use super::{
     args::run::{ProverComponent, ProverRunArgs, WitnessGeneratorArgs, WitnessGeneratorRound},
     utils::get_link_to_prover,
 };
-use crate::{
-    consts::BELLMAN_CUDA_DIR,
-    messages::{
-        MSG_BELLMAN_CUDA_DIR_ERR, MSG_CHAIN_NOT_FOUND_ERR, MSG_MISSING_COMPONENT_ERR,
-        MSG_RUNNING_COMPRESSOR, MSG_RUNNING_COMPRESSOR_ERR, MSG_RUNNING_PROVER,
-        MSG_RUNNING_PROVER_ERR, MSG_RUNNING_PROVER_GATEWAY, MSG_RUNNING_PROVER_GATEWAY_ERR,
-        MSG_RUNNING_WITNESS_GENERATOR, MSG_RUNNING_WITNESS_GENERATOR_ERR,
-        MSG_RUNNING_WITNESS_VECTOR_GENERATOR, MSG_RUNNING_WITNESS_VECTOR_GENERATOR_ERR,
-        MSG_WITNESS_GENERATOR_ROUND_ERR,
-    },
+use crate::messages::{
+    MSG_CHAIN_NOT_FOUND_ERR, MSG_MISSING_COMPONENT_ERR, MSG_RUNNING_COMPRESSOR,
+    MSG_RUNNING_COMPRESSOR_ERR, MSG_RUNNING_PROVER, MSG_RUNNING_PROVER_ERR,
+    MSG_RUNNING_PROVER_GATEWAY, MSG_RUNNING_PROVER_GATEWAY_ERR, MSG_RUNNING_WITNESS_GENERATOR,
+    MSG_RUNNING_WITNESS_GENERATOR_ERR, MSG_RUNNING_WITNESS_VECTOR_GENERATOR,
+    MSG_RUNNING_WITNESS_VECTOR_GENERATOR_ERR, MSG_WITNESS_GENERATOR_ROUND_ERR,
 };
 
 pub(crate) async fn run(args: ProverRunArgs, shell: &Shell) -> anyhow::Result<()> {
@@ -41,14 +34,15 @@ pub(crate) async fn run(args: ProverRunArgs, shell: &Shell) -> anyhow::Result<()
             run_witness_vector_generator(shell, &chain)?
         }
         Some(ProverComponent::Prover) => run_prover(shell, &chain)?,
-        Some(ProverComponent::Compressor) => run_compressor(shell, &chain)?,
+        Some(ProverComponent::Compressor) => run_compressor(shell, &chain, &ecosystem_config)?,
         None => anyhow::bail!(MSG_MISSING_COMPONENT_ERR),
     }
 
     Ok(())
 }
 
 fn run_gateway(shell: &Shell, chain: &ChainConfig) -> anyhow::Result<()> {
+    check_prover_prequisites(shell);
     logger::info(MSG_RUNNING_PROVER_GATEWAY);
     let config_path = chain.path_to_general_config();
     let secrets_path = chain.path_to_secrets_config();
@@ -104,15 +98,16 @@ fn run_prover(shell: &Shell, chain: &ChainConfig) -> anyhow::Result<()> {
     cmd.run().context(MSG_RUNNING_PROVER_ERR)
 }
 
-fn run_compressor(shell: &Shell, chain: &ChainConfig) -> anyhow::Result<()> {
+fn run_compressor(
+    shell: &Shell,
+    chain: &ChainConfig,
+    ecosystem: &EcosystemConfig,
+) -> anyhow::Result<()> {
     logger::info(MSG_RUNNING_COMPRESSOR);
     let config_path = chain.path_to_general_config();
     let secrets_path = chain.path_to_secrets_config();
 
-    let bellman_cuda_dir = PathBuf::from(BELLMAN_CUDA_DIR);
-    let bellman_cuda_dir = bellman_cuda_dir.absolutize()?;
-    let bellman_cuda_str = bellman_cuda_dir.to_str().expect(MSG_BELLMAN_CUDA_DIR_ERR);
-    shell.set_var("BELLMAN_CUDA_DIR", bellman_cuda_str);
+    shell.set_var("BELLMAN_CUDA_DIR", ecosystem.bellman_cuda_dir.clone());
 
     let mut cmd = Cmd::new(cmd!(shell, "cargo run --features gpu --release --bin zksync_proof_fri_compressor -- --config-path={config_path} --secrets-path={secrets_path}"));
     cmd = cmd.with_force_run();

zk_toolbox/crates/zk_inception/src/messages.rs

@@ -263,6 +263,10 @@ pub(super) const MSG_SETUP_KEY_PATH_ERROR: &str = "Failed to get setup key path"
 pub(super) const MSG_CLONING_BELLMAN_CUDA_SPINNER: &str = "Cloning bellman-cuda...";
 pub(super) const MSG_BUILDING_BELLMAN_CUDA_SPINNER: &str = "Building bellman-cuda...";
 pub(super) const MSG_BELLMAN_CUDA_DIR_ERR: &str = "Failed to get bellman-cuda directory";
+pub(super) const MSG_BELLMAN_CUDA_DIR_PROMPT: &str =
+    "Provide the path to the bellman-cuda directory:";
+pub(super) const MSG_BELLMAN_CUDA_INITIALIZED_MSG: &str =
+    "bellman-cuda has been initialized successfully";
 
 pub(super) fn msg_bucket_created(bucket_name: &str) -> String {
     format!("Bucket created successfully with url: gs://{bucket_name}")