Pub/Sub 구조로 Inference Stop 기능 구현 및 최적화하기 [Redis, Server Local Cache]

거의 모든 AI Chat 서비스는 아래와 같이 응답이 나오는 중에 중지가 가능하다.

우리는 원래 Inference Stop 기능을 제공하지 않았으나, 잘못 보냈을 때에도 응답을 다 기다려야 한다는 것이 불편하다는 사용자 피드백을 받고 Inference Stop 기능을 구현하기로 하였다.

아래와 같이 응답이 나오는 중에 Send 버튼이 Stop 으로 바뀌어서 누르면 채팅이 나오다가 멈춰야 한다.

하지만 구현은 생각보다 쉽지 않다는 것을 조금만 생각해보면 알 것이다.

1. [실패] 프론트엔드 단에서 응답을 강제로 막아버리면?

채팅 API를 끊어도 이미 실행되었으면 채팅 내용은 데이터베이스에 결국 저장된다.

그래서 프론트엔드 단에서 응답을 처리해버리면 멈췄을 당시에는 제대로 처리 되겠지만, 페이지를 리로딩하거나 나중에 그 채팅방에 들어갔을 때 DB에는 결국 응답이 저장돼 있기 때문에 Stop이 안 되었다고 느낄 수 있다.

따라서 아래와 같이 리로딩 했을 때에도 Stop이 정상적으로 적용되어야 하므로 이 방법은 적용할 수 없다.

2. [실패] Inference Stop 요청을 받아, 그 요청을 Inference 서버에 저장해두고 현재 응답 중인 Inference에 Stop 요청이 들어왔으면 Streaming 중단하기

결국 프론트엔드 단에서는 해결이 힘들다는 것을 파악했으니 백엔드 단에서 해결해야 한다.
스트리밍 중간에 Stop 요청이 오면 멈춰야 하므로 스트리밍 중에 주기적으로, 또는 토큰이 나올 때마다 확인해서 Streaming 중단 여부를 확인해야 한다.

위 해결책은 직관적이면서도 작동할 것처럼 보이는데, 한 가지를 추가로 고려해야 한다.
우리는 서버를 쿠버네티스 환경(EKS)에서 운영하고 있고, Inference Server Pod은 기본적으로 3개가 돌아가고 있다.

그래서 요청이 Pod에 도달하기 전 Load Balancer에 의해 요청이 분산되는데, 이때 Stop 요청이 반드시 Streaming 중인 Inference Pod에 간다는 보장이 없다.

물론 항상 Inference Pod이 1개라는 것이 보장되면 작동하는 방식이다.

[참고]

현재 우리는 Round Robin 방식으로 트래픽을 라우팅 하고 있는데, Load Balancer의 Routing 정책을 바꿔서 해결할 수도 있다.
네트워크 단에서 해결하는 것은 조금 불안했고, Deterministic Routing을 위한 설정이 오히려 더 복잡할 것이라고 판단했고, 라우팅 중 예상 못하는 문제가 발생할 수도 있을 것이라고 생각였다.

3. [70% 쯤 성공] Stop 요청을 받아 Redis에 저장하고, Inference 처리 중에 Stop 데이터가 들어있는지 확인 후 처리하기

2번 방법의 문제는 Stop Request가 어디로 갈지 모른다는 것이었다.
따라서 Redis 같은 가벼운 In-Memory DB를 사용해서 Stop 요청을 저장해두면 모든 Pod이 같은 데이터를 참조할 수 있어서 2번 문제를 해결할 수 있다.

우리는 Redis를 이미 사용하고 있었기 때문에 러닝 커브나 개발 공수 측면에서도 효율적이라고 생각하였다.

구현은 매우 간단했다.

POST /api/v1/inference/stop에서 요청을 받으면 Redis에 Key를inference:stop:{inference_id} 로 설정하여 저장한다.
Inference Streaming Processing 중, 토큰이 날라올 때마다 Redis에 쿼리를 해서 Stop 요청이 왔는지 확인한다.
Stop 데이터가 존재하면 바로 Streaming을 멈춘다.

일단 빠른 배포가 필요한 상황이었기에 이렇게 임시로 해결했다.
하지만 위 문제는 큰 문제가 하나 있다.

