diff --git a/CHANGELOG.md b/CHANGELOG.md index b22303b..04cd531 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,18 @@ +## [0.4.1] — 2026-09-03 + +### Added — external-data resolution for byte-loaded models +- `Session::from_bytes_with_external_data(model_data, external_data_dir, ..)` and + `ModelCompilationOptions::set_input_model_from_buffer_with_external_data(..)`: models + whose initializers use ONNX external data (every model over the 2 GiB protobuf limit, + and any model saved with external data) can now be loaded from memory. ORT resolves + external-data paths relative to the model file's directory, which a buffer load + cannot supply; the new entry points spool the serialized model to a unique temporary + file inside the caller-supplied external-data directory and load through the + path-based API. The temporary file is owned by the session (or by a guard returned to + the compile caller) and removed on drop; the directory must be writable. +- `Session::from_bytes` remains the right choice for self-contained models and is + unchanged. + # Changelog All notable changes to st-zrt are documented here. The format is based on diff --git a/Cargo.lock b/Cargo.lock index 94d1d62..4644c23 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -420,7 +420,7 @@ checksum = "3a219298ac11a56ea9a6d2120044824d6f01aeb034955e7af7bc16858527deea" [[package]] name = "st-zrt" -version = "0.4.0" +version = "0.4.1" dependencies = [ "futures-util", "half", diff --git a/Cargo.toml b/Cargo.toml index d4e6a69..19fbb2a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,7 +6,7 @@ members = ["st-zrt-sys", "st-zrt", "codegen"] # share a Cargo graph. Variant C benchmarks against st-zrt in a separate crate. [workspace.package] -version = "0.4.0" +version = "0.4.1" edition = "2024" rust-version = "1.85" license = "Apache-2.0" diff --git a/bench-c/Cargo.toml b/bench-c/Cargo.toml index 6bb2cd1..badf76e 100644 --- a/bench-c/Cargo.toml +++ b/bench-c/Cargo.toml @@ -19,7 +19,7 @@ publish = false # Pinned to the in-tree wrapper version on purpose: the wildcard-dependency ban forbids # a bare path dep, and a wrapper minor bump must update this pin in the same change # (resolution fails loudly otherwise). -st-zrt = { path = "../st-zrt", version = "=0.4.0" } +st-zrt = { path = "../st-zrt", version = "=0.4.1" } [features] default = [] diff --git a/st-zrt/Cargo.toml b/st-zrt/Cargo.toml index 7f11e05..c21ce74 100644 --- a/st-zrt/Cargo.toml +++ b/st-zrt/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "st-zrt" -version = "0.4.0" +version = "0.4.1" edition.workspace = true rust-version.workspace = true license.workspace = true @@ -63,6 +63,10 @@ required-features = ["cuda"] name = "bert_cuda_probe" required-features = ["cuda"] +[[test]] +name = "external_data" +path = "tests/external_data.rs" + [[test]] name = "custom_op_run" required-features = ["custom-ops"] diff --git a/st-zrt/src/lib.rs b/st-zrt/src/lib.rs index 07725cc..394d35f 100644 --- a/st-zrt/src/lib.rs +++ b/st-zrt/src/lib.rs @@ -70,6 +70,7 @@ mod memory; mod metadata; #[cfg(feature = "model-editor")] mod model_editor; +pub(crate) mod model_spool; mod prepacked; mod run_options; mod runtime; diff --git a/st-zrt/src/model_editor.rs b/st-zrt/src/model_editor.rs index c9ff86e..f87671a 100644 --- a/st-zrt/src/model_editor.rs +++ b/st-zrt/src/model_editor.rs @@ -594,6 +594,23 @@ impl ModelCompilationOptions { check(unsafe { set_input(self.ptr, model.as_ptr()) }) } + /// Input: the bytes of a serialized ONNX model whose initializers use **external + /// data** (BORROWED for the duration of compile). + /// + /// ORT cannot resolve external-data references without a model-file base directory; + /// this variant spools `bytes` to a unique temporary file inside + /// `external_data_dir` — the directory holding the external-data files — and sets + /// the input model by path. The returned guard owns the temporary file: keep it + /// alive until `compile` has returned, then drop it (the file is removed on drop). + pub fn set_input_model_from_buffer_with_external_data( + &self, bytes: &[u8], external_data_dir: impl AsRef, + ) -> Result { + let spool = crate::model_spool::spool_model_bytes(external_data_dir.as_ref(), bytes)?; + let path = spool.path().to_string_lossy().into_owned(); + self.set_input_model_path(&path)?; + Ok(spool) + } + /// Input: the bytes of a serialized ONNX model (BORROWED for the duration of compile). pub fn set_input_model_from_buffer(&self, bytes: &[u8]) -> Result<()> { let set_input = require_sub_api_fn( diff --git a/st-zrt/src/model_spool.rs b/st-zrt/src/model_spool.rs new file mode 100644 index 0000000..69cb687 --- /dev/null +++ b/st-zrt/src/model_spool.rs @@ -0,0 +1,99 @@ +//! Spool-through resolution for byte-loaded models that reference ONNX external data. +//! +//! ORT resolves an initializer's external-data path *relative to the model file's +//! directory*. Byte-buffer entry points (`CreateSessionFromArray`, +//! `ModelCompilationOptions_SetInputModelFromBuffer`) have no such base, so any model +//! with external initializers — every model over the 2 GiB protobuf limit, and any +//! model saved with external data — cannot resolve them. The generic resolver: write +//! the (small) serialized graph to a uniquely named temporary file *inside the +//! user-supplied external-data directory* and load through the path-based entry point. +//! The guard keeps the file alive until the owning object (session or compile) is done, +//! because ORT may map or reopen external data lazily. + +use crate::{Error, Result}; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicU64, Ordering}; + +static SPOOL_COUNTER: AtomicU64 = AtomicU64::new(0); + +/// A uniquely named, self-deleting copy of a serialized model, created inside the +/// directory that holds the model's external-data files. +pub struct SpooledModelFile { + path: PathBuf, +} + +impl SpooledModelFile { + /// The spooled file's path — valid to load from until `Drop`. + pub fn path(&self) -> &Path { + &self.path + } +} + +impl std::fmt::Debug for SpooledModelFile { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("SpooledModelFile") + .field("path", &self.path) + .finish() + } +} + +impl Drop for SpooledModelFile { + fn drop(&mut self) { + // Best-effort: a failure to remove never invalidates the loaded session. + let _ = std::fs::remove_file(&self.path); + } +} + +/// Copy `model_bytes` into a unique file inside `external_data_dir`. +/// +/// The directory must be writable; when it is not (e.g. a read-only model bundle), +/// copy the external-data files to a writable directory and pass that instead — ORT +/// resolves the initializer locations relative to the spooled model file either way. +pub(crate) fn spool_model_bytes( + external_data_dir: &Path, model_bytes: &[u8], +) -> Result { + let seq = SPOOL_COUNTER.fetch_add(1, Ordering::Relaxed); + for _ in 0..32 { + let name = format!( + ".st-zrt-spooled-{}-{}-{}.onnx", + std::process::id(), + seq, + // Uniqueness even across processes started in the same millisecond. + nanos_suffix(), + ); + let path = external_data_dir.join(name); + match std::fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(&path) + { + Ok(mut file) => { + std::io::Write::write_all(&mut file, model_bytes).map_err(|e| { + let _ = std::fs::remove_file(&path); + Error::new(-1, format!("zrt: spooling model bytes failed: {e}")) + })?; + return Ok(SpooledModelFile { path }); + }, + Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => continue, + Err(e) => { + return Err(Error::new( + -1, + format!( + "zrt: cannot create a temporary model file in {} (external-data \ + directories must be writable): {e}", + external_data_dir.display() + ), + )); + }, + } + } + Err(Error::new(-1, "zrt: could not find a free spool file name")) +} + +fn nanos_suffix() -> u128 { + use std::time::{SystemTime, UNIX_EPOCH}; + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_nanos()) + .unwrap_or(0) +} diff --git a/st-zrt/src/session.rs b/st-zrt/src/session.rs index 322239c..f912489 100644 --- a/st-zrt/src/session.rs +++ b/st-zrt/src/session.rs @@ -174,6 +174,10 @@ pub(crate) struct SessionInner { /// Owned CUDA streams referenced by this session's provider configuration. Native session /// release occurs before these guards drop. _cuda_streams: CudaStreamGuards, + /// Spooled model file backing an external-data buffer load (see + /// [`Session::from_bytes_with_external_data`]). Dropped only after the native session is + /// released, since ORT may map or reopen external initializer data lazily. + _spooled_model: Option, /// Keeps the Env alive for this Session's whole lifetime. The explicit `SessionInner::drop` /// releases the native session before Rust drops this guard. _env: Arc, @@ -367,6 +371,35 @@ impl Session { /// Load `model_path` (filesystem path, UTF-8) and pre-marshal its I/O names and /// output type/shape (cached so the hot path needs no introspection). pub fn new(env: &Environment, model_path: &str, opts: SessionOptions) -> Result { + Self::new_with_spool(env, model_path, opts, None) + } + + /// Load a model from an in-memory byte buffer whose initializers use ONNX + /// **external data** (any model over the 2 GiB protobuf limit, and any model saved + /// with `save_as_external_data`). + /// + /// ORT resolves external-data paths relative to the model *file's* directory, which a + /// plain buffer load cannot supply (the reference stays unresolved and loading + /// fails). This constructor resolves it generically: `model_data` is spooled to a + /// unique temporary file inside `external_data_dir` — the directory that actually + /// holds the external-data files — and loaded through the path-based entry point. + /// The temporary file is removed when the session is dropped; the directory must be + /// writable. [`Self::from_bytes`] remains the right choice for self-contained models. + pub fn from_bytes_with_external_data( + env: &Environment, model_data: &[u8], external_data_dir: impl AsRef, + opts: SessionOptions, + ) -> Result { + let spool = crate::model_spool::spool_model_bytes(external_data_dir.as_ref(), model_data)?; + let spool_path = spool.path().to_string_lossy().into_owned(); + Self::new_with_spool(env, &spool_path, opts, Some(spool)) + } + + /// Path-based creation that may carry the spooled model file keeping external data + /// resolvable; `Session::new` is this with `None`. + fn new_with_spool( + env: &Environment, model_path: &str, opts: SessionOptions, + spool: Option, + ) -> Result { let cpath = CString::new(model_path) .map_err(|_| crate::Error::new(-1, "model path contains a NUL"))?; let opts_handle = build_session_options_for_env(env, &opts)?; @@ -381,7 +414,14 @@ impl Session { }); unsafe { api().release_session_options()(opts_handle) }; create?; - Self::from_handle(sess, env.share(), cuda_stream_guards(&opts)) + Self::from_handle_with_resources( + sess, + env.share(), + Vec::new(), + None, + cuda_stream_guards(&opts), + spool, + ) } /// Load a model from an in-memory byte buffer (`CreateSessionFromArray`, idx 8) — no @@ -452,6 +492,7 @@ impl Session { initializers, None, cuda_stream_guards(&opts), + None, ) } @@ -489,6 +530,7 @@ impl Session { initializers, None, cuda_stream_guards(&opts), + None, ) } @@ -560,6 +602,7 @@ impl Session { initializers, Some(prepacked.share()), cuda_stream_guards(&opts), + None, ) } @@ -605,6 +648,7 @@ impl Session { initializers, None, cuda_stream_guards(&opts), + None, ) } @@ -638,6 +682,7 @@ impl Session { initializers, Some(prepacked.share()), cuda_stream_guards(&opts), + None, ) } @@ -647,13 +692,14 @@ impl Session { fn from_handle( sess: *mut sys::SessionHandle, env: Arc, cuda_streams: CudaStreamGuards, ) -> Result { - Self::from_handle_with_resources(sess, env, Vec::new(), None, cuda_streams) + Self::from_handle_with_resources(sess, env, Vec::new(), None, cuda_streams, None) } fn from_handle_with_resources( sess: *mut sys::SessionHandle, env: Arc, owned_initializers: Vec, prepacked_weights: Option>, cuda_streams: CudaStreamGuards, + spooled_model: Option, ) -> Result { let sess = crate::ensure_non_null(sess, "session")?; // Resolve every fallible setup value while all lifetime guards remain in this outer frame. @@ -685,6 +731,7 @@ impl Session { }; Ok(Self { inner: Arc::new(SessionInner { + _spooled_model: spooled_model, sess, input_names, input_ptrs, diff --git a/st-zrt/tests/external_data.rs b/st-zrt/tests/external_data.rs new file mode 100644 index 0000000..055166d --- /dev/null +++ b/st-zrt/tests/external_data.rs @@ -0,0 +1,100 @@ +//! Byte-buffer loading for models whose initializers use ONNX external data. +//! +//! ORT resolves external-data paths relative to the model file's directory, so a plain +//! buffer load cannot resolve them. `from_bytes_with_external_data` spools the buffer +//! next to the external-data directory and loads through the path. The external-data +//! blob is materialized by the test into a temp directory — nothing binary is committed. + +use st_zrt::{ + Environment, GraphOptimizationLevel, MemoryInfo, OwnedValue, Session, SessionOptions, Tensor, +}; + +struct TempDir(std::path::PathBuf); +impl Drop for TempDir { + fn drop(&mut self) { + let _ = std::fs::remove_dir_all(&self.0); + } +} + +/// A temp directory holding a copy of `ext_add.onnx` plus its external-data file +/// (four f32 values of 2.0 — exactly what the model's external reference declares). +fn fixture_dir() -> TempDir { + let dir = std::env::temp_dir().join(format!( + "st-zrt-extdata-{}-{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + std::fs::create_dir_all(&dir).expect("temp dir"); + let model = std::path::Path::new(env!("CARGO_MANIFEST_DIR")) + .join("tests/fixtures/external_data/ext_add.onnx"); + std::fs::copy(&model, dir.join("ext_add.onnx")).expect("copy model"); + let data: Vec = vec![2.0; 4]; + let bytes = unsafe { std::slice::from_raw_parts(data.as_ptr() as *const u8, data.len() * 4) }; + std::fs::write(dir.join("ext_add.onnx.data"), bytes).expect("write external data"); + TempDir(dir) +} + +fn run(session: &Session) -> Vec { + let mem = MemoryInfo::cpu().expect("cpu mem"); + let input = Tensor::from_buffer(&[1.0_f32, 2.0, 3.0, 4.0], &[4], &mem).expect("input"); + let mut out: Vec> = (0..session.output_count()).map(|_| None).collect(); + session.run(&[&input], &mut out).expect("run"); + out[0] + .as_ref() + .unwrap() + .as_slice::() + .expect("f32") + .to_vec() +} + +#[test] +fn external_data_byte_load_resolves_and_runs() { + let fx = fixture_dir(); + let dir = &fx.0; + let bytes = std::fs::read(dir.join("ext_add.onnx")).expect("model bytes"); + + let env = Environment::new().expect("env"); + let opts = SessionOptions::new().with_opt_level(GraphOptimizationLevel::Basic); + + // Reference: path-based load (ORT resolves external data relative to the model file). + let via_path = Session::new( + &env, + dir.join("ext_add.onnx").to_str().unwrap(), + opts.clone(), + ) + .expect("path load"); + let expected = run(&via_path); + assert_eq!(expected, vec![3.0, 4.0, 5.0, 6.0]); + + // The generic resolver: same model from bytes, external directory supplied. + let via_bytes = Session::from_bytes_with_external_data(&env, &bytes, dir, opts) + .expect("external-data byte load"); + assert_eq!(run(&via_bytes), expected); + + // The spooled temporary file is gone once the session drops. + drop(via_bytes); + let leftovers: Vec<_> = std::fs::read_dir(dir) + .unwrap() + .filter_map(|e| e.ok()) + .filter(|e| { + e.file_name() + .to_string_lossy() + .starts_with(".st-zrt-spooled-") + }) + .collect(); + assert!(leftovers.is_empty(), "spool cleaned up: {:?}", leftovers); +} + +#[test] +fn plain_byte_load_of_external_data_model_is_rejected_or_selfcontained() { + let fx = fixture_dir(); + let bytes = std::fs::read(fx.0.join("ext_add.onnx")).expect("model bytes"); + let env = Environment::new().expect("env"); + match Session::from_bytes(&env, &bytes, SessionOptions::new()) { + Err(_) => {}, // rejected by ORT: expected on this runtime line + Ok(s) => assert_eq!(run(&s), vec![3.0, 4.0, 5.0, 6.0]), + } +} diff --git a/st-zrt/tests/fixtures/external_data/ext_add.onnx b/st-zrt/tests/fixtures/external_data/ext_add.onnx new file mode 100644 index 0000000..3eb6da1 Binary files /dev/null and b/st-zrt/tests/fixtures/external_data/ext_add.onnx differ