Skip to content

Commit

Permalink
Add prover manager
Browse files Browse the repository at this point in the history
  • Loading branch information
bkolad committed Dec 19, 2023
1 parent 6fc0edf commit 5cc3219
Show file tree
Hide file tree
Showing 12 changed files with 304 additions and 112 deletions.
3 changes: 3 additions & 0 deletions examples/demo-rollup/celestia_rollup_config.toml
Original file line number Diff line number Diff line change
Expand Up @@ -21,3 +21,6 @@ start_height = 3
# the host and port to bind the rpc server for
bind_host = "127.0.0.1"
bind_port = 12345

[prover_service]
aggregated_proof_block_jump = 1
5 changes: 4 additions & 1 deletion examples/demo-rollup/mock_rollup_config.toml
Original file line number Diff line number Diff line change
Expand Up @@ -13,4 +13,7 @@ start_height = 0
[runner.rpc_config]
# the host and port to bind the rpc server for
bind_host = "127.0.0.1"
bind_port = 12345
bind_port = 12345

[prover_service]
aggregated_proof_block_jump = 1
6 changes: 5 additions & 1 deletion examples/demo-rollup/src/celestia_rollup.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,9 @@ use sov_risc0_adapter::host::Risc0Host;
use sov_rollup_interface::zk::ZkvmHost;
use sov_state::storage_manager::ProverStorageManager;
use sov_state::{DefaultStorageSpec, Storage, ZkStorage};
use sov_stf_runner::{ParallelProverService, RollupConfig, RollupProverConfig};
use sov_stf_runner::{
ParallelProverService, ProverServiceConfig, RollupConfig, RollupProverConfig,
};

use crate::{ROLLUP_BATCH_NAMESPACE, ROLLUP_PROOF_NAMESPACE};

Expand Down Expand Up @@ -101,6 +103,7 @@ impl RollupBlueprint for CelestiaDemoRollup {
async fn create_prover_service(
&self,
prover_config: RollupProverConfig,
rollup_config: &RollupConfig<Self::DaConfig>,
_da_service: &Self::DaService,
) -> Self::ProverService {
let vm = Risc0Host::new(risc0::ROLLUP_ELF);
Expand All @@ -117,6 +120,7 @@ impl RollupBlueprint for CelestiaDemoRollup {
da_verifier,
prover_config,
zk_storage,
rollup_config.prover_service,
)
}
}
Expand Down
6 changes: 5 additions & 1 deletion examples/demo-rollup/src/mock_rollup.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,9 @@ use sov_risc0_adapter::host::Risc0Host;
use sov_rollup_interface::zk::ZkvmHost;
use sov_state::storage_manager::ProverStorageManager;
use sov_state::{DefaultStorageSpec, Storage, ZkStorage};
use sov_stf_runner::{ParallelProverService, RollupConfig, RollupProverConfig};
use sov_stf_runner::{
ParallelProverService, ProverServiceConfig, RollupConfig, RollupProverConfig,
};

/// Rollup with MockDa
pub struct MockDemoRollup {}
Expand Down Expand Up @@ -92,6 +94,7 @@ impl RollupBlueprint for MockDemoRollup {
async fn create_prover_service(
&self,
prover_config: RollupProverConfig,
rollup_config: &RollupConfig<Self::DaConfig>,
_da_service: &Self::DaService,
) -> Self::ProverService {
let vm = Risc0Host::new(risc0::MOCK_DA_ELF);
Expand All @@ -105,6 +108,7 @@ impl RollupBlueprint for MockDemoRollup {
da_verifier,
prover_config,
zk_storage,
rollup_config.prover_service,
)
}
}
7 changes: 6 additions & 1 deletion examples/demo-rollup/tests/test_helpers.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,9 @@ use demo_stf::genesis_config::GenesisPaths;
use sov_demo_rollup::MockDemoRollup;
use sov_mock_da::{MockAddress, MockDaConfig};
use sov_modules_rollup_blueprint::RollupBlueprint;
use sov_stf_runner::{RollupConfig, RollupProverConfig, RpcConfig, RunnerConfig, StorageConfig};
use sov_stf_runner::{
ProverServiceConfig, RollupConfig, RollupProverConfig, RpcConfig, RunnerConfig, StorageConfig,
};
use tokio::sync::oneshot;

