fix: axiom intergration (#238)

* fix axiom

* fmt

Author: lightsing

crates/integration/src/axiom.rs

@@ -8,6 +8,7 @@ use chrono::DateTime;
 use openvm_sdk::commit::CommitBytes;
 use openvm_sdk::types::{EvmProof, VersionedVmStarkProof};
 use scroll_zkvm_types::ProvingTask as UniversalProvingTask;
+use scroll_zkvm_types::axiom::AxiomProgram;
 use scroll_zkvm_types::proof::ProofEnum;
 use scroll_zkvm_types::types_agg::ProgramCommitment;
 use scroll_zkvm_types::utils::serialize_vk;
@@ -23,18 +24,18 @@ struct TracingProgressCallback;
 
 impl AxiomProver {
     /// Create a new client
-    pub fn from_env(name: String, config_id: String, program_id: String) -> Self {
+    pub fn from_env(name: String, program: &AxiomProgram) -> Self {
         let api_key = env::var("AXIOM_API_KEY").expect("AXIOM_API_KEY env var is required");
         let config = AxiomConfig {
             api_key: Some(api_key),
-            config_id: Some(config_id),
+            config_id: Some(program.config_id().to_string()),
             ..Default::default()
         };
         let sdk = AxiomSdk::new(config).with_callback(TracingProgressCallback);
         Self {
             name,
             sdk,
-            program_id,
+            program_id: program.program_id().to_string(),
         }
     }
 
@@ -72,7 +73,7 @@ impl TaskProver for AxiomProver {
             program_id: Some(self.program_id.clone()),
             input: Some(Input::Value(input)),
             proof_type: Some(proof_type),
-            num_gpus: None,
+            num_gpus: Some(16),
             priority: None,
         })?;
 

crates/integration/src/lib.rs

@@ -7,6 +7,7 @@ use scroll_zkvm_prover::{
     setup::{read_app_config, read_app_exe},
     utils::{read_json, vm::ExecutionResult, write_json},
 };
+use scroll_zkvm_types::axiom::AxiomProgram;
 use scroll_zkvm_types::{
     ProvingTask as UniversalProvingTask,
     proof::{EvmProof, ProofEnum, StarkProof},
@@ -99,9 +100,9 @@ pub static PROGRAM_COMMITMENTS: LazyLock<HashMap<String, ProgramCommitment>> =
         commitments
     });
 
-pub static AXIOM_PROGRAM_IDS: LazyLock<HashMap<String, String>> = LazyLock::new(|| {
+pub static AXIOM_PROGRAM_IDS: LazyLock<HashMap<String, AxiomProgram>> = LazyLock::new(|| {
     let axiom_program_ids = ASSET_BASE_DIR.join("axiom_program_ids.json");
-    let mut program_ids: HashMap<String, String> =
+    let mut program_ids: HashMap<String, AxiomProgram> =
         read_json(&axiom_program_ids).expect("failed to read axiom program ids");
     program_ids.shrink_to_fit();
     eprintln!("AXIOM_PROGRAM_IDS = {program_ids:#?}");
@@ -223,15 +224,10 @@ pub trait ProverTester {
         let mut prover = Self::load_prover(false)?;
         let vk = prover.get_app_commitment();
         let vk = hex::encode(serialize_vk::serialize(&vk));
-        let program_id = AXIOM_PROGRAM_IDS
+        let program = AXIOM_PROGRAM_IDS
             .get(&vk)
-            .ok_or_else(|| eyre::eyre!("missing axiom program id for {}: {}", Self::NAME, vk))?
-            .to_string();
-        let prover = AxiomProver::from_env(
-            Self::NAME.to_string(),
-            scroll_zkvm_types::axiom::get_config_id(Self::NAME).to_string(),
-            program_id,
-        );
+            .ok_or_else(|| eyre::eyre!("missing axiom program id for {}: {}", Self::NAME, vk))?;
+        let prover = AxiomProver::from_env(Self::NAME.to_string(), program);
         Ok(prover)
     }