Session refactor and execution provider management (#390)

Co-authored-by: Raz Besaleli <besaleli@mozilla.ai>

Author: javiermtorres

encoderfile-runtime/Cargo.toml

@@ -4,6 +4,13 @@ version = "0.6.2"
 edition = "2024"
 publish = false
 
+[features]
+default = []
+coreml = ["encoderfile/coreml"]
+cuda = ["encoderfile/cuda"]
+bench = []
+dev-utils = []
+
 [dependencies]
 
 [dependencies.encoderfile]
@@ -23,7 +30,3 @@ workspace = true
 
 [dependencies.tokio]
 workspace = true
-
-[features]
-bench = []
-dev-utils = []

encoderfile-runtime/src/main.rs

@@ -1,78 +1,22 @@
-use parking_lot::Mutex;
-use std::{
-    fs::File,
-    io::{BufReader, Read, Seek},
-    sync::Arc,
-};
+use std::{fs::File, io::BufReader};
 
 use anyhow::Result;
 use clap::Parser;
-use encoderfile::{
-    common::{
-        ModelType,
-        model_type::{Embedding, SentenceEmbedding, SequenceClassification, TokenClassification},
-    },
-    runtime::{EncoderfileLoader, EncoderfileState, load_assets},
-    transport::cli::Cli,
-};
+use encoderfile::{runtime::load_assets, transport::cli::Cli};
 
 #[tokio::main]
 async fn main() -> Result<()> {
+    // parse CLI
+    let cli = Cli::parse();
+
     // open current executable
     let path = std::env::current_exe()?;
     let file = File::open(path)?;
     let mut file = BufReader::new(file);
+
     // load encoderfile
     let mut loader = load_assets(&mut file)?;
 
-    // entrypoint
-    entrypoint(&mut loader).await
-}
-
-macro_rules! run_cli {
-    ($model_type:ident, $cli:expr, $config:expr, $session:expr, $tokenizer:expr, $model_config:expr) => {{
-        let state = Arc::new(EncoderfileState::<$model_type>::new(
-            $config,
-            $session,
-            $tokenizer,
-            $model_config,
-        ));
-        $cli.command.execute(state).await
-    }};
-}
-
-async fn entrypoint<'a, R: Read + Seek>(loader: &mut EncoderfileLoader<'a, R>) -> Result<()> {
-    let cli = Cli::parse();
-    let session = Mutex::new(loader.session()?);
-    let model_config = loader.model_config()?;
-    let tokenizer = loader.tokenizer()?;
-    let config = loader.encoderfile_config()?;
-
-    match loader.model_type() {
-        ModelType::Embedding => run_cli!(Embedding, cli, config, session, tokenizer, model_config),
-        ModelType::SequenceClassification => run_cli!(
-            SequenceClassification,
-            cli,
-            config,
-            session,
-            tokenizer,
-            model_config
-        ),
-        ModelType::TokenClassification => run_cli!(
-            TokenClassification,
-            cli,
-            config,
-            session,
-            tokenizer,
-            model_config
-        ),
-        ModelType::SentenceEmbedding => run_cli!(
-            SentenceEmbedding,
-            cli,
-            config,
-            session,
-            tokenizer,
-            model_config
-        ),
-    }
+    // execute
+    cli.command.execute(&mut loader).await
 }

encoderfile/Cargo.toml

