feat: exposing deno_cache + ai cache interceptor (#446)

* chore: add `deno_cache` crate

* feat(sb_core): expose `deno_cache` to js land

**NOTE:** saving cache in `temp folder` consider move it to a better
location

* stamp: add example for cache api

* stamp: add `onnxruntime` cache adapter

- adding a cache adapter that intercepts `transformers-cache`
- fetch and caching models from rust land

* fix: tracing wrong import

* test(sb_ai): adding `ort` snapshots in a separated file

- using a separated file to store `ort` predictions snapshots for both
`x64` and `arm64`.

* stamp(test): add test command with debug support

- Checking for `RUST_LOG` env, then executing tests as `debug` with
tracing enabled

* fix(sb_ai): apply url check before `load_from_url()`

- Applying url check to prevent request errors, matching between `Url`
and `Model Bytes`.

* test(sb_ai): add integration tests for `ort cache`

- Adding tests scenarios for `env.useBrowserCache = true`

* stamp: clippy :)

* chore: rid all things that are related to model downloading

* chore: cleanup ort integration tests

* refactor(sb_ai): polishing and cleanup

* stamp: cleanup

* stamp: polishing

* chore: update dependencies

* chore: update `Cargo.lock`

* stamp(sb_ai): polishing

* chore: allow `install_onnx.sh` able to be used in another os

* chore: add x86_64 snapshots of ort integration tests

* stamp: attach the os prefix to snapshots

* stamp: should be cleanup unused sessions while running integration tests

* chore: add x86_64 (linux) snapshots of ort integration tests

* stamp: snapshots must not be created in CI

* stamp: disabling `Web Cache API` for general use

- blocking `Web Cache API` to only allow caching of `.onnx` files
comming from `transformers.js` lib.

* stamp: format

* stamp: allow cache dir to be specified through env var

* stamp: apply `cargo fmt`

* Revert "stamp(test): add test command with debug support"

This reverts commit d6bf8a7.

* stamp: solve merge conflicts

* chore: add aarch64 (osx) snapshots for integration tests

(cherry picked from commit 8f9a535)

---------

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

Author: 3 people

.devcontainer/Dockerfile

@@ -8,11 +8,9 @@ RUN apt-get update && apt-get install -y build-essential cmake libclang-dev lldb
 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
 
-ENV SB_AI_MODELS_DIR=/etc/sb_ai/models
 ENV DENO_INSTALL=/deno
 
 RUN /tmp/install.sh $TARGETPLATFORM

.devcontainer/install.sh

@@ -6,10 +6,8 @@ TARGETPLATFORM=$1
 export $(grep -v '^#' /tmp/.env | xargs)
 
 # ONNX Runtime
-/tmp/install_onnx.sh $ONNXRUNTIME_VERSION $TARGETPLATFORM /tmp/onnxruntime
+/tmp/install_onnx.sh $ONNXRUNTIME_VERSION linux $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

.dockerignore

