sp1-cuda 6.2.3

CUDA integration for SP1 GPU proving
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229

#![allow(unused)]

#[cfg(target_arch = "x86_64")]
mod native {
    use crate::CudaClientError;
    use std::{
        os::unix::process::CommandExt,
        path::{Path, PathBuf},
        process::Stdio,
    };
    use tokio::{
        io::AsyncWriteExt,
        process::{Child, Command},
    };

    /// Install a systemd unit for the given CUDA device id, and try to start it.
    ///
    /// Note: This method may cause race conditions, it should be called in a critical section.
    pub(crate) async fn start_server(cuda_id: u32) -> Result<Child, CudaClientError> {
        const PATH: &str = ".sp1/bin/sp1-gpu-server";

        // Get the path to where the server binary is located.
        let path = PathBuf::from(std::env::var("HOME").expect("$HOME is not set")).join(PATH);

        // Download the server binary if it doesn't exist.
        maybe_download_server(&path).await?;

        // Start the server binary.
        let child = start_binary(cuda_id, &path).await?;

        Ok(child)
    }

    /// Start the server binary, ideally with systemd-run. If systemd (--user) is not available,
    /// we will run the binary as a daemon.
    async fn start_binary(cuda_id: u32, path: &Path) -> Result<Child, CudaClientError> {
        let mut cmd = Command::new(path);
        cmd.env("CUDA_VISIBLE_DEVICES", cuda_id.to_string());

        cmd.kill_on_drop(true)
            .stdout(Stdio::inherit())
            .stderr(Stdio::inherit())
            .spawn()
            .map_err(|e| CudaClientError::new_connect(e, "Could not start `sp1-gpu-server`"))
    }

    // If the server binary is not found in the path, or if it the version is not compatible,
    // download the server binary from the release page.
    async fn maybe_download_server(path: &Path) -> Result<(), CudaClientError> {
        // Check if the server binary is in the path.
        let mut download = false;
        if !path.exists() {
            // If the path doesnt exist then there shouldnt be any instances of the server running.
            download = true;
        } else {
            let version = Command::new(path).arg("--version").output().await.map_err(|e| {
                CudaClientError::new_download_io(e, "Could not check `sp1-gpu-server` version")
            })?;

            let version = String::from_utf8_lossy(&version.stdout);
            tracing::debug!("sp1-gpu-server version: {}", version);

            // If the version is not compatible, stop all instances of the server
            // and download the new version.
            if version.trim() != sp1_primitives::SP1_CRATE_VERSION {
                download = true;

                // Stop *ALL* services, so we can replace it with a new version.
                //
                // NOTE: If a user is running a CUDA prover, across different versions,
                // on the same machine, this will cause other instances to crash!
                let mut cmd = Command::new("systemctl");
                cmd.arg("--user").arg("stop").arg(r#"sp1-gpu-server-\*"#);

                let _ = cmd.status().await.map_err(|e| {
                    CudaClientError::new_download_io(e, "Could not stop `sp1-gpu-server`")
                })?;
            }
        }

        if download {
            tracing::debug!("Downloading `sp1-gpu-server`");

            let version = format!("v{}", sp1_primitives::SP1_CRATE_VERSION);
            let repo = "succinctlabs/sp1";
            let static_url = format!("https://github.com/{repo}/releases/download");
            let asset_name = format!("sp1_gpu_server_{version}_x86_64.tar.gz");

            // Create the tar file were going to extract from.
            let tar_file = path.with_extension("tar.gz");

            // Ensure that the `.sp1` directory exists.
            tokio::fs::create_dir_all(path.parent().expect("path has no parent")).await.map_err(
                |e| CudaClientError::new_download_io(e, "Could not create `.sp1` directory"),
            )?;

            let mut file = tokio::fs::File::create(&tar_file).await.map_err(|e| {
                CudaClientError::new_download_io(e, "Could not create `sp1-gpu-server` tar file")
            })?;

            // Download the release, use a token if it exists for private releases.
            let bytes = match std::env::var("DEV_GITHUB_TOKEN").ok() {
                Some(token) => download_with_auth(&version, repo, &token, &asset_name).await,
                None => {
                    // Create the static url of the release that we expect to exist.
                    let url = format!("{static_url}/{version}/{asset_name}");

                    // Download the release.
                    let client = reqwest::Client::new();
                    let response =
                        client.get(url).send().await.map_err(CudaClientError::Download)?;

                    if !response.status().is_success() {
                        return Err(CudaClientError::Unexpected(format!(
                            "Failed to download CUDA server: {}",
                            response.text().await.expect("failed to read response text")
                        )));
                    }

                    response.bytes().await.map_err(CudaClientError::Download)
                }
            }?;

            file.write_all(&bytes).await.map_err(|e| {
                CudaClientError::new_download_io(e, "Could not write `sp1-gpu-server` tar file")
            })?;

            // Extract the tar file.
            let mut cmd = Command::new("tar");
            cmd.arg("-xzf")
                .arg(&tar_file)
                .arg("-C")
                .arg(path.parent().expect("path has no parent"));

            cmd.status().await.map_err(|e| {
                CudaClientError::new_download_io(e, "Could not extract `sp1-gpu-server` tar file")
            })?;

            // Remove the tar file.
            tokio::fs::remove_file(tar_file).await.map_err(|e| {
                CudaClientError::new_download_io(e, "Could not remove `sp1-gpu-server` tar file")
            })?;
        }

        Ok(())
    }

    async fn download_with_auth(
        tag: &str,
        repo: &str,
        token: &str,
        asset_name: &str,
    ) -> Result<bytes::Bytes, CudaClientError> {
        tracing::trace!("downloading with auth");

        // 1. Find the release by tag
        #[derive(serde::Deserialize)]
        struct Release {
            assets: Vec<Asset>,
        }
        #[derive(serde::Deserialize)]
        struct Asset {
            id: u64,
            name: String,
        }

        let api = format!("https://api.github.com/repos/{repo}");
        let client = reqwest::Client::builder()
            .user_agent("sp1-cuda-downloader")
            .build()
            .expect("failed to build reqwest client");

        let release: Release = client
            .get(format!("{api}/releases/tags/{tag}"))
            .bearer_auth(token)
            .send()
            .await
            .map_err(CudaClientError::Download)?
            .error_for_status()
            .map_err(CudaClientError::Download)?
            .json()
            .await
            .map_err(CudaClientError::Download)?;

        let asset_id = release
            .assets
            .into_iter()
            .find(|a| a.name == asset_name)
            .ok_or_else(|| {
                CudaClientError::Unexpected(format!(
                    "asset {asset_name} not found in release {tag}"
                ))
            })?
            .id;

        // 2. Download the binary content
        let bytes = client
            .get(format!("{api}/releases/assets/{asset_id}"))
            .bearer_auth(token)
            .header(reqwest::header::ACCEPT, "application/octet-stream")
            .send()
            .await
            .map_err(CudaClientError::Download)?
            .error_for_status()
            .map_err(CudaClientError::Download)?
            .bytes()
            .await
            .map_err(CudaClientError::Download)?;

        Ok(bytes)
    }

    /// The name of the systemd unit for the given CUDA device id.
    fn unit_name(cuda_id: u32) -> String {
        format!("sp1-gpu-server-{cuda_id}")
    }
}

#[cfg(not(target_arch = "x86_64"))]
mod native {
    use crate::CudaClientError;
    use tokio::process::Child;

    pub(crate) async fn start_server(cuda_id: u32) -> Result<Child, CudaClientError> {
        panic!("Unsupported architecture for CUDA server: {}", std::env::consts::ARCH);
    }
}

pub(crate) use native::start_server;