-
Notifications
You must be signed in to change notification settings - Fork 90
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
First working version of SP1 Distributed Prover
- Loading branch information
Showing
24 changed files
with
1,232 additions
and
135 deletions.
There are no files selected for viewing
Large diffs are not rendered by default.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,73 @@ | ||
| use crate::ProverState; | ||
| use raiko_lib::prover::{ProverError, WorkerError}; | ||
| use sp1_driver::{PartialProofRequest, WorkerProtocol, WorkerSocket}; | ||
| use tokio::net::TcpListener; | ||
| use tracing::{error, info, warn}; | ||
|
|
||
| async fn handle_worker_socket(mut socket: WorkerSocket) -> Result<(), ProverError> { | ||
| let protocol = socket.receive().await?; | ||
|
|
||
| info!("Received request from orchestrator: {}", protocol); | ||
|
|
||
| match protocol { | ||
| WorkerProtocol::Ping => { | ||
| socket.send(WorkerProtocol::Pong).await?; | ||
| } | ||
| WorkerProtocol::PartialProofRequest(data) => { | ||
| process_partial_proof_request(socket, data).await?; | ||
| } | ||
| _ => Err(WorkerError::InvalidRequest)?, | ||
| } | ||
|
|
||
| Ok(()) | ||
| } | ||
|
|
||
| async fn process_partial_proof_request( | ||
| mut socket: WorkerSocket, | ||
| data: PartialProofRequest, | ||
| ) -> Result<(), ProverError> { | ||
| let partial_proof = sp1_driver::Sp1DistributedProver::run_as_worker(data).await?; | ||
|
|
||
| socket | ||
| .send(WorkerProtocol::PartialProofResponse(partial_proof)) | ||
| .await?; | ||
|
|
||
| Ok(()) | ||
| } | ||
|
|
||
| async fn listen_worker(state: ProverState) { | ||
| info!( | ||
| "Listening as a SP1 worker on: {}", | ||
| state.opts.worker_address | ||
| ); | ||
|
|
||
| let listener = TcpListener::bind(state.opts.worker_address).await.unwrap(); | ||
|
|
||
| loop { | ||
| let Ok((socket, addr)) = listener.accept().await else { | ||
| error!("Error while accepting connection from orchestrator: Closing socket"); | ||
|
|
||
| return; | ||
| }; | ||
|
|
||
| if let Some(orchestrator_address) = &state.opts.orchestrator_address { | ||
| if addr.ip().to_string() != *orchestrator_address { | ||
| warn!("Unauthorized orchestrator connection from: {}", addr); | ||
|
|
||
| continue; | ||
| } | ||
| } | ||
|
|
||
| // We purposely don't spawn the task here, as we want to block to limit the number | ||
| // of concurrent connections to one. | ||
| if let Err(e) = handle_worker_socket(WorkerSocket::new(socket)).await { | ||
| error!("Error while handling worker socket: {:?}", e); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| pub async fn serve(state: ProverState) { | ||
| if state.opts.orchestrator_address.is_some() { | ||
| tokio::spawn(listen_worker(state)); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -71,4 +71,4 @@ std = [ | |
| sgx = [] | ||
| sp1 = [] | ||
| risc0 = [] | ||
| sp1-cycle-tracker = [] | ||
| sp1-cycle-tracker = [] | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,9 @@ | ||
| mod orchestrator; | ||
| mod partial_proof_request; | ||
| mod prover; | ||
| mod sp1_specifics; | ||
| mod worker; | ||
|
|
||
| pub use partial_proof_request::PartialProofRequest; | ||
| pub use prover::Sp1DistributedProver; | ||
| pub use worker::{WorkerEnvelope, WorkerProtocol, WorkerSocket}; |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,79 @@ | ||
| mod worker_client; | ||
|
|
||
| use raiko_lib::prover::WorkerError; | ||
| use sp1_core::{runtime::ExecutionState, stark::ShardProof, utils::BabyBearPoseidon2}; | ||
| use worker_client::WorkerClient; | ||
|
|
||
| use super::partial_proof_request::PartialProofRequest; | ||
|
|
||
| pub async fn distribute_work( | ||
| ip_list: Vec<String>, | ||
| checkpoints: Vec<ExecutionState>, | ||
| partial_proof_request: PartialProofRequest, | ||
| ) -> Result<Vec<ShardProof<BabyBearPoseidon2>>, WorkerError> { | ||
| let mut nb_workers = ip_list.len(); | ||
|
|
||
| let (queue_tx, queue_rx) = async_channel::bounded(nb_workers); | ||
| let (answer_tx, answer_rx) = async_channel::bounded(nb_workers); | ||
|
|
||
| // Spawn the workers | ||
| for (i, url) in ip_list.iter().enumerate() { | ||
| let worker = WorkerClient::new( | ||
| i, | ||
| url.clone(), | ||
| queue_rx.clone(), | ||
| answer_tx.clone(), | ||
| partial_proof_request.clone(), | ||
| ); | ||
|
|
||
| tokio::spawn(async move { | ||
| worker.run().await; | ||
| }); | ||
| } | ||
|
|
||
| // Send the checkpoints to the workers | ||
| for (i, checkpoint) in checkpoints.iter().enumerate() { | ||
| queue_tx.send((i, checkpoint.clone())).await.unwrap(); | ||
| } | ||
|
|
||
| let mut proofs = Vec::new(); | ||
|
|
||
| // Get the partial proofs from the workers | ||
| loop { | ||
| let (checkpoint_id, partial_proof_result) = answer_rx.recv().await.unwrap(); | ||
|
|
||
| match partial_proof_result { | ||
| Ok(partial_proof) => { | ||
| proofs.push((checkpoint_id as usize, partial_proof)); | ||
| } | ||
| Err(_e) => { | ||
| // Decrease the number of workers | ||
| nb_workers -= 1; | ||
|
|
||
| if nb_workers == 0 { | ||
| return Err(WorkerError::AllWorkersFailed); | ||
| } | ||
|
|
||
| // Push back the work for it to be done by another worker | ||
| queue_tx | ||
| .send((checkpoint_id, checkpoints[checkpoint_id as usize].clone())) | ||
| .await | ||
| .unwrap(); | ||
| } | ||
| } | ||
|
|
||
| if proofs.len() == checkpoints.len() { | ||
| break; | ||
| } | ||
| } | ||
|
|
||
| proofs.sort_by_key(|(checkpoint_id, _)| *checkpoint_id); | ||
|
|
||
| let proofs = proofs | ||
| .into_iter() | ||
| .map(|(_, proof)| proof) | ||
| .flatten() | ||
| .collect(); | ||
|
|
||
| Ok(proofs) | ||
| } |
Oops, something went wrong.