Fix Milvus ML tests

Author: aIbrahiim

sdks/python/apache_beam/ml/rag/enrichment/milvus_search_it_test.py

@@ -279,6 +279,7 @@ def test_invalid_query_on_non_existent_collection(self):
 
     with self.assertRaises(Exception) as context:
       with TestPipeline() as p:
+        p.not_use_test_runner_api = True
         _ = (p | beam.Create(test_chunks) | Enrichment(handler))
 
     expect_err_msg_contains = "collection not found"
@@ -307,6 +308,7 @@ def test_invalid_query_on_non_existent_field(self):
 
     with self.assertRaises(Exception) as context:
       with TestPipeline() as p:
+        p.not_use_test_runner_api = True
         _ = (p | beam.Create(test_chunks) | Enrichment(handler))
 
     expect_err_msg_contains = f"fieldName({non_existent_field}) not found"
@@ -330,6 +332,7 @@ def test_empty_input_chunks(self):
     expected_chunks = []
 
     with TestPipeline() as p:
+      p.not_use_test_runner_api = True
       result = (p | beam.Create(test_chunks) | Enrichment(handler))
       assert_that(
           result, lambda actual: MilvusTestHelpers.assert_chunks_equivalent(
@@ -458,6 +461,7 @@ def test_filtered_search_with_cosine_similarity_and_batching(self):
     ]
 
     with TestPipeline() as p:
+      p.not_use_test_runner_api = True
       result = (p | beam.Create(test_chunks) | Enrichment(handler))
       assert_that(
           result, lambda actual: MilvusTestHelpers.assert_chunks_equivalent(
@@ -563,6 +567,7 @@ def test_filtered_search_with_bm25_full_text_and_batching(self):
     ]
 
     with TestPipeline() as p:
+      p.not_use_test_runner_api = True
       result = (p | beam.Create(test_chunks) | Enrichment(handler))
       assert_that(
           result, lambda actual: MilvusTestHelpers.assert_chunks_equivalent(
@@ -704,6 +709,7 @@ def test_vector_search_with_euclidean_distance(self):
     ]
 
     with TestPipeline() as p:
+      p.not_use_test_runner_api = True
       result = (p | beam.Create(test_chunks) | Enrichment(handler))
       assert_that(
           result, lambda actual: MilvusTestHelpers.assert_chunks_equivalent(
@@ -844,6 +850,7 @@ def test_vector_search_with_inner_product_similarity(self):
     ]
 
     with TestPipeline() as p:
+      p.not_use_test_runner_api = True
       result = (p | beam.Create(test_chunks) | Enrichment(handler))
       assert_that(
           result, lambda actual: MilvusTestHelpers.assert_chunks_equivalent(
@@ -909,6 +916,7 @@ def test_keyword_search_with_inner_product_sparse_embedding(self):
     ]
 
     with TestPipeline() as p:
+      p.not_use_test_runner_api = True
       result = (p | beam.Create(test_chunks) | Enrichment(handler))
       assert_that(
           result, lambda actual: MilvusTestHelpers.assert_chunks_equivalent(
@@ -982,6 +990,7 @@ def test_hybrid_search(self):
     ]
 
     with TestPipeline() as p:
+      p.not_use_test_runner_api = True
       result = (p | beam.Create(test_chunks) | Enrichment(handler))
       assert_that(
           result, lambda actual: MilvusTestHelpers.assert_chunks_equivalent(

sdks/python/apache_beam/ml/rag/ingestion/milvus_search_it_test.py

@@ -192,7 +192,6 @@ def tearDownClass(cls):
 
   def setUp(self):
     self.write_test_pipeline = TestPipeline()
-    self.write_test_pipeline.not_use_test_runner_api = True
     self._collection_name = f"test_collection_{self._testMethodName}"
     self._partition_name = f"test_partition_{self._testMethodName}"
     config = unpack_dataclass_with_kwargs(self._connection_config)
@@ -232,6 +231,7 @@ def test_invalid_write_on_non_existent_collection(self):
     # Write pipeline.
     with self.assertRaises(Exception) as context:
       with TestPipeline() as p:
+        p.not_use_test_runner_api = True
         _ = (p | beam.Create(test_chunks) | config.create_write_transform())
 
     # Assert on what should happen.
@@ -252,6 +252,7 @@ def test_invalid_write_on_non_existent_partition(self):
     # Write pipeline.
     with self.assertRaises(Exception) as context:
       with TestPipeline() as p:
+        p.not_use_test_runner_api = True
         _ = (p | beam.Create(test_chunks) | config.create_write_transform())
 
     # Assert on what should happen.
@@ -287,6 +288,7 @@ def test_invalid_write_on_missing_primary_key_in_entity(self):
     # Write pipeline.
     with self.assertRaises(Exception) as context:
       with TestPipeline() as p:
+        p.not_use_test_runner_api = True
         _ = (p | beam.Create(test_chunks) | config.create_write_transform())
 
     # Assert on what should happen.
@@ -332,6 +334,7 @@ def test_write_on_auto_id_primary_key(self):
     config = MilvusVectorWriterConfig(
         connection_params=self._connection_config, write_config=write_config)
 
+    self.write_test_pipeline.not_use_test_runner_api = True
     with self.write_test_pipeline as p:
       _ = (p | beam.Create(test_chunks) | config.create_write_transform())
 
@@ -357,6 +360,7 @@ def test_write_on_existent_collection_with_default_schema(self):
     config = MilvusVectorWriterConfig(
         connection_params=self._connection_config, write_config=write_config)
 
+    self.write_test_pipeline.not_use_test_runner_api = True
     with self.write_test_pipeline as p:
       _ = (p | beam.Create(test_chunks) | config.create_write_transform())
 
@@ -422,6 +426,7 @@ def test_write_with_custom_column_specifications(self):
         write_config=write_config,
         column_specs=custom_column_specs)
 
+    self.write_test_pipeline.not_use_test_runner_api = True
     with self.write_test_pipeline as p:
       _ = (p | beam.Create(test_chunks) | config.create_write_transform())
 
@@ -474,6 +479,7 @@ def test_write_with_batching(self):
     config = MilvusVectorWriterConfig(
         connection_params=self._connection_config, write_config=write_config)
 
+    self.write_test_pipeline.not_use_test_runner_api = True
     with self.write_test_pipeline as p:
       _ = (p | beam.Create(test_chunks) | config.create_write_transform())
 

sdks/python/apache_beam/ml/rag/test_utils.py

@@ -87,7 +87,8 @@ class CustomMilvusContainer(MilvusContainer):
   Extends MilvusContainer to provide custom port binding and environment
   configuration for testing with standalone Milvus instances.
   """
-  def __init__(
+
+  def __init__(  # pylint: disable=bad-super-call
       self,
       image: str,
       service_container_port,
@@ -96,7 +97,11 @@ def __init__(
   ) -> None:
     # Skip the parent class's constructor and go straight to
     # GenericContainer.
-    super(MilvusContainer, self).__init__(image=image, **kwargs)
+    super(
+        MilvusContainer,
+        self,
+    ).__init__(
+        image=image, **kwargs)
     self.port = service_container_port
     self.healthcheck_port = healthcheck_container_port
     self.with_exposed_ports(service_container_port, healthcheck_container_port)
@@ -133,6 +138,27 @@ class MilvusTestHelpers:
   # https://milvus.io/docs/release_notes.md or PyPI at
   # https://pypi.org/project/pymilvus/ for version compatibility.
   # Example: Milvus v2.6.0 requires pymilvus==2.6.0 (exact match required).
+  @staticmethod
+  def _wait_for_milvus_grpc(uri: str) -> None:
+    """Wait until Milvus accepts RPCs.
+
+    Docker may report started before gRPC is ready.
+    """
+    def list_collections_probe():
+      client = MilvusClient(uri=uri)
+      try:
+        client.list_collections()
+      finally:
+        client.close()
+
+    retry_with_backoff(
+        list_collections_probe,
+        max_retries=25,
+        retry_delay=2.0,
+        retry_backoff_factor=1.2,
+        operation_name="Milvus client connection after container start",
+        exception_types=(MilvusException, ))
+
   @staticmethod
   def start_db_container(
       image="milvusdb/milvus:v2.5.10",
@@ -148,23 +174,30 @@ def start_db_container(
       if tc_max_retries is not None:
         testcontainers_config.max_tries = tc_max_retries
       for i in range(vector_client_max_retries):
+        vector_db_container: Optional[CustomMilvusContainer] = None
         try:
           vector_db_container = CustomMilvusContainer(
               image=image,
               service_container_port=service_container_port,
               healthcheck_container_port=healthcheck_container_port)
-          vector_db_container = vector_db_container.with_volume_mapping(
+          mapped_container = vector_db_container.with_volume_mapping(
               cfg, "/milvus/configs/user.yaml")
+          assert mapped_container is not None
+          vector_db_container = mapped_container
           vector_db_container.start()
           host = vector_db_container.get_container_host_ip()
           port = vector_db_container.get_exposed_port(service_container_port)
           info = VectorDBContainerInfo(vector_db_container, host, port)
+          MilvusTestHelpers._wait_for_milvus_grpc(info.uri)
           _LOGGER.info(
               "milvus db container started successfully on %s.", info.uri)
+          break
         except Exception as e:
-          stdout_logs, stderr_logs = vector_db_container.get_logs()
-          stdout_logs = stdout_logs.decode("utf-8")
-          stderr_logs = stderr_logs.decode("utf-8")
+          stdout_logs = stderr_logs = ""
+          if vector_db_container is not None:
+            raw_out, raw_err = vector_db_container.get_logs()
+            stdout_logs = raw_out.decode("utf-8")
+            stderr_logs = raw_err.decode("utf-8")
           _LOGGER.warning(
               "Retry %d/%d: Failed to start Milvus DB container. Reason: %s. "
               "STDOUT logs:\n%s\nSTDERR logs:\n%s",