Skip to content

feat: Use async bonsai client #2988

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rakanalh wants to merge 1 commit into
base: nightly
Choose a base branch
from
Open

feat: Use async bonsai client #2988

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,7 @@ tower = { version = "0.4.13", features = ["full"] }
url = { version = "2.5.4" }

# Risc0 dependencies
bonsai-sdk = { version = "1.4.1" }
bonsai-sdk = { version = "1.4.1", features = ["non_blocking"] }
boundless-market = { version = "1.0.0", default-features = false }
risc0-ethereum-contracts = { version = "3.0.1" }
risc0-binfmt = "3.0.2"
Expand Down
16 changes: 11 additions & 5 deletions crates/risc0/src/host/bonsai.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use std::time::Duration;

use anyhow::{anyhow, Context};
use bonsai_sdk::blocking::{Client, SessionId, SnarkId};
use bonsai_sdk::non_blocking::{Client, SessionId, SnarkId};
use bonsai_sdk::responses::SessionStats;
use metrics::gauge;
use risc0_zkvm::{compute_image_id, AssumptionReceipt, Digest, InnerAssumptionReceipt, Receipt};
Expand Down Expand Up @@ -34,7 +34,7 @@ impl BonsaiProver {
Self { client, ledger_db }
}

pub fn prove(
pub async fn prove(
&self,
job_id: Uuid,
elf: Vec<u8>,
Expand All @@ -47,12 +47,14 @@ impl BonsaiProver {
let image_id_hex = hex::encode(image_id);
self.client
.upload_img(&image_id_hex, elf)
.await
.context("Failed to upload img")?;

// Upload input
let input_id = self
.client
.upload_input(input)
.await
.context("Failed to upload input")?;

// Upload assumptions
Expand All @@ -63,6 +65,7 @@ impl BonsaiProver {
let receipt_id = self
.client
.upload_receipt(serialized_receipt)
.await
.context("Failed to upload receipt")?;
receipt_ids.push(receipt_id);
}
Expand All @@ -71,6 +74,7 @@ impl BonsaiProver {
let session = self
.client
.create_session(image_id_hex, input_id, receipt_ids, false)
.await
.context("Failed to create session")?;
info!(
"Started bonsai proving session, job_id={} session_id={}",
Expand Down Expand Up @@ -163,7 +167,7 @@ impl BonsaiProver {
return Ok(succinct_receipt);
}

let snark_session = self.client.create_snark(session.uuid.clone())?;
let snark_session = self.client.create_snark(session.uuid.clone()).await?;

let db_session = BonsaiSession {
kind: BonsaiSessionKind::SnarkSession(session.uuid, snark_session.uuid.clone()),
Expand All @@ -188,7 +192,7 @@ impl BonsaiProver {
) -> anyhow::Result<(Receipt, SessionStats)> {
let polling_interval = Duration::from_secs(1);
loop {
let res = session.status(&self.client)?;
let res = session.status(&self.client).await?;
match res.status.as_str() {
"RUNNING" => tokio::time::sleep(polling_interval).await,
"SUCCEEDED" => {
Expand All @@ -198,6 +202,7 @@ impl BonsaiProver {
let receipt_buf = self
.client
.download(&receipt_url)
.await
.context("Failed to download stark receipt")?;
let receipt: Receipt = bincode::deserialize(&receipt_buf)
.expect("Receipt deserialization cannot fail");
Expand All @@ -219,7 +224,7 @@ impl BonsaiProver {
async fn wait_snark_receipt(&self, session: &SnarkId) -> anyhow::Result<Receipt> {
let polling_interval = Duration::from_secs(1);
loop {
let res = session.status(&self.client)?;
let res = session.status(&self.client).await?;
match res.status.as_str() {
"RUNNING" => tokio::time::sleep(polling_interval).await,
"SUCCEEDED" => {
Expand All @@ -230,6 +235,7 @@ impl BonsaiProver {
let receipt_buf = self
.client
.download(&receipt_url)
.await
.context("Failed to download snark receipt")?;
let receipt: Receipt = bincode::deserialize(&receipt_buf)
.expect("Receipt deserialization cannot fail");
Expand Down
4 changes: 3 additions & 1 deletion crates/risc0/src/host/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -114,7 +114,9 @@ impl ZkvmHost for Risc0Host {
with_prove,
"Bonsai prover must always be run with prove set to true"
);
bonsai.prove(job_id, elf, input, assumptions, receipt_type)
bonsai
.prove(job_id, elf, input, assumptions, receipt_type)
.await
}
Prover::Boundless(boundless) => {
assert!(
Expand Down