From 66828924bac821716a68c7e1bd20863f42d372e9 Mon Sep 17 00:00:00 2001 From: Mateusz Chudkowski <120587768+chudkowsky@users.noreply.github.com> Date: Tue, 7 Jan 2025 15:02:46 +0100 Subject: [PATCH] error handling when proving fails (#75) * error handling when proving fails * error handling for bootloader --- bin/cairo-prove/src/main.rs | 3 +++ prover-sdk/src/errors.rs | 4 ++++ prover-sdk/src/sdk.rs | 6 ++++++ prover/src/threadpool/mod.rs | 29 ++++++++++++++++++++++++----- prover/src/threadpool/prove.rs | 3 ++- prover/src/threadpool/run.rs | 32 +++++++++++++++++++------------- 6 files changed, 58 insertions(+), 19 deletions(-) diff --git a/bin/cairo-prove/src/main.rs b/bin/cairo-prove/src/main.rs index 4d1e4ec..1475485 100644 --- a/bin/cairo-prove/src/main.rs +++ b/bin/cairo-prove/src/main.rs @@ -12,6 +12,9 @@ pub async fn main() -> Result<(), ProveErrors> { tracing_subscriber::fmt().init(); let args = Args::parse(); let access_key = ProverAccessKey::from_hex_string(&args.prover_access_key.clone())?; + if !args.layout.is_bootloadable() && args.bootload { + return Err(ProveErrors::Custom("Invalid layout for bootloading, supported layouts for bootloader: recursive, recursive_with_poseidon, starknet, starknet_with_keccak".to_string())); + } let sdk = ProverSDK::new(args.prover_url.clone(), access_key).await?; let job = prove(args.clone(), sdk.clone()).await?; if args.wait { diff --git a/prover-sdk/src/errors.rs b/prover-sdk/src/errors.rs index 561529c..03c6ebe 100644 --- a/prover-sdk/src/errors.rs +++ b/prover-sdk/src/errors.rs @@ -34,4 +34,8 @@ pub enum SdkErrors { VerifyResponseError(String), #[error("Invalid key")] InvalidKey, + #[error( + "Invalid layout for bootload, small and dex layouts are not supported with bootload option" + )] + BootloaderError, } diff --git a/prover-sdk/src/sdk.rs b/prover-sdk/src/sdk.rs index 4db25be..a2c160b 100644 --- a/prover-sdk/src/sdk.rs +++ b/prover-sdk/src/sdk.rs @@ -43,11 +43,17 @@ impl ProverSDK { } pub async fn prove_cairo0(&self, data: Cairo0ProverInput) -> Result { + if !data.layout.is_bootloadable() && data.bootload { + return Err(SdkErrors::BootloaderError); + } self.prove(ProverInput::Cairo0(data), self.prover_cairo0.clone()) .await } pub async fn prove_cairo(&self, data: CairoProverInput) -> Result { + if !data.layout.is_bootloadable() && data.bootload { + return Err(SdkErrors::BootloaderError); + } self.prove(ProverInput::Cairo(data), self.prover_cairo.clone()) .await } diff --git a/prover/src/threadpool/mod.rs b/prover/src/threadpool/mod.rs index 4c2864a..85822c0 100644 --- a/prover/src/threadpool/mod.rs +++ b/prover/src/threadpool/mod.rs @@ -7,7 +7,7 @@ use tokio::{ sync::{broadcast::Sender, mpsc, Mutex}, task::JoinHandle, }; -use tracing::trace; +use tracing::{error, trace}; pub mod prove; pub mod run; @@ -139,19 +139,38 @@ impl Worker { if let Err(e) = prove( job_id, - job_store, + job_store.clone(), dir, program_input, - sse_tx, + sse_tx.clone(), n_queries, pow_bits, bootload, ) .await { - eprintln!("Worker {id} encountered an error: {:?}", e); + job_store + .update_job_status( + job_id, + common::models::JobStatus::Failed, + Some(e.to_string()), + ) + .await; + let sender = sse_tx.clone(); + let sender = sender.lock().await; + if sender.receiver_count() > 0 { + sender + .send( + serde_json::to_string(&( + common::models::JobStatus::Failed, + job_id, + )) + .unwrap(), + ) + .unwrap(); + } + error!("Worker {id} encountered an error: {:?}", e); } - trace!("Worker {id} finished the job."); } None => { diff --git a/prover/src/threadpool/prove.rs b/prover/src/threadpool/prove.rs index 93cd156..945a84d 100644 --- a/prover/src/threadpool/prove.rs +++ b/prover/src/threadpool/prove.rs @@ -16,6 +16,7 @@ use tempfile::TempDir; use tokio::process::Command; use tokio::sync::broadcast::Sender; use tokio::sync::Mutex; +use tracing::trace; #[allow(clippy::too_many_arguments)] pub async fn prove( @@ -39,7 +40,7 @@ pub async fn prove( .await?; Template::generate_from_public_input_file(&paths.public_input_file, n_queries, pow_bits)? .save_to_file(&paths.params_file)?; - + trace!("Running prover"); let prove_status = paths.prove_command().spawn()?.wait().await?; let result = fs::read_to_string(&paths.proof_path)?; let proof: Value = serde_json::from_str(&result)?; diff --git a/prover/src/threadpool/run.rs b/prover/src/threadpool/run.rs index a9a2494..7cf7531 100644 --- a/prover/src/threadpool/run.rs +++ b/prover/src/threadpool/run.rs @@ -13,17 +13,19 @@ pub enum CairoVersionedInput { Cairo0(Cairo0ProverInput), } pub trait BootloaderPath { - fn path(&self) -> PathBuf; + fn path(&self) -> Result; } impl BootloaderPath for Layout { - fn path(&self) -> PathBuf { + fn path(&self) -> Result { match self { - Layout::Recursive => "bootloaders/recursive.json".into(), - Layout::RecursiveWithPoseidon => "bootloaders/recursive_with_poseidon.json".into(), - Layout::Starknet => "bootloaders/starknet.json".into(), - Layout::StarknetWithKeccak => "bootloaders/starknet_with_keccak.json".into(), - Layout::Dex | Layout::Small => panic!("Invalid layout"), //TODO: Handle this case + Layout::Recursive => Ok("bootloaders/recursive.json".into()), + Layout::RecursiveWithPoseidon => Ok("bootloaders/recursive_with_poseidon.json".into()), + Layout::Starknet => Ok("bootloaders/starknet.json".into()), + Layout::StarknetWithKeccak => Ok("bootloaders/starknet_with_keccak.json".into()), + Layout::Dex | Layout::Small => { + Err(ProverError::CustomError("Invalid layout".to_string())) + } } } } @@ -64,7 +66,7 @@ impl CairoVersionedInput { let pie_file_str = paths.pie_output.to_str().unwrap(); let program_input_file_str = paths.program_input_path.to_str().unwrap(); create_template(pie_file_str, program_input_file_str)?; - let command = paths.cairo0_run_command(input.layout.clone(), bootloader); + let command = paths.cairo0_run_command(input.layout.clone(), bootloader)?; command_run(command).await } else { let command = paths.cairo1_run_command(&input.layout.to_string()); @@ -79,10 +81,10 @@ impl CairoVersionedInput { let pie_file_str = paths.pie_output.to_str().unwrap(); let program_input_file_str = paths.program_input_path.to_str().unwrap(); create_template(pie_file_str, program_input_file_str)?; - let command = paths.cairo0_run_command(input.layout.clone(), bootloader); + let command = paths.cairo0_run_command(input.layout.clone(), bootloader)?; command_run(command).await } else { - let command = paths.cairo0_run_command(input.layout.clone(), bootloader); + let command = paths.cairo0_run_command(input.layout.clone(), bootloader)?; command_run(command).await } } @@ -120,9 +122,13 @@ impl RunPaths<'_> { .arg(self.program); command } - pub fn cairo0_run_command(&self, layout: Layout, bootloader: bool) -> Command { + pub fn cairo0_run_command( + &self, + layout: Layout, + bootloader: bool, + ) -> Result { let program = if bootloader && layout.is_bootloadable() { - layout.path() + layout.path()? } else { self.program.to_path_buf() }; @@ -145,7 +151,7 @@ impl RunPaths<'_> { .arg(self.program_input_path) .arg("--program") .arg(program); - command + Ok(command) } pub fn cairo0_pie_command(&self, layout: &str) -> Command { let mut command = Command::new("python");