Skip to content

Commit

Permalink
Remove witness
Browse files Browse the repository at this point in the history
  • Loading branch information
@bkolad
bkolad committed Dec 19, 2023
1 parent c3250c2 commit 1e80b4c
Show file tree
Hide file tree
Showing 2 changed files with 104 additions and 39 deletions.
18 changes: 10 additions & 8 deletions full-node/sov-stf-runner/src/prover_service/parallel/prover.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -118,17 +119,18 @@ where
block_header_hash: <Da::Spec as DaSpec>::SlotHash,
) -> Result<ProofSubmissionStatus, anyhow::Error> {
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
)),
Expand Down
125 changes: 94 additions & 31 deletions full-node/sov-stf-runner/src/prover_service/parallel/prover_manager.rs
View file
Original file line number Diff line number Diff line change
@@ -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<StateRoot, Witness, Da: DaSpec> {
WitnessSubmitted(StateTransitionData<StateRoot, Witness, Da>),
pub(crate) enum ProverStatus {
WitnessSubmitted,
ProvingInProgress,
Proved(Proof),
Err(anyhow::Error),
}

struct ProverState<StateRoot, Witness, Da: DaSpec> {
prover_status: HashMap<Da::SlotHash, ProverStatus<StateRoot, Witness, Da>>,
prover_status: HashMap<Da::SlotHash, ProverStatus>,
witness: HashMap<Da::SlotHash, StateTransitionData<StateRoot, Witness, Da>>,
pending_tasks_count: usize,
}

impl<StateRoot, Witness, Da: DaSpec> ProverState<StateRoot, Witness, Da> {
fn remove(&mut self, hash: &Da::SlotHash) -> Option<ProverStatus<StateRoot, Witness, Da>> {
fn remove(&mut self, hash: &Da::SlotHash) -> Option<ProverStatus> {
self.prover_status.remove(hash)
}

fn set_to_proving(
fn remove_witness(
&mut self,
hash: Da::SlotHash,
) -> Option<ProverStatus<StateRoot, Witness, Da>> {
hash: &Da::SlotHash,
) -> Option<StateTransitionData<StateRoot, Witness, Da>> {
self.witness.remove(hash)
}

fn set_to_proving(&mut self, hash: Da::SlotHash) -> Option<ProverStatus> {
self.prover_status
.insert(hash, ProverStatus::ProvingInProgress)
}
Expand All @@ -33,18 +38,15 @@ impl<StateRoot, Witness, Da: DaSpec> ProverState<StateRoot, Witness, Da> {
&mut self,
hash: Da::SlotHash,
proof: Result<Proof, anyhow::Error>,
) -> Option<ProverStatus<StateRoot, Witness, Da>> {
) -> Option<ProverStatus> {
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<StateRoot, Witness, Da>> {
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 {
Expand All @@ -65,6 +67,7 @@ impl<StateRoot, Witness, Da: DaSpec> ProverState<StateRoot, Witness, Da> {
#[derive(Default)]
struct AggregatedProofInfo<SlotHash> {
height_to_slot_hash: HashMap<u64, SlotHash>,
start_height: u64,
jump: u64,
}

Expand All @@ -81,33 +84,25 @@ impl<StateRoot, Witness, Da: DaSpec> ProverManager<StateRoot, Witness, Da> {
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<ProverStatus<StateRoot, Witness, Da>> {
self.prover_state.remove(hash)
}

pub(crate) fn set_to_proving(
&mut self,
hash: Da::SlotHash,
) -> Option<ProverStatus<StateRoot, Witness, Da>> {
pub(crate) fn set_to_proving(&mut self, hash: Da::SlotHash) -> Option<ProverStatus> {
self.prover_state.set_to_proving(hash)
}

pub(crate) fn set_to_proved(
&mut self,
hash: Da::SlotHash,
proof: Result<Proof, anyhow::Error>,
) -> Option<ProverStatus<StateRoot, Witness, Da>> {
) -> Option<ProverStatus> {
self.prover_state.set_to_proved(hash, proof)
}

Expand All @@ -119,14 +114,21 @@ impl<StateRoot, Witness, Da: DaSpec> ProverManager<StateRoot, Witness, Da> {
self.prover_state.dec_task_count()
}

pub(crate) fn get_witness(
&mut self,
hash: &Da::SlotHash,
) -> &StateTransitionData<StateRoot, Witness, Da> {
self.prover_state.witness.get(hash).unwrap()
}

pub(crate) fn submit_witness(
&mut self,
height: u64,
header_hash: Da::SlotHash,
state_transition_data: StateTransitionData<StateRoot, Witness, Da>,
) -> 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,
Expand All @@ -135,16 +137,77 @@ impl<StateRoot, Witness, Da: DaSpec> ProverManager<StateRoot, Witness, Da> {
// 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<StateRoot, Witness, Da>> {
hash: &Da::SlotHash,
) -> Option<(ProverStatus, StateTransitionData<StateRoot, Witness, Da>)> {
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<ProofSubmissionStatus, anyhow::Error> {
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<StateRoot, SlotHash> {
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 {}

0 comments on commit 1e80b4c

Please sign in to comment.