From 4cadc8113b2d93be6b59333e5ba49e0e6f4d906a Mon Sep 17 00:00:00 2001 From: Mubelotix Date: Fri, 20 Jun 2025 12:42:22 +0200 Subject: [PATCH 01/18] Add embedder stats in batches --- crates/benchmarks/benches/indexing.rs | 2 +- crates/benchmarks/benches/utils.rs | 2 +- crates/dump/src/lib.rs | 1 + .../src/scheduler/process_batch.rs | 1 + .../src/scheduler/process_index_operation.rs | 3 ++ crates/meilisearch-types/src/batches.rs | 10 ++++ crates/meilisearch/src/lib.rs | 3 +- crates/milli/src/progress.rs | 34 +++++++++++- .../milli/src/search/new/tests/integration.rs | 2 +- crates/milli/src/test_index.rs | 2 +- .../extract/extract_vector_points.rs | 8 ++- .../src/update/index_documents/extract/mod.rs | 5 ++ .../milli/src/update/index_documents/mod.rs | 9 +++- .../src/update/new/extract/vectors/mod.rs | 2 +- crates/milli/src/update/settings.rs | 9 ++-- crates/milli/src/vector/composite.rs | 32 ++++++----- crates/milli/src/vector/mod.rs | 28 +++++----- crates/milli/src/vector/ollama.rs | 15 ++++-- crates/milli/src/vector/openai.rs | 17 +++--- crates/milli/src/vector/rest.rs | 54 ++++++++++++++----- crates/milli/tests/search/distinct.rs | 2 +- .../milli/tests/search/facet_distribution.rs | 2 +- crates/milli/tests/search/mod.rs | 2 +- crates/milli/tests/search/phrase_search.rs | 2 +- crates/milli/tests/search/query_criteria.rs | 6 +-- crates/milli/tests/search/typo_tolerance.rs | 8 +-- 26 files changed, 188 insertions(+), 73 deletions(-) diff --git a/crates/benchmarks/benches/indexing.rs b/crates/benchmarks/benches/indexing.rs index 9199c3877..b882b598d 100644 --- a/crates/benchmarks/benches/indexing.rs +++ b/crates/benchmarks/benches/indexing.rs @@ -65,7 +65,7 @@ fn setup_settings<'t>( let sortable_fields = sortable_fields.iter().map(|s| s.to_string()).collect(); builder.set_sortable_fields(sortable_fields); - builder.execute(|_| (), || false).unwrap(); + builder.execute(|_| (), || false, None).unwrap(); } fn setup_index_with_settings( diff --git a/crates/benchmarks/benches/utils.rs b/crates/benchmarks/benches/utils.rs index aaa2d50a0..913807b45 100644 --- a/crates/benchmarks/benches/utils.rs +++ b/crates/benchmarks/benches/utils.rs @@ -90,7 +90,7 @@ pub fn base_setup(conf: &Conf) -> Index { (conf.configure)(&mut builder); - builder.execute(|_| (), || false).unwrap(); + builder.execute(|_| (), || false, None).unwrap(); wtxn.commit().unwrap(); let config = IndexerConfig::default(); diff --git a/crates/dump/src/lib.rs b/crates/dump/src/lib.rs index 285818a87..c48c68f62 100644 --- a/crates/dump/src/lib.rs +++ b/crates/dump/src/lib.rs @@ -328,6 +328,7 @@ pub(crate) mod test { progress_trace: Default::default(), write_channel_congestion: None, internal_database_sizes: Default::default(), + embeddings: Default::default(), }, enqueued_at: Some(BatchEnqueuedAt { earliest: datetime!(2022-11-11 0:00 UTC), diff --git a/crates/index-scheduler/src/scheduler/process_batch.rs b/crates/index-scheduler/src/scheduler/process_batch.rs index c349f90ad..71e423a58 100644 --- a/crates/index-scheduler/src/scheduler/process_batch.rs +++ b/crates/index-scheduler/src/scheduler/process_batch.rs @@ -242,6 +242,7 @@ impl IndexScheduler { .execute( |indexing_step| tracing::debug!(update = ?indexing_step), || must_stop_processing.get(), + Some(progress.embedder_stats), ) .map_err(|e| Error::from_milli(e, Some(index_uid.to_string())))?; index_wtxn.commit()?; diff --git a/crates/index-scheduler/src/scheduler/process_index_operation.rs b/crates/index-scheduler/src/scheduler/process_index_operation.rs index 093c6209d..92d13e7e7 100644 --- a/crates/index-scheduler/src/scheduler/process_index_operation.rs +++ b/crates/index-scheduler/src/scheduler/process_index_operation.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use bumpalo::collections::CollectIn; use bumpalo::Bump; use meilisearch_types::heed::RwTxn; @@ -472,6 +474,7 @@ impl IndexScheduler { .execute( |indexing_step| tracing::debug!(update = ?indexing_step), || must_stop_processing.get(), + Some(Arc::clone(&progress.embedder_stats)) ) .map_err(|err| Error::from_milli(err, Some(index_uid.clone())))?; diff --git a/crates/meilisearch-types/src/batches.rs b/crates/meilisearch-types/src/batches.rs index 4d40189db..2ef373eac 100644 --- a/crates/meilisearch-types/src/batches.rs +++ b/crates/meilisearch-types/src/batches.rs @@ -82,4 +82,14 @@ pub struct BatchStats { pub write_channel_congestion: Option>, #[serde(default, skip_serializing_if = "serde_json::Map::is_empty")] pub internal_database_sizes: serde_json::Map, + pub embeddings: BatchEmbeddingStats +} + +#[derive(Default, Debug, Clone, PartialEq, Eq, Serialize, Deserialize, ToSchema)] +#[serde(rename_all = "camelCase")] +#[schema(rename_all = "camelCase")] +pub struct BatchEmbeddingStats { + pub total_count: usize, + pub error_count: usize, + pub last_error: Option, } diff --git a/crates/meilisearch/src/lib.rs b/crates/meilisearch/src/lib.rs index 1e0c205d0..782d6172f 100644 --- a/crates/meilisearch/src/lib.rs +++ b/crates/meilisearch/src/lib.rs @@ -543,7 +543,7 @@ fn import_dump( let settings = index_reader.settings()?; apply_settings_to_builder(&settings, &mut builder); builder - .execute(|indexing_step| tracing::debug!("update: {:?}", indexing_step), || false)?; + .execute(|indexing_step| tracing::debug!("update: {:?}", indexing_step), || false, None)?; // 4.3 Import the documents. // 4.3.1 We need to recreate the grenad+obkv format accepted by the index. @@ -574,6 +574,7 @@ fn import_dump( }, |indexing_step| tracing::trace!("update: {:?}", indexing_step), || false, + None, )?; let builder = builder.with_embedders(embedders); diff --git a/crates/milli/src/progress.rs b/crates/milli/src/progress.rs index fa651e17f..ff795b220 100644 --- a/crates/milli/src/progress.rs +++ b/crates/milli/src/progress.rs @@ -1,7 +1,7 @@ use std::any::TypeId; use std::borrow::Cow; use std::marker::PhantomData; -use std::sync::atomic::{AtomicU32, Ordering}; +use std::sync::atomic::{AtomicU32, AtomicUsize, Ordering}; use std::sync::{Arc, RwLock}; use std::time::{Duration, Instant}; @@ -20,6 +20,13 @@ pub trait Step: 'static + Send + Sync { #[derive(Clone, Default)] pub struct Progress { steps: Arc>, + pub embedder_stats: Arc, +} + +#[derive(Default)] +pub struct EmbedderStats { + pub errors: Arc, u32)>>, + pub total_requests: AtomicUsize } #[derive(Default)] @@ -65,7 +72,19 @@ impl Progress { }); } - ProgressView { steps: step_view, percentage: percentage * 100.0 } + let embedder_view = { + let (last_error, error_count) = match self.embedder_stats.errors.read() { + Ok(guard) => (guard.0.clone(), guard.1), + Err(_) => (None, 0), + }; + EmbedderStatsView { + last_error, + request_count: self.embedder_stats.total_requests.load(Ordering::Relaxed) as u32, + error_count, + } + }; + + ProgressView { steps: step_view, percentage: percentage * 100.0, embedder: embedder_view } } pub fn accumulated_durations(&self) -> IndexMap { @@ -209,6 +228,7 @@ make_enum_progress! { pub struct ProgressView { pub steps: Vec, pub percentage: f32, + pub embedder: EmbedderStatsView, } #[derive(Debug, Serialize, Clone, ToSchema)] @@ -220,6 +240,16 @@ pub struct ProgressStepView { pub total: u32, } +#[derive(Debug, Serialize, Clone, ToSchema)] +#[serde(rename_all = "camelCase")] +#[schema(rename_all = "camelCase")] +pub struct EmbedderStatsView { + #[serde(skip_serializing_if = "Option::is_none")] + pub last_error: Option, + pub request_count: u32, + pub error_count: u32, +} + /// Used when the name can change but it's still the same step. /// To avoid conflicts on the `TypeId`, create a unique type every time you use this step: /// ```text diff --git a/crates/milli/src/search/new/tests/integration.rs b/crates/milli/src/search/new/tests/integration.rs index 4a6cc9b90..e7634a4eb 100644 --- a/crates/milli/src/search/new/tests/integration.rs +++ b/crates/milli/src/search/new/tests/integration.rs @@ -44,7 +44,7 @@ pub fn setup_search_index_with_criteria(criteria: &[Criterion]) -> Index { S("america") => vec![S("the united states")], }); builder.set_searchable_fields(vec![S("title"), S("description")]); - builder.execute(|_| (), || false).unwrap(); + builder.execute(|_| (), || false, None).unwrap(); wtxn.commit().unwrap(); // index documents diff --git a/crates/milli/src/test_index.rs b/crates/milli/src/test_index.rs index dfd570b96..634d45195 100644 --- a/crates/milli/src/test_index.rs +++ b/crates/milli/src/test_index.rs @@ -134,7 +134,7 @@ impl TempIndex { ) -> Result<(), crate::error::Error> { let mut builder = update::Settings::new(wtxn, &self.inner, &self.indexer_config); update(&mut builder); - builder.execute(drop, || false)?; + builder.execute(drop, || false, None)?; Ok(()) } diff --git a/crates/milli/src/update/index_documents/extract/extract_vector_points.rs b/crates/milli/src/update/index_documents/extract/extract_vector_points.rs index cb8c121ce..5e6bde53d 100644 --- a/crates/milli/src/update/index_documents/extract/extract_vector_points.rs +++ b/crates/milli/src/update/index_documents/extract/extract_vector_points.rs @@ -17,6 +17,7 @@ use crate::constants::RESERVED_VECTORS_FIELD_NAME; use crate::error::FaultSource; use crate::fields_ids_map::metadata::FieldIdMapWithMetadata; use crate::index::IndexEmbeddingConfig; +use crate::progress::EmbedderStats; use crate::prompt::Prompt; use crate::update::del_add::{DelAdd, KvReaderDelAdd, KvWriterDelAdd}; use crate::update::settings::InnerIndexSettingsDiff; @@ -682,6 +683,7 @@ pub fn extract_embeddings( embedder: Arc, embedder_name: &str, possible_embedding_mistakes: &PossibleEmbeddingMistakes, + embedder_stats: Option>, unused_vectors_distribution: &UnusedVectorsDistribution, request_threads: &ThreadPoolNoAbort, ) -> Result>> { @@ -724,6 +726,7 @@ pub fn extract_embeddings( std::mem::replace(&mut chunks, Vec::with_capacity(n_chunks)), embedder_name, possible_embedding_mistakes, + embedder_stats.clone(), unused_vectors_distribution, request_threads, )?; @@ -746,6 +749,7 @@ pub fn extract_embeddings( std::mem::take(&mut chunks), embedder_name, possible_embedding_mistakes, + embedder_stats.clone(), unused_vectors_distribution, request_threads, )?; @@ -764,6 +768,7 @@ pub fn extract_embeddings( vec![std::mem::take(&mut current_chunk)], embedder_name, possible_embedding_mistakes, + embedder_stats, unused_vectors_distribution, request_threads, )?; @@ -783,10 +788,11 @@ fn embed_chunks( text_chunks: Vec>, embedder_name: &str, possible_embedding_mistakes: &PossibleEmbeddingMistakes, + embedder_stats: Option>, unused_vectors_distribution: &UnusedVectorsDistribution, request_threads: &ThreadPoolNoAbort, ) -> Result>> { - match embedder.embed_index(text_chunks, request_threads) { + match embedder.embed_index(text_chunks, request_threads, embedder_stats) { Ok(chunks) => Ok(chunks), Err(error) => { if let FaultSource::Bug = error.fault { diff --git a/crates/milli/src/update/index_documents/extract/mod.rs b/crates/milli/src/update/index_documents/extract/mod.rs index 8cd664a2f..020b48f2c 100644 --- a/crates/milli/src/update/index_documents/extract/mod.rs +++ b/crates/milli/src/update/index_documents/extract/mod.rs @@ -31,6 +31,7 @@ use self::extract_word_position_docids::extract_word_position_docids; use super::helpers::{as_cloneable_grenad, CursorClonableMmap, GrenadParameters}; use super::{helpers, TypedChunk}; use crate::index::IndexEmbeddingConfig; +use crate::progress::EmbedderStats; use crate::update::settings::InnerIndexSettingsDiff; use crate::vector::error::PossibleEmbeddingMistakes; use crate::{FieldId, Result, ThreadPoolNoAbort, ThreadPoolNoAbortBuilder}; @@ -49,6 +50,7 @@ pub(crate) fn data_from_obkv_documents( settings_diff: Arc, max_positions_per_attributes: Option, possible_embedding_mistakes: Arc, + embedder_stats: Option>, ) -> Result<()> { let (original_pipeline_result, flattened_pipeline_result): (Result<_>, Result<_>) = rayon::join( || { @@ -62,6 +64,7 @@ pub(crate) fn data_from_obkv_documents( embedders_configs.clone(), settings_diff.clone(), possible_embedding_mistakes.clone(), + embedder_stats.clone(), ) }) .collect::>() @@ -231,6 +234,7 @@ fn send_original_documents_data( embedders_configs: Arc>, settings_diff: Arc, possible_embedding_mistakes: Arc, + embedder_stats: Option>, ) -> Result<()> { let original_documents_chunk = original_documents_chunk.and_then(|c| unsafe { as_cloneable_grenad(&c) })?; @@ -270,6 +274,7 @@ fn send_original_documents_data( embedder.clone(), &embedder_name, &possible_embedding_mistakes, + embedder_stats.clone(), &unused_vectors_distribution, request_threads(), ) { diff --git a/crates/milli/src/update/index_documents/mod.rs b/crates/milli/src/update/index_documents/mod.rs index f547c68d4..fad43bd30 100644 --- a/crates/milli/src/update/index_documents/mod.rs +++ b/crates/milli/src/update/index_documents/mod.rs @@ -32,7 +32,7 @@ use crate::database_stats::DatabaseStats; use crate::documents::{obkv_to_object, DocumentsBatchReader}; use crate::error::{Error, InternalError}; use crate::index::{PrefixSearch, PrefixSettings}; -use crate::progress::Progress; +use crate::progress::{EmbedderStats, Progress}; pub use crate::update::index_documents::helpers::CursorClonableMmap; use crate::update::{ IndexerConfig, UpdateIndexingStep, WordPrefixDocids, WordPrefixIntegerDocids, WordsPrefixesFst, @@ -81,6 +81,7 @@ pub struct IndexDocuments<'t, 'i, 'a, FP, FA> { added_documents: u64, deleted_documents: u64, embedders: EmbeddingConfigs, + embedder_stats: Option>, } #[derive(Default, Debug, Clone)] @@ -103,6 +104,7 @@ where config: IndexDocumentsConfig, progress: FP, should_abort: FA, + embedder_stats: Option>, ) -> Result> { let transform = Some(Transform::new( wtxn, @@ -123,6 +125,7 @@ where added_documents: 0, deleted_documents: 0, embedders: Default::default(), + embedder_stats, }) } @@ -292,6 +295,7 @@ where // Run extraction pipeline in parallel. let mut modified_docids = RoaringBitmap::new(); + let embedder_stats = self.embedder_stats.clone(); pool.install(|| { let settings_diff_cloned = settings_diff.clone(); rayon::spawn(move || { @@ -326,7 +330,8 @@ where embedders_configs.clone(), settings_diff_cloned, max_positions_per_attributes, - Arc::new(possible_embedding_mistakes) + Arc::new(possible_embedding_mistakes), + embedder_stats.clone() ) }); diff --git a/crates/milli/src/update/new/extract/vectors/mod.rs b/crates/milli/src/update/new/extract/vectors/mod.rs index 43647e786..5b6559d74 100644 --- a/crates/milli/src/update/new/extract/vectors/mod.rs +++ b/crates/milli/src/update/new/extract/vectors/mod.rs @@ -450,7 +450,7 @@ impl<'a, 'b, 'extractor> Chunks<'a, 'b, 'extractor> { return Err(crate::Error::UserError(crate::UserError::DocumentEmbeddingError(msg))); } - let res = match embedder.embed_index_ref(texts.as_slice(), threads) { + let res = match embedder.embed_index_ref(texts.as_slice(), threads, None) { Ok(embeddings) => { for (docid, embedding) in ids.into_iter().zip(embeddings) { sender.set_vector(*docid, embedder_id, embedding).unwrap(); diff --git a/crates/milli/src/update/settings.rs b/crates/milli/src/update/settings.rs index f396cd079..7c5a70aa3 100644 --- a/crates/milli/src/update/settings.rs +++ b/crates/milli/src/update/settings.rs @@ -27,6 +27,7 @@ use crate::index::{ DEFAULT_MIN_WORD_LEN_ONE_TYPO, DEFAULT_MIN_WORD_LEN_TWO_TYPOS, }; use crate::order_by_map::OrderByMap; +use crate::progress::EmbedderStats; use crate::prompt::{default_max_bytes, default_template_text, PromptData}; use crate::proximity::ProximityPrecision; use crate::update::index_documents::IndexDocumentsMethod; @@ -466,7 +467,7 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> { #[tracing::instrument( level = "trace" - skip(self, progress_callback, should_abort, settings_diff), + skip(self, progress_callback, should_abort, settings_diff, embedder_stats), target = "indexing::documents" )] fn reindex( @@ -474,6 +475,7 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> { progress_callback: &FP, should_abort: &FA, settings_diff: InnerIndexSettingsDiff, + embedder_stats: Option>, ) -> Result<()> where FP: Fn(UpdateIndexingStep) + Sync, @@ -505,6 +507,7 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> { IndexDocumentsConfig::default(), &progress_callback, &should_abort, + embedder_stats, )?; indexing_builder.execute_raw(output)?; @@ -1355,7 +1358,7 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> { } } - pub fn execute(mut self, progress_callback: FP, should_abort: FA) -> Result<()> + pub fn execute(mut self, progress_callback: FP, should_abort: FA, embedder_stats: Option>) -> Result<()> where FP: Fn(UpdateIndexingStep) + Sync, FA: Fn() -> bool + Sync, @@ -1413,7 +1416,7 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> { ); if inner_settings_diff.any_reindexing_needed() { - self.reindex(&progress_callback, &should_abort, inner_settings_diff)?; + self.reindex(&progress_callback, &should_abort, inner_settings_diff, embedder_stats)?; } Ok(()) diff --git a/crates/milli/src/vector/composite.rs b/crates/milli/src/vector/composite.rs index 9c5992bd3..daec50e4b 100644 --- a/crates/milli/src/vector/composite.rs +++ b/crates/milli/src/vector/composite.rs @@ -1,3 +1,4 @@ +use std::sync::Arc; use std::time::Instant; use arroy::Distance; @@ -7,6 +8,7 @@ use super::{ hf, manual, ollama, openai, rest, DistributionShift, EmbedError, Embedding, EmbeddingCache, NewEmbedderError, }; +use crate::progress::EmbedderStats; use crate::ThreadPoolNoAbort; #[derive(Debug)] @@ -81,6 +83,7 @@ impl Embedder { "This is a sample text. It is meant to compare similarity.".into(), ], None, + None, ) .map_err(|error| NewEmbedderError::composite_test_embedding_failed(error, "search"))?; @@ -92,6 +95,7 @@ impl Embedder { "This is a sample text. It is meant to compare similarity.".into(), ], None, + None, ) .map_err(|error| { NewEmbedderError::composite_test_embedding_failed(error, "indexing") @@ -150,13 +154,14 @@ impl SubEmbedder { &self, texts: Vec, deadline: Option, + embedder_stats: Option>, ) -> std::result::Result, EmbedError> { match self { SubEmbedder::HuggingFace(embedder) => embedder.embed(texts), - SubEmbedder::OpenAi(embedder) => embedder.embed(&texts, deadline), - SubEmbedder::Ollama(embedder) => embedder.embed(&texts, deadline), + SubEmbedder::OpenAi(embedder) => embedder.embed(&texts, deadline, embedder_stats), + SubEmbedder::Ollama(embedder) => embedder.embed(&texts, deadline, embedder_stats), SubEmbedder::UserProvided(embedder) => embedder.embed(&texts), - SubEmbedder::Rest(embedder) => embedder.embed(texts, deadline), + SubEmbedder::Rest(embedder) => embedder.embed(texts, deadline, embedder_stats), } } @@ -164,18 +169,19 @@ impl SubEmbedder { &self, text: &str, deadline: Option, + embedder_stats: Option>, ) -> std::result::Result { match self { SubEmbedder::HuggingFace(embedder) => embedder.embed_one(text), SubEmbedder::OpenAi(embedder) => { - embedder.embed(&[text], deadline)?.pop().ok_or_else(EmbedError::missing_embedding) + embedder.embed(&[text], deadline, embedder_stats)?.pop().ok_or_else(EmbedError::missing_embedding) } SubEmbedder::Ollama(embedder) => { - embedder.embed(&[text], deadline)?.pop().ok_or_else(EmbedError::missing_embedding) + embedder.embed(&[text], deadline, embedder_stats)?.pop().ok_or_else(EmbedError::missing_embedding) } SubEmbedder::UserProvided(embedder) => embedder.embed_one(text), SubEmbedder::Rest(embedder) => embedder - .embed_ref(&[text], deadline)? + .embed_ref(&[text], deadline, embedder_stats)? .pop() .ok_or_else(EmbedError::missing_embedding), } @@ -188,13 +194,14 @@ impl SubEmbedder { &self, text_chunks: Vec>, threads: &ThreadPoolNoAbort, + embedder_stats: Option>, ) -> std::result::Result>, EmbedError> { match self { SubEmbedder::HuggingFace(embedder) => embedder.embed_index(text_chunks), - SubEmbedder::OpenAi(embedder) => embedder.embed_index(text_chunks, threads), - SubEmbedder::Ollama(embedder) => embedder.embed_index(text_chunks, threads), + SubEmbedder::OpenAi(embedder) => embedder.embed_index(text_chunks, threads, embedder_stats), + SubEmbedder::Ollama(embedder) => embedder.embed_index(text_chunks, threads, embedder_stats), SubEmbedder::UserProvided(embedder) => embedder.embed_index(text_chunks), - SubEmbedder::Rest(embedder) => embedder.embed_index(text_chunks, threads), + SubEmbedder::Rest(embedder) => embedder.embed_index(text_chunks, threads, embedder_stats), } } @@ -203,13 +210,14 @@ impl SubEmbedder { &self, texts: &[&str], threads: &ThreadPoolNoAbort, + embedder_stats: Option>, ) -> std::result::Result, EmbedError> { match self { SubEmbedder::HuggingFace(embedder) => embedder.embed_index_ref(texts), - SubEmbedder::OpenAi(embedder) => embedder.embed_index_ref(texts, threads), - SubEmbedder::Ollama(embedder) => embedder.embed_index_ref(texts, threads), + SubEmbedder::OpenAi(embedder) => embedder.embed_index_ref(texts, threads, embedder_stats), + SubEmbedder::Ollama(embedder) => embedder.embed_index_ref(texts, threads, embedder_stats), SubEmbedder::UserProvided(embedder) => embedder.embed_index_ref(texts), - SubEmbedder::Rest(embedder) => embedder.embed_index_ref(texts, threads), + SubEmbedder::Rest(embedder) => embedder.embed_index_ref(texts, threads, embedder_stats), } } diff --git a/crates/milli/src/vector/mod.rs b/crates/milli/src/vector/mod.rs index c2978f5db..124e17cff 100644 --- a/crates/milli/src/vector/mod.rs +++ b/crates/milli/src/vector/mod.rs @@ -13,7 +13,7 @@ use serde::{Deserialize, Serialize}; use utoipa::ToSchema; use self::error::{EmbedError, NewEmbedderError}; -use crate::progress::Progress; +use crate::progress::{EmbedderStats, Progress}; use crate::prompt::{Prompt, PromptData}; use crate::ThreadPoolNoAbort; @@ -720,17 +720,17 @@ impl Embedder { let embedding = match self { Embedder::HuggingFace(embedder) => embedder.embed_one(text), Embedder::OpenAi(embedder) => { - embedder.embed(&[text], deadline)?.pop().ok_or_else(EmbedError::missing_embedding) + embedder.embed(&[text], deadline, None)?.pop().ok_or_else(EmbedError::missing_embedding) } Embedder::Ollama(embedder) => { - embedder.embed(&[text], deadline)?.pop().ok_or_else(EmbedError::missing_embedding) + embedder.embed(&[text], deadline, None)?.pop().ok_or_else(EmbedError::missing_embedding) } Embedder::UserProvided(embedder) => embedder.embed_one(text), Embedder::Rest(embedder) => embedder - .embed_ref(&[text], deadline)? + .embed_ref(&[text], deadline, None)? .pop() .ok_or_else(EmbedError::missing_embedding), - Embedder::Composite(embedder) => embedder.search.embed_one(text, deadline), + Embedder::Composite(embedder) => embedder.search.embed_one(text, deadline, None), }?; if let Some(cache) = self.cache() { @@ -747,14 +747,15 @@ impl Embedder { &self, text_chunks: Vec>, threads: &ThreadPoolNoAbort, + embedder_stats: Option>, ) -> std::result::Result>, EmbedError> { match self { Embedder::HuggingFace(embedder) => embedder.embed_index(text_chunks), - Embedder::OpenAi(embedder) => embedder.embed_index(text_chunks, threads), - Embedder::Ollama(embedder) => embedder.embed_index(text_chunks, threads), + Embedder::OpenAi(embedder) => embedder.embed_index(text_chunks, threads, embedder_stats), + Embedder::Ollama(embedder) => embedder.embed_index(text_chunks, threads, embedder_stats), Embedder::UserProvided(embedder) => embedder.embed_index(text_chunks), - Embedder::Rest(embedder) => embedder.embed_index(text_chunks, threads), - Embedder::Composite(embedder) => embedder.index.embed_index(text_chunks, threads), + Embedder::Rest(embedder) => embedder.embed_index(text_chunks, threads, embedder_stats), + Embedder::Composite(embedder) => embedder.index.embed_index(text_chunks, threads, embedder_stats), } } @@ -763,14 +764,15 @@ impl Embedder { &self, texts: &[&str], threads: &ThreadPoolNoAbort, + embedder_stats: Option>, ) -> std::result::Result, EmbedError> { match self { Embedder::HuggingFace(embedder) => embedder.embed_index_ref(texts), - Embedder::OpenAi(embedder) => embedder.embed_index_ref(texts, threads), - Embedder::Ollama(embedder) => embedder.embed_index_ref(texts, threads), + Embedder::OpenAi(embedder) => embedder.embed_index_ref(texts, threads, embedder_stats), + Embedder::Ollama(embedder) => embedder.embed_index_ref(texts, threads, embedder_stats), Embedder::UserProvided(embedder) => embedder.embed_index_ref(texts), - Embedder::Rest(embedder) => embedder.embed_index_ref(texts, threads), - Embedder::Composite(embedder) => embedder.index.embed_index_ref(texts, threads), + Embedder::Rest(embedder) => embedder.embed_index_ref(texts, threads, embedder_stats), + Embedder::Composite(embedder) => embedder.index.embed_index_ref(texts, threads, embedder_stats), } } diff --git a/crates/milli/src/vector/ollama.rs b/crates/milli/src/vector/ollama.rs index 8beae6205..b3ee925e6 100644 --- a/crates/milli/src/vector/ollama.rs +++ b/crates/milli/src/vector/ollama.rs @@ -1,3 +1,4 @@ +use std::sync::Arc; use std::time::Instant; use rayon::iter::{IntoParallelIterator as _, ParallelIterator as _}; @@ -7,6 +8,7 @@ use super::error::{EmbedError, EmbedErrorKind, NewEmbedderError, NewEmbedderErro use super::rest::{Embedder as RestEmbedder, EmbedderOptions as RestEmbedderOptions}; use super::{DistributionShift, EmbeddingCache, REQUEST_PARALLELISM}; use crate::error::FaultSource; +use crate::progress::EmbedderStats; use crate::vector::Embedding; use crate::ThreadPoolNoAbort; @@ -104,8 +106,9 @@ impl Embedder { &self, texts: &[S], deadline: Option, + embedder_stats: Option> ) -> Result, EmbedError> { - match self.rest_embedder.embed_ref(texts, deadline) { + match self.rest_embedder.embed_ref(texts, deadline, embedder_stats) { Ok(embeddings) => Ok(embeddings), Err(EmbedError { kind: EmbedErrorKind::RestOtherStatusCode(404, error), fault: _ }) => { Err(EmbedError::ollama_model_not_found(error)) @@ -118,15 +121,16 @@ impl Embedder { &self, text_chunks: Vec>, threads: &ThreadPoolNoAbort, + embedder_stats: Option>, ) -> Result>, EmbedError> { // This condition helps reduce the number of active rayon jobs // so that we avoid consuming all the LMDB rtxns and avoid stack overflows. if threads.active_operations() >= REQUEST_PARALLELISM { - text_chunks.into_iter().map(move |chunk| self.embed(&chunk, None)).collect() + text_chunks.into_iter().map(move |chunk| self.embed(&chunk, None, embedder_stats.clone())).collect() } else { threads .install(move || { - text_chunks.into_par_iter().map(move |chunk| self.embed(&chunk, None)).collect() + text_chunks.into_par_iter().map(move |chunk| self.embed(&chunk, None, embedder_stats.clone())).collect() }) .map_err(|error| EmbedError { kind: EmbedErrorKind::PanicInThreadPool(error), @@ -139,13 +143,14 @@ impl Embedder { &self, texts: &[&str], threads: &ThreadPoolNoAbort, + embedder_stats: Option> ) -> Result>, EmbedError> { // This condition helps reduce the number of active rayon jobs // so that we avoid consuming all the LMDB rtxns and avoid stack overflows. if threads.active_operations() >= REQUEST_PARALLELISM { let embeddings: Result>, _> = texts .chunks(self.prompt_count_in_chunk_hint()) - .map(move |chunk| self.embed(chunk, None)) + .map(move |chunk| self.embed(chunk, None, embedder_stats.clone())) .collect(); let embeddings = embeddings?; @@ -155,7 +160,7 @@ impl Embedder { .install(move || { let embeddings: Result>, _> = texts .par_chunks(self.prompt_count_in_chunk_hint()) - .map(move |chunk| self.embed(chunk, None)) + .map(move |chunk| self.embed(chunk, None, embedder_stats.clone())) .collect(); let embeddings = embeddings?; diff --git a/crates/milli/src/vector/openai.rs b/crates/milli/src/vector/openai.rs index df29f6916..384abe880 100644 --- a/crates/milli/src/vector/openai.rs +++ b/crates/milli/src/vector/openai.rs @@ -1,4 +1,5 @@ use std::fmt; +use std::sync::Arc; use std::time::Instant; use ordered_float::OrderedFloat; @@ -9,6 +10,7 @@ use super::error::{EmbedError, NewEmbedderError}; use super::rest::{Embedder as RestEmbedder, EmbedderOptions as RestEmbedderOptions}; use super::{DistributionShift, EmbeddingCache, REQUEST_PARALLELISM}; use crate::error::FaultSource; +use crate::progress::EmbedderStats; use crate::vector::error::EmbedErrorKind; use crate::vector::Embedding; use crate::ThreadPoolNoAbort; @@ -215,8 +217,9 @@ impl Embedder { &self, texts: &[S], deadline: Option, + embedder_stats: Option>, ) -> Result, EmbedError> { - match self.rest_embedder.embed_ref(texts, deadline) { + match self.rest_embedder.embed_ref(texts, deadline, embedder_stats) { Ok(embeddings) => Ok(embeddings), Err(EmbedError { kind: EmbedErrorKind::RestBadRequest(error, _), fault: _ }) => { tracing::warn!(error=?error, "OpenAI: received `BAD_REQUEST`. Input was maybe too long, retrying on tokenized version. For best performance, limit the size of your document template."); @@ -238,7 +241,7 @@ impl Embedder { let encoded = self.tokenizer.encode_ordinary(text); let len = encoded.len(); if len < max_token_count { - all_embeddings.append(&mut self.rest_embedder.embed_ref(&[text], deadline)?); + all_embeddings.append(&mut self.rest_embedder.embed_ref(&[text], deadline, None)?); continue; } @@ -255,15 +258,16 @@ impl Embedder { &self, text_chunks: Vec>, threads: &ThreadPoolNoAbort, + embedder_stats: Option>, ) -> Result>, EmbedError> { // This condition helps reduce the number of active rayon jobs // so that we avoid consuming all the LMDB rtxns and avoid stack overflows. if threads.active_operations() >= REQUEST_PARALLELISM { - text_chunks.into_iter().map(move |chunk| self.embed(&chunk, None)).collect() + text_chunks.into_iter().map(move |chunk| self.embed(&chunk, None, embedder_stats.clone())).collect() } else { threads .install(move || { - text_chunks.into_par_iter().map(move |chunk| self.embed(&chunk, None)).collect() + text_chunks.into_par_iter().map(move |chunk| self.embed(&chunk, None, embedder_stats.clone())).collect() }) .map_err(|error| EmbedError { kind: EmbedErrorKind::PanicInThreadPool(error), @@ -276,13 +280,14 @@ impl Embedder { &self, texts: &[&str], threads: &ThreadPoolNoAbort, + embedder_stats: Option>, ) -> Result>, EmbedError> { // This condition helps reduce the number of active rayon jobs // so that we avoid consuming all the LMDB rtxns and avoid stack overflows. if threads.active_operations() >= REQUEST_PARALLELISM { let embeddings: Result>, _> = texts .chunks(self.prompt_count_in_chunk_hint()) - .map(move |chunk| self.embed(chunk, None)) + .map(move |chunk| self.embed(chunk, None, embedder_stats.clone())) .collect(); let embeddings = embeddings?; Ok(embeddings.into_iter().flatten().collect()) @@ -291,7 +296,7 @@ impl Embedder { .install(move || { let embeddings: Result>, _> = texts .par_chunks(self.prompt_count_in_chunk_hint()) - .map(move |chunk| self.embed(chunk, None)) + .map(move |chunk| self.embed(chunk, None, embedder_stats.clone())) .collect(); let embeddings = embeddings?; diff --git a/crates/milli/src/vector/rest.rs b/crates/milli/src/vector/rest.rs index b87ac9f77..fc0ff308b 100644 --- a/crates/milli/src/vector/rest.rs +++ b/crates/milli/src/vector/rest.rs @@ -1,4 +1,5 @@ use std::collections::BTreeMap; +use std::sync::Arc; use std::time::Instant; use deserr::Deserr; @@ -14,6 +15,7 @@ use super::{ }; use crate::error::FaultSource; use crate::ThreadPoolNoAbort; +use crate::progress::EmbedderStats; // retrying in case of failure pub struct Retry { @@ -168,19 +170,21 @@ impl Embedder { &self, texts: Vec, deadline: Option, + embedder_stats: Option>, ) -> Result, EmbedError> { - embed(&self.data, texts.as_slice(), texts.len(), Some(self.dimensions), deadline) + embed(&self.data, texts.as_slice(), texts.len(), Some(self.dimensions), deadline, embedder_stats) } pub fn embed_ref( &self, texts: &[S], deadline: Option, + embedder_stats: Option>, ) -> Result, EmbedError> where S: AsRef + Serialize, { - embed(&self.data, texts, texts.len(), Some(self.dimensions), deadline) + embed(&self.data, texts, texts.len(), Some(self.dimensions), deadline, embedder_stats) } pub fn embed_tokens( @@ -188,7 +192,7 @@ impl Embedder { tokens: &[u32], deadline: Option, ) -> Result { - let mut embeddings = embed(&self.data, tokens, 1, Some(self.dimensions), deadline)?; + let mut embeddings = embed(&self.data, tokens, 1, Some(self.dimensions), deadline, None)?; // unwrap: guaranteed that embeddings.len() == 1, otherwise the previous line terminated in error Ok(embeddings.pop().unwrap()) } @@ -197,15 +201,16 @@ impl Embedder { &self, text_chunks: Vec>, threads: &ThreadPoolNoAbort, + embedder_stats: Option>, ) -> Result>, EmbedError> { // This condition helps reduce the number of active rayon jobs // so that we avoid consuming all the LMDB rtxns and avoid stack overflows. if threads.active_operations() >= REQUEST_PARALLELISM { - text_chunks.into_iter().map(move |chunk| self.embed(chunk, None)).collect() + text_chunks.into_iter().map(move |chunk| self.embed(chunk, None, embedder_stats.clone())).collect() } else { threads .install(move || { - text_chunks.into_par_iter().map(move |chunk| self.embed(chunk, None)).collect() + text_chunks.into_par_iter().map(move |chunk| self.embed(chunk, None, embedder_stats.clone())).collect() }) .map_err(|error| EmbedError { kind: EmbedErrorKind::PanicInThreadPool(error), @@ -218,13 +223,14 @@ impl Embedder { &self, texts: &[&str], threads: &ThreadPoolNoAbort, + embedder_stats: Option> ) -> Result, EmbedError> { // This condition helps reduce the number of active rayon jobs // so that we avoid consuming all the LMDB rtxns and avoid stack overflows. if threads.active_operations() >= REQUEST_PARALLELISM { let embeddings: Result>, _> = texts .chunks(self.prompt_count_in_chunk_hint()) - .map(move |chunk| self.embed_ref(chunk, None)) + .map(move |chunk| self.embed_ref(chunk, None, embedder_stats.clone())) .collect(); let embeddings = embeddings?; @@ -234,7 +240,7 @@ impl Embedder { .install(move || { let embeddings: Result>, _> = texts .par_chunks(self.prompt_count_in_chunk_hint()) - .map(move |chunk| self.embed_ref(chunk, None)) + .map(move |chunk| self.embed_ref(chunk, None, embedder_stats.clone())) .collect(); let embeddings = embeddings?; @@ -272,7 +278,7 @@ impl Embedder { } fn infer_dimensions(data: &EmbedderData) -> Result { - let v = embed(data, ["test"].as_slice(), 1, None, None) + let v = embed(data, ["test"].as_slice(), 1, None, None, None) .map_err(NewEmbedderError::could_not_determine_dimension)?; // unwrap: guaranteed that v.len() == 1, otherwise the previous line terminated in error Ok(v.first().unwrap().len()) @@ -284,6 +290,7 @@ fn embed( expected_count: usize, expected_dimension: Option, deadline: Option, + embedder_stats: Option>, ) -> Result, EmbedError> where S: Serialize, @@ -302,6 +309,9 @@ where let body = data.request.inject_texts(inputs); for attempt in 0..10 { + if let Some(embedder_stats) = &embedder_stats { + embedder_stats.as_ref().total_requests.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + } let response = request.clone().send_json(&body); let result = check_response(response, data.configuration_source).and_then(|response| { response_to_embedding(response, data, expected_count, expected_dimension) @@ -311,6 +321,12 @@ where Ok(response) => return Ok(response), Err(retry) => { tracing::warn!("Failed: {}", retry.error); + if let Some(embedder_stats) = &embedder_stats { + if let Ok(mut errors) = embedder_stats.errors.write() { + errors.0 = Some(retry.error.to_string()); + errors.1 += 1; + } + } if let Some(deadline) = deadline { let now = std::time::Instant::now(); if now > deadline { @@ -336,12 +352,26 @@ where std::thread::sleep(retry_duration); } + if let Some(embedder_stats) = &embedder_stats { + embedder_stats.as_ref().total_requests.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + } let response = request.send_json(&body); - let result = check_response(response, data.configuration_source); - result.map_err(Retry::into_error).and_then(|response| { + let result = check_response(response, data.configuration_source).and_then(|response| { response_to_embedding(response, data, expected_count, expected_dimension) - .map_err(Retry::into_error) - }) + }); + + match result { + Ok(response) => Ok(response), + Err(retry) => { + if let Some(embedder_stats) = &embedder_stats { + if let Ok(mut errors) = embedder_stats.errors.write() { + errors.0 = Some(retry.error.to_string()); + errors.1 += 1; + } + } + Err(retry.into_error()) + } + } } fn check_response( diff --git a/crates/milli/tests/search/distinct.rs b/crates/milli/tests/search/distinct.rs index fc890dfe8..55e43c8fa 100644 --- a/crates/milli/tests/search/distinct.rs +++ b/crates/milli/tests/search/distinct.rs @@ -19,7 +19,7 @@ macro_rules! test_distinct { let config = milli::update::IndexerConfig::default(); let mut builder = Settings::new(&mut wtxn, &index, &config); builder.set_distinct_field(S(stringify!($distinct))); - builder.execute(|_| (), || false).unwrap(); + builder.execute(|_| (), || false, None).unwrap(); wtxn.commit().unwrap(); let rtxn = index.read_txn().unwrap(); diff --git a/crates/milli/tests/search/facet_distribution.rs b/crates/milli/tests/search/facet_distribution.rs index 8934cbea4..588662735 100644 --- a/crates/milli/tests/search/facet_distribution.rs +++ b/crates/milli/tests/search/facet_distribution.rs @@ -25,7 +25,7 @@ fn test_facet_distribution_with_no_facet_values() { FilterableAttributesRule::Field(S("genres")), FilterableAttributesRule::Field(S("tags")), ]); - builder.execute(|_| (), || false).unwrap(); + builder.execute(|_| (), || false, None).unwrap(); wtxn.commit().unwrap(); // index documents diff --git a/crates/milli/tests/search/mod.rs b/crates/milli/tests/search/mod.rs index 906956716..1e0c24608 100644 --- a/crates/milli/tests/search/mod.rs +++ b/crates/milli/tests/search/mod.rs @@ -63,7 +63,7 @@ pub fn setup_search_index_with_criteria(criteria: &[Criterion]) -> Index { S("america") => vec![S("the united states")], }); builder.set_searchable_fields(vec![S("title"), S("description")]); - builder.execute(|_| (), || false).unwrap(); + builder.execute(|_| (), || false, None).unwrap(); wtxn.commit().unwrap(); // index documents diff --git a/crates/milli/tests/search/phrase_search.rs b/crates/milli/tests/search/phrase_search.rs index b7f792bfc..c5a95f7cd 100644 --- a/crates/milli/tests/search/phrase_search.rs +++ b/crates/milli/tests/search/phrase_search.rs @@ -10,7 +10,7 @@ fn set_stop_words(index: &Index, stop_words: &[&str]) { let mut builder = Settings::new(&mut wtxn, index, &config); let stop_words = stop_words.iter().map(|s| s.to_string()).collect(); builder.set_stop_words(stop_words); - builder.execute(|_| (), || false).unwrap(); + builder.execute(|_| (), || false, None).unwrap(); wtxn.commit().unwrap(); } diff --git a/crates/milli/tests/search/query_criteria.rs b/crates/milli/tests/search/query_criteria.rs index 1acc89484..b7614c215 100644 --- a/crates/milli/tests/search/query_criteria.rs +++ b/crates/milli/tests/search/query_criteria.rs @@ -236,7 +236,7 @@ fn criteria_mixup() { let mut wtxn = index.write_txn().unwrap(); let mut builder = Settings::new(&mut wtxn, &index, &config); builder.set_criteria(criteria.clone()); - builder.execute(|_| (), || false).unwrap(); + builder.execute(|_| (), || false, None).unwrap(); wtxn.commit().unwrap(); let rtxn = index.read_txn().unwrap(); @@ -276,7 +276,7 @@ fn criteria_ascdesc() { S("name"), S("age"), }); - builder.execute(|_| (), || false).unwrap(); + builder.execute(|_| (), || false, None).unwrap(); wtxn.commit().unwrap(); let mut wtxn = index.write_txn().unwrap(); @@ -358,7 +358,7 @@ fn criteria_ascdesc() { let mut wtxn = index.write_txn().unwrap(); let mut builder = Settings::new(&mut wtxn, &index, &config); builder.set_criteria(vec![criterion.clone()]); - builder.execute(|_| (), || false).unwrap(); + builder.execute(|_| (), || false, None).unwrap(); wtxn.commit().unwrap(); let rtxn = index.read_txn().unwrap(); diff --git a/crates/milli/tests/search/typo_tolerance.rs b/crates/milli/tests/search/typo_tolerance.rs index 3c0717063..bf9a730c9 100644 --- a/crates/milli/tests/search/typo_tolerance.rs +++ b/crates/milli/tests/search/typo_tolerance.rs @@ -46,7 +46,7 @@ fn test_typo_tolerance_one_typo() { let config = IndexerConfig::default(); let mut builder = Settings::new(&mut txn, &index, &config); builder.set_min_word_len_one_typo(4); - builder.execute(|_| (), || false).unwrap(); + builder.execute(|_| (), || false, None).unwrap(); // typo is now supported for 4 letters words let mut search = Search::new(&txn, &index); @@ -92,7 +92,7 @@ fn test_typo_tolerance_two_typo() { let config = IndexerConfig::default(); let mut builder = Settings::new(&mut txn, &index, &config); builder.set_min_word_len_two_typos(7); - builder.execute(|_| (), || false).unwrap(); + builder.execute(|_| (), || false, None).unwrap(); // typo is now supported for 4 letters words let mut search = Search::new(&txn, &index); @@ -180,7 +180,7 @@ fn test_typo_disabled_on_word() { // `zealand` doesn't allow typos anymore exact_words.insert("zealand".to_string()); builder.set_exact_words(exact_words); - builder.execute(|_| (), || false).unwrap(); + builder.execute(|_| (), || false, None).unwrap(); let mut search = Search::new(&txn, &index); search.query("zealand"); @@ -218,7 +218,7 @@ fn test_disable_typo_on_attribute() { let mut builder = Settings::new(&mut txn, &index, &config); // disable typos on `description` builder.set_exact_attributes(vec!["description".to_string()].into_iter().collect()); - builder.execute(|_| (), || false).unwrap(); + builder.execute(|_| (), || false, None).unwrap(); let mut search = Search::new(&txn, &index); search.query("antebelum"); From 4925b3019640107759c33fa99521b99c6c202277 Mon Sep 17 00:00:00 2001 From: Mubelotix Date: Mon, 23 Jun 2025 15:24:14 +0200 Subject: [PATCH 02/18] Move embedder stats out of progress --- crates/benchmarks/benches/indexing.rs | 33 ++++++++++++++- crates/benchmarks/benches/utils.rs | 3 +- crates/dump/src/lib.rs | 2 +- crates/fuzzers/src/bin/fuzz-indexing.rs | 1 + crates/index-scheduler/src/insta_snapshot.rs | 4 +- crates/index-scheduler/src/queue/batches.rs | 10 ++++- .../src/scheduler/process_batch.rs | 17 ++++++-- .../src/scheduler/process_index_operation.rs | 12 ++++-- crates/index-scheduler/src/utils.rs | 19 ++++++++- crates/meilisearch-types/src/batch_view.rs | 19 +++++++-- crates/meilisearch-types/src/batches.rs | 28 ++++++++++++- crates/meilisearch/src/lib.rs | 6 ++- crates/meilisearch/tests/vector/rest.rs | 40 +++++++++++++++++++ crates/milli/src/progress.rs | 37 ++++++----------- .../milli/src/search/new/tests/integration.rs | 3 +- crates/milli/src/test_index.rs | 5 ++- .../extract/extract_vector_points.rs | 2 + .../src/update/index_documents/extract/mod.rs | 6 +-- .../milli/src/update/index_documents/mod.rs | 17 +++++++- .../src/update/new/extract/vectors/mod.rs | 17 ++++++-- .../milli/src/update/new/indexer/extract.rs | 4 ++ crates/milli/src/update/new/indexer/mod.rs | 5 ++- crates/milli/src/update/settings.rs | 4 +- crates/milli/src/vector/rest.rs | 4 ++ crates/milli/tests/search/distinct.rs | 2 +- .../milli/tests/search/facet_distribution.rs | 3 +- crates/milli/tests/search/mod.rs | 3 +- crates/milli/tests/search/phrase_search.rs | 2 +- crates/milli/tests/search/query_criteria.rs | 7 ++-- crates/milli/tests/search/typo_tolerance.rs | 9 +++-- 30 files changed, 255 insertions(+), 69 deletions(-) diff --git a/crates/benchmarks/benches/indexing.rs b/crates/benchmarks/benches/indexing.rs index b882b598d..8241da9d2 100644 --- a/crates/benchmarks/benches/indexing.rs +++ b/crates/benchmarks/benches/indexing.rs @@ -65,7 +65,7 @@ fn setup_settings<'t>( let sortable_fields = sortable_fields.iter().map(|s| s.to_string()).collect(); builder.set_sortable_fields(sortable_fields); - builder.execute(|_| (), || false, None).unwrap(); + builder.execute(|_| (), || false, Default::default()).unwrap(); } fn setup_index_with_settings( @@ -169,6 +169,7 @@ fn indexing_songs_default(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), + Default::default(), ) .unwrap(); @@ -235,6 +236,7 @@ fn reindexing_songs_default(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), + Default::default(), ) .unwrap(); @@ -279,6 +281,7 @@ fn reindexing_songs_default(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), + Default::default(), ) .unwrap(); @@ -347,6 +350,7 @@ fn deleting_songs_in_batches_default(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), + Default::default(), ) .unwrap(); @@ -423,6 +427,7 @@ fn indexing_songs_in_three_batches_default(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), + Default::default(), ) .unwrap(); @@ -467,6 +472,7 @@ fn indexing_songs_in_three_batches_default(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), + Default::default(), ) .unwrap(); @@ -507,6 +513,7 @@ fn indexing_songs_in_three_batches_default(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), + Default::default(), ) .unwrap(); @@ -574,6 +581,7 @@ fn indexing_songs_without_faceted_numbers(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), + Default::default(), ) .unwrap(); @@ -640,6 +648,7 @@ fn indexing_songs_without_faceted_fields(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), + Default::default(), ) .unwrap(); @@ -706,6 +715,7 @@ fn indexing_wiki(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), + Default::default(), ) .unwrap(); @@ -771,6 +781,7 @@ fn reindexing_wiki(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), + Default::default(), ) .unwrap(); @@ -815,6 +826,7 @@ fn reindexing_wiki(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), + Default::default(), ) .unwrap(); @@ -882,6 +894,7 @@ fn deleting_wiki_in_batches_default(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), + Default::default(), ) .unwrap(); @@ -958,6 +971,7 @@ fn indexing_wiki_in_three_batches(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), + Default::default(), ) .unwrap(); @@ -1003,6 +1017,7 @@ fn indexing_wiki_in_three_batches(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), + Default::default(), ) .unwrap(); @@ -1044,6 +1059,7 @@ fn indexing_wiki_in_three_batches(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), + Default::default(), ) .unwrap(); @@ -1110,6 +1126,7 @@ fn indexing_movies_default(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), + Default::default(), ) .unwrap(); @@ -1175,6 +1192,7 @@ fn reindexing_movies_default(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), + Default::default(), ) .unwrap(); @@ -1219,6 +1237,7 @@ fn reindexing_movies_default(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), + Default::default(), ) .unwrap(); @@ -1286,6 +1305,7 @@ fn deleting_movies_in_batches_default(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), + Default::default(), ) .unwrap(); @@ -1334,6 +1354,7 @@ fn delete_documents_from_ids(index: Index, document_ids_to_delete: Vec Index { (conf.configure)(&mut builder); - builder.execute(|_| (), || false, None).unwrap(); + builder.execute(|_| (), || false, Default::default()).unwrap(); wtxn.commit().unwrap(); let config = IndexerConfig::default(); @@ -128,6 +128,7 @@ pub fn base_setup(conf: &Conf) -> Index { EmbeddingConfigs::default(), &|| false, &Progress::default(), + Default::default(), ) .unwrap(); diff --git a/crates/dump/src/lib.rs b/crates/dump/src/lib.rs index c48c68f62..b7a35ad5c 100644 --- a/crates/dump/src/lib.rs +++ b/crates/dump/src/lib.rs @@ -328,8 +328,8 @@ pub(crate) mod test { progress_trace: Default::default(), write_channel_congestion: None, internal_database_sizes: Default::default(), - embeddings: Default::default(), }, + embedder_stats: None, enqueued_at: Some(BatchEnqueuedAt { earliest: datetime!(2022-11-11 0:00 UTC), oldest: datetime!(2022-11-11 0:00 UTC), diff --git a/crates/fuzzers/src/bin/fuzz-indexing.rs b/crates/fuzzers/src/bin/fuzz-indexing.rs index 4df989b51..23c4cb9c2 100644 --- a/crates/fuzzers/src/bin/fuzz-indexing.rs +++ b/crates/fuzzers/src/bin/fuzz-indexing.rs @@ -144,6 +144,7 @@ fn main() { embedders, &|| false, &Progress::default(), + Default::default(), ) .unwrap(); diff --git a/crates/index-scheduler/src/insta_snapshot.rs b/crates/index-scheduler/src/insta_snapshot.rs index d01548319..06ec01b5e 100644 --- a/crates/index-scheduler/src/insta_snapshot.rs +++ b/crates/index-scheduler/src/insta_snapshot.rs @@ -1,7 +1,7 @@ use std::collections::BTreeSet; use std::fmt::Write; -use meilisearch_types::batches::{Batch, BatchEnqueuedAt, BatchStats}; +use meilisearch_types::batches::{Batch, BatchEmbeddingStats, BatchEnqueuedAt, BatchStats}; use meilisearch_types::heed::types::{SerdeBincode, SerdeJson, Str}; use meilisearch_types::heed::{Database, RoTxn}; use meilisearch_types::milli::{CboRoaringBitmapCodec, RoaringBitmapCodec, BEU32}; @@ -343,6 +343,7 @@ pub fn snapshot_batch(batch: &Batch) -> String { uid, details, stats, + embedder_stats, started_at, finished_at, progress: _, @@ -366,6 +367,7 @@ pub fn snapshot_batch(batch: &Batch) -> String { snap.push_str(&format!("uid: {uid}, ")); snap.push_str(&format!("details: {}, ", serde_json::to_string(details).unwrap())); snap.push_str(&format!("stats: {}, ", serde_json::to_string(&stats).unwrap())); + snap.push_str(&format!("embedder_stats: {}, ", serde_json::to_string(&embedder_stats).unwrap())); snap.push_str(&format!("stop reason: {}, ", serde_json::to_string(&stop_reason).unwrap())); snap.push('}'); snap diff --git a/crates/index-scheduler/src/queue/batches.rs b/crates/index-scheduler/src/queue/batches.rs index b5b63e1d7..b14601733 100644 --- a/crates/index-scheduler/src/queue/batches.rs +++ b/crates/index-scheduler/src/queue/batches.rs @@ -1,7 +1,7 @@ use std::collections::HashSet; use std::ops::{Bound, RangeBounds}; -use meilisearch_types::batches::{Batch, BatchId}; +use meilisearch_types::batches::{Batch, BatchEmbeddingStats, BatchId}; use meilisearch_types::heed::types::{DecodeIgnore, SerdeBincode, SerdeJson, Str}; use meilisearch_types::heed::{Database, Env, RoTxn, RwTxn, WithoutTls}; use meilisearch_types::milli::{CboRoaringBitmapCodec, RoaringBitmapCodec, BEU32}; @@ -92,7 +92,10 @@ impl BatchQueue { } pub(crate) fn get_batch(&self, rtxn: &RoTxn, batch_id: BatchId) -> Result> { - Ok(self.all_batches.get(rtxn, &batch_id)?) + println!("Got batch from db {batch_id:?}"); + let r = Ok(self.all_batches.get(rtxn, &batch_id)?); + println!("Got batch from db => {:?}", r); + r } /// Returns the whole set of batches that belongs to this index. @@ -171,6 +174,8 @@ impl BatchQueue { pub(crate) fn write_batch(&self, wtxn: &mut RwTxn, batch: ProcessingBatch) -> Result<()> { let old_batch = self.all_batches.get(wtxn, &batch.uid)?; + println!("Saving batch: {}", batch.embedder_stats.is_some()); + self.all_batches.put( wtxn, &batch.uid, @@ -179,6 +184,7 @@ impl BatchQueue { progress: None, details: batch.details, stats: batch.stats, + embedder_stats: batch.embedder_stats.as_ref().map(|s| BatchEmbeddingStats::from(s.as_ref())), started_at: batch.started_at, finished_at: batch.finished_at, enqueued_at: batch.enqueued_at, diff --git a/crates/index-scheduler/src/scheduler/process_batch.rs b/crates/index-scheduler/src/scheduler/process_batch.rs index 71e423a58..4e36b65b6 100644 --- a/crates/index-scheduler/src/scheduler/process_batch.rs +++ b/crates/index-scheduler/src/scheduler/process_batch.rs @@ -1,10 +1,11 @@ use std::collections::{BTreeSet, HashMap, HashSet}; use std::panic::{catch_unwind, AssertUnwindSafe}; use std::sync::atomic::Ordering; +use std::sync::Arc; use meilisearch_types::batches::{BatchEnqueuedAt, BatchId}; use meilisearch_types::heed::{RoTxn, RwTxn}; -use meilisearch_types::milli::progress::{Progress, VariableNameStep}; +use meilisearch_types::milli::progress::{EmbedderStats, Progress, VariableNameStep}; use meilisearch_types::milli::{self, ChannelCongestion}; use meilisearch_types::tasks::{Details, IndexSwap, Kind, KindWithContent, Status, Task}; use meilisearch_types::versioning::{VERSION_MAJOR, VERSION_MINOR, VERSION_PATCH}; @@ -163,7 +164,7 @@ impl IndexScheduler { let pre_commit_dabases_sizes = index.database_sizes(&index_wtxn)?; let (tasks, congestion) = - self.apply_index_operation(&mut index_wtxn, &index, op, &progress)?; + self.apply_index_operation(&mut index_wtxn, &index, op, &progress, current_batch.clone_embedder_stats())?; { progress.update_progress(FinalizingIndexStep::Committing); @@ -238,11 +239,21 @@ impl IndexScheduler { ); builder.set_primary_key(primary_key); let must_stop_processing = self.scheduler.must_stop_processing.clone(); + + let embedder_stats = match current_batch.embedder_stats { + Some(ref stats) => stats.clone(), + None => { + let embedder_stats: Arc = Default::default(); + current_batch.embedder_stats = Some(embedder_stats.clone()); + embedder_stats + }, + }; + builder .execute( |indexing_step| tracing::debug!(update = ?indexing_step), || must_stop_processing.get(), - Some(progress.embedder_stats), + embedder_stats, ) .map_err(|e| Error::from_milli(e, Some(index_uid.to_string())))?; index_wtxn.commit()?; diff --git a/crates/index-scheduler/src/scheduler/process_index_operation.rs b/crates/index-scheduler/src/scheduler/process_index_operation.rs index 92d13e7e7..b5338e511 100644 --- a/crates/index-scheduler/src/scheduler/process_index_operation.rs +++ b/crates/index-scheduler/src/scheduler/process_index_operation.rs @@ -4,7 +4,7 @@ use bumpalo::collections::CollectIn; use bumpalo::Bump; use meilisearch_types::heed::RwTxn; use meilisearch_types::milli::documents::PrimaryKey; -use meilisearch_types::milli::progress::Progress; +use meilisearch_types::milli::progress::{EmbedderStats, Progress}; use meilisearch_types::milli::update::new::indexer::{self, UpdateByFunction}; use meilisearch_types::milli::update::DocumentAdditionResult; use meilisearch_types::milli::{self, ChannelCongestion, Filter}; @@ -26,7 +26,7 @@ impl IndexScheduler { /// The list of processed tasks. #[tracing::instrument( level = "trace", - skip(self, index_wtxn, index, progress), + skip(self, index_wtxn, index, progress, embedder_stats), target = "indexing::scheduler" )] pub(crate) fn apply_index_operation<'i>( @@ -35,6 +35,7 @@ impl IndexScheduler { index: &'i Index, operation: IndexOperation, progress: &Progress, + embedder_stats: Arc, ) -> Result<(Vec, Option)> { let indexer_alloc = Bump::new(); let started_processing_at = std::time::Instant::now(); @@ -179,6 +180,7 @@ impl IndexScheduler { embedders, &|| must_stop_processing.get(), progress, + embedder_stats, ) .map_err(|e| Error::from_milli(e, Some(index_uid.clone())))?, ); @@ -290,6 +292,7 @@ impl IndexScheduler { embedders, &|| must_stop_processing.get(), progress, + embedder_stats, ) .map_err(|err| Error::from_milli(err, Some(index_uid.clone())))?, ); @@ -438,6 +441,7 @@ impl IndexScheduler { embedders, &|| must_stop_processing.get(), progress, + embedder_stats, ) .map_err(|err| Error::from_milli(err, Some(index_uid.clone())))?, ); @@ -474,7 +478,7 @@ impl IndexScheduler { .execute( |indexing_step| tracing::debug!(update = ?indexing_step), || must_stop_processing.get(), - Some(Arc::clone(&progress.embedder_stats)) + embedder_stats, ) .map_err(|err| Error::from_milli(err, Some(index_uid.clone())))?; @@ -494,6 +498,7 @@ impl IndexScheduler { tasks: cleared_tasks, }, progress, + embedder_stats.clone(), )?; let (settings_tasks, _congestion) = self.apply_index_operation( @@ -501,6 +506,7 @@ impl IndexScheduler { index, IndexOperation::Settings { index_uid, settings, tasks: settings_tasks }, progress, + embedder_stats, )?; let mut tasks = settings_tasks; diff --git a/crates/index-scheduler/src/utils.rs b/crates/index-scheduler/src/utils.rs index 67e8fc090..22e319580 100644 --- a/crates/index-scheduler/src/utils.rs +++ b/crates/index-scheduler/src/utils.rs @@ -2,8 +2,10 @@ use std::collections::{BTreeSet, HashSet}; use std::ops::Bound; +use std::sync::Arc; +use crate::milli::progress::EmbedderStats; -use meilisearch_types::batches::{Batch, BatchEnqueuedAt, BatchId, BatchStats}; +use meilisearch_types::batches::{Batch, BatchEmbeddingStats, BatchEnqueuedAt, BatchId, BatchStats}; use meilisearch_types::heed::{Database, RoTxn, RwTxn}; use meilisearch_types::milli::CboRoaringBitmapCodec; use meilisearch_types::task_view::DetailsView; @@ -27,6 +29,7 @@ pub struct ProcessingBatch { pub uid: BatchId, pub details: DetailsView, pub stats: BatchStats, + pub embedder_stats: Option>, pub statuses: HashSet, pub kinds: HashSet, @@ -48,6 +51,7 @@ impl ProcessingBatch { uid, details: DetailsView::default(), stats: BatchStats::default(), + embedder_stats: None, statuses, kinds: HashSet::default(), @@ -60,6 +64,17 @@ impl ProcessingBatch { } } + pub fn clone_embedder_stats(&mut self) -> Arc { + match self.embedder_stats { + Some(ref stats) => stats.clone(), + None => { + let embedder_stats: Arc = Default::default(); + self.embedder_stats = Some(embedder_stats.clone()); + embedder_stats + }, + } + } + /// Update itself with the content of the task and update the batch id in the task. pub fn processing<'a>(&mut self, tasks: impl IntoIterator) { for task in tasks.into_iter() { @@ -141,11 +156,13 @@ impl ProcessingBatch { } pub fn to_batch(&self) -> Batch { + println!("Converting to batch: {:?}", self.embedder_stats); Batch { uid: self.uid, progress: None, details: self.details.clone(), stats: self.stats.clone(), + embedder_stats: self.embedder_stats.as_ref().map(|s| BatchEmbeddingStats::from(s.as_ref())), started_at: self.started_at, finished_at: self.finished_at, enqueued_at: self.enqueued_at, diff --git a/crates/meilisearch-types/src/batch_view.rs b/crates/meilisearch-types/src/batch_view.rs index f0a5f364b..0a9b80f4e 100644 --- a/crates/meilisearch-types/src/batch_view.rs +++ b/crates/meilisearch-types/src/batch_view.rs @@ -3,7 +3,7 @@ use serde::Serialize; use time::{Duration, OffsetDateTime}; use utoipa::ToSchema; -use crate::batches::{Batch, BatchId, BatchStats}; +use crate::batches::{Batch, BatchEmbeddingStats, BatchId, BatchStats}; use crate::task_view::DetailsView; use crate::tasks::serialize_duration; @@ -14,7 +14,7 @@ pub struct BatchView { pub uid: BatchId, pub progress: Option, pub details: DetailsView, - pub stats: BatchStats, + pub stats: BatchStatsView, #[serde(serialize_with = "serialize_duration", default)] pub duration: Option, #[serde(with = "time::serde::rfc3339", default)] @@ -25,13 +25,26 @@ pub struct BatchView { pub batch_strategy: String, } +#[derive(Debug, Clone, Serialize, ToSchema)] +#[serde(rename_all = "camelCase")] +#[schema(rename_all = "camelCase")] +pub struct BatchStatsView { + #[serde(flatten)] + pub stats: BatchStats, + #[serde(skip_serializing_if = "BatchEmbeddingStats::skip_serializing")] + pub embedder: Option, +} + impl BatchView { pub fn from_batch(batch: &Batch) -> Self { Self { uid: batch.uid, progress: batch.progress.clone(), details: batch.details.clone(), - stats: batch.stats.clone(), + stats: BatchStatsView { + stats: batch.stats.clone(), + embedder: batch.embedder_stats.clone(), + }, duration: batch.finished_at.map(|finished_at| finished_at - batch.started_at), started_at: batch.started_at, finished_at: batch.finished_at, diff --git a/crates/meilisearch-types/src/batches.rs b/crates/meilisearch-types/src/batches.rs index 2ef373eac..24be75d1c 100644 --- a/crates/meilisearch-types/src/batches.rs +++ b/crates/meilisearch-types/src/batches.rs @@ -1,6 +1,7 @@ use std::collections::BTreeMap; +use std::sync::Arc; -use milli::progress::ProgressView; +use milli::progress::{EmbedderStats, ProgressView}; use serde::{Deserialize, Serialize}; use time::OffsetDateTime; use utoipa::ToSchema; @@ -19,6 +20,7 @@ pub struct Batch { pub progress: Option, pub details: DetailsView, pub stats: BatchStats, + pub embedder_stats: Option, #[serde(with = "time::serde::rfc3339")] pub started_at: OffsetDateTime, @@ -43,6 +45,7 @@ impl PartialEq for Batch { progress, details, stats, + embedder_stats, started_at, finished_at, enqueued_at, @@ -53,6 +56,7 @@ impl PartialEq for Batch { && progress.is_none() == other.progress.is_none() && details == &other.details && stats == &other.stats + && embedder_stats == &other.embedder_stats && started_at == &other.started_at && finished_at == &other.finished_at && enqueued_at == &other.enqueued_at @@ -82,7 +86,6 @@ pub struct BatchStats { pub write_channel_congestion: Option>, #[serde(default, skip_serializing_if = "serde_json::Map::is_empty")] pub internal_database_sizes: serde_json::Map, - pub embeddings: BatchEmbeddingStats } #[derive(Default, Debug, Clone, PartialEq, Eq, Serialize, Deserialize, ToSchema)] @@ -91,5 +94,26 @@ pub struct BatchStats { pub struct BatchEmbeddingStats { pub total_count: usize, pub error_count: usize, + #[serde(skip_serializing_if = "Option::is_none")] pub last_error: Option, } + +impl From<&EmbedderStats> for BatchEmbeddingStats { + fn from(stats: &EmbedderStats) -> Self { + let errors = stats.errors.read().unwrap(); + Self { + total_count: stats.total_requests.load(std::sync::atomic::Ordering::Relaxed), + error_count: errors.1 as usize, + last_error: errors.0.clone(), + } + } +} + +impl BatchEmbeddingStats { + pub fn skip_serializing(this: &Option) -> bool { + match this { + Some(stats) => stats.total_count == 0 && stats.error_count == 0 && stats.last_error.is_none(), + None => true, + } + } +} diff --git a/crates/meilisearch/src/lib.rs b/crates/meilisearch/src/lib.rs index 782d6172f..72be6aec9 100644 --- a/crates/meilisearch/src/lib.rs +++ b/crates/meilisearch/src/lib.rs @@ -37,6 +37,7 @@ use index_scheduler::{IndexScheduler, IndexSchedulerOptions}; use meilisearch_auth::{open_auth_store_env, AuthController}; use meilisearch_types::milli::constants::VERSION_MAJOR; use meilisearch_types::milli::documents::{DocumentsBatchBuilder, DocumentsBatchReader}; +use meilisearch_types::milli::progress::EmbedderStats; use meilisearch_types::milli::update::{ default_thread_pool_and_threads, IndexDocumentsConfig, IndexDocumentsMethod, IndexerConfig, }; @@ -542,8 +543,9 @@ fn import_dump( tracing::info!("Importing the settings."); let settings = index_reader.settings()?; apply_settings_to_builder(&settings, &mut builder); + let embedder_stats: Arc = Default::default(); // FIXME: this isn't linked to anything builder - .execute(|indexing_step| tracing::debug!("update: {:?}", indexing_step), || false, None)?; + .execute(|indexing_step| tracing::debug!("update: {:?}", indexing_step), || false, embedder_stats.clone())?; // 4.3 Import the documents. // 4.3.1 We need to recreate the grenad+obkv format accepted by the index. @@ -574,7 +576,7 @@ fn import_dump( }, |indexing_step| tracing::trace!("update: {:?}", indexing_step), || false, - None, + embedder_stats, )?; let builder = builder.with_embedders(embedders); diff --git a/crates/meilisearch/tests/vector/rest.rs b/crates/meilisearch/tests/vector/rest.rs index 82fc71b26..1ff2dd9fe 100644 --- a/crates/meilisearch/tests/vector/rest.rs +++ b/crates/meilisearch/tests/vector/rest.rs @@ -4,6 +4,8 @@ use meili_snap::{json_string, snapshot}; use reqwest::IntoUrl; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, Request, ResponseTemplate}; +use std::thread::sleep; +use std::time::Duration; use crate::common::Value; use crate::json; @@ -305,6 +307,7 @@ async fn create_mock_raw() -> (MockServer, Value) { Mock::given(method("POST")) .and(path("/")) .respond_with(move |req: &Request| { + println!("Sent!"); let req: String = match req.body_json() { Ok(req) => req, Err(error) => { @@ -2111,3 +2114,40 @@ async fn searchable_reindex() { } "###); } + + +#[actix_rt::test] +async fn observability() { + let (_mock, setting) = create_mock_raw().await; + let server = get_server_vector().await; + let index = server.index("doggo"); + + let (response, code) = index + .update_settings(json!({ + "embedders": { + "rest": setting, + }, + })) + .await; + snapshot!(code, @"202 Accepted"); + let task = server.wait_task(response.uid()).await; + snapshot!(task["status"], @r###""succeeded""###); + let documents = json!([ + {"id": 0, "name": "kefir"}, + {"id": 1, "name": "echo", "_vectors": { "rest": [1, 1, 1] }}, + {"id": 2, "name": "intel"}, + {"id": 3, "name": "missing"}, // Stuff that doesn't exist + {"id": 4, "name": "invalid"}, + {"id": 5, "name": "foobar"}, + ]); + let (value, code) = index.add_documents(documents, None).await; + snapshot!(code, @"202 Accepted"); + + let batches = index.filtered_batches(&[], &[], &[]).await; + println!("Batches: {batches:?}"); + + let task = index.wait_task(value.uid()).await; + let batches = index.filtered_batches(&[], &[], &[]).await; + println!("Batches: {batches:?}"); + +} diff --git a/crates/milli/src/progress.rs b/crates/milli/src/progress.rs index ff795b220..7026f0c11 100644 --- a/crates/milli/src/progress.rs +++ b/crates/milli/src/progress.rs @@ -20,7 +20,6 @@ pub trait Step: 'static + Send + Sync { #[derive(Clone, Default)] pub struct Progress { steps: Arc>, - pub embedder_stats: Arc, } #[derive(Default)] @@ -29,6 +28,17 @@ pub struct EmbedderStats { pub total_requests: AtomicUsize } +impl std::fmt::Debug for EmbedderStats { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let (error, count) = self.errors.read().unwrap().clone(); + f.debug_struct("EmbedderStats") + .field("errors", &error) + .field("total_requests", &self.total_requests.load(Ordering::Relaxed)) + .field("error_count", &count) + .finish() + } +} + #[derive(Default)] struct InnerProgress { /// The hierarchy of steps. @@ -72,19 +82,7 @@ impl Progress { }); } - let embedder_view = { - let (last_error, error_count) = match self.embedder_stats.errors.read() { - Ok(guard) => (guard.0.clone(), guard.1), - Err(_) => (None, 0), - }; - EmbedderStatsView { - last_error, - request_count: self.embedder_stats.total_requests.load(Ordering::Relaxed) as u32, - error_count, - } - }; - - ProgressView { steps: step_view, percentage: percentage * 100.0, embedder: embedder_view } + ProgressView { steps: step_view, percentage: percentage * 100.0 } } pub fn accumulated_durations(&self) -> IndexMap { @@ -228,7 +226,6 @@ make_enum_progress! { pub struct ProgressView { pub steps: Vec, pub percentage: f32, - pub embedder: EmbedderStatsView, } #[derive(Debug, Serialize, Clone, ToSchema)] @@ -240,16 +237,6 @@ pub struct ProgressStepView { pub total: u32, } -#[derive(Debug, Serialize, Clone, ToSchema)] -#[serde(rename_all = "camelCase")] -#[schema(rename_all = "camelCase")] -pub struct EmbedderStatsView { - #[serde(skip_serializing_if = "Option::is_none")] - pub last_error: Option, - pub request_count: u32, - pub error_count: u32, -} - /// Used when the name can change but it's still the same step. /// To avoid conflicts on the `TypeId`, create a unique type every time you use this step: /// ```text diff --git a/crates/milli/src/search/new/tests/integration.rs b/crates/milli/src/search/new/tests/integration.rs index e7634a4eb..0b7e1a292 100644 --- a/crates/milli/src/search/new/tests/integration.rs +++ b/crates/milli/src/search/new/tests/integration.rs @@ -44,7 +44,7 @@ pub fn setup_search_index_with_criteria(criteria: &[Criterion]) -> Index { S("america") => vec![S("the united states")], }); builder.set_searchable_fields(vec![S("title"), S("description")]); - builder.execute(|_| (), || false, None).unwrap(); + builder.execute(|_| (), || false, Default::default()).unwrap(); wtxn.commit().unwrap(); // index documents @@ -95,6 +95,7 @@ pub fn setup_search_index_with_criteria(criteria: &[Criterion]) -> Index { embedders, &|| false, &Progress::default(), + Default::default(), ) .unwrap(); diff --git a/crates/milli/src/test_index.rs b/crates/milli/src/test_index.rs index 634d45195..3546660b0 100644 --- a/crates/milli/src/test_index.rs +++ b/crates/milli/src/test_index.rs @@ -103,6 +103,7 @@ impl TempIndex { embedders, &|| false, &Progress::default(), + Default::default(), ) }) .unwrap()?; @@ -134,7 +135,7 @@ impl TempIndex { ) -> Result<(), crate::error::Error> { let mut builder = update::Settings::new(wtxn, &self.inner, &self.indexer_config); update(&mut builder); - builder.execute(drop, || false, None)?; + builder.execute(drop, || false, Default::default())?; Ok(()) } @@ -185,6 +186,7 @@ impl TempIndex { embedders, &|| false, &Progress::default(), + Default::default(), ) }) .unwrap()?; @@ -259,6 +261,7 @@ fn aborting_indexation() { embedders, &|| should_abort.load(Relaxed), &Progress::default(), + Default::default(), ) }) .unwrap() diff --git a/crates/milli/src/update/index_documents/extract/extract_vector_points.rs b/crates/milli/src/update/index_documents/extract/extract_vector_points.rs index 5e6bde53d..de91e9f10 100644 --- a/crates/milli/src/update/index_documents/extract/extract_vector_points.rs +++ b/crates/milli/src/update/index_documents/extract/extract_vector_points.rs @@ -687,6 +687,8 @@ pub fn extract_embeddings( unused_vectors_distribution: &UnusedVectorsDistribution, request_threads: &ThreadPoolNoAbort, ) -> Result>> { + println!("Extract embedder stats {}:", embedder_stats.is_some()); + let n_chunks = embedder.chunk_count_hint(); // chunk level parallelism let n_vectors_per_chunk = embedder.prompt_count_in_chunk_hint(); // number of vectors in a single chunk diff --git a/crates/milli/src/update/index_documents/extract/mod.rs b/crates/milli/src/update/index_documents/extract/mod.rs index 020b48f2c..f4f3ad22e 100644 --- a/crates/milli/src/update/index_documents/extract/mod.rs +++ b/crates/milli/src/update/index_documents/extract/mod.rs @@ -50,7 +50,7 @@ pub(crate) fn data_from_obkv_documents( settings_diff: Arc, max_positions_per_attributes: Option, possible_embedding_mistakes: Arc, - embedder_stats: Option>, + embedder_stats: Arc, ) -> Result<()> { let (original_pipeline_result, flattened_pipeline_result): (Result<_>, Result<_>) = rayon::join( || { @@ -234,7 +234,7 @@ fn send_original_documents_data( embedders_configs: Arc>, settings_diff: Arc, possible_embedding_mistakes: Arc, - embedder_stats: Option>, + embedder_stats: Arc, ) -> Result<()> { let original_documents_chunk = original_documents_chunk.and_then(|c| unsafe { as_cloneable_grenad(&c) })?; @@ -274,7 +274,7 @@ fn send_original_documents_data( embedder.clone(), &embedder_name, &possible_embedding_mistakes, - embedder_stats.clone(), + Some(embedder_stats.clone()), &unused_vectors_distribution, request_threads(), ) { diff --git a/crates/milli/src/update/index_documents/mod.rs b/crates/milli/src/update/index_documents/mod.rs index fad43bd30..f2e1783e4 100644 --- a/crates/milli/src/update/index_documents/mod.rs +++ b/crates/milli/src/update/index_documents/mod.rs @@ -81,7 +81,7 @@ pub struct IndexDocuments<'t, 'i, 'a, FP, FA> { added_documents: u64, deleted_documents: u64, embedders: EmbeddingConfigs, - embedder_stats: Option>, + embedder_stats: Arc, } #[derive(Default, Debug, Clone)] @@ -104,7 +104,7 @@ where config: IndexDocumentsConfig, progress: FP, should_abort: FA, - embedder_stats: Option>, + embedder_stats: Arc, ) -> Result> { let transform = Some(Transform::new( wtxn, @@ -2030,6 +2030,7 @@ mod tests { EmbeddingConfigs::default(), &|| false, &Progress::default(), + Default::default(), ) .unwrap(); wtxn.commit().unwrap(); @@ -2117,6 +2118,7 @@ mod tests { EmbeddingConfigs::default(), &|| false, &Progress::default(), + Default::default(), ) .unwrap(); wtxn.commit().unwrap(); @@ -2302,6 +2304,7 @@ mod tests { embedders, &|| false, &Progress::default(), + Default::default(), ) .unwrap(); wtxn.commit().unwrap(); @@ -2364,6 +2367,7 @@ mod tests { embedders, &|| false, &Progress::default(), + Default::default(), ) .unwrap(); wtxn.commit().unwrap(); @@ -2417,6 +2421,7 @@ mod tests { embedders, &|| false, &Progress::default(), + Default::default(), ) .unwrap(); wtxn.commit().unwrap(); @@ -2469,6 +2474,7 @@ mod tests { embedders, &|| false, &Progress::default(), + Default::default(), ) .unwrap(); wtxn.commit().unwrap(); @@ -2523,6 +2529,7 @@ mod tests { embedders, &|| false, &Progress::default(), + Default::default(), ) .unwrap(); wtxn.commit().unwrap(); @@ -2582,6 +2589,7 @@ mod tests { embedders, &|| false, &Progress::default(), + Default::default(), ) .unwrap(); wtxn.commit().unwrap(); @@ -2634,6 +2642,7 @@ mod tests { embedders, &|| false, &Progress::default(), + Default::default(), ) .unwrap(); wtxn.commit().unwrap(); @@ -2686,6 +2695,7 @@ mod tests { embedders, &|| false, &Progress::default(), + Default::default(), ) .unwrap(); wtxn.commit().unwrap(); @@ -2884,6 +2894,7 @@ mod tests { embedders, &|| false, &Progress::default(), + Default::default(), ) .unwrap(); wtxn.commit().unwrap(); @@ -2943,6 +2954,7 @@ mod tests { embedders, &|| false, &Progress::default(), + Default::default(), ) .unwrap(); wtxn.commit().unwrap(); @@ -2999,6 +3011,7 @@ mod tests { embedders, &|| false, &Progress::default(), + Default::default(), ) .unwrap(); wtxn.commit().unwrap(); diff --git a/crates/milli/src/update/new/extract/vectors/mod.rs b/crates/milli/src/update/new/extract/vectors/mod.rs index 5b6559d74..f720b81e2 100644 --- a/crates/milli/src/update/new/extract/vectors/mod.rs +++ b/crates/milli/src/update/new/extract/vectors/mod.rs @@ -1,4 +1,5 @@ -use std::cell::RefCell; +use std::f32::consts::E; +use std::{cell::RefCell, sync::Arc}; use bumpalo::collections::Vec as BVec; use bumpalo::Bump; @@ -6,6 +7,7 @@ use hashbrown::{DefaultHashBuilder, HashMap}; use super::cache::DelAddRoaringBitmap; use crate::error::FaultSource; +use crate::progress::EmbedderStats; use crate::prompt::Prompt; use crate::update::new::channel::EmbeddingSender; use crate::update::new::indexer::document_changes::{DocumentChangeContext, Extractor}; @@ -22,6 +24,7 @@ pub struct EmbeddingExtractor<'a, 'b> { embedders: &'a EmbeddingConfigs, sender: EmbeddingSender<'a, 'b>, possible_embedding_mistakes: PossibleEmbeddingMistakes, + embedder_stats: Option>, threads: &'a ThreadPoolNoAbort, } @@ -30,10 +33,11 @@ impl<'a, 'b> EmbeddingExtractor<'a, 'b> { embedders: &'a EmbeddingConfigs, sender: EmbeddingSender<'a, 'b>, field_distribution: &'a FieldDistribution, + embedder_stats: Option>, threads: &'a ThreadPoolNoAbort, ) -> Self { let possible_embedding_mistakes = PossibleEmbeddingMistakes::new(field_distribution); - Self { embedders, sender, threads, possible_embedding_mistakes } + Self { embedders, sender, threads, possible_embedding_mistakes, embedder_stats } } } @@ -75,6 +79,7 @@ impl<'extractor> Extractor<'extractor> for EmbeddingExtractor<'_, '_> { prompt, context.data, &self.possible_embedding_mistakes, + self.embedder_stats.clone(), self.threads, self.sender, &context.doc_alloc, @@ -307,6 +312,7 @@ struct Chunks<'a, 'b, 'extractor> { dimensions: usize, prompt: &'a Prompt, possible_embedding_mistakes: &'a PossibleEmbeddingMistakes, + embedder_stats: Option>, user_provided: &'a RefCell>, threads: &'a ThreadPoolNoAbort, sender: EmbeddingSender<'a, 'b>, @@ -322,6 +328,7 @@ impl<'a, 'b, 'extractor> Chunks<'a, 'b, 'extractor> { prompt: &'a Prompt, user_provided: &'a RefCell>, possible_embedding_mistakes: &'a PossibleEmbeddingMistakes, + embedder_stats: Option>, threads: &'a ThreadPoolNoAbort, sender: EmbeddingSender<'a, 'b>, doc_alloc: &'a Bump, @@ -336,6 +343,7 @@ impl<'a, 'b, 'extractor> Chunks<'a, 'b, 'extractor> { embedder, prompt, possible_embedding_mistakes, + embedder_stats, threads, sender, embedder_id, @@ -371,6 +379,7 @@ impl<'a, 'b, 'extractor> Chunks<'a, 'b, 'extractor> { self.embedder_id, self.embedder_name, self.possible_embedding_mistakes, + self.embedder_stats.clone(), unused_vectors_distribution, self.threads, self.sender, @@ -389,6 +398,7 @@ impl<'a, 'b, 'extractor> Chunks<'a, 'b, 'extractor> { self.embedder_id, self.embedder_name, self.possible_embedding_mistakes, + self.embedder_stats.clone(), unused_vectors_distribution, self.threads, self.sender, @@ -407,6 +417,7 @@ impl<'a, 'b, 'extractor> Chunks<'a, 'b, 'extractor> { embedder_id: u8, embedder_name: &str, possible_embedding_mistakes: &PossibleEmbeddingMistakes, + embedder_stats: Option>, unused_vectors_distribution: &UnusedVectorsDistributionBump, threads: &ThreadPoolNoAbort, sender: EmbeddingSender<'a, 'b>, @@ -450,7 +461,7 @@ impl<'a, 'b, 'extractor> Chunks<'a, 'b, 'extractor> { return Err(crate::Error::UserError(crate::UserError::DocumentEmbeddingError(msg))); } - let res = match embedder.embed_index_ref(texts.as_slice(), threads, None) { + let res = match embedder.embed_index_ref(texts.as_slice(), threads, embedder_stats) { Ok(embeddings) => { for (docid, embedding) in ids.into_iter().zip(embeddings) { sender.set_vector(*docid, embedder_id, embedding).unwrap(); diff --git a/crates/milli/src/update/new/indexer/extract.rs b/crates/milli/src/update/new/indexer/extract.rs index bb36ddc37..72c63b605 100644 --- a/crates/milli/src/update/new/indexer/extract.rs +++ b/crates/milli/src/update/new/indexer/extract.rs @@ -1,6 +1,7 @@ use std::collections::BTreeMap; use std::sync::atomic::AtomicBool; use std::sync::OnceLock; +use std::sync::Arc; use bumpalo::Bump; use roaring::RoaringBitmap; @@ -13,6 +14,7 @@ use super::super::thread_local::{FullySend, ThreadLocal}; use super::super::FacetFieldIdsDelta; use super::document_changes::{extract, DocumentChanges, IndexingContext}; use crate::index::IndexEmbeddingConfig; +use crate::progress::EmbedderStats; use crate::progress::MergingWordCache; use crate::proximity::ProximityPrecision; use crate::update::new::extract::EmbeddingExtractor; @@ -34,6 +36,7 @@ pub(super) fn extract_all<'pl, 'extractor, DC, MSP>( mut index_embeddings: Vec, document_ids: &mut RoaringBitmap, modified_docids: &mut RoaringBitmap, + embedder_stats: Arc, ) -> Result<(FacetFieldIdsDelta, Vec)> where DC: DocumentChanges<'pl>, @@ -245,6 +248,7 @@ where embedders, embedding_sender, field_distribution, + Some(embedder_stats), request_threads(), ); let mut datastore = ThreadLocal::with_capacity(rayon::current_num_threads()); diff --git a/crates/milli/src/update/new/indexer/mod.rs b/crates/milli/src/update/new/indexer/mod.rs index 2ea3c787e..52fd6cd0b 100644 --- a/crates/milli/src/update/new/indexer/mod.rs +++ b/crates/milli/src/update/new/indexer/mod.rs @@ -1,6 +1,7 @@ use std::sync::atomic::AtomicBool; use std::sync::{Once, RwLock}; use std::thread::{self, Builder}; +use std::sync::Arc; use big_s::S; use document_changes::{DocumentChanges, IndexingContext}; @@ -19,7 +20,7 @@ use super::steps::IndexingStep; use super::thread_local::ThreadLocal; use crate::documents::PrimaryKey; use crate::fields_ids_map::metadata::{FieldIdMapWithMetadata, MetadataBuilder}; -use crate::progress::Progress; +use crate::progress::{EmbedderStats, Progress}; use crate::update::GrenadParameters; use crate::vector::{ArroyWrapper, EmbeddingConfigs}; use crate::{FieldsIdsMap, GlobalFieldsIdsMap, Index, InternalError, Result, ThreadPoolNoAbort}; @@ -55,6 +56,7 @@ pub fn index<'pl, 'indexer, 'index, DC, MSP>( embedders: EmbeddingConfigs, must_stop_processing: &'indexer MSP, progress: &'indexer Progress, + embedder_stats: Arc, ) -> Result where DC: DocumentChanges<'pl>, @@ -158,6 +160,7 @@ where index_embeddings, document_ids, modified_docids, + embedder_stats, ) }) .unwrap() diff --git a/crates/milli/src/update/settings.rs b/crates/milli/src/update/settings.rs index 7c5a70aa3..98ee86978 100644 --- a/crates/milli/src/update/settings.rs +++ b/crates/milli/src/update/settings.rs @@ -475,7 +475,7 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> { progress_callback: &FP, should_abort: &FA, settings_diff: InnerIndexSettingsDiff, - embedder_stats: Option>, + embedder_stats: Arc, ) -> Result<()> where FP: Fn(UpdateIndexingStep) + Sync, @@ -1358,7 +1358,7 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> { } } - pub fn execute(mut self, progress_callback: FP, should_abort: FA, embedder_stats: Option>) -> Result<()> + pub fn execute(mut self, progress_callback: FP, should_abort: FA, embedder_stats: Arc) -> Result<()> where FP: Fn(UpdateIndexingStep) + Sync, FA: Fn() -> bool + Sync, diff --git a/crates/milli/src/vector/rest.rs b/crates/milli/src/vector/rest.rs index fc0ff308b..9aeb73f42 100644 --- a/crates/milli/src/vector/rest.rs +++ b/crates/milli/src/vector/rest.rs @@ -295,6 +295,10 @@ fn embed( where S: Serialize, { + use std::backtrace::Backtrace; + + println!("Embedder stats? {}", embedder_stats.is_some()); + let request = data.client.post(&data.url); let request = if let Some(bearer) = &data.bearer { request.set("Authorization", bearer) diff --git a/crates/milli/tests/search/distinct.rs b/crates/milli/tests/search/distinct.rs index 55e43c8fa..15fcf70a2 100644 --- a/crates/milli/tests/search/distinct.rs +++ b/crates/milli/tests/search/distinct.rs @@ -19,7 +19,7 @@ macro_rules! test_distinct { let config = milli::update::IndexerConfig::default(); let mut builder = Settings::new(&mut wtxn, &index, &config); builder.set_distinct_field(S(stringify!($distinct))); - builder.execute(|_| (), || false, None).unwrap(); + builder.execute(|_| (), || false, Default::default()).unwrap(); wtxn.commit().unwrap(); let rtxn = index.read_txn().unwrap(); diff --git a/crates/milli/tests/search/facet_distribution.rs b/crates/milli/tests/search/facet_distribution.rs index 588662735..5ed223400 100644 --- a/crates/milli/tests/search/facet_distribution.rs +++ b/crates/milli/tests/search/facet_distribution.rs @@ -25,7 +25,7 @@ fn test_facet_distribution_with_no_facet_values() { FilterableAttributesRule::Field(S("genres")), FilterableAttributesRule::Field(S("tags")), ]); - builder.execute(|_| (), || false, None).unwrap(); + builder.execute(|_| (), || false, Default::default()).unwrap(); wtxn.commit().unwrap(); // index documents @@ -74,6 +74,7 @@ fn test_facet_distribution_with_no_facet_values() { embedders, &|| false, &Progress::default(), + Default::default(), ) .unwrap(); diff --git a/crates/milli/tests/search/mod.rs b/crates/milli/tests/search/mod.rs index 1e0c24608..beee4ac54 100644 --- a/crates/milli/tests/search/mod.rs +++ b/crates/milli/tests/search/mod.rs @@ -63,7 +63,7 @@ pub fn setup_search_index_with_criteria(criteria: &[Criterion]) -> Index { S("america") => vec![S("the united states")], }); builder.set_searchable_fields(vec![S("title"), S("description")]); - builder.execute(|_| (), || false, None).unwrap(); + builder.execute(|_| (), || false, Default::default()).unwrap(); wtxn.commit().unwrap(); // index documents @@ -114,6 +114,7 @@ pub fn setup_search_index_with_criteria(criteria: &[Criterion]) -> Index { embedders, &|| false, &Progress::default(), + Default::default(), ) .unwrap(); diff --git a/crates/milli/tests/search/phrase_search.rs b/crates/milli/tests/search/phrase_search.rs index c5a95f7cd..180fcd176 100644 --- a/crates/milli/tests/search/phrase_search.rs +++ b/crates/milli/tests/search/phrase_search.rs @@ -10,7 +10,7 @@ fn set_stop_words(index: &Index, stop_words: &[&str]) { let mut builder = Settings::new(&mut wtxn, index, &config); let stop_words = stop_words.iter().map(|s| s.to_string()).collect(); builder.set_stop_words(stop_words); - builder.execute(|_| (), || false, None).unwrap(); + builder.execute(|_| (), || false, Default::default()).unwrap(); wtxn.commit().unwrap(); } diff --git a/crates/milli/tests/search/query_criteria.rs b/crates/milli/tests/search/query_criteria.rs index b7614c215..04b8374de 100644 --- a/crates/milli/tests/search/query_criteria.rs +++ b/crates/milli/tests/search/query_criteria.rs @@ -236,7 +236,7 @@ fn criteria_mixup() { let mut wtxn = index.write_txn().unwrap(); let mut builder = Settings::new(&mut wtxn, &index, &config); builder.set_criteria(criteria.clone()); - builder.execute(|_| (), || false, None).unwrap(); + builder.execute(|_| (), || false, Default::default()).unwrap(); wtxn.commit().unwrap(); let rtxn = index.read_txn().unwrap(); @@ -276,7 +276,7 @@ fn criteria_ascdesc() { S("name"), S("age"), }); - builder.execute(|_| (), || false, None).unwrap(); + builder.execute(|_| (), || false, Default::default()).unwrap(); wtxn.commit().unwrap(); let mut wtxn = index.write_txn().unwrap(); @@ -344,6 +344,7 @@ fn criteria_ascdesc() { embedders, &|| false, &Progress::default(), + Default::default(), ) .unwrap(); @@ -358,7 +359,7 @@ fn criteria_ascdesc() { let mut wtxn = index.write_txn().unwrap(); let mut builder = Settings::new(&mut wtxn, &index, &config); builder.set_criteria(vec![criterion.clone()]); - builder.execute(|_| (), || false, None).unwrap(); + builder.execute(|_| (), || false, Default::default()).unwrap(); wtxn.commit().unwrap(); let rtxn = index.read_txn().unwrap(); diff --git a/crates/milli/tests/search/typo_tolerance.rs b/crates/milli/tests/search/typo_tolerance.rs index bf9a730c9..e2cdab550 100644 --- a/crates/milli/tests/search/typo_tolerance.rs +++ b/crates/milli/tests/search/typo_tolerance.rs @@ -46,7 +46,7 @@ fn test_typo_tolerance_one_typo() { let config = IndexerConfig::default(); let mut builder = Settings::new(&mut txn, &index, &config); builder.set_min_word_len_one_typo(4); - builder.execute(|_| (), || false, None).unwrap(); + builder.execute(|_| (), || false, Default::default()).unwrap(); // typo is now supported for 4 letters words let mut search = Search::new(&txn, &index); @@ -92,7 +92,7 @@ fn test_typo_tolerance_two_typo() { let config = IndexerConfig::default(); let mut builder = Settings::new(&mut txn, &index, &config); builder.set_min_word_len_two_typos(7); - builder.execute(|_| (), || false, None).unwrap(); + builder.execute(|_| (), || false, Default::default()).unwrap(); // typo is now supported for 4 letters words let mut search = Search::new(&txn, &index); @@ -153,6 +153,7 @@ fn test_typo_disabled_on_word() { embedders, &|| false, &Progress::default(), + Default::default(), ) .unwrap(); @@ -180,7 +181,7 @@ fn test_typo_disabled_on_word() { // `zealand` doesn't allow typos anymore exact_words.insert("zealand".to_string()); builder.set_exact_words(exact_words); - builder.execute(|_| (), || false, None).unwrap(); + builder.execute(|_| (), || false, Default::default()).unwrap(); let mut search = Search::new(&txn, &index); search.query("zealand"); @@ -218,7 +219,7 @@ fn test_disable_typo_on_attribute() { let mut builder = Settings::new(&mut txn, &index, &config); // disable typos on `description` builder.set_exact_attributes(vec!["description".to_string()].into_iter().collect()); - builder.execute(|_| (), || false, None).unwrap(); + builder.execute(|_| (), || false, Default::default()).unwrap(); let mut search = Search::new(&txn, &index); search.query("antebelum"); From 2f82d945028fce52efaf704c2b0a4060093369fa Mon Sep 17 00:00:00 2001 From: Mubelotix Date: Mon, 23 Jun 2025 18:55:23 +0200 Subject: [PATCH 03/18] Fix the test and simplify types --- crates/dump/src/lib.rs | 2 +- crates/index-scheduler/src/queue/batches.rs | 4 +- .../src/scheduler/create_batch.rs | 7 +- .../src/scheduler/process_batch.rs | 13 +--- crates/index-scheduler/src/utils.rs | 26 +++---- crates/meilisearch-types/src/batch_view.rs | 4 +- crates/meilisearch-types/src/batches.rs | 10 ++- crates/meilisearch/tests/vector/rest.rs | 71 +++++++++++++++---- crates/milli/src/vector/rest.rs | 1 + 9 files changed, 87 insertions(+), 51 deletions(-) diff --git a/crates/dump/src/lib.rs b/crates/dump/src/lib.rs index b7a35ad5c..a84ec4ba5 100644 --- a/crates/dump/src/lib.rs +++ b/crates/dump/src/lib.rs @@ -329,7 +329,7 @@ pub(crate) mod test { write_channel_congestion: None, internal_database_sizes: Default::default(), }, - embedder_stats: None, + embedder_stats: Default::default(), enqueued_at: Some(BatchEnqueuedAt { earliest: datetime!(2022-11-11 0:00 UTC), oldest: datetime!(2022-11-11 0:00 UTC), diff --git a/crates/index-scheduler/src/queue/batches.rs b/crates/index-scheduler/src/queue/batches.rs index b14601733..c82d5acd2 100644 --- a/crates/index-scheduler/src/queue/batches.rs +++ b/crates/index-scheduler/src/queue/batches.rs @@ -174,7 +174,7 @@ impl BatchQueue { pub(crate) fn write_batch(&self, wtxn: &mut RwTxn, batch: ProcessingBatch) -> Result<()> { let old_batch = self.all_batches.get(wtxn, &batch.uid)?; - println!("Saving batch: {}", batch.embedder_stats.is_some()); + println!("Saving batch: {:?}", batch.embedder_stats); self.all_batches.put( wtxn, @@ -184,7 +184,7 @@ impl BatchQueue { progress: None, details: batch.details, stats: batch.stats, - embedder_stats: batch.embedder_stats.as_ref().map(|s| BatchEmbeddingStats::from(s.as_ref())), + embedder_stats: batch.embedder_stats.as_ref().into(), started_at: batch.started_at, finished_at: batch.finished_at, enqueued_at: batch.enqueued_at, diff --git a/crates/index-scheduler/src/scheduler/create_batch.rs b/crates/index-scheduler/src/scheduler/create_batch.rs index e3763881b..fc20b6fd5 100644 --- a/crates/index-scheduler/src/scheduler/create_batch.rs +++ b/crates/index-scheduler/src/scheduler/create_batch.rs @@ -437,8 +437,10 @@ impl IndexScheduler { #[cfg(test)] self.maybe_fail(crate::test_utils::FailureLocation::InsideCreateBatch)?; + println!("create next batch"); let batch_id = self.queue.batches.next_batch_id(rtxn)?; let mut current_batch = ProcessingBatch::new(batch_id); + println!("over"); let enqueued = &self.queue.tasks.get_status(rtxn, Status::Enqueued)?; let count_total_enqueued = enqueued.len(); @@ -454,6 +456,7 @@ impl IndexScheduler { kind: Kind::TaskCancelation, id: task_id, }); + println!("task cancelled"); return Ok(Some((Batch::TaskCancelation { task }, current_batch))); } @@ -524,7 +527,7 @@ impl IndexScheduler { } // 5. We make a batch from the unprioritised tasks. Start by taking the next enqueued task. - let task_id = if let Some(task_id) = enqueued.min() { task_id } else { return Ok(None) }; + let task_id = if let Some(task_id) = enqueued.min() { task_id } else { println!("return"); return Ok(None) }; let mut task = self.queue.tasks.get_task(rtxn, task_id)?.ok_or(Error::CorruptedTaskQueue)?; @@ -602,6 +605,7 @@ impl IndexScheduler { autobatcher::autobatch(enqueued, index_already_exists, primary_key.as_deref()) { current_batch.reason(autobatch_stop_reason.unwrap_or(stop_reason)); + println!("autobatch"); return Ok(self .create_next_batch_index( rtxn, @@ -615,6 +619,7 @@ impl IndexScheduler { // If we found no tasks then we were notified for something that got autobatched // somehow and there is nothing to do. + println!("nothing to do"); Ok(None) } } diff --git a/crates/index-scheduler/src/scheduler/process_batch.rs b/crates/index-scheduler/src/scheduler/process_batch.rs index 4e36b65b6..c5305cf21 100644 --- a/crates/index-scheduler/src/scheduler/process_batch.rs +++ b/crates/index-scheduler/src/scheduler/process_batch.rs @@ -164,7 +164,7 @@ impl IndexScheduler { let pre_commit_dabases_sizes = index.database_sizes(&index_wtxn)?; let (tasks, congestion) = - self.apply_index_operation(&mut index_wtxn, &index, op, &progress, current_batch.clone_embedder_stats())?; + self.apply_index_operation(&mut index_wtxn, &index, op, &progress, current_batch.embedder_stats.clone())?; { progress.update_progress(FinalizingIndexStep::Committing); @@ -240,20 +240,11 @@ impl IndexScheduler { builder.set_primary_key(primary_key); let must_stop_processing = self.scheduler.must_stop_processing.clone(); - let embedder_stats = match current_batch.embedder_stats { - Some(ref stats) => stats.clone(), - None => { - let embedder_stats: Arc = Default::default(); - current_batch.embedder_stats = Some(embedder_stats.clone()); - embedder_stats - }, - }; - builder .execute( |indexing_step| tracing::debug!(update = ?indexing_step), || must_stop_processing.get(), - embedder_stats, + current_batch.embedder_stats.clone(), ) .map_err(|e| Error::from_milli(e, Some(index_uid.to_string())))?; index_wtxn.commit()?; diff --git a/crates/index-scheduler/src/utils.rs b/crates/index-scheduler/src/utils.rs index 22e319580..455b6a2e7 100644 --- a/crates/index-scheduler/src/utils.rs +++ b/crates/index-scheduler/src/utils.rs @@ -29,7 +29,7 @@ pub struct ProcessingBatch { pub uid: BatchId, pub details: DetailsView, pub stats: BatchStats, - pub embedder_stats: Option>, + pub embedder_stats: Arc, pub statuses: HashSet, pub kinds: HashSet, @@ -47,11 +47,13 @@ impl ProcessingBatch { let mut statuses = HashSet::default(); statuses.insert(Status::Processing); + println!("Processing batch created: {}", uid); + Self { uid, details: DetailsView::default(), stats: BatchStats::default(), - embedder_stats: None, + embedder_stats: Default::default(), statuses, kinds: HashSet::default(), @@ -64,17 +66,6 @@ impl ProcessingBatch { } } - pub fn clone_embedder_stats(&mut self) -> Arc { - match self.embedder_stats { - Some(ref stats) => stats.clone(), - None => { - let embedder_stats: Arc = Default::default(); - self.embedder_stats = Some(embedder_stats.clone()); - embedder_stats - }, - } - } - /// Update itself with the content of the task and update the batch id in the task. pub fn processing<'a>(&mut self, tasks: impl IntoIterator) { for task in tasks.into_iter() { @@ -113,11 +104,14 @@ impl ProcessingBatch { } pub fn reason(&mut self, reason: BatchStopReason) { + println!("batch stopped: {:?}", reason); self.reason = reason; } /// Must be called once the batch has finished processing. pub fn finished(&mut self) { + println!("Batch finished: {}", self.uid); + self.details = DetailsView::default(); self.stats = BatchStats::default(); self.finished_at = Some(OffsetDateTime::now_utc()); @@ -132,6 +126,8 @@ impl ProcessingBatch { /// Update the timestamp of the tasks and the inner structure of this structure. pub fn update(&mut self, task: &mut Task) { + println!("Updating task: {} in batch: {}", task.uid, self.uid); + // We must re-set this value in case we're dealing with a task that has been added between // the `processing` and `finished` state // We must re-set this value in case we're dealing with a task that has been added between @@ -156,13 +152,13 @@ impl ProcessingBatch { } pub fn to_batch(&self) -> Batch { - println!("Converting to batch: {:?}", self.embedder_stats); + println!("Converting to batch: {:?} {:?}", self.uid, self.embedder_stats); Batch { uid: self.uid, progress: None, details: self.details.clone(), stats: self.stats.clone(), - embedder_stats: self.embedder_stats.as_ref().map(|s| BatchEmbeddingStats::from(s.as_ref())), + embedder_stats: self.embedder_stats.as_ref().into(), started_at: self.started_at, finished_at: self.finished_at, enqueued_at: self.enqueued_at, diff --git a/crates/meilisearch-types/src/batch_view.rs b/crates/meilisearch-types/src/batch_view.rs index 0a9b80f4e..bd56f5b1a 100644 --- a/crates/meilisearch-types/src/batch_view.rs +++ b/crates/meilisearch-types/src/batch_view.rs @@ -31,8 +31,8 @@ pub struct BatchView { pub struct BatchStatsView { #[serde(flatten)] pub stats: BatchStats, - #[serde(skip_serializing_if = "BatchEmbeddingStats::skip_serializing")] - pub embedder: Option, + #[serde(skip_serializing_if = "BatchEmbeddingStats::skip_serializing", default)] + pub embedder: BatchEmbeddingStats, } impl BatchView { diff --git a/crates/meilisearch-types/src/batches.rs b/crates/meilisearch-types/src/batches.rs index 24be75d1c..e1c9411b6 100644 --- a/crates/meilisearch-types/src/batches.rs +++ b/crates/meilisearch-types/src/batches.rs @@ -20,7 +20,8 @@ pub struct Batch { pub progress: Option, pub details: DetailsView, pub stats: BatchStats, - pub embedder_stats: Option, + #[serde(skip_serializing_if = "BatchEmbeddingStats::skip_serializing", default)] + pub embedder_stats: BatchEmbeddingStats, #[serde(with = "time::serde::rfc3339")] pub started_at: OffsetDateTime, @@ -110,10 +111,7 @@ impl From<&EmbedderStats> for BatchEmbeddingStats { } impl BatchEmbeddingStats { - pub fn skip_serializing(this: &Option) -> bool { - match this { - Some(stats) => stats.total_count == 0 && stats.error_count == 0 && stats.last_error.is_none(), - None => true, - } + pub fn skip_serializing(&self) -> bool { + self.total_count == 0 && self.error_count == 0 && self.last_error.is_none() } } diff --git a/crates/meilisearch/tests/vector/rest.rs b/crates/meilisearch/tests/vector/rest.rs index 1ff2dd9fe..156a2f07b 100644 --- a/crates/meilisearch/tests/vector/rest.rs +++ b/crates/meilisearch/tests/vector/rest.rs @@ -1,10 +1,12 @@ -use std::collections::BTreeMap; +use std::collections::{BTreeMap, BTreeSet}; use meili_snap::{json_string, snapshot}; use reqwest::IntoUrl; +use tokio::spawn; +use tokio::sync::mpsc; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, Request, ResponseTemplate}; -use std::thread::sleep; +use tokio::time::sleep; use std::time::Duration; use crate::common::Value; @@ -307,7 +309,6 @@ async fn create_mock_raw() -> (MockServer, Value) { Mock::given(method("POST")) .and(path("/")) .respond_with(move |req: &Request| { - println!("Sent!"); let req: String = match req.body_json() { Ok(req) => req, Err(error) => { @@ -337,6 +338,50 @@ async fn create_mock_raw() -> (MockServer, Value) { (mock_server, embedder_settings) } +/// A mock server that returns 500 errors, and sends a message once 5 requests are received +async fn create_faulty_mock_raw(mut sender: mpsc::Sender<()>) -> (MockServer, Value) { + let mock_server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/")) + .respond_with(move |req: &Request| { + let req: String = match req.body_json() { + Ok(req) => req, + Err(error) => { + return ResponseTemplate::new(400).set_body_json(json!({ + "error": format!("Invalid request: {error}") + })); + } + }; + + let sender = sender.clone(); + spawn(async move { + sender.send(()).await; + }); + + ResponseTemplate::new(500) + .set_delay(Duration::from_millis(500)) + .set_body_json(json!({ + "error": "Service Unavailable", + "text": req + })) + }) + .mount(&mock_server) + .await; + let url = mock_server.uri(); + + let embedder_settings = json!({ + "source": "rest", + "url": url, + "dimensions": 3, + "request": "{{text}}", + "response": "{{embedding}}", + "documentTemplate": "{{doc.name}}" + }); + + (mock_server, embedder_settings) +} + pub async fn post(url: T, text: &str) -> reqwest::Result { reqwest::Client::builder().build()?.post(url).json(&json!(text)).send().await } @@ -2118,7 +2163,8 @@ async fn searchable_reindex() { #[actix_rt::test] async fn observability() { - let (_mock, setting) = create_mock_raw().await; + let (sender, mut receiver) = mpsc::channel(10); + let (_mock, setting) = create_faulty_mock_raw(sender).await; let server = get_server_vector().await; let index = server.index("doggo"); @@ -2133,20 +2179,19 @@ async fn observability() { let task = server.wait_task(response.uid()).await; snapshot!(task["status"], @r###""succeeded""###); let documents = json!([ - {"id": 0, "name": "kefir"}, - {"id": 1, "name": "echo", "_vectors": { "rest": [1, 1, 1] }}, - {"id": 2, "name": "intel"}, - {"id": 3, "name": "missing"}, // Stuff that doesn't exist - {"id": 4, "name": "invalid"}, - {"id": 5, "name": "foobar"}, + {"id": 0, "name": "will_return_500"}, // Stuff that doesn't exist + {"id": 1, "name": "will_error"}, + {"id": 2, "name": "must_error"}, ]); let (value, code) = index.add_documents(documents, None).await; snapshot!(code, @"202 Accepted"); - let batches = index.filtered_batches(&[], &[], &[]).await; - println!("Batches: {batches:?}"); + // The task will eventually fail, so let's not wait for it. + // Let's just wait for 5 errors from the mock server. + for _errors in 0..5 { + receiver.recv().await; + } - let task = index.wait_task(value.uid()).await; let batches = index.filtered_batches(&[], &[], &[]).await; println!("Batches: {batches:?}"); diff --git a/crates/milli/src/vector/rest.rs b/crates/milli/src/vector/rest.rs index 9aeb73f42..706a411fb 100644 --- a/crates/milli/src/vector/rest.rs +++ b/crates/milli/src/vector/rest.rs @@ -316,6 +316,7 @@ where if let Some(embedder_stats) = &embedder_stats { embedder_stats.as_ref().total_requests.fetch_add(1, std::sync::atomic::Ordering::Relaxed); } + // TODO: also catch 403 errors let response = request.clone().send_json(&body); let result = check_response(response, data.configuration_source).and_then(|response| { response_to_embedding(response, data, expected_count, expected_dimension) From 59a1c5d9a7f25fc4b7113eb99c8fd11e7a19ce95 Mon Sep 17 00:00:00 2001 From: Mubelotix Date: Tue, 24 Jun 2025 11:08:06 +0200 Subject: [PATCH 04/18] Make test more reproducible --- crates/meilisearch/tests/vector/rest.rs | 45 ++++++++++++++----------- 1 file changed, 25 insertions(+), 20 deletions(-) diff --git a/crates/meilisearch/tests/vector/rest.rs b/crates/meilisearch/tests/vector/rest.rs index 156a2f07b..54ed52213 100644 --- a/crates/meilisearch/tests/vector/rest.rs +++ b/crates/meilisearch/tests/vector/rest.rs @@ -1,4 +1,5 @@ use std::collections::{BTreeMap, BTreeSet}; +use std::sync::atomic::AtomicUsize; use meili_snap::{json_string, snapshot}; use reqwest::IntoUrl; @@ -338,36 +339,43 @@ async fn create_mock_raw() -> (MockServer, Value) { (mock_server, embedder_settings) } -/// A mock server that returns 500 errors, and sends a message once 5 requests are received -async fn create_faulty_mock_raw(mut sender: mpsc::Sender<()>) -> (MockServer, Value) { +async fn create_faulty_mock_raw(sender: mpsc::Sender<()>) -> (MockServer, Value) { let mock_server = MockServer::start().await; - + let count = AtomicUsize::new(0); + Mock::given(method("POST")) .and(path("/")) .respond_with(move |req: &Request| { - let req: String = match req.body_json() { - Ok(req) => req, + let count = count.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + + let req_body = match req.body_json::() { + Ok(body) => body, Err(error) => { return ResponseTemplate::new(400).set_body_json(json!({ - "error": format!("Invalid request: {error}") + "error": format!("Invalid request: {error}") })); } }; - let sender = sender.clone(); - spawn(async move { - sender.send(()).await; - }); + if count >= 5 { + let _ = sender.try_send(()); + ResponseTemplate::new(500) + .set_delay(Duration::from_secs(u64::MAX)) + .set_body_json(json!({ + "error": "Service Unavailable", + "text": req_body + })) + } else { - ResponseTemplate::new(500) - .set_delay(Duration::from_millis(500)) - .set_body_json(json!({ + ResponseTemplate::new(500).set_body_json(json!({ "error": "Service Unavailable", - "text": req + "text": req_body })) + } }) .mount(&mock_server) .await; + let url = mock_server.uri(); let embedder_settings = json!({ @@ -2187,12 +2195,9 @@ async fn observability() { snapshot!(code, @"202 Accepted"); // The task will eventually fail, so let's not wait for it. - // Let's just wait for 5 errors from the mock server. - for _errors in 0..5 { - receiver.recv().await; - } + // Let's just wait for the server to block + receiver.recv().await; let batches = index.filtered_batches(&[], &[], &[]).await; - println!("Batches: {batches:?}"); - + snapshot!(task, @r###""###); } From 4a179fb3c064aed5320e6a0e0011447e1c17b125 Mon Sep 17 00:00:00 2001 From: Mubelotix Date: Tue, 24 Jun 2025 11:38:11 +0200 Subject: [PATCH 05/18] Improve code quality --- crates/index-scheduler/src/insta_snapshot.rs | 2 +- crates/index-scheduler/src/queue/batches.rs | 9 ++------- .../index-scheduler/src/scheduler/create_batch.rs | 7 +------ crates/index-scheduler/src/utils.rs | 10 +--------- crates/meilisearch-types/src/batch_view.rs | 6 +++--- crates/meilisearch-types/src/batches.rs | 14 +++++++------- crates/meilisearch/tests/vector/rest.rs | 10 +++++----- crates/milli/src/progress.rs | 6 +++--- crates/milli/src/vector/rest.rs | 9 ++------- 9 files changed, 25 insertions(+), 48 deletions(-) diff --git a/crates/index-scheduler/src/insta_snapshot.rs b/crates/index-scheduler/src/insta_snapshot.rs index 06ec01b5e..8e1fb1c2c 100644 --- a/crates/index-scheduler/src/insta_snapshot.rs +++ b/crates/index-scheduler/src/insta_snapshot.rs @@ -1,7 +1,7 @@ use std::collections::BTreeSet; use std::fmt::Write; -use meilisearch_types::batches::{Batch, BatchEmbeddingStats, BatchEnqueuedAt, BatchStats}; +use meilisearch_types::batches::{Batch, EmbedderStatsView, BatchEnqueuedAt, BatchStats}; use meilisearch_types::heed::types::{SerdeBincode, SerdeJson, Str}; use meilisearch_types::heed::{Database, RoTxn}; use meilisearch_types::milli::{CboRoaringBitmapCodec, RoaringBitmapCodec, BEU32}; diff --git a/crates/index-scheduler/src/queue/batches.rs b/crates/index-scheduler/src/queue/batches.rs index c82d5acd2..96a3940a5 100644 --- a/crates/index-scheduler/src/queue/batches.rs +++ b/crates/index-scheduler/src/queue/batches.rs @@ -1,7 +1,7 @@ use std::collections::HashSet; use std::ops::{Bound, RangeBounds}; -use meilisearch_types::batches::{Batch, BatchEmbeddingStats, BatchId}; +use meilisearch_types::batches::{Batch, EmbedderStatsView, BatchId}; use meilisearch_types::heed::types::{DecodeIgnore, SerdeBincode, SerdeJson, Str}; use meilisearch_types::heed::{Database, Env, RoTxn, RwTxn, WithoutTls}; use meilisearch_types::milli::{CboRoaringBitmapCodec, RoaringBitmapCodec, BEU32}; @@ -92,10 +92,7 @@ impl BatchQueue { } pub(crate) fn get_batch(&self, rtxn: &RoTxn, batch_id: BatchId) -> Result> { - println!("Got batch from db {batch_id:?}"); - let r = Ok(self.all_batches.get(rtxn, &batch_id)?); - println!("Got batch from db => {:?}", r); - r + Ok(self.all_batches.get(rtxn, &batch_id)?) } /// Returns the whole set of batches that belongs to this index. @@ -174,8 +171,6 @@ impl BatchQueue { pub(crate) fn write_batch(&self, wtxn: &mut RwTxn, batch: ProcessingBatch) -> Result<()> { let old_batch = self.all_batches.get(wtxn, &batch.uid)?; - println!("Saving batch: {:?}", batch.embedder_stats); - self.all_batches.put( wtxn, &batch.uid, diff --git a/crates/index-scheduler/src/scheduler/create_batch.rs b/crates/index-scheduler/src/scheduler/create_batch.rs index fc20b6fd5..e3763881b 100644 --- a/crates/index-scheduler/src/scheduler/create_batch.rs +++ b/crates/index-scheduler/src/scheduler/create_batch.rs @@ -437,10 +437,8 @@ impl IndexScheduler { #[cfg(test)] self.maybe_fail(crate::test_utils::FailureLocation::InsideCreateBatch)?; - println!("create next batch"); let batch_id = self.queue.batches.next_batch_id(rtxn)?; let mut current_batch = ProcessingBatch::new(batch_id); - println!("over"); let enqueued = &self.queue.tasks.get_status(rtxn, Status::Enqueued)?; let count_total_enqueued = enqueued.len(); @@ -456,7 +454,6 @@ impl IndexScheduler { kind: Kind::TaskCancelation, id: task_id, }); - println!("task cancelled"); return Ok(Some((Batch::TaskCancelation { task }, current_batch))); } @@ -527,7 +524,7 @@ impl IndexScheduler { } // 5. We make a batch from the unprioritised tasks. Start by taking the next enqueued task. - let task_id = if let Some(task_id) = enqueued.min() { task_id } else { println!("return"); return Ok(None) }; + let task_id = if let Some(task_id) = enqueued.min() { task_id } else { return Ok(None) }; let mut task = self.queue.tasks.get_task(rtxn, task_id)?.ok_or(Error::CorruptedTaskQueue)?; @@ -605,7 +602,6 @@ impl IndexScheduler { autobatcher::autobatch(enqueued, index_already_exists, primary_key.as_deref()) { current_batch.reason(autobatch_stop_reason.unwrap_or(stop_reason)); - println!("autobatch"); return Ok(self .create_next_batch_index( rtxn, @@ -619,7 +615,6 @@ impl IndexScheduler { // If we found no tasks then we were notified for something that got autobatched // somehow and there is nothing to do. - println!("nothing to do"); Ok(None) } } diff --git a/crates/index-scheduler/src/utils.rs b/crates/index-scheduler/src/utils.rs index 455b6a2e7..226ef9f06 100644 --- a/crates/index-scheduler/src/utils.rs +++ b/crates/index-scheduler/src/utils.rs @@ -5,7 +5,7 @@ use std::ops::Bound; use std::sync::Arc; use crate::milli::progress::EmbedderStats; -use meilisearch_types::batches::{Batch, BatchEmbeddingStats, BatchEnqueuedAt, BatchId, BatchStats}; +use meilisearch_types::batches::{Batch, EmbedderStatsView, BatchEnqueuedAt, BatchId, BatchStats}; use meilisearch_types::heed::{Database, RoTxn, RwTxn}; use meilisearch_types::milli::CboRoaringBitmapCodec; use meilisearch_types::task_view::DetailsView; @@ -47,8 +47,6 @@ impl ProcessingBatch { let mut statuses = HashSet::default(); statuses.insert(Status::Processing); - println!("Processing batch created: {}", uid); - Self { uid, details: DetailsView::default(), @@ -104,14 +102,11 @@ impl ProcessingBatch { } pub fn reason(&mut self, reason: BatchStopReason) { - println!("batch stopped: {:?}", reason); self.reason = reason; } /// Must be called once the batch has finished processing. pub fn finished(&mut self) { - println!("Batch finished: {}", self.uid); - self.details = DetailsView::default(); self.stats = BatchStats::default(); self.finished_at = Some(OffsetDateTime::now_utc()); @@ -126,8 +121,6 @@ impl ProcessingBatch { /// Update the timestamp of the tasks and the inner structure of this structure. pub fn update(&mut self, task: &mut Task) { - println!("Updating task: {} in batch: {}", task.uid, self.uid); - // We must re-set this value in case we're dealing with a task that has been added between // the `processing` and `finished` state // We must re-set this value in case we're dealing with a task that has been added between @@ -152,7 +145,6 @@ impl ProcessingBatch { } pub fn to_batch(&self) -> Batch { - println!("Converting to batch: {:?} {:?}", self.uid, self.embedder_stats); Batch { uid: self.uid, progress: None, diff --git a/crates/meilisearch-types/src/batch_view.rs b/crates/meilisearch-types/src/batch_view.rs index bd56f5b1a..aced97d7a 100644 --- a/crates/meilisearch-types/src/batch_view.rs +++ b/crates/meilisearch-types/src/batch_view.rs @@ -3,7 +3,7 @@ use serde::Serialize; use time::{Duration, OffsetDateTime}; use utoipa::ToSchema; -use crate::batches::{Batch, BatchEmbeddingStats, BatchId, BatchStats}; +use crate::batches::{Batch, EmbedderStatsView, BatchId, BatchStats}; use crate::task_view::DetailsView; use crate::tasks::serialize_duration; @@ -31,8 +31,8 @@ pub struct BatchView { pub struct BatchStatsView { #[serde(flatten)] pub stats: BatchStats, - #[serde(skip_serializing_if = "BatchEmbeddingStats::skip_serializing", default)] - pub embedder: BatchEmbeddingStats, + #[serde(skip_serializing_if = "EmbedderStatsView::skip_serializing", default)] + pub embedder: EmbedderStatsView, } impl BatchView { diff --git a/crates/meilisearch-types/src/batches.rs b/crates/meilisearch-types/src/batches.rs index e1c9411b6..45cc2d9f4 100644 --- a/crates/meilisearch-types/src/batches.rs +++ b/crates/meilisearch-types/src/batches.rs @@ -20,8 +20,8 @@ pub struct Batch { pub progress: Option, pub details: DetailsView, pub stats: BatchStats, - #[serde(skip_serializing_if = "BatchEmbeddingStats::skip_serializing", default)] - pub embedder_stats: BatchEmbeddingStats, + #[serde(skip_serializing_if = "EmbedderStatsView::skip_serializing", default)] + pub embedder_stats: EmbedderStatsView, #[serde(with = "time::serde::rfc3339")] pub started_at: OffsetDateTime, @@ -92,25 +92,25 @@ pub struct BatchStats { #[derive(Default, Debug, Clone, PartialEq, Eq, Serialize, Deserialize, ToSchema)] #[serde(rename_all = "camelCase")] #[schema(rename_all = "camelCase")] -pub struct BatchEmbeddingStats { +pub struct EmbedderStatsView { pub total_count: usize, pub error_count: usize, - #[serde(skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none", default)] pub last_error: Option, } -impl From<&EmbedderStats> for BatchEmbeddingStats { +impl From<&EmbedderStats> for EmbedderStatsView { fn from(stats: &EmbedderStats) -> Self { let errors = stats.errors.read().unwrap(); Self { - total_count: stats.total_requests.load(std::sync::atomic::Ordering::Relaxed), + total_count: stats.total_count.load(std::sync::atomic::Ordering::Relaxed), error_count: errors.1 as usize, last_error: errors.0.clone(), } } } -impl BatchEmbeddingStats { +impl EmbedderStatsView { pub fn skip_serializing(&self) -> bool { self.total_count == 0 && self.error_count == 0 && self.last_error.is_none() } diff --git a/crates/meilisearch/tests/vector/rest.rs b/crates/meilisearch/tests/vector/rest.rs index 54ed52213..1fdd18d28 100644 --- a/crates/meilisearch/tests/vector/rest.rs +++ b/crates/meilisearch/tests/vector/rest.rs @@ -2170,7 +2170,7 @@ async fn searchable_reindex() { #[actix_rt::test] -async fn observability() { +async fn last_error_stats() { let (sender, mut receiver) = mpsc::channel(10); let (_mock, setting) = create_faulty_mock_raw(sender).await; let server = get_server_vector().await; @@ -2187,7 +2187,7 @@ async fn observability() { let task = server.wait_task(response.uid()).await; snapshot!(task["status"], @r###""succeeded""###); let documents = json!([ - {"id": 0, "name": "will_return_500"}, // Stuff that doesn't exist + {"id": 0, "name": "will_return_500"}, {"id": 1, "name": "will_error"}, {"id": 2, "name": "must_error"}, ]); @@ -2195,9 +2195,9 @@ async fn observability() { snapshot!(code, @"202 Accepted"); // The task will eventually fail, so let's not wait for it. - // Let's just wait for the server to block + // Let's just wait for the server's signal receiver.recv().await; - let batches = index.filtered_batches(&[], &[], &[]).await; - snapshot!(task, @r###""###); + let (response, _code) = index.filtered_batches(&[], &[], &[]).await; + snapshot!(response["results"][0], @r###""###); } diff --git a/crates/milli/src/progress.rs b/crates/milli/src/progress.rs index 7026f0c11..8cd2c9336 100644 --- a/crates/milli/src/progress.rs +++ b/crates/milli/src/progress.rs @@ -25,15 +25,15 @@ pub struct Progress { #[derive(Default)] pub struct EmbedderStats { pub errors: Arc, u32)>>, - pub total_requests: AtomicUsize + pub total_count: AtomicUsize } impl std::fmt::Debug for EmbedderStats { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { let (error, count) = self.errors.read().unwrap().clone(); f.debug_struct("EmbedderStats") - .field("errors", &error) - .field("total_requests", &self.total_requests.load(Ordering::Relaxed)) + .field("last_error", &error) + .field("total_count", &self.total_count.load(Ordering::Relaxed)) .field("error_count", &count) .finish() } diff --git a/crates/milli/src/vector/rest.rs b/crates/milli/src/vector/rest.rs index 706a411fb..d8de89c6a 100644 --- a/crates/milli/src/vector/rest.rs +++ b/crates/milli/src/vector/rest.rs @@ -295,10 +295,6 @@ fn embed( where S: Serialize, { - use std::backtrace::Backtrace; - - println!("Embedder stats? {}", embedder_stats.is_some()); - let request = data.client.post(&data.url); let request = if let Some(bearer) = &data.bearer { request.set("Authorization", bearer) @@ -314,9 +310,8 @@ where for attempt in 0..10 { if let Some(embedder_stats) = &embedder_stats { - embedder_stats.as_ref().total_requests.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + embedder_stats.as_ref().total_count.fetch_add(1, std::sync::atomic::Ordering::Relaxed); } - // TODO: also catch 403 errors let response = request.clone().send_json(&body); let result = check_response(response, data.configuration_source).and_then(|response| { response_to_embedding(response, data, expected_count, expected_dimension) @@ -358,7 +353,7 @@ where } if let Some(embedder_stats) = &embedder_stats { - embedder_stats.as_ref().total_requests.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + embedder_stats.as_ref().total_count.fetch_add(1, std::sync::atomic::Ordering::Relaxed); } let response = request.send_json(&body); let result = check_response(response, data.configuration_source).and_then(|response| { From d7721fe607f4645bc805423d06fddf1f39f63f8d Mon Sep 17 00:00:00 2001 From: Mubelotix Date: Tue, 24 Jun 2025 12:20:22 +0200 Subject: [PATCH 06/18] Format --- crates/index-scheduler/src/insta_snapshot.rs | 7 ++-- crates/index-scheduler/src/queue/batches.rs | 2 +- .../src/scheduler/process_batch.rs | 12 ++++--- crates/index-scheduler/src/utils.rs | 4 +-- crates/meilisearch-types/src/batch_view.rs | 2 +- crates/meilisearch-types/src/batches.rs | 1 - crates/meilisearch/src/lib.rs | 7 ++-- crates/meilisearch/tests/vector/rest.rs | 20 +++++------ crates/milli/src/progress.rs | 2 +- .../milli/src/search/new/tests/integration.rs | 2 +- .../milli/src/update/new/indexer/extract.rs | 2 +- crates/milli/src/update/new/indexer/mod.rs | 2 +- crates/milli/src/update/settings.rs | 7 +++- crates/milli/src/vector/composite.rs | 34 +++++++++++++------ crates/milli/src/vector/mod.rs | 30 ++++++++++------ crates/milli/src/vector/ollama.rs | 14 +++++--- crates/milli/src/vector/openai.rs | 16 +++++++-- crates/milli/src/vector/rest.rs | 23 ++++++++++--- 18 files changed, 124 insertions(+), 63 deletions(-) diff --git a/crates/index-scheduler/src/insta_snapshot.rs b/crates/index-scheduler/src/insta_snapshot.rs index 8e1fb1c2c..d4504c246 100644 --- a/crates/index-scheduler/src/insta_snapshot.rs +++ b/crates/index-scheduler/src/insta_snapshot.rs @@ -1,7 +1,7 @@ use std::collections::BTreeSet; use std::fmt::Write; -use meilisearch_types::batches::{Batch, EmbedderStatsView, BatchEnqueuedAt, BatchStats}; +use meilisearch_types::batches::{Batch, BatchEnqueuedAt, BatchStats}; use meilisearch_types::heed::types::{SerdeBincode, SerdeJson, Str}; use meilisearch_types::heed::{Database, RoTxn}; use meilisearch_types::milli::{CboRoaringBitmapCodec, RoaringBitmapCodec, BEU32}; @@ -367,7 +367,10 @@ pub fn snapshot_batch(batch: &Batch) -> String { snap.push_str(&format!("uid: {uid}, ")); snap.push_str(&format!("details: {}, ", serde_json::to_string(details).unwrap())); snap.push_str(&format!("stats: {}, ", serde_json::to_string(&stats).unwrap())); - snap.push_str(&format!("embedder_stats: {}, ", serde_json::to_string(&embedder_stats).unwrap())); + snap.push_str(&format!( + "embedder_stats: {}, ", + serde_json::to_string(&embedder_stats).unwrap() + )); snap.push_str(&format!("stop reason: {}, ", serde_json::to_string(&stop_reason).unwrap())); snap.push('}'); snap diff --git a/crates/index-scheduler/src/queue/batches.rs b/crates/index-scheduler/src/queue/batches.rs index 96a3940a5..b96f65836 100644 --- a/crates/index-scheduler/src/queue/batches.rs +++ b/crates/index-scheduler/src/queue/batches.rs @@ -1,7 +1,7 @@ use std::collections::HashSet; use std::ops::{Bound, RangeBounds}; -use meilisearch_types::batches::{Batch, EmbedderStatsView, BatchId}; +use meilisearch_types::batches::{Batch, BatchId}; use meilisearch_types::heed::types::{DecodeIgnore, SerdeBincode, SerdeJson, Str}; use meilisearch_types::heed::{Database, Env, RoTxn, RwTxn, WithoutTls}; use meilisearch_types::milli::{CboRoaringBitmapCodec, RoaringBitmapCodec, BEU32}; diff --git a/crates/index-scheduler/src/scheduler/process_batch.rs b/crates/index-scheduler/src/scheduler/process_batch.rs index c5305cf21..5261692b6 100644 --- a/crates/index-scheduler/src/scheduler/process_batch.rs +++ b/crates/index-scheduler/src/scheduler/process_batch.rs @@ -1,11 +1,10 @@ use std::collections::{BTreeSet, HashMap, HashSet}; use std::panic::{catch_unwind, AssertUnwindSafe}; use std::sync::atomic::Ordering; -use std::sync::Arc; use meilisearch_types::batches::{BatchEnqueuedAt, BatchId}; use meilisearch_types::heed::{RoTxn, RwTxn}; -use meilisearch_types::milli::progress::{EmbedderStats, Progress, VariableNameStep}; +use meilisearch_types::milli::progress::{Progress, VariableNameStep}; use meilisearch_types::milli::{self, ChannelCongestion}; use meilisearch_types::tasks::{Details, IndexSwap, Kind, KindWithContent, Status, Task}; use meilisearch_types::versioning::{VERSION_MAJOR, VERSION_MINOR, VERSION_PATCH}; @@ -163,8 +162,13 @@ impl IndexScheduler { .set_currently_updating_index(Some((index_uid.clone(), index.clone()))); let pre_commit_dabases_sizes = index.database_sizes(&index_wtxn)?; - let (tasks, congestion) = - self.apply_index_operation(&mut index_wtxn, &index, op, &progress, current_batch.embedder_stats.clone())?; + let (tasks, congestion) = self.apply_index_operation( + &mut index_wtxn, + &index, + op, + &progress, + current_batch.embedder_stats.clone(), + )?; { progress.update_progress(FinalizingIndexStep::Committing); diff --git a/crates/index-scheduler/src/utils.rs b/crates/index-scheduler/src/utils.rs index 226ef9f06..ca37065ec 100644 --- a/crates/index-scheduler/src/utils.rs +++ b/crates/index-scheduler/src/utils.rs @@ -1,11 +1,11 @@ //! Utility functions on the DBs. Mainly getter and setters. +use crate::milli::progress::EmbedderStats; use std::collections::{BTreeSet, HashSet}; use std::ops::Bound; use std::sync::Arc; -use crate::milli::progress::EmbedderStats; -use meilisearch_types::batches::{Batch, EmbedderStatsView, BatchEnqueuedAt, BatchId, BatchStats}; +use meilisearch_types::batches::{Batch, BatchEnqueuedAt, BatchId, BatchStats}; use meilisearch_types::heed::{Database, RoTxn, RwTxn}; use meilisearch_types::milli::CboRoaringBitmapCodec; use meilisearch_types::task_view::DetailsView; diff --git a/crates/meilisearch-types/src/batch_view.rs b/crates/meilisearch-types/src/batch_view.rs index aced97d7a..ea027b74e 100644 --- a/crates/meilisearch-types/src/batch_view.rs +++ b/crates/meilisearch-types/src/batch_view.rs @@ -3,7 +3,7 @@ use serde::Serialize; use time::{Duration, OffsetDateTime}; use utoipa::ToSchema; -use crate::batches::{Batch, EmbedderStatsView, BatchId, BatchStats}; +use crate::batches::{Batch, BatchId, BatchStats, EmbedderStatsView}; use crate::task_view::DetailsView; use crate::tasks::serialize_duration; diff --git a/crates/meilisearch-types/src/batches.rs b/crates/meilisearch-types/src/batches.rs index 45cc2d9f4..cec74fb75 100644 --- a/crates/meilisearch-types/src/batches.rs +++ b/crates/meilisearch-types/src/batches.rs @@ -1,5 +1,4 @@ use std::collections::BTreeMap; -use std::sync::Arc; use milli::progress::{EmbedderStats, ProgressView}; use serde::{Deserialize, Serialize}; diff --git a/crates/meilisearch/src/lib.rs b/crates/meilisearch/src/lib.rs index 72be6aec9..cdecd520c 100644 --- a/crates/meilisearch/src/lib.rs +++ b/crates/meilisearch/src/lib.rs @@ -544,8 +544,11 @@ fn import_dump( let settings = index_reader.settings()?; apply_settings_to_builder(&settings, &mut builder); let embedder_stats: Arc = Default::default(); // FIXME: this isn't linked to anything - builder - .execute(|indexing_step| tracing::debug!("update: {:?}", indexing_step), || false, embedder_stats.clone())?; + builder.execute( + |indexing_step| tracing::debug!("update: {:?}", indexing_step), + || false, + embedder_stats.clone(), + )?; // 4.3 Import the documents. // 4.3.1 We need to recreate the grenad+obkv format accepted by the index. diff --git a/crates/meilisearch/tests/vector/rest.rs b/crates/meilisearch/tests/vector/rest.rs index 1fdd18d28..363931a86 100644 --- a/crates/meilisearch/tests/vector/rest.rs +++ b/crates/meilisearch/tests/vector/rest.rs @@ -1,14 +1,12 @@ -use std::collections::{BTreeMap, BTreeSet}; +use std::collections::BTreeMap; use std::sync::atomic::AtomicUsize; use meili_snap::{json_string, snapshot}; use reqwest::IntoUrl; -use tokio::spawn; +use std::time::Duration; use tokio::sync::mpsc; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, Request, ResponseTemplate}; -use tokio::time::sleep; -use std::time::Duration; use crate::common::Value; use crate::json; @@ -342,7 +340,7 @@ async fn create_mock_raw() -> (MockServer, Value) { async fn create_faulty_mock_raw(sender: mpsc::Sender<()>) -> (MockServer, Value) { let mock_server = MockServer::start().await; let count = AtomicUsize::new(0); - + Mock::given(method("POST")) .and(path("/")) .respond_with(move |req: &Request| { @@ -359,14 +357,13 @@ async fn create_faulty_mock_raw(sender: mpsc::Sender<()>) -> (MockServer, Value) if count >= 5 { let _ = sender.try_send(()); - ResponseTemplate::new(500) - .set_delay(Duration::from_secs(u64::MAX)) - .set_body_json(json!({ + ResponseTemplate::new(500).set_delay(Duration::from_secs(u64::MAX)).set_body_json( + json!({ "error": "Service Unavailable", "text": req_body - })) + }), + ) } else { - ResponseTemplate::new(500).set_body_json(json!({ "error": "Service Unavailable", "text": req_body @@ -2168,7 +2165,6 @@ async fn searchable_reindex() { "###); } - #[actix_rt::test] async fn last_error_stats() { let (sender, mut receiver) = mpsc::channel(10); @@ -2191,7 +2187,7 @@ async fn last_error_stats() { {"id": 1, "name": "will_error"}, {"id": 2, "name": "must_error"}, ]); - let (value, code) = index.add_documents(documents, None).await; + let (_value, code) = index.add_documents(documents, None).await; snapshot!(code, @"202 Accepted"); // The task will eventually fail, so let's not wait for it. diff --git a/crates/milli/src/progress.rs b/crates/milli/src/progress.rs index 8cd2c9336..7ecfcc095 100644 --- a/crates/milli/src/progress.rs +++ b/crates/milli/src/progress.rs @@ -25,7 +25,7 @@ pub struct Progress { #[derive(Default)] pub struct EmbedderStats { pub errors: Arc, u32)>>, - pub total_count: AtomicUsize + pub total_count: AtomicUsize, } impl std::fmt::Debug for EmbedderStats { diff --git a/crates/milli/src/search/new/tests/integration.rs b/crates/milli/src/search/new/tests/integration.rs index 0b7e1a292..c4e521a88 100644 --- a/crates/milli/src/search/new/tests/integration.rs +++ b/crates/milli/src/search/new/tests/integration.rs @@ -95,7 +95,7 @@ pub fn setup_search_index_with_criteria(criteria: &[Criterion]) -> Index { embedders, &|| false, &Progress::default(), - Default::default(), + Default::default(), ) .unwrap(); diff --git a/crates/milli/src/update/new/indexer/extract.rs b/crates/milli/src/update/new/indexer/extract.rs index 72c63b605..040886236 100644 --- a/crates/milli/src/update/new/indexer/extract.rs +++ b/crates/milli/src/update/new/indexer/extract.rs @@ -1,7 +1,7 @@ use std::collections::BTreeMap; use std::sync::atomic::AtomicBool; -use std::sync::OnceLock; use std::sync::Arc; +use std::sync::OnceLock; use bumpalo::Bump; use roaring::RoaringBitmap; diff --git a/crates/milli/src/update/new/indexer/mod.rs b/crates/milli/src/update/new/indexer/mod.rs index 52fd6cd0b..33774f892 100644 --- a/crates/milli/src/update/new/indexer/mod.rs +++ b/crates/milli/src/update/new/indexer/mod.rs @@ -1,7 +1,7 @@ use std::sync::atomic::AtomicBool; +use std::sync::Arc; use std::sync::{Once, RwLock}; use std::thread::{self, Builder}; -use std::sync::Arc; use big_s::S; use document_changes::{DocumentChanges, IndexingContext}; diff --git a/crates/milli/src/update/settings.rs b/crates/milli/src/update/settings.rs index 98ee86978..b3f70d1b6 100644 --- a/crates/milli/src/update/settings.rs +++ b/crates/milli/src/update/settings.rs @@ -1358,7 +1358,12 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> { } } - pub fn execute(mut self, progress_callback: FP, should_abort: FA, embedder_stats: Arc) -> Result<()> + pub fn execute( + mut self, + progress_callback: FP, + should_abort: FA, + embedder_stats: Arc, + ) -> Result<()> where FP: Fn(UpdateIndexingStep) + Sync, FA: Fn() -> bool + Sync, diff --git a/crates/milli/src/vector/composite.rs b/crates/milli/src/vector/composite.rs index daec50e4b..7d9497165 100644 --- a/crates/milli/src/vector/composite.rs +++ b/crates/milli/src/vector/composite.rs @@ -173,12 +173,14 @@ impl SubEmbedder { ) -> std::result::Result { match self { SubEmbedder::HuggingFace(embedder) => embedder.embed_one(text), - SubEmbedder::OpenAi(embedder) => { - embedder.embed(&[text], deadline, embedder_stats)?.pop().ok_or_else(EmbedError::missing_embedding) - } - SubEmbedder::Ollama(embedder) => { - embedder.embed(&[text], deadline, embedder_stats)?.pop().ok_or_else(EmbedError::missing_embedding) - } + SubEmbedder::OpenAi(embedder) => embedder + .embed(&[text], deadline, embedder_stats)? + .pop() + .ok_or_else(EmbedError::missing_embedding), + SubEmbedder::Ollama(embedder) => embedder + .embed(&[text], deadline, embedder_stats)? + .pop() + .ok_or_else(EmbedError::missing_embedding), SubEmbedder::UserProvided(embedder) => embedder.embed_one(text), SubEmbedder::Rest(embedder) => embedder .embed_ref(&[text], deadline, embedder_stats)? @@ -198,10 +200,16 @@ impl SubEmbedder { ) -> std::result::Result>, EmbedError> { match self { SubEmbedder::HuggingFace(embedder) => embedder.embed_index(text_chunks), - SubEmbedder::OpenAi(embedder) => embedder.embed_index(text_chunks, threads, embedder_stats), - SubEmbedder::Ollama(embedder) => embedder.embed_index(text_chunks, threads, embedder_stats), + SubEmbedder::OpenAi(embedder) => { + embedder.embed_index(text_chunks, threads, embedder_stats) + } + SubEmbedder::Ollama(embedder) => { + embedder.embed_index(text_chunks, threads, embedder_stats) + } SubEmbedder::UserProvided(embedder) => embedder.embed_index(text_chunks), - SubEmbedder::Rest(embedder) => embedder.embed_index(text_chunks, threads, embedder_stats), + SubEmbedder::Rest(embedder) => { + embedder.embed_index(text_chunks, threads, embedder_stats) + } } } @@ -214,8 +222,12 @@ impl SubEmbedder { ) -> std::result::Result, EmbedError> { match self { SubEmbedder::HuggingFace(embedder) => embedder.embed_index_ref(texts), - SubEmbedder::OpenAi(embedder) => embedder.embed_index_ref(texts, threads, embedder_stats), - SubEmbedder::Ollama(embedder) => embedder.embed_index_ref(texts, threads, embedder_stats), + SubEmbedder::OpenAi(embedder) => { + embedder.embed_index_ref(texts, threads, embedder_stats) + } + SubEmbedder::Ollama(embedder) => { + embedder.embed_index_ref(texts, threads, embedder_stats) + } SubEmbedder::UserProvided(embedder) => embedder.embed_index_ref(texts), SubEmbedder::Rest(embedder) => embedder.embed_index_ref(texts, threads, embedder_stats), } diff --git a/crates/milli/src/vector/mod.rs b/crates/milli/src/vector/mod.rs index 124e17cff..efa981694 100644 --- a/crates/milli/src/vector/mod.rs +++ b/crates/milli/src/vector/mod.rs @@ -719,12 +719,14 @@ impl Embedder { } let embedding = match self { Embedder::HuggingFace(embedder) => embedder.embed_one(text), - Embedder::OpenAi(embedder) => { - embedder.embed(&[text], deadline, None)?.pop().ok_or_else(EmbedError::missing_embedding) - } - Embedder::Ollama(embedder) => { - embedder.embed(&[text], deadline, None)?.pop().ok_or_else(EmbedError::missing_embedding) - } + Embedder::OpenAi(embedder) => embedder + .embed(&[text], deadline, None)? + .pop() + .ok_or_else(EmbedError::missing_embedding), + Embedder::Ollama(embedder) => embedder + .embed(&[text], deadline, None)? + .pop() + .ok_or_else(EmbedError::missing_embedding), Embedder::UserProvided(embedder) => embedder.embed_one(text), Embedder::Rest(embedder) => embedder .embed_ref(&[text], deadline, None)? @@ -751,11 +753,17 @@ impl Embedder { ) -> std::result::Result>, EmbedError> { match self { Embedder::HuggingFace(embedder) => embedder.embed_index(text_chunks), - Embedder::OpenAi(embedder) => embedder.embed_index(text_chunks, threads, embedder_stats), - Embedder::Ollama(embedder) => embedder.embed_index(text_chunks, threads, embedder_stats), + Embedder::OpenAi(embedder) => { + embedder.embed_index(text_chunks, threads, embedder_stats) + } + Embedder::Ollama(embedder) => { + embedder.embed_index(text_chunks, threads, embedder_stats) + } Embedder::UserProvided(embedder) => embedder.embed_index(text_chunks), Embedder::Rest(embedder) => embedder.embed_index(text_chunks, threads, embedder_stats), - Embedder::Composite(embedder) => embedder.index.embed_index(text_chunks, threads, embedder_stats), + Embedder::Composite(embedder) => { + embedder.index.embed_index(text_chunks, threads, embedder_stats) + } } } @@ -772,7 +780,9 @@ impl Embedder { Embedder::Ollama(embedder) => embedder.embed_index_ref(texts, threads, embedder_stats), Embedder::UserProvided(embedder) => embedder.embed_index_ref(texts), Embedder::Rest(embedder) => embedder.embed_index_ref(texts, threads, embedder_stats), - Embedder::Composite(embedder) => embedder.index.embed_index_ref(texts, threads, embedder_stats), + Embedder::Composite(embedder) => { + embedder.index.embed_index_ref(texts, threads, embedder_stats) + } } } diff --git a/crates/milli/src/vector/ollama.rs b/crates/milli/src/vector/ollama.rs index b3ee925e6..e26b7e1ea 100644 --- a/crates/milli/src/vector/ollama.rs +++ b/crates/milli/src/vector/ollama.rs @@ -106,7 +106,7 @@ impl Embedder { &self, texts: &[S], deadline: Option, - embedder_stats: Option> + embedder_stats: Option>, ) -> Result, EmbedError> { match self.rest_embedder.embed_ref(texts, deadline, embedder_stats) { Ok(embeddings) => Ok(embeddings), @@ -126,11 +126,17 @@ impl Embedder { // This condition helps reduce the number of active rayon jobs // so that we avoid consuming all the LMDB rtxns and avoid stack overflows. if threads.active_operations() >= REQUEST_PARALLELISM { - text_chunks.into_iter().map(move |chunk| self.embed(&chunk, None, embedder_stats.clone())).collect() + text_chunks + .into_iter() + .map(move |chunk| self.embed(&chunk, None, embedder_stats.clone())) + .collect() } else { threads .install(move || { - text_chunks.into_par_iter().map(move |chunk| self.embed(&chunk, None, embedder_stats.clone())).collect() + text_chunks + .into_par_iter() + .map(move |chunk| self.embed(&chunk, None, embedder_stats.clone())) + .collect() }) .map_err(|error| EmbedError { kind: EmbedErrorKind::PanicInThreadPool(error), @@ -143,7 +149,7 @@ impl Embedder { &self, texts: &[&str], threads: &ThreadPoolNoAbort, - embedder_stats: Option> + embedder_stats: Option>, ) -> Result>, EmbedError> { // This condition helps reduce the number of active rayon jobs // so that we avoid consuming all the LMDB rtxns and avoid stack overflows. diff --git a/crates/milli/src/vector/openai.rs b/crates/milli/src/vector/openai.rs index 384abe880..ca072d6e5 100644 --- a/crates/milli/src/vector/openai.rs +++ b/crates/milli/src/vector/openai.rs @@ -241,7 +241,11 @@ impl Embedder { let encoded = self.tokenizer.encode_ordinary(text); let len = encoded.len(); if len < max_token_count { - all_embeddings.append(&mut self.rest_embedder.embed_ref(&[text], deadline, None)?); + all_embeddings.append(&mut self.rest_embedder.embed_ref( + &[text], + deadline, + None, + )?); continue; } @@ -263,11 +267,17 @@ impl Embedder { // This condition helps reduce the number of active rayon jobs // so that we avoid consuming all the LMDB rtxns and avoid stack overflows. if threads.active_operations() >= REQUEST_PARALLELISM { - text_chunks.into_iter().map(move |chunk| self.embed(&chunk, None, embedder_stats.clone())).collect() + text_chunks + .into_iter() + .map(move |chunk| self.embed(&chunk, None, embedder_stats.clone())) + .collect() } else { threads .install(move || { - text_chunks.into_par_iter().map(move |chunk| self.embed(&chunk, None, embedder_stats.clone())).collect() + text_chunks + .into_par_iter() + .map(move |chunk| self.embed(&chunk, None, embedder_stats.clone())) + .collect() }) .map_err(|error| EmbedError { kind: EmbedErrorKind::PanicInThreadPool(error), diff --git a/crates/milli/src/vector/rest.rs b/crates/milli/src/vector/rest.rs index d8de89c6a..294b0ceda 100644 --- a/crates/milli/src/vector/rest.rs +++ b/crates/milli/src/vector/rest.rs @@ -14,8 +14,8 @@ use super::{ DistributionShift, EmbedError, Embedding, EmbeddingCache, NewEmbedderError, REQUEST_PARALLELISM, }; use crate::error::FaultSource; -use crate::ThreadPoolNoAbort; use crate::progress::EmbedderStats; +use crate::ThreadPoolNoAbort; // retrying in case of failure pub struct Retry { @@ -172,7 +172,14 @@ impl Embedder { deadline: Option, embedder_stats: Option>, ) -> Result, EmbedError> { - embed(&self.data, texts.as_slice(), texts.len(), Some(self.dimensions), deadline, embedder_stats) + embed( + &self.data, + texts.as_slice(), + texts.len(), + Some(self.dimensions), + deadline, + embedder_stats, + ) } pub fn embed_ref( @@ -206,11 +213,17 @@ impl Embedder { // This condition helps reduce the number of active rayon jobs // so that we avoid consuming all the LMDB rtxns and avoid stack overflows. if threads.active_operations() >= REQUEST_PARALLELISM { - text_chunks.into_iter().map(move |chunk| self.embed(chunk, None, embedder_stats.clone())).collect() + text_chunks + .into_iter() + .map(move |chunk| self.embed(chunk, None, embedder_stats.clone())) + .collect() } else { threads .install(move || { - text_chunks.into_par_iter().map(move |chunk| self.embed(chunk, None, embedder_stats.clone())).collect() + text_chunks + .into_par_iter() + .map(move |chunk| self.embed(chunk, None, embedder_stats.clone())) + .collect() }) .map_err(|error| EmbedError { kind: EmbedErrorKind::PanicInThreadPool(error), @@ -223,7 +236,7 @@ impl Embedder { &self, texts: &[&str], threads: &ThreadPoolNoAbort, - embedder_stats: Option> + embedder_stats: Option>, ) -> Result, EmbedError> { // This condition helps reduce the number of active rayon jobs // so that we avoid consuming all the LMDB rtxns and avoid stack overflows. From bc4d1530ee3ffbd8434946f26ac582ff46011c61 Mon Sep 17 00:00:00 2001 From: Mubelotix Date: Tue, 24 Jun 2025 14:50:23 +0200 Subject: [PATCH 07/18] Fix tests --- .gitignore | 3 ++ crates/index-scheduler/src/insta_snapshot.rs | 10 +++--- crates/meilisearch/tests/vector/rest.rs | 32 +++++++++++++++++++- 3 files changed, 40 insertions(+), 5 deletions(-) diff --git a/.gitignore b/.gitignore index 07453a58f..fc24b8306 100644 --- a/.gitignore +++ b/.gitignore @@ -18,5 +18,8 @@ ## ... unreviewed *.snap.new +# Database snapshot +crates/meilisearch/db.snapshot + # Fuzzcheck data for the facet indexing fuzz test crates/milli/fuzz/update::facet::incremental::fuzz::fuzz/ diff --git a/crates/index-scheduler/src/insta_snapshot.rs b/crates/index-scheduler/src/insta_snapshot.rs index d4504c246..a5bb1ea56 100644 --- a/crates/index-scheduler/src/insta_snapshot.rs +++ b/crates/index-scheduler/src/insta_snapshot.rs @@ -367,10 +367,12 @@ pub fn snapshot_batch(batch: &Batch) -> String { snap.push_str(&format!("uid: {uid}, ")); snap.push_str(&format!("details: {}, ", serde_json::to_string(details).unwrap())); snap.push_str(&format!("stats: {}, ", serde_json::to_string(&stats).unwrap())); - snap.push_str(&format!( - "embedder_stats: {}, ", - serde_json::to_string(&embedder_stats).unwrap() - )); + if !embedder_stats.skip_serializing() { + snap.push_str(&format!( + "embedder stats: {}, ", + serde_json::to_string(&embedder_stats).unwrap() + )); + } snap.push_str(&format!("stop reason: {}, ", serde_json::to_string(&stop_reason).unwrap())); snap.push('}'); snap diff --git a/crates/meilisearch/tests/vector/rest.rs b/crates/meilisearch/tests/vector/rest.rs index 363931a86..2c8d3ed7c 100644 --- a/crates/meilisearch/tests/vector/rest.rs +++ b/crates/meilisearch/tests/vector/rest.rs @@ -2195,5 +2195,35 @@ async fn last_error_stats() { receiver.recv().await; let (response, _code) = index.filtered_batches(&[], &[], &[]).await; - snapshot!(response["results"][0], @r###""###); + snapshot!(json_string!(response["results"][0], { ".progress" => "[ignored]", ".stats.embedder.totalCount" => "[ignored]", ".startedAt" => "[ignored]" }), @r#" + { + "uid": 1, + "progress": "[ignored]", + "details": { + "receivedDocuments": 3, + "indexedDocuments": null + }, + "stats": { + "totalNbTasks": 1, + "status": { + "processing": 1 + }, + "types": { + "documentAdditionOrUpdate": 1 + }, + "indexUids": { + "doggo": 1 + }, + "embedder": { + "totalCount": "[ignored]", + "errorCount": 5, + "lastError": "runtime error: received internal error HTTP 500 from embedding server\n - server replied with `{\"error\":\"Service Unavailable\",\"text\":\"will_error\"}`" + } + }, + "duration": null, + "startedAt": "[ignored]", + "finishedAt": null, + "batchStrategy": "batched all enqueued tasks" + } + "#); } From 695877043ae6af463915447a9319771c0ee2974c Mon Sep 17 00:00:00 2001 From: Mubelotix Date: Tue, 24 Jun 2025 14:53:39 +0200 Subject: [PATCH 08/18] Fix warnings --- .../src/update/index_documents/extract/extract_vector_points.rs | 1 + crates/milli/src/update/new/extract/vectors/mod.rs | 1 - 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/milli/src/update/index_documents/extract/extract_vector_points.rs b/crates/milli/src/update/index_documents/extract/extract_vector_points.rs index de91e9f10..e940e743b 100644 --- a/crates/milli/src/update/index_documents/extract/extract_vector_points.rs +++ b/crates/milli/src/update/index_documents/extract/extract_vector_points.rs @@ -675,6 +675,7 @@ fn compare_vectors(a: &[f32], b: &[f32]) -> Ordering { a.iter().copied().map(OrderedFloat).cmp(b.iter().copied().map(OrderedFloat)) } +#[allow(clippy::too_many_arguments)] #[tracing::instrument(level = "trace", skip_all, target = "indexing::extract")] pub fn extract_embeddings( // docid, prompt diff --git a/crates/milli/src/update/new/extract/vectors/mod.rs b/crates/milli/src/update/new/extract/vectors/mod.rs index f720b81e2..946fb00b5 100644 --- a/crates/milli/src/update/new/extract/vectors/mod.rs +++ b/crates/milli/src/update/new/extract/vectors/mod.rs @@ -1,4 +1,3 @@ -use std::f32::consts::E; use std::{cell::RefCell, sync::Arc}; use bumpalo::collections::Vec as BVec; From d08e89ea3d36026d4dcb554f7fb45170bc398d17 Mon Sep 17 00:00:00 2001 From: Mubelotix Date: Tue, 24 Jun 2025 15:10:15 +0200 Subject: [PATCH 09/18] Remove options --- .../index_documents/extract/extract_vector_points.rs | 6 ++---- .../milli/src/update/index_documents/extract/mod.rs | 2 +- crates/milli/src/update/new/extract/vectors/mod.rs | 10 +++++----- crates/milli/src/update/new/indexer/extract.rs | 2 +- crates/milli/src/vector/composite.rs | 4 ++-- crates/milli/src/vector/mod.rs | 4 ++-- crates/milli/src/vector/ollama.rs | 12 ++++++------ crates/milli/src/vector/openai.rs | 12 ++++++------ crates/milli/src/vector/rest.rs | 12 ++++++------ 9 files changed, 31 insertions(+), 33 deletions(-) diff --git a/crates/milli/src/update/index_documents/extract/extract_vector_points.rs b/crates/milli/src/update/index_documents/extract/extract_vector_points.rs index e940e743b..e6d874a69 100644 --- a/crates/milli/src/update/index_documents/extract/extract_vector_points.rs +++ b/crates/milli/src/update/index_documents/extract/extract_vector_points.rs @@ -684,12 +684,10 @@ pub fn extract_embeddings( embedder: Arc, embedder_name: &str, possible_embedding_mistakes: &PossibleEmbeddingMistakes, - embedder_stats: Option>, + embedder_stats: Arc, unused_vectors_distribution: &UnusedVectorsDistribution, request_threads: &ThreadPoolNoAbort, ) -> Result>> { - println!("Extract embedder stats {}:", embedder_stats.is_some()); - let n_chunks = embedder.chunk_count_hint(); // chunk level parallelism let n_vectors_per_chunk = embedder.prompt_count_in_chunk_hint(); // number of vectors in a single chunk @@ -791,7 +789,7 @@ fn embed_chunks( text_chunks: Vec>, embedder_name: &str, possible_embedding_mistakes: &PossibleEmbeddingMistakes, - embedder_stats: Option>, + embedder_stats: Arc, unused_vectors_distribution: &UnusedVectorsDistribution, request_threads: &ThreadPoolNoAbort, ) -> Result>> { diff --git a/crates/milli/src/update/index_documents/extract/mod.rs b/crates/milli/src/update/index_documents/extract/mod.rs index f4f3ad22e..1eeddcccb 100644 --- a/crates/milli/src/update/index_documents/extract/mod.rs +++ b/crates/milli/src/update/index_documents/extract/mod.rs @@ -274,7 +274,7 @@ fn send_original_documents_data( embedder.clone(), &embedder_name, &possible_embedding_mistakes, - Some(embedder_stats.clone()), + embedder_stats.clone(), &unused_vectors_distribution, request_threads(), ) { diff --git a/crates/milli/src/update/new/extract/vectors/mod.rs b/crates/milli/src/update/new/extract/vectors/mod.rs index 946fb00b5..c21dabf74 100644 --- a/crates/milli/src/update/new/extract/vectors/mod.rs +++ b/crates/milli/src/update/new/extract/vectors/mod.rs @@ -23,7 +23,7 @@ pub struct EmbeddingExtractor<'a, 'b> { embedders: &'a EmbeddingConfigs, sender: EmbeddingSender<'a, 'b>, possible_embedding_mistakes: PossibleEmbeddingMistakes, - embedder_stats: Option>, + embedder_stats: Arc, threads: &'a ThreadPoolNoAbort, } @@ -32,7 +32,7 @@ impl<'a, 'b> EmbeddingExtractor<'a, 'b> { embedders: &'a EmbeddingConfigs, sender: EmbeddingSender<'a, 'b>, field_distribution: &'a FieldDistribution, - embedder_stats: Option>, + embedder_stats: Arc, threads: &'a ThreadPoolNoAbort, ) -> Self { let possible_embedding_mistakes = PossibleEmbeddingMistakes::new(field_distribution); @@ -311,7 +311,7 @@ struct Chunks<'a, 'b, 'extractor> { dimensions: usize, prompt: &'a Prompt, possible_embedding_mistakes: &'a PossibleEmbeddingMistakes, - embedder_stats: Option>, + embedder_stats: Arc, user_provided: &'a RefCell>, threads: &'a ThreadPoolNoAbort, sender: EmbeddingSender<'a, 'b>, @@ -327,7 +327,7 @@ impl<'a, 'b, 'extractor> Chunks<'a, 'b, 'extractor> { prompt: &'a Prompt, user_provided: &'a RefCell>, possible_embedding_mistakes: &'a PossibleEmbeddingMistakes, - embedder_stats: Option>, + embedder_stats: Arc, threads: &'a ThreadPoolNoAbort, sender: EmbeddingSender<'a, 'b>, doc_alloc: &'a Bump, @@ -416,7 +416,7 @@ impl<'a, 'b, 'extractor> Chunks<'a, 'b, 'extractor> { embedder_id: u8, embedder_name: &str, possible_embedding_mistakes: &PossibleEmbeddingMistakes, - embedder_stats: Option>, + embedder_stats: Arc, unused_vectors_distribution: &UnusedVectorsDistributionBump, threads: &ThreadPoolNoAbort, sender: EmbeddingSender<'a, 'b>, diff --git a/crates/milli/src/update/new/indexer/extract.rs b/crates/milli/src/update/new/indexer/extract.rs index 040886236..c721a2563 100644 --- a/crates/milli/src/update/new/indexer/extract.rs +++ b/crates/milli/src/update/new/indexer/extract.rs @@ -248,7 +248,7 @@ where embedders, embedding_sender, field_distribution, - Some(embedder_stats), + embedder_stats, request_threads(), ); let mut datastore = ThreadLocal::with_capacity(rayon::current_num_threads()); diff --git a/crates/milli/src/vector/composite.rs b/crates/milli/src/vector/composite.rs index 7d9497165..87f05d4fe 100644 --- a/crates/milli/src/vector/composite.rs +++ b/crates/milli/src/vector/composite.rs @@ -196,7 +196,7 @@ impl SubEmbedder { &self, text_chunks: Vec>, threads: &ThreadPoolNoAbort, - embedder_stats: Option>, + embedder_stats: Arc, ) -> std::result::Result>, EmbedError> { match self { SubEmbedder::HuggingFace(embedder) => embedder.embed_index(text_chunks), @@ -218,7 +218,7 @@ impl SubEmbedder { &self, texts: &[&str], threads: &ThreadPoolNoAbort, - embedder_stats: Option>, + embedder_stats: Arc, ) -> std::result::Result, EmbedError> { match self { SubEmbedder::HuggingFace(embedder) => embedder.embed_index_ref(texts), diff --git a/crates/milli/src/vector/mod.rs b/crates/milli/src/vector/mod.rs index efa981694..481eb6c99 100644 --- a/crates/milli/src/vector/mod.rs +++ b/crates/milli/src/vector/mod.rs @@ -749,7 +749,7 @@ impl Embedder { &self, text_chunks: Vec>, threads: &ThreadPoolNoAbort, - embedder_stats: Option>, + embedder_stats: Arc, ) -> std::result::Result>, EmbedError> { match self { Embedder::HuggingFace(embedder) => embedder.embed_index(text_chunks), @@ -772,7 +772,7 @@ impl Embedder { &self, texts: &[&str], threads: &ThreadPoolNoAbort, - embedder_stats: Option>, + embedder_stats: Arc, ) -> std::result::Result, EmbedError> { match self { Embedder::HuggingFace(embedder) => embedder.embed_index_ref(texts), diff --git a/crates/milli/src/vector/ollama.rs b/crates/milli/src/vector/ollama.rs index e26b7e1ea..045b65b72 100644 --- a/crates/milli/src/vector/ollama.rs +++ b/crates/milli/src/vector/ollama.rs @@ -121,21 +121,21 @@ impl Embedder { &self, text_chunks: Vec>, threads: &ThreadPoolNoAbort, - embedder_stats: Option>, + embedder_stats: Arc, ) -> Result>, EmbedError> { // This condition helps reduce the number of active rayon jobs // so that we avoid consuming all the LMDB rtxns and avoid stack overflows. if threads.active_operations() >= REQUEST_PARALLELISM { text_chunks .into_iter() - .map(move |chunk| self.embed(&chunk, None, embedder_stats.clone())) + .map(move |chunk| self.embed(&chunk, None, Some(embedder_stats.clone()))) .collect() } else { threads .install(move || { text_chunks .into_par_iter() - .map(move |chunk| self.embed(&chunk, None, embedder_stats.clone())) + .map(move |chunk| self.embed(&chunk, None, Some(embedder_stats.clone()))) .collect() }) .map_err(|error| EmbedError { @@ -149,14 +149,14 @@ impl Embedder { &self, texts: &[&str], threads: &ThreadPoolNoAbort, - embedder_stats: Option>, + embedder_stats: Arc, ) -> Result>, EmbedError> { // This condition helps reduce the number of active rayon jobs // so that we avoid consuming all the LMDB rtxns and avoid stack overflows. if threads.active_operations() >= REQUEST_PARALLELISM { let embeddings: Result>, _> = texts .chunks(self.prompt_count_in_chunk_hint()) - .map(move |chunk| self.embed(chunk, None, embedder_stats.clone())) + .map(move |chunk| self.embed(chunk, None, Some(embedder_stats.clone()))) .collect(); let embeddings = embeddings?; @@ -166,7 +166,7 @@ impl Embedder { .install(move || { let embeddings: Result>, _> = texts .par_chunks(self.prompt_count_in_chunk_hint()) - .map(move |chunk| self.embed(chunk, None, embedder_stats.clone())) + .map(move |chunk| self.embed(chunk, None, Some(embedder_stats.clone()))) .collect(); let embeddings = embeddings?; diff --git a/crates/milli/src/vector/openai.rs b/crates/milli/src/vector/openai.rs index ca072d6e5..b64e3d467 100644 --- a/crates/milli/src/vector/openai.rs +++ b/crates/milli/src/vector/openai.rs @@ -262,21 +262,21 @@ impl Embedder { &self, text_chunks: Vec>, threads: &ThreadPoolNoAbort, - embedder_stats: Option>, + embedder_stats: Arc, ) -> Result>, EmbedError> { // This condition helps reduce the number of active rayon jobs // so that we avoid consuming all the LMDB rtxns and avoid stack overflows. if threads.active_operations() >= REQUEST_PARALLELISM { text_chunks .into_iter() - .map(move |chunk| self.embed(&chunk, None, embedder_stats.clone())) + .map(move |chunk| self.embed(&chunk, None, Some(embedder_stats.clone()))) .collect() } else { threads .install(move || { text_chunks .into_par_iter() - .map(move |chunk| self.embed(&chunk, None, embedder_stats.clone())) + .map(move |chunk| self.embed(&chunk, None, Some(embedder_stats.clone()))) .collect() }) .map_err(|error| EmbedError { @@ -290,14 +290,14 @@ impl Embedder { &self, texts: &[&str], threads: &ThreadPoolNoAbort, - embedder_stats: Option>, + embedder_stats: Arc, ) -> Result>, EmbedError> { // This condition helps reduce the number of active rayon jobs // so that we avoid consuming all the LMDB rtxns and avoid stack overflows. if threads.active_operations() >= REQUEST_PARALLELISM { let embeddings: Result>, _> = texts .chunks(self.prompt_count_in_chunk_hint()) - .map(move |chunk| self.embed(chunk, None, embedder_stats.clone())) + .map(move |chunk| self.embed(chunk, None, Some(embedder_stats.clone()))) .collect(); let embeddings = embeddings?; Ok(embeddings.into_iter().flatten().collect()) @@ -306,7 +306,7 @@ impl Embedder { .install(move || { let embeddings: Result>, _> = texts .par_chunks(self.prompt_count_in_chunk_hint()) - .map(move |chunk| self.embed(chunk, None, embedder_stats.clone())) + .map(move |chunk| self.embed(chunk, None, Some(embedder_stats.clone()))) .collect(); let embeddings = embeddings?; diff --git a/crates/milli/src/vector/rest.rs b/crates/milli/src/vector/rest.rs index 294b0ceda..409284b65 100644 --- a/crates/milli/src/vector/rest.rs +++ b/crates/milli/src/vector/rest.rs @@ -208,21 +208,21 @@ impl Embedder { &self, text_chunks: Vec>, threads: &ThreadPoolNoAbort, - embedder_stats: Option>, + embedder_stats: Arc, ) -> Result>, EmbedError> { // This condition helps reduce the number of active rayon jobs // so that we avoid consuming all the LMDB rtxns and avoid stack overflows. if threads.active_operations() >= REQUEST_PARALLELISM { text_chunks .into_iter() - .map(move |chunk| self.embed(chunk, None, embedder_stats.clone())) + .map(move |chunk| self.embed(chunk, None, Some(embedder_stats.clone()))) .collect() } else { threads .install(move || { text_chunks .into_par_iter() - .map(move |chunk| self.embed(chunk, None, embedder_stats.clone())) + .map(move |chunk| self.embed(chunk, None, Some(embedder_stats.clone()))) .collect() }) .map_err(|error| EmbedError { @@ -236,14 +236,14 @@ impl Embedder { &self, texts: &[&str], threads: &ThreadPoolNoAbort, - embedder_stats: Option>, + embedder_stats: Arc, ) -> Result, EmbedError> { // This condition helps reduce the number of active rayon jobs // so that we avoid consuming all the LMDB rtxns and avoid stack overflows. if threads.active_operations() >= REQUEST_PARALLELISM { let embeddings: Result>, _> = texts .chunks(self.prompt_count_in_chunk_hint()) - .map(move |chunk| self.embed_ref(chunk, None, embedder_stats.clone())) + .map(move |chunk| self.embed_ref(chunk, None, Some(embedder_stats.clone()))) .collect(); let embeddings = embeddings?; @@ -253,7 +253,7 @@ impl Embedder { .install(move || { let embeddings: Result>, _> = texts .par_chunks(self.prompt_count_in_chunk_hint()) - .map(move |chunk| self.embed_ref(chunk, None, embedder_stats.clone())) + .map(move |chunk| self.embed_ref(chunk, None, Some(embedder_stats.clone()))) .collect(); let embeddings = embeddings?; From 1d3b18f774029fe8c3710ae632aa981b25ababa9 Mon Sep 17 00:00:00 2001 From: Mubelotix Date: Wed, 25 Jun 2025 14:58:21 +0200 Subject: [PATCH 10/18] Update test to be more reproducible --- crates/meilisearch/tests/vector/rest.rs | 33 +++++++++---------------- 1 file changed, 11 insertions(+), 22 deletions(-) diff --git a/crates/meilisearch/tests/vector/rest.rs b/crates/meilisearch/tests/vector/rest.rs index 2c8d3ed7c..7e2245223 100644 --- a/crates/meilisearch/tests/vector/rest.rs +++ b/crates/meilisearch/tests/vector/rest.rs @@ -343,31 +343,16 @@ async fn create_faulty_mock_raw(sender: mpsc::Sender<()>) -> (MockServer, Value) Mock::given(method("POST")) .and(path("/")) - .respond_with(move |req: &Request| { + .respond_with(move |_req: &Request| { let count = count.fetch_add(1, std::sync::atomic::Ordering::SeqCst); - let req_body = match req.body_json::() { - Ok(body) => body, - Err(error) => { - return ResponseTemplate::new(400).set_body_json(json!({ - "error": format!("Invalid request: {error}") - })); - } - }; - if count >= 5 { let _ = sender.try_send(()); - ResponseTemplate::new(500).set_delay(Duration::from_secs(u64::MAX)).set_body_json( - json!({ - "error": "Service Unavailable", - "text": req_body - }), - ) + ResponseTemplate::new(500) + .set_delay(Duration::from_secs(u64::MAX)) + .set_body_string("Service Unavailable") } else { - ResponseTemplate::new(500).set_body_json(json!({ - "error": "Service Unavailable", - "text": req_body - })) + ResponseTemplate::new(500).set_body_string("Service Unavailable") } }) .mount(&mock_server) @@ -2195,7 +2180,11 @@ async fn last_error_stats() { receiver.recv().await; let (response, _code) = index.filtered_batches(&[], &[], &[]).await; - snapshot!(json_string!(response["results"][0], { ".progress" => "[ignored]", ".stats.embedder.totalCount" => "[ignored]", ".startedAt" => "[ignored]" }), @r#" + snapshot!(json_string!(response["results"][0], { + ".progress" => "[ignored]", + ".stats.embedder.totalCount" => "[ignored]", + ".startedAt" => "[ignored]" + }), @r#" { "uid": 1, "progress": "[ignored]", @@ -2217,7 +2206,7 @@ async fn last_error_stats() { "embedder": { "totalCount": "[ignored]", "errorCount": 5, - "lastError": "runtime error: received internal error HTTP 500 from embedding server\n - server replied with `{\"error\":\"Service Unavailable\",\"text\":\"will_error\"}`" + "lastError": "runtime error: received internal error HTTP 500 from embedding server\n - server replied with `Service Unavailable`" } }, "duration": null, From 9422b6d654cf44a60d540444310e9c96b173b18a Mon Sep 17 00:00:00 2001 From: Mubelotix Date: Thu, 26 Jun 2025 10:58:27 +0200 Subject: [PATCH 11/18] Update crates/meilisearch/src/lib.rs Co-authored-by: Louis Dureuil --- crates/meilisearch/src/lib.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/meilisearch/src/lib.rs b/crates/meilisearch/src/lib.rs index cdecd520c..5acfb4bc9 100644 --- a/crates/meilisearch/src/lib.rs +++ b/crates/meilisearch/src/lib.rs @@ -543,7 +543,7 @@ fn import_dump( tracing::info!("Importing the settings."); let settings = index_reader.settings()?; apply_settings_to_builder(&settings, &mut builder); - let embedder_stats: Arc = Default::default(); // FIXME: this isn't linked to anything + let embedder_stats: Arc = Default::default(); builder.execute( |indexing_step| tracing::debug!("update: {:?}", indexing_step), || false, From 3fc16c627dfdf033f9521264d2517108532d9b4c Mon Sep 17 00:00:00 2001 From: Mubelotix Date: Thu, 26 Jun 2025 11:11:38 +0200 Subject: [PATCH 12/18] Comment the delay --- crates/meilisearch/tests/vector/rest.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/meilisearch/tests/vector/rest.rs b/crates/meilisearch/tests/vector/rest.rs index 7e2245223..e80dfeb0a 100644 --- a/crates/meilisearch/tests/vector/rest.rs +++ b/crates/meilisearch/tests/vector/rest.rs @@ -349,7 +349,7 @@ async fn create_faulty_mock_raw(sender: mpsc::Sender<()>) -> (MockServer, Value) if count >= 5 { let _ = sender.try_send(()); ResponseTemplate::new(500) - .set_delay(Duration::from_secs(u64::MAX)) + .set_delay(Duration::from_secs(u64::MAX)) // Make the response hang forever .set_body_string("Service Unavailable") } else { ResponseTemplate::new(500).set_body_string("Service Unavailable") From ef007d547df6fcc48480a8593ce135f46863a293 Mon Sep 17 00:00:00 2001 From: Mubelotix Date: Thu, 26 Jun 2025 11:12:01 +0200 Subject: [PATCH 13/18] Remove panics --- crates/meilisearch-types/src/batches.rs | 2 +- crates/milli/src/progress.rs | 4 +++- crates/milli/src/vector/rest.rs | 19 ++++++++++--------- 3 files changed, 14 insertions(+), 11 deletions(-) diff --git a/crates/meilisearch-types/src/batches.rs b/crates/meilisearch-types/src/batches.rs index cec74fb75..c8d98655f 100644 --- a/crates/meilisearch-types/src/batches.rs +++ b/crates/meilisearch-types/src/batches.rs @@ -100,7 +100,7 @@ pub struct EmbedderStatsView { impl From<&EmbedderStats> for EmbedderStatsView { fn from(stats: &EmbedderStats) -> Self { - let errors = stats.errors.read().unwrap(); + let errors = stats.errors.read().unwrap_or_else(|p| p.into_inner()); Self { total_count: stats.total_count.load(std::sync::atomic::Ordering::Relaxed), error_count: errors.1 as usize, diff --git a/crates/milli/src/progress.rs b/crates/milli/src/progress.rs index 7ecfcc095..61c61cd49 100644 --- a/crates/milli/src/progress.rs +++ b/crates/milli/src/progress.rs @@ -30,7 +30,9 @@ pub struct EmbedderStats { impl std::fmt::Debug for EmbedderStats { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - let (error, count) = self.errors.read().unwrap().clone(); + let guard = self.errors.read().unwrap_or_else(|p| p.into_inner()); + let (error, count) = (guard.0.clone(), guard.1); + std::mem::drop(guard); f.debug_struct("EmbedderStats") .field("last_error", &error) .field("total_count", &self.total_count.load(Ordering::Relaxed)) diff --git a/crates/milli/src/vector/rest.rs b/crates/milli/src/vector/rest.rs index 409284b65..dd08c6a5e 100644 --- a/crates/milli/src/vector/rest.rs +++ b/crates/milli/src/vector/rest.rs @@ -335,10 +335,11 @@ where Err(retry) => { tracing::warn!("Failed: {}", retry.error); if let Some(embedder_stats) = &embedder_stats { - if let Ok(mut errors) = embedder_stats.errors.write() { - errors.0 = Some(retry.error.to_string()); - errors.1 += 1; - } + let stringified_error = retry.error.to_string(); + let mut errors = + embedder_stats.errors.write().unwrap_or_else(|p| p.into_inner()); + errors.0 = Some(stringified_error); + errors.1 += 1; } if let Some(deadline) = deadline { let now = std::time::Instant::now(); @@ -377,11 +378,11 @@ where Ok(response) => Ok(response), Err(retry) => { if let Some(embedder_stats) = &embedder_stats { - if let Ok(mut errors) = embedder_stats.errors.write() { - errors.0 = Some(retry.error.to_string()); - errors.1 += 1; - } - } + let stringified_error = retry.error.to_string(); + let mut errors = embedder_stats.errors.write().unwrap_or_else(|p| p.into_inner()); + errors.0 = Some(stringified_error); + errors.1 += 1; + }; Err(retry.into_error()) } } From 29f6eeff8fc82b55799e1c958249cb53349e603e Mon Sep 17 00:00:00 2001 From: Mubelotix Date: Thu, 26 Jun 2025 12:07:48 +0200 Subject: [PATCH 14/18] Remove lots of Arcs --- crates/benchmarks/benches/indexing.rs | 62 +++++++++---------- crates/benchmarks/benches/utils.rs | 2 +- crates/fuzzers/src/bin/fuzz-indexing.rs | 2 +- .../src/scheduler/process_index_operation.rs | 10 +-- .../milli/src/search/new/tests/integration.rs | 2 +- crates/milli/src/test_index.rs | 6 +- .../extract/extract_vector_points.rs | 8 +-- .../src/update/index_documents/extract/mod.rs | 6 +- .../milli/src/update/index_documents/mod.rs | 30 ++++----- .../src/update/new/extract/vectors/mod.rs | 18 +++--- .../milli/src/update/new/indexer/extract.rs | 3 +- crates/milli/src/update/new/indexer/mod.rs | 3 +- crates/milli/src/update/settings.rs | 4 +- crates/milli/src/vector/composite.rs | 9 ++- crates/milli/src/vector/mod.rs | 4 +- crates/milli/src/vector/ollama.rs | 15 +++-- crates/milli/src/vector/openai.rs | 15 +++-- crates/milli/src/vector/rest.rs | 23 ++++--- .../milli/tests/search/facet_distribution.rs | 2 +- crates/milli/tests/search/mod.rs | 2 +- crates/milli/tests/search/query_criteria.rs | 2 +- crates/milli/tests/search/typo_tolerance.rs | 2 +- 22 files changed, 112 insertions(+), 118 deletions(-) diff --git a/crates/benchmarks/benches/indexing.rs b/crates/benchmarks/benches/indexing.rs index 8241da9d2..3afad8ee5 100644 --- a/crates/benchmarks/benches/indexing.rs +++ b/crates/benchmarks/benches/indexing.rs @@ -169,7 +169,7 @@ fn indexing_songs_default(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); @@ -236,7 +236,7 @@ fn reindexing_songs_default(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); @@ -281,7 +281,7 @@ fn reindexing_songs_default(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); @@ -350,7 +350,7 @@ fn deleting_songs_in_batches_default(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); @@ -427,7 +427,7 @@ fn indexing_songs_in_three_batches_default(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); @@ -472,7 +472,7 @@ fn indexing_songs_in_three_batches_default(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); @@ -513,7 +513,7 @@ fn indexing_songs_in_three_batches_default(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); @@ -581,7 +581,7 @@ fn indexing_songs_without_faceted_numbers(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); @@ -648,7 +648,7 @@ fn indexing_songs_without_faceted_fields(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); @@ -715,7 +715,7 @@ fn indexing_wiki(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); @@ -781,7 +781,7 @@ fn reindexing_wiki(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); @@ -826,7 +826,7 @@ fn reindexing_wiki(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); @@ -894,7 +894,7 @@ fn deleting_wiki_in_batches_default(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); @@ -971,7 +971,7 @@ fn indexing_wiki_in_three_batches(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); @@ -1017,7 +1017,7 @@ fn indexing_wiki_in_three_batches(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); @@ -1059,7 +1059,7 @@ fn indexing_wiki_in_three_batches(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); @@ -1126,7 +1126,7 @@ fn indexing_movies_default(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); @@ -1192,7 +1192,7 @@ fn reindexing_movies_default(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); @@ -1237,7 +1237,7 @@ fn reindexing_movies_default(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); @@ -1305,7 +1305,7 @@ fn deleting_movies_in_batches_default(c: &mut Criterion) { EmbeddingConfigs::default(), &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); @@ -1354,7 +1354,7 @@ fn delete_documents_from_ids(index: Index, document_ids_to_delete: Vec Index { EmbeddingConfigs::default(), &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); diff --git a/crates/fuzzers/src/bin/fuzz-indexing.rs b/crates/fuzzers/src/bin/fuzz-indexing.rs index 23c4cb9c2..0632b7846 100644 --- a/crates/fuzzers/src/bin/fuzz-indexing.rs +++ b/crates/fuzzers/src/bin/fuzz-indexing.rs @@ -144,7 +144,7 @@ fn main() { embedders, &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); diff --git a/crates/index-scheduler/src/scheduler/process_index_operation.rs b/crates/index-scheduler/src/scheduler/process_index_operation.rs index b5338e511..14b07aea0 100644 --- a/crates/index-scheduler/src/scheduler/process_index_operation.rs +++ b/crates/index-scheduler/src/scheduler/process_index_operation.rs @@ -35,7 +35,7 @@ impl IndexScheduler { index: &'i Index, operation: IndexOperation, progress: &Progress, - embedder_stats: Arc, + embedder_stats: Arc, // Cant change ) -> Result<(Vec, Option)> { let indexer_alloc = Bump::new(); let started_processing_at = std::time::Instant::now(); @@ -180,7 +180,7 @@ impl IndexScheduler { embedders, &|| must_stop_processing.get(), progress, - embedder_stats, + &embedder_stats, ) .map_err(|e| Error::from_milli(e, Some(index_uid.clone())))?, ); @@ -292,7 +292,7 @@ impl IndexScheduler { embedders, &|| must_stop_processing.get(), progress, - embedder_stats, + &embedder_stats, ) .map_err(|err| Error::from_milli(err, Some(index_uid.clone())))?, ); @@ -441,7 +441,7 @@ impl IndexScheduler { embedders, &|| must_stop_processing.get(), progress, - embedder_stats, + &embedder_stats, ) .map_err(|err| Error::from_milli(err, Some(index_uid.clone())))?, ); @@ -478,7 +478,7 @@ impl IndexScheduler { .execute( |indexing_step| tracing::debug!(update = ?indexing_step), || must_stop_processing.get(), - embedder_stats, + embedder_stats.clone(), ) .map_err(|err| Error::from_milli(err, Some(index_uid.clone())))?; diff --git a/crates/milli/src/search/new/tests/integration.rs b/crates/milli/src/search/new/tests/integration.rs index c4e521a88..36917c10e 100644 --- a/crates/milli/src/search/new/tests/integration.rs +++ b/crates/milli/src/search/new/tests/integration.rs @@ -95,7 +95,7 @@ pub fn setup_search_index_with_criteria(criteria: &[Criterion]) -> Index { embedders, &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); diff --git a/crates/milli/src/test_index.rs b/crates/milli/src/test_index.rs index 3546660b0..d218bb3a6 100644 --- a/crates/milli/src/test_index.rs +++ b/crates/milli/src/test_index.rs @@ -103,7 +103,7 @@ impl TempIndex { embedders, &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) }) .unwrap()?; @@ -186,7 +186,7 @@ impl TempIndex { embedders, &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) }) .unwrap()?; @@ -261,7 +261,7 @@ fn aborting_indexation() { embedders, &|| should_abort.load(Relaxed), &Progress::default(), - Default::default(), + &Default::default(), ) }) .unwrap() diff --git a/crates/milli/src/update/index_documents/extract/extract_vector_points.rs b/crates/milli/src/update/index_documents/extract/extract_vector_points.rs index e6d874a69..e1981a615 100644 --- a/crates/milli/src/update/index_documents/extract/extract_vector_points.rs +++ b/crates/milli/src/update/index_documents/extract/extract_vector_points.rs @@ -684,7 +684,7 @@ pub fn extract_embeddings( embedder: Arc, embedder_name: &str, possible_embedding_mistakes: &PossibleEmbeddingMistakes, - embedder_stats: Arc, + embedder_stats: &EmbedderStats, unused_vectors_distribution: &UnusedVectorsDistribution, request_threads: &ThreadPoolNoAbort, ) -> Result>> { @@ -727,7 +727,7 @@ pub fn extract_embeddings( std::mem::replace(&mut chunks, Vec::with_capacity(n_chunks)), embedder_name, possible_embedding_mistakes, - embedder_stats.clone(), + embedder_stats, unused_vectors_distribution, request_threads, )?; @@ -750,7 +750,7 @@ pub fn extract_embeddings( std::mem::take(&mut chunks), embedder_name, possible_embedding_mistakes, - embedder_stats.clone(), + embedder_stats, unused_vectors_distribution, request_threads, )?; @@ -789,7 +789,7 @@ fn embed_chunks( text_chunks: Vec>, embedder_name: &str, possible_embedding_mistakes: &PossibleEmbeddingMistakes, - embedder_stats: Arc, + embedder_stats: &EmbedderStats, unused_vectors_distribution: &UnusedVectorsDistribution, request_threads: &ThreadPoolNoAbort, ) -> Result>> { diff --git a/crates/milli/src/update/index_documents/extract/mod.rs b/crates/milli/src/update/index_documents/extract/mod.rs index 1eeddcccb..3af665c67 100644 --- a/crates/milli/src/update/index_documents/extract/mod.rs +++ b/crates/milli/src/update/index_documents/extract/mod.rs @@ -50,7 +50,7 @@ pub(crate) fn data_from_obkv_documents( settings_diff: Arc, max_positions_per_attributes: Option, possible_embedding_mistakes: Arc, - embedder_stats: Arc, + embedder_stats: Arc, // Cant change ) -> Result<()> { let (original_pipeline_result, flattened_pipeline_result): (Result<_>, Result<_>) = rayon::join( || { @@ -234,7 +234,7 @@ fn send_original_documents_data( embedders_configs: Arc>, settings_diff: Arc, possible_embedding_mistakes: Arc, - embedder_stats: Arc, + embedder_stats: Arc, // Cant change ) -> Result<()> { let original_documents_chunk = original_documents_chunk.and_then(|c| unsafe { as_cloneable_grenad(&c) })?; @@ -274,7 +274,7 @@ fn send_original_documents_data( embedder.clone(), &embedder_name, &possible_embedding_mistakes, - embedder_stats.clone(), + &embedder_stats, &unused_vectors_distribution, request_threads(), ) { diff --git a/crates/milli/src/update/index_documents/mod.rs b/crates/milli/src/update/index_documents/mod.rs index f2e1783e4..2bddf1b17 100644 --- a/crates/milli/src/update/index_documents/mod.rs +++ b/crates/milli/src/update/index_documents/mod.rs @@ -81,7 +81,7 @@ pub struct IndexDocuments<'t, 'i, 'a, FP, FA> { added_documents: u64, deleted_documents: u64, embedders: EmbeddingConfigs, - embedder_stats: Arc, + embedder_stats: Arc, // Cant change } #[derive(Default, Debug, Clone)] @@ -104,7 +104,7 @@ where config: IndexDocumentsConfig, progress: FP, should_abort: FA, - embedder_stats: Arc, + embedder_stats: Arc, // Cant change ) -> Result> { let transform = Some(Transform::new( wtxn, @@ -2030,7 +2030,7 @@ mod tests { EmbeddingConfigs::default(), &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); wtxn.commit().unwrap(); @@ -2118,7 +2118,7 @@ mod tests { EmbeddingConfigs::default(), &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); wtxn.commit().unwrap(); @@ -2304,7 +2304,7 @@ mod tests { embedders, &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); wtxn.commit().unwrap(); @@ -2367,7 +2367,7 @@ mod tests { embedders, &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); wtxn.commit().unwrap(); @@ -2421,7 +2421,7 @@ mod tests { embedders, &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); wtxn.commit().unwrap(); @@ -2474,7 +2474,7 @@ mod tests { embedders, &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); wtxn.commit().unwrap(); @@ -2529,7 +2529,7 @@ mod tests { embedders, &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); wtxn.commit().unwrap(); @@ -2589,7 +2589,7 @@ mod tests { embedders, &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); wtxn.commit().unwrap(); @@ -2642,7 +2642,7 @@ mod tests { embedders, &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); wtxn.commit().unwrap(); @@ -2695,7 +2695,7 @@ mod tests { embedders, &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); wtxn.commit().unwrap(); @@ -2894,7 +2894,7 @@ mod tests { embedders, &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); wtxn.commit().unwrap(); @@ -2954,7 +2954,7 @@ mod tests { embedders, &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); wtxn.commit().unwrap(); @@ -3011,7 +3011,7 @@ mod tests { embedders, &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); wtxn.commit().unwrap(); diff --git a/crates/milli/src/update/new/extract/vectors/mod.rs b/crates/milli/src/update/new/extract/vectors/mod.rs index c21dabf74..85398aa99 100644 --- a/crates/milli/src/update/new/extract/vectors/mod.rs +++ b/crates/milli/src/update/new/extract/vectors/mod.rs @@ -1,4 +1,4 @@ -use std::{cell::RefCell, sync::Arc}; +use std::cell::RefCell; use bumpalo::collections::Vec as BVec; use bumpalo::Bump; @@ -23,7 +23,7 @@ pub struct EmbeddingExtractor<'a, 'b> { embedders: &'a EmbeddingConfigs, sender: EmbeddingSender<'a, 'b>, possible_embedding_mistakes: PossibleEmbeddingMistakes, - embedder_stats: Arc, + embedder_stats: &'a EmbedderStats, threads: &'a ThreadPoolNoAbort, } @@ -32,7 +32,7 @@ impl<'a, 'b> EmbeddingExtractor<'a, 'b> { embedders: &'a EmbeddingConfigs, sender: EmbeddingSender<'a, 'b>, field_distribution: &'a FieldDistribution, - embedder_stats: Arc, + embedder_stats: &'a EmbedderStats, threads: &'a ThreadPoolNoAbort, ) -> Self { let possible_embedding_mistakes = PossibleEmbeddingMistakes::new(field_distribution); @@ -78,7 +78,7 @@ impl<'extractor> Extractor<'extractor> for EmbeddingExtractor<'_, '_> { prompt, context.data, &self.possible_embedding_mistakes, - self.embedder_stats.clone(), + self.embedder_stats, self.threads, self.sender, &context.doc_alloc, @@ -311,7 +311,7 @@ struct Chunks<'a, 'b, 'extractor> { dimensions: usize, prompt: &'a Prompt, possible_embedding_mistakes: &'a PossibleEmbeddingMistakes, - embedder_stats: Arc, + embedder_stats: &'a EmbedderStats, user_provided: &'a RefCell>, threads: &'a ThreadPoolNoAbort, sender: EmbeddingSender<'a, 'b>, @@ -327,7 +327,7 @@ impl<'a, 'b, 'extractor> Chunks<'a, 'b, 'extractor> { prompt: &'a Prompt, user_provided: &'a RefCell>, possible_embedding_mistakes: &'a PossibleEmbeddingMistakes, - embedder_stats: Arc, + embedder_stats: &'a EmbedderStats, threads: &'a ThreadPoolNoAbort, sender: EmbeddingSender<'a, 'b>, doc_alloc: &'a Bump, @@ -378,7 +378,7 @@ impl<'a, 'b, 'extractor> Chunks<'a, 'b, 'extractor> { self.embedder_id, self.embedder_name, self.possible_embedding_mistakes, - self.embedder_stats.clone(), + self.embedder_stats, unused_vectors_distribution, self.threads, self.sender, @@ -397,7 +397,7 @@ impl<'a, 'b, 'extractor> Chunks<'a, 'b, 'extractor> { self.embedder_id, self.embedder_name, self.possible_embedding_mistakes, - self.embedder_stats.clone(), + self.embedder_stats, unused_vectors_distribution, self.threads, self.sender, @@ -416,7 +416,7 @@ impl<'a, 'b, 'extractor> Chunks<'a, 'b, 'extractor> { embedder_id: u8, embedder_name: &str, possible_embedding_mistakes: &PossibleEmbeddingMistakes, - embedder_stats: Arc, + embedder_stats: &EmbedderStats, unused_vectors_distribution: &UnusedVectorsDistributionBump, threads: &ThreadPoolNoAbort, sender: EmbeddingSender<'a, 'b>, diff --git a/crates/milli/src/update/new/indexer/extract.rs b/crates/milli/src/update/new/indexer/extract.rs index c721a2563..97ffc8624 100644 --- a/crates/milli/src/update/new/indexer/extract.rs +++ b/crates/milli/src/update/new/indexer/extract.rs @@ -1,6 +1,5 @@ use std::collections::BTreeMap; use std::sync::atomic::AtomicBool; -use std::sync::Arc; use std::sync::OnceLock; use bumpalo::Bump; @@ -36,7 +35,7 @@ pub(super) fn extract_all<'pl, 'extractor, DC, MSP>( mut index_embeddings: Vec, document_ids: &mut RoaringBitmap, modified_docids: &mut RoaringBitmap, - embedder_stats: Arc, + embedder_stats: &EmbedderStats, ) -> Result<(FacetFieldIdsDelta, Vec)> where DC: DocumentChanges<'pl>, diff --git a/crates/milli/src/update/new/indexer/mod.rs b/crates/milli/src/update/new/indexer/mod.rs index 33774f892..bb6ba0102 100644 --- a/crates/milli/src/update/new/indexer/mod.rs +++ b/crates/milli/src/update/new/indexer/mod.rs @@ -1,5 +1,4 @@ use std::sync::atomic::AtomicBool; -use std::sync::Arc; use std::sync::{Once, RwLock}; use std::thread::{self, Builder}; @@ -56,7 +55,7 @@ pub fn index<'pl, 'indexer, 'index, DC, MSP>( embedders: EmbeddingConfigs, must_stop_processing: &'indexer MSP, progress: &'indexer Progress, - embedder_stats: Arc, + embedder_stats: &'indexer EmbedderStats, ) -> Result where DC: DocumentChanges<'pl>, diff --git a/crates/milli/src/update/settings.rs b/crates/milli/src/update/settings.rs index b3f70d1b6..71cedf456 100644 --- a/crates/milli/src/update/settings.rs +++ b/crates/milli/src/update/settings.rs @@ -475,7 +475,7 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> { progress_callback: &FP, should_abort: &FA, settings_diff: InnerIndexSettingsDiff, - embedder_stats: Arc, + embedder_stats: Arc, // Cant change ) -> Result<()> where FP: Fn(UpdateIndexingStep) + Sync, @@ -1362,7 +1362,7 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> { mut self, progress_callback: FP, should_abort: FA, - embedder_stats: Arc, + embedder_stats: Arc, // Cant change ) -> Result<()> where FP: Fn(UpdateIndexingStep) + Sync, diff --git a/crates/milli/src/vector/composite.rs b/crates/milli/src/vector/composite.rs index 87f05d4fe..8314b8649 100644 --- a/crates/milli/src/vector/composite.rs +++ b/crates/milli/src/vector/composite.rs @@ -1,4 +1,3 @@ -use std::sync::Arc; use std::time::Instant; use arroy::Distance; @@ -154,7 +153,7 @@ impl SubEmbedder { &self, texts: Vec, deadline: Option, - embedder_stats: Option>, + embedder_stats: Option<&EmbedderStats>, ) -> std::result::Result, EmbedError> { match self { SubEmbedder::HuggingFace(embedder) => embedder.embed(texts), @@ -169,7 +168,7 @@ impl SubEmbedder { &self, text: &str, deadline: Option, - embedder_stats: Option>, + embedder_stats: Option<&EmbedderStats>, ) -> std::result::Result { match self { SubEmbedder::HuggingFace(embedder) => embedder.embed_one(text), @@ -196,7 +195,7 @@ impl SubEmbedder { &self, text_chunks: Vec>, threads: &ThreadPoolNoAbort, - embedder_stats: Arc, + embedder_stats: &EmbedderStats, ) -> std::result::Result>, EmbedError> { match self { SubEmbedder::HuggingFace(embedder) => embedder.embed_index(text_chunks), @@ -218,7 +217,7 @@ impl SubEmbedder { &self, texts: &[&str], threads: &ThreadPoolNoAbort, - embedder_stats: Arc, + embedder_stats: &EmbedderStats, ) -> std::result::Result, EmbedError> { match self { SubEmbedder::HuggingFace(embedder) => embedder.embed_index_ref(texts), diff --git a/crates/milli/src/vector/mod.rs b/crates/milli/src/vector/mod.rs index 481eb6c99..065beb5fb 100644 --- a/crates/milli/src/vector/mod.rs +++ b/crates/milli/src/vector/mod.rs @@ -749,7 +749,7 @@ impl Embedder { &self, text_chunks: Vec>, threads: &ThreadPoolNoAbort, - embedder_stats: Arc, + embedder_stats: &EmbedderStats, ) -> std::result::Result>, EmbedError> { match self { Embedder::HuggingFace(embedder) => embedder.embed_index(text_chunks), @@ -772,7 +772,7 @@ impl Embedder { &self, texts: &[&str], threads: &ThreadPoolNoAbort, - embedder_stats: Arc, + embedder_stats: &EmbedderStats, ) -> std::result::Result, EmbedError> { match self { Embedder::HuggingFace(embedder) => embedder.embed_index_ref(texts), diff --git a/crates/milli/src/vector/ollama.rs b/crates/milli/src/vector/ollama.rs index 045b65b72..d4329a2de 100644 --- a/crates/milli/src/vector/ollama.rs +++ b/crates/milli/src/vector/ollama.rs @@ -1,4 +1,3 @@ -use std::sync::Arc; use std::time::Instant; use rayon::iter::{IntoParallelIterator as _, ParallelIterator as _}; @@ -106,7 +105,7 @@ impl Embedder { &self, texts: &[S], deadline: Option, - embedder_stats: Option>, + embedder_stats: Option<&EmbedderStats>, ) -> Result, EmbedError> { match self.rest_embedder.embed_ref(texts, deadline, embedder_stats) { Ok(embeddings) => Ok(embeddings), @@ -121,21 +120,21 @@ impl Embedder { &self, text_chunks: Vec>, threads: &ThreadPoolNoAbort, - embedder_stats: Arc, + embedder_stats: &EmbedderStats, ) -> Result>, EmbedError> { // This condition helps reduce the number of active rayon jobs // so that we avoid consuming all the LMDB rtxns and avoid stack overflows. if threads.active_operations() >= REQUEST_PARALLELISM { text_chunks .into_iter() - .map(move |chunk| self.embed(&chunk, None, Some(embedder_stats.clone()))) + .map(move |chunk| self.embed(&chunk, None, Some(embedder_stats))) .collect() } else { threads .install(move || { text_chunks .into_par_iter() - .map(move |chunk| self.embed(&chunk, None, Some(embedder_stats.clone()))) + .map(move |chunk| self.embed(&chunk, None, Some(embedder_stats))) .collect() }) .map_err(|error| EmbedError { @@ -149,14 +148,14 @@ impl Embedder { &self, texts: &[&str], threads: &ThreadPoolNoAbort, - embedder_stats: Arc, + embedder_stats: &EmbedderStats, ) -> Result>, EmbedError> { // This condition helps reduce the number of active rayon jobs // so that we avoid consuming all the LMDB rtxns and avoid stack overflows. if threads.active_operations() >= REQUEST_PARALLELISM { let embeddings: Result>, _> = texts .chunks(self.prompt_count_in_chunk_hint()) - .map(move |chunk| self.embed(chunk, None, Some(embedder_stats.clone()))) + .map(move |chunk| self.embed(chunk, None, Some(embedder_stats))) .collect(); let embeddings = embeddings?; @@ -166,7 +165,7 @@ impl Embedder { .install(move || { let embeddings: Result>, _> = texts .par_chunks(self.prompt_count_in_chunk_hint()) - .map(move |chunk| self.embed(chunk, None, Some(embedder_stats.clone()))) + .map(move |chunk| self.embed(chunk, None, Some(embedder_stats))) .collect(); let embeddings = embeddings?; diff --git a/crates/milli/src/vector/openai.rs b/crates/milli/src/vector/openai.rs index b64e3d467..0159d5c76 100644 --- a/crates/milli/src/vector/openai.rs +++ b/crates/milli/src/vector/openai.rs @@ -1,5 +1,4 @@ use std::fmt; -use std::sync::Arc; use std::time::Instant; use ordered_float::OrderedFloat; @@ -217,7 +216,7 @@ impl Embedder { &self, texts: &[S], deadline: Option, - embedder_stats: Option>, + embedder_stats: Option<&EmbedderStats>, ) -> Result, EmbedError> { match self.rest_embedder.embed_ref(texts, deadline, embedder_stats) { Ok(embeddings) => Ok(embeddings), @@ -262,21 +261,21 @@ impl Embedder { &self, text_chunks: Vec>, threads: &ThreadPoolNoAbort, - embedder_stats: Arc, + embedder_stats: &EmbedderStats, ) -> Result>, EmbedError> { // This condition helps reduce the number of active rayon jobs // so that we avoid consuming all the LMDB rtxns and avoid stack overflows. if threads.active_operations() >= REQUEST_PARALLELISM { text_chunks .into_iter() - .map(move |chunk| self.embed(&chunk, None, Some(embedder_stats.clone()))) + .map(move |chunk| self.embed(&chunk, None, Some(embedder_stats))) .collect() } else { threads .install(move || { text_chunks .into_par_iter() - .map(move |chunk| self.embed(&chunk, None, Some(embedder_stats.clone()))) + .map(move |chunk| self.embed(&chunk, None, Some(embedder_stats))) .collect() }) .map_err(|error| EmbedError { @@ -290,14 +289,14 @@ impl Embedder { &self, texts: &[&str], threads: &ThreadPoolNoAbort, - embedder_stats: Arc, + embedder_stats: &EmbedderStats, ) -> Result>, EmbedError> { // This condition helps reduce the number of active rayon jobs // so that we avoid consuming all the LMDB rtxns and avoid stack overflows. if threads.active_operations() >= REQUEST_PARALLELISM { let embeddings: Result>, _> = texts .chunks(self.prompt_count_in_chunk_hint()) - .map(move |chunk| self.embed(chunk, None, Some(embedder_stats.clone()))) + .map(move |chunk| self.embed(chunk, None, Some(embedder_stats))) .collect(); let embeddings = embeddings?; Ok(embeddings.into_iter().flatten().collect()) @@ -306,7 +305,7 @@ impl Embedder { .install(move || { let embeddings: Result>, _> = texts .par_chunks(self.prompt_count_in_chunk_hint()) - .map(move |chunk| self.embed(chunk, None, Some(embedder_stats.clone()))) + .map(move |chunk| self.embed(chunk, None, Some(embedder_stats))) .collect(); let embeddings = embeddings?; diff --git a/crates/milli/src/vector/rest.rs b/crates/milli/src/vector/rest.rs index dd08c6a5e..fbe3c1129 100644 --- a/crates/milli/src/vector/rest.rs +++ b/crates/milli/src/vector/rest.rs @@ -1,5 +1,4 @@ use std::collections::BTreeMap; -use std::sync::Arc; use std::time::Instant; use deserr::Deserr; @@ -170,7 +169,7 @@ impl Embedder { &self, texts: Vec, deadline: Option, - embedder_stats: Option>, + embedder_stats: Option<&EmbedderStats>, ) -> Result, EmbedError> { embed( &self.data, @@ -186,7 +185,7 @@ impl Embedder { &self, texts: &[S], deadline: Option, - embedder_stats: Option>, + embedder_stats: Option<&EmbedderStats>, ) -> Result, EmbedError> where S: AsRef + Serialize, @@ -208,21 +207,21 @@ impl Embedder { &self, text_chunks: Vec>, threads: &ThreadPoolNoAbort, - embedder_stats: Arc, + embedder_stats: &EmbedderStats, ) -> Result>, EmbedError> { // This condition helps reduce the number of active rayon jobs // so that we avoid consuming all the LMDB rtxns and avoid stack overflows. if threads.active_operations() >= REQUEST_PARALLELISM { text_chunks .into_iter() - .map(move |chunk| self.embed(chunk, None, Some(embedder_stats.clone()))) + .map(move |chunk| self.embed(chunk, None, Some(embedder_stats))) .collect() } else { threads .install(move || { text_chunks .into_par_iter() - .map(move |chunk| self.embed(chunk, None, Some(embedder_stats.clone()))) + .map(move |chunk| self.embed(chunk, None, Some(embedder_stats))) .collect() }) .map_err(|error| EmbedError { @@ -236,14 +235,14 @@ impl Embedder { &self, texts: &[&str], threads: &ThreadPoolNoAbort, - embedder_stats: Arc, + embedder_stats: &EmbedderStats, ) -> Result, EmbedError> { // This condition helps reduce the number of active rayon jobs // so that we avoid consuming all the LMDB rtxns and avoid stack overflows. if threads.active_operations() >= REQUEST_PARALLELISM { let embeddings: Result>, _> = texts .chunks(self.prompt_count_in_chunk_hint()) - .map(move |chunk| self.embed_ref(chunk, None, Some(embedder_stats.clone()))) + .map(move |chunk| self.embed_ref(chunk, None, Some(embedder_stats))) .collect(); let embeddings = embeddings?; @@ -253,7 +252,7 @@ impl Embedder { .install(move || { let embeddings: Result>, _> = texts .par_chunks(self.prompt_count_in_chunk_hint()) - .map(move |chunk| self.embed_ref(chunk, None, Some(embedder_stats.clone()))) + .map(move |chunk| self.embed_ref(chunk, None, Some(embedder_stats))) .collect(); let embeddings = embeddings?; @@ -303,7 +302,7 @@ fn embed( expected_count: usize, expected_dimension: Option, deadline: Option, - embedder_stats: Option>, + embedder_stats: Option<&EmbedderStats>, ) -> Result, EmbedError> where S: Serialize, @@ -323,7 +322,7 @@ where for attempt in 0..10 { if let Some(embedder_stats) = &embedder_stats { - embedder_stats.as_ref().total_count.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + embedder_stats.total_count.fetch_add(1, std::sync::atomic::Ordering::Relaxed); } let response = request.clone().send_json(&body); let result = check_response(response, data.configuration_source).and_then(|response| { @@ -367,7 +366,7 @@ where } if let Some(embedder_stats) = &embedder_stats { - embedder_stats.as_ref().total_count.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + embedder_stats.total_count.fetch_add(1, std::sync::atomic::Ordering::Relaxed); } let response = request.send_json(&body); let result = check_response(response, data.configuration_source).and_then(|response| { diff --git a/crates/milli/tests/search/facet_distribution.rs b/crates/milli/tests/search/facet_distribution.rs index 5ed223400..8548f0d01 100644 --- a/crates/milli/tests/search/facet_distribution.rs +++ b/crates/milli/tests/search/facet_distribution.rs @@ -74,7 +74,7 @@ fn test_facet_distribution_with_no_facet_values() { embedders, &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); diff --git a/crates/milli/tests/search/mod.rs b/crates/milli/tests/search/mod.rs index beee4ac54..4098af736 100644 --- a/crates/milli/tests/search/mod.rs +++ b/crates/milli/tests/search/mod.rs @@ -114,7 +114,7 @@ pub fn setup_search_index_with_criteria(criteria: &[Criterion]) -> Index { embedders, &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); diff --git a/crates/milli/tests/search/query_criteria.rs b/crates/milli/tests/search/query_criteria.rs index 04b8374de..b72978330 100644 --- a/crates/milli/tests/search/query_criteria.rs +++ b/crates/milli/tests/search/query_criteria.rs @@ -344,7 +344,7 @@ fn criteria_ascdesc() { embedders, &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); diff --git a/crates/milli/tests/search/typo_tolerance.rs b/crates/milli/tests/search/typo_tolerance.rs index e2cdab550..9aacbf82a 100644 --- a/crates/milli/tests/search/typo_tolerance.rs +++ b/crates/milli/tests/search/typo_tolerance.rs @@ -153,7 +153,7 @@ fn test_typo_disabled_on_word() { embedders, &|| false, &Progress::default(), - Default::default(), + &Default::default(), ) .unwrap(); From 0f6dd133b2bcb54d1489fa748f8d8dedadca4125 Mon Sep 17 00:00:00 2001 From: Mubelotix Date: Thu, 26 Jun 2025 12:11:43 +0200 Subject: [PATCH 15/18] Turn to references --- crates/meilisearch/src/lib.rs | 2 +- crates/milli/src/update/index_documents/extract/mod.rs | 4 ++-- crates/milli/src/update/index_documents/mod.rs | 6 +++--- crates/milli/src/update/settings.rs | 6 +++--- 4 files changed, 9 insertions(+), 9 deletions(-) diff --git a/crates/meilisearch/src/lib.rs b/crates/meilisearch/src/lib.rs index 5acfb4bc9..c902f4e60 100644 --- a/crates/meilisearch/src/lib.rs +++ b/crates/meilisearch/src/lib.rs @@ -579,7 +579,7 @@ fn import_dump( }, |indexing_step| tracing::trace!("update: {:?}", indexing_step), || false, - embedder_stats, + &embedder_stats, )?; let builder = builder.with_embedders(embedders); diff --git a/crates/milli/src/update/index_documents/extract/mod.rs b/crates/milli/src/update/index_documents/extract/mod.rs index 3af665c67..9c1971356 100644 --- a/crates/milli/src/update/index_documents/extract/mod.rs +++ b/crates/milli/src/update/index_documents/extract/mod.rs @@ -50,7 +50,7 @@ pub(crate) fn data_from_obkv_documents( settings_diff: Arc, max_positions_per_attributes: Option, possible_embedding_mistakes: Arc, - embedder_stats: Arc, // Cant change + embedder_stats: &Arc, ) -> Result<()> { let (original_pipeline_result, flattened_pipeline_result): (Result<_>, Result<_>) = rayon::join( || { @@ -234,7 +234,7 @@ fn send_original_documents_data( embedders_configs: Arc>, settings_diff: Arc, possible_embedding_mistakes: Arc, - embedder_stats: Arc, // Cant change + embedder_stats: Arc, ) -> Result<()> { let original_documents_chunk = original_documents_chunk.and_then(|c| unsafe { as_cloneable_grenad(&c) })?; diff --git a/crates/milli/src/update/index_documents/mod.rs b/crates/milli/src/update/index_documents/mod.rs index 2bddf1b17..6e56ad155 100644 --- a/crates/milli/src/update/index_documents/mod.rs +++ b/crates/milli/src/update/index_documents/mod.rs @@ -81,7 +81,7 @@ pub struct IndexDocuments<'t, 'i, 'a, FP, FA> { added_documents: u64, deleted_documents: u64, embedders: EmbeddingConfigs, - embedder_stats: Arc, // Cant change + embedder_stats: &'t Arc, } #[derive(Default, Debug, Clone)] @@ -104,7 +104,7 @@ where config: IndexDocumentsConfig, progress: FP, should_abort: FA, - embedder_stats: Arc, // Cant change + embedder_stats: &'t Arc, ) -> Result> { let transform = Some(Transform::new( wtxn, @@ -331,7 +331,7 @@ where settings_diff_cloned, max_positions_per_attributes, Arc::new(possible_embedding_mistakes), - embedder_stats.clone() + &embedder_stats ) }); diff --git a/crates/milli/src/update/settings.rs b/crates/milli/src/update/settings.rs index 71cedf456..06c2d0cc2 100644 --- a/crates/milli/src/update/settings.rs +++ b/crates/milli/src/update/settings.rs @@ -475,7 +475,7 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> { progress_callback: &FP, should_abort: &FA, settings_diff: InnerIndexSettingsDiff, - embedder_stats: Arc, // Cant change + embedder_stats: &Arc, // Cant change ) -> Result<()> where FP: Fn(UpdateIndexingStep) + Sync, @@ -507,7 +507,7 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> { IndexDocumentsConfig::default(), &progress_callback, &should_abort, - embedder_stats, + &embedder_stats, )?; indexing_builder.execute_raw(output)?; @@ -1421,7 +1421,7 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> { ); if inner_settings_diff.any_reindexing_needed() { - self.reindex(&progress_callback, &should_abort, inner_settings_diff, embedder_stats)?; + self.reindex(&progress_callback, &should_abort, inner_settings_diff, &embedder_stats)?; } Ok(()) From 2ff382c023c25d65fb255bc388223789b395be0a Mon Sep 17 00:00:00 2001 From: Mubelotix Date: Thu, 26 Jun 2025 12:14:56 +0200 Subject: [PATCH 16/18] Remove useless clone --- crates/index-scheduler/src/scheduler/process_index_operation.rs | 2 +- crates/milli/src/update/settings.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/index-scheduler/src/scheduler/process_index_operation.rs b/crates/index-scheduler/src/scheduler/process_index_operation.rs index 14b07aea0..84554849f 100644 --- a/crates/index-scheduler/src/scheduler/process_index_operation.rs +++ b/crates/index-scheduler/src/scheduler/process_index_operation.rs @@ -478,7 +478,7 @@ impl IndexScheduler { .execute( |indexing_step| tracing::debug!(update = ?indexing_step), || must_stop_processing.get(), - embedder_stats.clone(), + embedder_stats, ) .map_err(|err| Error::from_milli(err, Some(index_uid.clone())))?; diff --git a/crates/milli/src/update/settings.rs b/crates/milli/src/update/settings.rs index 06c2d0cc2..99736f971 100644 --- a/crates/milli/src/update/settings.rs +++ b/crates/milli/src/update/settings.rs @@ -507,7 +507,7 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> { IndexDocumentsConfig::default(), &progress_callback, &should_abort, - &embedder_stats, + embedder_stats, )?; indexing_builder.execute_raw(output)?; From 4d26e9c6f2b64f0b1f5afbeafdad242271cecc69 Mon Sep 17 00:00:00 2001 From: Mubelotix Date: Thu, 26 Jun 2025 12:21:34 +0200 Subject: [PATCH 17/18] Remove my comments --- .../index-scheduler/src/scheduler/process_index_operation.rs | 2 +- crates/milli/src/update/settings.rs | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/crates/index-scheduler/src/scheduler/process_index_operation.rs b/crates/index-scheduler/src/scheduler/process_index_operation.rs index 84554849f..4c0db9ce4 100644 --- a/crates/index-scheduler/src/scheduler/process_index_operation.rs +++ b/crates/index-scheduler/src/scheduler/process_index_operation.rs @@ -35,7 +35,7 @@ impl IndexScheduler { index: &'i Index, operation: IndexOperation, progress: &Progress, - embedder_stats: Arc, // Cant change + embedder_stats: Arc, ) -> Result<(Vec, Option)> { let indexer_alloc = Bump::new(); let started_processing_at = std::time::Instant::now(); diff --git a/crates/milli/src/update/settings.rs b/crates/milli/src/update/settings.rs index 99736f971..05dbb4784 100644 --- a/crates/milli/src/update/settings.rs +++ b/crates/milli/src/update/settings.rs @@ -475,7 +475,7 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> { progress_callback: &FP, should_abort: &FA, settings_diff: InnerIndexSettingsDiff, - embedder_stats: &Arc, // Cant change + embedder_stats: &Arc, ) -> Result<()> where FP: Fn(UpdateIndexingStep) + Sync, @@ -1362,7 +1362,7 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> { mut self, progress_callback: FP, should_abort: FA, - embedder_stats: Arc, // Cant change + embedder_stats: Arc, ) -> Result<()> where FP: Fn(UpdateIndexingStep) + Sync, From 44d6430bae887c11bc9f866684bf857204137d57 Mon Sep 17 00:00:00 2001 From: Mubelotix Date: Thu, 26 Jun 2025 12:30:08 +0200 Subject: [PATCH 18/18] Rename fields --- crates/meilisearch-types/src/batch_view.rs | 4 ++-- crates/meilisearch-types/src/batches.rs | 10 +++++----- crates/meilisearch/tests/vector/rest.rs | 8 ++++---- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/crates/meilisearch-types/src/batch_view.rs b/crates/meilisearch-types/src/batch_view.rs index ea027b74e..297b10ba1 100644 --- a/crates/meilisearch-types/src/batch_view.rs +++ b/crates/meilisearch-types/src/batch_view.rs @@ -32,7 +32,7 @@ pub struct BatchStatsView { #[serde(flatten)] pub stats: BatchStats, #[serde(skip_serializing_if = "EmbedderStatsView::skip_serializing", default)] - pub embedder: EmbedderStatsView, + pub embedder_requests: EmbedderStatsView, } impl BatchView { @@ -43,7 +43,7 @@ impl BatchView { details: batch.details.clone(), stats: BatchStatsView { stats: batch.stats.clone(), - embedder: batch.embedder_stats.clone(), + embedder_requests: batch.embedder_stats.clone(), }, duration: batch.finished_at.map(|finished_at| finished_at - batch.started_at), started_at: batch.started_at, diff --git a/crates/meilisearch-types/src/batches.rs b/crates/meilisearch-types/src/batches.rs index c8d98655f..e1cc2b7c7 100644 --- a/crates/meilisearch-types/src/batches.rs +++ b/crates/meilisearch-types/src/batches.rs @@ -92,8 +92,8 @@ pub struct BatchStats { #[serde(rename_all = "camelCase")] #[schema(rename_all = "camelCase")] pub struct EmbedderStatsView { - pub total_count: usize, - pub error_count: usize, + pub total: usize, + pub failed: usize, #[serde(skip_serializing_if = "Option::is_none", default)] pub last_error: Option, } @@ -102,8 +102,8 @@ impl From<&EmbedderStats> for EmbedderStatsView { fn from(stats: &EmbedderStats) -> Self { let errors = stats.errors.read().unwrap_or_else(|p| p.into_inner()); Self { - total_count: stats.total_count.load(std::sync::atomic::Ordering::Relaxed), - error_count: errors.1 as usize, + total: stats.total_count.load(std::sync::atomic::Ordering::Relaxed), + failed: errors.1 as usize, last_error: errors.0.clone(), } } @@ -111,6 +111,6 @@ impl From<&EmbedderStats> for EmbedderStatsView { impl EmbedderStatsView { pub fn skip_serializing(&self) -> bool { - self.total_count == 0 && self.error_count == 0 && self.last_error.is_none() + self.total == 0 && self.failed == 0 && self.last_error.is_none() } } diff --git a/crates/meilisearch/tests/vector/rest.rs b/crates/meilisearch/tests/vector/rest.rs index e80dfeb0a..6e781e525 100644 --- a/crates/meilisearch/tests/vector/rest.rs +++ b/crates/meilisearch/tests/vector/rest.rs @@ -2182,7 +2182,7 @@ async fn last_error_stats() { let (response, _code) = index.filtered_batches(&[], &[], &[]).await; snapshot!(json_string!(response["results"][0], { ".progress" => "[ignored]", - ".stats.embedder.totalCount" => "[ignored]", + ".stats.embedderRequests.total" => "[ignored]", ".startedAt" => "[ignored]" }), @r#" { @@ -2203,9 +2203,9 @@ async fn last_error_stats() { "indexUids": { "doggo": 1 }, - "embedder": { - "totalCount": "[ignored]", - "errorCount": 5, + "embedderRequests": { + "total": "[ignored]", + "failed": 5, "lastError": "runtime error: received internal error HTTP 500 from embedding server\n - server replied with `Service Unavailable`" } },