Skip to content

Commit

Permalink
dynamic C lib now linking to onnxruntime
Browse files Browse the repository at this point in the history
  • Loading branch information
maxwellflitton committed Jan 21, 2025
1 parent f4145ae commit 6e12b7e
Show file tree
Hide file tree
Showing 22 changed files with 296 additions and 239 deletions.
1 change: 1 addition & 0 deletions modules/c-wrapper/.gitignore
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
onnx_lib/
1 change: 1 addition & 0 deletions modules/c-wrapper/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -21,3 +21,4 @@ reqwest = { version = "0.12.12", features = ["blocking", "json"] }
# tokio = { version = "1", features = ["full"] } # Required for reqwest
tar = "0.4" # For extracting tar files
flate2 = "1.0" # For handling gzip
zip = "2.2.2"
4 changes: 3 additions & 1 deletion modules/c-wrapper/Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,9 @@ COPY . .
RUN apt-get update && apt-get install -y python3 python3-pip

# Clean and build the Rust project
RUN cd c-wrapper/scripts && bash prep_tests.sh
# RUN cd c-wrapper/scripts && bash prep_tests.sh
# RUN cd c-wrapper && cargo build --verbose > build_log.txt 2>&1
RUN cd c-wrapper && cargo build && bash scripts/copy_over_lib.sh

# RUN rm /onnxruntime

Expand Down
4 changes: 3 additions & 1 deletion modules/c-wrapper/build-context/Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,9 @@ COPY . .
RUN apt-get update && apt-get install -y python3 python3-pip

# Clean and build the Rust project
RUN cd c-wrapper/scripts && bash prep_tests.sh
# RUN cd c-wrapper/scripts && bash prep_tests.sh
# RUN cd c-wrapper && cargo build --verbose > build_log.txt 2>&1
RUN cd c-wrapper && cargo build && bash scripts/copy_over_lib.sh

# RUN rm /onnxruntime

Expand Down
11 changes: 9 additions & 2 deletions modules/c-wrapper/build-context/c-wrapper/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,8 @@ version = "0.1.0"
edition = "2021"

[dependencies]
surrealml-core = { path = "../core" }
uuid = { version = "1.4.1", features = ["v4"] }
surrealml-core = { path = "../core", features = ["dynamic"] }
uuid = { version = "1.11.1", features = ["v4"] }
ndarray = "0.16.1"

# for the uploading the model to the server
Expand All @@ -15,3 +15,10 @@ base64 = "0.13"

[lib]
crate-type = ["cdylib"]

[build-dependencies]
reqwest = { version = "0.12.12", features = ["blocking", "json"] }
# tokio = { version = "1", features = ["full"] } # Required for reqwest
tar = "0.4" # For extracting tar files
flate2 = "1.0" # For handling gzip
zip = "2.2.2"
Original file line number Diff line number Diff line change
Expand Up @@ -6,26 +6,29 @@ cd $SCRIPTPATH

cd ..

# wipe and build the build context
BUILD_DIR="build-context"

if [ -d "$BUILD_DIR" ]; then
echo "Cleaning up existing build directory..."
rm -rf "$BUILD_DIR"
fi

mkdir "$BUILD_DIR"
mkdir "$BUILD_DIR"/c-wrapper

# copy over the code to be built
cp -r src "$BUILD_DIR"/c-wrapper/src
cp -r tests "$BUILD_DIR"/c-wrapper/tests
cp -r scripts "$BUILD_DIR"/c-wrapper/scripts
cp Cargo.toml "$BUILD_DIR"/c-wrapper/Cargo.toml

cp build.rs "$BUILD_DIR"/c-wrapper/build.rs
cp -r ../core "$BUILD_DIR"/core
cp Dockerfile "$BUILD_DIR"/Dockerfile

# remove unnecessary files
rm -rf "$BUILD_DIR"/core/.git
rm -rf "$BUILD_DIR"/core/target/

cp Dockerfile "$BUILD_DIR"/Dockerfile
# build the docker image
cd "$BUILD_DIR"
docker build --no-cache -t c-wrapper-tests .

Expand Down
1 change: 1 addition & 0 deletions modules/c-wrapper/build-context/c-wrapper/src/api/mod.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
//! C API for interacting with the SurML file storage and executing models.
pub mod execution;
pub mod storage;
pub mod ml_sys;
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import platform
from pathlib import Path
import os
from test_utils.return_structs import EmptyReturn


