셀렉트스타의 Amazon API Gateway WebSocket 을 활용한 AI Red teaming API 스트림 처리 방법

Data-centric AI company, 셀렉트스타(Selectstar)

학습 데이터 품질이 인공지능 성능을 결정하고 좋은 인공지능은 실서비스 환경에서 수집되는 다양하고 방대한 데이터와 상호작용하며 끊임없이 발전합니다. 특히, 좋은 품질의 데이터를 수집하고 관리하는 작업은 여러 가지 측면에서 복잡하고 매우 중요합니다. 고품질의 데이터를 확보하기 위해서는 데이터의 정확성, 완전성, 관련성, 그리고 시기적절성을 모두 충족시켜야 하는데 이를 위해 데이터 수집 과정에서 발생할수 있는 오류를 최소화하고, 데이터의 다양성과 대표성을 확보하는 것이 필수적입니다.

이러한 과정들은 모두 시간과 자원을 상당히 많이 소모하는 작업이며, 전문 도메인 지식을 필요로 합니다. 따라서, 데이터의 품질을 지속적으로 유지하고 향상시키기 위한 전략적 접근과 체계적인 관리가 중요합니다. 이는 데이터 기반의 의사결정이 점점 더 중요해지는 현대 사회에서 기업과 연구 기관이 성공적으로 경쟁하기 위한 필수 요소가 되었습니다. 셀렉트스타(Selectstar)는 데이터 기획부터 구축 / 판매 운영 관리까지, AI 라이프사이클을 함께하는 올인원 데이터 컴퍼니(All-in-one Data company)입니다.

[영상 1] 데이터는 셀렉트스타 | Data-centric AI company, Selectstar
– 2024년 3월 기준 데이터 처리량

Gen AI Korea 2024 생성형 AI 레드팀 챌린지

전세계적으로 초거대언어모델 (Large Language Model, LLM)에 대한 관심이 급증하고 있습니다. ChatGPT, Claude 3, Llama 3 등 대규모 언어 모델(LLM)은 놀라울 정도로 현실적인 텍스트를 생성할 수 있는 능력으로 주목받고 있지만, 동시에 편향, 허위 정보, 증오 발언 등 부정적인 측면도 지니고 있습니다. 이러한 문제점들을 파악하고 해결하기 위해 특정 AI 시스템의 결함이나 취약점을 발견하기 위해 전문가들이 모의 공격이나 체계적인 테스트를 수행하는 활동을 AI 레드팀(AI Red teaming) 이라고 합니다. 이러한 테스트는 시스템의 보안을 강화하고, 잠재적인 위험을 사전에 파악하여 해결함으로써 AI의 안정성과 신뢰성을 높이는 데 중요한 역할을 합니다. 따라서 AI 레드팀은 AI의 신뢰성과 성능을 검증하는데 매우 중요합니다.

오픈AI, 구글, 네이버클라우드 등 여러 기업이 AI 레드팀을 운영하고 있으며, 셀렉트스타는 과학기술정보통신부와 한국정보통신기술협회(TTA)의 주관하는 ‘2024년 Gen AI Korea, 생성형 AI 레드팀 챌린지‘를 준비하였습니다. Gen AI Korea 2024 생성형 AI 레드팀 챌린지는 참가자들이 생성 인공지능(AI) 기술을 활용하여 현실적이고 독창적인 문제 해결 방안을 제시하고, 이를 통해 AI의 보안 문제를 파악하고 개선 방안을 모색하는 것을 목표로 합니다. 이 행사는 AI의 기술적 한계를 시험하고, 참가자들에게 AI 시스템의 취약점을 이해하고 대응하는 기회를 제공함으로써 AI 기술의 안정성을 강화하려는 의도에서 개최되었습니다.

셀렉트스타는 이번 챌린지에 주요 기술 파트너로 참여하여, 참가자들이 사용할 수 있는 AI 플랫폼을 제공하는 역할을 맡았습니다. 셀렉트스타는 네이버, SKT, Upstage, 42마루 등의 참여 기업에서 자체적으로 개발한 AI 모델을 가지고, 참가자들이 AI 레드팀 챌린지를 진행할 수있는 AI 레드팀 공급자(AI Red team Provider) 플랫폼을 개발하였습니다. 또한, 행사 기간 동안 발생할 수 있는 기술적 문제에 대응하기위해 기술 지원팀을 운영하여 원활한 행사 진행을 도왔습니다.