@@ -6,5 +6,4 @@ assets/
 .github/
 
 scripts/*
-!scripts/install_onnx.sh
-!scripts/download_models.sh
+!scripts/install_onnx.sh

.github/workflows/ci.yml

@@ -57,7 +57,7 @@ jobs:
           envFile: ".env"
 
       - name: Install ONNX Runtime Library
-        run: ./scripts/install_onnx.sh ${{ env.ONNXRUNTIME_VERSION }} x64 /tmp/onnxruntime
+        run: ./scripts/install_onnx.sh ${{ env.ONNXRUNTIME_VERSION }} linux x64 /tmp/onnxruntime
 
       - name: Install S3 credentials for testing
         if: github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository

Cargo.lock

@@ -26,6 +26,7 @@ deno_ast = { version = "=0.40.0", features = ["transpiling"] }
 deno_core = "0.293.0"
 deno_broadcast_channel = "0.155.0"
 deno_canvas = "0.30.0"
+deno_cache = "0.93.0"
 deno_console = "0.161.0"
 deno_crypto = "0.175.0"
 deno_fetch = "0.185.0"
@@ -156,4 +157,4 @@ debug-assertions = false
 
 [profile.release]
 debug = true
-lto = "thin"
+lto = "thin"

Cargo.toml

@@ -40,17 +40,13 @@ COPY --from=builder /root/edge-runtime.debug /usr/local/bin/edge-runtime.debug
 # ONNX Runtime provider
 # Application runtime with ONNX
 FROM builder as ort
-RUN ./scripts/install_onnx.sh $ONNXRUNTIME_VERSION $TARGETPLATFORM /root/onnxruntime
+RUN ./scripts/install_onnx.sh $ONNXRUNTIME_VERSION linux $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
+RUN ./scripts/install_onnx.sh $ONNXRUNTIME_VERSION linux $TARGETPLATFORM /root/onnxruntime --gpu
 
 
 # With CUDA
@@ -59,7 +55,6 @@ 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/lib/libonnxruntime.so* /usr/lib
-COPY --from=preload-models /usr/src/edge-runtime/models /etc/sb_ai/models
 
 ENV NVIDIA_VISIBLE_DEVICES=all
 ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility
@@ -70,6 +65,5 @@ ENTRYPOINT ["edge-runtime"]
 # Base
 FROM edge-runtime-base as edge-runtime
 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

@@ -8,6 +8,7 @@ deno_ast.workspace = true
 deno_fs.workspace = true
 deno_io.workspace = true
 deno_core.workspace = true
+deno_cache.workspace = true
 deno_console.workspace = true
 deno_config.workspace = true
 deno_crypto.workspace = true
@@ -43,6 +44,7 @@ sb_node = { version = "0.1.0", path = "../node" }
 sb_ai = { version = "0.1.0", path = "../sb_ai" }
 sb_fs = { version = "0.1.0", path = "../sb_fs" }
 
+libc.workspace = true
 async-trait.workspace = true
 thiserror.workspace = true
 monch.workspace = true
@@ -77,6 +79,7 @@ tracing.workspace = true
 reqwest_v011.workspace = true
 tracing-subscriber = { workspace = true, optional = true, features = ["env-filter", "tracing-log"] }
 num-traits.workspace = true
+tempfile.workspace = true
 
 tls-listener = { version = "0.10", features = ["rustls"] }
 strum = { version = "0.25", features = ["derive"] }
@@ -85,6 +88,9 @@ cooked-waker = "5"
 tokio-rustls = "0.25.0"
 arc-swap = "1.7"
 
+[target.'cfg(windows)'.dependencies]
+winapi = { workspace = true, features = ["knownfolders", "mswsock", "objbase", "shlobj", "tlhelp32", "winbase", "winerror", "winsock2"] }
+
 [dev-dependencies]
 tokio-util = { workspace = true, features = ["rt", "compat"] }
 tracing-subscriber = { workspace = true, features = ["env-filter", "tracing-log"] }
@@ -100,6 +106,7 @@ deno_io.workspace = true
 deno_console.workspace = true
 deno_crypto.workspace = true
 deno_fetch.workspace = true
+deno_cache.workspace = true
 deno_http.workspace = true
 deno_net.workspace = true
 deno_url.workspace = true
@@ -134,4 +141,4 @@ url.workspace = true
 
 [features]
 tracing = ["dep:tracing-subscriber"]
-termination-signal-ext = []
+termination-signal-ext = []

crates/base/Cargo.toml

@@ -4,6 +4,7 @@ use std::path::PathBuf;
 
 mod supabase_startup_snapshot {
     use super::*;
+    use deno_cache::SqliteBackedCache;
     use deno_core::error::AnyError;
     use deno_core::snapshot::{create_snapshot, CreateSnapshotOptions};
     use deno_core::Extension;
@@ -225,6 +226,10 @@ mod supabase_startup_snapshot {
             sb_core_http::init_ops_and_esm(),
             sb_core_http_start::init_ops_and_esm(),
             deno_node::init_ops_and_esm::<Permissions>(None, None, fs),
+            // NOTE(kallebysantos):
+            // Full `Web Cache API` via `SqliteBackedCache` is disabled. Cache flow is
+            // handled by `sb_ai: Cache Adapter`
+            deno_cache::deno_cache::init_ops_and_esm::<SqliteBackedCache>(None),
             sb_core_runtime::init_ops_and_esm(None),
         ];
 

crates/base/build.rs

@@ -13,6 +13,7 @@ use base_rt::DenoRuntimeDropToken;
 use base_rt::{get_current_cpu_time_ns, BlockingScopeCPUUsage};
 use cooked_waker::{IntoWaker, WakeRef};
 use ctor::ctor;
+use deno_cache::SqliteBackedCache;
 use deno_core::error::{AnyError, JsError};
 use deno_core::url::Url;
 use deno_core::v8::{self, GCCallbackFlags, GCType, HeapStatistics, Isolate};
@@ -526,6 +527,30 @@ where
             Arc::new(DenoCompileFileSystem::from_rc(vfs))
         })?;
 
+        /* NOTE(kallebysantos): Cache via SqliteBackedCache is disabled.
+         *
+         * ```
+         * let cache_base_dir = dirs::cache_dir()
+         *     .context("could not resolve cache directory")?
+         *     .join("web_caches");
+         *
+         * tokio::fs::create_dir_all(cache_base_dir.as_path())
+         *     .await
+         *     .context("could not make cache directory")?;
+         *
+         * struct CacheStorageDir(TempDir);
+         *
+         * let cache_storage_dir = CacheStorageDir(
+         *     tempfile::tempdir_in(cache_base_dir).context("could not make cache directory")?,
+         * );
+         *
+         * let cache_backend = CreateCache(Arc::new({
+         *     let dir = cache_storage_dir.0.path().to_path_buf();
+         *     move || SqliteBackedCache::new(dir.clone())
+         * }));
+         * ```
+         */
+
         let mod_code = module_code;
         let extensions = vec![
             sb_core_permissions::init_ops(net_access_disabled, allow_net),
@@ -575,6 +600,7 @@ where
                 Some(npm_resolver),
                 file_system,
             ),
+            deno_cache::deno_cache::init_ops::<SqliteBackedCache>(None),
             sb_core_runtime::init_ops(Some(main_module_url.clone())),
         ];
 
@@ -694,11 +720,12 @@ where
         let version: Option<&str> = option_env!("GIT_V_TAG");
 
         {
-            // @andreespirela : We do this because "NODE_DEBUG" is trying to be read during
-            // initialization, But we need the gotham state to be up-to-date
             let op_state_rc = js_runtime.op_state();
             let mut op_state = op_state_rc.borrow_mut();
-            op_state.put::<sb_env::EnvVars>(sb_env::EnvVars::new());
+
+            // NOTE(Andreespirela): We do this because "NODE_DEBUG" is trying to be read during
+            // initialization, But we need the gotham state to be up-to-date.
+            op_state.put(sb_env::EnvVars::default());
         }
 
         // Bootstrapping stage
@@ -805,8 +832,8 @@ where
                 }
             }
 
-            op_state.put::<sb_env::EnvVars>(env_vars);
-            op_state.put(DenoRuntimeDropToken(drop_token.clone()))
+            op_state.put(sb_env::EnvVars(env_vars));
+            op_state.put(DenoRuntimeDropToken(drop_token.clone()));
         }
 
         let main_module_id = {