Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

make pool size configurable #644

Merged
merged 2 commits into from
Sep 19, 2024
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 13 additions & 3 deletions zero/src/prover.rs
Original file line number Diff line number Diff line change
Expand Up @@ -31,9 +31,8 @@ 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);
const DEFAULT_PARALLEL_BLOCK_PROVING_PERMIT_POOL_SIZE: usize = 16;
static PARALLEL_BLOCK_PROVING_PERMIT_POOL: Semaphore = Semaphore::const_new(0);

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

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

if prover_config.block_pool_size > 0 {
PARALLEL_BLOCK_PROVING_PERMIT_POOL.add_permits(prover_config.block_pool_size);
} else {
anyhow::bail!(
"block_pool_size should be greater than 0, value passed from cli is {}",
temaniarpit27 marked this conversation as resolved.
Show resolved Hide resolved
prover_config.block_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.
temaniarpit27 marked this conversation as resolved.
Show resolved Hide resolved
#[arg(long, default_value_t = 16)]
block_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,
block_pool_size: cli.block_pool_size,
}
}
}
Loading