Improve test performance

tests/_utils/kafka.py

@@ -42,11 +42,15 @@
 
 @pytest.fixture(name='kafka_topics', scope='function')
 def kafka_topics_fixture():
-    yield KafkaTopics(f'morpheus_unittest_input_{time.time()}', f'morpheus_unittest_output_{time.time()}')
+    """
+    Every test receives a unique pair of Kafka topics
+    """
+    ts = time.time()
+    yield KafkaTopics(f'morpheus_unittest_input_{ts}', f'morpheus_unittest_output_{ts}')
 
 
 @pytest.fixture(name='kafka_bootstrap_servers', scope="function")
-def kafka_bootstrap_servers_fixture(kafka_server: (subprocess.Popen, int)):  # pylint: disable=redefined-outer-name
+def kafka_bootstrap_servers_fixture(kafka_server: tuple[subprocess.Popen, int]):  # pylint: disable=redefined-outer-name
     """
     Used by tests that require both an input and an output topic
     """

tests/conftest.py

@@ -502,14 +502,12 @@ def disable_gc():
     gc.enable()
 
 
-def wait_for_camouflage(host="localhost", port=8000, timeout=30):
+def wait_for_server(url: str, timeout: int, parse_fn: typing.Callable[[requests.Response], bool]) -> bool:
 
     start_time = time.time()
     cur_time = start_time
     end_time = start_time + timeout
 
-    url = f"http://{host}:{port}/ping"
-
     while cur_time - start_time <= timeout:
         timeout_epoch = min(cur_time + 2.0, end_time)
 
@@ -518,13 +516,9 @@ def wait_for_camouflage(host="localhost", port=8000, timeout=30):
             resp = requests.get(url, timeout=request_timeout)
 
             if (resp.status_code == 200):
-                if (resp.json()['message'] == 'I am alive.'):
+                if parse_fn(resp):
                     return True
 
-                warnings.warn(("Camoflage returned status 200 but had incorrect response JSON. Continuing to wait. "
-                               "Response JSON:\n%s"),
-                              resp.json())
-
         except Exception:
             pass
 
@@ -537,6 +531,31 @@ def wait_for_camouflage(host="localhost", port=8000, timeout=30):
     return False
 
 
