feat(prover): Witness generator service (#4219)

## What ❔
The Witness generator Service is a lib for running Witness generator.
The primitive exported by this lib is witness_generator_runner which is
based on
`ProverJobProcessor`.

<!-- What are the changes this PR brings about? -->
<!-- Example: This PR adds a PR template to the repo. -->
<!-- (For bigger PRs adding more context is appreciated) -->

## Why ❔
The lib makes it easier for external teams to implement their own
Artifacts Manager, Keystore and DB to run Witness Generator.

<!-- Why are these changes done? What goal do they contribute to? What
are the principles behind them? -->
<!-- The `Why` has to be clear to non-Matter Labs entities running their
own ZK Chain -->
<!-- Example: PR templates ensure PR reviewers, observers, and future
iterators are in context about the evolution of repos. -->

## Is this a breaking change?
- [ ] Yes
- [ ] No

## Operational changes
<!-- Any config changes? Any new flags? Any changes to any scripts? -->
<!-- Please add anything that non-Matter Labs entities running their own
ZK Chain may need to know -->

## Checklist

<!-- Check your PR fulfills the following items. -->
<!-- For draft PRs check the boxes as you complete them. -->

- [ ] PR title corresponds to the body of PR (we generate changelog
entries from PRs).
- [ ] Tests for the changes have been added / updated.
- [ ] Documentation comments have been added / updated.
- [ ] Code has been formatted via `zkstack dev fmt` and `zkstack dev
lint`.

Author: AlexOnTheStorm

prover/Cargo.lock

@@ -113,6 +113,7 @@ zksync_prover_job_processor = { version = "=21.2.0", path = "crates/lib/prover_j
 zksync_circuit_prover_service = { version = "=21.2.0", path = "crates/lib/circuit_prover_service" }
 zksync_prover_job_monitor = { version = "=21.2.0", path = "crates/bin/prover_job_monitor" }
 zksync_proof_fri_compressor_service = { version = "=21.2.0", path = "crates/lib/proof_fri_compressor_service" }
+zksync_witness_generator_service = { version = "=21.2.0", path = "crates/lib/witness_generator_service" }
 
 # for `perf` profiling
 [profile.perf]

prover/Cargo.toml

@@ -1,10 +1,10 @@
 use std::{collections::HashMap, sync::Arc};
 
+use zksync_circuit_prover_service::types::setup_data::GoldilocksGpuProverSetupData;
 use zksync_prover_fri_types::{
     circuit_definitions::boojum::cs::implementations::setup::FinalizationHintsForProver,
     ProverServiceDataKey,
 };
-use zksync_prover_keystore::GoldilocksGpuProverSetupData;
 
 // TODO: To be moved to circuit_prover_service lib & adjusted to new type idiom
 // cache types

prover/crates/bin/circuit_prover/src/types.rs

@@ -14,7 +14,7 @@ use tokio::fs;
 use zksync_circuit_prover_service::{
     gpu_circuit_prover::GpuCircuitProverExecutor,
     types::{
-        circuit_prover_payload::GpuCircuitProverPayload,
+        circuit_prover_payload::GpuCircuitProverPayload, setup_data::GoldilocksGpuProverSetupData,
         witness_vector_generator_payload::WitnessVectorGeneratorPayload,
     },
     witness_vector_generator::WitnessVectorGeneratorExecutor,
@@ -34,10 +34,7 @@ use zksync_prover_fri_types::{
     ProverServiceDataKey,
 };
 use zksync_prover_job_processor::Executor;
-use zksync_prover_keystore::{
-    keystore::{Keystore, ProverServiceDataType},
-    GoldilocksGpuProverSetupData,
-};
+use zksync_prover_keystore::keystore::{Keystore, ProverServiceDataType};
 use zksync_types::{
     basic_fri_types::AggregationRound, prover_dal::FriProverJobMetadata, L1BatchId, L1BatchNumber,
     L2ChainId,

prover/crates/bin/gpu_checker/src/main.rs

@@ -25,13 +25,15 @@ zksync_prover_keystore.workspace = true
 zksync_prover_fri_types.workspace = true
 zksync_prover_fri_utils.workspace = true
 zksync_circuit_prover_service.workspace = true
+zksync_witness_generator_service.workspace = true
 
 zkevm_test_harness = { workspace = true }
 circuit_definitions = { workspace = true, features = ["log_tracing"] }
 
 anyhow.workspace = true
 tracing.workspace = true
 tokio = { workspace = true, features = ["time", "macros"] }
+tokio-util.workspace = true
 futures = { workspace = true, features = ["compat"] }
 serde = { workspace = true, features = ["derive"] }
 async-trait.workspace = true

prover/crates/bin/witness_generator/Cargo.toml

@@ -1,12 +1,4 @@
 #![allow(incomplete_features)] // We have to use generic const exprs.
 #![feature(generic_const_exprs)]
 
-pub mod artifacts;
 pub mod metrics;
-pub mod precalculated_merkle_paths_provider;
-pub mod rounds;
-mod storage_oracle;
-#[cfg(test)]
-mod tests;
-pub mod utils;
-mod witness;

prover/crates/bin/witness_generator/src/lib.rs

@@ -1,14 +1,17 @@
 #![allow(incomplete_features)] // We have to use generic const exprs.
 #![feature(generic_const_exprs)]
 
-use std::time::{Duration, Instant};
+use std::{
+    sync::Arc,
+    time::{Duration, Instant},
+};
 
 use anyhow::{anyhow, Context as _};
-use futures::{channel::mpsc, executor::block_on, SinkExt, StreamExt};
 #[cfg(not(target_env = "msvc"))]
 use jemallocator::Jemalloc;
 use structopt::StructOpt;
-use tokio::sync::watch;
+use tokio::sync::oneshot;
+use tokio_util::sync::CancellationToken;
 use zksync_config::{
     configs::{GeneralConfig, PostgresSecrets},
     full_config_schema,
@@ -18,16 +21,16 @@ use zksync_object_store::ObjectStoreFactory;
 use zksync_prover_dal::{ConnectionPool, Prover, ProverDal};
 use zksync_prover_fri_types::PROVER_PROTOCOL_SEMANTIC_VERSION;
 use zksync_prover_keystore::keystore::Keystore;
-use zksync_queued_job_processor::JobProcessor;
 use zksync_task_management::ManagedTasks;
 use zksync_types::{basic_fri_types::AggregationRound, protocol_version::ProtocolSemanticVersion};
-use zksync_witness_generator::{
-    metrics::SERVER_METRICS,
-    rounds::{
-        BasicCircuits, LeafAggregation, NodeAggregation, RecursionTip, Scheduler, WitnessGenerator,
-    },
+use zksync_witness_generator::metrics::SERVER_METRICS;
+use zksync_witness_generator_service::{
+    rounds::{BasicCircuits, LeafAggregation, NodeAggregation, RecursionTip, Scheduler},
+    witness_generator_runner,
 };
 
+const GRACEFUL_SHUTDOWN_DURATION: Duration = Duration::from_secs(20);
+
 #[cfg(not(target_env = "msvc"))]
 #[global_allocator]
 static GLOBAL: Jemalloc = Jemalloc;
@@ -94,7 +97,42 @@ async fn ensure_protocol_alignment(
 
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
-    let started_at = Instant::now();
+    let (stop_signal_sender, stop_signal_receiver) = oneshot::channel();
+    let mut stop_signal_sender = Some(stop_signal_sender);
+    ctrlc::set_handler(move || {
+        if let Some(sender) = stop_signal_sender.take() {
+            sender.send(()).ok();
+        }
+    })
+    .context("Error setting Ctrl+C handler")?;
+
+    let cancellation_token = CancellationToken::new();
+    let mut managed_tasks = ManagedTasks::new(vec![]);
+    let (metrics_stop_sender, metrics_stop_receiver) = tokio::sync::watch::channel(false);
+
+    tokio::select! {
+        res = run_inner(cancellation_token.clone(), metrics_stop_receiver, &mut managed_tasks) => {
+            res.context("Failed to run witness generator")?;
+        },
+        _ = stop_signal_receiver => {
+            tracing::info!("Stop request received, shutting down");
+        }
+    }
+    let shutdown_time = Instant::now();
+    cancellation_token.cancel();
+    metrics_stop_sender
+        .send(true)
+        .context("failed to stop metrics")?;
+    managed_tasks.complete(GRACEFUL_SHUTDOWN_DURATION).await;
+    tracing::info!("Tasks completed in {:?}.", shutdown_time.elapsed());
+    Ok(())
+}
+
+async fn run_inner(
+    cancellation_token: CancellationToken,
+    metrics_stop_receiver: tokio::sync::watch::Receiver<bool>,
+    managed_tasks: &mut ManagedTasks,
+) -> anyhow::Result<()> {
     let opt = Opt::from_args();
     let schema = full_config_schema();
     let config_file_paths = ConfigFilePaths {
@@ -109,6 +147,7 @@ async fn main() -> anyhow::Result<()> {
     let mut repo = config_sources.build_repository(&schema);
     let general_config: GeneralConfig = repo.parse()?;
     let database_secrets: PostgresSecrets = repo.parse()?;
+    let started_at = Instant::now();
 
     let prover_config = general_config.prover_config.context("prover config")?;
     let object_store_config = prover_config.prover_object_store;
@@ -129,7 +168,6 @@ async fn main() -> anyhow::Result<()> {
         .build()
         .await
         .context("failed to build a prover_connection_pool")?;
-    let (stop_sender, stop_receiver) = watch::channel(false);
 
     let protocol_version = PROVER_PROTOCOL_SEMANTIC_VERSION;
 
@@ -159,9 +197,11 @@ async fn main() -> anyhow::Result<()> {
     };
 
     let mut tasks = vec![tokio::spawn(
-        prometheus_exporter_config.run(stop_receiver.clone()),
+        prometheus_exporter_config.run(metrics_stop_receiver),
     )];
 
+    let keystore = Arc::new(keystore);
+
     for round in rounds {
         tracing::info!(
             "initializing the {:?} witness generator, batch size: {:?} with protocol_version: {:?}",
@@ -170,60 +210,65 @@ async fn main() -> anyhow::Result<()> {
             &protocol_version
         );
 
-        let witness_generator_task = match round {
+        let witness_generator_tasks = match round {
             AggregationRound::BasicCircuits => {
-                let generator = WitnessGenerator::<BasicCircuits>::new(
-                    config.clone(),
+                let runner = witness_generator_runner::<BasicCircuits>(
+                    config.max_circuits_in_flight,
                     store_factory.create_store().await?,
                     connection_pool.clone(),
                     protocol_version,
                     keystore.clone(),
+                    cancellation_token.clone(),
                 );
-                generator.run(stop_receiver.clone(), opt.batch_size)
+                runner.run()
             }
             AggregationRound::LeafAggregation => {
-                let generator = WitnessGenerator::<LeafAggregation>::new(
-                    config.clone(),
+                let runner = witness_generator_runner::<LeafAggregation>(
+                    config.max_circuits_in_flight,
                     store_factory.create_store().await?,
                     connection_pool.clone(),
                     protocol_version,
                     keystore.clone(),
+                    cancellation_token.clone(),
                 );
-                generator.run(stop_receiver.clone(), opt.batch_size)
+                runner.run()
             }
             AggregationRound::NodeAggregation => {
-                let generator = WitnessGenerator::<NodeAggregation>::new(
-                    config.clone(),
+                let runner = witness_generator_runner::<NodeAggregation>(
+                    config.max_circuits_in_flight,
                     store_factory.create_store().await?,
                     connection_pool.clone(),
                     protocol_version,
                     keystore.clone(),
+                    cancellation_token.clone(),
                 );
-                generator.run(stop_receiver.clone(), opt.batch_size)
+                runner.run()
             }
             AggregationRound::RecursionTip => {
-                let generator = WitnessGenerator::<RecursionTip>::new(
-                    config.clone(),
+                let runner = witness_generator_runner::<RecursionTip>(
+                    config.max_circuits_in_flight,
                     store_factory.create_store().await?,
                     connection_pool.clone(),
                     protocol_version,
                     keystore.clone(),
+                    cancellation_token.clone(),
                 );
-                generator.run(stop_receiver.clone(), opt.batch_size)
+                runner.run()
             }
             AggregationRound::Scheduler => {
-                let generator = WitnessGenerator::<Scheduler>::new(
-                    config.clone(),
+                let runner = witness_generator_runner::<Scheduler>(
+                    config.max_circuits_in_flight,
                     store_factory.create_store().await?,
                     connection_pool.clone(),
                     protocol_version,
                     keystore.clone(),
+                    cancellation_token.clone(),
                 );
-                generator.run(stop_receiver.clone(), opt.batch_size)
+                runner.run()
             }
         };
 
-        tasks.push(tokio::spawn(witness_generator_task));
+        tasks.extend(witness_generator_tasks);
 
         tracing::info!(
             "initialized {:?} witness generator in {:?}",
@@ -233,21 +278,7 @@ async fn main() -> anyhow::Result<()> {
         SERVER_METRICS.init_latency[&round.into()].set(started_at.elapsed());
     }
 
-    let (mut stop_signal_sender, mut stop_signal_receiver) = mpsc::channel(256);
-    ctrlc::set_handler(move || {
-        block_on(stop_signal_sender.send(true)).expect("Ctrl+C signal send");
-    })
-    .expect("Error setting Ctrl+C handler");
-    let mut tasks = ManagedTasks::new(tasks).allow_tasks_to_finish();
-    tokio::select! {
-        _ = tasks.wait_single() => {},
-        _ = stop_signal_receiver.next() => {
-            tracing::info!("Stop request received, shutting down");
-        }
-    }
-
-    stop_sender.send_replace(true);
-    tasks.complete(Duration::from_secs(5)).await;
-    tracing::info!("Finished witness generation");
+    *managed_tasks = ManagedTasks::new(tasks);
+    managed_tasks.wait_single().await;
     Ok(())
 }

prover/crates/bin/witness_generator/src/main.rs

@@ -1,25 +1,8 @@
 use std::time::Duration;
 
-use vise::{Buckets, Family, Gauge, Histogram, LabeledFamily, Metrics};
+use vise::{Gauge, LabeledFamily, Metrics};
 use zksync_prover_fri_utils::metrics::StageLabel;
 
-#[derive(Debug, Metrics)]
-#[metrics(prefix = "prover_fri_witness_generator")]
-pub(crate) struct WitnessGeneratorMetrics {
-    #[metrics(buckets = Buckets::LATENCIES)]
-    pub blob_fetch_time: Family<StageLabel, Histogram<Duration>>,
-    #[metrics(buckets = Buckets::LATENCIES)]
-    pub prepare_job_time: Family<StageLabel, Histogram<Duration>>,
-    #[metrics(buckets = Buckets::exponential(60.0..=61440.0, 2.0))]
-    pub witness_generation_time: Family<StageLabel, Histogram<Duration>>,
-    #[metrics(buckets = Buckets::LATENCIES)]
-    pub blob_save_time: Family<StageLabel, Histogram<Duration>>,
-}
-
-#[vise::register]
-pub(crate) static WITNESS_GENERATOR_METRICS: vise::Global<WitnessGeneratorMetrics> =
-    vise::Global::new();
-
 #[derive(Debug, Metrics)]
 #[metrics(prefix = "prover")]
 pub struct ServerMetrics {