error handling when proving fails (#75)

* error handling when proving fails

* error handling for bootloader

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 {

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,
 }

prover-sdk/src/sdk.rs

@@ -43,11 +43,17 @@ impl ProverSDK {
     }
 
     pub async fn prove_cairo0(&self, data: Cairo0ProverInput) -> Result<u64, SdkErrors> {
+        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<u64, SdkErrors> {
+        if !data.layout.is_bootloadable() && data.bootload {
+            return Err(SdkErrors::BootloaderError);
+        }
         self.prove(ProverInput::Cairo(data), self.prover_cairo.clone())
             .await
     }

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 => {

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)?;

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<PathBuf, ProverError>;
 }
 
 impl BootloaderPath for Layout {
-    fn path(&self) -> PathBuf {
+    fn path(&self) -> Result<PathBuf, ProverError> {
         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<Command, ProverError> {
         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");