From 0b86aa0e500ef9da085ca93c1e210aceaf9581f4 Mon Sep 17 00:00:00 2001 From: matias-gonz Date: Thu, 27 Jun 2024 13:17:17 +0200 Subject: [PATCH] Add file based config for witness vector generator --- prover/Cargo.lock | 2 + prover/witness_vector_generator/Cargo.toml | 2 + prover/witness_vector_generator/src/main.rs | 63 ++++++++++++--------- 3 files changed, 41 insertions(+), 26 deletions(-) diff --git a/prover/Cargo.lock b/prover/Cargo.lock index 7de9254ed2ee..8719e133ed74 100644 --- a/prover/Cargo.lock +++ b/prover/Cargo.lock @@ -8481,6 +8481,7 @@ dependencies = [ "anyhow", "async-trait", "bincode", + "clap 4.5.4", "ctrlc", "futures 0.3.30", "queues", @@ -8493,6 +8494,7 @@ dependencies = [ "zksync_config", "zksync_env_config", "zksync_object_store", + "zksync_prover_config", "zksync_prover_dal", "zksync_prover_fri_types", "zksync_prover_fri_utils", diff --git a/prover/witness_vector_generator/Cargo.toml b/prover/witness_vector_generator/Cargo.toml index 2b95d81d49e9..cf218ed8ae30 100644 --- a/prover/witness_vector_generator/Cargo.toml +++ b/prover/witness_vector_generator/Cargo.toml @@ -19,6 +19,7 @@ zksync_object_store.workspace = true zksync_prover_fri_utils.workspace = true zksync_utils.workspace = true zksync_prover_fri_types.workspace = true +zksync_prover_config.workspace = true zksync_queued_job_processor.workspace = true vk_setup_data_generator_server_fri.workspace = true zksync_vlog.workspace = true @@ -30,6 +31,7 @@ tokio = { workspace = true, features = ["time", "macros"] } futures = { workspace = true, features = ["compat"] } ctrlc = { workspace = true, features = ["termination"] } serde = { workspace = true, features = ["derive"] } +clap = { workspace = true, features = ["derive"] } async-trait.workspace = true queues.workspace = true bincode.workspace = true diff --git a/prover/witness_vector_generator/src/main.rs b/prover/witness_vector_generator/src/main.rs index 212abf1cb4e2..1649c8e82acd 100644 --- a/prover/witness_vector_generator/src/main.rs +++ b/prover/witness_vector_generator/src/main.rs @@ -3,14 +3,11 @@ use std::time::Duration; use anyhow::Context as _; -use structopt::StructOpt; +use clap::Parser; use tokio::sync::{oneshot, watch}; -use zksync_config::configs::{ - fri_prover_group::FriProverGroupConfig, DatabaseSecrets, FriProverConfig, - FriWitnessVectorGeneratorConfig, ObservabilityConfig, -}; -use zksync_env_config::{object_store::ProverObjectStoreConfig, FromEnv}; +use zksync_env_config::object_store::ProverObjectStoreConfig; use zksync_object_store::ObjectStoreFactory; +use zksync_prover_config::{load_database_secrets, load_general_config}; use zksync_prover_dal::ConnectionPool; use zksync_prover_fri_types::PROVER_PROTOCOL_SEMANTIC_VERSION; use zksync_prover_fri_utils::{get_all_circuit_id_round_tuples_for, region_fetcher::get_zone}; @@ -23,21 +20,29 @@ use crate::generator::WitnessVectorGenerator; mod generator; mod metrics; -#[derive(Debug, StructOpt)] -#[structopt( - name = "zksync_witness_vector_generator", - about = "Tool for generating witness vectors for circuits" -)] -struct Opt { +#[derive(Debug, Parser)] +#[command(author = "Matter Labs", version)] +struct Cli { /// Number of times `witness_vector_generator` should be run. - #[structopt(short = "n", long = "n_iterations")] - number_of_iterations: Option, + #[arg(long)] + #[arg(short)] + n_iterations: Option, + #[arg(long)] + pub(crate) config_path: Option, + #[arg(long)] + pub(crate) secrets_path: Option, } #[tokio::main] async fn main() -> anyhow::Result<()> { - let observability_config = - ObservabilityConfig::from_env().context("ObservabilityConfig::from_env()")?; + let opt = Cli::parse(); + + let general_config = load_general_config(opt.config_path).context("general config")?; + let database_secrets = load_database_secrets(opt.secrets_path).context("database secrets")?; + + let observability_config = general_config + .observability + .context("observability config")?; let log_format: zksync_vlog::LogFormat = observability_config .log_format .parse() @@ -61,29 +66,35 @@ async fn main() -> anyhow::Result<()> { } let _guard = builder.build(); - let opt = Opt::from_args(); - let config = FriWitnessVectorGeneratorConfig::from_env() - .context("FriWitnessVectorGeneratorConfig::from_env()")?; + let config = general_config + .witness_vector_generator + .context("witness vector generator config")?; let specialized_group_id = config.specialized_group_id; let exporter_config = PrometheusExporterConfig::pull(config.prometheus_listener_port); - let database_secrets = DatabaseSecrets::from_env().context("DatabaseSecrets::from_env()")?; let pool = ConnectionPool::singleton(database_secrets.prover_url()?) .build() .await .context("failed to build a connection pool")?; - let object_store_config = - ProverObjectStoreConfig::from_env().context("ProverObjectStoreConfig::from_env()")?; + let object_store_config = ProverObjectStoreConfig( + general_config + .prover_config + .clone() + .context("prover config")? + .prover_object_store + .context("object store")?, + ); let object_store = ObjectStoreFactory::new(object_store_config.0) .create_store() .await?; - let circuit_ids_for_round_to_be_proven = FriProverGroupConfig::from_env() - .context("FriProverGroupConfig::from_env()")? + let circuit_ids_for_round_to_be_proven = general_config + .prover_group_config + .expect("prover_group_config") .get_circuit_ids_for_group_id(specialized_group_id) .unwrap_or_default(); let circuit_ids_for_round_to_be_proven = get_all_circuit_id_round_tuples_for(circuit_ids_for_round_to_be_proven); - let fri_prover_config = FriProverConfig::from_env().context("FriProverConfig::from_env()")?; + let fri_prover_config = general_config.prover_config.context("prover config")?; let zone_url = &fri_prover_config.zone_read_url; let zone = get_zone(zone_url).await.context("get_zone()")?; @@ -114,7 +125,7 @@ async fn main() -> anyhow::Result<()> { let tasks = vec![ tokio::spawn(exporter_config.run(stop_receiver.clone())), - tokio::spawn(witness_vector_generator.run(stop_receiver, opt.number_of_iterations)), + tokio::spawn(witness_vector_generator.run(stop_receiver, opt.n_iterations)), ]; let mut tasks = ManagedTasks::new(tasks);