LLM 서버 아껴쓰고 나눠쓰고 바꿔쓰고 다시쓰자
문제
로컬 모델을 이용해서 서버 구동하기까지 하고 보니까, 생각보다 생성이 느리다. 물론 전문 서비스 업체에서 열심히 최적화한 시스템보다야 느린건 어찌보면 당연하지만, 이정도는 하겠지 기대했던 성능에 미치지 못해서 못쓸 물건이 되어버리는 건 곤란하다.
그중 가장 신경쓰였던 문제는 문장을 다다닥 입력하고 다음에 모델이 뭔 소리 하나 하고 생성을 기다리는데 수십초 이상 소요되는 문제였다. 파일을 닫았다 다시 열어서 같은 지점에서 생성을 기다리는 경우 2~3초 내에 뭔가 나오는데, 파일을 열어놓고 한참 작업을 하다가 생성을 시키는 경우 딜레이가 생긴다? 그사이에 GPU가 100%를 계속 친다? 여러모로 좋은 시그널은 아니었다.
뭐가 문제일까?
누구나 생각해볼 수 있는 가설은 문장을 입력하는 중간에도 계속해서 생성 요청이 들어가고, 서버에서 그걸 처리하느라 바빠서 막상 내가 필요한 요청은 대기열을 기다리느라 시간이 오래 걸리는 거라고 생각해볼 수 있다.
생성 요청 하나 날릴 때마다 로그를 찍어봄으로써 이 가설이 실제로 발생하고 있음을 확인할 수 있었다.
왜 요청을 여러개 날리게 되나?
에디터에서 문서 변경되는 동안 LSP 서버로 계속 기존 요청을 취소하고, 거의 곧바로 새 요청을 보내게 되는데, 기본적으로 LLM 서비스가 무제한 동시 요청이 가능하다고 가정하고 있기 때문에 요청들이 제한없이 마구 실행된다.
해결방안: 그래서 어떻게 해야하나?
요청을 날리는 작업이 값비싼 작업이니까, 이걸 동시에 한 작업만 들어갈 수 있게 명시적으로 처리하면 되지 않을까?
큰 그림: 리버스 프록시 서버
물론 서버를 새로 짜는게 더 확실한 방법일 수는 있겠지만, 수고가 많이 드는 작업이 될 것 같아서 최대한 적은 노력으로 구현할 수 있는 방법으로 프록시를 생각해 봤다. 모든 생성 요청을 먼저 프록시에서 받아서 들고 있으면서, LLM 서버는 동시에 하나의 요청만 처리할 수 있도록 중간에서 대기열 관리를 하는 방식을 생각해 봤다.
첫 시도 - Mutex를 이용한 명시적인 Lock
단순히 요청을 처리하는 과정을 Lock으로 감싸서 하나의 쓰레드만 들어갈 수 있게 하면 되지 않을까 시도해 봤다. 대충 이런 모습이다.
#[cfg(test)]
mod tests {
use std::{sync::Arc, time::Duration};
use futures::future;
use tokio::sync::Mutex;
#[tokio::test]
async fn without_spawn() -> anyhow::Result<()> {
let mutex = Arc::new(Mutex::new(1));
let m1 = mutex.clone();
let h1 = async move {
let guard = m1.lock().await;
dbg!(&guard); // try to get value
tokio::time::sleep(Duration::from_secs(1)).await;
dbg!(&guard); // check lock is still there
};
let (h11, flag1) = future::abortable(h1);
let m2 = mutex.clone();
let h2 = async move {
// flag1.abort(); // 안돼~~~~
let guard = m2.lock().await;
dbg!(&guard); // try to get value
tokio::time::sleep(Duration::from_secs(1)).await;
dbg!(&guard); // check lock is still there
};
tokio::join!(h11, h2);
Ok(())
}
}
m1이 먼저 실행되고 5ms가 지난 다음 m2가 실행되지만, m1이 1초동안 Lock을 잡고 있기 때문에 m2는 m1이 끝난 뒤에 실행되었으면 좋겠다는 작은 바램을 담고 있다.
그치만 이 희망은 금방 산산조각 나고 말았는데, 사실 요청들이 취소되고 있었기 때문.
요청이 취소되면 어떤 일이 일어나나요?
간단히 말하자면 모든 취소 작업이 거의 자동으로 아키텍쳐를 따라 전파되어 요청을 처리하지 않았던 것처럼 되돌아가게 되는데, 그 흐름이 딱 llama.cpp 서버 직전까지 전달되는게 모든 문제의 원인이라고 할 수 있다.
대기중인 요청이 취소되는 시점에 프록시 서버에서도 요청 취소를 인지할 수 있게 되고, 이때 프록시 서버에서도 진행중인 작업을 취소하게 된다. 아예 요청이 LLM 서버로 날아가지 않은 상황이었으면 참 좋았겠는데, 현실은… 에디터에서 변경사항이 발생할 때마다, 기존 자동완성 요청을 취소하고 새 자동완성 요청을 실행하게 되어있고, 기존 자동완성 요청을 취소하는 시점에 LSP 서버에서도 LLM 프록시에 보냈던 요청 연결을 끊어버리고, 프록시에서는 그 연결 끊김을 인지해서 자기가 가지고 있던 Future의 실행을 중단시키게 된다. Future를 실행하는 과정에서 가지고 있던 MutexGuard도 해제되면서 자연스럽게 Mutex가 다시 사용가능한 상태로 되돌아가게 되고 다음에 요청이 다시 들어오면 아직 llama.cpp 서버가 작업을 진행중임에도 불구하고 요청이 전달되는 문제가 발생하게 된다.
(헉;; 설명이 이렇게 길어질 줄이야!)
시도 2 - 어떻게 하면 Future를 취소하지 못하게 할 수 있을까?
그냥 이리저리 Future를 다루는 기능들을 가지고 찔러보다가 tokio::spawn
에 이르게 되었다. spawn을 통해 만들어진 JoinHandle은 자체적인 abort 기능을 갖는
대신에 future::abortable로 감싸서 만든 abort 기능이 동작하지 않는 것 같더라구.
내 입장에선 요청이 취소되어도 계속 실행되는 future가 필요했기 때문에 바로 이렇게 spawn하는게 맞는 방향이라고 생각하게 되었고, 거기에 따라 코드를 작성해 봤는데…
#[tokio::test]
async fn pass_the_mutexguard() -> anyhow::Result<()> {
let mutex = Arc::new(Mutex::new(1));
let m1 = mutex.clone();
let h1 = async move {
let guard = m1.lock().await;
tokio::spawn(async move {
dbg!(&guard);
tokio::time::sleep(Duration::from_secs(1)).await;
dbg!(&guard);
}).await
};
let (h1, flag1) = future::abortable(h1);
let m2 = mutex.clone();
let h2 = async move {
flag1.abort();
let guard = m2.lock().await;
dbg!(&guard); // try to get value
tokio::time::sleep(Duration::from_secs(1)).await;
dbg!(&guard); // check lock is still there
};
tokio::join!(h1, h2);
Ok(())
}
대충 이런 느낌으로 코드를 작성해 봤지만, 잘 안된다.
가장 큰 문제는 m1에서 생성한 guard가 tokio::spawn으로 옮겨지는 과정에서 lifetime이 애매~ 해지는 문제가 있다는 거다. spawn하는 시점에서 작업이 언제 끝날 지 모르니까 guard를 영원히 가지고 가야 하는데, guard는 m1을 잠그는 시점에서 만들어진 거라서 m1이 guard만큼 오래 살아남아야 하고, 그러면 그것을 가지는 h1까지 다함께 영원히 오래 살아남아야만 말이 되는데 실제론 h1은 필멸자라서 영원히 모두 행복하게 살 방법은 없다…는 것이다.
이걸 피하려면 guard가 spawn 블록 안에서 선언되어서 다함께 불멸자가 되는 방법이 하나 있긴 한데 그러려면 lock을 잡기 전에 h1이 취소되는 경우에도 이미 불멸자 블록에서 돌아가고 있는 작업은 여전히 lock을 기다리고, 요청을 처리하게 되기 때문에 쓸데없이 LLM을 낭비하는 결과를 가져오게 된다.
여기까지 와서 보니까 이게 그냥 lock을 조금 세심하게 잡는 문제가 아니라 그때그때 엿가락처럼 왔다갔다하는 lifetime을 관리해야 하는 문제고, Rust가 잘 관리하는 컴파일 타임에 lifetime 분석하기랑은 잘 맞지 않는 영역이 된다는 것을 깨닫게 되었다. 아… unsafe로 가야 하나? ManuallyDrop?
시도 3 - Mutex가 아니라도 괜찮아. Semaphore가 있잖아
생각해 보니 LLM 서버로 요청을 보내는 것 자체는 멀티쓰레드를 잘 지원하기 때문에 굳이 Mutex를 고집할 필요는 없었다. 단지 lock을 잡고, 잡는 녀석들이 대기를 하고 언젠가 되돌려주기만 하면 된다. 괜히 MutexGuard의 lifetime을 억지로 늘리려고 애쓰지 않아도 대충 Atomic counter 역할만 있으면 되는 것 아닌가.
tokio Semaphore 문서 에서 아주 비슷한 상황에 대해 이야기하고 있다. Semaphore가 단지 요청할 수 있는 토큰을 담고, 토큰을 가져갈 때 Semaphore를 이용해서 대기를 하고, 시간에 따라 다시 토큰을 생성하는 식으로 rate limit을 구현하는 방식인데, 우리는 단지 LLM 서버에서 응답이 돌아올 때 토큰을 생성하는 식으로 살짝만 변형하면 되는 것이다.
일반적인 사용의 경우 MutexGuard와 비슷한 SemaphorePermit이란 구조체가 semaphore를 돌려놓는 역할을 하는데, 내 경우 lifetime을 이용할 수 없기 때문에 문서에서 하듯 그냥 돌려놓는 것을 잊어버리게 하고, 대신 LLM이 사용가능해진 시점에 Semaphore에 알려주는 방식으로 했다.
결론: 이젠 진짜 생성 자체는 그럴듯한 latency에 되긴 하는데…
문제 자체는 아주 심플한데 간단하게 해결하기가 의외로 어려웠다. 아직 Rust에 대한 이해가 부족했던 때문일까 아니면 맞지 않는 도구를 어거지로 우겨넣으려다가 생긴 문제일까… javascript였으면 그냥 Promise chain 같은걸로 대충 큐 구현하고 했겠지? 간만에 포스팅도 엄청 늘어지게 되었다.
CodeLlama 13b는 생성 퀄리티가 괜찮은데 4090에서나 적당한 레이턴시가 나오고, 맥북에선 너무 느리다. 7b 모델도 여전히 느릴 것 같아서 stablecode-3b나 deepseek-coder-1.3b 같은 아주 작은 모델로 변경해 볼 생각.