Spaces:
Build error
Build error
File size: 4,274 Bytes
84d2a97 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 |
use std::num::NonZeroU32;
use std::path::Path;
use std::sync::Arc;
use collection::collection::{Collection, RequestShardTransfer};
use collection::config::{CollectionConfigInternal, CollectionParams, WalConfig};
use collection::operations::types::CollectionError;
use collection::operations::vector_params_builder::VectorParamsBuilder;
use collection::optimizers_builder::OptimizersConfig;
use collection::shards::channel_service::ChannelService;
use collection::shards::collection_shard_distribution::CollectionShardDistribution;
use collection::shards::replica_set::{AbortShardTransfer, ChangePeerFromState, ReplicaState};
use collection::shards::CollectionId;
use common::cpu::CpuBudget;
use segment::types::Distance;
/// Test collections for this upper bound of shards.
/// Testing with more shards is problematic due to `number of open files problem`
/// See https://github.com/qdrant/qdrant/issues/379
pub const N_SHARDS: u32 = 3;
pub const REST_PORT: u16 = 6333;
pub const TEST_OPTIMIZERS_CONFIG: OptimizersConfig = OptimizersConfig {
deleted_threshold: 0.9,
vacuum_min_vector_number: 1000,
default_segment_number: 2,
max_segment_size: None,
memmap_threshold: None,
indexing_threshold: Some(50_000),
flush_interval_sec: 30,
max_optimization_threads: Some(2),
};
#[cfg(test)]
pub async fn simple_collection_fixture(collection_path: &Path, shard_number: u32) -> Collection {
let wal_config = WalConfig {
wal_capacity_mb: 1,
wal_segments_ahead: 0,
};
let collection_params = CollectionParams {
vectors: VectorParamsBuilder::new(4, Distance::Dot).build().into(),
shard_number: NonZeroU32::new(shard_number).expect("Shard number can not be zero"),
..CollectionParams::empty()
};
let collection_config = CollectionConfigInternal {
params: collection_params,
optimizer_config: TEST_OPTIMIZERS_CONFIG.clone(),
wal_config,
hnsw_config: Default::default(),
quantization_config: Default::default(),
strict_mode_config: Default::default(),
uuid: None,
};
let snapshot_path = collection_path.join("snapshots");
// Default to a collection with all the shards local
new_local_collection(
"test".to_string(),
collection_path,
&snapshot_path,
&collection_config,
)
.await
.unwrap()
}
pub fn dummy_on_replica_failure() -> ChangePeerFromState {
Arc::new(move |_peer_id, _shard_id, _from_state| {})
}
pub fn dummy_request_shard_transfer() -> RequestShardTransfer {
Arc::new(move |_transfer| {})
}
pub fn dummy_abort_shard_transfer() -> AbortShardTransfer {
Arc::new(|_transfer, _reason| {})
}
/// Default to a collection with all the shards local
#[cfg(test)]
pub async fn new_local_collection(
id: CollectionId,
path: &Path,
snapshots_path: &Path,
config: &CollectionConfigInternal,
) -> Result<Collection, CollectionError> {
let collection = Collection::new(
id,
0,
path,
snapshots_path,
config,
Default::default(),
CollectionShardDistribution::all_local(Some(config.params.shard_number.into()), 0),
ChannelService::new(REST_PORT, None),
dummy_on_replica_failure(),
dummy_request_shard_transfer(),
dummy_abort_shard_transfer(),
None,
None,
CpuBudget::default(),
None,
)
.await;
let collection = collection?;
let local_shards = collection.get_local_shards().await;
for shard_id in local_shards {
collection
.set_shard_replica_state(shard_id, 0, ReplicaState::Active, None)
.await?;
}
Ok(collection)
}
/// Default to a collection with all the shards local
#[cfg(test)]
pub async fn load_local_collection(
id: CollectionId,
path: &Path,
snapshots_path: &Path,
) -> Collection {
Collection::load(
id,
0,
path,
snapshots_path,
Default::default(),
ChannelService::new(REST_PORT, None),
dummy_on_replica_failure(),
dummy_request_shard_transfer(),
dummy_abort_shard_transfer(),
None,
None,
CpuBudget::default(),
None,
)
.await
}
|