feat: use abort signal for proving tasks (#748)

* feat: use abort signal for proving tasks

* fix: abort signal

* fix: update paladin dependency

* fix: cargo fmt

Cargo.toml

@@ -70,7 +70,8 @@ num-bigint = "0.4.5"
 num-traits = "0.2.19"
 nunny = "0.2.1"
 once_cell = "1.19.0"
-paladin-core = "0.4.3"
+# TODO: update when paladin is released to 0.4.4 with abort signal support
+paladin-core = { git = "https://github.com/0xPolygonZero/paladin.git", branch = "main" }
 parking_lot = "0.12.3"
 pest = "2.7.10"
 pest_derive = "2.7.10"

zero/src/bin/leader.rs

@@ -27,7 +27,7 @@ mod leader {
 
 const HEAVY_PROOF_ROUTING_KEY: &str = "heavy-proof";
 const LIGHT_PROOF_ROUTING_KEY: &str = "light-proof";
-const DEFAULT_ROUTING_KEY: &str = paladin::runtime::DEFAULT_ROUTING_KEY;
+const TASK_IPC_ROUTING_KEY: &str = paladin::runtime::TASK_IPC_ROUTING_KEY;
 
 #[tokio::main]
 async fn main() -> Result<()> {
@@ -41,8 +41,8 @@ async fn main() -> Result<()> {
         return zero::prover_state::persistence::delete_all();
     }
 
-    let mut light_proof_routing_key = DEFAULT_ROUTING_KEY.to_string();
-    let mut heavy_proof_routing_key = DEFAULT_ROUTING_KEY.to_string();
+    let mut light_proof_routing_key = TASK_IPC_ROUTING_KEY.to_string();
+    let mut heavy_proof_routing_key = TASK_IPC_ROUTING_KEY.to_string();
     if args.worker_run_mode == cli::WorkerRunMode::Affinity {
         // If we're running in affinity mode, we need to set the routing key for the
         // heavy proof and light proof.

zero/src/ops.rs

@@ -8,7 +8,7 @@ use evm_arithmetization::{prover::testing::simulate_execution_all_segments, Gene
 use evm_arithmetization::{Field, ProofWithPublicValues, PublicValues, TrimmedGenerationInputs};
 use paladin::{
     operation::{FatalError, FatalStrategy, Monoid, Operation, Result},
-    registry, RemoteExecute,
+    registry, AbortSignal, RemoteExecute,
 };
 use serde::{Deserialize, Serialize};
 use tracing::error;
@@ -30,7 +30,7 @@ impl Operation for SegmentProof {
     type Input = evm_arithmetization::AllData;
     type Output = SegmentAggregatableProof;
 
-    fn execute(&self, all_data: Self::Input) -> Result<Self::Output> {
+    fn execute(&self, all_data: Self::Input, abort: AbortSignal) -> Result<Self::Output> {
         let all_data =
             all_data.map_err(|e| FatalError::from_str(&e.to_string(), FatalStrategy::Terminate))?;
 
@@ -39,7 +39,7 @@ impl Operation for SegmentProof {
         let _span = SegmentProofSpan::new(&input, all_data.1.segment_index());
         let proof = if self.save_inputs_on_error {
             crate::prover_state::p_manager()
-                .generate_segment_proof(all_data)
+                .generate_segment_proof(all_data, abort)
                 .map_err(|e| {
                     if let Err(write_err) = save_inputs_to_disk(
                         format!(
@@ -58,7 +58,7 @@ impl Operation for SegmentProof {
                 })?
         } else {
             crate::prover_state::p_manager()
-                .generate_segment_proof(all_data)
+                .generate_segment_proof(all_data, abort)
                 .map_err(|e| FatalError::from_str(&e.to_string(), FatalStrategy::Terminate))?
         };
 
@@ -78,7 +78,7 @@ impl Operation for SegmentProofTestOnly {
     type Input = (GenerationInputs, usize, usize);
     type Output = ();
 
-    fn execute(&self, inputs: Self::Input) -> Result<Self::Output> {
+    fn execute(&self, inputs: Self::Input, _abort: AbortSignal) -> Result<Self::Output> {
         if self.save_inputs_on_error || self.save_tries_on_error {
             simulate_execution_all_segments::<Field>(inputs.0.clone(), inputs.1).map_err(|err| {
                 let block_number = inputs.0.block_metadata.block_number.low_u64();
@@ -259,7 +259,7 @@ pub fn generate_segment_agg_proof(
 impl Monoid for SegmentAggProof {
     type Elem = SegmentAggregatableProof;
 
-    fn combine(&self, a: Self::Elem, b: Self::Elem) -> Result<Self::Elem> {
+    fn combine(&self, a: Self::Elem, b: Self::Elem, _abort: AbortSignal) -> Result<Self::Elem> {
         let proof = generate_segment_agg_proof(p_state(), &a, &b, false).map_err(|e| {
             if self.save_inputs_on_error {
                 let pv = vec![
@@ -304,7 +304,7 @@ fn get_agg_proof_public_values(elem: BatchAggregatableProof) -> PublicValues {
 impl Monoid for BatchAggProof {
     type Elem = BatchAggregatableProof;
 
-    fn combine(&self, a: Self::Elem, b: Self::Elem) -> Result<Self::Elem> {
+    fn combine(&self, a: Self::Elem, b: Self::Elem, _abort: AbortSignal) -> Result<Self::Elem> {
         let lhs = match a {
             BatchAggregatableProof::Segment(segment) => BatchAggregatableProof::SegmentAgg(
                 generate_segment_agg_proof(
@@ -378,7 +378,7 @@ impl Operation for BlockProof {
     type Input = ProofWithPublicValues;
     type Output = GeneratedBlockProof;
 
-    fn execute(&self, input: Self::Input) -> Result<Self::Output> {
+    fn execute(&self, input: Self::Input, _abort: AbortSignal) -> Result<Self::Output> {
         let b_height = input.public_values.block_metadata.block_number.low_u64();
         let parent_intern = self.prev.as_ref().map(|p| &p.intern);
 

zero/src/prover_state/mod.rs

@@ -12,6 +12,8 @@
 //! - Global prover state management via the `P_STATE` static and the
 //!   [`p_state`] function.
 use std::borrow::Borrow;
+use std::sync::atomic::AtomicBool;
+use std::sync::Arc;
 use std::{fmt::Display, sync::OnceLock};
 
 use clap::ValueEnum;
@@ -209,6 +211,7 @@ impl ProverStateManager {
         input: TrimmedGenerationInputs,
         segment_data: &mut GenerationSegmentData,
         config: &StarkConfig,
+        abort_signal: Option<Arc<AtomicBool>>,
     ) -> anyhow::Result<ProofWithPublicValues> {
         let all_stark = AllStark::default();
 
@@ -218,15 +221,16 @@ impl ProverStateManager {
             input,
             segment_data,
             &mut TimingTree::default(),
-            None,
+            abort_signal.clone(),
         )?;
 
         let table_circuits = self.load_table_circuits(config, &all_proof)?;
 
-        let proof_with_pvs =
-            p_state()
-                .state
-                .prove_segment_after_initial_stark(all_proof, &table_circuits, None)?;
+        let proof_with_pvs = p_state().state.prove_segment_after_initial_stark(
+            all_proof,
+            &table_circuits,
+            abort_signal,
+        )?;
 
         Ok(proof_with_pvs)
     }
@@ -238,14 +242,15 @@ impl ProverStateManager {
         input: TrimmedGenerationInputs,
         segment_data: &mut GenerationSegmentData,
         config: &StarkConfig,
+        abort_signal: Option<Arc<AtomicBool>>,
     ) -> anyhow::Result<ProofWithPublicValues> {
         let p_out = p_state().state.prove_segment(
             &AllStark::default(),
             config,
             input,
             segment_data,
             &mut TimingTree::default(),
-            None,
+            abort_signal,
         )?;
 
         let ProverOutputData {
@@ -269,6 +274,7 @@ impl ProverStateManager {
     pub fn generate_segment_proof(
         &self,
         input: (TrimmedGenerationInputs, GenerationSegmentData),
+        abort_signal: Option<Arc<AtomicBool>>,
     ) -> anyhow::Result<ProofWithPublicValues> {
         let (generation_inputs, mut segment_data) = input;
         let config = if self.circuit_config.use_test_config {
@@ -280,11 +286,21 @@ impl ProverStateManager {
         match self.persistence {
             CircuitPersistence::None | CircuitPersistence::Disk(TableLoadStrategy::Monolithic) => {
                 info!("using monolithic circuit {:?}", self);
-                self.segment_proof_monolithic(generation_inputs, &mut segment_data, &config)
+                self.segment_proof_monolithic(
+                    generation_inputs,
+                    &mut segment_data,
+                    &config,
+                    abort_signal,
+                )
             }
             CircuitPersistence::Disk(TableLoadStrategy::OnDemand) => {
                 info!("using on demand circuit {:?}", self);
-                self.segment_proof_on_demand(generation_inputs, &mut segment_data, &config)
+                self.segment_proof_on_demand(
+                    generation_inputs,
+                    &mut segment_data,
+                    &config,
+                    abort_signal,
+                )
             }
         }
     }