feat(zk_toolbox): Ease requirements, add option to download setup keys

zk_toolbox/crates/common/src/lib.rs

@@ -14,6 +14,9 @@ pub mod git;
 pub mod server;
 pub mod wallets;
 
-pub use prerequisites::{check_general_prerequisites, check_prover_prequisites};
+pub use prerequisites::{
+    check_general_prerequisites, check_prerequisites, GCLOUD_PREREQUISITES, GPU_PREREQUISITES,
+    WGET_PREREQUISITES,
+};
 pub use prompt::{init_prompt_theme, Prompt, PromptConfirm, PromptSelect};
 pub use term::{error, logger, spinner};

zk_toolbox/crates/common/src/prerequisites.rs

@@ -30,15 +30,7 @@ const DOCKER_COMPOSE_PREREQUISITE: Prerequisite = Prerequisite {
     download_link: "https://docs.docker.com/compose/install/",
 };
 
-const PROVER_PREREQUISITES: [Prerequisite; 5] = [
-    Prerequisite {
-        name: "gcloud",
-        download_link: "https://cloud.google.com/sdk/docs/install",
-    },
-    Prerequisite {
-        name: "wget",
-        download_link: "https://www.gnu.org/software/wget/",
-    },
+pub const GPU_PREREQUISITES: [Prerequisite; 3] = [
     Prerequisite {
         name: "cmake",
         download_link: "https://cmake.org/download/",
@@ -53,7 +45,17 @@ const PROVER_PREREQUISITES: [Prerequisite; 5] = [
     }, // CUDA GPU driver
 ];
 
-struct Prerequisite {
+pub const WGET_PREREQUISITES: [Prerequisite; 1] = [Prerequisite {
+    name: "wget",
+    download_link: "https://www.gnu.org/software/wget/",
+}];
+
+pub const GCLOUD_PREREQUISITES: [Prerequisite; 1] = [Prerequisite {
+    name: "gcloud",
+    download_link: "https://cloud.google.com/sdk/docs/install",
+}];
+
+pub struct Prerequisite {
     name: &'static str,
     download_link: &'static str,
 }
@@ -62,11 +64,7 @@ pub fn check_general_prerequisites(shell: &Shell) {
     check_prerequisites(shell, &PREREQUISITES, true);
 }
 
-pub fn check_prover_prequisites(shell: &Shell) {
-    check_prerequisites(shell, &PROVER_PREREQUISITES, false);
-}
-
-fn check_prerequisites(shell: &Shell, prerequisites: &[Prerequisite], check_compose: bool) {
+pub fn check_prerequisites(shell: &Shell, prerequisites: &[Prerequisite], check_compose: bool) {
     let mut missing_prerequisites = vec![];
 
     for prerequisite in prerequisites {

zk_toolbox/crates/zk_inception/README.md

@@ -21,7 +21,7 @@ This document contains the help content for the `zk_inception` command-line prog
 - [`zk_inception chain update-token-multiplier-setter`↴](#zk_inception-chain-update-token-multiplier-setter)
 - [`zk_inception prover`↴](#zk_inception-prover)
 - [`zk_inception prover init`↴](#zk_inception-prover-init)
-- [`zk_inception prover generate-sk`↴](#zk_inception-prover-generate-sk)
+- [`zk_inception prover setup-keys`↴](#zk_inception-prover-setup-keys)
 - [`zk_inception prover run`↴](#zk_inception-prover-run)
 - [`zk_inception prover init-bellman-cuda`↴](#zk_inception-prover-init-bellman-cuda)
 - [`zk_inception server`↴](#zk_inception-server)
@@ -475,11 +475,21 @@ Initialize prover
 
   Possible values: `gcp`, `local`
 
-## `zk_inception prover generate-sk`
+## `zk_inception prover setup-keys`
 
-Generate setup keys
+Setup keys
 
-**Usage:** `zk_inception prover generate-sk`
+**Usage:** `zk_inception prover setup-keys`
+
+###### **Options:**
+
+- `--mode`
+
+  Possible valuess: `download`, `generate`
+
+- `--region`
+
+  Possible values: `asia`, `europe`, `us`
 
 ## `zk_inception prover run`
 

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

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

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

@@ -0,0 +1,53 @@
+use clap::{Parser, ValueEnum};
+use common::PromptSelect;
+use strum::{EnumIter, IntoEnumIterator};
+
+use crate::messages::{MSG_SETUP_KEYS_DOWNLOAD_HELP, MSG_SETUP_KEYS_REGION_PROMPT};
+
+#[derive(Debug, Clone, Parser, Default)]
+pub struct SetupKeysArgs {
+    #[clap(long)]
+    pub region: Option<Region>,
+    #[clap(long)]
+    pub mode: Option<Mode>,
+}
+
+#[derive(Debug, Clone)]
+pub struct SetupKeysArgsFinal {
+    pub region: Option<Region>,
+    pub mode: Mode,
+}
+
+#[derive(Debug, Clone, ValueEnum, strum::EnumString, EnumIter, PartialEq, Eq, strum::Display)]
+pub enum Mode {
+    Download,
+    Generate,
+}
+
+#[derive(Debug, Clone, ValueEnum, strum::EnumString, EnumIter, PartialEq, Eq, strum::Display)]
+pub enum Region {
+    Us,
+    Europe,
+    Asia,
+}
+
+impl SetupKeysArgs {
+    pub fn fill_values_with_prompt(self) -> SetupKeysArgsFinal {
+        let mode = self
+            .mode
+            .unwrap_or_else(|| PromptSelect::new(MSG_SETUP_KEYS_DOWNLOAD_HELP, Mode::iter()).ask());
+
+        if mode == Mode::Download {
+            let region = self.region.unwrap_or_else(|| {
+                PromptSelect::new(MSG_SETUP_KEYS_REGION_PROMPT, Region::iter()).ask()
+            });
+
+            SetupKeysArgsFinal {
+                region: Some(region),
+                mode,
+            }
+        } else {
+            SetupKeysArgsFinal { region: None, mode }
+        }
+    }
+}

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

@@ -1,4 +1,4 @@
-use common::{cmd::Cmd, logger, spinner::Spinner};
+use common::{check_prerequisites, cmd::Cmd, logger, spinner::Spinner, GCLOUD_PREREQUISITES};
 use xshell::{cmd, Shell};
 use zksync_config::{configs::object_store::ObjectStoreMode, ObjectStoreConfig};
 
@@ -14,6 +14,8 @@ pub(crate) fn create_gcs_bucket(
     shell: &Shell,
     config: ProofStorageGCSCreateBucket,
 ) -> anyhow::Result<ObjectStoreConfig> {
+    check_prerequisites(shell, &GCLOUD_PREREQUISITES, false);
+
     let bucket_name = config.bucket_name;
     let location = config.location;
     let project_id = config.project_id;

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

@@ -2,12 +2,13 @@ use std::path::PathBuf;
 
 use anyhow::Context;
 use common::{
-    check_prover_prequisites,
+    check_prerequisites,
     cmd::Cmd,
     config::global_config,
     db::{drop_db_if_exists, init_db, migrate_db, DatabaseConfig},
     logger,
     spinner::Spinner,
+    WGET_PREREQUISITES,
 };
 use config::{copy_configs, set_prover_database, traits::SaveConfigWithBasePath, EcosystemConfig};
 use xshell::{cmd, Shell};
@@ -34,8 +35,6 @@ use crate::{
 };
 
 pub(crate) async fn run(args: ProverInitArgs, shell: &Shell) -> anyhow::Result<()> {
-    check_prover_prequisites(shell);
-
     let ecosystem_config = EcosystemConfig::from_file(shell)?;
 
     let setup_key_path = get_default_setup_key_path(&ecosystem_config)?;
@@ -115,6 +114,7 @@ fn download_setup_key(
     general_config: &GeneralConfig,
     path: &str,
 ) -> anyhow::Result<()> {
+    check_prerequisites(shell, &WGET_PREREQUISITES, false);
     let spinner = Spinner::new(MSG_DOWNLOADING_SETUP_KEY_SPINNER);
     let compressor_config: zksync_config::configs::FriProofCompressorConfig = general_config
         .proof_compressor_config

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

@@ -1,5 +1,5 @@
 use anyhow::Context;
-use common::{check_prover_prequisites, cmd::Cmd, git, logger, spinner::Spinner};
+use common::{check_prerequisites, cmd::Cmd, git, logger, spinner::Spinner, GPU_PREREQUISITES};
 use config::{traits::SaveConfigWithBasePath, EcosystemConfig};
 use xshell::{cmd, Shell};
 
@@ -13,7 +13,7 @@ use crate::{
 };
 
 pub(crate) async fn run(shell: &Shell, args: InitBellmanCudaArgs) -> anyhow::Result<()> {
-    check_prover_prequisites(shell);
+    check_prerequisites(shell, &GPU_PREREQUISITES, false);
 
     let mut ecosystem_config = EcosystemConfig::from_file(shell)?;
 

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

@@ -2,12 +2,14 @@ use args::{init::ProverInitArgs, init_bellman_cuda::InitBellmanCudaArgs, run::Pr
 use clap::Subcommand;
 use xshell::Shell;
 
+use crate::commands::prover::args::setup_keys::SetupKeysArgs;
+
 mod args;
 mod gcs;
-mod generate_sk;
 mod init;
 mod init_bellman_cuda;
 mod run;
+mod setup_keys;
 mod utils;
 
 #[derive(Subcommand, Debug)]
@@ -16,7 +18,7 @@ pub enum ProverCommands {
     Init(Box<ProverInitArgs>),
     /// Generate setup keys
     #[command(alias = "sk")]
-    GenerateSK,
+    SetupKeys(SetupKeysArgs),
     /// Run prover
     Run(ProverRunArgs),
     /// Initialize bellman-cuda
@@ -27,7 +29,7 @@ pub enum ProverCommands {
 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::SetupKeys(args) => setup_keys::run(args, 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,5 +1,5 @@
 use anyhow::Context;
-use common::{check_prover_prequisites, cmd::Cmd, config::global_config, logger};
+use common::{check_prerequisites, cmd::Cmd, config::global_config, logger, GPU_PREREQUISITES};
 use config::{ChainConfig, EcosystemConfig};
 use xshell::{cmd, Shell};
 
@@ -20,7 +20,6 @@ use crate::messages::{
 };
 
 pub(crate) async fn run(args: ProverRunArgs, shell: &Shell) -> anyhow::Result<()> {
-    check_prover_prequisites(shell);
     let args = args.fill_values_with_prompt()?;
     let ecosystem_config = EcosystemConfig::from_file(shell)?;
     let chain = ecosystem_config
@@ -97,6 +96,7 @@ fn run_witness_vector_generator(
 }
 
 fn run_prover(shell: &Shell, chain: &ChainConfig) -> anyhow::Result<()> {
+    check_prerequisites(shell, &GPU_PREREQUISITES, false);
     logger::info(MSG_RUNNING_PROVER);
     let config_path = chain.path_to_general_config();
     let secrets_path = chain.path_to_secrets_config();
@@ -113,6 +113,7 @@ fn run_compressor(
     chain: &ChainConfig,
     ecosystem: &EcosystemConfig,
 ) -> anyhow::Result<()> {
+    check_prerequisites(shell, &GPU_PREREQUISITES, false);
     logger::info(MSG_RUNNING_COMPRESSOR);
     let config_path = chain.path_to_general_config();
     let secrets_path = chain.path_to_secrets_config();

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

@@ -0,0 +1,83 @@
+use anyhow::Ok;
+use common::{
+    check_prerequisites, cmd::Cmd, logger, spinner::Spinner, GCLOUD_PREREQUISITES,
+    GPU_PREREQUISITES,
+};
+use config::EcosystemConfig;
+use xshell::{cmd, Shell};
+
+use super::utils::get_link_to_prover;
+use crate::{
+    commands::prover::args::setup_keys::{Mode, Region, SetupKeysArgs},
+    messages::{MSG_GENERATING_SK_SPINNER, MSG_SK_GENERATED},
+};
+
+pub(crate) async fn run(args: SetupKeysArgs, shell: &Shell) -> anyhow::Result<()> {
+    let args = args.fill_values_with_prompt();
+    let ecosystem_config = EcosystemConfig::from_file(shell)?;
+
+    if args.mode == Mode::Generate {
+        check_prerequisites(shell, &GPU_PREREQUISITES, false);
+        let link_to_prover = get_link_to_prover(&ecosystem_config);
+        shell.change_dir(&link_to_prover);
+
+        let spinner = Spinner::new(MSG_GENERATING_SK_SPINNER);
+        let cmd = Cmd::new(cmd!(
+            shell,
+            "cargo run --features gpu --release --bin key_generator -- 
+            generate-sk-gpu all --recompute-if-missing
+            --setup-path=data/keys
+            --path={link_to_prover}/data/keys"
+        ));
+        cmd.run()?;
+        spinner.finish();
+        logger::outro(MSG_SK_GENERATED);
+    } else {
+        check_prerequisites(shell, &GCLOUD_PREREQUISITES, false);
+
+        let link_to_setup_keys = get_link_to_prover(&ecosystem_config).join("data/keys");
+        let path_to_keys_buckets =
+            get_link_to_prover(&ecosystem_config).join("setup-data-gpu-keys.json");
+
+        let region = args.region.expect("Region is not provided");
+
+        let file = shell
+            .read_file(path_to_keys_buckets)
+            .expect("Could not find commitments file in zksync-era");
+        let json: serde_json::Value =
+            serde_json::from_str(&file).expect("Could not parse commitments.json");
+
+        let bucket = &match region {
+            Region::Us => json
+                .get("us")
+                .expect("Could not find link to US bucket")
+                .to_string(),
+            Region::Europe => json
+                .get("europe")
+                .expect("Could not find link to Europe bucket")
+                .to_string(),
+            Region::Asia => json
+                .get("asia")
+                .expect("Could not find link to Asia bucket")
+                .to_string(),
+        };
+
+        let len = bucket.len() - 2usize;
+        let bucket = &bucket[1..len];
+
+        let spinner = Spinner::new(&format!(
+            "Downloading keys from bucket: {} to {:?}",
+            bucket, link_to_setup_keys
+        ));
+
+        let cmd = Cmd::new(cmd!(
+            shell,
+            "gsutil -m rsync -r {bucket} {link_to_setup_keys}"
+        ));
+        cmd.run()?;
+        spinner.finish();
+        logger::outro("Keys are downloaded");
+    }
+
+    Ok(())
+}