Streaming Response Chunk가 넘어올 때마다 Redis에 쿼리를 날린다는 것이었다.
Langfuse에서 분석한 결과, 한 Inference Request 당 평균 Chunk의 개수와 Response Time는 약 80개 / 10초이다.

그러면 약 1초 당 8개의 Chunk가 날라오는 것이고, 초당 100회의 Inference 요청이 있다면 Redis 쿼리는 초당 800회가 일어나는 것이다.

그리고 Stop을 하는 경우가 전체 Request의 3%도 안 되는데, 3% 때문에 Chunk 마다 쿼리를 날리는 것은 매우 비효율적이라고 생각하였다.

[70% 쯤 성공] 이라고 한 이유도 동작에서는 성공했으나, 성능에서는 실패했기 때문이다.
그래서 당장 최적화 방안을 강구하기 시작했다.

4. [최적화 성공] Stop 요청을 받아, Stop Event를 발행하고, 모든 Inference 서버가 이 이벤트를 받아 처리하기

우리는 Stop 이벤트를 발행하고, 모든 Pod들이 이 이벤트를 즉시 받아서 처리하는 방식으로 구현하기로 하였다.

모든 Pod들이 Stop 이벤트를 구독하고 있기 때문에, Inference를 실제로 처리 중인 Pod이 확인 후 처리만 하면 된다.
처리 중이지 않은 Pod은 Stop 이벤트를 받아도 본인의 Stop 요청이 아니기 때문에 자연스럽게 무시된다.

Pub/Sub을 제공하는 솔루션 중 가장 유명한 것은 Kakfa다.

하지만 우리는 Redis를 사전에 사용하고 있었고, Redis에서 Pub/Sub을 제공한다.
그리고 우리 규모와 상황에서 Kakfa를 새로 도입해서 개발하는 것보다 기존에 사용 중인 Redis를 활용하는 것이 개발 공수와 비용 측면에서 더 효율적이라고 판단했다.

이렇게 하면 이벤트가 매번 있는지 쿼리를 날리는 것이 아니라, 요청이 날라오면 받아서 처리하는 구조이기 때문에 더 효율적이었다.

하지만 그렇게 해도 Subscribe 하는 Channel과 Inference를 처리하는 Flow를 연결하는 중간 클래스가 필요했기 떄문에 우리는 InferenceStopSignalHolder라는 클래스를 개발했다.

Redis Subscriber는 Stop 이벤트를 구독하고 있다가 이벤트를 감지하면 InferenceStopSignalHolder로 보낸다.
Inference 처리 중 Chunk가 올 때마다 InferenceStopSignalHolder를 확인하여 Stop Event가 있는지 확인한다.
Stop Event가 있으면 즉시 Streaming을 멈추고 Stop Response로 반환한다.

[참고]

InferenceStopSignalHolder는 내부적으로 threading.Lock과 Dictionary로 구현되어 있다.
혹시 모를 동시에 Stop 이벤트가 발생되었을 때 이벤트가 누락되는 현상을 방지하려고 하였다.
Lookup을 상수 시간으로(O(1)) Stop Event Checking을 처리하려고 하였다.
삭제 이벤트는 따로 발행하지 않고, Redis에 저장되는 데이터에 TTL을 설정하여 자동으로 사라지도록 하였다.
InferenceStopSignalHolder에서도 저장될 때 Event 발행 시간과 함께 관리하여 TTL과 함께 서버에서도 사라지도록 하였다. (OOM 방지)
식별자는 Inference 요청마다 생겨나는 inference_id를 사용하여 실행 중인 Inference Request만 정확히 Stop 할 수 있도록 하였다.
user_id나 chatroom_id를 식별자로 사용하면, TTL로 인해 사라지기 전에 채팅 요청을 날리면 같은 식별자를 가지고 있기 때문에 또 Stop이 될 수 있다.

InferenceStopSignalHolder는 서버 프로세스 내부에서 상태를 관리하는 클래스다. 따라서 Stop 요청이 있는지를 확인하는 연산은 외부 스토리지를 조회하지 않고, 단순한 메모리 접근으로 처리된다.
이로 인해 Stop 여부를 확인하는 데 걸리는 시간은 나노초(ns) 단위 수준으로 사실상 무시해도 될 정도다.

