feat: onnx runtime shared sessions

.devcontainer/devcontainer.json

@@ -4,7 +4,7 @@
     "dockerfile": "Dockerfile",
     "context": "..",
     "args": {
-      "ONNXRUNTIME_VERSION": "1.17.0",
+      "ONNXRUNTIME_VERSION": "1.19.2",
       "DENO_VERSION": "1.45.2"
     }
   },
@@ -39,4 +39,4 @@
       ]
     }
   }
-}
+}

Cargo.toml

@@ -100,6 +100,7 @@ ctor = "0.2.6"
 fastwebsockets = { version = "0.6", features = ["upgrade", "unstable-split"] }
 percent-encoding = "2.3.0"
 scopeguard = "1.2.0"
+clap = { version = "4.0.29", features = ["cargo", "string", "env"] }
 glob = "0.3.1"
 httparse = "1.8.0"
 http = "1.0"
@@ -110,7 +111,6 @@ faster-hex = "0.9"
 num-bigint = { version = "0.4", features = ["rand"] }
 notify = "=5.0.0"
 parking_lot = "0.12.0"
-pin-project = "1.0.11"
 rustls = "0.22.4"
 rustls-pemfile = "2"
 rustls-tokio-stream = "=0.2.23"
@@ -130,10 +130,13 @@ rand = "=0.8.5"
 signature = "2.1"
 spki = "0.7.2"
 urlencoding = "2.1.2"
+pin-project = "1"
+fxhash = "0.2"
 tracing = "0.1"
-tracing-subscriber = "0.3"
+tracing-subscriber = { version = "0.3", features = ["env-filter", "tracing-log"] }
 rkyv = "0.7"
 tempfile = "3"
+xxhash-rust = "0.8"
 
 [patch.crates-io]
 # If the PR is merged upstream, remove the line below.

Dockerfile

@@ -4,7 +4,7 @@ FROM rust:1.79.0-bookworm as builder
 
 ARG TARGETPLATFORM
 ARG GIT_V_VERSION
-ARG ONNXRUNTIME_VERSION=1.17.0
+ARG ONNXRUNTIME_VERSION=1.19.2
 ARG PROFILE=release
 ARG FEATURES
 
@@ -26,20 +26,53 @@ RUN objcopy --strip-debug \
     --add-gnu-debuglink=/root/edge-runtime.debug \
     /root/edge-runtime
 
-RUN ./scripts/install_onnx.sh $ONNXRUNTIME_VERSION $TARGETPLATFORM /root/libonnxruntime.so
-RUN ./scripts/download_models.sh
 
-FROM debian:bookworm-slim
+# Application runtime without ONNX
+FROM debian:bookworm-slim as edge-runtime-base
 
 RUN apt-get update && apt-get install -y libssl-dev && rm -rf /var/lib/apt/lists/*
 RUN apt-get remove -y perl && apt-get autoremove -y
 
 COPY --from=builder /root/edge-runtime /usr/local/bin/edge-runtime
 COPY --from=builder /root/edge-runtime.debug /usr/local/bin/edge-runtime.debug
-COPY --from=builder /root/libonnxruntime.so /usr/local/bin/libonnxruntime.so
-COPY --from=builder /usr/src/edge-runtime/models /etc/sb_ai/models
 
-ENV ORT_DYLIB_PATH=/usr/local/bin/libonnxruntime.so
-ENV SB_AI_MODELS_DIR=/etc/sb_ai/models
+ENV ORT_DYLIB_PATH=/usr/local/bin/onnxruntime/lib/libonnxruntime.so
+
+
+# ONNX Runtime provider
+# Application runtime with ONNX
+FROM builder as ort
+RUN ./scripts/install_onnx.sh $ONNXRUNTIME_VERSION $TARGETPLATFORM /root/onnxruntime
+
+
+# ONNX Runtime CUDA provider
+# Application runtime with ONNX CUDA
+FROM builder as ort-cuda
+RUN ./scripts/install_onnx.sh $ONNXRUNTIME_VERSION $TARGETPLATFORM /root/onnxruntime --gpu
+
+
+FROM builder as preload-models
+RUN ./scripts/download_models.sh
+
+
+# With CUDA
+FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04 as edge-runtime-cuda
+
+COPY --from=edge-runtime-base /usr/local/bin/edge-runtime /usr/local/bin/edge-runtime
+COPY --from=builder /root/edge-runtime.debug /usr/local/bin/edge-runtime.debug
+COPY --from=ort-cuda /root/onnxruntime /usr/local/bin/onnxruntime
+COPY --from=preload-models /usr/src/edge-runtime/models /etc/sb_ai/models
+
+ENV ORT_DYLIB_PATH=/usr/local/bin/onnxruntime/lib/libonnxruntime.so
+ENV NVIDIA_VISIBLE_DEVICES=all
+ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility
+
+ENTRYPOINT ["edge-runtime"]
+
+
+# Base
+FROM edge-runtime-base as edge-runtime
+COPY --from=ort /root/onnxruntime /usr/local/bin/onnxruntime
+COPY --from=preload-models /usr/src/edge-runtime/models /etc/sb_ai/models
 
 ENTRYPOINT ["edge-runtime"]

crates/base/src/commands.rs

@@ -26,7 +26,7 @@ pub async fn start_server(
     inspector_option: Option<InspectorOption>,
     jsx_specifier: Option<String>,
     jsx_module: Option<String>,
-) -> Result<(), Error> {
+) -> Result<Option<i32>, Error> {
     let mut server = Server::new(
         ip,
         port,