Compare commits

..

1 Commits

Author SHA1 Message Date
2373adf9b7 Add ollama url in settings GET 2024-05-02 15:36:13 +02:00
9 changed files with 29 additions and 238 deletions

10
Cargo.lock generated
View File

@ -2506,14 +2506,6 @@ dependencies = [
"generic-array",
]
[[package]]
name = "inspecting-allocator"
version = "1.8.0"
dependencies = [
"tracing",
"tracing-error",
]
[[package]]
name = "insta"
version = "1.34.0"
@ -3324,7 +3316,6 @@ dependencies = [
"http 0.2.11",
"index-scheduler",
"indexmap",
"inspecting-allocator",
"insta",
"is-terminal",
"itertools 0.11.0",
@ -3374,7 +3365,6 @@ dependencies = [
"toml",
"tracing",
"tracing-actix-web",
"tracing-error",
"tracing-subscriber",
"tracing-trace",
"url",

View File

@ -18,7 +18,7 @@ members = [
"fuzzers",
"tracing-trace",
"xtask",
"build-info", "inspecting-allocator",
"build-info",
]
[workspace.package]

View File

@ -1,15 +0,0 @@
[package]
name = "inspecting-allocator"
version.workspace = true
authors.workspace = true
description.workspace = true
homepage.workspace = true
readme.workspace = true
edition.workspace = true
license.workspace = true
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies]
tracing = "0.1.40"
tracing-error = { version = "0.2.0", default-features = false }

View File

@ -1,160 +0,0 @@
use std::alloc::GlobalAlloc;
use std::cell::{Cell, RefCell};
use std::collections::HashMap;
use std::sync::atomic::AtomicU64;
use tracing_error::SpanTrace;
#[derive(Debug, Clone)]
pub struct AllocEntry {
generation: u64,
span: SpanTrace,
}
impl AllocEntry {
pub fn generation(&self) -> u64 {
self.generation
}
}
impl std::fmt::Display for AllocEntry {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let mut res = Ok(());
let mut depth = 0;
self.span.with_spans(|metadata, fields| {
let name_with_module_name: Vec<&str> = metadata
.module_path()
.into_iter()
.chain(std::iter::once(metadata.name()))
.collect();
let name_with_module_name = name_with_module_name.join("::");
let location = format!(
"{}:{}",
metadata.file().unwrap_or_default(),
metadata.line().unwrap_or_default()
);
if let Err(error) =
writeln!(f, "[{depth}]{name_with_module_name}({fields}) at {location}")
{
res = Err(error);
return false;
}
depth += 1;
true
});
res
}
}
struct AllocatorState {
is_allocating: Cell<bool>,
state: RefCell<HashMap<*mut u8, AllocEntry>>,
}
thread_local! {
static ALLOCATOR_STATE: AllocatorState = AllocatorState { is_allocating: Cell::new(false), state: RefCell::new(Default::default()) };
}
pub struct InspectingAllocator<InnerAllocator> {
inner: InnerAllocator,
current_generation: AtomicU64,
}
impl AllocatorState {
fn handle_alloc(&self, allocated: *mut u8, current_generation: u64) -> *mut u8 {
if self.is_allocating.get() {
return allocated;
}
self.is_allocating.set(true);
{
self.state.borrow_mut().insert(
allocated,
AllocEntry { generation: current_generation, span: SpanTrace::capture() },
);
}
self.is_allocating.set(false);
allocated
}
fn handle_dealloc(&self, allocated: *mut u8) {
if self.is_allocating.get() {
return;
}
self.is_allocating.set(true);
{
self.state.borrow_mut().remove(&allocated);
}
self.is_allocating.set(false);
}
fn find_older_generations(&self, older_generation: u64) -> Vec<(*mut u8, AllocEntry)> {
if self.is_allocating.get() {
return Vec::new();
}
self.is_allocating.set(true);
let mut entries = Vec::new();
self.state.borrow_mut().retain(|k, v| {
if v.generation > older_generation {
return true;
}
entries.push((*k, v.clone()));
false
});
self.is_allocating.set(false);
entries
}
}
impl<A> InspectingAllocator<A> {
pub const fn wrap(inner: A) -> Self {
Self { inner, current_generation: AtomicU64::new(0) }
}
pub fn next_generation(&self) {
self.current_generation.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
}
pub fn find_older_generations(&self, older_than: u64) -> Vec<(*mut u8, AllocEntry)> {
let current_generation = self.current_generation.load(std::sync::atomic::Ordering::Relaxed);
if current_generation < older_than {
return Vec::new();
}
ALLOCATOR_STATE.with(|allocator_state| {
allocator_state.find_older_generations(current_generation - older_than)
})
}
}
unsafe impl<InnerAllocator: GlobalAlloc> GlobalAlloc for InspectingAllocator<InnerAllocator> {
unsafe fn alloc(&self, layout: std::alloc::Layout) -> *mut u8 {
let allocated = self.inner.alloc(layout);
let current_generation = self.current_generation.load(std::sync::atomic::Ordering::Relaxed);
ALLOCATOR_STATE
.with(|allocator_state| allocator_state.handle_alloc(allocated, current_generation))
}
unsafe fn dealloc(&self, ptr: *mut u8, layout: std::alloc::Layout) {
self.inner.dealloc(ptr, layout);
ALLOCATOR_STATE.with(|allocator_state| allocator_state.handle_dealloc(ptr))
}
unsafe fn alloc_zeroed(&self, layout: std::alloc::Layout) -> *mut u8 {
let allocated = self.inner.alloc_zeroed(layout);
let current_generation = self.current_generation.load(std::sync::atomic::Ordering::Relaxed);
ALLOCATOR_STATE
.with(|allocator_state| allocator_state.handle_alloc(allocated, current_generation))
}
unsafe fn realloc(&self, ptr: *mut u8, layout: std::alloc::Layout, new_size: usize) -> *mut u8 {
let reallocated = self.inner.realloc(ptr, layout, new_size);
if reallocated == ptr {
return reallocated;
}
let current_generation = self.current_generation.load(std::sync::atomic::Ordering::Relaxed);
ALLOCATOR_STATE.with(|allocator_state| allocator_state.handle_dealloc(ptr));
ALLOCATOR_STATE
.with(|allocator_state| allocator_state.handle_alloc(reallocated, current_generation))
}
}

