Skip to content

Commit

Permalink
feat(witness_vector_generator): Make it possible to run multiple wvg …
Browse files Browse the repository at this point in the history
…instances in one binary (#2493)

## 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`.
  • Loading branch information
popzxc authored Jul 26, 2024
1 parent b12d01b commit 572ad40
Showing 1 changed file with 29 additions and 16 deletions.
45 changes: 29 additions & 16 deletions prover/crates/bin/witness_vector_generator/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down Expand Up @@ -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();
Expand All @@ -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! {
Expand Down

0 comments on commit 572ad40

Please sign in to comment.