From f87d3f0e0c85a53066df368fd32ea86198d175f5 Mon Sep 17 00:00:00 2001 From: matias-gonz Date: Thu, 5 Sep 2024 11:02:40 -0300 Subject: [PATCH] Add setup keys step to init --- .../src/commands/prover/args/init.rs | 65 ++++++++++++------- .../src/commands/prover/args/setup_keys.rs | 8 +-- .../zk_inception/src/commands/prover/init.rs | 11 +++- .../crates/zk_inception/src/messages.rs | 8 ++- 4 files changed, 60 insertions(+), 32 deletions(-) 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 e8c9cf1888d5..6a18cc0c38d7 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 @@ -8,7 +8,7 @@ use url::Url; use xshell::Shell; use zksync_config::configs::fri_prover::CloudConnectionMode; -use super::init_bellman_cuda::InitBellmanCudaArgs; +use super::{init_bellman_cuda::InitBellmanCudaArgs, setup_keys::SetupKeysArgs}; use crate::{ commands::prover::gcs::get_project_ids, consts::{DEFAULT_CREDENTIALS_FILE, DEFAULT_PROOF_STORE_DIR}, @@ -18,25 +18,24 @@ use crate::{ 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_DOWNLOAD_SETUP_COMPRESSOR_KEY_PROMPT, MSG_GETTING_PROOF_STORE_CONFIG, MSG_GETTING_PUBLIC_STORE_CONFIG, MSG_PROOF_STORE_CONFIG_PROMPT, MSG_PROOF_STORE_DIR_PROMPT, MSG_PROOF_STORE_GCS_BUCKET_BASE_URL_ERR, MSG_PROOF_STORE_GCS_BUCKET_BASE_URL_PROMPT, MSG_PROOF_STORE_GCS_CREDENTIALS_FILE_PROMPT, MSG_PROVER_DB_NAME_HELP, - MSG_PROVER_DB_URL_HELP, MSG_SAVE_TO_PUBLIC_BUCKET_PROMPT, MSG_SETUP_KEY_PATH_PROMPT, + MSG_PROVER_DB_URL_HELP, MSG_SAVE_TO_PUBLIC_BUCKET_PROMPT, + MSG_SETUP_COMPRESSOR_KEY_PATH_PROMPT, MSG_SETUP_KEYS_PROMPT, MSG_USE_DEFAULT_DATABASES_HELP, }, }; -#[derive(Debug, Clone, Serialize, Deserialize, Parser, Default)] +#[derive(Debug, Clone, Parser, Default)] pub struct ProverInitArgs { // Proof store object #[clap(long)] pub proof_store_dir: Option, #[clap(flatten)] - #[serde(flatten)] pub proof_store_gcs_config: ProofStorageGCSTmp, #[clap(flatten)] - #[serde(flatten)] pub create_gcs_bucket_config: ProofStorageGCSCreateBucketTmp, // Public store object @@ -45,20 +44,21 @@ pub struct ProverInitArgs { #[clap(long)] pub public_store_dir: Option, #[clap(flatten)] - #[serde(flatten)] pub public_store_gcs_config: PublicStorageGCSTmp, #[clap(flatten)] - #[serde(flatten)] pub public_create_gcs_bucket_config: PublicStorageGCSCreateBucketTmp, // Bellman cuda #[clap(flatten)] - #[serde(flatten)] pub bellman_cuda_config: InitBellmanCudaArgs, #[clap(flatten)] - #[serde(flatten)] - pub setup_key_config: SetupKeyConfigTmp, + pub setup_compressor_key_config: SetupCompressorKeyConfigTmp, + + #[clap(flatten)] + pub setup_keys_args: SetupKeysArgs, + #[clap(long, default_missing_value = "true", num_args = 0..=1)] + pub setup_keys: Option, #[clap(long)] pub setup_database: Option, @@ -137,7 +137,7 @@ pub struct PublicStorageGCSCreateBucketTmp { } #[derive(Clone, Debug, Serialize, Deserialize, Parser, Default)] -pub struct SetupKeyConfigTmp { +pub struct SetupCompressorKeyConfigTmp { #[clap(long)] pub download_key: Option, #[clap(long)] @@ -172,7 +172,7 @@ pub enum ProofStorageConfig { } #[derive(Debug, Clone)] -pub struct SetupKeyConfig { +pub struct SetupCompressorKeyConfig { pub download_key: bool, pub setup_key_path: String, } @@ -187,7 +187,8 @@ pub struct ProverDatabaseConfig { pub struct ProverInitArgsFinal { pub proof_store: ProofStorageConfig, pub public_store: Option, - pub setup_key_config: SetupKeyConfig, + pub setup_compressor_key_config: SetupCompressorKeyConfig, + pub setup_keys: Option, pub bellman_cuda_config: InitBellmanCudaArgs, pub cloud_type: CloudConnectionMode, pub database_config: Option, @@ -202,15 +203,18 @@ impl ProverInitArgs { ) -> anyhow::Result { let proof_store = self.fill_proof_storage_values_with_prompt(shell)?; 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 setup_compressor_key_config = + self.fill_setup_compressor_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(); let database_config = self.fill_database_values_with_prompt(chain_config); + let setup_keys = self.fill_setup_keys_values_with_prompt(); Ok(ProverInitArgsFinal { proof_store, public_store, - setup_key_config, + setup_compressor_key_config, + setup_keys, bellman_cuda_config, cloud_type, database_config, @@ -336,32 +340,49 @@ impl ProverInitArgs { } } - fn fill_setup_key_values_with_prompt(&self, setup_key_path: &str) -> SetupKeyConfig { + fn fill_setup_compressor_key_values_with_prompt( + &self, + setup_key_path: &str, + ) -> SetupCompressorKeyConfig { let download_key = self .clone() - .setup_key_config + .setup_compressor_key_config .download_key .unwrap_or_else(|| { - PromptConfirm::new(MSG_DOWNLOAD_SETUP_KEY_PROMPT) + PromptConfirm::new(MSG_DOWNLOAD_SETUP_COMPRESSOR_KEY_PROMPT) .default(true) .ask() }); let setup_key_path = self .clone() - .setup_key_config + .setup_compressor_key_config .setup_key_path .unwrap_or_else(|| { - Prompt::new(MSG_SETUP_KEY_PATH_PROMPT) + Prompt::new(MSG_SETUP_COMPRESSOR_KEY_PATH_PROMPT) .default(setup_key_path) .ask() }); - SetupKeyConfig { + SetupCompressorKeyConfig { download_key, setup_key_path, } } + fn fill_setup_keys_values_with_prompt(&self) -> Option { + let args = self.setup_keys_args.clone(); + + if self.setup_keys.unwrap_or_else(|| { + PromptConfirm::new(MSG_SETUP_KEYS_PROMPT) + .default(true) + .ask() + }) { + Some(args) + } else { + None + } + } + fn partial_create_gcs_bucket_config_provided( &self, bucket_name: Option, diff --git a/zk_toolbox/crates/zk_inception/src/commands/prover/args/setup_keys.rs b/zk_toolbox/crates/zk_inception/src/commands/prover/args/setup_keys.rs index 4839c03eb429..155977b8812a 100644 --- a/zk_toolbox/crates/zk_inception/src/commands/prover/args/setup_keys.rs +++ b/zk_toolbox/crates/zk_inception/src/commands/prover/args/setup_keys.rs @@ -2,7 +2,7 @@ use clap::{Parser, ValueEnum}; use common::PromptSelect; use strum::{EnumIter, IntoEnumIterator}; -use crate::messages::{MSG_SETUP_KEYS_DOWNLOAD_HELP, MSG_SETUP_KEYS_REGION_PROMPT}; +use crate::messages::{MSG_SETUP_KEYS_DOWNLOAD_SELECTION_PROMPT, MSG_SETUP_KEYS_REGION_PROMPT}; #[derive(Debug, Clone, Parser, Default)] pub struct SetupKeysArgs { @@ -33,9 +33,9 @@ pub enum Region { 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()); + let mode = self.mode.unwrap_or_else(|| { + PromptSelect::new(MSG_SETUP_KEYS_DOWNLOAD_SELECTION_PROMPT, Mode::iter()).ask() + }); if mode == Mode::Download { let region = self.region.unwrap_or_else(|| { 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 7aadd04bf6b7..17fecf8a8443 100644 --- a/zk_toolbox/crates/zk_inception/src/commands/prover/init.rs +++ b/zk_toolbox/crates/zk_inception/src/commands/prover/init.rs @@ -21,6 +21,7 @@ use super::{ args::init::{ProofStorageConfig, ProverInitArgs}, gcs::create_gcs_bucket, init_bellman_cuda::run as init_bellman_cuda, + setup_keys, utils::get_link_to_prover, }; use crate::{ @@ -55,14 +56,18 @@ pub(crate) async fn run(args: ProverInitArgs, shell: &Shell) -> anyhow::Result<( let proof_object_store_config = get_object_store_config(shell, Some(args.proof_store))?; let public_object_store_config = get_object_store_config(shell, args.public_store)?; - if args.setup_key_config.download_key { + if args.setup_compressor_key_config.download_key { download_setup_key( shell, &general_config, - &args.setup_key_config.setup_key_path, + &args.setup_compressor_key_config.setup_key_path, )?; } + if let Some(args) = args.setup_keys { + setup_keys::run(args, &shell).await?; + } + let mut prover_config = general_config .prover_config .expect(MSG_PROVER_CONFIG_NOT_FOUND_ERR); @@ -81,7 +86,7 @@ pub(crate) async fn run(args: ProverInitArgs, shell: &Shell) -> anyhow::Result<( let mut proof_compressor_config = general_config .proof_compressor_config .expect(MSG_PROOF_COMPRESSOR_CONFIG_NOT_FOUND_ERR); - proof_compressor_config.universal_setup_path = args.setup_key_config.setup_key_path; + proof_compressor_config.universal_setup_path = args.setup_compressor_key_config.setup_key_path; general_config.proof_compressor_config = Some(proof_compressor_config); chain_config.save_general_config(&general_config)?; diff --git a/zk_toolbox/crates/zk_inception/src/messages.rs b/zk_toolbox/crates/zk_inception/src/messages.rs index 25933d39db30..04c082b5f849 100644 --- a/zk_toolbox/crates/zk_inception/src/messages.rs +++ b/zk_toolbox/crates/zk_inception/src/messages.rs @@ -5,7 +5,7 @@ use ethers::{ utils::format_ether, }; -pub(super) const MSG_SETUP_KEYS_DOWNLOAD_HELP: &str = +pub(super) const MSG_SETUP_KEYS_DOWNLOAD_SELECTION_PROMPT: &str = "Do you want to download the setup keys or generate them?"; pub(super) const MSG_SETUP_KEYS_REGION_PROMPT: &str = "From which region you want setup keys to be downloaded?"; @@ -302,8 +302,9 @@ pub(super) const MSG_CREATE_GCS_BUCKET_LOCATION_PROMPT: &str = "What location do pub(super) const MSG_PROOF_COMPRESSOR_CONFIG_NOT_FOUND_ERR: &str = "Proof compressor config not found"; pub(super) const MSG_DOWNLOADING_SETUP_KEY_SPINNER: &str = "Downloading setup key..."; -pub(super) const MSG_DOWNLOAD_SETUP_KEY_PROMPT: &str = "Do you want to download the setup key?"; -pub(super) const MSG_SETUP_KEY_PATH_PROMPT: &str = "Provide the path to the setup key:"; +pub(super) const MSG_DOWNLOAD_SETUP_COMPRESSOR_KEY_PROMPT: &str = + "Do you want to download the setup key for compressor?"; +pub(super) const MSG_SETUP_COMPRESSOR_KEY_PATH_PROMPT: &str = "Provide the path to the setup key:"; pub(super) const MSG_GETTING_GCP_PROJECTS_SPINNER: &str = "Getting GCP projects..."; pub(super) const MSG_GETTING_PROOF_STORE_CONFIG: &str = "Getting proof store configuration..."; pub(super) const MSG_GETTING_PUBLIC_STORE_CONFIG: &str = "Getting public store configuration..."; @@ -325,6 +326,7 @@ pub(super) const MSG_BELLMAN_CUDA_SELECTION_CLONE: &str = "Clone for me (recomme 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 connection mode:"; pub(super) const MSG_THREADS_PROMPT: &str = "Provide the number of threads:"; +pub(super) const MSG_SETUP_KEYS_PROMPT: &str = "Do you want to setup keys?"; pub(super) fn msg_bucket_created(bucket_name: &str) -> String { format!("Bucket created successfully with url: gs://{bucket_name}")