Trace gen endpoint

Author: chudkowsky

bin/prover-cli/src/prove.rs

@@ -4,7 +4,7 @@ use std::path::PathBuf;
 use clap::Parser;
 use prover_sdk::{
     access_key::ProverAccessKey, sdk::ProverSDK, Cairo0ProverInput, CairoCompiledProgram,
-    CairoProverInput, JobResult, Layout, ProverResult,
+    CairoProverInput, JobResult, Layout, ProverResult, RunMode,
 };
 use serde_json::Value;
 use url::Url;
@@ -39,15 +39,15 @@ pub struct Prove {
     pub n_queries: Option<u32>,
     #[arg(long, env)]
     pub pow_bits: Option<u32>,
-    #[arg(long, env, default_value = "false")]
-    pub bootload: bool,
+    #[arg(long, env, default_value = "trace")]
+    pub run_mode: RunMode,
     #[arg(long, env, default_value = "false")]
     pub full_output: bool,
 }
 impl Prove {
     pub async fn run(self) {
         let access_key = ProverAccessKey::from_hex_string(&self.prover_access_key.clone()).unwrap();
-        if self.bootload {
+        if matches!(self.run_mode, RunMode::Bootload) {
             assert!(self.layout.is_bootloadable(),"Invalid layout for bootloading, supported layouts for bootloader: recursive, recursive_with_poseidon, starknet, starknet_with_keccak")
         }
         let sdk = ProverSDK::new(self.prover_url.clone(), access_key)
@@ -94,7 +94,7 @@ pub async fn prove(args: Prove, sdk: ProverSDK) -> u64 {
                 program_input,
                 pow_bits: args.pow_bits,
                 n_queries: args.n_queries,
-                bootload: args.bootload,
+                run_mode: args.run_mode,
             };
             sdk.prove_cairo0(data).await.unwrap()
         }
@@ -109,7 +109,7 @@ pub async fn prove(args: Prove, sdk: ProverSDK) -> u64 {
                 program_input: input,
                 pow_bits: args.pow_bits,
                 n_queries: args.n_queries,
-                bootload: args.bootload,
+                run_mode: args.run_mode,
             };
             sdk.prove_cairo(data).await.unwrap()
         }

bin/prover-cli/src/run.rs

@@ -3,7 +3,7 @@ use std::path::PathBuf;
 use clap::Parser;
 use prover_sdk::{
     access_key::ProverAccessKey, sdk::ProverSDK, Cairo0ProverInput, CairoCompiledProgram,
-    CairoProverInput, JobResult, Layout, RunResult,
+    CairoProverInput, JobResult, Layout, RunMode, RunResult,
 };
 use serde_json::Value;
 use tokio::fs;
@@ -35,8 +35,8 @@ pub struct CairoRunner {
     pub sse: bool,
     #[arg(long, env)]
     pub proof_dir: PathBuf,
-    #[arg(long, env, default_value = "false")]
-    pub bootload: bool,
+    #[arg(long, env, default_value = "trace")]
+    pub run_mode: RunMode,
 }
 impl CairoRunner {
     pub async fn run(self) {
@@ -55,33 +55,44 @@ impl CairoRunner {
             if !(self.proof_dir).exists() {
                 fs::create_dir(&self.proof_dir).await.unwrap();
             }
-            let path: std::path::PathBuf = self.proof_dir;
+            match result {
+                RunResult::Pie(pie) => {
+                    fs::write("pie.zip", pie).await.unwrap();
+                }
+                RunResult::Trace(trace_files) => {
+                    let path: std::path::PathBuf = self.proof_dir;
 
-            let trace_path = path.join("trace");
-            let memory_path = path.join("memory");
-            fs::write(trace_path.clone(), result.trace).await.unwrap();
-            fs::write(memory_path.clone(), result.memory).await.unwrap();
+                    let trace_path = path.join("trace");
+                    let memory_path = path.join("memory");
+                    fs::write(trace_path.clone(), trace_files.trace)
+                        .await
+                        .unwrap();
+                    fs::write(memory_path.clone(), trace_files.memory)
+                        .await
+                        .unwrap();
 
-            let canonical_trace_path = fs::canonicalize(&trace_path).await.unwrap();
-            let canonical_memory_path = fs::canonicalize(&memory_path).await.unwrap();
+                    let canonical_trace_path = fs::canonicalize(&trace_path).await.unwrap();
+                    let canonical_memory_path = fs::canonicalize(&memory_path).await.unwrap();
 
-            let public_inputs_path = path.join("public_inputs");
-            let private_inputs_path = path.join("private_inputs");
+                    let public_inputs_path = path.join("public_inputs");
+                    let private_inputs_path = path.join("private_inputs");
 
-            fs::write(public_inputs_path, result.public_input)
-                .await
-                .unwrap();
-            fs::write(private_inputs_path, result.private_input)
-                .await
-                .unwrap();
+                    fs::write(public_inputs_path, trace_files.public_input)
+                        .await
+                        .unwrap();
+                    fs::write(private_inputs_path, trace_files.private_input)
+                        .await
+                        .unwrap();
 
-            println!(
-                "Proof files are saved, change the paths in private input for those:
-            trace: {}
-            memory: {}",
-                canonical_trace_path.to_str().unwrap(),
-                canonical_memory_path.to_str().unwrap()
-            );
+                    println!(
+                        "Proof files are saved, change the paths in private input for those:
+                    trace: {}
+                    memory: {}",
+                        canonical_trace_path.to_str().unwrap(),
+                        canonical_memory_path.to_str().unwrap()
+                    );
+                }
+            }
         }
     }
 }
@@ -98,7 +109,7 @@ pub async fn cairo_runner(args: CairoRunner, sdk: ProverSDK) -> u64 {
                 program_input,
                 pow_bits: None,
                 n_queries: None,
-                bootload: args.bootload,
+                run_mode: args.run_mode,
             };
             sdk.run_cairo0(data).await.unwrap()
         }