pub async fn start_rollup(
Expand All @@ -29,6 +31,9 @@ pub async fn start_rollup(
da: MockDaConfig {
sender_address: MockAddress::from([0; 32]),
},
prover_service: ProverServiceConfig {
aggregated_proof_block_jump: 1,
},
};

let mock_demo_rollup = MockDemoRollup {};
Expand Down
14 changes: 14 additions & 0 deletions full-node/sov-stf-runner/src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,13 @@ pub struct StorageConfig {
pub path: PathBuf,
}

///TODO
#[derive(Debug, Clone, PartialEq, Deserialize, Copy)]
pub struct ProverServiceConfig {
///TODO
pub aggregated_proof_block_jump: u64,
}

/// Rollup Configuration
#[derive(Debug, Clone, PartialEq, Deserialize)]
pub struct RollupConfig<DaServiceConfig> {
Expand All @@ -39,6 +46,8 @@ pub struct RollupConfig<DaServiceConfig> {
pub runner: RunnerConfig,
/// Data Availability service configuration.
pub da: DaServiceConfig,
/// Prover service configuration.
pub prover_service: ProverServiceConfig,
}

/// Reads toml file as a specific type.
Expand Down Expand Up @@ -85,6 +94,8 @@ mod tests {
[runner.rpc_config]
bind_host = "127.0.0.1"
bind_port = 12345
[prover_service]
aggregated_proof_block_jump = 22
"#;

let config_file = create_config_from(config);
Expand All @@ -109,6 +120,9 @@ mod tests {
storage: StorageConfig {
path: PathBuf::from("/tmp"),
},
prover_service: ProverServiceConfig {
aggregated_proof_block_jump: 22,
},
};
assert_eq!(config, expected);
}
Expand Down
2 changes: 1 addition & 1 deletion full-node/sov-stf-runner/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ pub use prover_service::*;
#[cfg(feature = "native")]
mod runner;
#[cfg(feature = "native")]
pub use config::{from_toml_path, RollupConfig, RunnerConfig, StorageConfig};
pub use config::{from_toml_path, ProverServiceConfig, RollupConfig, RunnerConfig, StorageConfig};
#[cfg(feature = "native")]
pub use runner::*;
use serde::de::DeserializeOwned;
Expand Down
24 changes: 20 additions & 4 deletions full-node/sov-stf-runner/src/prover_service/parallel/mod.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
mod prover;
mod prover_manager;
use std::sync::Arc;

use async_trait::async_trait;
Expand All @@ -13,8 +14,8 @@ use sov_rollup_interface::zk::ZkvmHost;
use super::{ProverService, ProverServiceError};
use crate::verifier::StateTransitionVerifier;
use crate::{
ProofGenConfig, ProofProcessingStatus, ProofSubmissionStatus, RollupProverConfig,
StateTransitionData, WitnessSubmissionStatus,
ProofGenConfig, ProofProcessingStatus, ProofSubmissionStatus, ProverServiceConfig,
RollupProverConfig, StateTransitionData, WitnessSubmissionStatus,
};

/// Prover service that generates proofs in parallel.
Expand All @@ -28,6 +29,7 @@ where
{
vm: Vm,
prover_config: Arc<ProofGenConfig<V, Da, Vm>>,
prover_service_config: ProverServiceConfig,
zk_storage: V::PreState,
prover_state: Prover<StateRoot, Witness, Da>,
}
Expand All @@ -49,6 +51,7 @@ where
config: RollupProverConfig,
zk_storage: V::PreState,
num_threads: usize,
prover_service_config: ProverServiceConfig,
) -> Self {
let stf_verifier =
StateTransitionVerifier::<V, Da::Verifier, Vm::Guest>::new(zk_stf, da_verifier);
Expand All @@ -65,7 +68,11 @@ where
Self {
vm,
prover_config,
prover_state: Prover::new(num_threads),
prover_state: Prover::new(
num_threads,
prover_service_config.aggregated_proof_block_jump,
),
prover_service_config,
zk_storage,
}
}
Expand All @@ -77,11 +84,20 @@ where
da_verifier: Da::Verifier,
config: RollupProverConfig,
zk_storage: V::PreState,
prover_service_config: ProverServiceConfig,
) -> Self {
let num_cpus = num_cpus::get();
assert!(num_cpus > 1, "Unable to create parallel prover service");

Self::new(vm, zk_stf, da_verifier, config, zk_storage, num_cpus - 1)
Self::new(
vm,
zk_stf,
da_verifier,
config,
zk_storage,
num_cpus - 1,
prover_service_config,
)
}
}

Expand Down
Loading

0 comments on commit 5cc3219

Please sign in to comment.