feat: Revive continue-on-error mode (#11)

This PR enables correct processing of panics for the continue-on-error
mode, and also adds logic to re-instantiate prover in case of panic.
Some minor fixes here and there as well.

Author: popzxc

ethereum_prover/configs/ethproofs_prod.yaml

@@ -1,7 +1,7 @@
 # Use this config to submit proofs to EthProofs staging
 
 eth_prover:
-  app_bin_path: "../../artifacts/app.bin"
+  app_bin_path: "../artifacts/app.bin"
   mode: "gpu_prove"
   cache_policy: "on_failure"
   ethproofs_submission: "prod"

ethereum_prover/configs/ethproofs_staging.yaml

@@ -1,7 +1,7 @@
 # Use this config to submit proofs to EthProofs staging
 
 eth_prover:
-  app_bin_path: "../../artifacts/app.bin"
+  app_bin_path: "../artifacts/app.bin"
   mode: "gpu_prove"
   cache_policy: "on_failure"
   ethproofs_submission: "staging"

ethereum_prover/src/lib.rs

@@ -20,6 +20,7 @@ pub mod metrics;
 pub mod prover;
 pub(crate) mod tasks;
 pub(crate) mod types;
+pub(crate) mod utils;
 
 #[derive(Debug, Default)]
 pub struct Runner {}

ethereum_prover/src/prover/cpu_witness.rs

@@ -4,6 +4,7 @@ use std::path::PathBuf;
 use alloy::providers::DynProvider;
 use alloy::providers::Provider;
 use alloy::rpc::types::Transaction;
+use anyhow::Context as _;
 use basic_bootloader::bootloader::BasicBootloader;
 use basic_bootloader::bootloader::config::BasicBootloaderForwardETHLikeConfig;
 use forward_system::run::InvalidTransaction;
@@ -68,7 +69,7 @@ impl CpuWitnessGenerator {
         oracle: ZkEENonDeterminismSource,
     ) -> anyhow::Result<Vec<u32>> {
         let app_bin_path = self.app_bin_path.clone();
-        tokio::task::spawn_blocking(move || {
+        match tokio::task::spawn_blocking(move || {
             let copy_source = ReadWitnessSource::new(oracle);
             let items = copy_source.get_read_items();
 
@@ -80,7 +81,18 @@ impl CpuWitnessGenerator {
             let witness = items.borrow().clone();
             Ok(witness)
         })
-        .await?
+        .await
+        {
+            Ok(Ok(witness)) => Ok(witness),
+            Ok(Err(err)) => Err(err).context("Failed to generate witness using zksync_os_runner"),
+            Err(err) => {
+                let panic_msg = crate::utils::extract_panic_message(err);
+                Err(anyhow::anyhow!(
+                    "Witness generation task panicked: {}",
+                    panic_msg
+                ))
+            }
+        }
     }
 }
 

ethereum_prover/src/prover/gpu_prover.rs

@@ -1,5 +1,6 @@
+use anyhow::Context as _;
 use oracle_provider::ZkEENonDeterminismSource;
-use std::path::Path;
+use std::path::{Path, PathBuf};
 use std::sync::{Arc, Mutex};
 use std::time::Instant;
 
@@ -11,6 +12,8 @@ pub struct ProofResult {
 }
 
 pub struct Prover {
+    app_bin_path: PathBuf,
+    worker_threads: Option<usize>,
     inner: Arc<Mutex<execution_utils::unrolled_gpu::UnrolledProver>>,
 }
 
@@ -22,19 +25,15 @@ impl std::fmt::Debug for Prover {
 
 impl Prover {
     pub fn new(app_bin_path: &Path, worker_threads: Option<usize>) -> anyhow::Result<Self> {
-        let base_path = strip_bin_suffix(app_bin_path)?;
-        let mut configuration =
-            execution_utils::gpu_prover::execution::prover::ExecutionProverConfiguration::default();
-        if let Some(threads) = worker_threads {
-            configuration.max_thread_pool_threads = Some(threads);
-            configuration.replay_worker_threads_count = threads;
-        }
-        let inner = execution_utils::unrolled_gpu::UnrolledProver::new(
-            &base_path,
-            configuration,
-            execution_utils::unrolled_gpu::UnrolledProverLevel::RecursionUnified,
-        );
+        let inner = create_unrolled_prover(app_bin_path, worker_threads).with_context(|| {
+            format!(
+                "failed to create unrolled prover with app binary at {:?}",
+                app_bin_path
+            )
+        })?;
         Ok(Self {
+            app_bin_path: app_bin_path.to_path_buf(),
+            worker_threads,
             inner: Arc::new(Mutex::new(inner)),
         })
     }
@@ -48,11 +47,33 @@ impl Prover {
 
         let inner = self.inner.clone();
 
-        let (proof, cycles) = tokio::task::spawn_blocking(move || {
+        let future_result = tokio::task::spawn_blocking(move || {
             let prover = inner.lock().unwrap();
             prover.prove(block_number, oracle)
         })
-        .await?;
+        .await;
+        let (proof, cycles) = match future_result {
+            Ok(result) => result,
+            Err(err) => {
+                let panic_msg = crate::utils::extract_panic_message(err);
+
+                // If prover panics, it is not safe to use it again, since some of threads may be poisoned/dead.
+                // We need to re-instantiate it.
+                {
+                    let mut inner = self.inner.lock().unwrap();
+                    // If we cannot reinstantiate the prover for some reason, we cannot do much -- better to panic.
+                    *inner = create_unrolled_prover(self.app_bin_path.as_path(), self.worker_threads).expect(
+                        "failed to re-instantiate prover after panic; prover app binary path is invalid",
+                    );
+                }
+
+                return Err(anyhow::anyhow!(
+                    "Prover task panicked for the block {}: {}",
+                    block_number,
+                    panic_msg
+                ));
+            }
+        };
 
         let proving_time_secs = start.elapsed().as_secs_f64();
         let proof_bytes = bincode::serde::encode_to_vec(&proof, bincode::config::standard())?;
@@ -74,3 +95,23 @@ fn strip_bin_suffix(path: &Path) -> anyhow::Result<String> {
         Ok(path_str.to_string())
     }
 }
+
+fn create_unrolled_prover(
+    app_bin_path: &Path,
+    worker_threads: Option<usize>,
+) -> anyhow::Result<execution_utils::unrolled_gpu::UnrolledProver> {
+    let base_path = strip_bin_suffix(app_bin_path)?;
+    let mut configuration =
+        execution_utils::gpu_prover::execution::prover::ExecutionProverConfiguration::default();
+    if let Some(threads) = worker_threads {
+        configuration.max_thread_pool_threads = Some(threads);
+        configuration.replay_worker_threads_count = threads;
+    }
+
+    let unrolled_prover = execution_utils::unrolled_gpu::UnrolledProver::new(
+        &base_path,
+        configuration,
+        execution_utils::unrolled_gpu::UnrolledProverLevel::RecursionUnified,
+    );
+    Ok(unrolled_prover)
+}

ethereum_prover/src/tasks/block_stream/continuous.rs

@@ -27,7 +27,8 @@ impl ContinuousBlockStream {
         cache: CacheStorage,
         cache_policy: CachePolicy,
     ) -> (Self, Receiver<EthBlockInput>) {
-        let (sender, receiver) = channel(10);
+        // Capacity of 1 ensures that we won't be too far behind in case proving takes more time than expected.
+        let (sender, receiver) = channel(1);
 
         let provider = alloy::providers::ProviderBuilder::new().connect_http(rpc_url);
         let provider = DynProvider::new(provider);

ethereum_prover/src/tasks/cpu_witness.rs

@@ -68,45 +68,31 @@ impl CpuWitnessTask {
                         })
                         .await?;
                 }
-                Err(err) => match self.on_failure {
-                    OnFailure::Exit => {
-                        METRICS.witness_failure_total.inc();
-                        latency.observe();
-                        sentry::with_scope(
-                            |scope| {
-                                scope.set_level(Some(sentry::Level::Error));
-                                scope.set_tag("mode", "cpu_witness");
-                                scope.set_tag("block_number", block_number.to_string());
-                            },
-                            || sentry_anyhow::capture_anyhow(&err),
-                        );
-                        return Err(err).with_context(|| {
-                            format!("Failed to generate witness for the block {block_number}")
-                        });
-                    }
-                    OnFailure::Continue => {
-                        METRICS.witness_failure_total.inc();
-                        latency.observe();
-                        sentry::with_scope(
-                            |scope| {
-                                scope.set_level(Some(sentry::Level::Error));
-                                scope.set_tag("mode", "cpu_witness");
-                                scope.set_tag("block_number", block_number.to_string());
-                            },
-                            || {
-                                sentry_anyhow::capture_anyhow(&err);
-                            },
-                        );
-                        tracing::error!(
-                            "Failed to generate witness for the block {block_number}: {err}"
-                        );
-                        // continue;
-                        // TODO: `continue` mode is not supported yet, as some of the panics are not unwind-safe.
-                        return Err(err).with_context(|| {
-                            format!("Failed to generate witness for the block {block_number}")
-                        });
+                Err(err) => {
+                    METRICS.witness_failure_total.inc();
+                    latency.observe();
+                    sentry::with_scope(
+                        |scope| {
+                            scope.set_level(Some(sentry::Level::Error));
+                            scope.set_tag("mode", "cpu_witness");
+                            scope.set_tag("block_number", block_number.to_string());
+                        },
+                        || sentry_anyhow::capture_anyhow(&err),
+                    );
+                    match self.on_failure {
+                        OnFailure::Exit => {
+                            return Err(err).with_context(|| {
+                                format!("Failed to generate witness for the block {block_number}")
+                            });
+                        }
+                        OnFailure::Continue => {
+                            tracing::error!(
+                                "Failed to generate witness for the block {block_number}: {err}"
+                            );
+                            continue;
+                        }
                     }
-                },
+                }
             }
 
             // We don't any commands as we're not generating proofs

ethereum_prover/src/tasks/gpu_prove.rs

@@ -55,7 +55,12 @@ impl GpuProveTask {
 
             match self.process_block(witness).await {
                 Ok(proof_result) => {
-                    tracing::info!("Generated GPU proof for block {}", block_number);
+                    tracing::info!(
+                        "Generated GPU proof for block {}. Number of cycles: {}, proving time: {}s",
+                        block_number,
+                        proof_result.cycles,
+                        proof_result.proving_time_secs
+                    );
                     METRICS.proof_success_total.inc();
                     latency.observe();
                     self.command_sender
@@ -66,48 +71,33 @@ impl GpuProveTask {
                         .await
                         .ok();
                 }
-                Err(err) => match self.on_failure {
-                    OnFailure::Exit => {
-                        METRICS.proof_failure_total.inc();
-                        latency.observe();
-                        sentry::with_scope(
-                            |scope| {
-                                scope.set_level(Some(sentry::Level::Error));
-                                scope.set_tag("mode", "gpu_prove");
-                                scope.set_tag("block_number", block_number.to_string());
-                            },
-                            || {
-                                sentry_anyhow::capture_anyhow(&err);
-                            },
-                        );
-                        return Err(err).with_context(|| {
-                            format!("Failed to generate proof for the block {block_number}")
-                        });
-                    }
-                    OnFailure::Continue => {
-                        METRICS.proof_failure_total.inc();
-                        latency.observe();
-                        sentry::with_scope(
-                            |scope| {
-                                scope.set_level(Some(sentry::Level::Error));
-                                scope.set_tag("mode", "gpu_prove");
-                                scope.set_tag("block_number", block_number.to_string());
-                            },
-                            || {
-                                sentry_anyhow::capture_anyhow(&err);
-                            },
-                        );
-                        tracing::error!(
-                            "Failed to generate proof for the block {block_number}: {err}"
-                        );
-                        // continue;
-                        // TODO: `continue` mode is not supported yet, as some of the panics are not unwind-safe.
-                        // TODO: On failure in the prover itself, we need to re-instantiate the prover.
-                        return Err(err).with_context(|| {
-                            format!("Failed to generate proof for the block {block_number}")
-                        });
+                Err(err) => {
+                    METRICS.proof_failure_total.inc();
+                    latency.observe();
+                    sentry::with_scope(
+                        |scope| {
+                            scope.set_level(Some(sentry::Level::Error));
+                            scope.set_tag("mode", "gpu_prove");
+                            scope.set_tag("block_number", block_number.to_string());
+                        },
+                        || {
+                            sentry_anyhow::capture_anyhow(&err);
+                        },
+                    );
+                    match self.on_failure {
+                        OnFailure::Exit => {
+                            return Err(err).with_context(|| {
+                                format!("Failed to generate proof for the block {block_number}")
+                            });
+                        }
+                        OnFailure::Continue => {
+                            tracing::error!(
+                                "Failed to generate proof for the block {block_number}: {err}"
+                            );
+                            continue;
+                        }
                     }
-                },
+                }
             }
         }
 

ethereum_prover/src/utils.rs

@@ -0,0 +1,14 @@
+pub(crate) fn extract_panic_message(err: tokio::task::JoinError) -> String {
+    if err.is_panic() {
+        let panic = err.into_panic();
+        if let Some(s) = panic.downcast_ref::<String>() {
+            s.clone()
+        } else if let Some(s) = panic.downcast_ref::<&str>() {
+            s.to_string()
+        } else {
+            "Unknown panic".to_string()
+        }
+    } else {
+        "Task was cancelled".to_string()
+    }
+}