@@ -113,7 +124,7 @@ pub async fn cairo_runner(args: CairoRunner, sdk: ProverSDK) -> u64 {
                 program_input: input,
                 pow_bits: None,
                 n_queries: None,
-                bootload: args.bootload,
+                run_mode: args.run_mode,
             };
             sdk.run_cairo(data).await.unwrap()
         }

common/src/models.rs

@@ -29,12 +29,18 @@ pub struct ProverResult {
     pub program_output_hash: Felt,
 }
 #[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct RunResult {
+pub struct TraceFiles {
     pub private_input: String,
     pub public_input: String,
     pub memory: Vec<u8>,
     pub trace: Vec<u8>,
-    pub pie: Option<String>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+#[serde(untagged)]
+pub enum RunResult {
+    Pie(Vec<u8>),
+    Trace(TraceFiles),
 }
 #[derive(Serialize, Deserialize)]
 #[serde(untagged)]

common/src/prover_input/cairo.rs

@@ -1,14 +1,16 @@
 use serde::{Deserialize, Serialize};
 use starknet_types_core::felt::Felt;
 
+use super::RunMode;
+
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 pub struct CairoProverInput {
     pub program: CairoCompiledProgram,
     pub program_input: Vec<Felt>,
     pub layout: super::Layout,
     pub n_queries: Option<u32>,
     pub pow_bits: Option<u32>,
-    pub bootload: bool,
+    pub run_mode: RunMode,
 }
 
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]

common/src/prover_input/cairo0.rs

@@ -1,4 +1,4 @@
-use super::Layout;
+use super::{Layout, RunMode};
 use serde::{Deserialize, Serialize};
 
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
@@ -8,7 +8,7 @@ pub struct Cairo0ProverInput {
     pub layout: Layout,
     pub n_queries: Option<u32>,
     pub pow_bits: Option<u32>,
-    pub bootload: bool,
+    pub run_mode: RunMode,
 }
 
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]

common/src/prover_input/mod.rs

