feat(Prover CLI): delete cmd

prover/Cargo.toml

@@ -37,6 +37,7 @@ clap = "4.4.6"
 colored = "2.0"
 const-decoder = "0.3.0"
 ctrlc = "3.1"
+dialoguer = "0.11"
 futures = "0.3"
 hex = "0.4"
 itertools = "0.10.5"

prover/prover_cli/Cargo.toml

@@ -10,6 +10,7 @@ keywords.workspace = true
 categories.workspace = true
 
 [dependencies]
+dialoguer.workspace = true
 tokio = { workspace = true, features = ["rt-multi-thread", "macros"] }
 clap = { workspace = true, features = ["derive"] }
 tracing-subscriber = { workspace = true, features = ["env-filter"] }

prover/prover_cli/src/cli.rs

@@ -1,7 +1,7 @@
 use clap::{command, Args, Parser, Subcommand};
 use zksync_types::url::SensitiveUrl;
 
-use crate::commands::{self, get_file_info};
+use crate::commands::{self, delete, get_file_info, restart};
 
 pub const VERSION_STRING: &str = env!("CARGO_PKG_VERSION");
 
@@ -29,15 +29,19 @@ pub struct ProverCLIConfig {
 #[derive(Subcommand)]
 enum ProverCommand {
     FileInfo(get_file_info::Args),
+    Delete(delete::Args),
     #[command(subcommand)]
     Status(commands::StatusCommand),
+    Restart(restart::Args),
 }
 
 pub async fn start() -> anyhow::Result<()> {
     let ProverCLI { command, config } = ProverCLI::parse();
     match command {
         ProverCommand::FileInfo(args) => get_file_info::run(args).await?,
+        ProverCommand::Delete(args) => delete::run(args).await?,
         ProverCommand::Status(cmd) => cmd.run(config).await?,
+        ProverCommand::Restart(args) => restart::run(args).await?,
     };
 
     Ok(())

prover/prover_cli/src/commands/delete.rs

@@ -0,0 +1,93 @@
+use anyhow::Context;
+use clap::Args as ClapArgs;
+use dialoguer::{theme::ColorfulTheme, Input};
+use prover_dal::{Connection, ConnectionPool, Prover, ProverDal};
+use zksync_config::PostgresConfig;
+use zksync_env_config::FromEnv;
+use zksync_types::L1BatchNumber;
+
+#[derive(ClapArgs)]
+pub(crate) struct Args {
+    /// Delete data from all batches
+    #[clap(
+        short,
+        long,
+        required_unless_present = "batch",
+        conflicts_with = "batch",
+        default_value_t = false
+    )]
+    all: bool,
+    /// Batch number to delete
+    #[clap(short, long, required_unless_present = "all", conflicts_with = "all", default_value_t = L1BatchNumber(0))]
+    batch: L1BatchNumber,
+}
+
+pub(crate) async fn run(args: Args) -> anyhow::Result<()> {
+    let confirmation = Input::<String>::with_theme(&ColorfulTheme::default())
+        .with_prompt("Are you sure you want to delete the data?")
+        .default("no".to_owned())
+        .interact_text()?;
+
+    if confirmation != "yes" {
+        println!("Aborted");
+        return Ok(());
+    }
+
+    let config = PostgresConfig::from_env()?;
+    let prover_connection_pool = ConnectionPool::<Prover>::singleton(config.prover_url()?)
+        .build()
+        .await
+        .context("failed to build a prover_connection_pool")?;
+    let conn = prover_connection_pool.connection().await.unwrap();
+
+    if args.all {
+        delete_prover_db(conn).await?;
+    } else {
+        delete_batch_data(conn, args.batch).await?;
+    }
+
+    Ok(())
+}
+
+async fn delete_prover_db(mut conn: Connection<'_, Prover>) -> anyhow::Result<()> {
+    conn.fri_gpu_prover_queue_dal()
+        .delete()
+        .await
+        .context("failed to delete gpu prover queue")?;
+    conn.fri_prover_jobs_dal()
+        .delete()
+        .await
+        .context("failed to delete prover jobs")?;
+    conn.fri_protocol_versions_dal()
+        .delete()
+        .await
+        .context("failed to delete protocol versions")?;
+    conn.fri_proof_compressor_dal()
+        .delete()
+        .await
+        .context("failed to delete proof compressor")?;
+    conn.fri_witness_generator_dal()
+        .delete()
+        .await
+        .context("failed to delete witness generator")?;
+    Ok(())
+}
+
+async fn delete_batch_data(
+    mut conn: Connection<'_, Prover>,
+    block_number: L1BatchNumber,
+) -> anyhow::Result<()> {
+    conn.fri_proof_compressor_dal()
+        .delete_batch_data(block_number)
+        .await
+        .context("failed to delete proof compressor data")?;
+    conn.fri_prover_jobs_dal()
+        .delete_batch_data(block_number)
+        .await
+        .context("failed to delete prover jobs data")?;
+    conn.fri_witness_generator_dal()
+        .delete_batch_data(block_number)
+        .await
+        .context("failed to delete witness generator data")?;
+    Ok(())
+}

