add worker run modes

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/main.rs

@@ -38,6 +38,7 @@ fn get_previous_proof(path: Option<PathBuf>) -> Result<Option<GeneratedBlockProo
 
 const SEGMENT_PROOF_ROUTING_KEY: &str = "segment_proof";
 const BLOCK_PROOF_ROUTING_KEY: &str = "block_proof";
+const DEFAULT_ROUTING_KEY: &str = "task";
 
 #[tokio::main]
 async fn main() -> Result<()> {
@@ -58,11 +59,21 @@ async fn main() -> Result<()> {
 
     let args = cli::Cli::parse();
 
+    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.to_string());
+    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.to_string());
+    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 =

zero_bin/tools/prove_rpc.sh

@@ -102,7 +102,7 @@ fi
 if [[ $8 == "test_only" ]]; then
     # test only run
     echo "Proving blocks ${BLOCK_INTERVAL} in a test_only mode now... (Total: ${TOT_BLOCKS})"
-    command='cargo r --release --bin leader -- --test-only --runtime in-memory --load-strategy on-demand rpc --rpc-type "$NODE_RPC_TYPE" --rpc-url "$NODE_RPC_URL" --block-interval $BLOCK_INTERVAL --proof-output-dir $PROOF_OUTPUT_DIR $PREV_PROOF_EXTRA_ARG --backoff "$BACKOFF" --max-retries "$RETRIES" '
+    command='cargo r --release --bin leader -- --runtime in-memory --test-only --load-strategy on-demand rpc --rpc-type "$NODE_RPC_TYPE" --rpc-url "$NODE_RPC_URL" --block-interval $BLOCK_INTERVAL --proof-output-dir $PROOF_OUTPUT_DIR $PREV_PROOF_EXTRA_ARG --backoff "$BACKOFF" --max-retries "$RETRIES" '
     if [ "$OUTPUT_TO_TERMINAL" = true ]; then
         eval $command
         retVal=$?
@@ -125,7 +125,7 @@ if [[ $8 == "test_only" ]]; then
 else
     # normal run
     echo "Proving blocks ${BLOCK_INTERVAL} now... (Total: ${TOT_BLOCKS})"
-    command='cargo r --release --bin leader -- --runtime in-memory --load-strategy on-demand rpc --rpc-type "$NODE_RPC_TYPE" --rpc-url "$3" --block-interval $BLOCK_INTERVAL --proof-output-dir $PROOF_OUTPUT_DIR $PREV_PROOF_EXTRA_ARG --backoff "$BACKOFF" --max-retries "$RETRIES" '
+    command='cargo r --release --bin leader -- --runtime in-memory --load-strategy on-demand rpc --rpc-type "$NODE_RPC_TYPE" --rpc-url "$3" --block-interval $BLOCK_INTERVAL --proof-output-dir $PROOF_OUTPUT_DIR $PREV_PROOF_EXTRA_ARG --backoff "$BACKOFF" --max-retries "$RETRIES"'
     if [ "$OUTPUT_TO_TERMINAL" = true ]; then
         eval $command
         echo -e "Proof generation finished with result: $?"