Merge pull request #1830 from matter-labs/prover_cli_status_command_lwg

feat(Prover CLI): status command lwg

core/lib/basic_types/Cargo.toml

@@ -20,3 +20,4 @@ serde_json.workspace = true
 chrono.workspace = true
 strum = { workspace = true, features = ["derive"] }
 num_enum.workspace = true
+sqlx = { workspace = true, feature= ["derive"]}

core/lib/basic_types/src/prover_dal.rs

@@ -124,7 +124,7 @@ pub struct ProverJobStatusInProgress {
     pub started_at: DateTime<Utc>,
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct WitnessJobStatusSuccessful {
     pub started_at: DateTime<Utc>,
     pub time_taken: Duration,
@@ -139,7 +139,7 @@ impl Default for WitnessJobStatusSuccessful {
     }
 }
 
-#[derive(Debug, Default)]
+#[derive(Debug, Default, Clone)]
 pub struct WitnessJobStatusFailed {
     pub started_at: DateTime<Utc>,
     pub error: String,
@@ -161,7 +161,7 @@ pub enum ProverJobStatus {
     Ignored,
 }
 
-#[derive(Debug, strum::Display, strum::EnumString, strum::AsRefStr)]
+#[derive(Debug, Clone, strum::Display, strum::EnumString, strum::AsRefStr)]
 pub enum WitnessJobStatus {
     #[strum(serialize = "failed")]
     Failed(WitnessJobStatusFailed),
@@ -273,6 +273,56 @@ pub struct BasicWitnessGeneratorJobInfo {
     pub eip_4844_blobs: Option<Eip4844Blobs>,
 }
 
+pub struct LeafWitnessGeneratorJobInfo {
+    pub id: u32,
+    pub l1_batch_number: L1BatchNumber,
+    pub circuit_id: u32,
+    pub closed_form_inputs_blob_url: Option<String>,
+    pub attempts: u32,
+    pub status: WitnessJobStatus,
+    pub error: Option<String>,
+    pub created_at: NaiveDateTime,
+    pub updated_at: NaiveDateTime,
+    pub processing_started_at: Option<NaiveDateTime>,
+    pub time_taken: Option<NaiveTime>,
+    pub is_blob_cleaned: Option<bool>,
+    pub number_of_basic_circuits: Option<i32>,
+    pub protocol_version: Option<i32>,
+    pub picked_by: Option<String>,
+}
+
+pub struct NodeWitnessGeneratorJobInfo {
+    pub id: u32,
+    pub l1_batch_number: L1BatchNumber,
+    pub circuit_id: u32,
+    pub depth: u32,
+    pub status: WitnessJobStatus,
+    pub attempts: u32,
+    pub aggregations_url: Option<String>,
+    pub processing_started_at: Option<NaiveDateTime>,
+    pub time_taken: Option<NaiveTime>,
+    pub error: Option<String>,
+    pub created_at: NaiveDateTime,
+    pub updated_at: NaiveDateTime,
+    pub number_of_dependent_jobs: Option<i32>,
+    pub protocol_version: Option<i32>,
+    pub picked_by: Option<String>,
+}
+
+pub struct SchedulerWitnessGeneratorJobInfo {
+    pub l1_batch_number: L1BatchNumber,
+    pub scheduler_partial_input_blob_url: String,
+    pub status: WitnessJobStatus,
+    pub processing_started_at: Option<NaiveDateTime>,
+    pub time_taken: Option<NaiveTime>,
+    pub error: Option<String>,
+    pub created_at: NaiveDateTime,
+    pub updated_at: NaiveDateTime,
+    pub attempts: u32,
+    pub protocol_version: Option<i32>,
+    pub picked_by: Option<String>,
+}
+
 #[derive(Debug, EnumString, Display)]
 pub enum ProofCompressionJobStatus {
     #[strum(serialize = "queued")]

prover/prover_cli/src/commands/status/batch.rs

@@ -1,7 +1,9 @@
 use anyhow::{ensure, Context as _};
 use clap::Args as ClapArgs;
 use prover_dal::{Connection, ConnectionPool, Prover, ProverDal};
-use zksync_types::{basic_fri_types::AggregationRound, L1BatchNumber};
+use zksync_types::{
+    basic_fri_types::AggregationRound, prover_dal::WitnessJobStatus, L1BatchNumber,
+};
 
 use super::utils::{AggregationRoundInfo, BatchData, Task, TaskStatus};
 use crate::commands::status::utils::postgres_config;
@@ -38,21 +40,52 @@ async fn get_batches_data(batches: Vec<L1BatchNumber>) -> anyhow::Result<Vec<Bat
             .await
             .context("failed to build a prover_connection_pool")?;
 
-    let mut conn = prover_connection_pool.connection().await.unwrap();
+    let mut conn = prover_connection_pool
+        .connection()
+        .await
+        .context("failed to get a connection")?;
 
     let mut batches_data = Vec::new();
     for batch in batches {
         let current_batch_data = BatchData {
             batch_number: batch,
             basic_witness_generator: Task::BasicWitnessGenerator {
                 status: get_proof_basic_witness_generator_status_for_batch(batch, &mut conn).await,
-                aggregation_round_info: get_prover_jobs_data_for_batch(
+                aggregation_round_info: get_aggregation_round_info_for_batch(
                     batch,
                     AggregationRound::BasicCircuits,
                     &mut conn,
                 )
                 .await,
             },
+            leaf_witness_generator: Task::LeafWitnessGenerator {
+                status: get_proof_leaf_witness_generator_status_for_batch(batch, &mut conn).await,
+                aggregation_round_info: get_aggregation_round_info_for_batch(
+                    batch,
+                    AggregationRound::LeafAggregation,
+                    &mut conn,
+                )
+                .await,
+            },
+            node_witness_generator: Task::NodeWitnessGenerator {
+                status: get_proof_node_witness_generator_status_for_batch(batch, &mut conn).await,
+                aggregation_round_info: get_aggregation_round_info_for_batch(
+                    batch,
+                    AggregationRound::NodeAggregation,
+                    &mut conn,
+                )
+                .await,
+            },
+            scheduler_witness_generator: Task::SchedulerWitnessGenerator {
+                status: get_proof_scheduler_witness_generator_status_for_batch(batch, &mut conn)
+                    .await,
+                aggregation_round_info: get_aggregation_round_info_for_batch(
+                    batch,
+                    AggregationRound::Scheduler,
+                    &mut conn,
+                )
+                .await,
+            },
             compressor: Task::Compressor(
                 get_proof_compression_job_status_for_batch(batch, &mut conn).await,
             ),
@@ -64,7 +97,7 @@ async fn get_batches_data(batches: Vec<L1BatchNumber>) -> anyhow::Result<Vec<Bat
     Ok(batches_data)
 }
 
-async fn get_prover_jobs_data_for_batch<'a>(
+async fn get_aggregation_round_info_for_batch<'a>(
     batch_number: L1BatchNumber,
     aggregation_round: AggregationRound,
     conn: &mut Connection<'a, Prover>,
@@ -76,7 +109,7 @@ async fn get_prover_jobs_data_for_batch<'a>(
         .into();
 
     AggregationRoundInfo {
-        round: AggregationRound::BasicCircuits,
+        round: aggregation_round,
         prover_jobs_status: status,
     }
 }
@@ -92,6 +125,45 @@ async fn get_proof_basic_witness_generator_status_for_batch<'a>(
         .unwrap_or_default()
 }
 
+async fn get_proof_leaf_witness_generator_status_for_batch<'a>(
+    batch_number: L1BatchNumber,
+    conn: &mut Connection<'a, Prover>,
+) -> TaskStatus {
+    conn.fri_witness_generator_dal()
+        .get_leaf_witness_generator_jobs_for_batch(batch_number)
+        .await
+        .iter()
+        .map(|s| s.status.clone())
+        .collect::<Vec<WitnessJobStatus>>()
+        .into()
+}
+
+async fn get_proof_node_witness_generator_status_for_batch<'a>(
+    batch_number: L1BatchNumber,
+    conn: &mut Connection<'a, Prover>,
+) -> TaskStatus {
+    conn.fri_witness_generator_dal()
+        .get_node_witness_generator_jobs_for_batch(batch_number)
+        .await
+        .iter()
+        .map(|s| s.status.clone())
+        .collect::<Vec<WitnessJobStatus>>()
+        .into()
+}
+
+async fn get_proof_scheduler_witness_generator_status_for_batch<'a>(
+    batch_number: L1BatchNumber,
+    conn: &mut Connection<'a, Prover>,
+) -> TaskStatus {
+    conn.fri_witness_generator_dal()
+        .get_scheduler_witness_generator_jobs_for_batch(batch_number)
+        .await
+        .iter()
+        .map(|s| s.status.clone())
+        .collect::<Vec<WitnessJobStatus>>()
+        .into()
+}
+
 async fn get_proof_compression_job_status_for_batch<'a>(
     batch_number: L1BatchNumber,
     conn: &mut Connection<'a, Prover>,

prover/prover_cli/src/commands/status/utils.rs

@@ -27,7 +27,7 @@ pub struct BatchData {
     /// The recursion tip data.
     pub recursion_tip: Task,
     /// The scheduler data.
-    pub scheduler: Task,
+    pub scheduler_witness_generator: Task,
     /// The compressor data.
     pub compressor: Task,
 }
@@ -45,7 +45,7 @@ impl Debug for BatchData {
         writeln!(f, "{:?}", self.leaf_witness_generator)?;
         writeln!(f, "{:?}", self.node_witness_generator)?;
         writeln!(f, "{:?}", self.recursion_tip)?;
-        writeln!(f, "{:?}", self.scheduler)?;
+        writeln!(f, "{:?}", self.scheduler_witness_generator)?;
         writeln!(f, "{:?}", self.compressor)
     }
 }
