From 1e80b4c99a0434a4159ec032c56bb09746820a7f Mon Sep 17 00:00:00 2001 From: bkolad Date: Tue, 19 Dec 2023 23:25:35 +0100 Subject: [PATCH] Remove witness --- .../src/prover_service/parallel/prover.rs | 18 +-- .../prover_service/parallel/prover_manager.rs | 125 +++++++++++++----- 2 files changed, 104 insertions(+), 39 deletions(-) diff --git a/full-node/sov-stf-runner/src/prover_service/parallel/prover.rs b/full-node/sov-stf-runner/src/prover_service/parallel/prover.rs index 9ecbcc5d8..8dcd8ac63 100644 --- a/full-node/sov-stf-runner/src/prover_service/parallel/prover.rs +++ b/full-node/sov-stf-runner/src/prover_service/parallel/prover.rs @@ -67,21 +67,22 @@ where V::PreState: Send + Sync + 'static, { let prover_manager_clone = self.prover_manager.clone(); - let mut prover_state = self.prover_manager.write().expect("Lock was poisoned"); + let mut prover_manager = self.prover_manager.write().expect("Lock was poisoned"); - let prover_status = prover_state + let (prover_status, state_transition_data) = prover_manager .remove(&block_header_hash) .ok_or_else(|| anyhow::anyhow!("Missing witness for block: {:?}", block_header_hash))?; match prover_status { - ProverStatus::WitnessSubmitted(state_transition_data) => { - let start_prover = prover_state.inc_task_count_if_not_busy(self.num_threads); + ProverStatus::WitnessSubmitted => { + let start_prover = prover_manager.inc_task_count_if_not_busy(self.num_threads); // Initiate a new proving job only if the prover is not busy. if start_prover { - prover_state.set_to_proving(block_header_hash.clone()); vm.add_hint(state_transition_data); + prover_manager.set_to_proving(block_header_hash.clone()); + self.pool.spawn(move || { tracing::info_span!("guest_execution").in_scope(|| { let proof = make_proof(vm, config, zk_storage); @@ -118,17 +119,18 @@ where block_header_hash: ::SlotHash, ) -> Result { let mut prover_manager = self.prover_manager.write().unwrap(); - let status = prover_manager.get_prover_status(block_header_hash.clone()); + let status = prover_manager.get_prover_status(&block_header_hash.clone()); match status { Some(ProverStatus::ProvingInProgress) => { Ok(ProofSubmissionStatus::ProofGenerationInProgress) } Some(ProverStatus::Proved(_)) => { - prover_manager.remove(&block_header_hash); + //TODO + //prover_manager.remove(&block_header_hash); Ok(ProofSubmissionStatus::Success) } - Some(ProverStatus::WitnessSubmitted(_)) => Err(anyhow::anyhow!( + Some(ProverStatus::WitnessSubmitted) => Err(anyhow::anyhow!( "Witness for {:?} was submitted, but the proof generation is not triggered.", block_header_hash )), diff --git a/full-node/sov-stf-runner/src/prover_service/parallel/prover_manager.rs b/full-node/sov-stf-runner/src/prover_service/parallel/prover_manager.rs index 51ed92e22..62c7f04ca 100644 --- a/full-node/sov-stf-runner/src/prover_service/parallel/prover_manager.rs +++ b/full-node/sov-stf-runner/src/prover_service/parallel/prover_manager.rs @@ -1,30 +1,35 @@ -use crate::{StateTransitionData, WitnessSubmissionStatus}; +use crate::{ProofSubmissionStatus, StateTransitionData, WitnessSubmissionStatus}; use sov_rollup_interface::da::DaSpec; use sov_rollup_interface::zk::Proof; use std::collections::hash_map::Entry; use std::collections::HashMap; -pub(crate) enum ProverStatus { - WitnessSubmitted(StateTransitionData), +pub(crate) enum ProverStatus { + WitnessSubmitted, ProvingInProgress, Proved(Proof), Err(anyhow::Error), } struct ProverState { - prover_status: HashMap>, + prover_status: HashMap, + witness: HashMap>, pending_tasks_count: usize, } impl ProverState { - fn remove(&mut self, hash: &Da::SlotHash) -> Option> { + fn remove(&mut self, hash: &Da::SlotHash) -> Option { self.prover_status.remove(hash) } - fn set_to_proving( + fn remove_witness( &mut self, - hash: Da::SlotHash, - ) -> Option> { + hash: &Da::SlotHash, + ) -> Option> { + self.witness.remove(hash) + } + + fn set_to_proving(&mut self, hash: Da::SlotHash) -> Option { self.prover_status .insert(hash, ProverStatus::ProvingInProgress) } @@ -33,18 +38,15 @@ impl ProverState { &mut self, hash: Da::SlotHash, proof: Result, - ) -> Option> { + ) -> Option { match proof { Ok(p) => self.prover_status.insert(hash, ProverStatus::Proved(p)), Err(e) => self.prover_status.insert(hash, ProverStatus::Err(e)), } } - fn get_prover_status( - &self, - hash: Da::SlotHash, - ) -> Option<&ProverStatus> { - self.prover_status.get(&hash) + fn get_prover_status(&self, hash: &Da::SlotHash) -> Option<&ProverStatus> { + self.prover_status.get(hash) } fn inc_task_count_if_not_busy(&mut self, num_threads: usize) -> bool { @@ -65,6 +67,7 @@ impl ProverState { #[derive(Default)] struct AggregatedProofInfo { height_to_slot_hash: HashMap, + start_height: u64, jump: u64, } @@ -81,25 +84,17 @@ impl ProverManager { prover_state: ProverState { prover_status: Default::default(), pending_tasks_count: Default::default(), + witness: Default::default(), }, aggregated_proof_info: AggregatedProofInfo { height_to_slot_hash: Default::default(), + start_height: 0, jump, }, } } - pub(crate) fn remove( - &mut self, - hash: &Da::SlotHash, - ) -> Option> { - self.prover_state.remove(hash) - } - - pub(crate) fn set_to_proving( - &mut self, - hash: Da::SlotHash, - ) -> Option> { + pub(crate) fn set_to_proving(&mut self, hash: Da::SlotHash) -> Option { self.prover_state.set_to_proving(hash) } @@ -107,7 +102,7 @@ impl ProverManager { &mut self, hash: Da::SlotHash, proof: Result, - ) -> Option> { + ) -> Option { self.prover_state.set_to_proved(hash, proof) } @@ -119,6 +114,13 @@ impl ProverManager { self.prover_state.dec_task_count() } + pub(crate) fn get_witness( + &mut self, + hash: &Da::SlotHash, + ) -> &StateTransitionData { + self.prover_state.witness.get(hash).unwrap() + } + pub(crate) fn submit_witness( &mut self, height: u64, @@ -126,7 +128,7 @@ impl ProverManager { state_transition_data: StateTransitionData, ) -> WitnessSubmissionStatus { let entry = self.prover_state.prover_status.entry(header_hash.clone()); - let data = ProverStatus::WitnessSubmitted(state_transition_data); + let data = ProverStatus::WitnessSubmitted; match entry { Entry::Occupied(_) => WitnessSubmissionStatus::WitnessExist, @@ -135,16 +137,77 @@ impl ProverManager { // TODO assert first insertion self.aggregated_proof_info .height_to_slot_hash - .insert(height, header_hash); + .insert(height, header_hash.clone()); + + self.prover_state + .witness + .insert(header_hash, state_transition_data); + WitnessSubmissionStatus::SubmittedForProving } } } - pub(crate) fn get_prover_status( + // TODO change name + pub(crate) fn remove( &mut self, - hash: Da::SlotHash, - ) -> Option<&ProverStatus> { + hash: &Da::SlotHash, + ) -> Option<(ProverStatus, StateTransitionData)> { + let status = self.prover_state.remove(hash)?; + let witness = self.prover_state.remove_witness(hash)?; + Some((status, witness)) + } + + pub(crate) fn get_prover_status(&mut self, hash: &Da::SlotHash) -> Option<&ProverStatus> { self.prover_state.get_prover_status(hash) } + + fn get_aggregated_proof(&mut self) -> Result { + let jump = self.aggregated_proof_info.jump; + let start_height = self.aggregated_proof_info.start_height; + + let mut proofs_data = Vec::default(); + + for height in start_height..start_height + jump { + let hash = self + .aggregated_proof_info + .height_to_slot_hash + .get(&height) + .unwrap(); + + let state = self.prover_state.get_prover_status(hash).unwrap(); + match state { + ProverStatus::WitnessSubmitted => { + return Err(anyhow::anyhow!( + "Witness for {:?} was submitted, but the proof generation is not triggered.", + hash + )) + } + ProverStatus::ProvingInProgress => { + return Ok(ProofSubmissionStatus::ProofGenerationInProgress) + } + ProverStatus::Proved(proof) => proofs_data.push(proof), + ProverStatus::Err(e) => return Err(anyhow::anyhow!(e.to_string())), + } + } + + todo!() + } +} + +struct AggregatedProofWitness { + proof: Proof, + pre_state: StateRoot, + post_state_root: StateRoot, + da_block_hash: SlotHash, + height: u64, +} + +struct AggregatedProofPublicInput { + initial_state: u64, + final_state_root: u64, + initial_height: u64, + final_height: u64, } + +struct AggrgatedProof {}