+def wait_for_camouflage(host: str = "localhost", port: int = 8000, timeout: int = 30):
+    url = f"http://{host}:{port}/ping"
+
+    def parse_fn(resp: requests.Response) -> bool:
+        if (resp.json()['message'] == 'I am alive.'):
+            return True
+
+        warnings.warn(("Camoflage returned status 200 but had incorrect response JSON. Continuing to wait. "
+                       "Response JSON:\n%s"),
+                      resp.json())
+        return False
+
+    return wait_for_server(url, timeout=timeout, parse_fn=parse_fn)
+
+
+def wait_for_milvus(host: str = "localhost", port: int = 19530, timeout: int = 180):
+    url = f'http://{host}:{port}/healthz'
+
+    def parse_fn(resp: requests.Response) -> bool:
+        content = resp.content.decode('utf-8')
+        return 'OK' in content
+
+    return wait_for_server(url, timeout=timeout, parse_fn=parse_fn)
+
+
 def _set_pdeathsig(sig=signal.SIGTERM):
     """
     Helper function to ensure once parent process exits, its child processes will automatically die
@@ -954,18 +973,21 @@ def milvus_server_uri(tmp_path_factory):
         yield uri
 
     else:
-        from milvus import default_server
+        from milvus import MilvusServer
 
-        # Milvus checks for already bound ports but it doesnt seem to work for webservice_port. Use a random one
-        default_server.webservice_port = _get_random_port()
-        with default_server:
-            default_server.set_base_dir(tmp_path_factory.mktemp("milvus_store"))
+        milvus_server = MilvusServer(wait_for_started=False)
 
-            host = default_server.server_address
-            port = default_server.listen_port
+        # Milvus checks for already bound ports but it doesnt seem to work for webservice_port. Use a random one
+        webservice_port = _get_random_port()
+        milvus_server.webservice_port = webservice_port
+        milvus_server.set_base_dir(tmp_path_factory.mktemp("milvus_store"))
+        with milvus_server:
+            host = milvus_server.server_address
+            port = milvus_server.listen_port
             uri = f"http://{host}:{port}"
 
             logger.info("Started Milvus at: %s", uri)
+            wait_for_milvus(host=host, port=webservice_port, timeout=180)
 
             yield uri
 
@@ -976,6 +998,13 @@ def milvus_data_fixture():
     yield inital_data
 
 
+@pytest.fixture(scope="session", name="milvus_service")
+def milvus_service_fixture(milvus_server_uri: str):
+    from morpheus_llm.service.vdb.milvus_vector_db_service import MilvusVectorDBService
+    service = MilvusVectorDBService(uri=milvus_server_uri)
+    yield service
+
+
 @pytest.fixture(scope="session", name="idx_part_collection_config")
 def idx_part_collection_config_fixture():
     from _utils import load_json_file

tests/morpheus/common/test_http_server.py

@@ -38,15 +38,14 @@ def make_parse_fn(status: HTTPStatus = HTTPStatus.OK,
 
 @pytest.mark.slow
 @pytest.mark.parametrize("endpoints", [("/t1", "/t2", "/t3"), ("test/", "123/", "a1d/"), ("/a", "/a/b", "/a/b/c/d")])
-@pytest.mark.parametrize("port", [8088, 9090])
-@pytest.mark.parametrize("method", ["GET", "POST", "PUT"])
+@pytest.mark.parametrize("port", [9090])
+@pytest.mark.parametrize("method", ["GET", "POST"])
 @pytest.mark.parametrize("use_callback", [True, False])
 @pytest.mark.parametrize("use_context_mgr", [True, False])
-@pytest.mark.parametrize("num_threads", [1, 2, min(8, len(os.sched_getaffinity(0)))])
+@pytest.mark.parametrize("num_threads", [1, min(8, len(os.sched_getaffinity(0)))])
 @pytest.mark.parametrize("status,content_type,content",
                          [(HTTPStatus.OK, MimeTypes.TEXT.value, "OK"),
                           (HTTPStatus.OK, MimeTypes.JSON.value, '{"test": "OK"}'),
-                          (HTTPStatus.NOT_FOUND, MimeTypes.TEXT.value, "NOT FOUND"),
                           (HTTPStatus.INTERNAL_SERVER_ERROR, MimeTypes.TEXT.value, "Unexpected error")])
 def test_simple_request(port: int,
                         endpoints: typing.Tuple[str, str, str],

tests/morpheus/stages/test_kafka_source_stage_pipe.py

@@ -30,6 +30,7 @@
 from morpheus.config import Config
 from morpheus.pipeline.linear_pipeline import LinearPipeline
 from morpheus.stages.general.trigger_stage import TriggerStage
+from morpheus.stages.input.kafka_source_stage import AutoOffsetReset
 from morpheus.stages.input.kafka_source_stage import KafkaSourceStage
 from morpheus.stages.output.compare_dataframe_stage import CompareDataFrameStage
 from morpheus.stages.postprocess.serialize_stage import SerializeStage
@@ -52,7 +53,7 @@ def test_kafka_source_stage_pipe(config: Config, kafka_bootstrap_servers: str, k
         KafkaSourceStage(config,
                          bootstrap_servers=kafka_bootstrap_servers,
                          input_topic=kafka_topics.input_topic,
-                         auto_offset_reset="earliest",
+                         auto_offset_reset=AutoOffsetReset.EARLIEST,
                          poll_interval="1seconds",
                          client_id='morpheus_kafka_source_stage_pipe',
                          stop_after=num_records))
@@ -85,7 +86,7 @@ def test_multi_topic_kafka_source_stage_pipe(config: Config, kafka_bootstrap_ser
         KafkaSourceStage(config,
                          bootstrap_servers=kafka_bootstrap_servers,
                          input_topic=input_topics,
-                         auto_offset_reset="earliest",
+                         auto_offset_reset=AutoOffsetReset.EARLIEST,
                          poll_interval="1seconds",
                          client_id='test_multi_topic_kafka_source_stage_pipe',
                          stop_after=num_records))
@@ -100,38 +101,41 @@ def test_multi_topic_kafka_source_stage_pipe(config: Config, kafka_bootstrap_ser
 @pytest.mark.gpu_and_cpu_mode
 @pytest.mark.kafka
 @pytest.mark.parametrize('async_commits', [True, False])
-@pytest.mark.parametrize('num_records', [10, 100, 1000])
-def test_kafka_source_commit(num_records: int,
-                             async_commits: bool,
-                             config: Config,
-                             kafka_bootstrap_servers: str,
-                             kafka_topics: KafkaTopics,
-                             kafka_consumer: "KafkaConsumer") -> None:
+@pytest.mark.parametrize('num_records', [10, 1000])
+def test_kafka_source_batch_pipe(num_records: int,
+                                 async_commits: bool,
+                                 config: Config,
+                                 kafka_bootstrap_servers: str,
+                                 kafka_topics: KafkaTopics,
+                                 kafka_consumer: "KafkaConsumer") -> None:
     group_id = 'morpheus'
 
     data = [{'v': i} for i in range(num_records)]
     num_written = write_data_to_kafka(kafka_bootstrap_servers, kafka_topics.input_topic, data)
     assert num_written == num_records
 
+    expected_length = config.pipeline_batch_size
+    num_exact = num_records // expected_length
+
     kafka_consumer.subscribe([kafka_topics.input_topic])
     seek_to_beginning(kafka_consumer)
     partitions = kafka_consumer.assignment()
 
-    # This method does not advance the consumer, and even if it did, this consumer has a different group_id than the
-    # source stage
+    # This method does not advance the consumer, and this consumer has a different group_id than the source stage
     expected_offsets = kafka_consumer.end_offsets(partitions)
 
     pipe = LinearPipeline(config)
     pipe.set_source(
         KafkaSourceStage(config,
                          bootstrap_servers=kafka_bootstrap_servers,
                          input_topic=kafka_topics.input_topic,
-                         auto_offset_reset="earliest",
+                         auto_offset_reset=AutoOffsetReset.EARLIEST,
                          poll_interval="1seconds",
                          group_id=group_id,
                          client_id='morpheus_kafka_source_commit',
                          stop_after=num_records,
                          async_commits=async_commits))
+    pipe.add_stage(DFPLengthChecker(config, expected_length=expected_length, num_exact=num_exact))
     pipe.add_stage(TriggerStage(config))
     pipe.add_stage(DeserializeStage(config))
     pipe.add_stage(SerializeStage(config))
@@ -148,39 +152,6 @@ def test_kafka_source_commit(num_records: int,
     # The broker may have created additional partitions, offsets should be a superset of expected_offsets
     for (topic_partition, expected_offset) in expected_offsets.items():
         # The value of the offsets dict being returned is a tuple of (offset, metadata), while the value of the
-        #  expected_offsets is just the offset.
+        # expected_offsets is just the offset.
         actual_offset = offsets[topic_partition][0]
         assert actual_offset == expected_offset
-
-
-@pytest.mark.gpu_and_cpu_mode
-@pytest.mark.kafka
-@pytest.mark.parametrize('num_records', [1000])
-def test_kafka_source_batch_pipe(config: Config,
-                                 kafka_bootstrap_servers: str,
-                                 kafka_topics: KafkaTopics,
-                                 num_records: int) -> None:
-    data = [{'v': i} for i in range(num_records)]
-    num_written = write_data_to_kafka(kafka_bootstrap_servers, kafka_topics.input_topic, data)
-    assert num_written == num_records
-
-    expected_length = config.pipeline_batch_size
-    num_exact = num_records // expected_length
-
-    pipe = LinearPipeline(config)
-    pipe.set_source(
-        KafkaSourceStage(config,
-                         bootstrap_servers=kafka_bootstrap_servers,
-                         input_topic=kafka_topics.input_topic,
-                         auto_offset_reset="earliest",
-                         poll_interval="1seconds",
-                         client_id='morpheus_kafka_source_stage_pipe',
-                         stop_after=num_records))
-    pipe.add_stage(DFPLengthChecker(config, expected_length=expected_length, num_exact=num_exact))
-    pipe.add_stage(DeserializeStage(config))
-    pipe.add_stage(SerializeStage(config))
-    comp_stage = pipe.add_stage(
-        CompareDataFrameStage(config, pd.DataFrame(data=data), include=[r'^v$'], reset_index=True))
-    pipe.run()
-
-    assert_results(comp_stage.get_results())

tests/morpheus_llm/llm/nodes/test_llm_retriever_node_pipe.py

@@ -34,12 +34,6 @@
 from morpheus_llm.stages.llm.llm_engine_stage import LLMEngineStage
 
 
-@pytest.fixture(scope="module", name="milvus_service")
-def milvus_service_fixture(milvus_server_uri: str):
-    service = MilvusVectorDBService(uri=milvus_server_uri)
-    yield service
-
-
 def _build_engine(vdb_service, **similarity_search_kwargs) -> LLMEngine:
     mock_embedding = mock.AsyncMock(return_value=[[1.2, 2.3, 3.4], [4.5, 5.6, 6.7]])
     engine = LLMEngine()