Added tests to get started

morpheus/service/milvus_vector_db_service.py

@@ -103,7 +103,6 @@ def _create_schema_field(self, field_conf: dict) -> pymilvus.FieldSchema:
 
         return field_schema
 
-    @with_mutex("_mutex")
     def create(self, name: str, overwrite: bool = False, **kwargs: dict[str, typing.Any]):
         collection_conf = kwargs.get("collection_conf")
         auto_id = collection_conf.get("auto_id", False)
@@ -268,15 +267,15 @@ def search(self, name: str, query: typing.Union[str, dict] = None, **kwargs: dic
             self._handler.release_collection(collection_name=name)
 
     @with_mutex("_mutex")
-    def update(self, name: str, data: typing.Any, **kwargs: dict[str, typing.Any]) -> typing.Any:
+    def update(self, name: str, data: list[dict], **kwargs: dict[str, typing.Any]) -> typing.Any:
         """
         Update data in the vector database.
 
         Parameters
         ----------
         name : str
             Name of the resource.
-        data : typing.Any
+        data : list[dict]
             Data to be updated in the resource.
         **kwargs : dict[str, typing.Any]
             Extra keyword arguments specific to upsert operation.

morpheus/service/vector_db_service.py

@@ -53,7 +53,8 @@ def insert(self, name: str, data: typing.Any, **kwargs: dict[str, typing.Any]) -
         pass
 
     @abstractmethod
-    def insert_dataframe(self, name: str,
+    def insert_dataframe(self,
+                         name: str,
                          df: typing.Union[cudf.DataFrame, pd.DataFrame],
                          **kwargs: dict[str, typing.Any]):
         """

morpheus/stages/output/write_to_vector_db.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
 import logging
 import typing
 
@@ -41,6 +40,8 @@ class WriteToVectorDBStage(SinglePortStage):
         Pipeline configuration instance.
     resource_name : str
         The name of the resource managed by this instance.
+    resource_conf : dict
+        Additional resource configuration when performing vector database writes.
     vdb_service : typing.Union[str, VectorDBService]
         Either the name of the vector database service to use or an instance of VectorDBService
         for managing the resource.
@@ -52,22 +53,30 @@ class WriteToVectorDBStage(SinglePortStage):
     ValueError
         If `vdb_service` is not a valid string (service name) or an instance of VectorDBService.
     """
-    def __init__(self, config: Config, resource_name: str,
+
+    def __init__(self,
+                 config: Config,
+                 resource_name: str,
                  vdb_service: typing.Union[str, VectorDBService],
                  **kwargs: dict[str, typing.Any]):
 
         super().__init__(config)
+
         self._resource_name = resource_name
+        self._resource_conf = {}
+
+        if "resource_conf" in kwargs:
+            self._resource_conf = kwargs.pop("resource_conf")
 
         if isinstance(vdb_service, str):
-            # If vdb_input is a string, assume it's the service name
+            # If vdb_service is a string, assume it's the service name
             self._vdb_service: VectorDBService = VectorDBServiceFactory.create_instance(service_name=vdb_service,
                                                                                         **kwargs)
         elif isinstance(vdb_service, VectorDBService):
-            # If vdb_input is an instance of VectorDBService, use it directly
+            # If vdb_service is an instance of VectorDBService, use it directly
             self._vdb_service: VectorDBService = vdb_service
         else:
-            raise ValueError("vdb_input must be a string (service name) or an instance of VectorDBService")
+            raise ValueError("vdb_service must be a string (service name) or an instance of VectorDBService")
 
     @property
     def name(self) -> str:
@@ -89,17 +98,24 @@ def supports_cpp_node(self):
         """Indicates whether this stage supports a C++ node."""
         return False
 
+    def on_completed(self):
+        # Close vector database service connection
+        self._vdb_service.close()
+
     def _build_single(self, builder: mrc.Builder, input_stream: StreamPair) -> StreamPair:
 
         stream = input_stream[0]
 
         def on_data(ctrl_msg: ControlMessage) -> ControlMessage:
-            result = self._vdb_service.insert_dataframe(name=self._resource_name, df=ctrl_msg.payload().df)
+            # Insert entries in the dataframe to vector database.
+            result = self._vdb_service.insert_dataframe(name=self._resource_name,
+                                                        df=ctrl_msg.payload().df,
+                                                        **self._resource_conf)
             ctrl_msg.set_metadata("insert_response", result)
 
             return ctrl_msg
 
-        to_vector_db = builder.make_node(self.unique_name, ops.map(on_data))
+        to_vector_db = builder.make_node(self.unique_name, ops.map(on_data), ops.on_completed(self.on_completed))
 
         builder.make_edge(stream, to_vector_db)
         stream = to_vector_db

tests/test_write_to_vector_db_stage_pipe.py

@@ -32,63 +32,177 @@
 from morpheus.utils.module_ids import TO_CONTROL_MESSAGE
 
 
+@pytest.fixture(scope="function", name="milvus_vdb_serivce_fixture")
+def milvus_vdb_serivce():
+    vdb_service = MilvusVectorDBService(uri="http://localhost:19530")
+    return vdb_service
+
+
+def create_milvus_collection(vdb_service: MilvusVectorDBService):
+    collection_config = {
+        "collection_conf": {
+            "shards": 2,
+            "auto_id": False,
+            "consistency_level": "Strong",
+            "description": "Test collection",
+            "schema_conf": {
+                "enable_dynamic_field": True,
+                "schema_fields": [
+                    {
+                        "name": "id",
+                        "dtype": "int64",
+                        "description": "Primary key for the collection",
+                        "is_primary": True,
+                    },
+                    {
+                        "name": "embedding",
+                        "dtype": "float_vector",
+                        "description": "Embedding vectors",
+                        "is_primary": False,
+                        "dim": 10,
+                    },
+                    {
+                        "name": "age",
+                        "dtype": "int64",
+                        "description ": "Age",
+                        "is_primary": False,
+                    },
+                ],
+                "description": "Test collection schema",
+            },
+        }
+    }
+    vdb_service.create(name="test", overwrite=True, **collection_config)
+
+
+def create_milvus_collection_idx_part(vdb_service: MilvusVectorDBService):
+    collection_config = {
+        "collection_conf": {
+            "shards": 2,
+            "auto_id": False,
+            "consistency_level": "Strong",
+            "description": "Test collection with partition and index",
+            "index_conf": {
+                "field_name": "embedding", "metric_type": "L2"
+            },
+            "partition_conf": {
+                "timeout": 1, "partitions": [{
+                    "name": "age_partition", "description": "Partition by age"
+                }]
+            },
+            "schema_conf": {
+                "enable_dynamic_field": True,
+                "schema_fields": [
+                    {
+                        "name": "id",
+                        "dtype": "int64",
+                        "description": "Primary key for the collection",
+                        "is_primary": True,
+                    },
+                    {
+                        "name": "embedding",
+                        "dtype": "float_vector",
+                        "description": "Embedding vectors",
+                        "is_primary": False,
+                        "dim": 10,
+                    },
+                    {
+                        "name": "age",
+                        "dtype": "int64",
+                        "description ": "Age",
+                        "is_primary": False,
+                    },
+                ],
+                "description": "Test collection schema",
+            },
+        }
+    }
+    vdb_service.create(name="test_idx_part", overwrite=True, **collection_config)
+
+
 @pytest.mark.use_cpp
-def test_write_to_vector_db_stage_with_instance_pipe(config: Config, pipeline_batch_size: int = 256):
+def test_write_to_vector_db_stage_with_instance_pipe(milvus_vdb_serivce_fixture,
+                                                     config: Config,
+                                                     pipeline_batch_size: int = 256):
     config.pipeline_batch_size = pipeline_batch_size
 
     rows_count = 5
     dimensions = 10
+    collection_name = "test"
+
+    create_milvus_collection(milvus_vdb_serivce_fixture)
+
     df = cudf.DataFrame({
-            "id": [i for i in range(rows_count)],
-            "age": [random.randint(20, 40) for i in range(rows_count)],
-            "embedding": [[random.random() for _ in range(dimensions)] for _ in range(rows_count)]
-        })
+        "id": [i for i in range(rows_count)],
+        "age": [random.randint(20, 40) for i in range(rows_count)],
+        "embedding": [[random.random() for _ in range(dimensions)] for _ in range(rows_count)]
+    })
 
     to_cm_module_config = {
-        "module_id": TO_CONTROL_MESSAGE,
-        "module_name": "to_control_message",
-        "namespace": MORPHEUS_MODULE_NAMESPACE
+        "module_id": TO_CONTROL_MESSAGE, "module_name": "to_control_message", "namespace": MORPHEUS_MODULE_NAMESPACE
     }
     vdb_service = MilvusVectorDBService(uri="http://localhost:19530")
 
     pipe = LinearPipeline(config)
     pipe.set_source(InMemorySourceStage(config, [df]))
-    pipe.add_stage(LinearModulesStage(config, to_cm_module_config,
-                                      input_port_name="input",
-                                      output_port_name="output",
-                                      output_type=ControlMessage))
-    pipe.add_stage(WriteToVectorDBStage(config, resource_name="test", vdb_service=vdb_service))
+    pipe.add_stage(
+        LinearModulesStage(config,
+                           to_cm_module_config,
+                           input_port_name="input",
+                           output_port_name="output",
+                           output_type=ControlMessage))
+    pipe.add_stage(WriteToVectorDBStage(config, resource_name=collection_name, vdb_service=vdb_service))
 
     pipe.run()
 
+    actual_count = milvus_vdb_serivce_fixture.count(name=collection_name)
+    milvus_vdb_serivce_fixture.close()
+
+    assert actual_count == 5
+
 
 @pytest.mark.use_cpp
-def test_write_to_vector_db_stage_with_name_pipe(config: Config, pipeline_batch_size: int = 256):
+def test_write_to_vector_db_stage_with_name_pipe(milvus_vdb_serivce_fixture,
+                                                 config: Config,
+                                                 pipeline_batch_size: int = 256):
     config.pipeline_batch_size = pipeline_batch_size
 
+    create_milvus_collection_idx_part(milvus_vdb_serivce_fixture)
+
     rows_count = 5
     dimensions = 10
+    collection_name = "test_idx_part"
+
+    resource_conf = {"collection_conf": {"partition_name": "age_partition"}}
+
     df = cudf.DataFrame({
-            "id": [i for i in range(rows_count)],
-            "age": [random.randint(20, 40) for i in range(rows_count)],
-            "embedding": [[random.random() for _ in range(dimensions)] for _ in range(rows_count)]
-        })
+        "id": [i for i in range(rows_count)],
+        "age": [random.randint(20, 40) for i in range(rows_count)],
+        "embedding": [[random.random() for _ in range(dimensions)] for _ in range(rows_count)]
+    })
 
     to_cm_module_config = {
-        "module_id": TO_CONTROL_MESSAGE,
-        "module_name": "to_control_message",
-        "namespace": MORPHEUS_MODULE_NAMESPACE
+        "module_id": TO_CONTROL_MESSAGE, "module_name": "to_control_message", "namespace": MORPHEUS_MODULE_NAMESPACE
     }
 
     pipe = LinearPipeline(config)
-    pipe.set_source(InMemorySourceStage(config, [df]))
-    pipe.add_stage(LinearModulesStage(config, to_cm_module_config,
-                                      input_port_name="input",
-                                      output_port_name="output",
-                                      output_type=ControlMessage))
-    pipe.add_stage(WriteToVectorDBStage(config,
-                                        resource_name="test_collection",
-                                        vdb_service="milvus",
-                                        uri="http://localhost:19530"))
+    pipe.set_source(InMemorySourceStage(config, [df], repeat=2))
+    pipe.add_stage(
+        LinearModulesStage(config,
+                           to_cm_module_config,
+                           input_port_name="input",
+                           output_port_name="output",
+                           output_type=ControlMessage))
+    pipe.add_stage(
+        WriteToVectorDBStage(config,
+                             resource_name=collection_name,
+                             vdb_service="milvus",
+                             uri="http://localhost:19530",
+                             resource_conf=resource_conf))
 
     pipe.run()
+
+    actual_count = milvus_vdb_serivce_fixture.count(name=collection_name)
+    milvus_vdb_serivce_fixture.close()
+
+    assert actual_count == 10