def load_library(lib_name: str = "libc_wrapper") -> ctypes.CDLL:
Expand Down Expand Up @@ -38,9 +39,18 @@ def load_library(lib_name: str = "libc_wrapper") -> ctypes.CDLL:
# os.environ["LD_LIBRARY_PATH"] = f"{onnx_lib_path}:{current_ld_library_path}"
# os.environ["ORT_LIB_LOCATION"] = str(onnx_lib_path)

ctypes.CDLL(str(onnx_path), mode=ctypes.RTLD_GLOBAL)
# ctypes.CDLL(str(onnx_path), mode=ctypes.RTLD_GLOBAL)
onnx_path = current_dir.joinpath("onnxruntime")

if not lib_path.exists():
raise FileNotFoundError(f"Shared library not found at: {lib_path}")

loaded_lib = ctypes.CDLL(str(lib_path))
loaded_lib.link_onnx.argtypes = [ctypes.c_char_p]
loaded_lib.link_onnx.restype = EmptyReturn
c_string = str(onnx_path).encode('utf-8')
load_info = loaded_lib.link_onnx(c_string)
if load_info.error_message:
raise OSError(f"Failed to load onnxruntime: {load_info.error_message.decode('utf-8')}")

return ctypes.CDLL(str(lib_path))
Original file line number Diff line number Diff line change
Expand Up @@ -9,3 +9,4 @@
TEST_SURML_PATH = ASSETS_PATH.joinpath("test.surml")
SHOULD_BREAK_FILE = ASSETS_PATH.joinpath("should_break.txt")
TEST_ONNX_FILE_PATH = ASSETS_PATH.joinpath("linear_test.onnx")
ONNX_LIB = UTILS_PATH.joinpath("..").joinpath("..").joinpath("onnx_lib").joinpath("onnxruntime")
2 changes: 1 addition & 1 deletion modules/c-wrapper/build-context/core/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -18,10 +18,10 @@ onnx-tests = []
torch-tests = []
tensorflow-tests = []
gpu = []
dynamic = ["ort/load-dynamic"]

[dependencies]
regex = "1.9.3"
# ort = { version = "1.16.2", features = ["load-dynamic"], default-features = false }
ort = { version = "2.0.0-rc.9", features = [ "cuda", "ndarray" ]}
ndarray = "0.16.1"
once_cell = "1.18.0"
Expand Down
111 changes: 1 addition & 110 deletions modules/c-wrapper/build-context/core/build.rs
Original file line number Diff line number Diff line change
@@ -1,113 +1,4 @@
// use std::env;
// use std::fs;
// use std::path::Path;

// /// works out where the `onnxruntime` library is in the build target and copies the library to the root
// /// of the crate so the core library can find it and load it into the binary using `include_bytes!()`.
// ///
// /// # Notes
// /// This is a workaround for the fact that `onnxruntime` doesn't support `cargo` yet. This build step
// /// is reliant on the `ort` crate downloading and building the `onnxruntime` library. This is
// /// why the following dependency is required in `Cargo.toml`:
// /// ```toml
// /// [build-dependencies]
// /// ort = { version = "1.16.2", default-features = true }
// /// ```
// /// Here we can see that the `default-features` is set to `true`. This is because the `ort` crate will download
// /// the correct package and build it for the target platform by default. In the main part of our dependencies
// /// we have the following:
// /// ```toml
// /// [dependencies]
// /// ort = { version = "1.16.2", features = ["load-dynamic"], default-features = false }
// /// ```
// /// Here we can see that the `default-features` is set to `false`. This is because we don't want the `ort` crate
// /// to download and build the `onnxruntime` library again. Instead we want to use the one that was built in the
// /// build step. We also set the `load-dynamic` feature to `true` so that the `ort` crate will load the `onnxruntime`
// /// library dynamically at runtime. This is because we don't want to statically link the `onnxruntime`. Our `onnxruntime`
// /// is embedded into the binary using `include_bytes!()` and we want to load it dynamically at runtime. This means that
// /// we do not need to move the `onnxruntime` library around with the binary, and there is no complicated setup required
// /// or linking.
// fn unpack_onnx() -> std::io::Result<()> {
// let out_dir = env::var("OUT_DIR").expect("OUT_DIR not set");
// let out_path = Path::new(&out_dir);
// let build_dir = out_path
// .ancestors() // This gives an iterator over all ancestors of the path
// .nth(3) // 'nth(3)' gets the fourth ancestor (counting from 0), which should be the debug directory
// .expect("Failed to find debug directory");

