Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(zk_toolbox): Add setup keys step to prover init #2811

Merged
merged 9 commits into from
Sep 9, 2024
Merged

feat(zk_toolbox): Add setup keys step to prover init #2811

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
f87d3f0
Add setup keys step to init
matias-gonz Sep 5, 2024
a0abf3b
Update messages
matias-gonz Sep 5, 2024
9142719
lint
matias-gonz Sep 5, 2024
f49b471
Merge branch 'main' into matias-zki-prover-setup-keys
matias-gonz Sep 6, 2024
f8a8276
make download compressor key default false
matias-gonz Sep 6, 2024
682ec99
Make init bellman cuda optional
matias-gonz Sep 6, 2024
1b8293a
Separate compressor key step
matias-gonz Sep 9, 2024
2574ddb
fmt
matias-gonz Sep 9, 2024
2503be9
Merge branch 'main' into matias-zki-prover-setup-keys
matias-gonz Sep 9, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
65 changes: 43 additions & 22 deletions zk_toolbox/crates/zk_inception/src/commands/prover/args/init.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand All @@ -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<String>,
#[clap(flatten)]
#[serde(flatten)]
pub proof_store_gcs_config: ProofStorageGCSTmp,
#[clap(flatten)]
#[serde(flatten)]
pub create_gcs_bucket_config: ProofStorageGCSCreateBucketTmp,

// Public store object
Expand All @@ -45,20 +44,21 @@ pub struct ProverInitArgs {
#[clap(long)]
pub public_store_dir: Option<String>,
#[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<bool>,

#[clap(long)]
pub setup_database: Option<bool>,
Expand Down Expand Up @@ -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<bool>,
#[clap(long)]
Expand Down Expand Up @@ -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,
}
Expand All @@ -187,7 +187,8 @@ pub struct ProverDatabaseConfig {
pub struct ProverInitArgsFinal {
pub proof_store: ProofStorageConfig,
pub public_store: Option<ProofStorageConfig>,
pub setup_key_config: SetupKeyConfig,
pub setup_compressor_key_config: SetupCompressorKeyConfig,
pub setup_keys: Option<SetupKeysArgs>,
pub bellman_cuda_config: InitBellmanCudaArgs,
pub cloud_type: CloudConnectionMode,
pub database_config: Option<ProverDatabaseConfig>,
Expand All @@ -202,15 +203,18 @@ impl ProverInitArgs {
) -> anyhow::Result<ProverInitArgsFinal> {
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,
Expand Down Expand Up @@ -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<SetupKeysArgs> {
let args = self.setup_keys_args.clone();

if self.setup_keys.unwrap_or_else(|| {
PromptConfirm::new(MSG_SETUP_KEYS_PROMPT)
.default(false)
.ask()
}) {
Some(args)
} else {
None
}
}

fn partial_create_gcs_bucket_config_provided(
&self,
bucket_name: Option<String>,
Expand Down
8 changes: 4 additions & 4 deletions zk_toolbox/crates/zk_inception/src/commands/prover/args/setup_keys.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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(|| {
Expand Down
15 changes: 10 additions & 5 deletions zk_toolbox/crates/zk_inception/src/commands/prover/init.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,12 +21,13 @@ 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::{
consts::{PROVER_MIGRATIONS, PROVER_STORE_MAX_RETRIES},
messages::{
MSG_CHAIN_NOT_FOUND_ERR, MSG_DOWNLOADING_SETUP_KEY_SPINNER,
MSG_CHAIN_NOT_FOUND_ERR, MSG_DOWNLOADING_SETUP_COMPRESSOR_KEY_SPINNER,
MSG_FAILED_TO_DROP_PROVER_DATABASE_ERR, MSG_GENERAL_CONFIG_NOT_FOUND_ERR,
MSG_INITIALIZING_DATABASES_SPINNER, MSG_INITIALIZING_PROVER_DATABASE,
MSG_PROOF_COMPRESSOR_CONFIG_NOT_FOUND_ERR, MSG_PROVER_CONFIG_NOT_FOUND_ERR,
Expand Down Expand Up @@ -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);
Expand All @@ -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)?;
Expand Down Expand Up @@ -115,7 +120,7 @@ fn download_setup_key(
path: &str,
) -> anyhow::Result<()> {
check_prerequisites(shell, &WGET_PREREQUISITES, false);
let spinner = Spinner::new(MSG_DOWNLOADING_SETUP_KEY_SPINNER);
let spinner = Spinner::new(MSG_DOWNLOADING_SETUP_COMPRESSOR_KEY_SPINNER);
let compressor_config: zksync_config::configs::FriProofCompressorConfig = general_config
.proof_compressor_config
.as_ref()
Expand Down
11 changes: 7 additions & 4 deletions zk_toolbox/crates/zk_inception/src/messages.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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?";
Expand Down Expand Up @@ -301,9 +301,11 @@ pub(super) const MSG_CREATE_GCS_BUCKET_NAME_PROMTP: &str = "What do you want to
pub(super) const MSG_CREATE_GCS_BUCKET_LOCATION_PROMPT: &str = "What location do you want to use? Find available locations at https://cloud.google.com/storage/docs/locations";
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_DOWNLOADING_SETUP_COMPRESSOR_KEY_SPINNER: &str =
"Downloading compressor 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...";
Expand All @@ -325,6 +327,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}")
Expand Down
Loading