wandb

1
use std::fs;
2
use std::io::Result;
3
use std::path::Path;
4
use tempfile::tempdir;
5

6
fn main() -> Result<()> {
7
    pyo3_build_config::add_extension_module_link_args();
8
    // proto magic
9
    let protos = [
10
        "../wandb/proto/wandb_base.proto",
11
        "../wandb/proto/wandb_settings.proto",
12
        "../wandb/proto/wandb_telemetry.proto",
13
        "../wandb/proto/wandb_internal.proto",
14
        "../wandb/proto/wandb_server.proto",
15
    ];
16
    let temp_dir = tempdir().expect("Could not create temp dir");
17
    let mut temp_files = Vec::new();
18

19
    for proto in &protos {
20
        let content = fs::read_to_string(proto).expect("Could not read proto file");
21
        let modified_content = content.replace("wandb/proto/", "");
22

23
        let file_name = Path::new(proto).file_name().unwrap();
24
        let temp_file_path = temp_dir.path().join(file_name);
25

26
        fs::write(&temp_file_path, modified_content).expect("Could not write to temp file");
27
        temp_files.push(temp_file_path);
28
    }
29

30
    let temp_paths: Vec<_> = temp_files.iter().map(|f| f.to_str().unwrap()).collect();
31
    let includes = [temp_dir.path().to_str().unwrap()];
32

33
    let mut config = prost_build::Config::new();
34
    config.out_dir("src");
35
    config.compile_protos(&temp_paths, &includes).unwrap();
36

37
    // TODO: build wandb-core here and
38
    //  - either place it under wandb/wandb-core and use the env var to point to it like we do now
39
    //  - or embed as in https://zameermanji.com/blog/2021/6/17/embedding-a-rust-binary-in-another-rust-binary
40

41
    Ok(())
42
}
43