AI 모델 공급자 API의 안정적인 제공

Gen AI Korea 2024 생성형 AI 레드팀 챌린지는 2,000명 이상의 참가자와 심사자가 참여하는 대규모 행사였기 때문에 생성형 AI 모델공급자의 API들을 안전하게 제공하는 것이 중요한 과제였습니다. 기획부터 AI 레드팀 챌린지 대회까지 2개월이 안되는 짧은 시간안에 소규모의 개발팀이 프로젝트를 진행 해야만 하는 상황이었습니다. 그래서, 아래와 같은 기술적인 요구 사항을 최소한의 개발 리소스를 이용해서 빠르게 구현할 수 있는 기술 스택과 아키텍처를 사용했습니다.

대용량 트래픽을 안정적으로 처리 할 수 있는 AI 모델 공급사의 API 제공 방법 필요

수천명의 참가자가 동시에 접속하는 환경에서 원활한 데이터 흐름과 빠른 응답 속도를 보장된 인프라를 구축해야 했습니다. 특히 중요한 이슈는 생성형 AI 레드팀 챌린지에 참여한 AI 모델 공급사의 API 제공 방법이었습니다. 참가자들이 다양한 AI 모델 공급사들의 API들을 동일한 방식으로 사용할 수 있는 API를 제공해야 했습니다. 또한, 동시에 참가자들이 직접 AI 모델 공급사들의 API들을 호출할 수 없도록 보호해야 했습니다. 이러한 API에 대한 다양한 기능 요구 사항에 따라서 API 제공 방식에 대한 복잡성이 증가하게 되었습니다.

AI 모델 공급사의 요구 사항과 API 제공 방식의 복잡성을 줄이기 위해서 WebSocket을 이용해서 AI 모델 공급사의 API를 사용할 수 있도록 구현했습니다. WebSocket은 클라이언트와 서버 간의 양방향 통신을 위한 프로토콜로써 HTTP와 다르게 서버와 클라이언트 간에 지속적인 연결을 유지하면서 실시간으로 데이터를 교환할 수 있습니다. 이러한 WebSocket의 기능은 웹이나 모바일 애플리케이션에 적합하며, 채팅 애플리케이션, 실시간 알림 서비스, 또는 실시간 대시보드 같은 실시간 애플리케이션을 구현할 때 매우 유용합니다.

셀렉스타에서는 WebSocket을 이용해서 참가자들의 프롬프트를 받아서, AI 모델 공급사의 API를 대신 호출할 수 있는 API 서버를 구현하기로 했습니다. 그래서, AI 모델 제공사의 다양한 요구사항에 쉽게 대응하면서, 동시에 AI 모델 제공사의 LLM(대규모 언어 모델) API가 생성한 답변을 안정적으로 최종 사용자에게 제공하고자 했습니다.

초기 아키텍처 설계

초기에는 Spring Framework에 제공하는 WbeSocket Protocol인 STOMP과 Apache Kafka를 Message Bus로 이용한 아키텍처 패턴으로 구현하려고 했습니다. 하지만 STOMP 서버들을 Scale-out 확장하는 데 많은 어려움이 있었습니다. STOMP 서버들을 Scale-out 방식으로 확장 시키기 위해서는 각 서버들이 STOMP 세션을 공유할 수 있어야합니다. 그래서, 서버들 간의 STOMP 세션 공유를 위해서 Message Bus가 필요합니다. 셀렉스타에서는 Apache Kafka를 활용해서 Message Bus를 구축하는 방법을 고려했습니다.

하지만 Apache Kafka를 이용해서 Message Bus 시스템을 구축하는 데는 생각 보다 많은 개발 리소스가 필요했습니다. (1) Apache Kafka 클러스터 구축 및 운영, (2) Publish/Subscribe 방식으로 서버들 간에 STOMP 세션 공유할 수 있는 로직 구현, (3) 서비스 적용을 위하 부하 및 안정성 테스트 등을 소규모의 개발팀에서 짧은 기간 내에 진행해야 하는 상황이었습니다. 특히, 만약 주어진 개발 기간 내에 개발을 완료했다고 하더라도, Message Bus용 Apache Kafka 클러스터 관리 및 운영에 많은 리소스가 필요한 상황이었습니다.

