From 9ffa2367f0ac6aad1f2c0226e8cb4d5a2a50e919 Mon Sep 17 00:00:00 2001 From: Marco Inacio Date: Wed, 1 May 2024 21:27:53 +0100 Subject: [PATCH] implement BridgeStan download and model compilation in Rust (#212) * implement BridgeStan download and module compilation on Rust * rust: add feature compile-stan-model * rust: allow user defined stanc_args and make_args when compiling model * rust: add model_compiling test * rust: updating readme * rust: update documentation * rust: fix tests * rust: skip model_compiling() test on windows * rust: fix race condition in tests * rust: make example path portable * rust: fix windows absolute path resolution * Delete rust/.vscode/settings.json * rust: mark model_compiling test as ignored * rust: use mingw32-make to compile model on windows * rust: change println! to info! * Update README.md * Update Cargo.toml * rust: single compile error message * rust: run tests without feature compile-stan-model * rust: adding comments about std::fs::canonicalize * rust: fix --include-paths to point to model dir * rust: disable enum variant feature gating * rust: fix macos build * rust: make bridgestan src download more explicit * rust: only bridgestan_download_src is to be feature gated * unify .gitignore * test improvements * remove asref generic * fix tests * Update model.rs * Clean up Rust doc, tests --------- Co-authored-by: Brian Ward --- .github/workflows/main.yaml | 8 ++-- .gitignore | 1 + docs/languages/rust.rst | 8 +--- rust/Cargo.toml | 13 ++++++ rust/README.md | 52 ++++++++++----------- rust/examples/example.rs | 27 +++++++++-- rust/src/bs_safe.rs | 9 +++- rust/src/compile.rs | 91 +++++++++++++++++++++++++++++++++++++ rust/src/download.rs | 62 +++++++++++++++++++++++++ rust/src/lib.rs | 9 ++++ rust/tests/model.rs | 58 ++++++++++++++++++++--- 11 files changed, 291 insertions(+), 47 deletions(-) create mode 100644 rust/src/compile.rs create mode 100644 rust/src/download.rs diff --git a/.github/workflows/main.yaml b/.github/workflows/main.yaml index ef949ec0..a834c6a8 100644 --- a/.github/workflows/main.yaml +++ b/.github/workflows/main.yaml @@ -3,7 +3,7 @@ name: bridgestan tests on: push: branches: - - 'main' + - "main" pull_request: workflow_dispatch: {} @@ -128,7 +128,6 @@ jobs: env: BRIDGESTAN: ${{ github.workspace }} - julia: needs: [build] runs-on: ${{matrix.os}} @@ -282,4 +281,7 @@ jobs: cargo clippy cargo fmt --check cargo run --example=example - cargo test --verbose + + # run all tests with feature download-bridgestan-src + cargo test --verbose --all-features + cargo test --verbose model_compiling -- --ignored diff --git a/.gitignore b/.gitignore index d52c9093..ac21f96d 100644 --- a/.gitignore +++ b/.gitignore @@ -12,6 +12,7 @@ c-example/example_static # Rust rust/target/ rust/Cargo.lock +rust/.vscode notes.org diff --git a/docs/languages/rust.rst b/docs/languages/rust.rst index 8ccbe0ef..6c24ecc3 100644 --- a/docs/languages/rust.rst +++ b/docs/languages/rust.rst @@ -14,13 +14,9 @@ The BridgeStan Rust client is available on `crates.io ` -or use the Rust client in tandem with an interface such as :doc:`Python <./python>` -which automates this process. +The first time you compile a model, the BridgeStan source code will be downloaded to `~/.bridgestan`. If you prefer to use a source distribution of BridgeStan, you can pass its path as the `bs_path` argument to `compile_model`. -``STAN_THREADS=true`` needs to be specified when compiling a model, for more -details see the `API reference `__. +Note that the system pre-requisites from the [Getting Started Guide](../getting-started.rst) are still required and will not be automatically installed by this method. Example Program --------------- diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 7586bfe9..899da5c7 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -11,6 +11,15 @@ homepage = "https://roualdes.github.io/bridgestan/latest/" [dependencies] libloading = "0.8.0" thiserror = "1.0.40" +path-absolutize = { version = "3.1" } +log = { version = "0.4" } +ureq = { version = "2.7", optional = true } +tar = { version = "0.4", optional = true } +flate2 = { version = "1.0", optional = true } +dirs = { version = "5.0", optional = true } + +[features] +download-bridgestan-src = ["dep:ureq", "dep:tar", "dep:flate2", "dep:dirs"] [build-dependencies] bindgen = "0.69.1" @@ -18,3 +27,7 @@ bindgen = "0.69.1" [dev-dependencies] approx = "0.5.1" rand = "0.8.5" +env_logger = "0.11" + +[[example]] +name = "example" diff --git a/rust/README.md b/rust/README.md index c5a9fede..567aaea7 100644 --- a/rust/README.md +++ b/rust/README.md @@ -11,42 +11,40 @@ Internally, it relies on [`bindgen`](https://docs.rs/bindgen/) and ## Compiling the model -The Rust wrapper does not currently have any functionality to compile Stan models. -Compiled shared libraries need to be built manually using `make` or with the Julia -or Python bindings. +The Rust wrapper has the ability to compile Stan models by invoking the `make` command through the [`compile_model`] function. -For safety reasons all Stan models need to be installed with `STAN_THREADS=true`. -When compiling a model using `make`, set the environment variable: +This requires a C++ toolchain and a copy of the BridgeStan source code. The source code can be downloaded automatically by enabling the `download-bridgestan-src` feature and calling [`download_bridgestan_src`]. Alternatively, the path to the BridgeStan source code can be provided manually. -```bash -STAN_THREADS=true make some_model -``` - -When compiling a Stan model in python, this has to be specified in the `make_args` -argument: - -```python -path = bridgestan.compile_model("stan_model.stan", make_args=["STAN_THREADS=true"]) -``` +For safety reasons all Stan models need to be built with `STAN_THREADS=true`. This is the default behavior in the `compile_model` function, +but may need to be set manually when compiling the model in other contexts. If `STAN_THREADS` was not specified while building the model, the Rust wrapper will throw an error when loading the model. -## Usage: +## Usage Run this example with `cargo run --example=example`. ```rust use std::ffi::CString; -use std::path::Path; -use bridgestan::{BridgeStanError, Model, open_library}; +use std::path::{Path, PathBuf}; +use bridgestan::{BridgeStanError, Model, open_library, compile_model}; -// The path to the compiled model. -// Get for instance from python `bridgestan.compile_model` +// The path to the Stan model let path = Path::new(env!["CARGO_MANIFEST_DIR"]) .parent() .unwrap() - .join("test_models/simple/simple_model.so"); + .join("test_models/simple/simple.stan"); + +// You can manually set the BridgeStan src path or +// automatically download it (but remember to +// enable the download-bridgestan-src feature first) +let bs_path: PathBuf = "..".into(); +// let bs_path = bridgestan::download_bridgestan_src().unwrap(); + +// The path to the compiled model +let path = compile_model(&bs_path, &path, &[], &[]).expect("Could not compile Stan model."); +println!("Compiled model: {:?}", path); let lib = open_library(path).expect("Could not load compiled Stan model."); @@ -59,11 +57,13 @@ let data = CString::new(data.to_string().into_bytes()).unwrap(); let seed = 42; let model = match Model::new(&lib, Some(data), seed) { -Ok(model) => { model }, -Err(BridgeStanError::ConstructFailed(msg)) => { - panic!("Model initialization failed. Error message from Stan was {}", msg) -}, -_ => { panic!("Unexpected error") }, + Ok(model) => model, + Err(BridgeStanError::ConstructFailed(msg)) => { + panic!("Model initialization failed. Error message from Stan was {msg}") + } + Err(e) => { + panic!("Unexpected error:\n{e}") + } }; let n_dim = model.param_unc_num(); diff --git a/rust/examples/example.rs b/rust/examples/example.rs index defda712..b1cac0b9 100644 --- a/rust/examples/example.rs +++ b/rust/examples/example.rs @@ -1,14 +1,31 @@ -use bridgestan::{open_library, BridgeStanError, Model}; +use bridgestan::{compile_model, open_library, BridgeStanError, Model}; use std::ffi::CString; -use std::path::Path; +use std::path::{Path, PathBuf}; fn main() { - // The path to the compiled model. - // Get for instance from python `bridgestan.compile_model` + // Set up logging - optional + if std::env::var("RUST_LOG").is_err() { + std::env::set_var("RUST_LOG", "bridgestan=info"); + } + env_logger::init(); + + // The path to the Stan model let path = Path::new(env!["CARGO_MANIFEST_DIR"]) .parent() .unwrap() - .join("test_models/simple/simple_model.so"); + .join("test_models") + .join("simple") + .join("simple.stan"); + + // You can manually set the BridgeStan src path or + // automatically download it (but remember to + // enable the download-bridgestan-src feature first) + let bs_path: PathBuf = "..".into(); + // let bs_path = bridgestan::download_bridgestan_src().unwrap(); + + // The path to the compiled model + let path = compile_model(&bs_path, &path, &[], &[]).expect("Could not compile Stan model."); + println!("Compiled model: {:?}", path); let lib = open_library(path).expect("Could not load compiled Stan model."); diff --git a/rust/src/bs_safe.rs b/rust/src/bs_safe.rs index 7515ee69..b71c46c4 100644 --- a/rust/src/bs_safe.rs +++ b/rust/src/bs_safe.rs @@ -1,4 +1,5 @@ use crate::ffi; +use crate::VERSION; use std::borrow::Borrow; use std::collections::hash_map::DefaultHasher; use std::ffi::c_char; @@ -101,9 +102,15 @@ pub enum BridgeStanError { /// Setting a print-callback failed. #[error("Failed to set a print-callback: {0}")] SetCallbackFailed(String), + /// Compilation of the Stan model shared object failed. + #[error("Failed to compile Stan model: {0}")] + ModelCompilingFailed(String), + /// Downloading BridgeStan's C++ source code from GitHub failed. + #[error("Failed to download BridgeStan {VERSION} from github.com: {0}")] + DownloadFailed(String), } -type Result = std::result::Result; +pub(crate) type Result = std::result::Result; /// Open a compiled Stan library. /// diff --git a/rust/src/compile.rs b/rust/src/compile.rs new file mode 100644 index 00000000..7ba3c715 --- /dev/null +++ b/rust/src/compile.rs @@ -0,0 +1,91 @@ +use crate::bs_safe::{BridgeStanError, Result}; +use log::info; +use path_absolutize::Absolutize; +use std::path::{Path, PathBuf}; + +const MAKE: &str = if cfg!(target_os = "windows") { + "mingw32-make" +} else { + "make" +}; + +/// Compile a Stan Model. Requires a path to the BridgeStan sources (can be +/// downloaded with [`download_bridgestan_src`](crate::download_bridgestan_src) if that feature +/// is enabled), a path to the `.stan` file, and additional arguments +/// for the Stan compiler and the make command. +pub fn compile_model( + bs_path: &Path, + stan_file: &Path, + stanc_args: &[&str], + make_args: &[&str], +) -> Result { + // using path_absolutize crate for now since + // std::fs::canonicalize doesn't behave well on windows + // we may switch to std::path::absolute once it stabilizes, see + // https://github.com/roualdes/bridgestan/pull/212#discussion_r1513375667 + let stan_file = stan_file + .absolutize() + .map_err(|e| BridgeStanError::ModelCompilingFailed(e.to_string()))?; + + // get --include-paths=model_dir + let includir_stan_file_dir = stan_file + .parent() + .and_then(Path::to_str) + .map(|x| format!("--include-paths={x}")) + .unwrap_or_default(); + + let includir_stan_file_dir = includir_stan_file_dir.as_str(); + + if stan_file.extension().unwrap_or_default() != "stan" { + return Err(BridgeStanError::ModelCompilingFailed( + "File must be a .stan file".to_owned(), + )); + } + + // add _model suffix and change extension to .so + let output = stan_file.with_extension(""); + let output = output.with_file_name(format!( + "{}_model", + output.file_name().unwrap_or_default().to_string_lossy() + )); + let output = output.with_extension("so"); + + let stanc_args = [&[includir_stan_file_dir], stanc_args].concat(); + let stanc_args = stanc_args.join(" "); + let stanc_args = format!("STANCFLAGS={}", stanc_args); + let stanc_args = [stanc_args.as_str()]; + + let cmd = [ + &[output.to_str().unwrap_or_default()], + make_args, + stanc_args.as_slice(), + ] + .concat(); + + info!( + "Compiling model with command: {} \"{}\"", + MAKE, + cmd.join("\" \"") + ); + std::process::Command::new(MAKE) + .args(cmd) + .current_dir(bs_path) + .env("STAN_THREADS", "true") + .output() + .map_err(|e| e.to_string()) + .and_then(|proc| { + if !proc.status.success() { + Err(format!( + "{} {}", + String::from_utf8_lossy(proc.stdout.as_slice()).into_owned(), + String::from_utf8_lossy(proc.stderr.as_slice()).into_owned(), + )) + } else { + Ok(()) + } + }) + .map_err(|e| BridgeStanError::ModelCompilingFailed(e.to_string()))?; + info!("Finished compiling model"); + + Ok(output) +} diff --git a/rust/src/download.rs b/rust/src/download.rs new file mode 100644 index 00000000..29b02bf6 --- /dev/null +++ b/rust/src/download.rs @@ -0,0 +1,62 @@ +use crate::bs_safe::{BridgeStanError, Result}; +use crate::VERSION; +use flate2::read::GzDecoder; +use log::info; +use std::{env::temp_dir, fs, path::PathBuf}; +use tar::Archive; + +/// Download and unzip the BridgeStan source distribution for this version +/// to `~/.bridgestan/bridgestan-$VERSION`. +pub fn download_bridgestan_src() -> Result { + let homedir = dirs::home_dir().unwrap_or(temp_dir()); + + let bs_path_download_temp = homedir.join(".bridgestan_tmp_dir"); + let bs_path_download = homedir.join(".bridgestan"); + + let bs_path_download_temp_join_version = + bs_path_download_temp.join(format!("bridgestan-{VERSION}")); + let bs_path_download_join_version = bs_path_download.join(format!("bridgestan-{VERSION}")); + + if !bs_path_download_join_version.exists() { + info!("Downloading BridgeStan"); + + fs::remove_dir_all(&bs_path_download_temp).unwrap_or_default(); + fs::create_dir(&bs_path_download_temp).unwrap_or_default(); + fs::create_dir(&bs_path_download).unwrap_or_default(); + + let url = "https://github.com/roualdes/bridgestan/releases/download/".to_owned() + + format!("v{VERSION}/bridgestan-{VERSION}.tar.gz").as_str(); + + let response = ureq::get(url.as_str()) + .call() + .map_err(|e| BridgeStanError::DownloadFailed(e.to_string()))?; + let len = response + .header("Content-Length") + .and_then(|s| s.parse::().ok()) + .unwrap_or(50_000_000); + + let mut bytes: Vec = Vec::with_capacity(len); + response + .into_reader() + .read_to_end(&mut bytes) + .map_err(|e| BridgeStanError::DownloadFailed(e.to_string()))?; + + let tar = GzDecoder::new(bytes.as_slice()); + let mut archive = Archive::new(tar); + archive + .unpack(&bs_path_download_temp) + .map_err(|e| BridgeStanError::DownloadFailed(e.to_string()))?; + + fs::rename( + bs_path_download_temp_join_version, + &bs_path_download_join_version, + ) + .map_err(|e| BridgeStanError::DownloadFailed(e.to_string()))?; + + fs::remove_dir(bs_path_download_temp).unwrap_or_default(); + + info!("Finished downloading BridgeStan"); + } + + Ok(bs_path_download_join_version) +} diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 6e0e75bf..9d545db5 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -1,6 +1,15 @@ #![doc = include_str!("../README.md")] mod bs_safe; +mod compile; +#[cfg(feature = "download-bridgestan-src")] +mod download; pub(crate) mod ffi; pub use bs_safe::{open_library, BridgeStanError, Model, Rng, StanLibrary}; +pub use compile::compile_model; + +#[cfg(feature = "download-bridgestan-src")] +pub use download::download_bridgestan_src; + +pub const VERSION: &str = env!("CARGO_PKG_VERSION"); diff --git a/rust/tests/model.rs b/rust/tests/model.rs index d28c5822..c859496f 100644 --- a/rust/tests/model.rs +++ b/rust/tests/model.rs @@ -25,8 +25,8 @@ fn throw_data() { fn bad_arglength() { let (lib, data) = get_model("stdnormal"); let model = Model::new(&lib, data, 42).unwrap(); - let theta = vec![]; - let mut grad = vec![]; + let theta = []; + let mut grad = []; let _ = model.log_density_gradient(&theta[..], true, true, &mut grad[..]); } @@ -34,8 +34,8 @@ fn bad_arglength() { fn logp_gradient() { let (lib, data) = get_model("stdnormal"); let model = Model::new(&lib, data, 42).unwrap(); - let theta = vec![1f64]; - let mut grad = vec![0f64]; + let theta = [1f64]; + let mut grad = [0f64]; let logp = model .log_density_gradient(&theta[..], false, true, &mut grad[..]) .unwrap(); @@ -47,8 +47,8 @@ fn logp_gradient() { fn logp_hessian() { let (lib, data) = get_model("stdnormal"); let model = Model::new(&lib, data, 42).unwrap(); - let theta = vec![1f64]; - let mut grad = vec![0f64]; + let theta = [1f64]; + let mut grad = [0f64]; let mut hessian = vec![0f64]; let logp = model .log_density_hessian(&theta[..], false, true, &mut grad[..], &mut hessian) @@ -170,3 +170,49 @@ fn test_params() { .unwrap(); assert_eq!(theta_unc[0], 0.); } + +#[cfg(feature = "download-bridgestan-src")] +#[test] +fn model_downloading() { + use bridgestan::download_bridgestan_src; + + let bs_path = download_bridgestan_src().unwrap(); + let stan_path = bs_path.join("stan"); + assert!(stan_path.is_dir()); + let makefile_path = bs_path.join("Makefile"); + assert!(makefile_path.is_file()); +} + +// ignore-d to prevent overwriting the model on disk +// while other tests are running +#[test] +#[ignore] +fn model_compiling() { + use bridgestan::compile_model; + use common::model_dir; + use std::fs::remove_file; + use std::path::PathBuf; + + let name = "stdnormal"; + let mut base = model_dir(); + base.push(name); + let lib_path = base.join(format!("{}_model.so", name)); + let stan_path = base.join(format!("{}.stan", name)); + remove_file(lib_path).unwrap_or_default(); + + let bs_path: PathBuf = std::env::var("BRIDGESTAN") + .unwrap_or("..".to_string()) + .into(); + + compile_model(&bs_path, &stan_path, &[], &[]).unwrap(); + + let (lib, data) = get_model(name); + let model = Model::new(&lib, data, 42).unwrap(); + let theta = [1f64]; + let mut grad = [0f64]; + let logp = model + .log_density_gradient(&theta[..], false, true, &mut grad[..]) + .unwrap(); + assert_ulps_eq!(logp, (2. * PI).sqrt().recip().ln() - 0.5); + assert_ulps_eq!(grad[0], -1f64); +}