[Locust] 대시보드 커스텀 (feat. websocket 부하테스트)

AI 솔루션의 한 종류인 speech to text (stt) 의 경우 websocket을 통해 응답을 주고 받는 구조로 구성되어 있었다.

대시보드가 구성되지 않는 현상이 발생하였고, 직접 구성하는 방법에 대해 공유해 보려한다.

 

문제 현상

 

• 파일 사이즈가 1mb이상 넘어가면 수신 측(stt) payload가 크다는 오류 발생

• chunk로 분류하면 utf-8 디코딩 불가하다는 오류 발생

• websocket 수신을 정의하지 않아 오류 발생 (영향은 없음)

• 파일 자체를 보내 (file, stream, chunk 등) 이슈 발생 -> text로 decode하여 전송으로 해결

• 오류는 잡았으나, dashboard가 나오지 않는 현상 발생

• event에 데이터를 직접 삽입하여 dashboard 구성 
https://docs.locust.io/en/stable/api.html#locust.event.Events.request

 

from locust import User, task, between, events
import websocket
import base64
import time

class DigitalTextbookUser(User):
   wait_time = between(10, 20)

   audioFilePath = "audio.wav"

   successful_requests = 0
   failed_requests = 0
   total_response_time = 0

   @task
   def stt(self):
        start_time = time.time()
        async with websockets.connect("wss://domain.com/api/stt") as ws:
        
        # 오디오 파일 읽기 밎 전송
        file = open(audioFilePath, "rb")
        while True:
        	data = file.read(int(NUMBYTESINSEC * LENBLOCKINSEC))
            if data == '' or len(data) <=0:
            	break
            audioContents = base64.b64encode(file.read()).decode("utf8")
            await ws.send(audioContents);
        await ws.send(b'%f0000');    
        
		
        # 서버 응답 받기
        while True:
            response = await ws.recv()
            json_data = data.decode('utf-8')
            json_data = json.loads(json_data)
            if json_data['message'] == 'server send final flag':
            	response_time = (time.time() - start_time) * 1000
               	self.successful_requests += 1
               	self.total_response_time += response_time

                # 성공적인 요청 Locust 이벤트 시스템에 전달
                events.request.fire(
                   request_type="WebSocket",
                   name="STT Request",
                   response_time=response_time,
                   response_length=len(response)
                )
            	break
       

       except Exception as e:
           # 실패 시
           self.failed_requests += 1
           events.request.fire(
               request_type="WebSocket",
               name="STT Request",
               response_time=(time.time() - start_time) * 1000,
               exception=str(e)
           )