@@ -107,6 +107,8 @@ transport = [
 dev-utils = ["transport"]
 cli = ["dep:console"]
 internal = []
+cuda = ["ort/cuda", "ort/tensorrt"]
+coreml = ["ort/coreml"]
 
 [dev-dependencies]
 tempfile = "3.23.0"

encoderfile/proto/manifest.proto

@@ -14,54 +14,51 @@ import "proto/metadata.proto";
 // This schema is NOT a general-purpose manifest and MUST NOT be extended
 // to support dynamic loading, headless execution, or runtime configuration.
 message EncoderfileManifest {
-	// ------------------------------------------------------------------
-	// 1–49: General model metadata
-	// ------------------------------------------------------------------
-
-	// Logical model type (e.g. embedding, sentence-embedding, classification).
-	// Used to validate compatibility with the runtime binary.
-	encoderfile.metadata.ModelType model_type = 1;
-
-	// User-declared model name.
-	// This value is informational and has no behavioral impact.
-	string name = 2;
-
-	// User-declared model version.
-	// This value is informational and has no behavioral impact.
-	string version = 3;
-
-    // Backend.
-    Backend backend = 4;
-
-	// ------------------------------------------------------------------
-	// 50–99: Transforms and preprocessing
-	// ------------------------------------------------------------------
-
-	// ------------------------------------------------------------------
-	// 100–999: Embedded artifacts
-	// ------------------------------------------------------------------
-	// Artifacts are referenced via explicit fields, not a generic table.
-	// Each artifact reserves its own numeric range for future expansion.
-	// Any protobuf field ≥100 that participates in layout or runtime loading
-	// MUST have value-independent encoded size.
-
-	// Model weights blob.
-	optional Artifact weights = 100;
-
-	// Optional transform blob applied during model initialization.
-	// If present, the runtime must support the declared TransformType.
-	optional Artifact transform = 110;
-
-	// Model configuration (e.g. architecture, hyperparameters).
-	optional Artifact model_config = 120;
-
-	// Tokenizer data (vocab, merges, config).
-	// Serialized runtime::tokenizer::TokenizerService
-	optional Artifact tokenizer = 130;
+  // ------------------------------------------------------------------
+  // 1–49: General model metadata
+  // ------------------------------------------------------------------
+
+  // Logical model type (e.g. embedding, sentence-embedding, classification).
+  // Used to validate compatibility with the runtime binary.
+  encoderfile.metadata.ModelType model_type = 1;
+
+  // User-declared model name.
+  // This value is informational and has no behavioral impact.
+  string name = 2;
+
+  // User-declared model version.
+  // This value is informational and has no behavioral impact.
+  string version = 3;
+
+  // ------------------------------------------------------------------
+  // 50–99: Transforms and preprocessing
+  // ------------------------------------------------------------------
+
+  // ------------------------------------------------------------------
+  // 100–999: Embedded artifacts
+  // ------------------------------------------------------------------
+  // Artifacts are referenced via explicit fields, not a generic table.
+  // Each artifact reserves its own numeric range for future expansion.
+  // Any protobuf field ≥100 that participates in layout or runtime loading
+  // MUST have value-independent encoded size.
+
+  // Model weights blob.
+  optional Artifact weights = 100;
+
+  // Optional transform blob applied during model initialization.
+  // If present, the runtime must support the declared TransformType.
+  optional Artifact transform = 110;
+
+  // Model configuration (e.g. architecture, hyperparameters).
+  optional Artifact model_config = 120;
+
+  // Tokenizer data (vocab, merges, config).
+  // Serialized runtime::tokenizer::TokenizerService
+  optional Artifact tokenizer = 130;
 }
 
 message LuaLibs {
-	repeated string libs = 1;
+  repeated string libs = 1;
 }
 
 // Transform describes embedded preprocessing/postprocessing steps applied to inputs
@@ -70,52 +67,37 @@ message LuaLibs {
 // Transforms are embedded source code and executed by the runtime.
 // Only explicitly supported TransformTypes are allowed.
 message Transform {
-	// Declares the execution environment for the transform.
-	TransformType transform_type = 1;
+  // Declares the execution environment for the transform.
+  TransformType transform_type = 1;
 
-	// Transform source code.
-	// Interpretation is defined by the TransformType.
-	string transform = 2;
+  // Transform source code.
+  // Interpretation is defined by the TransformType.
+  string transform = 2;
 
-	optional LuaLibs lua_libs = 3;
+  optional LuaLibs lua_libs = 3;
 }
 
 // Artifact describes a contiguous byte range within the embedded payload.
 //
 // Byte offsets are relative to manifest offset found in footer, not the start of the binary.
 message Artifact {
-	// Byte offset relative to manifest offset found in footer
-	fixed64 offset = 1;
+  // Byte offset relative to manifest offset found in footer
+  fixed64 offset = 1;
 
-	// Length of the artifact in bytes.
-	fixed64 length = 2;
+  // Length of the artifact in bytes.
+  fixed64 length = 2;
 
-	// SHA-256 checksum of the artifact contents, hex-encoded.
-	// Used for integrity verification at runtime.
-	// MUST always be 32 bytes.
-	bytes sha256 = 3;
+  // SHA-256 checksum of the artifact contents, hex-encoded.
+  // Used for integrity verification at runtime.
+  // MUST always be 32 bytes.
+  bytes sha256 = 3;
 }
 
 // Supported transform types.
 enum TransformType {
-	// Unspecified transform type. This should not be used.
-	TRANSFORM_TYPE_UNSPECIFIED = 0;
+  // Unspecified transform type. This should not be used.
+  TRANSFORM_TYPE_UNSPECIFIED = 0;
 
-	// Lua source executed by the runtime.
-	LUA = 1;
-}
-
-// Execution backend required by the embedded runtime binary.
-//
-// This enum is informational here; backend compatibility is enforced
-// by the runtime binary itself, not by the payload.
-enum Backend {
-	// CPU-only execution (no hardware acceleration).
-	CPU = 0;
-
-	// NVIDIA CUDA backend.
-	CUDA = 1;
-
-	// Apple Metal backend.
-	METAL = 2;
+  // Lua source executed by the runtime.
+  LUA = 1;
 }

encoderfile/src/builder/builder.rs

@@ -16,7 +16,6 @@ use crate::{
         assets::{AssetKind, AssetPlan, AssetSource, PlannedAsset},
         codec::EncoderfileCodec,
     },
-    generated::manifest::Backend,
 };
 use anyhow::{Context, Result};
 use serde::{Deserialize, Serialize};
@@ -129,7 +128,6 @@ impl EncoderfileBuilder {
             self.config.encoderfile.name.clone(),
             self.config.encoderfile.version.clone(),
             self.config.encoderfile.model_type.clone(),
-            Backend::Cpu,
             &asset_plan,
             &mut out,
         )?;

encoderfile/src/builder/model.rs

@@ -1,4 +1,7 @@
-use crate::format::assets::{AssetKind, AssetSource, PlannedAsset};
+use crate::{
+    format::assets::{AssetKind, AssetSource, PlannedAsset},
+    runtime::ORTSessionBuilder,
+};
 use anyhow::{Result, bail};
 use ort::{
     session::{Output, Session},
@@ -12,7 +15,7 @@ pub trait ModelTypeExt {
 
 impl ModelTypeExt for crate::common::ModelType {
     fn validate_model<'a>(&self, path: &'a Path) -> Result<PlannedAsset<'a>> {
-        let model = load_model(path)?;
+        let model = ORTSessionBuilder::default().from_file(path)?;
 
         match self {
             Self::Embedding => validate_embedding_model(model),
@@ -74,7 +77,3 @@ fn get_outp_dim<'a>(outputs: &'a [Output], outp_name: &str) -> Result<&'a Shape>
         .tensor_shape()
         .ok_or(anyhow::anyhow!("Model must return tensor"))
 }
-
-fn load_model(file: &Path) -> Result<Session> {
-    Ok(Session::builder()?.commit_from_file(file)?)
-}

encoderfile/src/dev_utils/mod.rs

@@ -3,7 +3,7 @@ use crate::{
         Config, ModelConfig, TokenizerConfig,
         model_type::{self, ModelTypeSpec},
     },
-    runtime::{AppState, EncoderfileState},
+    runtime::{AppState, EncoderfileState, ORTSessionBuilder},
 };
 use ort::session::Session;
 use parking_lot::Mutex;
@@ -62,12 +62,10 @@ fn get_tokenizer(dir: &str) -> crate::runtime::TokenizerService {
 }
 
 fn get_model(dir: &str) -> Mutex<Session> {
-    Mutex::new(
-        ort::session::Session::builder()
-            .expect("Failed to load session")
-            .commit_from_file(format!("{}/{}", dir, "model.onnx"))
-            .expect("Failed to load model"),
-    )
+    ORTSessionBuilder::default()
+        .from_file(format!("{}/{}", dir, "model.onnx"))
+        .expect("Failed to load model")
+        .into()
 }
 
 fn get_tokenizer_from_string(s: &str) -> crate::runtime::TokenizerService {

encoderfile/src/format/codec/decoder.rs

@@ -49,7 +49,6 @@ mod tests {
     use crate::{
         common::model_type::ModelType,
         format::assets::{AssetKind, AssetPlan, AssetSource, PlannedAsset},
-        generated::manifest::Backend,
     };
     use std::{
         borrow::Cow,
@@ -81,7 +80,6 @@ mod tests {
                 "test-model".to_string(),
                 "0.1.0".to_string(),
                 ModelType::Embedding,
-                Backend::Cpu,
                 &plan,
                 &mut buf,
             )
@@ -113,7 +111,6 @@ mod tests {
                 "offset-test".to_string(),
                 "1.0.0".to_string(),
                 ModelType::Embedding,
-                Backend::Cpu,
                 &plan,
                 &mut buf,
             )