multi-threaded prover service (#1160)

* Add ParallelProver * Implement ParallelProverService * Merge * Reduce da submission wait_time * Update full-node/sov-stf-runner/src/prover_service/parallel/prover.rs Co-authored-by: Cem Özer <[email protected]> * Add skip variant to RollupProverConfig * Fix lint * Fix build * lint * Add comment --------- Co-authored-by: Cem Özer <[email protected]>
Sovereign-Labs · Nov 20, 2023 · b29e9c2 · b29e9c2
1 parent ba8a25d
commit b29e9c2
Show file tree

Hide file tree

Showing 15 changed files with 355 additions and 162 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/examples/demo-rollup/src/celestia_rollup.rs b/examples/demo-rollup/src/celestia_rollup.rs
@@ -12,7 +12,7 @@ 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::{BlockingProver, RollupConfig, RollupProverConfig};
+use sov_stf_runner::{ParallelProverService, RollupConfig, RollupProverConfig};
 
 use crate::ROLLUP_NAMESPACE;
 
@@ -37,7 +37,7 @@ impl RollupBlueprint for CelestiaDemoRollup {
     type NativeKernel = BasicKernel<Self::NativeContext>;
     type ZkKernel = BasicKernel<Self::ZkContext>;
 
-    type ProverService = BlockingProver<
+    type ProverService = ParallelProverService<
         <<Self::NativeContext as Spec>::Storage as Storage>::Root,
         <<Self::NativeContext as Spec>::Storage as Storage>::Witness,
         Self::DaService,
@@ -99,7 +99,7 @@ impl RollupBlueprint for CelestiaDemoRollup {
 
     async fn create_prover_service(
         &self,
-        prover_config: Option<RollupProverConfig>,
+        prover_config: RollupProverConfig,
         _da_service: &Self::DaService,
     ) -> Self::ProverService {
         let vm = Risc0Host::new(risc0::ROLLUP_ELF);
@@ -110,7 +110,7 @@ impl RollupBlueprint for CelestiaDemoRollup {
             rollup_namespace: ROLLUP_NAMESPACE,
         };
 
-        BlockingProver::new(vm, zk_stf, da_verifier, prover_config, zk_storage)
+        ParallelProverService::new(vm, zk_stf, da_verifier, prover_config, zk_storage)
     }
 }
 

diff --git a/examples/demo-rollup/src/main.rs b/examples/demo-rollup/src/main.rs
@@ -46,7 +46,7 @@ async fn main() -> Result<(), anyhow::Error> {
             let rollup = new_rollup_with_mock_da(
                 &GenesisPaths::from_dir("../test-data/genesis/integration-tests"),
                 rollup_config_path,
-                Some(RollupProverConfig::Execute),
+                RollupProverConfig::Execute,
             )
             .await?;
             rollup.run().await
@@ -55,7 +55,7 @@ async fn main() -> Result<(), anyhow::Error> {
             let rollup = new_rollup_with_celestia_da(
                 &GenesisPaths::from_dir("../test-data/genesis/demo-tests"),
                 rollup_config_path,
-                Some(RollupProverConfig::Execute),
+                RollupProverConfig::Execute,
             )
             .await?;
             rollup.run().await
@@ -67,7 +67,7 @@ async fn main() -> Result<(), anyhow::Error> {
 async fn new_rollup_with_celestia_da(
     genesis_paths: &GenesisPaths,
     rollup_config_path: &str,
-    prover_config: Option<RollupProverConfig>,
+    prover_config: RollupProverConfig,
 ) -> Result<Rollup<CelestiaDemoRollup>, anyhow::Error> {
     debug!(
         "Starting celestia rollup with config {}",
@@ -86,7 +86,7 @@ async fn new_rollup_with_celestia_da(
 async fn new_rollup_with_mock_da(
     genesis_paths: &GenesisPaths,
     rollup_config_path: &str,
-    prover_config: Option<RollupProverConfig>,
+    prover_config: RollupProverConfig,
 ) -> Result<Rollup<MockDemoRollup>, anyhow::Error> {
     debug!("Starting mock rollup with config {}", rollup_config_path);
 

diff --git a/examples/demo-rollup/src/mock_rollup.rs b/examples/demo-rollup/src/mock_rollup.rs
@@ -12,7 +12,7 @@ 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::{BlockingProver, RollupConfig, RollupProverConfig};
+use sov_stf_runner::{ParallelProverService, RollupConfig, RollupProverConfig};
 
 /// Rollup with MockDa
 pub struct MockDemoRollup {}
@@ -35,7 +35,7 @@ impl RollupBlueprint for MockDemoRollup {
     type NativeKernel = BasicKernel<Self::NativeContext>;
     type ZkKernel = BasicKernel<Self::ZkContext>;
 
-    type ProverService = BlockingProver<
+    type ProverService = ParallelProverService<
         <<Self::NativeContext as Spec>::Storage as Storage>::Root,
         <<Self::NativeContext as Spec>::Storage as Storage>::Witness,
         Self::DaService,
@@ -91,14 +91,14 @@ impl RollupBlueprint for MockDemoRollup {
 
     async fn create_prover_service(
         &self,
-        prover_config: Option<RollupProverConfig>,
+        prover_config: RollupProverConfig,
         _da_service: &Self::DaService,
     ) -> Self::ProverService {
         let vm = Risc0Host::new(risc0::MOCK_DA_ELF);
         let zk_stf = StfBlueprint::new();
         let zk_storage = ZkStorage::new();
         let da_verifier = Default::default();
 
-        BlockingProver::new(vm, zk_stf, da_verifier, prover_config, zk_storage)
+        ParallelProverService::new(vm, zk_stf, da_verifier, prover_config, zk_storage)
     }
 }
diff --git a/examples/demo-rollup/tests/bank/mod.rs b/examples/demo-rollup/tests/bank/mod.rs
@@ -26,7 +26,7 @@ async fn bank_tx_tests() -> Result<(), anyhow::Error> {
         start_rollup(
             port_tx,
             GenesisPaths::from_dir("../test-data/genesis/integration-tests"),
-            Some(RollupProverConfig::Execute),
+            RollupProverConfig::Execute,
         )
         .await;
     });

diff --git a/examples/demo-rollup/tests/evm/mod.rs b/examples/demo-rollup/tests/evm/mod.rs
@@ -7,6 +7,7 @@ use demo_stf::genesis_config::GenesisPaths;
 use ethers_core::abi::Address;
 use ethers_signers::{LocalWallet, Signer};
 use sov_evm::SimpleStorageContract;
+use sov_stf_runner::RollupProverConfig;
 use test_client::TestClient;
 
 use crate::test_helpers::start_rollup;
@@ -21,7 +22,7 @@ async fn evm_tx_tests() -> Result<(), anyhow::Error> {
         start_rollup(
             port_tx,
             GenesisPaths::from_dir("../test-data/genesis/integration-tests"),
-            None,
+            RollupProverConfig::Skip,
         )
         .await;
     });

diff --git a/examples/demo-rollup/tests/test_helpers.rs b/examples/demo-rollup/tests/test_helpers.rs
@@ -10,7 +10,7 @@ use tokio::sync::oneshot;
 pub async fn start_rollup(
     rpc_reporting_channel: oneshot::Sender<SocketAddr>,
     genesis_paths: GenesisPaths,
-    rollup_prover_config: Option<RollupProverConfig>,
+    rollup_prover_config: RollupProverConfig,
 ) {
     let temp_dir = tempfile::tempdir().unwrap();
     let temp_path = temp_dir.path();

diff --git a/full-node/sov-stf-runner/Cargo.toml b/full-node/sov-stf-runner/Cargo.toml
@@ -23,9 +23,11 @@ hex = { workspace = true }
 tracing = { workspace = true, optional = true }
 futures = { workspace = true, optional = true }
 async-trait = { workspace = true, optional = true }
+rayon = { workspace = true, optional = true }
 sov-db = { path = "../db/sov-db", version = "0.3", optional = true }
 sov-rollup-interface = { path = "../../rollup-interface", version = "0.3" }
 
+
 [dev-dependencies]
 tempfile = { workspace = true }
 
@@ -50,5 +52,6 @@ native = [
     "tracing",
     "futures",
     "async-trait",
+    "rayon",
     "thiserror",
 ]
diff --git a/full-node/sov-stf-runner/src/prover_service/blocking_prover.rs b/full-node/sov-stf-runner/src/prover_service/blocking_prover.rs
diff --git a/full-node/sov-stf-runner/src/prover_service/mod.rs b/full-node/sov-stf-runner/src/prover_service/mod.rs
@@ -1,6 +1,6 @@
-mod blocking_prover;
+mod parallel;
 use async_trait::async_trait;
-pub use blocking_prover::BlockingProver;
+pub use parallel::ParallelProverService;
 use serde::Serialize;
 use sov_rollup_interface::services::da::DaService;
 use thiserror::Error;
@@ -11,6 +11,8 @@ pub(crate) type Hash = [u8; 32];
 
 /// The possible configurations of the prover.
 pub enum RollupProverConfig {
+    /// Skip proving.
+    Skip,
     /// Run the rollup verification logic inside the current process
     Simulate,
     /// Run the rollup verifier in a zkVM executor
@@ -19,6 +21,16 @@ pub enum RollupProverConfig {
     Prove,
 }
 
+/// Indicates the status of the DA proof submission.
+pub enum ProofSubmissionStatus {
+    /// Proof was submitted to the DA.
+    Success,
+    /// Proof generation is still in progress.
+    ProvingInProgress,
+    /// Proof submission failed
+    Err(anyhow::Error),
+}
+
 /// An error that occurred during ZKP proving.
 #[derive(Error, Debug)]
 pub enum ProverServiceError {
@@ -55,4 +67,7 @@ pub trait ProverService {
 
     /// Creates ZKP prove for a block corresponding to `block_header_hash`.
     async fn prove(&self, block_header_hash: Hash) -> Result<(), ProverServiceError>;
+
+    /// Sends the ZK proof to the DA create by the `prove`.
+    async fn send_proof_to_da(&self, block_header_hash: Hash) -> ProofSubmissionStatus;
 }