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.

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(Prover CLI): Configuration (with flag) #1861

Merged
merged 10 commits into from
May 8, 2024
21 changes: 18 additions & 3 deletions prover/prover_cli/src/cli.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use clap::{command, Parser, Subcommand};
use clap::{command, Args, Parser, Subcommand};
use zksync_types::url::SensitiveUrl;

use crate::commands::{self, get_file_info};

Expand All @@ -9,6 +10,20 @@ pub const VERSION_STRING: &str = env!("CARGO_PKG_VERSION");
struct ProverCLI {
#[command(subcommand)]
command: ProverCommand,
#[clap(flatten)]
config: ProverCLIConfig,
}

// Note: This is a temporary solution for the configuration of the CLI. In the
// future, we should have an `config` command to set the configuration in a
// `.config` file.
#[derive(Args)]
pub struct ProverCLIConfig {
#[clap(
long,
default_value = "postgres://postgres:notsecurepassword@localhost/prover_local"
)]
pub db_url: SensitiveUrl,
}

#[derive(Subcommand)]
Expand All @@ -19,10 +34,10 @@ enum ProverCommand {
}

pub async fn start() -> anyhow::Result<()> {
let ProverCLI { command } = ProverCLI::parse();
let ProverCLI { command, config } = ProverCLI::parse();
match command {
ProverCommand::FileInfo(args) => get_file_info::run(args).await?,
ProverCommand::Status(cmd) => cmd.run().await?,
ProverCommand::Status(cmd) => cmd.run(config).await?,
};

Ok(())
Expand Down
25 changes: 13 additions & 12 deletions prover/prover_cli/src/commands/status/batch.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,12 @@ use anyhow::{ensure, Context as _};
use clap::Args as ClapArgs;
use prover_dal::{Connection, ConnectionPool, Prover, ProverDal};
use zksync_types::{
basic_fri_types::AggregationRound, prover_dal::WitnessJobStatus, L1BatchNumber,
basic_fri_types::AggregationRound, prover_dal::WitnessJobStatus, url::SensitiveUrl,
L1BatchNumber,
};

use super::utils::{AggregationRoundInfo, BatchData, Task, TaskStatus};
use crate::commands::status::utils::postgres_config;
use crate::cli::ProverCLIConfig;

#[derive(ClapArgs)]
pub struct Args {
Expand All @@ -16,13 +17,13 @@ pub struct Args {
verbose: bool,
}

pub(crate) async fn run(args: Args) -> anyhow::Result<()> {
pub(crate) async fn run(args: Args, config: ProverCLIConfig) -> anyhow::Result<()> {
ensure!(
!args.batches.is_empty(),
"At least one batch number should be provided"
);

let batches_data = get_batches_data(args.batches).await?;
let batches_data = get_batches_data(args.batches, config.db_url).await?;

for batch_data in batches_data {
println!("{batch_data:?}");
Expand All @@ -31,14 +32,14 @@ pub(crate) async fn run(args: Args) -> anyhow::Result<()> {
Ok(())
}

async fn get_batches_data(batches: Vec<L1BatchNumber>) -> anyhow::Result<Vec<BatchData>> {
let config = postgres_config()?;

let prover_connection_pool =
ConnectionPool::<Prover>::builder(config.prover_url()?, config.max_connections()?)
.build()
.await
.context("failed to build a prover_connection_pool")?;
async fn get_batches_data(
batches: Vec<L1BatchNumber>,
db_url: SensitiveUrl,
) -> anyhow::Result<Vec<BatchData>> {
let prover_connection_pool = ConnectionPool::<Prover>::singleton(db_url)
.build()
.await
.context("failed to build a prover_connection_pool")?;

let mut conn = prover_connection_pool
.connection()
Expand Down
6 changes: 4 additions & 2 deletions prover/prover_cli/src/commands/status/mod.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
use clap::Subcommand;

use crate::cli::ProverCLIConfig;

pub(crate) mod batch;
mod utils;

Expand All @@ -9,9 +11,9 @@ pub enum StatusCommand {
}

impl StatusCommand {
pub(crate) async fn run(self) -> anyhow::Result<()> {
pub(crate) async fn run(self, config: ProverCLIConfig) -> anyhow::Result<()> {
match self {
StatusCommand::Batch(args) => batch::run(args).await,
StatusCommand::Batch(args) => batch::run(args, config).await,
}
}
}
6 changes: 0 additions & 6 deletions prover/prover_cli/src/commands/status/utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,18 +2,12 @@ use std::fmt::Debug;

use colored::*;
use strum::{Display, EnumString};
use zksync_config::PostgresConfig;
use zksync_env_config::FromEnv;
use zksync_types::{
basic_fri_types::AggregationRound,
prover_dal::{ProofCompressionJobStatus, ProverJobFriInfo, ProverJobStatus, WitnessJobStatus},
L1BatchNumber,
};

pub fn postgres_config() -> anyhow::Result<PostgresConfig> {
PostgresConfig::from_env()
}

/// Represents the proving data of a batch.
pub struct BatchData {
/// The number of the batch.
Expand Down
Loading