diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 656573d..1b93162 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -40,9 +40,14 @@ jobs: - os: linux arch: arm64 - runner: ubuntu-latest + runner: ubuntu-24.04-arm catboost_version: "1.2.7" + - os: linux + arch: arm64 + runner: ubuntu-24.04-arm + catboost_version: "1.2.10" + # Windows x86_64 - os: windows arch: x86_64 diff --git a/README.md b/README.md index 0770879..a53c637 100644 --- a/README.md +++ b/README.md @@ -121,6 +121,20 @@ let features = ObjectsOrderFeatures::new() let predictions = model.predict(features)?; ``` +### Zero-Copy Buffer Loading (Recommended) + +`Model::load_buffer_zero_copy` is the recommended way to load models from memory. Unlike `load_buffer`, it avoids copying the model data, resulting in lower memory usage, faster loading, and no internal memory pool leaks. Requires CatBoost v1.2.9+ (the default). + +```rust +use catboost_rust::Model; +use std::fs; + +let buffer = fs::read("model.cbm")?; +let model = Model::load_buffer_zero_copy(buffer)?; +``` + +The buffer is owned by the `Model` and freed automatically when it is dropped. + ## Configuration ### CatBoost Version diff --git a/build.rs b/build.rs index 1d94bfe..e4d044a 100644 --- a/build.rs +++ b/build.rs @@ -260,6 +260,7 @@ fn main() { println!("cargo::rustc-check-cfg=cfg(catboost_text_count)"); println!("cargo::rustc-check-cfg=cfg(catboost_staged_prediction)"); println!("cargo::rustc-check-cfg=cfg(catboost_feature_indices)"); + println!("cargo::rustc-check-cfg=cfg(catboost_zero_copy)"); // Parse version for feature detection let version = get_catboost_version(); @@ -294,6 +295,11 @@ fn main() { println!("cargo:rustc-cfg=catboost_feature_indices"); } + // v1.2.9+: Zero-copy buffer loading + if major > 1 || (major == 1 && minor > 2) || (major == 1 && minor == 2 && patch >= 9) { + println!("cargo:rustc-cfg=catboost_zero_copy"); + } + // Download the model interface headers if let Err(e) = download_model_interface_headers(&out_dir) { eprintln!("Failed to download model interface headers: {}", e); diff --git a/examples/advanced_usage.rs b/examples/advanced_usage.rs index 7a32291..8365194 100644 --- a/examples/advanced_usage.rs +++ b/examples/advanced_usage.rs @@ -15,9 +15,9 @@ fn main() -> Result<(), CatBoostError> { return Ok(()); } - // Load the model + // Load the model (prefer zero-copy when available) println!("Loading model from {}...", model_path); - let model = Model::load(model_path)?; + let model = load_model(model_path)?; println!("Model loaded successfully!"); @@ -114,6 +114,21 @@ fn main() -> Result<(), CatBoostError> { Ok(()) } +#[cfg(catboost_zero_copy)] +fn load_model(path: &str) -> Result { + println!(" (using zero-copy buffer loading)"); + let buffer = fs::read(path).map_err(|e| CatBoostError { + description: format!("could not read file into memory: {}", e), + })?; + Model::load_buffer_zero_copy(buffer) +} + +#[cfg(not(catboost_zero_copy))] +fn load_model(path: &str) -> Result { + println!(" (using file loading - zero-copy not available in this CatBoost version)"); + Model::load(path) +} + fn display_model_info(model: &Model) -> Result<(), CatBoostError> { println!("Model Information:"); println!( diff --git a/examples/basic_usage.rs b/examples/basic_usage.rs index feaf181..a9d7619 100644 --- a/examples/basic_usage.rs +++ b/examples/basic_usage.rs @@ -13,12 +13,14 @@ fn main() -> Result<(), CatBoostError> { model_path ); create_simple_example()?; - return Ok(()); + return Err(CatBoostError { + description: "No model file found at {}. Creating a simple example..".to_string(), + }); } - // Load the model + // Load the model (prefer zero-copy when available) println!("Loading model from {}...", model_path); - let model = Model::load(model_path)?; + let model = load_model(model_path)?; println!("Model loaded successfully!"); println!("Model info:"); @@ -88,6 +90,21 @@ fn main() -> Result<(), CatBoostError> { Ok(()) } +#[cfg(catboost_zero_copy)] +fn load_model(path: &str) -> Result { + println!(" (using zero-copy buffer loading)"); + let buffer = fs::read(path).map_err(|e| CatBoostError { + description: format!("could not read file into memory: {}", e), + })?; + Model::load_buffer_zero_copy(buffer) +} + +#[cfg(not(catboost_zero_copy))] +fn load_model(path: &str) -> Result { + println!(" (using file loading - zero-copy not available in this CatBoost version)"); + Model::load(path) +} + fn create_simple_example() -> Result<(), CatBoostError> { println!("Since no model file is available, here's how you would use the library:"); println!(); diff --git a/examples/gpu_usage.rs b/examples/gpu_usage.rs index 3948fda..5d36cee 100644 --- a/examples/gpu_usage.rs +++ b/examples/gpu_usage.rs @@ -1,12 +1,25 @@ use catboost_rust::{Model, ObjectsOrderFeatures}; +#[cfg(catboost_zero_copy)] +fn load_model(path: &str) -> Result> { + println!(" (using zero-copy buffer loading)"); + let buffer = std::fs::read(path)?; + Ok(Model::load_buffer_zero_copy(buffer)?) +} + +#[cfg(not(catboost_zero_copy))] +fn load_model(path: &str) -> Result> { + println!(" (using file loading - zero-copy not available in this CatBoost version)"); + Ok(Model::load(path)?) +} + fn main() -> Result<(), Box> { println!("CatBoost Rust Example - GPU Usage"); println!("=================================="); - // Load a model + // Load a model (prefer zero-copy when available) println!("Loading model from tmp/model.bin..."); - let model = Model::load("tmp/model.bin")?; + let model = load_model("tmp/model.bin")?; println!("Model loaded successfully!"); // Display model information diff --git a/src/model.rs b/src/model.rs index eb4a76b..cbd106b 100644 --- a/src/model.rs +++ b/src/model.rs @@ -7,6 +7,9 @@ use std::path::Path; pub struct Model { handle: *mut sys::ModelCalcerHandle, + /// Buffer owner for zero-copy loading - keeps the buffer alive for model's lifetime + #[cfg(catboost_zero_copy)] + _buffer_owner: Option>, } unsafe impl Send for Model {} @@ -17,6 +20,8 @@ impl Model { let model_handle = unsafe { sys::ModelCalcerCreate() }; Model { handle: model_handle, + #[cfg(catboost_zero_copy)] + _buffer_owner: None, } } @@ -30,6 +35,37 @@ impl Model { Ok(model) } + /// Load a model from a buffer using zero-copy approach + /// + /// This method uses LoadFullModelZeroCopy which does NOT copy the model data. + /// Instead, the model keeps a reference to the buffer and reads from it directly. + /// + /// Requires CatBoost v1.2.9+. + /// + /// # Example + /// ```no_run + /// use catboost_rust::Model; + /// use std::fs; + /// + /// let buffer = fs::read("model.cbm").unwrap(); + /// let model = Model::load_buffer_zero_copy(buffer).unwrap(); + /// ``` + #[cfg(catboost_zero_copy)] + pub fn load_buffer_zero_copy(buffer: Vec) -> CatBoostResult { + let mut model = Model::new(); + + CatBoostError::check_return_value(unsafe { + sys::LoadFullModelZeroCopy( + model.handle, + buffer.as_ptr() as *const std::os::raw::c_void, + buffer.len(), + ) + })?; + + model._buffer_owner = Some(buffer); + Ok(model) + } + /// Load a model from a buffer pub fn load_buffer>>(buffer: P) -> CatBoostResult { let model = Model::new();