risc0_zkvm/host/client/prove/
bonsai.rs

1// Copyright 2025 RISC Zero, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::time::Duration;
16
17use anyhow::{anyhow, bail, ensure, Context, Result};
18use bonsai_sdk::blocking::Client;
19
20use super::Prover;
21use crate::{
22    compute_image_id, is_dev_mode, AssumptionReceipt, ExecutorEnv, InnerAssumptionReceipt,
23    InnerReceipt, ProveInfo, ProverOpts, Receipt, ReceiptKind, SessionStats, VerifierContext,
24    VERSION,
25};
26
27/// An implementation of a [Prover] that runs proof workloads via Bonsai.
28///
29/// Requires `BONSAI_API_URL` and `BONSAI_API_KEY` environment variables to
30/// submit proving sessions to Bonsai.
31pub struct BonsaiProver {
32    name: String,
33}
34
35impl BonsaiProver {
36    /// Construct a [BonsaiProver].
37    pub fn new(name: &str) -> Self {
38        Self {
39            name: name.to_string(),
40        }
41    }
42}
43
44// Only proven assumptions that are succinct are supported by Bonsai.
45fn get_inner_assumption_receipt(assumption: &AssumptionReceipt) -> Result<&InnerAssumptionReceipt> {
46    match assumption {
47        AssumptionReceipt::Proven(receipt) => {
48            if !matches!(receipt, InnerAssumptionReceipt::Succinct(_)) {
49                bail!(
50                    "Bonsai only supports succinct assumption receipts. \
51                    Use `ProverOpts::succinct()` when proving any assumptions."
52                );
53            };
54            Ok(receipt)
55        }
56        AssumptionReceipt::Unresolved(_) => {
57            bail!("only proven assumptions can be uploaded to Bonsai.")
58        }
59    }
60}
61
62impl Prover for BonsaiProver {
63    fn get_name(&self) -> String {
64        self.name.clone()
65    }
66
67    fn prove_with_ctx(
68        &self,
69        env: ExecutorEnv<'_>,
70        ctx: &VerifierContext,
71        elf: &[u8],
72        opts: &ProverOpts,
73    ) -> Result<ProveInfo> {
74        let client = Client::from_env(VERSION)?;
75
76        // Compute the ImageID and upload the ELF binary
77        let image_id = compute_image_id(elf)?;
78        let image_id_hex = hex::encode(image_id);
79        client.upload_img(&image_id_hex, elf.to_vec())?;
80
81        // upload input data
82        let input_id = client.upload_input(env.input)?;
83
84        // upload receipts
85        let mut receipts_ids = vec![];
86        for assumption in env.assumptions.borrow().0.iter() {
87            let inner_receipt = get_inner_assumption_receipt(assumption)?;
88            let serialized_receipt = bincode::serialize(inner_receipt)?;
89            let receipt_id = client.upload_receipt(serialized_receipt)?;
90            receipts_ids.push(receipt_id);
91        }
92
93        // While this is the executor, we want to start a session on the bonsai prover.
94        // By doing so, we can return a session ID so that the prover can use it to
95        // retrieve the receipt.
96        let session = client.create_session_with_limit(
97            image_id_hex,
98            input_id,
99            receipts_ids,
100            false,
101            env.session_limit,
102        )?;
103        tracing::debug!("Bonsai proving SessionID: {}", session.uuid);
104
105        // TODO(#1759): Improve upon this polling solution.
106        let polling_interval = if let Ok(ms) = std::env::var("BONSAI_POLL_INTERVAL_MS") {
107            Duration::from_millis(ms.parse().context("invalid bonsai poll interval")?)
108        } else {
109            Duration::from_secs(1)
110        };
111        let succinct_prove_info = loop {
112            // The session has already been started in the executor. Poll bonsai to check if
113            // the proof request succeeded.
114            let res = session.status(&client)?;
115            if res.status == "RUNNING" {
116                std::thread::sleep(polling_interval);
117                continue;
118            }
119            if res.status == "SUCCEEDED" {
120                // Download the receipt, containing the output
121                let receipt_url = res
122                    .receipt_url
123                    .ok_or_else(|| anyhow!("API error, missing receipt on completed session"))?;
124
125                let stats = res
126                    .stats
127                    .context("Missing stats object on Bonsai status res")?;
128                tracing::debug!(
129                    "Bonsai usage: cycles: {} total_cycles: {}",
130                    stats.cycles,
131                    stats.total_cycles
132                );
133
134                let receipt_buf = client.download(&receipt_url)?;
135                let receipt: Receipt = bincode::deserialize(&receipt_buf)?;
136
137                if opts.prove_guest_errors {
138                    receipt.verify_integrity_with_context(ctx)?;
139                } else {
140                    receipt.verify_with_context(ctx, image_id)?;
141                }
142                break ProveInfo {
143                    receipt,
144                    stats: SessionStats {
145                        segments: stats.segments,
146                        total_cycles: stats.total_cycles,
147                        user_cycles: stats.cycles,
148                        // These are currently unavailable from Bonsai
149                        paging_cycles: 0,
150                        reserved_cycles: 0,
151                    },
152                };
153            } else {
154                bail!(
155                    "Bonsai prover workflow [{}] exited: {} err: {}",
156                    session.uuid,
157                    res.status,
158                    res.error_msg
159                        .unwrap_or("Bonsai workflow missing error_msg".into()),
160                );
161            }
162        };
163        match opts.receipt_kind {
164            // If the caller requested a composite or succinct receipt, we are done.
165            ReceiptKind::Composite | ReceiptKind::Succinct => {
166                return Ok(succinct_prove_info);
167            }
168            // If they requested a groth16 receipts, we need to continue.
169            ReceiptKind::Groth16 => {}
170        }
171
172        // Request that Bonsai compress further, to Groth16.
173        let snark_session = client.create_snark(session.uuid)?;
174        let snark_receipt_url = loop {
175            let res = snark_session.status(&client)?;
176            match res.status.as_str() {
177                "RUNNING" => {
178                    std::thread::sleep(polling_interval);
179                    continue;
180                }
181                "SUCCEEDED" => {
182                    break res.output.with_context(|| {
183                        format!(
184                            "Bonsai prover workflow [{}] reported success, but provided no receipt",
185                            snark_session.uuid
186                        )
187                    })?;
188                }
189                _ => {
190                    bail!(
191                        "Bonsai prover workflow [{}] exited: {} err: {}",
192                        snark_session.uuid,
193                        res.status,
194                        res.error_msg
195                            .unwrap_or("Bonsai workflow missing error_msg".into()),
196                    );
197                }
198            }
199        };
200
201        // Assemble the groth16 receipt, and verify it.
202        // TODO(bonsai-alpha#461): If the Groth16 parameters used by Bonsai do not match, this
203        // verification will fail. When Bonsai returned ReceiptMetadata, we'll be able to detect
204        // this error condition and report a better message. Constructing the receipt here will all
205        // the verifier parameters for this version of risc0-zkvm, which may be different than
206        // Bonsai. By verifying the receipt though, we at least know the proving key used on Bonsai
207        // matches the verifying key used here.
208        let receipt_buf = client.download(&snark_receipt_url)?;
209        let groth16_receipt: Receipt = bincode::deserialize(&receipt_buf)?;
210        groth16_receipt
211            .verify_integrity_with_context(ctx)
212            .context("failed to verify Groth16Receipt returned by Bonsai")?;
213
214        // Return the groth16 receipt, with the stats collected earlier.
215        Ok(ProveInfo {
216            receipt: groth16_receipt,
217            stats: succinct_prove_info.stats,
218        })
219    }
220
221    fn compress(&self, opts: &ProverOpts, receipt: &Receipt) -> Result<Receipt> {
222        match (&receipt.inner, opts.receipt_kind) {
223            // Compression is a no-op when the requested kind is at least as large as the current.
224            (InnerReceipt::Composite(_), ReceiptKind::Composite)
225            | (InnerReceipt::Succinct(_), ReceiptKind::Composite | ReceiptKind::Succinct)
226            | (
227                InnerReceipt::Groth16(_),
228                ReceiptKind::Composite | ReceiptKind::Succinct | ReceiptKind::Groth16,
229            ) => Ok(receipt.clone()),
230            // Compression is always a no-op in dev mode
231            (InnerReceipt::Fake { .. }, _) => {
232                ensure!(
233                    is_dev_mode(),
234                    "dev mode must be enabled to compress fake receipts"
235                );
236                Ok(receipt.clone())
237            }
238            // NOTE: Bonsai always returns a SuccinctReceipt, and does not currently support
239            // compression of existing receipts uploaded by clients.
240            (_, ReceiptKind::Succinct) => {
241                bail!("BonsaiProver does not support compression on existing receipts");
242            }
243            (_, ReceiptKind::Groth16) => {
244                // Caller is requesting a Groth16Receipt. Provide a hint on how to get one.
245                bail!([
246                    "BonsaiProver does not support compression on existing receipts",
247                    "Set receipt_kind on ProverOpts in initial prove request to get a Groth16Receipt."
248                ].join(""));
249            }
250        }
251    }
252}