아래는 Redis Query Count를 실제 운영 환경에서 확인해본 결과이다. (Peak Time 기준)

최적화 전 Stop Exists 쿼리: 분당 약 1,300회
최적화 후 Stop 이벤트 발행: 분당 평균 0.4회

위에서 말했던대로 Stop 요청은 거의 없는데 이걸 위해서 상당히 많은 조회를 하고 있었다는 것이 사실이었으며, 이벤트 기반 구조로 바꿔서 Stop Exists 쿼리를 전부 없애고 요청이 있을 때만 가져와서 효율적으로 처리하는데 성공하였다.

한계 및 고려사항

1. 서버 프로세스가 중단되는 경우

Redis는 이벤트를 발행한 이후, 해당 이벤트가 실제로 소비되었는지까지는 보장하지 않는다. 따라서 Stop 이벤트가 발행된 직후 서버 프로세스가 종료된다면, Stop 신호가 정상적으로 처리되지 않을 가능성이 존재한다.

다만 이 한계는 실제 서비스 환경에서는 큰 문제가 되지 않는다고 판단했다.
Stop 이벤트가 의미를 가지는 대상은 해당 Inference 요청을 실제로 처리 중인 Pod인데, 그 Pod 자체가 종료되는 상황이라면 Stop 이벤트가 소비되지 않더라도 Chat Streaming Response 역시 즉시 중단된다.
 이 문제는 Stop Inference 이전 문제이기 때문에 서버 종료 시 Stop 이벤트 미처리 케이스는 Inference Stop 기능 설계 단계에서 의도적으로 고려 대상에서 제외했다.

2. Stop 요청 이후 일부 토큰이 추가로 전달되는 문제

Stop 요청을 전송한 이후에도 네트워크 지연과 Publish → Event Consume 사이의 Latency로 인해 3~4개 정도의 토큰이 추가로 전송되는 경우가 발생한다.

이는 스트리밍 기반 구조에서 완전히 제거하기 어려운 특성에 가깝고, 추가 공수를 들여서 해결할 만큼 큰 문제가 아니라고 판단하였다. 실제로 ChatGPT를 포함한 대부분의 LLM 스트리밍 서비스에서도 스트리밍 중지 후 몇 개의 토큰이 더 날라오는 것을 확인하였다.

물론 Chunk 단위에 chunk_id를 부여하고 Stop 요청 시 해당 ID와 비교하여 정확히 특정 시점에서 스트리밍을 중단하는 방식으로 구현은 가능하다고 생각했으나, 불필요한 공수라고 판단하여 하지 않았다.

따라서 현 구조에서는 즉시 완벽한 정지보다 충분히 빠르고 예측 가능한 정지를 선택했다.
비록 일부 한계는 존재하지만, 운영 환경과 사용자 경험을 기준으로 봤을 때 충분히 합리적인 Trade-off였다고 판단하고 있다.

Pub/Sub 구조 기반의 Stop 처리 구조는 분산 환경에서 발생하는 이벤트를 어떻게 효율적으로 처리할 것인가에 대한 고민을 해본 좋은 기능이었다.

단순 Redis 조회로 해결했을 때 비효율적인 부분을 금방 발견했기 때문에 더 좋은 구조로 발전시킬 수 있었다는 생각이 든다.

Event 기반 처리 도입
불필요한 데이터베이스 쿼리 제거
Stop 여부 확인을 In-memory 연산으로 효율적으로 치환

최종적으로 Inference Stop 기능은 Streaming 성능에 거의 영향을 주지 않으면서도 사용자 경험을 확실하게 개선하는 기능이 되었다.
 서버 종료 시 이벤트 미처리, Stop 요청 이후 소량의 토큰이 추가로 전달되는 문제 같은 한계도 존재하지만 이 상황에서 Trade-off를 고려해 적절한 선에서 구현하는 것도 개발자의 중요한 역할이다.

그런 의미에서 이번 Pub/Sub 기반 Stop 처리는 기술적인 고민부터 의사결정 과정까지 전부 의미있고 재밌는 기능이었다.