prover/prover_cli/src/commands/mod.rs

@@ -1,4 +1,6 @@
+pub(crate) mod delete;
 pub(crate) mod get_file_info;
+pub(crate) mod restart;
 pub(crate) mod status;
 
 pub(crate) use status::StatusCommand;

prover/prover_cli/src/commands/restart.rs

@@ -0,0 +1,78 @@
+use anyhow::Context;
+use clap::Args as ClapArgs;
+use prover_dal::{
+    fri_witness_generator_dal::FriWitnessJobStatus, Connection, ConnectionPool, Prover, ProverDal,
+};
+use zksync_config::PostgresConfig;
+use zksync_env_config::FromEnv;
+use zksync_types::{basic_fri_types::AggregationRound, L1BatchNumber};
+
+#[derive(ClapArgs)]
+pub(crate) struct Args {
+    /// Batch number to restart
+    #[clap(
+        short,
+        long,
+        required_unless_present = "prover_job",
+        conflicts_with = "prover_job"
+    )]
+    batch: Option<L1BatchNumber>,
+    /// Prover job to restart
+    #[clap(short, long, required_unless_present = "batch")]
+    prover_job: Option<u32>,
+}
+
+pub(crate) async fn run(args: Args) -> anyhow::Result<()> {
+    let config = PostgresConfig::from_env()?;
+    let prover_connection_pool = ConnectionPool::<Prover>::singleton(config.prover_url()?)
+        .build()
+        .await
+        .context("failed to build a prover_connection_pool")?;
+    let mut conn = prover_connection_pool.connection().await.unwrap();
+
+    if let Some(batch_number) = args.batch {
+        restart_batch(batch_number, &mut conn).await?;
+    } else if let Some(id) = args.prover_job {
+        restart_prover_job(id, &mut conn).await;
+    }
+
+    Ok(())
+}
+
+async fn restart_batch(
+    batch_number: L1BatchNumber,
+    conn: &mut Connection<'_, Prover>,
+) -> anyhow::Result<()> {
+    conn.fri_proof_compressor_dal()
+        .delete_batch_data(batch_number)
+        .await
+        .context("failed to delete proof compression job for batch")?;
+    conn.fri_prover_jobs_dal()
+        .delete_batch_data(batch_number)
+        .await
+        .context("failed to delete prover jobs for batch")?;
+    conn.fri_witness_generator_dal()
+        .delete_witness_generator_data_for_batch(batch_number, AggregationRound::LeafAggregation)
+        .await
+        .context("failed to restart batch: fri_witness_generator_dal()")?;
+    conn.fri_witness_generator_dal()
+        .delete_witness_generator_data_for_batch(batch_number, AggregationRound::NodeAggregation)
+        .await
+        .context("failed to restart batch: fri_witness_generator_dal()")?;
+    conn.fri_witness_generator_dal()
+        .delete_witness_generator_data_for_batch(batch_number, AggregationRound::RecursionTip)
+        .await
+        .context("failed to restart batch: fri_witness_generator_dal()")?;
+    conn.fri_witness_generator_dal()
+        .delete_witness_generator_data_for_batch(batch_number, AggregationRound::Scheduler)
+        .await
+        .context("failed to restart batch: fri_witness_generator_dal()")?;
+    conn.fri_witness_generator_dal()
+        .mark_witness_job(FriWitnessJobStatus::Queued, batch_number)
+        .await;
+    Ok(())
+}
+
+async fn restart_prover_job(id: u32, conn: &mut Connection<'_, Prover>) {
+    conn.fri_prover_jobs_dal().update_status(id, "queued").await;
+}