Skip to content

Commit

Permalink
make pool size configurable
Browse files Browse the repository at this point in the history
  • Loading branch information
temaniarpit27 committed Sep 19, 2024
1 parent 6ef8827 commit 92294bc
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 3 deletions.
9 changes: 6 additions & 3 deletions zero/src/prover.rs
Original file line number Diff line number Diff line change
Expand Up @@ -31,9 +31,7 @@ use crate::ops;
//
// While proving a block interval, we will output proofs corresponding to block
// batches as soon as they are generated.
const PARALLEL_BLOCK_PROVING_PERMIT_POOL_SIZE: usize = 16;
static PARALLEL_BLOCK_PROVING_PERMIT_POOL: Semaphore =
Semaphore::const_new(PARALLEL_BLOCK_PROVING_PERMIT_POOL_SIZE);
static PARALLEL_BLOCK_PROVING_PERMIT_POOL: Semaphore = Semaphore::const_new(0);

#[derive(Debug, Clone)]
pub struct ProverConfig {
Expand All @@ -44,6 +42,7 @@ pub struct ProverConfig {
pub proof_output_dir: PathBuf,
pub keep_intermediate_proofs: bool,
pub block_batch_size: usize,
pub parallel_block_proving_permit_pool_size: usize,
}

#[derive(Clone, Debug, Deserialize, Serialize)]
Expand Down Expand Up @@ -242,6 +241,10 @@ pub async fn prove(
let mut task_set: JoinSet<
std::result::Result<std::result::Result<u64, anyhow::Error>, anyhow::Error>,
> = JoinSet::new();

PARALLEL_BLOCK_PROVING_PERMIT_POOL
.add_permits(prover_config.parallel_block_proving_permit_pool_size);

while let Some((block_prover_input, is_last_block)) = block_receiver.recv().await {
block_counter += 1;
let (tx, rx) = oneshot::channel::<GeneratedBlockProof>();
Expand Down
4 changes: 4 additions & 0 deletions zero/src/prover/cli.rs
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,9 @@ pub struct CliProverConfig {
/// generate one proof file.
#[arg(long, default_value_t = 8)]
block_batch_size: usize,
/// The maximum number of block proving tasks that can run in parallel.
#[arg(long, default_value_t = 16)]
parallel_block_proving_permit_pool_size: usize,
}

impl From<CliProverConfig> for super::ProverConfig {
Expand All @@ -55,6 +58,7 @@ impl From<CliProverConfig> for super::ProverConfig {
proof_output_dir: cli.proof_output_dir,
keep_intermediate_proofs: cli.keep_intermediate_proofs,
block_batch_size: cli.block_batch_size,
parallel_block_proving_permit_pool_size: cli.parallel_block_proving_permit_pool_size,
}
}
}

0 comments on commit 92294bc

Please sign in to comment.