st-zrt 0.4.1: external-data resolution for byte-loaded models

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

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"

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 = []

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"]

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;

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<std::path::Path>,
+    ) -> Result<crate::model_spool::SpooledModelFile> {
+        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(

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<SpooledModelFile> {
+    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)
+}

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<crate::model_spool::SpooledModelFile>,
     /// 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<EnvInner>,
@@ -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> {
+        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<std::path::Path>,
+        opts: SessionOptions,
+    ) -> Result<Self> {
+        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<crate::model_spool::SpooledModelFile>,
+    ) -> Result<Self> {
         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<EnvInner>, cuda_streams: CudaStreamGuards,
     ) -> Result<Self> {
-        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<EnvInner>,
         owned_initializers: Vec<OwnedInitializer>,
         prepacked_weights: Option<Arc<PrepackedWeightsInner>>, cuda_streams: CudaStreamGuards,
+        spooled_model: Option<crate::model_spool::SpooledModelFile>,
     ) -> Result<Self> {
         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,