diff --git a/Cargo.lock b/Cargo.lock index 5ca1c9fd8..7dd048eac 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3164,7 +3164,6 @@ dependencies = [ "faster-hex", "futures", "futures-util", - "fxhash", "log", "ndarray", "once_cell", @@ -3175,12 +3174,12 @@ dependencies = [ "reqwest 0.12.7", "scopeguard", "serde", + "sha2", "tokenizers", "tokio", "tokio-util", "tracing", "tracing-subscriber", - "xxhash-rust", ] [[package]] @@ -3781,15 +3780,6 @@ dependencies = [ "slab", ] -[[package]] -name = "fxhash" -version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c31b6d751ae2c7f11320402d34e41349dd1016f8d5d45e48c4312bc8625af50c" -dependencies = [ - "byteorder", -] - [[package]] name = "generic-array" version = "0.14.7" diff --git a/ext/ai/Cargo.toml b/ext/ai/Cargo.toml index fb4c2b48d..e2a681d98 100644 --- a/ext/ai/Cargo.toml +++ b/ext/ai/Cargo.toml @@ -20,17 +20,16 @@ ctor.workspace = true faster-hex.workspace = true futures.workspace = true futures-util = { workspace = true, features = ["io"] } -fxhash.workspace = true log.workspace = true once_cell.workspace = true reqwest.workspace = true scopeguard.workspace = true serde.workspace = true +sha2.workspace = true tokio.workspace = true tokio-util = { workspace = true, features = ["compat"] } tracing.workspace = true tracing-subscriber.workspace = true -xxhash-rust = { workspace = true, features = ["std", "xxh3"] } blas-src = { version = "0.8", features = ['openblas'] } convert_case = "0.6" diff --git a/ext/ai/onnxruntime/mod.rs b/ext/ai/onnxruntime/mod.rs index b1dd136f8..1a327c8fd 100644 --- a/ext/ai/onnxruntime/mod.rs +++ b/ext/ai/onnxruntime/mod.rs @@ -50,7 +50,7 @@ pub async fn op_ai_ort_init_session( } Err(_) => { trace!(kind = "bytes", len = model_bytes.len()); - Model::from_bytes(&model_bytes).await? + Model::from_bytes(model_bytes).await? } }; diff --git a/ext/ai/onnxruntime/model.rs b/ext/ai/onnxruntime/model.rs index f3a17e6e4..7c6740c3b 100644 --- a/ext/ai/onnxruntime/model.rs +++ b/ext/ai/onnxruntime/model.rs @@ -75,7 +75,7 @@ impl Model { load_session_from_url(model_url).await.map(Self::new) } - pub async fn from_bytes(model_bytes: &[u8]) -> Result { + pub async fn from_bytes(model_bytes: Box<[u8]>) -> Result { load_session_from_bytes(model_bytes).await.map(Self::new) } } diff --git a/ext/ai/onnxruntime/session.rs b/ext/ai/onnxruntime/session.rs index 6205e8550..6cad53823 100644 --- a/ext/ai/onnxruntime/session.rs +++ b/ext/ai/onnxruntime/session.rs @@ -1,16 +1,14 @@ use deno_core::error::AnyError; -use futures::io::AllowStdIo; use once_cell::sync::Lazy; use reqwest::Url; +use sha2::Digest; +use sha2::Sha256; use std::collections::HashMap; -use std::hash::Hasher; use std::sync::Arc; use tokio::sync::Mutex; -use tokio_util::compat::FuturesAsyncWriteCompatExt; use tracing::debug; use tracing::instrument; use tracing::trace; -use xxhash_rust::xxh3::Xxh3; use anyhow::anyhow; use anyhow::Context; @@ -125,18 +123,16 @@ fn create_session(model_bytes: &[u8]) -> Result, Error> { #[instrument(level = "debug", skip_all, fields(model_bytes = model_bytes.len()), err)] pub(crate) async fn load_session_from_bytes( - model_bytes: &[u8], + model_bytes: Box<[u8]>, ) -> Result { - let session_id = { - let mut model_bytes = tokio::io::BufReader::new(model_bytes); - let mut hasher = AllowStdIo::new(Xxh3::new()).compat_write(); - let _ = tokio::io::copy(&mut model_bytes, &mut hasher) - .await - .context("failed to calculate checksum")?; - - let hasher = hasher.into_inner().into_inner(); - faster_hex::hex_string(&hasher.finish().to_be_bytes()) - }; + // Hash off the isolate thread: this runs on every create(bytes) call, even + // on a cache hit, and would otherwise count toward the worker's CPU budget. + let (session_id, model_bytes) = tokio::task::spawn_blocking(move || { + let digest = Sha256::digest(&model_bytes); + (faster_hex::hex_string(&digest), model_bytes) + }) + .await + .context("failed to calculate checksum")?; let mut sessions = SESSIONS.lock().await; @@ -145,7 +141,7 @@ pub(crate) async fn load_session_from_bytes( } trace!(session_id, "new session"); - let session = create_session(model_bytes)?; + let session = create_session(&model_bytes)?; sessions.insert(session_id.clone(), session.clone()); @@ -156,7 +152,7 @@ pub(crate) async fn load_session_from_bytes( pub(crate) async fn load_session_from_url( model_url: Url, ) -> Result { - let session_id = fxhash::hash(model_url.as_str()).to_string(); + let session_id = faster_hex::hex_string(&Sha256::digest(model_url.as_str())); let mut sessions = SESSIONS.lock().await; diff --git a/ext/ai/utils.rs b/ext/ai/utils.rs index f9e188ba3..2080c3652 100644 --- a/ext/ai/utils.rs +++ b/ext/ai/utils.rs @@ -1,4 +1,3 @@ -use std::hash::Hasher; use std::path::PathBuf; use anyhow::Context; @@ -6,6 +5,8 @@ use deno_core::error::AnyError; use futures::io::AllowStdIo; use futures::StreamExt; use reqwest::Url; +use sha2::Digest; +use sha2::Sha256; use tokio::io::AsyncWriteExt; use tokio_util::compat::FuturesAsyncWriteCompatExt; use tracing::info; @@ -13,7 +14,6 @@ use tracing::info_span; use tracing::instrument; use tracing::trace; use tracing::Instrument; -use xxhash_rust::xxh3::Xxh3; #[instrument(fields(%kind, url = %url))] pub async fn fetch_and_cache_from_url( @@ -21,7 +21,8 @@ pub async fn fetch_and_cache_from_url( url: Url, cache_id: Option, ) -> Result { - let cache_id = cache_id.unwrap_or(fxhash::hash(url.as_str()).to_string()); + let cache_id = cache_id + .unwrap_or_else(|| faster_hex::hex_string(&Sha256::digest(url.as_str()))); let download_dir = std::env::var("EXT_AI_CACHE_DIR") .map_err(AnyError::from) .map(PathBuf::from) @@ -63,10 +64,10 @@ pub async fn fetch_and_cache_from_url( let filepath = filepath.clone(); let checksum = tokio::task::spawn_blocking(move || { let mut file = std::fs::File::open(filepath).ok()?; - let mut hasher = Xxh3::new(); + let mut hasher = Sha256::new(); let _ = std::io::copy(&mut file, &mut hasher).ok()?; - Some(hasher.finish().to_be_bytes()) + Some(hasher.finalize()) }) .await; @@ -106,7 +107,7 @@ pub async fn fetch_and_cache_from_url( let mut stream = resp.bytes_stream(); let mut writer = tokio::io::BufWriter::new(file); - let mut hasher = AllowStdIo::new(Xxh3::new()).compat_write(); + let mut hasher = AllowStdIo::new(Sha256::new()).compat_write(); let mut written = 0; while let Some(chunk) = stream.next().await { @@ -125,7 +126,7 @@ pub async fn fetch_and_cache_from_url( let checksum_str = { let hasher = hasher.into_inner().into_inner(); - faster_hex::hex_string(&hasher.finish().to_be_bytes()) + faster_hex::hex_string(&hasher.finalize()) }; info!({ bytes_written = written, checksum = &checksum_str }, "done");