[아키텍처 1] Spring STOMP와 Kafka를 Message Bus로 사용한 아키텍처

그래서, 단기간에 적은 개발 리소스를 사용해서 개발할 수 있는 대안을 찾던 중에 Amazon API Gateway의 WebSocket API를 알게 되었습니다.

Amazon API Gateway WebSocket 이란?

Amazon API Gateway는 클라이언트와 백엔드 서비스 사이의 요청과 응답을 관리하고 Amazon API Gateway를 통해 사용자는 REST API와 WebSocket API를 생성하고 관리할 수 있습니다. WebSocket API를 사용하면 실시간 메시지 교환을 비롯해 다양한 실시간 통신기능을 구현할 수 있습니다.

Amazon API Gateway WebSocket은 아래와 같은 주요 기능을 가지고 있습니다.

양방향 통신: Amazon API Gateway를 통해 WebSocket 연결을 설정하고 관리할 수 있으며, 이를 통해 서버와 클라이언트 간 양방향 실시간 통신을구현할 수 있습니다.
확장성: AWS 클라우드의 일부로서, Amazon API Gateway WebSocket은 높은 확장성을 제공하고 사용자는 수백만의 동시 연결을 관리할 수있으며, AWS의 글로벌 인프라를 활용하여 전 세계 어디서나 빠르고 안정적인 서비스를 제공할 수 있습니다.
통합 및 연결: AWS Lambda와 같은 서버리스 컴퓨팅 서비스와 직접 통합하여, 서버를 관리하지 않고도 복잡한 백엔드 로직을 실행할 수 있습니다. 또한, 다른 AWS 서비스와의 연동을 통해 데이터베이스, 인증, 로깅 등 다양한 서비스를 손쉽게 추가할 수 있습니다.
보안: Amazon API Gateway는 AWS IAM(Identity and Access Management)을 사용하여 WebSocket API를 보호합니다. 사용자는 권한을 세밀하게 제어하여 API 접근을 안전하게 관리할 수 있습니다.

Amazon API Gateway WebSocket 를 이용한 아키텍처 개선안

Amazon API Gateway WebSocket API를 사용함으로써, Apache Kafka와 같은 Message Bus를 사용하지 않고, 실시간으로 사용자의 프롬프트를 AI 모델 공급사들의 LLM API 서버에 보내고, LLM에서 생성된 답변을 주고 받을 수 있는 아래와 같은 WebSocket 서버를 구축할 수 있었습니다.

[아키텍처 2] Amazon API Gateway WebSocket을 사용한 아키텍처

Amazon API Gateway WebSocket을 사용해서 AI 모델 공급자 API의 QPM(Query Per Minute) 보장을 위한 티켓 시스템 구축

생성형 AI 레드팀 챌린지 대회 진행을 위해서 AI 모델 공급사들은 평균 100QPM 이내로 API가 호출되기를 요구했습니다. QPM (Query Per Minute)은 인공지능(AI)이나 다른 종류의 컴퓨터 시스템에서 분당 처리할 수 있는 쿼리(질의)의 수를 의미합니다. AI 모델 공급사의 API를 안정적으로 제공하기 위해서 Amazon API Gateway WebSocket을 이용해서 분당 100개 이내 쿼리를 호출하도록 제한하는 티켓시스템이 필요했습니다. 티켓시스템은 AI 모델별로 전체 티켓수가 있고, 사용자가 LLM 호출 티켓을 확보하고 LLM 답변이 완성되면 티켓을 다시 반납하는 방식으로 동작합니다.

