Disable the cache by default and make it experimental

crates/index-scheduler/src/lib.rs

@@ -125,6 +125,10 @@ pub struct IndexSchedulerOptions {
     pub instance_features: InstanceTogglableFeatures,
     /// The experimental features enabled for this instance.
     pub auto_upgrade: bool,
+    /// The maximal number of entries in the search query cache of an embedder.
+    ///
+    /// 0 disables the cache.
+    pub embedding_cache_cap: usize,
 }
 
 /// Structure which holds meilisearch's indexes and schedules the tasks
@@ -156,6 +160,11 @@ pub struct IndexScheduler {
     /// The Authorization header to send to the webhook URL.
     pub(crate) webhook_authorization_header: Option<String>,
 
+    /// A map to retrieve the runtime representation of an embedder depending on its configuration.
+    ///
+    /// This map may return the same embedder object for two different indexes or embedder settings,
+    /// but it will only do this if the embedder configuration options are the same, leading
+    /// to the same embeddings for the same input text.
     embedders: Arc<RwLock<HashMap<EmbedderOptions, Arc<Embedder>>>>,
 
     // ================= test
@@ -818,7 +827,7 @@ impl IndexScheduler {
 
                     // add missing embedder
                     let embedder = Arc::new(
-                        Embedder::new(embedder_options.clone())
+                        Embedder::new(embedder_options.clone(), self.scheduler.embedding_cache_cap)
                             .map_err(meilisearch_types::milli::vector::Error::from)
                             .map_err(|err| {
                                 Error::from_milli(err.into(), Some(index_uid.clone()))

crates/index-scheduler/src/scheduler/mod.rs

@@ -76,6 +76,11 @@ pub struct Scheduler {
 
     /// The path to the version file of Meilisearch.
     pub(crate) version_file_path: PathBuf,
+
+    /// The maximal number of entries in the search query cache of an embedder.
+    ///
+    /// 0 disables the cache.
+    pub(crate) embedding_cache_cap: usize,
 }
 
 impl Scheduler {
@@ -90,6 +95,7 @@ impl Scheduler {
             snapshots_path: self.snapshots_path.clone(),
             auth_env: self.auth_env.clone(),
             version_file_path: self.version_file_path.clone(),
+            embedding_cache_cap: self.embedding_cache_cap,
         }
     }
 
@@ -105,6 +111,7 @@ impl Scheduler {
             snapshots_path: options.snapshots_path.clone(),
             auth_env,
             version_file_path: options.version_file_path.clone(),
+            embedding_cache_cap: options.embedding_cache_cap,
         }
     }
 }

crates/index-scheduler/src/test_utils.rs

@@ -112,6 +112,7 @@ impl IndexScheduler {
             batched_tasks_size_limit: u64::MAX,
             instance_features: Default::default(),
             auto_upgrade: true, // Don't cost much and will ensure the happy path works
+            embedding_cache_cap: 10,
         };
         let version = configuration(&mut options).unwrap_or_else(|| {
             (

crates/meilisearch/src/analytics/segment_analytics.rs

@@ -199,6 +199,7 @@ struct Infos {
     experimental_network: bool,
     experimental_get_task_documents_route: bool,
     experimental_composite_embedders: bool,
+    experimental_embedding_cache_entries: usize,
     gpu_enabled: bool,
     db_path: bool,
     import_dump: bool,
@@ -246,6 +247,7 @@ impl Infos {
             experimental_reduce_indexing_memory_usage,
             experimental_max_number_of_batched_tasks,
             experimental_limit_batched_tasks_total_size,
+            experimental_embedding_cache_entries,
             http_addr,
             master_key: _,
             env,
@@ -312,6 +314,7 @@ impl Infos {
             experimental_network: network,
             experimental_get_task_documents_route: get_task_documents_route,
             experimental_composite_embedders: composite_embedders,
+            experimental_embedding_cache_entries,
             gpu_enabled: meilisearch_types::milli::vector::is_cuda_enabled(),
             db_path: db_path != PathBuf::from("./data.ms"),
             import_dump: import_dump.is_some(),

crates/meilisearch/src/lib.rs

@@ -233,6 +233,7 @@ pub fn setup_meilisearch(opt: &Opt) -> anyhow::Result<(Arc<IndexScheduler>, Arc<
         index_count: DEFAULT_INDEX_COUNT,
         instance_features: opt.to_instance_features(),
         auto_upgrade: opt.experimental_dumpless_upgrade,
+        embedding_cache_cap: opt.experimental_embedding_cache_entries,
     };
     let bin_major: u32 = VERSION_MAJOR.parse().unwrap();
     let bin_minor: u32 = VERSION_MINOR.parse().unwrap();

crates/meilisearch/src/option.rs

@@ -63,7 +63,8 @@ const MEILI_EXPERIMENTAL_MAX_NUMBER_OF_BATCHED_TASKS: &str =
     "MEILI_EXPERIMENTAL_MAX_NUMBER_OF_BATCHED_TASKS";
 const MEILI_EXPERIMENTAL_LIMIT_BATCHED_TASKS_TOTAL_SIZE: &str =
     "MEILI_EXPERIMENTAL_LIMIT_BATCHED_TASKS_SIZE";
-
+const MEILI_EXPERIMENTAL_EMBEDDING_CACHE_ENTRIES: &str =
+    "MEILI_EXPERIMENTAL_EMBEDDING_CACHE_ENTRIES";
 const DEFAULT_CONFIG_FILE_PATH: &str = "./config.toml";
 const DEFAULT_DB_PATH: &str = "./data.ms";
 const DEFAULT_HTTP_ADDR: &str = "localhost:7700";
@@ -446,6 +447,14 @@ pub struct Opt {
     #[serde(default = "default_limit_batched_tasks_total_size")]
     pub experimental_limit_batched_tasks_total_size: u64,
 
+    /// Enables experimental caching of search query embeddings. The value represents the maximal number of entries in the cache of each
+    /// distinct embedder.
+    ///
+    /// For more information, see <https://github.com/orgs/meilisearch/discussions/818>.
+    #[clap(long, env = MEILI_EXPERIMENTAL_EMBEDDING_CACHE_ENTRIES, default_value_t = default_embedding_cache_entries())]
+    #[serde(default = "default_embedding_cache_entries")]
+    pub experimental_embedding_cache_entries: usize,
+
     #[serde(flatten)]
     #[clap(flatten)]
     pub indexer_options: IndexerOpts,
@@ -549,6 +558,7 @@ impl Opt {
             experimental_reduce_indexing_memory_usage,
             experimental_max_number_of_batched_tasks,
             experimental_limit_batched_tasks_total_size,
+            experimental_embedding_cache_entries,
         } = self;
         export_to_env_if_not_present(MEILI_DB_PATH, db_path);
         export_to_env_if_not_present(MEILI_HTTP_ADDR, http_addr);
@@ -641,6 +651,10 @@ impl Opt {
             MEILI_EXPERIMENTAL_LIMIT_BATCHED_TASKS_TOTAL_SIZE,
             experimental_limit_batched_tasks_total_size.to_string(),
         );
+        export_to_env_if_not_present(
+            MEILI_EXPERIMENTAL_EMBEDDING_CACHE_ENTRIES,
+            experimental_embedding_cache_entries.to_string(),
+        );
         indexer_options.export_to_env();
     }
 
@@ -948,6 +962,10 @@ fn default_limit_batched_tasks_total_size() -> u64 {
     u64::MAX
 }
 
+fn default_embedding_cache_entries() -> usize {
+    0
+}
+
 fn default_snapshot_dir() -> PathBuf {
     PathBuf::from(DEFAULT_SNAPSHOT_DIR)
 }

crates/milli/src/update/index_documents/mod.rs

@@ -2806,8 +2806,9 @@ mod tests {
             embedding_configs.pop().unwrap();
         insta::assert_snapshot!(embedder_name, @"manual");
         insta::assert_debug_snapshot!(user_provided, @"RoaringBitmap<[0, 1, 2]>");
-        let embedder =
+        let embedder = std::sync::Arc::new(
-            std::sync::Arc::new(crate::vector::Embedder::new(embedder.embedder_options).unwrap());
+            crate::vector::Embedder::new(embedder.embedder_options, 0).unwrap(),
+        );
         let res = index
             .search(&rtxn)
             .semantic(embedder_name, embedder, false, Some([0.0, 1.0, 2.0].to_vec()))

crates/milli/src/update/settings.rs

@@ -1628,7 +1628,8 @@ fn embedders(embedding_configs: Vec<IndexEmbeddingConfig>) -> Result<EmbeddingCo
                 let prompt = Arc::new(prompt.try_into().map_err(crate::Error::from)?);
 
                 let embedder = Arc::new(
-                    Embedder::new(embedder_options.clone())
+                    // cache_cap: no cache needed for indexing purposes
+                    Embedder::new(embedder_options.clone(), 0)
                         .map_err(crate::vector::Error::from)
                         .map_err(crate::Error::from)?,
                 );

crates/milli/src/vector/composite.rs

@@ -59,9 +59,11 @@ pub struct EmbedderOptions {
 impl Embedder {
     pub fn new(
         EmbedderOptions { search, index }: EmbedderOptions,
+        cache_cap: usize,
     ) -> Result<Self, NewEmbedderError> {
-        let search = SubEmbedder::new(search)?;
+        let search = SubEmbedder::new(search, cache_cap)?;
-        let index = SubEmbedder::new(index)?;
+        // cache is only used at search
+        let index = SubEmbedder::new(index, 0)?;
 
         // check dimensions
         if search.dimensions() != index.dimensions() {
@@ -119,19 +121,28 @@ impl Embedder {
 }
 
 impl SubEmbedder {
-    pub fn new(options: SubEmbedderOptions) -> std::result::Result<Self, NewEmbedderError> {
+    pub fn new(
+        options: SubEmbedderOptions,
+        cache_cap: usize,
+    ) -> std::result::Result<Self, NewEmbedderError> {
         Ok(match options {
             SubEmbedderOptions::HuggingFace(options) => {
-                Self::HuggingFace(hf::Embedder::new(options)?)
+                Self::HuggingFace(hf::Embedder::new(options, cache_cap)?)
+            }
+            SubEmbedderOptions::OpenAi(options) => {
+                Self::OpenAi(openai::Embedder::new(options, cache_cap)?)
+            }
+            SubEmbedderOptions::Ollama(options) => {
+                Self::Ollama(ollama::Embedder::new(options, cache_cap)?)
             }
-            SubEmbedderOptions::OpenAi(options) => Self::OpenAi(openai::Embedder::new(options)?),
-            SubEmbedderOptions::Ollama(options) => Self::Ollama(ollama::Embedder::new(options)?),
             SubEmbedderOptions::UserProvided(options) => {
                 Self::UserProvided(manual::Embedder::new(options))
             }
-            SubEmbedderOptions::Rest(options) => {
+            SubEmbedderOptions::Rest(options) => Self::Rest(rest::Embedder::new(
-                Self::Rest(rest::Embedder::new(options, rest::ConfigurationSource::User)?)
+                options,
-            }
+                cache_cap,
+                rest::ConfigurationSource::User,
+            )?),
         })
     }
 

crates/milli/src/vector/hf.rs

@@ -150,7 +150,10 @@ impl From<PoolingConfig> for Pooling {
 }
 
 impl Embedder {
-    pub fn new(options: EmbedderOptions) -> std::result::Result<Self, NewEmbedderError> {
+    pub fn new(
+        options: EmbedderOptions,
+        cache_cap: usize,
+    ) -> std::result::Result<Self, NewEmbedderError> {
         let device = match candle_core::Device::cuda_if_available(0) {
             Ok(device) => device,
             Err(error) => {
@@ -252,7 +255,7 @@ impl Embedder {
             options,
             dimensions: 0,
             pooling,
-            cache: EmbeddingCache::new(super::CACHE_CAP),
+            cache: EmbeddingCache::new(cache_cap),
         };
 
         let embeddings = this

crates/milli/src/vector/mod.rs

@@ -560,8 +560,8 @@ struct EmbeddingCache {
 impl EmbeddingCache {
     const MAX_TEXT_LEN: usize = 2000;
 
-    pub fn new(cap: u16) -> Self {
+    pub fn new(cap: usize) -> Self {
-        let data = NonZeroUsize::new(cap.into()).map(lru::LruCache::new).map(Mutex::new);
+        let data = NonZeroUsize::new(cap).map(lru::LruCache::new).map(Mutex::new);
         Self { data }
     }
 
@@ -584,14 +584,14 @@ impl EmbeddingCache {
         if text.len() > Self::MAX_TEXT_LEN {
             return;
         }
+        tracing::trace!(text, "embedding added to cache");
+
         let mut cache = data.lock().unwrap();
 
         cache.put(text, embedding);
     }
 }
 
-pub const CACHE_CAP: u16 = 150;
-
 /// Configuration for an embedder.
 #[derive(Debug, Clone, Default, serde::Deserialize, serde::Serialize)]
 pub struct EmbeddingConfig {
@@ -670,19 +670,30 @@ impl Default for EmbedderOptions {
 
 impl Embedder {
     /// Spawns a new embedder built from its options.
-    pub fn new(options: EmbedderOptions) -> std::result::Result<Self, NewEmbedderError> {
+    pub fn new(
+        options: EmbedderOptions,
+        cache_cap: usize,
+    ) -> std::result::Result<Self, NewEmbedderError> {
         Ok(match options {
-            EmbedderOptions::HuggingFace(options) => Self::HuggingFace(hf::Embedder::new(options)?),
+            EmbedderOptions::HuggingFace(options) => {
-            EmbedderOptions::OpenAi(options) => Self::OpenAi(openai::Embedder::new(options)?),
+                Self::HuggingFace(hf::Embedder::new(options, cache_cap)?)
-            EmbedderOptions::Ollama(options) => Self::Ollama(ollama::Embedder::new(options)?),
+            }
+            EmbedderOptions::OpenAi(options) => {
+                Self::OpenAi(openai::Embedder::new(options, cache_cap)?)
+            }
+            EmbedderOptions::Ollama(options) => {
+                Self::Ollama(ollama::Embedder::new(options, cache_cap)?)
+            }
             EmbedderOptions::UserProvided(options) => {
                 Self::UserProvided(manual::Embedder::new(options))
             }
-            EmbedderOptions::Rest(options) => {
+            EmbedderOptions::Rest(options) => Self::Rest(rest::Embedder::new(
-                Self::Rest(rest::Embedder::new(options, rest::ConfigurationSource::User)?)
+                options,
-            }
+                cache_cap,
+                rest::ConfigurationSource::User,
+            )?),
             EmbedderOptions::Composite(options) => {
-                Self::Composite(composite::Embedder::new(options)?)
+                Self::Composite(composite::Embedder::new(options, cache_cap)?)
             }
         })
     }
@@ -718,7 +729,6 @@ impl Embedder {
         }?;
 
         if let Some(cache) = self.cache() {
-            tracing::trace!(text, "embedding added to cache");
             cache.put(text.to_owned(), embedding.clone());
         }
 

crates/milli/src/vector/ollama.rs

@@ -75,9 +75,10 @@ impl EmbedderOptions {
 }
 
 impl Embedder {
-    pub fn new(options: EmbedderOptions) -> Result<Self, NewEmbedderError> {
+    pub fn new(options: EmbedderOptions, cache_cap: usize) -> Result<Self, NewEmbedderError> {
         let rest_embedder = match RestEmbedder::new(
             options.into_rest_embedder_config()?,
+            cache_cap,
             super::rest::ConfigurationSource::Ollama,
         ) {
             Ok(embedder) => embedder,

crates/milli/src/vector/openai.rs

@@ -176,7 +176,7 @@ pub struct Embedder {
 }
 
 impl Embedder {
-    pub fn new(options: EmbedderOptions) -> Result<Self, NewEmbedderError> {
+    pub fn new(options: EmbedderOptions, cache_cap: usize) -> Result<Self, NewEmbedderError> {
         let mut inferred_api_key = Default::default();
         let api_key = options.api_key.as_ref().unwrap_or_else(|| {
             inferred_api_key = infer_api_key();
@@ -201,6 +201,7 @@ impl Embedder {
                 }),
                 headers: Default::default(),
             },
+            cache_cap,
             super::rest::ConfigurationSource::OpenAi,
         )?;
 

crates/milli/src/vector/rest.rs

@@ -10,8 +10,7 @@ use serde::{Deserialize, Serialize};
 use super::error::EmbedErrorKind;
 use super::json_template::ValueTemplate;
 use super::{
-    DistributionShift, EmbedError, Embedding, EmbeddingCache, NewEmbedderError, CACHE_CAP,
+    DistributionShift, EmbedError, Embedding, EmbeddingCache, NewEmbedderError, REQUEST_PARALLELISM,
-    REQUEST_PARALLELISM,
 };
 use crate::error::FaultSource;
 use crate::ThreadPoolNoAbort;
@@ -127,6 +126,7 @@ enum InputType {
 impl Embedder {
     pub fn new(
         options: EmbedderOptions,
+        cache_cap: usize,
         configuration_source: ConfigurationSource,
     ) -> Result<Self, NewEmbedderError> {
         let bearer = options.api_key.as_deref().map(|api_key| format!("Bearer {api_key}"));
@@ -160,7 +160,7 @@ impl Embedder {
             data,
             dimensions,
             distribution: options.distribution,
-            cache: EmbeddingCache::new(CACHE_CAP),
+            cache: EmbeddingCache::new(cache_cap),
         })
     }