View File

@ -108,8 +108,6 @@ tracing-subscriber = { version = "0.3.18", features = ["json"] }
tracing-trace = { version = "0.1.0", path = "../tracing-trace" }
tracing-actix-web = "0.7.9"
build-info = { version = "1.7.0", path = "../build-info" }
inspecting-allocator = { version = "1.8.0", path = "../inspecting-allocator" }
tracing-error = { version = "0.2.0", default-features = false }
[dev-dependencies]
actix-rt = "2.9.0"

View File

@ -16,11 +16,15 @@ use meilisearch::{
LogStderrType, Opt, SubscriberForSecondLayer,
};
use meilisearch_auth::{generate_master_key, AuthController, MASTER_KEY_MIN_SIZE};
use mimalloc::MiMalloc;
use termcolor::{Color, ColorChoice, ColorSpec, StandardStream, WriteColor};
use tracing::level_filters::LevelFilter;
use tracing_subscriber::layer::SubscriberExt as _;
use tracing_subscriber::Layer;
#[global_allocator]
static ALLOC: MiMalloc = MiMalloc;
fn default_log_route_layer() -> LogRouteType {
None.with_filter(tracing_subscriber::filter::Targets::new().with_target("", LevelFilter::OFF))
}
@ -52,10 +56,8 @@ fn setup(opt: &Opt) -> anyhow::Result<(LogRouteHandle, LogStderrHandle)> {
let (stderr_layer, stderr_layer_handle) =
tracing_subscriber::reload::Layer::new(default_log_stderr_layer(opt));
let route_layer: tracing_subscriber::reload::Layer<_, _> = route_layer;
let error_layer = tracing_error::ErrorLayer::default();
let subscriber =
tracing_subscriber::registry().with(route_layer).with(stderr_layer).with(error_layer);
let subscriber = tracing_subscriber::registry().with(route_layer).with(stderr_layer);
// set the subscriber as the default for the application
tracing::subscriber::set_global_default(subscriber).unwrap();

View File

@ -8,7 +8,6 @@ use deserr::actix_web::{AwebJson, AwebQueryParameter};
use deserr::Deserr;
use futures::StreamExt;
use index_scheduler::{IndexScheduler, TaskId};
use inspecting_allocator::InspectingAllocator;
use meilisearch_types::deserr::query_params::Param;
use meilisearch_types::deserr::{DeserrJsonError, DeserrQueryParamError};
use meilisearch_types::document_formats::{read_csv, read_json, read_ndjson, PayloadType};
@ -21,7 +20,6 @@ use meilisearch_types::milli::DocumentId;
use meilisearch_types::star_or::OptionStarOrList;
use meilisearch_types::tasks::KindWithContent;
use meilisearch_types::{milli, Document, Index};
use mimalloc::MiMalloc;
use mime::Mime;
use once_cell::sync::Lazy;
use serde::Deserialize;
@ -48,9 +46,6 @@ static ACCEPTED_CONTENT_TYPE: Lazy<Vec<String>> = Lazy::new(|| {
vec!["application/json".to_string(), "application/x-ndjson".to_string(), "text/csv".to_string()]
});
#[global_allocator]
static ALLOC: InspectingAllocator<MiMalloc> = InspectingAllocator::wrap(MiMalloc);
/// Extracts the mime type from the content type and return
/// a meilisearch error if anything bad happen.
fn extract_mime_type(req: &HttpRequest) -> Result<Option<Mime>, MeilisearchHttpError> {
@ -473,14 +468,6 @@ async fn document_addition(
};
let scheduler = index_scheduler.clone();
ALLOC.next_generation();
for (address, entry) in ALLOC.find_older_generations(5) {
println!(
"Found allocation older than 5 generations: {address:p} in generation {}. Span trace",
entry.generation()
);
println!("{entry}")
}
let task = match tokio::task::spawn_blocking(move || scheduler.register(task, task_id, dry_run))
.await?
{

View File

@ -367,6 +367,12 @@ async fn get_version(
})
}
#[derive(Serialize)]
struct KeysResponse {
private: Option<String>,
public: Option<String>,
}
pub async fn get_health(
index_scheduler: Data<IndexScheduler>,
auth_controller: Data<AuthController>,

View File

@ -301,14 +301,10 @@ impl From<EmbeddingConfig> for EmbeddingSettings {
fn from(value: EmbeddingConfig) -> Self {
let EmbeddingConfig { embedder_options, prompt } = value;
match embedder_options {
super::EmbedderOptions::HuggingFace(super::hf::EmbedderOptions {
model,
revision,
distribution,
}) => Self {
super::EmbedderOptions::HuggingFace(options) => Self {
source: Setting::Set(EmbedderSource::HuggingFace),
model: Setting::Set(model),
revision: revision.map(Setting::Set).unwrap_or_default(),
model: Setting::Set(options.model),
revision: options.revision.map(Setting::Set).unwrap_or_default(),
api_key: Setting::NotSet,
dimensions: Setting::NotSet,
document_template: Setting::Set(prompt.template),
@ -318,19 +314,14 @@ impl From<EmbeddingConfig> for EmbeddingSettings {
path_to_embeddings: Setting::NotSet,
embedding_object: Setting::NotSet,
input_type: Setting::NotSet,
distribution: distribution.map(Setting::Set).unwrap_or_default(),
distribution: options.distribution.map(Setting::Set).unwrap_or_default(),
},
super::EmbedderOptions::OpenAi(super::openai::EmbedderOptions {
api_key,
embedding_model,
dimensions,
distribution,
}) => Self {
super::EmbedderOptions::OpenAi(options) => Self {
source: Setting::Set(EmbedderSource::OpenAi),
model: Setting::Set(embedding_model.name().to_owned()),
model: Setting::Set(options.embedding_model.name().to_owned()),
revision: Setting::NotSet,
api_key: api_key.map(Setting::Set).unwrap_or_default(),
dimensions: dimensions.map(Setting::Set).unwrap_or_default(),
api_key: options.api_key.map(Setting::Set).unwrap_or_default(),
dimensions: options.dimensions.map(Setting::Set).unwrap_or_default(),
document_template: Setting::Set(prompt.template),
url: Setting::NotSet,
query: Setting::NotSet,
@ -338,37 +329,29 @@ impl From<EmbeddingConfig> for EmbeddingSettings {
path_to_embeddings: Setting::NotSet,
embedding_object: Setting::NotSet,
input_type: Setting::NotSet,
distribution: distribution.map(Setting::Set).unwrap_or_default(),
distribution: options.distribution.map(Setting::Set).unwrap_or_default(),
},
super::EmbedderOptions::Ollama(super::ollama::EmbedderOptions {
embedding_model,
url,
api_key,
distribution,
}) => Self {
super::EmbedderOptions::Ollama(options) => Self {
source: Setting::Set(EmbedderSource::Ollama),
model: Setting::Set(embedding_model),
model: Setting::Set(options.embedding_model.to_owned()),
revision: Setting::NotSet,
api_key: api_key.map(Setting::Set).unwrap_or_default(),
api_key: options.api_key.map(Setting::Set).unwrap_or_default(),
dimensions: Setting::NotSet,
document_template: Setting::Set(prompt.template),
url: url.map(Setting::Set).unwrap_or_default(),
url: options.url.map(Setting::Set).unwrap_or_default(),
query: Setting::NotSet,
input_field: Setting::NotSet,
path_to_embeddings: Setting::NotSet,
embedding_object: Setting::NotSet,
input_type: Setting::NotSet,
distribution: distribution.map(Setting::Set).unwrap_or_default(),
distribution: options.distribution.map(Setting::Set).unwrap_or_default(),
},
super::EmbedderOptions::UserProvided(super::manual::EmbedderOptions {
dimensions,
distribution,
}) => Self {
super::EmbedderOptions::UserProvided(options) => Self {
source: Setting::Set(EmbedderSource::UserProvided),
model: Setting::NotSet,
revision: Setting::NotSet,
api_key: Setting::NotSet,
dimensions: Setting::Set(dimensions),
dimensions: Setting::Set(options.dimensions),
document_template: Setting::NotSet,
url: Setting::NotSet,
query: Setting::NotSet,
@ -376,7 +359,7 @@ impl From<EmbeddingConfig> for EmbeddingSettings {
path_to_embeddings: Setting::NotSet,
embedding_object: Setting::NotSet,
input_type: Setting::NotSet,
distribution: distribution.map(Setting::Set).unwrap_or_default(),
distribution: options.distribution.map(Setting::Set).unwrap_or_default(),
},
super::EmbedderOptions::Rest(super::rest::EmbedderOptions {
api_key,