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 = "main" }
 parking_lot = "0.12.3"
 paste = "1.0.15"
 pest = "2.7.10"

zero_bin/common/Cargo.toml

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

zero_bin/common/src/lib.rs

@@ -6,6 +6,7 @@ pub mod env;
 pub mod fs;
 pub mod parsing;
 pub mod pre_checks;
+pub mod proof_runtime;
 pub mod prover_state;
 pub mod provider;
 pub mod tracing;

zero_bin/common/src/proof_runtime.rs

@@ -0,0 +1,6 @@
+use paladin::runtime::Runtime;
+
+pub struct ProofRuntime {
+    pub light_proof_runtime: Runtime,
+    pub heavy_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 (affinity or default)
+    #[arg(long = "worker-run-mode", help_heading = WORKER_HELP_HEADING, value_enum, default_value = "default")]
+    pub(crate) worker_run_mode: WorkerRunMode,
+}
+
+#[derive(ValueEnum, Clone, PartialEq, Debug)]
+pub enum WorkerRunMode {
+    Affinity,
+    Default,
 }
 
 #[derive(Subcommand)]

zero_bin/leader/src/client.rs

@@ -3,14 +3,14 @@ use std::sync::Arc;
 use alloy::rpc::types::{BlockId, BlockNumberOrTag};
 use alloy::transports::http::reqwest::Url;
 use anyhow::{anyhow, Result};
-use paladin::runtime::Runtime;
 use proof_gen::proof_types::GeneratedBlockProof;
 use prover::{BlockProverInput, ProverConfig};
 use rpc::{retry::build_http_retry_provider, RpcType};
 use tokio::sync::mpsc;
 use tracing::info;
 use zero_bin_common::block_interval::{BlockInterval, BlockIntervalStream};
 use zero_bin_common::pre_checks::check_previous_proof_and_checkpoint;
+use zero_bin_common::proof_runtime::ProofRuntime;
 
 #[derive(Debug)]
 pub struct RpcParams {
@@ -30,7 +30,7 @@ pub struct LeaderConfig {
 
 /// The main function for the client.
 pub(crate) async fn client_main(
-    runtime: Arc<Runtime>,
+    proof_runtime: Arc<ProofRuntime>,
     rpc_params: RpcParams,
     block_interval: BlockInterval,
     mut leader_config: LeaderConfig,
@@ -58,10 +58,10 @@ pub(crate) async fn client_main(
     let test_only = leader_config.prover_config.test_only;
 
     // Run proving task
-    let runtime_ = runtime.clone();
+    let proof_runtime_ = proof_runtime.clone();
     let proving_task = tokio::spawn(prover::prove(
         block_rx,
-        runtime_,
+        proof_runtime_,
         leader_config.previous_proof.take(),
         Arc::new(leader_config.prover_config),
     ));
@@ -107,7 +107,8 @@ pub(crate) async fn client_main(
         }
     }
 
-    runtime.close().await?;
+    proof_runtime.light_proof_runtime.close().await?;
+    proof_runtime.heavy_proof_runtime.close().await?;
 
     if test_only {
         info!("All proof witnesses have been generated successfully.");