Skip to content

Commit

Permalink
feat: exposing deno_cache + ai cache interceptor (#446)
Browse files Browse the repository at this point in the history
* 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]>
  • Loading branch information
3 people authored Nov 24, 2024
1 parent 997f2d9 commit 13ee2a2
Show file tree
Hide file tree
Showing 132 changed files with 2,004 additions and 700 deletions.
2 changes: 0 additions & 2 deletions .devcontainer/Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -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
4 changes: 1 addition & 3 deletions .devcontainer/install.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
3 changes: 1 addition & 2 deletions .dockerignore
Original file line number Diff line number Diff line change
Expand Up @@ -6,5 +6,4 @@ assets/
.github/

scripts/*
!scripts/install_onnx.sh
!scripts/download_models.sh
!scripts/install_onnx.sh
2 changes: 1 addition & 1 deletion .github/workflows/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
39 changes: 32 additions & 7 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

3 changes: 2 additions & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -156,4 +157,4 @@ debug-assertions = false

[profile.release]
debug = true
lto = "thin"
lto = "thin"
10 changes: 2 additions & 8 deletions Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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"]
9 changes: 8 additions & 1 deletion crates/base/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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"] }
Expand All @@ -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"] }
Expand All @@ -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
Expand Down Expand Up @@ -134,4 +141,4 @@ url.workspace = true

[features]
tracing = ["dep:tracing-subscriber"]
termination-signal-ext = []
termination-signal-ext = []
5 changes: 5 additions & 0 deletions crates/base/build.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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),
];

Expand Down
37 changes: 32 additions & 5 deletions crates/base/src/deno_runtime.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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),
Expand Down Expand Up @@ -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())),
];

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 = {
Expand Down
1 change: 1 addition & 0 deletions crates/base/src/rt_worker/utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ type WorkerCoreConfig = (
String,
);

// TODO: Refactor this
pub fn parse_worker_conf(conf: &WorkerRuntimeOpts) -> WorkerCoreConfig {
let worker_core: WorkerCoreConfig = match conf {
WorkerRuntimeOpts::UserWorker(worker_opts) => (
Expand Down
2 changes: 1 addition & 1 deletion crates/base/src/rt_worker/worker.rs
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,7 @@ impl Worker {

#[allow(clippy::too_many_arguments)]
pub fn start(
&self,
self: Arc<Self>,
mut opts: WorkerContextInitOpts,
duplex_stream_pair: (
UnboundedSender<DuplexStreamEntry>,
Expand Down
Loading

0 comments on commit 13ee2a2

Please sign in to comment.