postgres/typesense
1
0
Fork 0
mirror of https://github.com/typesense/typesense.git synced 2025-05-21 14:12:27 +08:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

Merge pull request #1282 from ozanarmagan/v0.25-join

Browse Source 
Unload model from RAM when it is unused
This commit is contained in:
Kishore Nallan 2023-10-09 19:10:50 +05:30 committed by GitHub
commit 37b3663ee1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 334 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
include/collection.h
View File

@ -278,6 +278,8 @@ private:
static void hide_credential(nlohmann::json& json, const std::string& credential_name);
void remove_embedding_field(const std::string& field_name);
public:
enum {MAX_ARRAY_MATCHES = 5};
@ -352,6 +354,8 @@ public:
tsl::htrie_map<char, field> get_embedding_fields();
tsl::htrie_map<char, field> get_embedding_fields_unsafe();
std::string get_default_sorting_field();
Option<doc_seq_id_t> to_doc(const std::string& json_str, nlohmann::json& document,

2
include/collection_manager.h
View File

@ -201,4 +201,6 @@ public:
Option<bool> upsert_preset(const std::string & preset_name, const nlohmann::json& preset_config);
Option<bool> delete_preset(const std::string & preset_name);
void process_embedding_field_delete(const std::string& model_name);
};

28
src/collection.cpp
View File