@@ -82,7 +82,7 @@ impl Default for BatchData {
                     prover_jobs_status: TaskStatus::default(),
                 },
             },
-            scheduler: Task::Scheduler {
+            scheduler_witness_generator: Task::SchedulerWitnessGenerator {
                 status: TaskStatus::default(),
                 aggregation_round_info: AggregationRoundInfo {
                     round: AggregationRound::Scheduler,
@@ -157,6 +157,31 @@ impl From<ProofCompressionJobStatus> for TaskStatus {
     }
 }
 
+impl From<Vec<WitnessJobStatus>> for TaskStatus {
+    fn from(status_vector: Vec<WitnessJobStatus>) -> Self {
+        if status_vector.is_empty() {
+            TaskStatus::JobsNotFound
+        } else if status_vector
+            .iter()
+            .all(|job| matches!(job, WitnessJobStatus::Queued))
+        {
+            TaskStatus::Queued
+        } else if status_vector
+            .iter()
+            .all(|job| matches!(job, WitnessJobStatus::WaitingForProofs))
+        {
+            TaskStatus::WaitingForProofs
+        } else if status_vector
+            .iter()
+            .all(|job| matches!(job, WitnessJobStatus::InProgress))
+        {
+            TaskStatus::Successful
+        } else {
+            TaskStatus::InProgress
+        }
+    }
+}
+
 #[derive(EnumString, Clone, Display)]
 pub enum Task {
     /// Represents the basic witness generator task and its status.
@@ -185,7 +210,7 @@ pub enum Task {
     },
     /// Represents the scheduler task and its status.
     #[strum(to_string = "Scheduler")]
-    Scheduler {
+    SchedulerWitnessGenerator {
         status: TaskStatus,
         aggregation_round_info: AggregationRoundInfo,
     },
@@ -201,7 +226,7 @@ impl Task {
             | Task::LeafWitnessGenerator { status, .. }
             | Task::NodeWitnessGenerator { status, .. }
             | Task::RecursionTip { status, .. }
-            | Task::Scheduler { status, .. }
+            | Task::SchedulerWitnessGenerator { status, .. }
             | Task::Compressor(status) => status.clone(),
         }
     }
@@ -224,7 +249,7 @@ impl Task {
                 aggregation_round_info,
                 ..
             }
-            | Task::Scheduler {
+            | Task::SchedulerWitnessGenerator {
                 aggregation_round_info,
                 ..
             } => Some(aggregation_round_info.round),
@@ -254,7 +279,7 @@ impl Task {
                 status,
                 aggregation_round_info,
             }
-            | Task::Scheduler {
+            | Task::SchedulerWitnessGenerator {
                 status,
                 aggregation_round_info,
             } => match status {