Add Vespa engine support for multi-engine benchmarking

osbenchmark/benchmark.py

@@ -598,6 +598,10 @@ def add_workload_source(subparser):
         help=f"Define a comma-separated list of client options to use. The options will be passed to the OpenSearch "
              f"Python client (default: {opts.ClientOptions.DEFAULT_CLIENT_OPTIONS}).",
         default=opts.ClientOptions.DEFAULT_CLIENT_OPTIONS)
+    test_run_parser.add_argument(
+        "--database-type",
+        help="Define the target database type (default: opensearch). Supported types: opensearch, vespa, milvus.",
+        default="opensearch")
     test_run_parser.add_argument("--on-error",
                              choices=["continue", "abort"],
                              help="Controls how OSB behaves on response errors (default: continue).",
@@ -1141,6 +1145,7 @@ def configure_test(arg_parser, args, cfg):
     cfg.add(config.Scope.applicationOverride, "workload", "randomization.alpha", args.randomization_alpha)
     cfg.add(config.Scope.applicationOverride, "workload", "visualize", args.visualize)
     cfg.add(config.Scope.applicationOverride, "workload", "visualize.output.path", args.visualize_output_path)
+    cfg.add(config.Scope.applicationOverride, "database", "type", args.database_type)
     configure_workload_params(arg_parser, args, cfg)
     configure_connection_params(arg_parser, args, cfg)
     configure_telemetry_params(args, cfg)

osbenchmark/client.py

@@ -30,10 +30,6 @@
 import urllib3
 from urllib3.util.ssl_ import is_ipaddress
 
-import grpc
-from opensearch.protobufs.services.document_service_pb2_grpc import DocumentServiceStub
-from opensearch.protobufs.services.search_service_pb2_grpc import SearchServiceStub
-
 from osbenchmark.kafka_client import KafkaMessageProducer
 from osbenchmark import exceptions, doc_link, async_connection
 from osbenchmark.context import RequestContextHolder
@@ -303,6 +299,14 @@ def create_grpc_stubs(self):
         Create gRPC service stubs.
         Returns a dict of {cluster_name: {service_name: stub}} structure.
         """
+        # Import grpc and protobuf stubs lazily — loading the grpc C extension
+        # initializes background threads. If that happens in the main process
+        # before Thespian forks actor children, the forked processes inherit
+        # broken gRPC thread state and may deadlock.
+        import grpc  # pylint: disable=import-outside-toplevel
+        from opensearch.protobufs.services.document_service_pb2_grpc import DocumentServiceStub  # pylint: disable=import-outside-toplevel
+        from opensearch.protobufs.services.search_service_pb2_grpc import SearchServiceStub  # pylint: disable=import-outside-toplevel
+
         stubs = {}
 
         if len(self.grpc_hosts.all_hosts.items()) > 1: