Assign specific jobs to dedicated workers

Cargo.toml

@@ -78,7 +78,7 @@ num-bigint = "0.4.5"
 num-traits = "0.2.19"
 nunny = "0.2.1"
 once_cell = "1.19.0"
-paladin-core = "0.4.2"
+paladin-core = { git = "https://github.com/0xPolygonZero/paladin.git", branch = "arpit/507-2" }
 parking_lot = "0.12.3"
 paste = "1.0.15"
 pest = "2.7.10"

zero_bin/common/Cargo.toml

@@ -18,6 +18,7 @@ evm_arithmetization = { workspace = true }
 futures = { workspace = true }
 lru = { workspace = true }
 once_cell = { workspace = true }
+paladin-core = { workspace = true }
 plonky2 = { workspace = true }
 proof_gen = { workspace = true }
 serde = { workspace = true }

zero_bin/common/src/lib.rs

@@ -3,6 +3,7 @@ pub mod debug_utils;
 pub mod fs;
 pub mod parsing;
 pub mod pre_checks;
+pub mod proof_runtime;
 pub mod prover_state;
 pub mod provider;
 pub mod version;

zero_bin/common/src/proof_runtime.rs

@@ -0,0 +1,6 @@
+use paladin::runtime::Runtime;
+
+pub struct ProofRuntime {
+    pub block_proof_runtime: Runtime,
+    pub segment_proof_runtime: Runtime,
+}

zero_bin/leader/src/cli.rs

@@ -1,11 +1,13 @@
 use std::path::PathBuf;
 
 use alloy::transports::http::reqwest::Url;
-use clap::{Parser, Subcommand, ValueHint};
+use clap::{Parser, Subcommand, ValueEnum, ValueHint};
 use prover::cli::CliProverConfig;
 use rpc::RpcType;
 use zero_bin_common::prover_state::cli::CliProverStateConfig;
 
+const WORKER_HELP_HEADING: &str = "Worker Config options";
+
 /// zero-bin leader config
 #[derive(Parser)]
 pub(crate) struct Cli {
@@ -22,6 +24,16 @@ pub(crate) struct Cli {
     // mode.
     #[clap(flatten)]
     pub(crate) prover_state_config: CliProverStateConfig,
+
+    // Mode to use for worker for setup (split or unified)
+    #[arg(long = "worker-run-mode", help_heading = WORKER_HELP_HEADING, value_enum, default_value = "unified")]
+    pub(crate) worker_run_mode: WorkerRunMode,
+}
+
+#[derive(ValueEnum, Clone, PartialEq, Debug)]
+pub enum WorkerRunMode {
+    Split,
+    Unified,
 }
 
 #[derive(Subcommand)]

zero_bin/leader/src/client.rs

@@ -5,14 +5,14 @@ use std::sync::Arc;
 use alloy::rpc::types::{BlockId, BlockNumberOrTag, BlockTransactionsKind};
 use alloy::transports::http::reqwest::Url;
 use anyhow::Result;
-use paladin::runtime::Runtime;
 use proof_gen::proof_types::GeneratedBlockProof;
 use prover::ProverConfig;
 use rpc::{retry::build_http_retry_provider, RpcType};
 use tracing::{error, info, warn};
 use zero_bin_common::block_interval::BlockInterval;
 use zero_bin_common::fs::generate_block_proof_file_name;
 use zero_bin_common::pre_checks::check_previous_proof_and_checkpoint;
+use zero_bin_common::proof_runtime::ProofRuntime;
 
 #[derive(Debug)]
 pub struct RpcParams {
@@ -33,7 +33,7 @@ pub struct ProofParams {
 
 /// The main function for the client.
 pub(crate) async fn client_main(
-    runtime: Runtime,
+    proof_runtime: ProofRuntime,
     rpc_params: RpcParams,
     block_interval: BlockInterval,
     mut params: ProofParams,
@@ -82,13 +82,14 @@ pub(crate) async fn client_main(
     // verify the whole sequence.
     let proved_blocks = prover::prove(
         block_prover_inputs,
-        &runtime,
+        &proof_runtime,
         params.previous_proof.take(),
         params.prover_config,
         params.proof_output_dir.clone(),
     )
     .await;
-    runtime.close().await?;
+    proof_runtime.block_proof_runtime.close().await?;
+    proof_runtime.segment_proof_runtime.close().await?;
     let proved_blocks = proved_blocks?;
 
     if params.prover_config.test_only {

zero_bin/leader/src/http.rs

@@ -3,30 +3,27 @@ use std::{net::SocketAddr, path::PathBuf, sync::Arc};
 use alloy::primitives::U256;
 use anyhow::{bail, Result};
 use axum::{http::StatusCode, routing::post, Json, Router};
-use paladin::runtime::Runtime;
 use proof_gen::proof_types::GeneratedBlockProof;
 use prover::{BlockProverInput, ProverConfig};
 use serde::{Deserialize, Serialize};
 use serde_json::to_writer;
 use tracing::{debug, error, info};
+use zero_bin_common::proof_runtime::ProofRuntime;
 
 /// The main function for the HTTP mode.
 pub(crate) async fn http_main(
-    runtime: Runtime,
+    proof_runtime: ProofRuntime,
     port: u16,
     output_dir: PathBuf,
     prover_config: ProverConfig,
 ) -> Result<()> {
     let addr = SocketAddr::from(([0, 0, 0, 0], port));
     debug!("listening on {}", addr);
 
-    let runtime = Arc::new(runtime);
+    let proof_runtime = Arc::new(proof_runtime);
     let app = Router::new().route(
         "/prove",
-        post({
-            let runtime = runtime.clone();
-            move |body| prove(body, runtime, output_dir.clone(), prover_config)
-        }),
+        post(move |body| prove(body, proof_runtime, output_dir.clone(), prover_config)),
     );
     let listener = tokio::net::TcpListener::bind(&addr).await?;
     Ok(axum::serve(listener, app).await?)
@@ -63,7 +60,7 @@ struct HttpProverInput {
 
 async fn prove(
     Json(payload): Json<HttpProverInput>,
-    runtime: Arc<Runtime>,
+    proof_runtime: Arc<ProofRuntime>,
     output_dir: PathBuf,
     prover_config: ProverConfig,
 ) -> StatusCode {
@@ -75,7 +72,7 @@ async fn prove(
         payload
             .prover_input
             .prove_test(
-                &runtime,
+                &proof_runtime,
                 payload.previous.map(futures::future::ok),
                 prover_config,
             )
@@ -84,7 +81,7 @@ async fn prove(
         payload
             .prover_input
             .prove(
-                &runtime,
+                &proof_runtime,
                 payload.previous.map(futures::future::ok),
                 prover_config,
             )

zero_bin/leader/src/main.rs

@@ -12,7 +12,8 @@ use proof_gen::proof_types::GeneratedBlockProof;
 use prover::ProverConfig;
 use tracing::{info, warn};
 use zero_bin_common::{
-    block_interval::BlockInterval, prover_state::persistence::set_circuit_cache_dir_env_if_not_set,
+    block_interval::BlockInterval, proof_runtime::ProofRuntime,
+    prover_state::persistence::set_circuit_cache_dir_env_if_not_set,
 };
 use zero_bin_common::{prover_state::persistence::CIRCUIT_VERSION, version};
 
@@ -36,6 +37,10 @@ fn get_previous_proof(path: Option<PathBuf>) -> Result<Option<GeneratedBlockProo
     Ok(Some(proof))
 }
 
+const SEGMENT_PROOF_ROUTING_KEY: &str = "segment_proof";
+const BLOCK_PROOF_ROUTING_KEY: &str = "block_proof";
+const DEFAULT_ROUTING_KEY: &str = paladin::runtime::DEFAULT_ROUTING_KEY;
+
 #[tokio::main]
 async fn main() -> Result<()> {
     load_dotenvy_vars_if_present();
@@ -55,7 +60,25 @@ async fn main() -> Result<()> {
 
     let args = cli::Cli::parse();
 
-    let runtime = Runtime::from_config(&args.paladin, register()).await?;
+    let mut block_proof_routing_key = DEFAULT_ROUTING_KEY.to_string();
+    let mut segment_proof_routing_key = DEFAULT_ROUTING_KEY.to_string();
+    if args.worker_run_mode == cli::WorkerRunMode::Split {
+        // If we're running in split mode, we need to set the routing key for the
+        // block proof and segment proof.
+        info!("Workers running in split mode");
+        block_proof_routing_key = BLOCK_PROOF_ROUTING_KEY.to_string();
+        segment_proof_routing_key = SEGMENT_PROOF_ROUTING_KEY.to_string();
+    }
+
+    let mut block_proof_paladin_args = args.paladin.clone();
+    block_proof_paladin_args.task_bus_routing_key = Some(block_proof_routing_key);
+
+    let mut segment_proof_paladin_args = args.paladin.clone();
+    segment_proof_paladin_args.task_bus_routing_key = Some(segment_proof_routing_key);
+
+    let block_proof_runtime = Runtime::from_config(&block_proof_paladin_args, register()).await?;
+    let segment_proof_runtime =
+        Runtime::from_config(&segment_proof_paladin_args, register()).await?;
 
     let prover_config: ProverConfig = args.prover_config.into();
 
@@ -69,11 +92,16 @@ async fn main() -> Result<()> {
         }
     }
 
+    let proof_runtime = ProofRuntime {
+        block_proof_runtime,
+        segment_proof_runtime,
+    };
+
     match args.command {
         Command::Clean => zero_bin_common::prover_state::persistence::delete_all()?,
         Command::Stdio { previous_proof } => {
             let previous_proof = get_previous_proof(previous_proof)?;
-            stdio::stdio_main(runtime, previous_proof, prover_config).await?;
+            stdio::stdio_main(proof_runtime, previous_proof, prover_config).await?;
         }
         Command::Http { port, output_dir } => {
             // check if output_dir exists, is a directory, and is writable
@@ -85,7 +113,7 @@ async fn main() -> Result<()> {
                 panic!("output-dir is not a writable directory");
             }
 
-            http::http_main(runtime, port, output_dir, prover_config).await?;
+            http::http_main(proof_runtime, port, output_dir, prover_config).await?;
         }
         Command::Rpc {
             rpc_url,
@@ -99,7 +127,6 @@ async fn main() -> Result<()> {
             backoff,
             max_retries,
         } => {
-            let runtime = Runtime::from_config(&args.paladin, register()).await?;
             let previous_proof = get_previous_proof(previous_proof)?;
             let mut block_interval = BlockInterval::new(&block_interval)?;
 
@@ -113,7 +140,7 @@ async fn main() -> Result<()> {
 
             info!("Proving interval {block_interval}");
             client_main(
-                runtime,
+                proof_runtime,
                 RpcParams {
                     rpc_url,
                     rpc_type,

zero_bin/leader/src/stdio.rs

@@ -1,14 +1,14 @@
 use std::io::{Read, Write};
 
 use anyhow::Result;
-use paladin::runtime::Runtime;
 use proof_gen::proof_types::GeneratedBlockProof;
 use prover::{BlockProverInput, BlockProverInputFuture, ProverConfig};
 use tracing::info;
+use zero_bin_common::proof_runtime::ProofRuntime;
 
 /// The main function for the stdio mode.
 pub(crate) async fn stdio_main(
-    runtime: Runtime,
+    proof_runtime: ProofRuntime,
     previous: Option<GeneratedBlockProof>,
     prover_config: ProverConfig,
 ) -> Result<()> {
@@ -21,9 +21,16 @@ pub(crate) async fn stdio_main(
         .map(Into::into)
         .collect::<Vec<BlockProverInputFuture>>();
 
-    let proved_blocks =
-        prover::prove(block_prover_inputs, &runtime, previous, prover_config, None).await;
-    runtime.close().await?;
+    let proved_blocks = prover::prove(
+        block_prover_inputs,
+        &proof_runtime,
+        previous,
+        prover_config,
+        None,
+    )
+    .await;
+    proof_runtime.block_proof_runtime.close().await?;
+    proof_runtime.segment_proof_runtime.close().await?;
     let proved_blocks = proved_blocks?;
 
     if prover_config.test_only {