[backend-comparison] Upload benchmarks to server (#1381)

Uploading is enabled with already implemented --share argument
of the burnbench command line tool.

The burnbench binary passes the URL of the server and the auth
token to the cargo bench process using the additional arguments
--sharing-url and --sharing-token respectively.

The persistence module then upload the results when a --sharing-url
is provided.

The URL is for now hardcoded. The endpoint is production when
compiling in release mode and it is localhost otherwise.

Author: syl20bnr

Cargo.lock

@@ -42,6 +42,7 @@ strum = { workspace = true }
 strum_macros = { workspace = true }
 
 [dev-dependencies]
+rstest = { workspace = true }
 serial_test = { workspace = true }
 
 [[bench]]

backend-comparison/Cargo.toml

@@ -36,13 +36,13 @@ impl<B: Backend, const D: usize> Benchmark for BinaryBenchmark<B, D> {
 }
 
 #[allow(dead_code)]
-fn bench<B: Backend>(device: &B::Device) {
+fn bench<B: Backend>(device: &B::Device, url: Option<&str>, token: Option<&str>) {
     let benchmark = BinaryBenchmark::<B, 3> {
         shape: [32, 512, 1024].into(),
         device: device.clone(),
     };
 
-    save::<B>(vec![run_benchmark(benchmark)], device).unwrap();
+    save::<B>(vec![run_benchmark(benchmark)], device, url, token).unwrap();
 }
 
 fn main() {

backend-comparison/benches/binary.rs

@@ -114,7 +114,7 @@ fn erf_positive<B: Backend, const D: usize>(x: Tensor<B, D>) -> Tensor<B, D> {
 }
 
 #[allow(dead_code)]
-fn bench<B: Backend>(device: &B::Device) {
+fn bench<B: Backend>(device: &B::Device, url: Option<&str>, token: Option<&str>) {
     const D: usize = 3;
     let shape: Shape<D> = [32, 512, 2048].into();
 
@@ -145,6 +145,8 @@ fn bench<B: Backend>(device: &B::Device) {
                 run_benchmark(custom_erf_gelu),
             ],
             device,
+            url,
+            token,
         )
         .unwrap();
     };

backend-comparison/benches/custom_gelu.rs

@@ -71,7 +71,7 @@ impl<B: Backend, const D: usize> Benchmark for FromDataBenchmark<B, D> {
 }
 
 #[allow(dead_code)]
-fn bench<B: Backend>(device: &B::Device) {
+fn bench<B: Backend>(device: &B::Device, url: Option<&str>, token: Option<&str>) {
     const D: usize = 3;
     let shape: Shape<D> = [32, 512, 1024].into();
 
@@ -81,6 +81,8 @@ fn bench<B: Backend>(device: &B::Device) {
     save::<B>(
         vec![run_benchmark(to_benchmark), run_benchmark(from_benchmark)],
         device,
+        url,
+        token,
     )
     .unwrap();
 }

backend-comparison/benches/data.rs

@@ -42,7 +42,7 @@ impl<B: Backend, const D: usize> Benchmark for MatmulBenchmark<B, D> {
 }
 
 #[allow(dead_code)]
-fn bench<B: Backend>(device: &B::Device) {
+fn bench<B: Backend>(device: &B::Device, url: Option<&str>, token: Option<&str>) {
     const D: usize = 3;
     let batch_size = 3;
     let m = 1024;
@@ -53,7 +53,7 @@ fn bench<B: Backend>(device: &B::Device) {
 
     let benchmark = MatmulBenchmark::<B, D>::new(shape_lhs, shape_rhs, device.clone());
 
-    save::<B>(vec![run_benchmark(benchmark)], device).unwrap();
+    save::<B>(vec![run_benchmark(benchmark)], device, url, token).unwrap();
 }
 
 fn main() {

backend-comparison/benches/matmul.rs

@@ -35,13 +35,13 @@ impl<B: Backend, const D: usize> Benchmark for UnaryBenchmark<B, D> {
 }
 
 #[allow(dead_code)]
-fn bench<B: Backend>(device: &B::Device) {
+fn bench<B: Backend>(device: &B::Device, url: Option<&str>, token: Option<&str>) {
     const D: usize = 3;
     let shape: Shape<D> = [32, 512, 1024].into();
 
     let benchmark = UnaryBenchmark::<B, D>::new(shape, device.clone());
 
-    save::<B>(vec![run_benchmark(benchmark)], device).unwrap();
+    save::<B>(vec![run_benchmark(benchmark)], device, url, token).unwrap();
 }
 
 fn main() {

backend-comparison/benches/unary.rs

@@ -16,13 +16,26 @@ use super::{
 };
 
 const FIVE_SECONDS: time::Duration = time::Duration::new(5, 0);
+const USER_BENCHMARK_SERVER_URL: &str = if cfg!(debug_assertions) {
+    // development
+    "http://localhost:8000/benchmarks"
+} else {
+    // production
+    "https://user-benchmark-server-gvtbw64teq-nn.a.run.app/benchmarks"
+};
 
 /// Base trait to define an application
 pub(crate) trait Application {
     fn init(&mut self) {}
 
     #[allow(unused)]
-    fn run(&mut self, benches: &[BenchmarkValues], backends: &[BackendValues]) {}
+    fn run(
+        &mut self,
+        benches: &[BenchmarkValues],
+        backends: &[BackendValues],
+        token: Option<&str>,
+    ) {
+    }
 
     fn cleanup(&mut self) {}
 }
@@ -150,17 +163,17 @@ fn command_list() {
 }
 
 fn command_run(run_args: RunArgs) {
+    let token = get_token_from_cache();
     if run_args.share {
         // Verify if a token is saved
-        let token = get_token_from_cache();
         if token.is_none() {
             eprintln!("You need to be authenticated to be able to share benchmark results.");
             eprintln!("Run the command 'burnbench auth' to authenticate.");
             return;
         }
         // TODO refresh the token when it is expired
         // Check for the validity of the saved token
-        if !verify_token(&token.unwrap()) {
+        if !verify_token(token.as_deref().unwrap()) {
             eprintln!("Your access token is no longer valid.");
             eprintln!("Run the command 'burnbench auth' again to get a new token.");
             return;
@@ -179,13 +192,12 @@ fn command_run(run_args: RunArgs) {
     let mut app = App::new();
     app.init();
     println!("Running benchmarks...");
-    app.run(&run_args.benches, &run_args.backends);
+    app.run(
+        &run_args.benches,
+        &run_args.backends,
+        token.as_deref().filter(|_| run_args.share),
+    );
     app.cleanup();
-    println!("Cleanup completed. Benchmark run(s) finished.");
-    if run_args.share {
-        println!("Sharing results...");
-        // TODO Post the results once backend can verify the GitHub access token
-    }
 }
 
 #[allow(unused)] // for tui as this is WIP
@@ -203,3 +215,33 @@ pub(crate) fn run_cargo(command: &str, params: &[&str]) {
         std::process::exit(status.code().unwrap_or(1));
     }
 }
+
+pub(crate) fn run_backend_comparison_benchmarks(
+    benches: &[BenchmarkValues],
+    backends: &[BackendValues],
+    token: Option<&str>,
+) {
+    // Iterate over each combination of backend and bench
+    for backend in backends.iter() {
+        for bench in benches.iter() {
+            let bench_str = bench.to_string();
+            let backend_str = backend.to_string();
+            let mut args = vec![
+                "-p",
+                "backend-comparison",
+                "--bench",
+                &bench_str,
+                "--features",
+                &backend_str,
+            ];
+            if let Some(t) = token {
+                args.push("--");
+                args.push("--sharing-url");
+                args.push(USER_BENCHMARK_SERVER_URL);
+                args.push("--sharing-token");
+                args.push(t);
+            }
+            run_cargo("bench", &args);
+        }
+    }
+}

backend-comparison/src/burnbenchapp/base.rs

@@ -1,4 +1,6 @@
-use crate::burnbenchapp::{run_cargo, Application, BackendValues, BenchmarkValues};
+use crate::burnbenchapp::{
+    run_backend_comparison_benchmarks, Application, BackendValues, BenchmarkValues,
+};
 
 use derive_new::new;
 
@@ -8,21 +10,13 @@ pub struct TermApplication;
 impl Application for TermApplication {
     fn init(&mut self) {}
 
-    fn run(&mut self, benches: &[BenchmarkValues], backends: &[BackendValues]) {
-        // Iterate over each combination of backend and bench
-        for backend in backends.iter() {
-            for bench in benches.iter() {
-                run_cargo(
-                    "bench",
-                    &[
-                        "--bench",
-                        &bench.to_string(),
-                        "--features",
-                        &backend.to_string(),
-                    ],
-                );
-            }
-        }
+    fn run(
+        &mut self,
+        benches: &[BenchmarkValues],
+        backends: &[BackendValues],
+        token: Option<&str>,
+    ) {
+        run_backend_comparison_benchmarks(benches, backends, token)
     }
 
     fn cleanup(&mut self) {}

backend-comparison/src/burnbenchapp/term/base.rs

@@ -32,7 +32,12 @@ impl Application for TuiApplication {
     fn init(&mut self) {}
 
     #[allow(unused)]
-    fn run(&mut self, benches: &[BenchmarkValues], backends: &[BackendValues]) {
+    fn run(
+        &mut self,
+        benches: &[BenchmarkValues],
+        backends: &[BackendValues],
+        token: Option<&str>,
+    ) {
         // TODO initialize widgets given passed benches and backends on the command line
         loop {
             self.terminal

backend-comparison/src/burnbenchapp/tui/base.rs

@@ -1,14 +1,44 @@
 pub mod burnbenchapp;
 pub mod persistence;
 
+/// Simple parse to retrieve additional argument passed to cargo bench command
+/// We cannot use clap here as clap parser does not allow to have unknown arguments.
+pub fn get_argument<'a>(args: &'a [String], arg_name: &'a str) -> Option<&'a str> {
+    let mut i = 0;
+    while i < args.len() {
+        match args[i].as_str() {
+            arg if arg == arg_name && i + 1 < args.len() => {
+                return Some(&args[i + 1]);
+            }
+            _ => i += 1,
+        }
+    }
+    None
+}
+
+/// Specialized function to retrieve the sharing token
+pub fn get_sharing_token(args: &[String]) -> Option<&str> {
+    get_argument(args, "--sharing-token")
+}
+
+/// Specialized function to retrieve the sharing URL
+pub fn get_sharing_url(args: &[String]) -> Option<&str> {
+    get_argument(args, "--sharing-url")
+}
+
 #[macro_export]
 macro_rules! bench_on_backend {
     () => {
+        use std::env;
+        let args: Vec<String> = env::args().collect();
+        let url = backend_comparison::get_sharing_url(&args);
+        let token = backend_comparison::get_sharing_token(&args);
+
         #[cfg(feature = "wgpu")]
         {
             use burn::backend::wgpu::{AutoGraphicsApi, Wgpu, WgpuDevice};
 
-            bench::<Wgpu<AutoGraphicsApi, f32, i32>>(&WgpuDevice::default());
+            bench::<Wgpu<AutoGraphicsApi, f32, i32>>(&WgpuDevice::default(), url, token);
         }
 
         #[cfg(feature = "tch-gpu")]
@@ -19,15 +49,15 @@ macro_rules! bench_on_backend {
             let device = LibTorchDevice::Cuda(0);
             #[cfg(target_os = "macos")]
             let device = LibTorchDevice::Mps;
-            bench::<LibTorch>(&device);
+            bench::<LibTorch>(&device, url, token);
         }
 
         #[cfg(feature = "tch-cpu")]
         {
             use burn::backend::{libtorch::LibTorchDevice, LibTorch};
 
             let device = LibTorchDevice::Cpu;
-            bench::<LibTorch>(&device);
+            bench::<LibTorch>(&device, url, token);
         }
 
         #[cfg(any(
@@ -41,7 +71,7 @@ macro_rules! bench_on_backend {
             use burn::backend::NdArray;
 
             let device = NdArrayDevice::Cpu;
-            bench::<NdArray>(&device);
+            bench::<NdArray>(&device, url, token);
         }
 
         #[cfg(feature = "candle-cpu")]
@@ -50,7 +80,7 @@ macro_rules! bench_on_backend {
             use burn::backend::Candle;
 
             let device = CandleDevice::Cpu;
-            bench::<Candle>(&device);
+            bench::<Candle>(&device, url, token);
         }
 
         #[cfg(feature = "candle-cuda")]
@@ -59,7 +89,7 @@ macro_rules! bench_on_backend {
             use burn::backend::Candle;
 
             let device = CandleDevice::Cuda(0);
-            bench::<Candle>(&device);
+            bench::<Candle>(&device, url, token);
         }
 
         #[cfg(feature = "candle-metal")]
@@ -68,7 +98,35 @@ macro_rules! bench_on_backend {
             use burn::backend::Candle;
 
             let device = CandleDevice::Metal(0);
-            bench::<Candle>(&device);
+            bench::<Candle>(&device, url, token);
         }
     };
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use rstest::rstest;
+
+    #[rstest]
+    #[case::sharing_token_argument_with_value(&["--sharing-token", "token123"], Some("token123"))]
+    #[case::sharing_token_argument_no_value(&["--sharing-token"], None)]
+    #[case::sharing_token_argument_with_additional_arguments(&["--other-arg", "value", "--sharing-token", "token789"], Some("token789"))]
+    #[case::other_argument(&["--other-arg", "value"], None)]
+    #[case::no_argument(&[], None)]
+    fn test_get_sharing_token(#[case] args: &[&str], #[case] expected: Option<&str>) {
+        let args = args.iter().map(|s| s.to_string()).collect::<Vec<String>>();
+        assert_eq!(get_sharing_token(&args), expected);
+    }
+
+    #[rstest]
+    #[case::sharing_url_argument_with_value(&["--sharing-url", "url123"], Some("url123"))]
+    #[case::sharing_url_argument_no_value(&["--sharing-url"], None)]
+    #[case::sharing_url_argument_with_additional_arguments(&["--other-arg", "value", "--sharing-url", "url789"], Some("url789"))]
+    #[case::other_argument(&["--other-arg", "value"], None)]
+    #[case::no_argument(&[], None)]
+    fn test_get_sharing_url(#[case] args: &[&str], #[case] expected: Option<&str>) {
+        let args = args.iter().map(|s| s.to_string()).collect::<Vec<String>>();
+        assert_eq!(get_sharing_url(&args), expected);
+    }
+}