@@ -15,10 +15,10 @@ pub enum ProverInput {
     Cairo(CairoProverInput),
 }
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ValueEnum)]
-pub enum LayoutBridgeOrBootload {
+pub enum RunMode {
     Bootload,
-    LayoutBridge,
-    None,
+    Pie,
+    Trace,
 }
 
 impl ProverInput {

k8s/prover.yaml

@@ -17,7 +17,7 @@ spec:
       containers:
       - name: http-prover
         # image: docker.io/chudas/http-prover:v0.4
-        image: docker.io/chudas/http-prover:v0.5.2
+        image: docker.io/chudas/http-prover:v0.5.5
         ports:
           - containerPort: 3000
         envFrom:

prover-sdk/src/sdk.rs

@@ -46,15 +46,19 @@ impl ProverSDK {
     }
 
     pub async fn prove_cairo0(&self, data: Cairo0ProverInput) -> Result<u64, SdkErrors> {
-        if !data.layout.is_bootloadable() && data.bootload {
+        if !data.layout.is_bootloadable()
+            && matches!(data.run_mode, common::prover_input::RunMode::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 {
+        if !data.layout.is_bootloadable()
+            && matches!(data.run_mode, common::prover_input::RunMode::Bootload)
+        {
             return Err(SdkErrors::BootloaderError);
         }
         self.prove(ProverInput::Cairo(data), self.prover_cairo.clone())
@@ -79,15 +83,19 @@ impl ProverSDK {
         Ok(job.job_id)
     }
     pub async fn run_cairo0(&self, data: Cairo0ProverInput) -> Result<u64, SdkErrors> {
-        if !data.layout.is_bootloadable() && data.bootload {
+        if !data.layout.is_bootloadable()
+            && matches!(data.run_mode, common::prover_input::RunMode::Bootload)
+        {
             return Err(SdkErrors::BootloaderError);
         }
         self.run(ProverInput::Cairo0(data), self.run_cairo0.clone())
             .await
     }
 
     pub async fn run_cairo(&self, data: CairoProverInput) -> Result<u64, SdkErrors> {
-        if !data.layout.is_bootloadable() && data.bootload {
+        if !data.layout.is_bootloadable()
+            && matches!(data.run_mode, common::prover_input::RunMode::Bootload)
+        {
             return Err(SdkErrors::BootloaderError);
         }
         self.prove(ProverInput::Cairo(data), self.run_cairo.clone())

prover-sdk/tests/prove_test.rs

@@ -29,7 +29,7 @@ async fn test_cairo_prove_bootloader() {
         program_input,
         n_queries: Some(16),
         pow_bits: Some(20),
-        bootload: true,
+        run_mode: RunMode::Bootload,
     };
     let job = sdk.prove_cairo(data).await.unwrap();
     let result = fetch_job(sdk.clone(), job).await;
@@ -75,7 +75,7 @@ async fn test_cairo_prove() {
         program_input,
         n_queries: Some(16),
         pow_bits: Some(20),
-        bootload: false,
+        run_mode: RunMode::Trace,
     };
     let job = sdk.prove_cairo(data).await.unwrap();
     let result = fetch_job(sdk.clone(), job).await;
@@ -116,7 +116,7 @@ async fn test_cairo0_prove_bootloader() {
         program_input,
         n_queries: Some(16),
         pow_bits: Some(20),
-        bootload: true,
+        run_mode: RunMode::Bootload,
     };
     let job = sdk.prove_cairo0(data).await.unwrap();
     let result = fetch_job(sdk.clone(), job).await;
@@ -157,7 +157,7 @@ async fn test_cairo0_prove() {
         program_input,
         n_queries: Some(16),
         pow_bits: Some(20),
-        bootload: false,
+        run_mode: RunMode::Trace,
     };
     let job = sdk.prove_cairo0(data).await.unwrap();
     let result = fetch_job(sdk.clone(), job).await;
@@ -203,7 +203,7 @@ async fn test_cairo_multi_prove() {
         program_input,
         n_queries: Some(16),
         pow_bits: Some(20),
-        bootload: false,
+        run_mode: RunMode::Trace,
     };
     let job1 = sdk.prove_cairo(data.clone()).await.unwrap();
     let job2 = sdk.prove_cairo(data.clone()).await.unwrap();

prover-sdk/tests/run_test.rs

@@ -4,6 +4,7 @@ use common::models::JobResult;
 use helpers::fetch_job;
 use prover_sdk::{
     access_key::ProverAccessKey, sdk::ProverSDK, CairoCompiledProgram, CairoProverInput, Layout,
+    RunResult,
 };
 use starknet_types_core::felt::Felt;
 use url::Url;
@@ -29,7 +30,7 @@ async fn test_cairo_run() {
         program_input,
         n_queries: Some(16),
         pow_bits: Some(20),
-        bootload: false,
+        run_mode: prover_sdk::RunMode::Trace,
     };
     let job = sdk.run_cairo(data).await.unwrap();
     let result = fetch_job(sdk.clone(), job).await;
@@ -40,18 +41,21 @@ async fn test_cairo_run() {
             panic!("Expected run result, got prove result");
         }
         JobResult::Run(run_result) => {
-            // Validate private input
-            assert!(
-                !run_result.private_input.is_empty(),
-                "Private input is empty"
-            );
-            // Validate public input
-            assert!(!run_result.public_input.is_empty(), "Public input is empty");
+            match run_result {
+                RunResult::Pie(_pie) => {
+                    panic!("Expected run result, got pie");
+                }
+                RunResult::Trace(trace) => {
+                    assert!(!trace.private_input.is_empty(), "Private input is empty");
+                    // Validate public input
+                    assert!(!trace.public_input.is_empty(), "Public input is empty");
 
-            // Validate memory
-            assert!(!run_result.memory.is_empty(), "Memory is empty");
-            // Validate trace
-            assert!(!run_result.trace.is_empty(), "Trace is empty");
+                    // Validate memory
+                    assert!(!trace.memory.is_empty(), "Memory is empty");
+                    // Validate trace
+                    assert!(!trace.trace.is_empty(), "Trace is empty");
+                }
+            }
         }
     }
 }