GRPC deadlock on startup about locust

[zeionara]

Using GprcUser instead of HttpUser as written here helped to avoid deadlock. The resulting test/load/demo.py:

from time import sleep

from locust import task
from langchain.vectorstores.milvus import Milvus
from pymilvus.grpc_gen.milvus_pb2_grpc import MilvusServiceStub

from villm.index.ChromaCollector import HuggingFaceEmbedder

from test.load.example import HOST, PORT
from test.load.quick import EMBEDDING_MODEL, COLLECTION
from test.load.GrpcUser import GrpcUser


class DemoTestUser(GrpcUser):
    host = HOST
    port = PORT

    stub_class = MilvusServiceStub

    def on_start(self):
        print('creating embedder...')
        embedder = HuggingFaceEmbedder(model_name = EMBEDDING_MODEL)

        print('initializing milvus client...')
        self.client = Milvus(
            embedder.get_embedder(),
            connection_args = {'host': HOST, 'port': PORT},
            collection_name = COLLECTION
        )

        print('ready to make queries')

    @task
    def query(self):
        print('starting task...')

        sleep(2)

        print('ending task...')

The contents of test/load/GrpcUser.py:

import time
from typing import Any, Callable
import grpc
import grpc.experimental.gevent as grpc_gevent
from grpc_interceptor import ClientInterceptor
from locust import User
from locust.exception import LocustError

# patch grpc so that it uses gevent instead of asyncio
grpc_gevent.init_gevent()


class LocustInterceptor(ClientInterceptor):
    def __init__(self, environment, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.env = environment

    def intercept(
        self,
        method: Callable,
        request_or_iterator: Any,
        call_details: grpc.ClientCallDetails,
    ):
        response = None
        exception = None
        start_perf_counter = time.perf_counter()
        response_length = 0
        try:
            response = method(request_or_iterator, call_details)
            response_length = response.result().ByteSize()
        except grpc.RpcError as e:
            exception = e

        self.env.events.request.fire(
            request_type="grpc",
            name=call_details.method,
            response_time=(time.perf_counter() - start_perf_counter) * 1000,
            response_length=response_length,
            response=response,
            context=None,
            exception=exception,
        )
        return response


class GrpcUser(User):
    abstract = True
    stub_class = None

    def __init__(self, environment):
        super().__init__(environment)
        for attr_value, attr_name in ((self.host, "host"), (self.stub_class, "stub_class")):
            if attr_value is None:
                raise LocustError(f"You must specify the {attr_name}.")

        self._channel = grpc.insecure_channel(self.host)
        interceptor = LocustInterceptor(environment=environment)
        self._channel = grpc.intercept_channel(self._channel, interceptor)

        self.stub = self.stub_class(self._channel)