// match std::env::var("ONNXRUNTIME_LIB_PATH") {
// Ok(onnx_path) => {
// println!("Surrealml Core Debug: ONNXRUNTIME_LIB_PATH set at: {}", onnx_path);
// println!("cargo:rustc-cfg=onnx_runtime_env_var_set");
// }
// Err(_) => {
// println!("Surrealml Core Debug: ONNXRUNTIME_LIB_PATH not set");
// let target_lib = match env::var("CARGO_CFG_TARGET_OS").unwrap() {
// ref s if s.contains("linux") => "libonnxruntime.so",
// ref s if s.contains("macos") => "libonnxruntime.dylib",
// ref s if s.contains("windows") => "onnxruntime.dll",
// // ref s if s.contains("android") => "android", => not building for android
// _ => panic!("Unsupported target os"),
// };

// let lib_path = build_dir.join(target_lib);
// let lib_path = lib_path.to_str().unwrap();
// println!("Surrealml Core Debug: lib_path={}", lib_path);

// // Check if the path exists
// if fs::metadata(lib_path).is_ok() {
// println!("Surrealml Core Debug: lib_path exists");
// } else {
// println!("Surrealml Core Debug: lib_path does not exist");
// // Extract the directory path
// if let Some(parent) = std::path::Path::new(lib_path).parent() {
// // Print the contents of the directory
// match fs::read_dir(parent) {
// Ok(entries) => {
// println!("Surrealml Core Debug: content of directory {}", parent.display());
// for entry in entries {
// if let Ok(entry) = entry {
// println!("{}", entry.path().display());
// }
// }
// }
// Err(e) => {
// println!("Surrealml Core Debug: Failed to read directory {}: {}", parent.display(), e);
// }
// }
// } else {
// println!("Surrealml Core Debug: Could not determine the parent directory of the path.");
// }
// }

// // put it next to the file of the embedding
// let destination = Path::new(target_lib);
// fs::copy(lib_path, destination)?;
// println!("Surrealml Core Debug: onnx lib copied from {} to {}", lib_path, destination.display());
// }
// }
// Ok(())
// }

// fn main() -> std::io::Result<()> {
// if std::env::var("DOCS_RS").is_ok() {
// // we are not going to be anything here for docs.rs, because we are merely building the docs. When we are just building
// // the docs, the onnx environment variable will not look for the `onnxruntime` library, so we don't need to unpack it.
// return Ok(());
// }

// if env::var("ORT_STRATEGY").as_deref() == Ok("system") {
// // If the ORT crate is built with the `system` strategy, then the crate will take care of statically linking the library.
// // No need to do anything here.
// println!("cargo:rustc-cfg=onnx_statically_linked");

// return Ok(());
// }

// unpack_onnx()?;
// Ok(())
// }

fn main() {

}
}
31 changes: 30 additions & 1 deletion modules/c-wrapper/build-context/core/src/execution/session.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,15 @@ use ort::session::Session;
use crate::errors::error::{SurrealError, SurrealErrorStatus};
use crate::safe_eject;

#[cfg(feature = "dynamic")]
use once_cell::sync::Lazy;
#[cfg(feature = "dynamic")]
use ort::environment::{EnvironmentBuilder, Environment};
#[cfg(feature = "dynamic")]
use std::sync::{Arc, Mutex};

use std::sync::LazyLock;


/// Creates a session for a model.
///
Expand All @@ -26,4 +35,24 @@ pub fn get_session(model_bytes: Vec<u8>) -> Result<Session, SurrealError> {
let session: Session = safe_eject!(builder
.commit_from_memory(&model_bytes), SurrealErrorStatus::Unknown);
Ok(session)
}
}


// #[cfg(feature = "dynamic")]
// pub static ORT_ENV: LazyLock<Arc<Mutex<Option<Arc<Environment>>>>> = LazyLock::new(|| Arc::new(Mutex::new(None)));


#[cfg(feature = "dynamic")]
pub fn set_environment(dylib_path: String) -> Result<(), SurrealError> {

let outcome: EnvironmentBuilder = ort::init_from(dylib_path);
match outcome.commit() {
Ok(env) => {
// ORT_ENV.lock().unwrap().replace(env);
},
Err(e) => {
return Err(SurrealError::new(e.to_string(), SurrealErrorStatus::Unknown));
}
}
Ok(())
}
Loading

0 comments on commit 6e12b7e

Please sign in to comment.