From 572ad40175784b62d61aa46dcf901a7e2f6ff8f5 Mon Sep 17 00:00:00 2001
From: Igor Aleksanov <popzxc@yandex.ru>
Date: Fri, 26 Jul 2024 08:09:44 +0400
Subject: [PATCH] feat(witness_vector_generator): Make it possible to run
 multiple wvg instances in one binary (#2493)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## What ❔

WVG is single-threaded, so running one "job" per binary is a waste (kind
of).
Also, it's not very convenient when developing locally.

This PR adds a `--threads` CLI parameter which can be used to manipulate
how much WVG jobs will be launched.

## Why ❔

Efficiency & convenience

## 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 `zk fmt` and `zk lint`.
---
 .../bin/witness_vector_generator/src/main.rs  | 45 ++++++++++++-------
 1 file changed, 29 insertions(+), 16 deletions(-)

diff --git a/prover/crates/bin/witness_vector_generator/src/main.rs b/prover/crates/bin/witness_vector_generator/src/main.rs
index 58db6d6d5eb..4451788ca9a 100644
--- a/prover/crates/bin/witness_vector_generator/src/main.rs
+++ b/prover/crates/bin/witness_vector_generator/src/main.rs
@@ -32,6 +32,10 @@ struct Cli {
     pub(crate) config_path: Option<std::path::PathBuf>,
     #[arg(long)]
     pub(crate) secrets_path: Option<std::path::PathBuf>,
+    /// Number of WVG jobs to run in parallel.
+    /// Default value is 1.
+    #[arg(long, default_value_t = 1)]
+    pub(crate) threads: usize,
 }
 
 #[tokio::main]
@@ -106,17 +110,6 @@ async fn main() -> anyhow::Result<()> {
 
     let protocol_version = PROVER_PROTOCOL_SEMANTIC_VERSION;
 
-    let witness_vector_generator = WitnessVectorGenerator::new(
-        object_store,
-        pool,
-        circuit_ids_for_round_to_be_proven.clone(),
-        zone.clone(),
-        config,
-        protocol_version,
-        prover_config.max_attempts,
-        Some(prover_config.setup_data_path.clone()),
-    );
-
     let (stop_sender, stop_receiver) = watch::channel(false);
 
     let (stop_signal_sender, stop_signal_receiver) = oneshot::channel();
@@ -128,12 +121,32 @@ async fn main() -> anyhow::Result<()> {
     })
     .expect("Error setting Ctrl+C handler");
 
-    tracing::info!("Starting witness vector generation for group: {} with circuits: {:?} in zone: {} with protocol_version: {:?}", specialized_group_id, circuit_ids_for_round_to_be_proven, zone, protocol_version);
+    tracing::info!(
+        "Starting {} witness vector generation jobs for group: {} with circuits: {:?} in zone: {} with protocol_version: {:?}",
+        opt.threads,
+        specialized_group_id,
+        circuit_ids_for_round_to_be_proven,
+        zone,
+        protocol_version
+    );
 
-    let tasks = vec![
-        tokio::spawn(exporter_config.run(stop_receiver.clone())),
-        tokio::spawn(witness_vector_generator.run(stop_receiver, opt.n_iterations)),
-    ];
+    let mut tasks = vec![tokio::spawn(exporter_config.run(stop_receiver.clone()))];
+
+    for _ in 0..opt.threads {
+        let witness_vector_generator = WitnessVectorGenerator::new(
+            object_store.clone(),
+            pool.clone(),
+            circuit_ids_for_round_to_be_proven.clone(),
+            zone.clone(),
+            config.clone(),
+            protocol_version,
+            prover_config.max_attempts,
+            Some(prover_config.setup_data_path.clone()),
+        );
+        tasks.push(tokio::spawn(
+            witness_vector_generator.run(stop_receiver.clone(), opt.n_iterations),
+        ));
+    }
 
     let mut tasks = ManagedTasks::new(tasks);
     tokio::select! {