snowykami/meilisearch
1
0
Fork 0
mirror of https://github.com/meilisearch/meilisearch.git synced 2025-02-22 02:25:32 +08:00
Code Issues Actions 6 Packages Projects Releases Wiki Activity

Fix the way we spawn the rayon threadpool

Browse Source
This commit is contained in:
Clément Renault 2024-11-27 16:05:44 +01:00
parent a514ce472a
commit 98d4a2909e
No known key found for this signature in database
GPG Key ID: F250A4C4E3AE5F5F
3 changed files with 313 additions and 273 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

40
crates/index-scheduler/src/batch.rs
View File

@ -1351,7 +1351,10 @@ impl IndexScheduler {
let pool = match &indexer_config.thread_pool {
Some(pool) => pool,
None => {
local_pool = ThreadPoolNoAbortBuilder::new().build().unwrap();
local_pool = ThreadPoolNoAbortBuilder::new()
.thread_name(|i| format!("indexing-thread-{i}"))
.build()
.unwrap();
&local_pool
}
};
@ -1399,10 +1402,10 @@ impl IndexScheduler {
}
if tasks.iter().any(|res| res.error.is_none()) {
pool.install(|| {
indexer::index(
index_wtxn,
index,
pool,
indexer_config.grenad_parameters(),
&db_fields_ids_map,
new_fields_ids_map,
@ -1411,9 +1414,7 @@ impl IndexScheduler {
embedders,
&|| must_stop_processing.get(),
&send_progress,
)
})
.unwrap()?;
)?;
tracing::info!(indexing_result = ?addition, processed_in = ?started_processing_at.elapsed(), "document indexing done");
}
@ -1489,21 +1490,25 @@ impl IndexScheduler {
let pool = match &indexer_config.thread_pool {
Some(pool) => pool,
None => {
local_pool = ThreadPoolNoAbortBuilder::new().build().unwrap();
local_pool = ThreadPoolNoAbortBuilder::new()
.thread_name(|i| format!("indexing-thread-{i}"))
.build()
.unwrap();
&local_pool
}
};
pool.install(|| {
let indexer =
UpdateByFunction::new(candidates, context.clone(), code.clone());
let document_changes = indexer.into_changes(&primary_key)?;
let indexer = UpdateByFunction::new(candidates, context.clone(), code.clone());
let document_changes =
pool.install(|| indexer.into_changes(&primary_key)).unwrap()?;
let embedders = index.embedding_configs(index_wtxn)?;
let embedders = self.embedders(embedders)?;
indexer::index(
index_wtxn,
index,
pool,
indexer_config.grenad_parameters(),
&db_fields_ids_map,
new_fields_ids_map,
@ -1514,10 +1519,6 @@ impl IndexScheduler {
&send_progress,
)?;
Result::Ok(())
})
.unwrap()?;
// tracing::info!(indexing_result = ?addition, processed_in = ?started_processing_at.elapsed(), "document indexing done");
}
@ -1641,7 +1642,10 @@ impl IndexScheduler {
let pool = match &indexer_config.thread_pool {
Some(pool) => pool,
None => {
local_pool = ThreadPoolNoAbortBuilder::new().build().unwrap();
local_pool = ThreadPoolNoAbortBuilder::new()
.thread_name(|i| format!("indexing-thread-{i}"))
.build()
.unwrap();
&local_pool
}
};
@ -1652,10 +1656,10 @@ impl IndexScheduler {
let embedders = index.embedding_configs(index_wtxn)?;
let embedders = self.embedders(embedders)?;
pool.install(|| {
indexer::index(
index_wtxn,
index,
pool,
indexer_config.grenad_parameters(),
&db_fields_ids_map,
new_fields_ids_map,
@ -1664,9 +1668,7 @@ impl IndexScheduler {
embedders,
&|| must_stop_processing.get(),
&send_progress,
)
})
.unwrap()?;
)?;
// tracing::info!(indexing_result = ?addition, processed_in = ?started_processing_at.elapsed(), "document indexing done");
}

36
crates/milli/src/update/new/channel.rs
View File

@ -55,6 +55,12 @@ pub fn extractor_writer_bbqueue(
let producers = ThreadLocal::with_capacity(bbbuffers.len());
let consumers = rayon::broadcast(|bi| {
eprintln!(
"hello thread #{:?} (#{:?}, #{:?})",
bi.index(),
std::thread::current().name(),
std::thread::current().id(),
);
let bbqueue = &bbbuffers[bi.index()];
let (producer, consumer) = bbqueue.try_split_framed().unwrap();
producers.get_or(|| FullySend(RefCell::new(producer)));
@ -399,7 +405,15 @@ impl<'b> ExtractorBbqueueSender<'b> {
fn delete_vector(&self, docid: DocumentId) -> crate::Result<()> {
let capacity = self.capacity;
let refcell = self.producers.get().unwrap();
let refcell = match self.producers.get() {
Some(refcell) => refcell,
None => panic!(
"hello thread #{:?} (#{:?}, #{:?})",
rayon::current_thread_index(),
std::thread::current().name(),
std::thread::current().id()
),
};
let mut producer = refcell.0.borrow_mut_or_yield();
let payload_header = EntryHeader::ArroyDeleteVector(ArroyDeleteVector { docid });
@ -438,7 +452,15 @@ impl<'b> ExtractorBbqueueSender<'b> {
embedding: &[f32],
) -> crate::Result<()> {
let capacity = self.capacity;
let refcell = self.producers.get().unwrap();
let refcell = match self.producers.get() {
Some(refcell) => refcell,
None => panic!(
"hello thread #{:?} (#{:?}, #{:?})",
rayon::current_thread_index(),
std::thread::current().name(),
std::thread::current().id()
),
};
let mut producer = refcell.0.borrow_mut_or_yield();
let payload_header =
@ -496,7 +518,15 @@ impl<'b> ExtractorBbqueueSender<'b> {
F: FnOnce(&mut [u8]) -> crate::Result<()>,
{
let capacity = self.capacity;
let refcell = self.producers.get().unwrap();
let refcell = match self.producers.get() {
Some(refcell) => refcell,
None => panic!(
"hello thread #{:?} (#{:?}, #{:?})",
rayon::current_thread_index(),
std::thread::current().name(),
std::thread::current().id()
),
};
let mut producer = refcell.0.borrow_mut_or_yield();
let operation = DbOperation { database, key_length: Some(key_length) };

14
crates/milli/src/update/new/indexer/mod.rs
View File

@ -62,6 +62,7 @@ mod update_by_function;
pub fn index<'pl, 'indexer, 'index, DC, MSP, SP>(
wtxn: &mut RwTxn,
index: &'index Index,
pool: &ThreadPoolNoAbort,
grenad_parameters: GrenadParameters,
db_fields_ids_map: &'indexer FieldsIdsMap,
new_fields_ids_map: FieldsIdsMap,
@ -77,10 +78,15 @@ where
SP: Fn(Progress) + Sync,
{
/// TODO restrict memory and remove this memory from the extractors bump allocators
let bbbuffers: Vec<_> = (0..rayon::current_num_threads())
let bbbuffers: Vec<_> = pool
.install(|| {
(0..rayon::current_num_threads())
.map(|_| bbqueue::BBBuffer::new(100 * 1024 * 1024)) // 100 MiB by thread
.collect();
let (extractor_sender, mut writer_receiver) = extractor_writer_bbqueue(&bbbuffers, 1000);
.collect()
})
.unwrap();
let (extractor_sender, mut writer_receiver) =
pool.install(|| extractor_writer_bbqueue(&bbbuffers, 1000)).unwrap();
let finished_extraction = AtomicBool::new(false);
let metadata_builder = MetadataBuilder::from_index(index, wtxn)?;
@ -112,6 +118,7 @@ where
let field_distribution = &mut field_distribution;
let document_ids = &mut document_ids;
let extractor_handle = Builder::new().name(S("indexer-extractors")).spawn_scoped(s, move || {
pool.install(move || {
let span = tracing::trace_span!(target: "indexing::documents", parent: &indexer_span, "extract");
let _entered = span.enter();
@ -359,6 +366,7 @@ where
finished_extraction.store(true, std::sync::atomic::Ordering::Relaxed);
Result::Ok((facet_field_ids_delta, index_embeddings))
}).unwrap()
})?;
let global_fields_ids_map = GlobalFieldsIdsMap::new(&new_fields_ids_map);