feat: exposing onnx backend to JS land (#436)

* feat: integrate `trasnformers.js` with rust backend

Simple integration between `transformers.js` and rust ort backend by
exposing js owned API

* stamp: refactoring tensors ser/de to try zero-copy

- Implementing v8 traits to cast tensors `ToV8` and `FromV8`.
- Resolving `ort::Tensor` type based on Js tensor's.

* stamp: refactoring to use `serde_v8`

- Since `serde_v8` allows zero-copy we use it to handle the model
inputs and send back the outputs

* fix(ai): seq2seq models causing null pointer error

- Solved the "`GetMutableData` should not be a null pointer" error while
executing seq2seq models.
- Ref.: pykeio/ort#185

* test(sb_ai): implementing tests for ort backend

- Applying integration tests over all NLP pipelines with ort backend

* stamp(sb_ai): example for generate image embeddings

* test(sb_ai): implementing computer vision tests for ort backend

- Applying integration tests over VISION pipelines with ort backend

* stamp: clippy

* fix(ci): makes share common env vars from dotenv file

* fix(ci): update `ORT_DYLIB_PATH`

* fix(ci): makes share common env vars from dotenv file

* chore(sb_ai): update dependencies

* chore(event_worker): add a dependency

* chore(event_worker): install a tracing macro

* chore(base): update `Cargo.toml`

* chore(base): trace `malloced_mb` more precisely

* chore: update an integration test case script

* chore: install tracing subscriber when `base/tracing` feature is enabled

* chore: update `Cargo.lock`

* stamp: add `docker build` script with shared envs

* fix(devcontainer): shared `.env` file path

---------

Co-authored-by: Nyannyacha <[email protected]>

Author: kallebysantos

.devcontainer/Dockerfile

@@ -1,28 +1,18 @@
 FROM mcr.microsoft.com/devcontainers/rust:dev-1-bookworm
 
 ARG TARGETPLATFORM
-ARG ONNXRUNTIME_VERSION
-ARG DENO_VERSION
 
 RUN apt-get update && apt-get install -y build-essential cmake libclang-dev lldb \
     nodejs npm hyperfine
 
+COPY .env /tmp/.env
+COPY .devcontainer/install.sh /tmp/install.sh
 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
-
-# Ollama
-RUN curl -fsSL https://ollama.com/install.sh | sh
-
-# Deno
 ENV DENO_INSTALL=/deno
-RUN mkdir -p /deno \
-    && curl -fsSL https://deno.land/install.sh | bash -s -- v$DENO_VERSION \
-    && chown -R vscode /deno
+
+RUN /tmp/install.sh $TARGETPLATFORM

.devcontainer/devcontainer.json

@@ -2,11 +2,7 @@
   "name": "Rust",
   "build": {
     "dockerfile": "Dockerfile",
-    "context": "..",
-    "args": {
-      "ONNXRUNTIME_VERSION": "1.19.2",
-      "DENO_VERSION": "1.45.2"
-    }
+    "context": ".."
   },
   "containerEnv": {
     "PATH": "${localEnv:PATH}:/deno/bin"
@@ -16,6 +12,8 @@
     "ghcr.io/jungaretti/features/make:1": {}
   },
   "runArgs": [
+    "--env-file",
+    ".env",
     "--rm",
     "--privileged",
     "--security-opt",

.devcontainer/install.sh

@@ -0,0 +1,20 @@
+#!/usr/bin/env bash
+set -e
+
+TARGETPLATFORM=$1
+
+export $(grep -v '^#' /tmp/.env | xargs)
+
+# ONNX Runtime
+/tmp/install_onnx.sh $ONNXRUNTIME_VERSION $TARGETPLATFORM /tmp/onnxruntime
+mv /tmp/onnxruntime/lib/libonnxruntime.so* /usr/lib
+/tmp/download_models.sh
+mkdir -p /etc/sb_ai && cp -r /tmp/models /etc/sb_ai/models
+
+# Ollama
+curl -fsSL https://ollama.com/install.sh | sh
+
+# Deno
+mkdir -p /deno
+curl -fsSL https://deno.land/install.sh | bash -s -- v$DENO_VERSION
+chown -R vscode /deno

.env

@@ -0,0 +1,4 @@
+ONNXRUNTIME_VERSION=1.19.2
+DENO_VERSION=1.45.2
+EDGE_RUNTIME_PORT=9998
+AI_INFERENCE_API_HOST=http://localhost:11434

.github/workflows/ci.yml

@@ -12,6 +12,7 @@ env:
   CARGO_NET_RETRY: 10
   CARGO_TERM_COLOR: always
   RUSTUP_MAX_RETRIES: 10
+  ORT_DYLIB_PATH: /tmp/onnxruntime/lib/libonnxruntime.so
 
 jobs:
   cargo-fmt:
@@ -49,4 +50,9 @@ jobs:
       - uses: actions/checkout@v4
       - run: rustup show
       - uses: Swatinem/rust-cache@v2
+      - uses: cardinalby/export-env-action@v2
+        with:
+          envFile: ".env"
+      - name: Install ONNX Runtime Library
+        run: ./scripts/install_onnx.sh ${{ env.ONNXRUNTIME_VERSION }} x64 /tmp/onnxruntime
       - run: ./scripts/test.sh

.github/workflows/release.yml

@@ -52,6 +52,9 @@ jobs:
           username: ${{ secrets.DOCKER_USERNAME }}
           password: ${{ secrets.DOCKER_PASSWORD }}
 
+      - uses: cardinalby/export-env-action@v2
+        with:
+          envFile: ".env"
       - id: build
         uses: docker/build-push-action@v3
         with:
@@ -61,7 +64,8 @@ jobs:
           cache-from: type=gha
           cache-to: type=gha,mode=max
           build-args: |
-            GIT_V_VERSION=${{ needs.release.outputs.version }}
+            GIT_V_TAG=${{ needs.release.outputs.version }}
+            ONNXRUNTIME_VERSION=${{ env.ONNXRUNTIME_VERSION }}
 
   publish_arm:
     needs:
@@ -95,6 +99,9 @@ jobs:
             image=moby/buildkit:master
             network=host
 
+      - uses: cardinalby/export-env-action@v2
+        with:
+          envFile: ".env"
       - id: build
         uses: docker/build-push-action@v3
         with:
@@ -104,7 +111,8 @@ jobs:
           tags: ${{ steps.meta.outputs.tags }}
           no-cache: true
           build-args: |
-            GIT_V_VERSION=${{ needs.release.outputs.version }}
+            GIT_V_TAG=${{ needs.release.outputs.version }}
+            ONNXRUNTIME_VERSION=${{ env.ONNXRUNTIME_VERSION }}
 
   merge_manifest:
     needs: [release, publish_x86, publish_arm]

Cargo.lock

@@ -3,8 +3,8 @@
 FROM rust:1.79.0-bookworm as builder
 
 ARG TARGETPLATFORM
-ARG GIT_V_VERSION
-ARG ONNXRUNTIME_VERSION=1.19.2
+ARG ONNXRUNTIME_VERSION
+ARG GIT_V_TAG
 ARG PROFILE=release
 ARG FEATURES
 
@@ -15,7 +15,7 @@ WORKDIR /usr/src/edge-runtime
 COPY . .
 
 RUN --mount=type=cache,target=/usr/local/cargo/registry,id=${TARGETPLATFORM} --mount=type=cache,target=/usr/src/edge-runtime/target,id=${TARGETPLATFORM} \
-    GIT_V_TAG=${GIT_V_VERSION} cargo build --profile ${PROFILE} --features "${FEATURES}" && \
+    ${GIT_V_TAG} cargo build --profile ${PROFILE} --features "${FEATURES}" && \
     mv /usr/src/edge-runtime/target/${PROFILE}/edge-runtime /root
 
 RUN objcopy --compress-debug-sections \
@@ -36,8 +36,6 @@ 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
 
-ENV ORT_DYLIB_PATH=/usr/local/bin/onnxruntime/lib/libonnxruntime.so
-
 
 # ONNX Runtime provider
 # Application runtime with ONNX
@@ -60,10 +58,9 @@ 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=ort-cuda /root/onnxruntime/lib/libonnxruntime.so* /usr/lib
 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
 
@@ -72,7 +69,7 @@ ENTRYPOINT ["edge-runtime"]
 
 # Base
 FROM edge-runtime-base as edge-runtime
-COPY --from=ort /root/onnxruntime /usr/local/bin/onnxruntime
+COPY --from=ort /root/onnxruntime/lib/libonnxruntime.so* /usr/lib
 COPY --from=preload-models /usr/src/edge-runtime/models /etc/sb_ai/models
 
 ENTRYPOINT ["edge-runtime"]

Dockerfile

@@ -74,6 +74,7 @@ notify.workspace = true
 pin-project.workspace = true
 rustls-pemfile.workspace = true
 tracing.workspace = true
+tracing-subscriber = { workspace = true, optional = true, features = ["env-filter", "tracing-log"] }
 
 reqwest_v011 = { package = "reqwest", version = "0.11", features = ["stream", "json", "multipart"] }
 tls-listener = { version = "0.10", features = ["rustls"] }
@@ -129,4 +130,5 @@ tokio.workspace = true
 url.workspace = true
 
 [features]
+tracing = ["dep:tracing-subscriber"]
 termination-signal-ext = []

crates/base/Cargo.toml

@@ -190,6 +190,7 @@ impl MemCheck {
             }
         }
 
+        trace!(malloced_mb = bytes_to_display(total_bytes as u64));
         total_bytes
     }
 }
@@ -965,11 +966,9 @@ where
 
             if is_user_worker {
                 let mem_state = mem_check_state.as_ref().unwrap();
-                let total_malloced_bytes = mem_state.check(js_runtime.v8_isolate().as_mut());
 
+                mem_state.check(js_runtime.v8_isolate().as_mut());
                 mem_state.waker.register(waker);
-
-                trace!(malloced_mb = bytes_to_display(total_malloced_bytes as u64));
             }
 
             // NOTE(Nyannyacha): If tasks are empty or V8 is not evaluating the