Skip to main content

sp1_prover/worker/prover/
engine.rs

1use std::sync::Arc;
2
3use slop_futures::pipeline::SubmitError;
4use sp1_core_executor::SP1CoreOpts;
5use sp1_hypercube::prover::ProverSemaphore;
6use sp1_prover_types::{Artifact, ArtifactClient};
7
8use crate::{
9    worker::{
10        CoreProveSubmitHandle, RawTaskRequest, RecursionVkWorker, ReduceSubmitHandle,
11        SP1CoreProver, SP1CoreProverConfig, SP1DeferredProver, SP1DeferredProverConfig,
12        SP1DeferredSubmitHandle, SP1RecursionProver, SP1RecursionProverConfig, SetupSubmitHandle,
13        SetupTask, TaskError, TaskId, WorkerClient,
14    },
15    SP1ProverComponents, WrapProverBuilder,
16};
17
18#[derive(Clone)]
19pub struct SP1ProverConfig {
20    pub core_prover_config: SP1CoreProverConfig,
21    pub recursion_prover_config: SP1RecursionProverConfig,
22    pub deferred_prover_config: SP1DeferredProverConfig,
23}
24
25pub struct SP1ProverEngine<A, W, C: SP1ProverComponents> {
26    pub core_prover: SP1CoreProver<A, W, C>,
27    pub recursion_prover: SP1RecursionProver<A, C>,
28    pub deferred_prover: SP1DeferredProver<A, C>,
29    pub vk_worker: RecursionVkWorker<C>,
30}
31
32pub struct WrapAirProverInit<C: SP1ProverComponents> {
33    builder: Arc<C::WrapProverBuilder>,
34    permits: ProverSemaphore,
35}
36
37impl<C: SP1ProverComponents> WrapAirProverInit<C> {
38    pub(crate) fn new(builder: C::WrapProverBuilder, permits: ProverSemaphore) -> Self {
39        Self { builder: Arc::new(builder), permits }
40    }
41
42    pub(crate) fn permits(&self) -> ProverSemaphore {
43        self.permits.clone()
44    }
45
46    pub(crate) fn build(&self) -> Arc<C::WrapProver> {
47        self.builder.build()
48    }
49}
50
51impl<A: ArtifactClient, W: WorkerClient, C: SP1ProverComponents> SP1ProverEngine<A, W, C> {
52    #[allow(clippy::too_many_arguments)]
53    pub async fn new(
54        config: SP1ProverConfig,
55        opts: SP1CoreOpts,
56        artifact_client: A,
57        worker_client: W,
58        core_prover_and_permits: (Arc<C::CoreProver>, ProverSemaphore),
59        recursion_prover_and_permits: (Arc<C::RecursionProver>, ProverSemaphore),
60        shrink_air_prover_and_permits: (Arc<C::RecursionProver>, ProverSemaphore),
61        wrap_air_prover_init: WrapAirProverInit<C>,
62    ) -> Self {
63        let recursion_prover = SP1RecursionProver::new(
64            config.recursion_prover_config,
65            artifact_client.clone(),
66            recursion_prover_and_permits.clone(),
67            shrink_air_prover_and_permits,
68            wrap_air_prover_init,
69        )
70        .await;
71
72        let core_prover = SP1CoreProver::new(
73            config.core_prover_config,
74            opts,
75            artifact_client.clone(),
76            worker_client,
77            core_prover_and_permits.0,
78            core_prover_and_permits.1,
79            recursion_prover.clone(),
80        );
81
82        let deferred_prover = SP1DeferredProver::new(
83            config.deferred_prover_config,
84            recursion_prover.clone(),
85            artifact_client,
86        );
87
88        let vk_worker = RecursionVkWorker {
89            recursion_permits: recursion_prover_and_permits.1,
90            recursion_prover: recursion_prover_and_permits.0,
91            shrink_prover: recursion_prover.shrink_prover.clone(),
92        };
93
94        Self { core_prover, vk_worker, recursion_prover, deferred_prover }
95    }
96
97    pub async fn submit_prove_core_shard(
98        &self,
99        request: RawTaskRequest,
100    ) -> Result<CoreProveSubmitHandle<A, W, C>, TaskError> {
101        self.core_prover.submit_prove_shard(request).await
102    }
103
104    pub async fn submit_setup(
105        &self,
106        id: TaskId,
107        elf: Artifact,
108        output: Artifact,
109    ) -> Result<SetupSubmitHandle<A, C>, SubmitError> {
110        let handle = self.core_prover.submit_setup(SetupTask { id, elf, output }).await?;
111        Ok(handle)
112    }
113
114    pub async fn submit_recursion_reduce(
115        &self,
116        request: RawTaskRequest,
117    ) -> Result<ReduceSubmitHandle<A, C>, TaskError> {
118        self.recursion_prover.submit_recursion_reduce(request).await
119    }
120
121    pub async fn submit_prove_deferred(
122        &self,
123        request: RawTaskRequest,
124    ) -> Result<SP1DeferredSubmitHandle<A, C>, TaskError> {
125        self.deferred_prover.submit(request).await
126    }
127
128    pub async fn run_shrink_wrap(&self, request: RawTaskRequest) -> Result<(), TaskError> {
129        self.recursion_prover.run_shrink_wrap(request).await
130    }
131
132    pub async fn run_plonk(&self, request: RawTaskRequest) -> Result<(), TaskError> {
133        self.recursion_prover.run_plonk(request).await
134    }
135
136    pub async fn run_groth16(&self, request: RawTaskRequest) -> Result<(), TaskError> {
137        self.recursion_prover.run_groth16(request).await
138    }
139}