Skip to main content

sp1_cuda/
pk.rs

1use std::sync::Arc;
2
3use sp1_primitives::Elf;
4use sp1_prover::SP1VerifyingKey;
5
6use crate::client::CudaClient;
7
8#[derive(Clone)]
9pub struct CudaProvingKey {
10    /// The inner session key type, tells the server
11    /// to drop the key it holds when the last reference to this session is dropped.
12    inner: Arc<SessionKey>,
13}
14
15impl CudaProvingKey {
16    pub(crate) fn id(&self) -> [u8; 32] {
17        self.inner.id
18    }
19
20    pub fn elf(&self) -> &Elf {
21        &self.inner.elf
22    }
23
24    pub fn verifying_key(&self) -> &SP1VerifyingKey {
25        &self.inner.vk
26    }
27}
28
29impl CudaProvingKey {
30    pub(crate) fn new(id: [u8; 32], elf: Elf, vk: SP1VerifyingKey, client: CudaClient) -> Self {
31        Self { inner: Arc::new(SessionKey::new(id, elf, vk, client)) }
32    }
33}
34
35/// A "reference" to a key held by the CUDA server.
36pub(crate) struct SessionKey {
37    /// THe ID of the actual proving key stored in the server.
38    id: [u8; 32],
39    /// The ELF of the program.
40    elf: Elf,
41    /// The verifying key of the program.
42    vk: SP1VerifyingKey,
43    /// A client to the server that created this key.
44    client: CudaClient,
45}
46
47impl SessionKey {
48    pub(crate) const fn new(
49        id: [u8; 32],
50        elf: Elf,
51        vk: SP1VerifyingKey,
52        client: CudaClient,
53    ) -> Self {
54        Self { id, elf, vk, client }
55    }
56}
57
58impl Drop for SessionKey {
59    fn drop(&mut self) {
60        let client = self.client.clone();
61        let id = std::mem::take(&mut self.id);
62
63        tokio::spawn(async move {
64            if let Err(e) = client.destroy(id).await {
65                tracing::error!("Failed to destroy session key: {}", e);
66            }
67        });
68    }
69}