fix(sb_ai): intermittent segfault when creating a Ort::Session

.devcontainer/Dockerfile

@@ -1,3 +1,18 @@
 FROM mcr.microsoft.com/devcontainers/rust:dev-1-bookworm
+
+ARG TARGETPLATFORM
+ARG ONNXRUNTIME_VERSION
+
 RUN apt-get update && apt-get install -y build-essential cmake libclang-dev lldb \
     nodejs npm hyperfine
+
+COPY scripts/install_onnx.sh /tmp/install_onnx.sh
+COPY scripts/download_models.sh /tmp/download_models.sh
+
+WORKDIR /tmp
+RUN ./install_onnx.sh $ONNXRUNTIME_VERSION $TARGETPLATFORM /usr/local/bin/libonnxruntime.so
+RUN ./download_models.sh
+RUN mkdir -p /etc/sb_ai && cp -r /tmp/models /etc/sb_ai/models
+
+ENV ORT_DYLIB_PATH=/usr/local/bin/libonnxruntime.so
+ENV SB_AI_MODELS_DIR=/etc/sb_ai/models

.devcontainer/devcontainer.json

@@ -1,7 +1,11 @@
 {
   "name": "Rust",
   "build": {
-    "dockerfile": "Dockerfile"
+    "dockerfile": "Dockerfile",
+    "context": "..",
+    "args": {
+      "ONNXRUNTIME_VERSION": "1.17.0"
+    }
   },
   "features": {
     "ghcr.io/jungaretti/features/make:1": {},
@@ -21,7 +25,8 @@
       "extensions": [
         "rust-lang.rust-analyzer",
         "eamodio.gitlens",
-        "ms-azuretools.vscode-docker"
+        "ms-azuretools.vscode-docker",
+        "ms-vscode.hexeditor"
       ]
     }
   }

crates/sb_ai/Cargo.toml

@@ -18,5 +18,6 @@ ort = { git = "https://github.com/pykeio/ort", default-features = false, feature
 ndarray = "0.15"
 ndarray-linalg = "0.15"
 tokenizers = { version = ">=0.13.4", default-features = false, features = [ "onig" ] }
-tokio.workspace = true
 rand = "0.8"
+tokio.workspace = true
+once_cell.workspace = true

crates/sb_ai/lib.rs

@@ -1,10 +1,12 @@
+use anyhow::anyhow;
 use anyhow::{bail, Error};
 use deno_core::error::AnyError;
 use deno_core::op2;
 use deno_core::OpState;
 use log::error;
 use ndarray::{Array1, Array2, ArrayView3, Axis, Ix3};
 use ndarray_linalg::norm::{normalize, NormalizeAxis};
+use once_cell::sync::Lazy;
 use ort::{inputs, GraphOptimizationLevel, Session};
 use std::cell::RefCell;
 use std::path::{Path, PathBuf};
@@ -45,8 +47,19 @@ fn mean_pool(last_hidden_states: ArrayView3<f32>, attention_mask: ArrayView3<i64
 }
 
 fn init_gte(state: &mut OpState) -> Result<(), Error> {
-    // Create the ONNX Runtime environment, for all sessions created in this process.
-    ort::init().with_name("GTE").commit()?;
+    static ONNX_ENV_INIT: Lazy<Option<ort::Error>> = Lazy::new(|| {
+        // Create the ONNX Runtime environment, for all sessions created in this process.
+        if let Err(err) = ort::init().with_name("GTE").commit() {
+            error!("sb_ai: failed to create environment - {}", err);
+            return Some(err);
+        }
+
+        None
+    });
+
+    if let Some(err) = &*ONNX_ENV_INIT {
+        return Err(anyhow!("failed to create onnx environment: {err}"));
+    }
 
     let models_dir = std::env::var("SB_AI_MODELS_DIR").unwrap_or("/etc/sb_ai/models".to_string());
 

examples/k6-gte/index.ts

@@ -0,0 +1,16 @@
+const session = new Supabase.ai.Session('gte-small');
+
+Deno.serve(async req => {
+    const payload = await req.json();
+    const text_for_embedding = payload.text_for_embedding;
+
+    // Generate embedding
+    const embedding = await session.run(text_for_embedding, {
+        mean_pool: true,
+        normalize: true
+    });
+
+    return Response.json({
+        length: embedding.length
+    });
+});

k6/specs/gte.ts

@@ -0,0 +1,42 @@
+/*
+./scripts/run.sh
+
+#!/usr/bin/env bash
+
+GIT_V_TAG=0.1.1 cargo build --features cli/tracing && \
+EDGE_RUNTIME_WORKER_POOL_SIZE=8 \
+EDGE_RUNTIME_PORT=9998 RUST_BACKTRACE=full ./target/debug/edge-runtime "$@" start \
+    --main-service ./examples/main \
+    --event-worker ./examples/event-manager
+
+*/
+
+import http from "k6/http";
+
+import { check } from "k6";
+import { Options } from "k6/options";
+
+import { target } from "../config";
+
+export const options: Options = {
+    scenarios: {
+        simple: {
+            executor: "constant-vus",
+            vus: 12,
+            duration: "3m",
+        }
+    }
+};
+
+export default function gte() {
+    const res = http.post(
+        `${target}/k6-gte`,
+        JSON.stringify({
+            "text_for_embedding": "meow"
+        })
+    );
+
+    check(res, {
+        "status is 200": r => r.status === 200
+    });
+}