Amazon API Gateway WebSocket API를 활용하면 메시지를 주고받을 때 connectionId를 사용합니다. 이 기능을 이용해서 LLM 서버 호출에 필요한 티켓 시스템을 구현할 수 있었습니다. 먼저, WebSocket에 연결(onConnect)될 때와 연결을 해제(onDisconnect)할 때마다 connectionId를 RDBMS에 저장(onConnect) 또는 삭제(onDisconnect)합니다. 이렇게 connectionId를 관리함으로써 각각의 WebSocket 연결을 고유하게 식별할 수 있습니다. 그리고, RDBMS 테이블에 저장된 connectionId의 상태를 변경하는 방식으로 티켓 시스템을 구현할 수 있습니다. 예를 들어, 특정 connectionId가 활성화되어 있는지 여부를 확인하거나, 티켓이 사용된 후에는 상태를 업데이트하는 방식으로 티켓을 관리합니다. 또한, RDBMS의 테이블 락(Lock)을 이용해서 여러 사용자들이 동시에 티켓을 요청하는 동시성 문제를 해결하고, 티켓의 상태 데이터에 대한 무결성을 유지합니다.

티켓시스템의 동작 설명

LLM API 호출 준비 – Amazon API Gateway WebSocket 연결 시도 (onConnect) (아키텍처 2의 Step 1, 2)
1. access_token 을 활용하여 유효한 사용자인지 검사.
2. 현재 사용가능한 티켓이 존재하는지 확인
3. 티켓이 존재하지않을 경우, 대기
4. 티켓이 존재하면 티켓을 확보.
LLM API 호출 – Amazon API Gateway WebSocket 메시지 전송 (sendMessage) (아키텍처 2의 Step 3, 4)
1. 확보된 티켓을 활용하여 현재 알고리즘에 맞는 LLM API 호출
2. 전달 받는 응답 토큰 데이터를 스트림(stream) 방식으로 사용자에게 전송
LLM API 호출 종료 – Amazon API Gateway WebSocket 연결 종료 (onDisconnect) (아키텍처 2의 Step 5)
1. LLM 호출 종료
2. 사용중인 티켓을 반환

AI 모델 공급사의 API 오류 발생 시, 대응 전략

AI 모델 공급사별로 QPM을 보장하기 것 외에도 LLM API 자체에서 오류가 발생해서 참가자들의 요청을 처리하기 못하는 상황에 대한 대응 전략이 필요했습니다.

AI 모델 공급사의 API 서버에서 다양한 원인에 의해서 AI 모델 공급사별로 API 응답 에러(Response Error Status)가 발생 할 수 있습니다. 이렇게 AI 모델 추론 API 에러가 발생하면, 에러 원인에 따라서 미리 정의된 규칙에 따라서 자동으로 적절한 조치를 취한 후에 API를 다시 호출합니다. 일정한 횟수 만큼 API 호출을 재시도 한 후에도 여전히 API 호출 에러가 발생하는 경우에는 최종적으로 사용자에게 에러 메시지를 전송합니다. 예를 들어, max_input token 초과 에러가 발생하면, 사용자가 입력한 프롬프트의 토큰 수를 줄여가면서, API를 다시 호출하는 방식으로 API 에러를 처리합니다. AI 모델 공급사의 API 호출 제약 조건과 에러 발생 원인이 다르기 때문에, AI 모델 공급사의 API별로 비즈니스 로직을 구현해서 처리했습니다.

이와 같이 AI 모델 공급사의 API 호출 에러에 대해서 유연하게 대응함으로써, AI 레드팀 서비스의 안정성을 높일 수 있었습니다. 또한, 다양한 AI 모델 공급사의 API 에러 메시지들이 비슷하지만, 서로 다르기 때문에 이러한 메시지를 최종 사용자들이 이해하기 쉽도록 정리해서 안내함으로써, 사용자들이 마치 하나의 AI 모델 공급사의 API를 사용하는 것처럼 여겨지도록 했습니다.

1명의 참여자가 제출한 대화 셋을 여러 심사자들에게 Broadcast 하기

생성형 AI 레드팀 챌린지 대회에서는 참가자 1명이 제출한 LLM과의 대화셋을 여러 심사자들이 평가합니다. 이틀 동안 진행되는 생성형 AI 레드팀 챌린지에서는 참여자가 4시간동안 생성한 LLM과의 대화셋을 심사자들이 12시간이내에 모든 평가를 진행합니다. 그래서, 참가자가 대화셋을 제출하면, 대화셋 타입에 따라 심사자들이 배정되고, 심사가 이루어집니다.