@ -3797,6 +3797,15 @@ Option<bool> Collection::batch_alter_data(const std::vector<field>& alter_fields
if(f.embed.count(fields::from) != 0) {
found_embedding_field = true;
const auto& text_embedders = TextEmbedderManager::get_instance()._get_text_embedders();
const auto& model_name = f.embed[fields::model_config][fields::model_name].get<std::string>();
if(text_embedders.count(model_name) == 0) {
size_t dummy_num_dim = 0;
auto validate_model_res = TextEmbedderManager::get_instance().validate_and_init_model(f.embed[fields::model_config], dummy_num_dim);
if(!validate_model_res.ok()) {
return Option<bool>(validate_model_res.code(), validate_model_res.error());
}
}
embedding_fields.emplace(f.name, f);
}
@ -3908,7 +3917,7 @@ Option<bool> Collection::batch_alter_data(const std::vector<field>& alter_fields
}
if(del_field.embed.count(fields::from) != 0) {
embedding_fields.erase(del_field.name);
remove_embedding_field(del_field.name);
}
if(del_field.name == ".*") {
@ -4983,7 +4992,7 @@ void Collection::process_remove_field_for_embedding_fields(const field& del_fiel
}
for(auto& garbage_field: garbage_embed_fields) {
embedding_fields.erase(garbage_field.name);
remove_embedding_field(garbage_field.name);
search_schema.erase(garbage_field.name);
fields.erase(std::remove_if(fields.begin(), fields.end(), [&garbage_field](const auto &f) {
return f.name == garbage_field.name;
@ -5025,3 +5034,18 @@ Option<bool> Collection::truncate_after_top_k(const string &field_name, size_t k
return Option<bool>(true);
}
void Collection::remove_embedding_field(const std::string& field_name) {
if(embedding_fields.find(field_name) == embedding_fields.end()) {
return;
}
const auto& del_field = embedding_fields[field_name];
const auto& model_name = del_field.embed[fields::model_config]["model_name"].get<std::string>();
embedding_fields.erase(field_name);
CollectionManager::get_instance().process_embedding_field_delete(model_name);
}
tsl::htrie_map<char, field> Collection::get_embedding_fields_unsafe() {
return embedding_fields;
}

34
src/collection_manager.cpp
View File

@ -531,7 +531,15 @@ Option<nlohmann::json> CollectionManager::drop_collection(const std::string& col
std::unique_lock u_lock(mutex);
collections.erase(actual_coll_name);
collection_id_names.erase(collection->get_collection_id());
const auto& embedding_fields = collection->get_embedding_fields();
u_lock.unlock();
for(const auto& embedding_field : embedding_fields) {
const auto& model_name = embedding_field.embed[fields::model_config]["model_name"].get<std::string>();
process_embedding_field_delete(model_name);
}
// don't hold any collection manager locks here, since this can take some time
delete collection;
@ -1555,3 +1563,29 @@ Option<Collection*> CollectionManager::clone_collection(const string& existing_n
return Option<Collection*>(new_coll);
}
void CollectionManager::process_embedding_field_delete(const std::string& model_name) {
std::shared_lock lock(mutex);
bool found = false;
for(const auto& collection: collections) {
// will be deadlock if we try to acquire lock on collection here
// caller of this function should have already acquired lock on collection
const auto& embedding_fields = collection.second->get_embedding_fields_unsafe();
for(const auto& embedding_field: embedding_fields) {
if(embedding_field.embed.count(fields::model_config) != 0) {
const auto& model_config = embedding_field.embed[fields::model_config];
if(model_config["model_name"].get<std::string>() == model_name) {
found = true;
break;
}
}
}
}
if(!found) {
LOG(INFO) << "Deleting text embedder: " << model_name;
TextEmbedderManager::get_instance().delete_text_embedder(model_name);
}
}

268
test/collection_vector_search_test.cpp
View File

@ -31,6 +31,7 @@ protected:
virtual void TearDown() {
collectionManager.dispose();
TextEmbedderManager::get_instance().delete_all_text_embedders();
delete store;
}
};
@ -2229,4 +2230,271 @@ TEST_F(CollectionVectorTest, QueryByNotAutoEmbeddingVectorField) {
ASSERT_FALSE(search_res.ok());
ASSERT_EQ("Vector field `embedding` is not an auto-embedding field, do not use `query_by` with it, use `vector_query` instead.", search_res.error());
}
TEST_F(CollectionVectorTest, TestUnloadingModelsOnCollectionDelete) {
nlohmann::json actual_schema = R"({
"name": "test",
"fields": [
{
"name": "title",
"type": "string"
},
{
"name": "title_vec",
"type": "float[]",
"embed": {
"from": [
"title"
],
"model_config": {
"model_name": "ts/e5-small"
}
}
}
]
})"_json;
TextEmbedderManager::set_model_dir("/tmp/typesense_test/models");
auto schema = actual_schema;
auto collection_create_op = collectionManager.create_collection(schema);
ASSERT_TRUE(collection_create_op.ok());
auto coll = collection_create_op.get();
auto text_embedders = TextEmbedderManager::get_instance()._get_text_embedders();
ASSERT_EQ(1, text_embedders.size());
auto delete_op = collectionManager.drop_collection("test", true);
ASSERT_TRUE(delete_op.ok());
text_embedders = TextEmbedderManager::get_instance()._get_text_embedders();
ASSERT_EQ(0, text_embedders.size());
// create another collection
schema = actual_schema;
collection_create_op = collectionManager.create_collection(schema);
ASSERT_TRUE(collection_create_op.ok());
coll = collection_create_op.get();
text_embedders = TextEmbedderManager::get_instance()._get_text_embedders();
ASSERT_EQ(1, text_embedders.size());
// create second collection
schema = actual_schema;
schema["name"] = "test2";
collection_create_op = collectionManager.create_collection(schema);
ASSERT_TRUE(collection_create_op.ok());
auto coll2 = collection_create_op.get();
text_embedders = TextEmbedderManager::get_instance()._get_text_embedders();
ASSERT_EQ(1, text_embedders.size());
delete_op = collectionManager.drop_collection("test", true);
ASSERT_TRUE(delete_op.ok());
text_embedders = TextEmbedderManager::get_instance()._get_text_embedders();
ASSERT_EQ(1, text_embedders.size());
delete_op = collectionManager.drop_collection("test2", true);
ASSERT_TRUE(delete_op.ok());
text_embedders = TextEmbedderManager::get_instance()._get_text_embedders();
ASSERT_EQ(0, text_embedders.size());
}
TEST_F(CollectionVectorTest, TestUnloadingModelsOnDrop) {
nlohmann::json actual_schema = R"({
"name": "test",
"fields": [
{
"name": "title",
"type": "string"
},
{
"name": "title_vec",
"type": "float[]",
"embed": {
"from": [
"title"
],
"model_config": {
"model_name": "ts/e5-small"
}
}
}
]
})"_json;
TextEmbedderManager::set_model_dir("/tmp/typesense_test/models");
auto schema = actual_schema;
auto collection_create_op = collectionManager.create_collection(schema);
ASSERT_TRUE(collection_create_op.ok());
auto coll = collection_create_op.get();
auto text_embedders = TextEmbedderManager::get_instance()._get_text_embedders();
ASSERT_EQ(1, text_embedders.size());
nlohmann::json drop_schema = R"({
"fields": [
{
"name": "title_vec",
"drop": true
}
]
})"_json;
auto drop_op = coll->alter(drop_schema);
ASSERT_TRUE(drop_op.ok());
text_embedders = TextEmbedderManager::get_instance()._get_text_embedders();
ASSERT_EQ(0, text_embedders.size());
// create another collection
schema = actual_schema;
schema["name"] = "test2";
collection_create_op = collectionManager.create_collection(schema);
ASSERT_TRUE(collection_create_op.ok());
auto coll2 = collection_create_op.get();
nlohmann::json alter_schema = R"({
"fields": [
{
"name": "title_vec",
"type": "float[]",
"embed": {
"from": [
"title"
],
"model_config": {
"model_name": "ts/e5-small"
}
}
}
]
})"_json;
auto alter_op = coll->alter(alter_schema);
ASSERT_TRUE(alter_op.ok());
text_embedders = TextEmbedderManager::get_instance()._get_text_embedders();
ASSERT_EQ(1, text_embedders.size());
drop_op = coll2->alter(drop_schema);
ASSERT_TRUE(drop_op.ok());
text_embedders = TextEmbedderManager::get_instance()._get_text_embedders();
ASSERT_EQ(1, text_embedders.size());
drop_op = coll->alter(drop_schema);
ASSERT_TRUE(drop_op.ok());
text_embedders = TextEmbedderManager::get_instance()._get_text_embedders();
ASSERT_EQ(0, text_embedders.size());
}
TEST_F(CollectionVectorTest, TestUnloadModelsCollectionHaveTwoEmbeddingField) {
nlohmann::json actual_schema = R"({
"name": "test",
"fields": [
{
"name": "title",
"type": "string"
},
{
"name": "title_vec",
"type": "float[]",
"embed": {
"from": [
"title"
],
"model_config": {
"model_name": "ts/e5-small"
}
}
},
{
"name": "title_vec2",
"type": "float[]",
"embed": {
"from": [
"title"
],
"model_config": {
"model_name": "ts/e5-small"
}
}
}
]
})"_json;
TextEmbedderManager::set_model_dir("/tmp/typesense_test/models");
auto schema = actual_schema;
auto collection_create_op = collectionManager.create_collection(schema);
ASSERT_TRUE(collection_create_op.ok());
auto coll = collection_create_op.get();
auto text_embedders = TextEmbedderManager::get_instance()._get_text_embedders();
ASSERT_EQ(1, text_embedders.size());
nlohmann::json drop_schema = R"({
"fields": [
{
"name": "title_vec",
"drop": true
}
]
})"_json;
auto alter_op = coll->alter(drop_schema);
ASSERT_TRUE(alter_op.ok());
text_embedders = TextEmbedderManager::get_instance()._get_text_embedders();
ASSERT_EQ(1, text_embedders.size());
drop_schema = R"({
"fields": [
{
"name": "title_vec2",
"drop": true
}
]
})"_json;
alter_op = coll->alter(drop_schema);
ASSERT_TRUE(alter_op.ok());
text_embedders = TextEmbedderManager::get_instance()._get_text_embedders();
ASSERT_EQ(0, text_embedders.size());
// create another collection
schema = actual_schema;
schema["name"] = "test2";
collection_create_op = collectionManager.create_collection(schema);
ASSERT_TRUE(collection_create_op.ok());
auto coll2 = collection_create_op.get();
text_embedders = TextEmbedderManager::get_instance()._get_text_embedders();
ASSERT_EQ(1, text_embedders.size());
// drop collection
auto drop_op = collectionManager.drop_collection("test2", true);
ASSERT_TRUE(drop_op.ok());
text_embedders = TextEmbedderManager::get_instance()._get_text_embedders();
ASSERT_EQ(0, text_embedders.size());
}