Enable resetting max_optimization_threads to automatic (#5634)

* Allow max-optimization-thread config to be set to null on update

* add new field. Support both, but prepare for deprecation.

* better openapi

* only introduce change in OptimizersConfigDiff

* move to `api::rest::schema`

* update openapi

* improve test, fix diff to config conversion

* upd grpc docs

* clippy

* remove schemars from common common

---------

Co-authored-by: Gulshan Kumar <kumargu@amazon.com>
This commit is contained in:
Luis Cossío
2024-12-18 17:17:47 -03:00
committed by timvisee
parent f2b142620f
commit 8b4e798948
9 changed files with 332 additions and 36 deletions

View File

@@ -47,6 +47,7 @@
- [ListCollectionsRequest](#qdrant-ListCollectionsRequest)
- [ListCollectionsResponse](#qdrant-ListCollectionsResponse)
- [LocalShardInfo](#qdrant-LocalShardInfo)
- [MaxOptimizationThreads](#qdrant-MaxOptimizationThreads)
- [MoveShard](#qdrant-MoveShard)
- [MultiVectorConfig](#qdrant-MultiVectorConfig)
- [OptimizerStatus](#qdrant-OptimizerStatus)
@@ -89,6 +90,7 @@
- [CompressionRatio](#qdrant-CompressionRatio)
- [Datatype](#qdrant-Datatype)
- [Distance](#qdrant-Distance)
- [MaxOptimizationThreads.Setting](#qdrant-MaxOptimizationThreads-Setting)
- [Modifier](#qdrant-Modifier)
- [MultiVectorComparator](#qdrant-MultiVectorComparator)
- [PayloadSchemaType](#qdrant-PayloadSchemaType)
@@ -1018,6 +1020,22 @@
<a name="qdrant-MaxOptimizationThreads"></a>
### MaxOptimizationThreads
| Field | Type | Label | Description |
| ----- | ---- | ----- | ----------- |
| value | [uint64](#uint64) | | |
| setting | [MaxOptimizationThreads.Setting](#qdrant-MaxOptimizationThreads-Setting) | | |
<a name="qdrant-MoveShard"></a>
### MoveShard
@@ -1101,7 +1119,8 @@ To disable vector indexing, set to `0`.
Note: 1kB = 1 vector of size 256. |
| flush_interval_sec | [uint64](#uint64) | optional | Interval between forced flushes. |
| max_optimization_threads | [uint64](#uint64) | optional | 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. |
| deprecated_max_optimization_threads | [uint64](#uint64) | optional | Deprecated in favor of `max_optimization_threads` |
| max_optimization_threads | [MaxOptimizationThreads](#qdrant-MaxOptimizationThreads) | optional | 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 &#34;auto&#34; - have no limit and choose dynamically to saturate CPU. If 0 - no optimization threads, optimizations will be disabled. |
@@ -1751,6 +1770,17 @@ Note: 1kB = 1 vector of size 256. |
<a name="qdrant-MaxOptimizationThreads-Setting"></a>
### MaxOptimizationThreads.Setting
| Name | Number | Description |
| ---- | ------ | ----------- |
| Auto | 0 | |
<a name="qdrant-Modifier"></a>
### Modifier

View File

@@ -9449,14 +9449,36 @@
"nullable": true
},
"max_optimization_threads": {
"description": "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.",
"type": "integer",
"format": "uint",
"minimum": 0,
"nullable": true
"description": "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 \"auto\" - have no limit and choose dynamically to saturate CPU. If 0 - no optimization threads, optimizations will be disabled.",
"anyOf": [
{
"$ref": "#/components/schemas/MaxOptimizationThreads"
},
{
"nullable": true
}
]
}
}
},
"MaxOptimizationThreads": {
"anyOf": [
{
"$ref": "#/components/schemas/MaxOptimizationThreadsSetting"
},
{
"type": "integer",
"format": "uint",
"minimum": 0
}
]
},
"MaxOptimizationThreadsSetting": {
"type": "string",
"enum": [
"auto"
]
},
"InitFrom": {
"description": "Operation for creating new collection and (optionally) specify index params",
"type": "object",

View File

@@ -22,10 +22,10 @@ use super::qdrant::{
raw_query, start_from, BinaryQuantization, BoolIndexParams, CompressionRatio,
DatetimeIndexParams, DatetimeRange, Direction, FacetHit, FacetHitInternal, FacetValue,
FacetValueInternal, FieldType, FloatIndexParams, GeoIndexParams, GeoLineString, GroupId,
HardwareUsage, HasVectorCondition, KeywordIndexParams, LookupLocation, MultiVectorComparator,
MultiVectorConfig, OrderBy, OrderValue, Range, RawVector, RecommendStrategy, RetrievedPoint,
SearchMatrixPair, SearchPointGroups, SearchPoints, ShardKeySelector, SparseIndices, StartFrom,
UuidIndexParams, VectorsOutput, WithLookup,
HardwareUsage, HasVectorCondition, KeywordIndexParams, LookupLocation, MaxOptimizationThreads,
MultiVectorComparator, MultiVectorConfig, OrderBy, OrderValue, Range, RawVector,
RecommendStrategy, RetrievedPoint, SearchMatrixPair, SearchPointGroups, SearchPoints,
ShardKeySelector, SparseIndices, StartFrom, UuidIndexParams, VectorsOutput, WithLookup,
};
use crate::conversions::json;
use crate::grpc::qdrant::condition::ConditionOneOf;
@@ -1609,6 +1609,76 @@ impl From<segment::data_types::order_by::StartFrom> for StartFrom {
}
}
impl TryFrom<MaxOptimizationThreads> for rest::MaxOptimizationThreads {
type Error = Status;
fn try_from(value: MaxOptimizationThreads) -> Result<Self, Self::Error> {
use crate::grpc::qdrant::max_optimization_threads::{Setting, Variant};
let variant = value
.variant
.ok_or_else(|| Status::invalid_argument("Malformed MaxOptimizationThreads"))?;
let converted = match variant {
Variant::Setting(setting_int) => {
let setting = Setting::try_from(setting_int).map_err(|err| {
Status::invalid_argument(format!(
"Invalid MaxOptimizationThreads setting: {err}"
))
})?;
match setting {
Setting::Auto => Self::Setting(rest::MaxOptimizationThreadsSetting::Auto),
}
}
Variant::Value(num_threads) => Self::Threads(num_threads as usize),
};
Ok(converted)
}
}
impl TryFrom<MaxOptimizationThreads> for Option<usize> {
type Error = Status;
fn try_from(value: MaxOptimizationThreads) -> Result<Self, Self::Error> {
use crate::grpc::qdrant::max_optimization_threads::{Setting, Variant};
let variant = value
.variant
.ok_or_else(|| Status::invalid_argument("Malformed MaxOptimizationThreads"))?;
Ok(match variant {
Variant::Setting(setting_int) => {
let setting = Setting::try_from(setting_int).map_err(|err| {
Status::invalid_argument(format!(
"Invalid MaxOptimizationThreads setting: {err}"
))
})?;
match setting {
Setting::Auto => None,
}
}
Variant::Value(num_threads) => Some(num_threads as usize),
})
}
}
impl From<Option<usize>> for MaxOptimizationThreads {
fn from(value: Option<usize>) -> Self {
use crate::grpc::qdrant::max_optimization_threads::{Setting, Variant};
let variant = match value {
None => Variant::Setting(Setting::Auto.into()),
Some(n) => Variant::Value(n as u64),
};
Self {
variant: Some(variant),
}
}
}
impl From<HnswConfigDiff> for segment::types::HnswConfig {
fn from(hnsw_config: HnswConfigDiff) -> Self {
Self {

View File

@@ -145,6 +145,17 @@ enum CompressionRatio {
x64 = 4;
}
message MaxOptimizationThreads {
enum Setting {
Auto = 0;
}
oneof variant {
uint64 value = 1;
Setting setting = 2;
}
}
message OptimizerStatus {
bool ok = 1;
string error = 2;
@@ -260,13 +271,17 @@ message OptimizersConfigDiff {
Interval between forced flushes.
*/
optional uint64 flush_interval_sec = 7;
// Deprecated in favor of `max_optimization_threads`
optional uint64 deprecated_max_optimization_threads = 8;
/*
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 "auto" - have no limit and choose dynamically to saturate CPU.
If 0 - no optimization threads, optimizations will be disabled.
*/
optional uint64 max_optimization_threads = 8;
optional MaxOptimizationThreads max_optimization_threads = 9;
}
message ScalarQuantization {

View File

@@ -210,6 +210,59 @@ pub struct ListCollectionsResponse {
#[derive(serde::Serialize)]
#[allow(clippy::derive_partial_eq_without_eq)]
#[derive(Clone, PartialEq, ::prost::Message)]
pub struct MaxOptimizationThreads {
#[prost(oneof = "max_optimization_threads::Variant", tags = "1, 2")]
pub variant: ::core::option::Option<max_optimization_threads::Variant>,
}
/// Nested message and enum types in `MaxOptimizationThreads`.
pub mod max_optimization_threads {
#[derive(serde::Serialize)]
#[derive(
Clone,
Copy,
Debug,
PartialEq,
Eq,
Hash,
PartialOrd,
Ord,
::prost::Enumeration
)]
#[repr(i32)]
pub enum Setting {
Auto = 0,
}
impl Setting {
/// String value of the enum field names used in the ProtoBuf definition.
///
/// The values are not transformed in any way and thus are considered stable
/// (if the ProtoBuf definition does not change) and safe for programmatic use.
pub fn as_str_name(&self) -> &'static str {
match self {
Setting::Auto => "Auto",
}
}
/// Creates an enum from field names used in the ProtoBuf definition.
pub fn from_str_name(value: &str) -> ::core::option::Option<Self> {
match value {
"Auto" => Some(Self::Auto),
_ => None,
}
}
}
#[derive(serde::Serialize)]
#[allow(clippy::derive_partial_eq_without_eq)]
#[derive(Clone, PartialEq, ::prost::Oneof)]
pub enum Variant {
#[prost(uint64, tag = "1")]
Value(u64),
#[prost(enumeration = "Setting", tag = "2")]
Setting(i32),
}
}
#[derive(serde::Serialize)]
#[allow(clippy::derive_partial_eq_without_eq)]
#[derive(Clone, PartialEq, ::prost::Message)]
pub struct OptimizerStatus {
#[prost(bool, tag = "1")]
pub ok: bool,
@@ -330,12 +383,15 @@ pub struct OptimizersConfigDiff {
/// Interval between forced flushes.
#[prost(uint64, optional, tag = "7")]
pub flush_interval_sec: ::core::option::Option<u64>,
/// Deprecated in favor of `max_optimization_threads`
#[prost(uint64, optional, tag = "8")]
pub deprecated_max_optimization_threads: ::core::option::Option<u64>,
/// 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 "auto" - have no limit and choose dynamically to saturate CPU.
/// If 0 - no optimization threads, optimizations will be disabled.
#[prost(uint64, optional, tag = "8")]
pub max_optimization_threads: ::core::option::Option<u64>,
#[prost(message, optional, tag = "9")]
pub max_optimization_threads: ::core::option::Option<MaxOptimizationThreads>,
}
#[derive(validator::Validate)]
#[derive(serde::Serialize)]

View File

@@ -1050,3 +1050,32 @@ impl PointInsertOperations {
self.len() == 0
}
}
#[derive(Debug, Deserialize, Serialize, Clone, Copy, PartialEq, Hash, Default, JsonSchema)]
#[serde(rename_all = "snake_case")]
pub enum MaxOptimizationThreadsSetting {
#[default]
Auto,
}
#[derive(Debug, Deserialize, Serialize, Clone, Copy, PartialEq, Hash, JsonSchema)]
#[serde(untagged)]
pub enum MaxOptimizationThreads {
Setting(MaxOptimizationThreadsSetting),
Threads(usize),
}
impl Default for MaxOptimizationThreads {
fn default() -> Self {
MaxOptimizationThreads::Setting(MaxOptimizationThreadsSetting::Auto)
}
}
impl From<MaxOptimizationThreads> for Option<usize> {
fn from(value: MaxOptimizationThreads) -> Self {
match value {
MaxOptimizationThreads::Setting(MaxOptimizationThreadsSetting::Auto) => None,
MaxOptimizationThreads::Threads(threads) => Some(threads),
}
}
}

View File

@@ -1,6 +1,7 @@
use std::hash::Hash;
use std::num::NonZeroU32;
use api::rest::MaxOptimizationThreads;
use merge::Merge;
use schemars::JsonSchema;
use segment::types::{
@@ -162,9 +163,9 @@ pub struct OptimizersConfigDiff {
pub flush_interval_sec: Option<u64>,
/// 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 "auto" - have no limit and choose dynamically to saturate CPU.
/// If 0 - no optimization threads, optimizations will be disabled.
pub max_optimization_threads: Option<usize>,
pub max_optimization_threads: Option<MaxOptimizationThreads>,
}
impl std::hash::Hash for OptimizersConfigDiff {
@@ -200,7 +201,29 @@ impl DiffConfig<HnswConfig> for HnswConfigDiff {}
impl DiffConfig<HnswConfigDiff> for HnswConfigDiff {}
impl DiffConfig<OptimizersConfig> for OptimizersConfigDiff {}
impl DiffConfig<OptimizersConfig> for OptimizersConfigDiff {
fn update(self, config: &OptimizersConfig) -> CollectionResult<OptimizersConfig>
where
Self: Sized + Serialize + DeserializeOwned + Merge,
{
Ok(OptimizersConfig {
deleted_threshold: self.deleted_threshold.unwrap_or(config.deleted_threshold),
vacuum_min_vector_number: self
.vacuum_min_vector_number
.unwrap_or(config.vacuum_min_vector_number),
default_segment_number: self
.default_segment_number
.unwrap_or(config.default_segment_number),
max_segment_size: self.max_segment_size.or(config.max_segment_size),
memmap_threshold: self.memmap_threshold.or(config.memmap_threshold),
indexing_threshold: self.indexing_threshold.or(config.indexing_threshold),
flush_interval_sec: self.flush_interval_sec.unwrap_or(config.flush_interval_sec),
max_optimization_threads: self
.max_optimization_threads
.map_or(config.max_optimization_threads, From::from),
})
}
}
impl DiffConfig<WalConfig> for WalConfigDiff {}
@@ -333,6 +356,7 @@ impl Validate for QuantizationConfigDiff {
#[cfg(test)]
mod tests {
use rstest::rstest;
use segment::types::{Distance, HnswConfig};
use super::*;
@@ -388,6 +412,31 @@ mod tests {
assert_eq!(new_config.indexing_threshold, Some(10000))
}
#[rstest]
#[case::number(r#"{ "max_optimization_threads": 5 }"#, Some(5))]
#[case::auto(r#"{ "max_optimization_threads": "auto" }"#, None)]
#[case::null(r#"{ "max_optimization_threads": null }"#, Some(1))] // no effect
#[case::nothing("{ }", Some(1))] // no effect
#[should_panic]
#[case::other(r#"{ "max_optimization_threads": "other" }"#, Some(1))]
fn test_set_optimizer_threads(#[case] json_diff: &str, #[case] expected: Option<usize>) {
let base_config = OptimizersConfig {
deleted_threshold: 0.9,
vacuum_min_vector_number: 1000,
default_segment_number: 10,
max_segment_size: None,
memmap_threshold: None,
indexing_threshold: Some(50_000),
flush_interval_sec: 30,
max_optimization_threads: Some(1),
};
let update: OptimizersConfigDiff = serde_json::from_str(json_diff).unwrap();
let new_config = update.update(&base_config).unwrap();
assert_eq!(new_config.max_optimization_threads, expected);
}
#[test]
fn test_wal_config() {
let base_config = WalConfig::default();

View File

@@ -13,7 +13,7 @@ use api::grpc::qdrant::update_collection_cluster_setup_request::{
};
use api::grpc::qdrant::{CreateShardKey, Vectors};
use api::rest::schema::ShardKeySelector;
use api::rest::BaseGroupRequest;
use api::rest::{BaseGroupRequest, MaxOptimizationThreads};
use common::types::ScoreType;
use itertools::Itertools;
use segment::common::operation_error::OperationError;
@@ -295,9 +295,11 @@ impl TryFrom<api::grpc::qdrant::CollectionParamsDiff> for CollectionParamsDiff {
}
}
impl From<api::grpc::qdrant::OptimizersConfigDiff> for OptimizersConfigDiff {
fn from(value: api::grpc::qdrant::OptimizersConfigDiff) -> Self {
Self {
impl TryFrom<api::grpc::qdrant::OptimizersConfigDiff> for OptimizersConfigDiff {
type Error = Status;
fn try_from(value: api::grpc::qdrant::OptimizersConfigDiff) -> Result<Self, Self::Error> {
Ok(Self {
deleted_threshold: value.deleted_threshold,
vacuum_min_vector_number: value.vacuum_min_vector_number.map(|v| v as usize),
default_segment_number: value.default_segment_number.map(|v| v as usize),
@@ -305,8 +307,15 @@ impl From<api::grpc::qdrant::OptimizersConfigDiff> for OptimizersConfigDiff {
memmap_threshold: value.memmap_threshold.map(|v| v as usize),
indexing_threshold: value.indexing_threshold.map(|v| v as usize),
flush_interval_sec: value.flush_interval_sec,
max_optimization_threads: value.max_optimization_threads.map(|v| v as usize),
}
// TODO: remove deprecated field in a later version
max_optimization_threads: value
.deprecated_max_optimization_threads
.map(|v| MaxOptimizationThreads::Threads(v as usize))
.or(value
.max_optimization_threads
.map(TryFrom::try_from)
.transpose()?),
})
}
}
@@ -426,10 +435,13 @@ impl From<CollectionInfo> for api::grpc::qdrant::CollectionInfo {
.indexing_threshold
.map(|x| x as u64),
flush_interval_sec: Some(config.optimizer_config.flush_interval_sec),
max_optimization_threads: config
deprecated_max_optimization_threads: config
.optimizer_config
.max_optimization_threads
.map(|n| n as u64),
.map(|x| x as u64),
max_optimization_threads: Some(From::from(
config.optimizer_config.max_optimization_threads,
)),
}),
wal_config: config
.wal_config
@@ -484,9 +496,18 @@ impl TryFrom<i32> for CollectionStatus {
}
}
impl From<api::grpc::qdrant::OptimizersConfigDiff> for OptimizersConfig {
fn from(optimizer_config: api::grpc::qdrant::OptimizersConfigDiff) -> Self {
Self {
impl TryFrom<api::grpc::qdrant::OptimizersConfigDiff> for OptimizersConfig {
type Error = Status;
fn try_from(
optimizer_config: api::grpc::qdrant::OptimizersConfigDiff,
) -> Result<Self, Self::Error> {
debug_assert!(
optimizer_config.max_optimization_threads.is_some(),
"This conversion is for CollectionInfo, max_optimization_threads should always have a value"
);
Ok(Self {
deleted_threshold: optimizer_config.deleted_threshold.unwrap_or_default(),
vacuum_min_vector_number: optimizer_config
.vacuum_min_vector_number
@@ -497,10 +518,11 @@ impl From<api::grpc::qdrant::OptimizersConfigDiff> for OptimizersConfig {
memmap_threshold: optimizer_config.memmap_threshold.map(|x| x as usize),
indexing_threshold: optimizer_config.indexing_threshold.map(|x| x as usize),
flush_interval_sec: optimizer_config.flush_interval_sec.unwrap_or_default(),
max_optimization_threads: optimizer_config
.max_optimization_threads
.map(|n| n as usize),
}
max_optimization_threads: match optimizer_config.max_optimization_threads {
None => return Err(Status::invalid_argument("Malformed OptimizersConfig")),
Some(max_optimization_threads) => TryFrom::try_from(max_optimization_threads)?,
},
})
}
}
@@ -1795,7 +1817,7 @@ impl TryFrom<api::grpc::qdrant::CollectionConfig> for CollectionConfig {
},
optimizer_config: match config.optimizer_config {
None => return Err(Status::invalid_argument("Malformed OptimizerConfig type")),
Some(optimizer_config) => OptimizersConfig::from(optimizer_config),
Some(optimizer_config) => OptimizersConfig::try_from(optimizer_config)?,
},
wal_config: match config.wal_config {
None => return Err(Status::invalid_argument("Malformed WalConfig type")),

View File

@@ -52,7 +52,7 @@ impl TryFrom<api::grpc::qdrant::CreateCollection> for CollectionMetaOperations {
.transpose()?,
hnsw_config: value.hnsw_config.map(|v| v.into()),
wal_config: value.wal_config.map(|v| v.into()),
optimizers_config: value.optimizers_config.map(|v| v.into()),
optimizers_config: value.optimizers_config.map(TryFrom::try_from).transpose()?,
shard_number: value.shard_number,
on_disk_payload: value.on_disk_payload,
replication_factor: value.replication_factor,
@@ -114,7 +114,10 @@ impl TryFrom<api::grpc::qdrant::UpdateCollection> for CollectionMetaOperations {
.params
.map(CollectionParamsDiff::try_from)
.transpose()?,
optimizers_config: value.optimizers_config.map(OptimizersConfigDiff::from),
optimizers_config: value
.optimizers_config
.map(OptimizersConfigDiff::try_from)
.transpose()?,
quantization_config: value
.quantization_config
.map(QuantizationConfigDiff::try_from)