From f53b4f3317e0f57354d430d48a1bbe314e2aed7c Mon Sep 17 00:00:00 2001 From: timvisee Date: Wed, 15 Jan 2025 11:25:12 +0100 Subject: [PATCH] Add max optimization threads builder --- .../max_optimization_threads_builder.rs | 72 +++++++++++++++++++ src/builders/mod.rs | 1 + .../optimizers_config_diff_builder.rs | 37 ++++++++-- 3 files changed, 104 insertions(+), 6 deletions(-) diff --git a/src/builders/max_optimization_threads_builder.rs b/src/builders/max_optimization_threads_builder.rs index a3cc130..2957628 100644 --- a/src/builders/max_optimization_threads_builder.rs +++ b/src/builders/max_optimization_threads_builder.rs @@ -1,5 +1,63 @@ use crate::qdrant::*; +/// Max number of threads (jobs) for running optimizations per shard. +/// Each optimization job will also use `max_indexing_threads` threads by itself for index building. +/// +/// - If `auto` - have no limit and choose dynamically to saturate CPU. +/// - If `disabled` or `0` - no optimization threads, optimizations will be disabled. +pub struct MaxOptimizationThreadsBuilder { + pub(crate) inner: MaxOptimizationThreads, +} + +impl MaxOptimizationThreadsBuilder { + /// Use specific number of optimization threads. + /// + /// - If `0` - no optimization threads, optimizations will be disabled. + #[allow(unused_mut)] + #[inline] + pub fn threads(threads: u64) -> Self { + Self { + inner: MaxOptimizationThreads::from(threads), + } + } + + /// No optimization threads, optimizations will be disabled. + #[allow(unused_mut)] + #[inline] + pub fn disabled() -> Self { + Self::threads(0) + } + + /// Have no limit and choose dynamically to saturate CPU. + #[allow(unused_mut)] + #[inline] + pub fn auto() -> Self { + Self { + inner: MaxOptimizationThreads::from(max_optimization_threads::Variant::Setting( + max_optimization_threads::Setting::Auto as i32, + )), + } + } +} + +impl From for MaxOptimizationThreads { + fn from(value: MaxOptimizationThreadsBuilder) -> Self { + value.build() + } +} + +impl MaxOptimizationThreadsBuilder { + pub fn build(self) -> MaxOptimizationThreads { + self.inner + } +} + +impl Default for MaxOptimizationThreadsBuilder { + fn default() -> Self { + Self::auto() + } +} + impl From for MaxOptimizationThreads { fn from(threads: u64) -> Self { MaxOptimizationThreads { @@ -8,8 +66,22 @@ impl From for MaxOptimizationThreads { } } +impl From for MaxOptimizationThreads { + fn from(setting: max_optimization_threads::Variant) -> Self { + MaxOptimizationThreads { + variant: Some(setting), + } + } +} + impl From for max_optimization_threads::Variant { fn from(threads: u64) -> Self { Self::Value(threads) } } + +impl From for max_optimization_threads::Variant { + fn from(setting: max_optimization_threads::Setting) -> Self { + Self::Setting(setting as i32) + } +} diff --git a/src/builders/mod.rs b/src/builders/mod.rs index 4c8f335..422da5c 100644 --- a/src/builders/mod.rs +++ b/src/builders/mod.rs @@ -71,6 +71,7 @@ mod rename_alias_builder; pub use rename_alias_builder::RenameAliasBuilder; mod max_optimization_threads_builder; +pub use max_optimization_threads_builder::MaxOptimizationThreadsBuilder; mod move_shard_builder; pub use move_shard_builder::MoveShardBuilder; diff --git a/src/builders/optimizers_config_diff_builder.rs b/src/builders/optimizers_config_diff_builder.rs index 521928e..1f63a1d 100644 --- a/src/builders/optimizers_config_diff_builder.rs +++ b/src/builders/optimizers_config_diff_builder.rs @@ -51,9 +51,10 @@ pub struct OptimizersConfigDiffBuilder { pub(crate) flush_interval_sec: Option>, /// /// Max number of threads (jobs) for running optimizations per shard. - /// Note: each optimization job will also use `max_indexing_threads` threads by itself for index building. - /// If null - have no limit and choose dynamically to saturate CPU. - /// If 0 - no optimization threads, optimizations will be disabled. + /// Each optimization job will also use `max_indexing_threads` threads by itself for index building. + /// + /// - If `auto` - have no limit and choose dynamically to saturate CPU. + /// - If `disabled` or `0` - no optimization threads, optimizations will be disabled. pub(crate) max_optimization_threads: Option>, } @@ -143,9 +144,33 @@ impl OptimizersConfigDiffBuilder { } /// /// Max number of threads (jobs) for running optimizations per shard. - /// Note: each optimization job will also use `max_indexing_threads` threads by itself for index building. - /// If null - have no limit and choose dynamically to saturate CPU. - /// If 0 - no optimization threads, optimizations will be disabled. + /// Each optimization job will also use `max_indexing_threads` threads by itself for index building. + /// + /// - If `auto` - have no limit and choose dynamically to saturate CPU. + /// - If `disabled` or `0` - no optimization threads, optimizations will be disabled. + /// + /// ```no_run + ///# use qdrant_client::{Qdrant, QdrantError}; + /// use qdrant_client::qdrant::{OptimizersConfigDiffBuilder, UpdateCollectionBuilder, MaxOptimizationThreadsBuilder}; + /// + ///# async fn create_collection(client: &Qdrant) + ///# -> Result<(), QdrantError> { + /// let optimizers_config = OptimizersConfigDiffBuilder::default() + /// // Use exactly 8 threads + /// .max_optimization_threads(8) + /// // Or automatically choose + /// .max_optimization_threads(MaxOptimizationThreadsBuilder::auto()) + /// // Or disable + /// .max_optimization_threads(MaxOptimizationThreadsBuilder::disabled()); + /// + /// client + /// .update_collection( + /// UpdateCollectionBuilder::new("my_collection").optimizers_config(optimizers_config), + /// ) + /// .await?; + ///# Ok(()) + ///# } + /// ``` #[allow(unused_mut)] pub fn max_optimization_threads>( self,