diff --git a/include/collection.h b/include/collection.h index 6e20beb6..90586fc1 100644 --- a/include/collection.h +++ b/include/collection.h @@ -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 get_embedding_fields(); + tsl::htrie_map get_embedding_fields_unsafe(); + std::string get_default_sorting_field(); Option to_doc(const std::string& json_str, nlohmann::json& document, diff --git a/include/collection_manager.h b/include/collection_manager.h index 03ef6b43..b92f29e5 100644 --- a/include/collection_manager.h +++ b/include/collection_manager.h @@ -201,4 +201,6 @@ public: Option upsert_preset(const std::string & preset_name, const nlohmann::json& preset_config); Option delete_preset(const std::string & preset_name); + + void process_embedding_field_delete(const std::string& model_name); }; \ No newline at end of file diff --git a/src/collection.cpp b/src/collection.cpp index 0ddd0bcf..adda2fc5 100644 --- a/src/collection.cpp +++ b/src/collection.cpp @@ -3797,6 +3797,15 @@ Option Collection::batch_alter_data(const std::vector& 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(); + 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(validate_model_res.code(), validate_model_res.error()); + } + } embedding_fields.emplace(f.name, f); } @@ -3908,7 +3917,7 @@ Option Collection::batch_alter_data(const std::vector& 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 Collection::truncate_after_top_k(const string &field_name, size_t k return Option(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(); + embedding_fields.erase(field_name); + CollectionManager::get_instance().process_embedding_field_delete(model_name); +} + +tsl::htrie_map Collection::get_embedding_fields_unsafe() { + return embedding_fields; +} \ No newline at end of file diff --git a/src/collection_manager.cpp b/src/collection_manager.cpp index 657a906a..98cbc50b 100644 --- a/src/collection_manager.cpp +++ b/src/collection_manager.cpp @@ -531,7 +531,15 @@ Option 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(); + 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 CollectionManager::clone_collection(const string& existing_n return Option(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() == model_name) { + found = true; + break; + } + } + } + } + + if(!found) { + LOG(INFO) << "Deleting text embedder: " << model_name; + TextEmbedderManager::get_instance().delete_text_embedder(model_name); + } +} \ No newline at end of file diff --git a/test/collection_vector_search_test.cpp b/test/collection_vector_search_test.cpp index 4eb84853..55049641 100644 --- a/test/collection_vector_search_test.cpp +++ b/test/collection_vector_search_test.cpp @@ -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()); } \ No newline at end of file