참가자가 제출한 대화셋을 여러 심사자들에게 동시에 전송하기 위해 Amazon API Gateway WebSocket을 활용했습니다. 1명의 참가자와 여러 심사자들이 서로 메시지를 주고 받을 수 있도록 WebSocket에 접속하면, connectionId가 생성됩니다. 생성된 connectionId는 참가자와 심사자들 사이의 세션이 종료될 때까지 RDBMS에 저장됩니다.

AI 레드팀 작업을 완료한 참가자가 대화셋을 제출하면, 제출 타입에 따라서 대화셋을 평가할 심사자들이 배정됩니다. 그리고, RDBMS에 저장된 connectionId를 이용해서 평가를 진행할 심사자들에게 참가자의 대화셋을 전달합니다. 마찬가지 방법으로 심사자들 역시 평가를 완료하면, 심사 결과를 참가자들에게 실시간으로 브로캐스팅(Broadcast) 합니다. Amazon API Gateway WebSocket을 이용해서 짧은 시간 안에 심사 완료 및 최종 결과 발표를 할 수 있었습니다.

스트레스 테스트

AI 레드팀 공급자(AI Red team Provider) 플랫폼 개발을 완료 후, 생성형 AI 레드팀 챌린지 대회 운영을 위한 서버 준비 및 성능 테스트를 대회 시나리오와 최대한 유사하게 진행하였습니다.

테스트 시나리오

사용자 3,000명을 4그룹으로 나눈다.
4그룹은 각각 하나의 LLM 모델을 호출하도록 한다.
각 그룹의 사용자들은 [ LLM 호출 → 응답 → 10초 대기 ] 작업을 지속적으로 반복한다.
1, 2, 3번 과정을 N 분 동안 지속적으로 호출하고 성공여부를 확인한다.

테스트 결과 요약

Amazon API Gateway 오류 0건
Amazon API Gateway 데이터 전송 중 오류 197건 – LLM API 서버에 QPM 이상으로 API를 호출 했을 경우 발생함

테스트 결과 상세 지표

[그림 1] 스트레스 테스트 Message Count 결과

[그림 2] 스트레스 테스트 Connection Count 결과

[그림 3] 스트레스 테스트 Execution Error 결과
Execution Error 의 발생은 LLM API 호출 시 QPM 이상의 값을 호출 했을 경우, 일부 서버 에러 발생

생성형 AI 레드팀 챌린지 대회 기간 동안 모니터링 결과 (2024.04.11 ~ 12)

스트레스 테스트 결과를 바탕으로 서버 준비를 완료 한 후에 Gen AI Korea 2024 생성형 AI 레드팀 챌린지 대회 당일에는 아래와 같이 아무 문제 없이 행사를 진행 할 수 있었습니다.

[그림 4] Gen AI Korea 2024 생성형 AI 레드팀 챌린지 기간 동안 Message Count 결과

[그림 5] Gen AI Korea 2024 생성형 AI 레드팀 챌린지 기간 동안 Connection Count 결과

마무리

“Gen AI Korea 2024 생성형 AI 레드팀 챌린지”의 기획부터 개발, 대회 운영까지 2개월이라는 짧은 기간 동안 소수의 개발 인원으로 이 프로젝트를 성공적으로 마칠 수 있었던 큰 요인 중 하나는 AWS 관리형 또는 서버리스 서비스들을 최대한 활용하는 것이었습니다. AWS 서비스들을 이용해서 운영 부담을 줄이고, 비즈니스 로직 개발에 집중할 수 있었으며, 대회의 테스트 및 운영 모니터링을 안정적으로 수행할 수있었습니다.

AI 모델 공급사들이 요구하는 다양한 API 관련 요구사항을 Amazon API Gateway를 이용해서 효율적으로 구현할 수 있었습니다. Amazon API Gateway 덕분에 개발의 유연성을 높이는 동시에, 운영의 효율성을 극대화할 수 있었습니다. 또한 QPM (Query Per Minute) 제어 로직, 사용자 그룹화, 최대 메시지 제한 등의 기능들 역시 AWS 서비스를 이용해서 안전하고, 관리하기 쉽도록 구현할 있었습니다.

Amazon API Gateway의 WebSocket API 사용법에 대해서 궁금하신 분들은 아래 참고 자료를 확인해 보세요.

AWS 기술 블로그