Redis 키스페이스 알림과 임대를 이용한 GPU 워커 풀 로드 밸런싱

GPU 풀은 상태 비저장(stateless) 웹 서버 풀과 같지 않으며, 이를 웹 서버처럼 로드 밸런싱하는 것은 여러분이 소유한 가장 비싼 하드웨어를 낭비하는 일입니다. 무거운 추론 작업을 실행하는 GPU는 완전히 점유됩니다. 여기에 두 번째 작업을 보내면 공유하지 않고 큐에 넣으며, 이제 다른 GPU는 유휴 상태로 있는 동안 한 요청은 다른 사람의 작업 뒤에서 몇 분간 기다리게 됩니다. 웹 로드 밸런싱에서 반사적으로 사용하는 라운드 로빈 및 최소 연결 방식은 GPU에게 '바쁨'이 무엇을 의미하는지 모르기 때문에 여기서는 둘 다 실패합니다. 이 게시물은 두 가지 Redis 기본 요소, 즉 GPU를 예약하는 임대(lease)와 GPU가 해제되는 즉시 알려주는 키스페이스 알림을 사용하여 바쁨을 인지하는 스케줄러를 구축하고, 이러한 설계를 강제하는 실패 모드를 다룹니다.

일반적인 로드 밸런싱이 GPU에서 실패하는 이유

일반적인 로드 밸런서의 가정은 GPU 작업에는 적용되지 않습니다. 웹 서버는 저렴하고 짧은 다수의 동시 요청을 처리하므로, 요청을 서버에 고르게 분산하는 것이 최적에 가까우며 오래된 부하 정보는 거의 문제가 되지 않습니다. GPU는 한 번에 하나씩, 비싸고 긴 무거운 작업을 처리하므로, 중요한 것은 특정 GPU가 지금 유휴 상태인지 여부뿐이며, 이에 대한 오래된 정보는 즉시 시작과 수 분간의 지연이라는 차이를 만듭니다.

라운드 로빈 방식은 바쁜 상태를 완전히 무시하고 작업 N을 아직 처리 중인 GPU에 아무렇지 않게 작업 N+1을 할당할 것입니다. 최소 연결 방식은 더 가깝지만 여전히 틀렸습니다. GPU 점유율이 아닌 연결 수를 세기 때문이며, 추론 중인 GPU는 연결이 전혀 없어도 완전히 사용 불가능한 상태일 수 있습니다. 필요한 것은 "이 특정 GPU가 유휴 상태인가"를 진실의 단위로 삼고, 상태가 변경되는 즉시 업데이트되며, 두 요청이 동일한 유휴 GPU를 동시에 점유할 수 없도록 강력한 예약 기능을 갖춘 스케줄러입니다. Redis는 이 두 가지를 모두 저렴하게 제공합니다.

리스는 예약입니다

핵심 상태는 GPU별 리스입니다. 즉, GPU가 사용 중일 때 존재하고 유휴 상태가 되면 사라지는 Redis 키입니다. GPU를 확보하는 것은 해당 리스를 가져가는 것이며, 이 확보는 원자적이어야 하므로 동일한 유휴 GPU를 두고 경쟁하는 두 요청이 모두 이길 수 없습니다.

// Claim GPU n only if its lease is absent. SET NX is atomic, so exactly one
// caller wins the race. The TTL is a safety net, covered below.
ok, err := rdb.SetNX(ctx, "gpu:lease:"+gpuID, workerID, leaseTTL).Result()
if err != nil {
	return "", err
}
if !ok {
	return "", errBusy // someone else holds this GPU
}

SetNX는 단일 명령으로 compare-and-swap을 수행합니다. 즉, 키가 아직 존재하지 않는 경우에만 키를 설정합니다. 경쟁에서 이긴 GPU는 리스를 얻고 작업을 시작합니다. 다른 모든 GPU는 ok가 false를 반환하는 것을 보고 다른 GPU를 시도하기 위해 넘어갑니다. 작업이 끝나면 워커는 리스를 삭제하고 GPU는 다시 유휴 상태가 됩니다. 만약 이것뿐이라면, 유휴 GPU를 찾기 위해 여전히 폴링해야 할 것이고, 폴링은 우리가 피하고자 하는 바로 그 지연 시간과 낭비의 원인이 됩니다. 두 번째 프리미티브는 폴링을 제거합니다.

키스페이스 알림은 "유휴"를 이벤트로 전환합니다

Redis는 키가 만료되거나 삭제될 때를 포함하여 키가 변경될 때마다 이벤트를 게시할 수 있습니다. 그 기능이 바로 키스페이스 알림이며, 이를 통해 대기 중인 요청은 타이머를 사용하여 주기적으로 확인하는 대신 GPU가 유휴 상태가 되는 즉시 알 수 있습니다. 관심 있는 이벤트 클래스인 일반 키 이벤트와 만료 이벤트에 대해 서버에서 한 번 활성화한 다음, 임대 키에 대한 패턴을 구독합니다.

// Subscribe to deletions and expirations of GPU lease keys. Each message
// means a specific GPU just became free.
pubsub := rdb.PSubscribe(ctx,
	"__keyevent@0__:del",
	"__keyevent@0__:expired",
)
for msg := range pubsub.Channel() {
	freedGPU := parseGPUID(msg.Payload) // the lease key that just vanished
	scheduler.notifyFree(freedGPU)
}

이제 흐름은 엔드투엔드 이벤트 기반이 됩니다. 요청은 SetNX를 사용하여 사용 가능한 GPU를 확보하려고 시도합니다. 모든 임대가 점유 중이면, 스피닝하는 대신 대기합니다. 어느 곳에서든 작업이 완료되어 임대를 삭제하거나 임대가 만료되는 순간, Redis는 이벤트를 게시하고, 대기 중인 요청이 깨어나 새로 사용 가능해진 GPU를 확보하려고 시도합니다. 고가의 하드웨어는 안정 상태에서 폴링 없이 Redis 메시지 하나가 이동하는 시간 안에 사용 중 상태에서 재할당 상태로 전환됩니다.

TTL이 존재하는 이유인 실패 모드

정상 경로에는 구멍이 있으며, 이는 가장 비싼 리소스를 발이 묶이게 만들기 때문에 가장 중요한 문제입니다. 워커가 GPU를 점유하고 작업을 시작한 다음 충돌하거나, 네트워크가 분할되거나, 프로세스가 실행 중에 종료된다고 가정해 봅시다. 임대를 삭제하는 코드 라인에 절대 도달하지 못합니다. 임대 키는 Redis에 영원히 남아 있고, 해당 GPU는 영원히 사용 중으로 표시되며, 더 이상 존재하지 않는 프로세스에서 키를 삭제할 것이 없으므로 키스페이스 이벤트가 절대 발생하지 않습니다. 한 번의 충돌로 풀에서 GPU가 영구적으로 제거됩니다.

이것이 임대가 TTL, 즉 위 SetNX 호출의 leaseTTL를 가지는 이유입니다. TTL은 워커가 만료되기 훨씬 전에 완료 시 임대를 삭제하는 정상적인 경우를 위한 것이 아닙니다. 이것은 비정상적인 경우를 위한 복구 메커니즘입니다. 보유자가 중단되면 임대는 저절로 만료되고, Redis는 만료 키스페이스 이벤트를 발생시키며, GPU는 사람의 개입 없이 자동으로 풀에 다시 합류합니다. 이 임대는 이전의 단일 실행기 작업 논의와 같은 의미의 진정한 임대입니다. 즉, 살아있는 보유자는 갱신하고 중단된 보유자는 상실하는 예약입니다.

이는 긴 작업이 하트비트처럼 TTL이 경과하기 전에 임대를 갱신해야 한다는 것을 의미합니다. 그렇지 않으면 복구 메커니즘에 의해 실행 중에 GPU를 잃게 됩니다. 가장 길 것으로 예상되는 작업보다 TTL을 여유 있게 길게 설정하거나, 길이를 한정할 수 없는 작업의 경우 임대를 주기적으로 하트비트하십시오. 트레이드오프는 익숙한 것입니다. 짧은 TTL은 충돌한 GPU를 빠르게 복구하지만 느린 작업을 빼앗을 위험이 있고, 긴 TTL은 느린 작업에는 안전하지만 충돌한 GPU를 더 오래 발이 묶인 상태로 둡니다. 기본값에서가 아니라 작업 길이 분포를 보고 선택하십시오.

점유와 대기 사이의 경쟁

사용 중 인지 스케줄러에는 작성하기는 쉽지만 재현하기는 어려운 미묘한 순서 버그가 있어 언급할 가치가 있습니다. 한 요청이 모든 GPU를 확인하고, 모두 사용 중임을 발견한 다음, 유휴 이벤트 피드를 구독한다고 상상해 보세요. 만약 확인과 구독 사이의 아주 짧은 시간 동안 GPU가 유휴 상태가 되면, 구독이 존재하기 전에 이벤트가 발생하고, 요청은 이를 수신하지 못하며, 몇 분 후에나 발생할 수 있는 다음 유휴 이벤트를 기다리게 됩니다. GPU는 유휴 상태이고 요청은 대기 중이지만, 서로의 상태를 알지 못합니다.

해결책은 순서를 뒤집어 구독한 후 다시 확인하는 것입니다. 먼저 유휴 피드를 구독한 다음, GPU 점유를 시도합니다. 점유에 성공하면 구독을 취소하고 실행합니다. 모든 리소스가 사용 중이라 실패하면 이미 구독한 피드에서 대기하고, 중요한 점은 이벤트가 도착했을 때 유휴 상태가 된 GPU가 자신의 것이라고 가정하는 대신 다시 점유를 시도해야 한다는 것입니다. 다른 대기자가 먼저 점유했을 수도 있기 때문입니다. 구독하고, 확인한 다음, 대기 후 재시도하는 것입니다. 이 순서는 그 간극을 없애고, 깨어날 때 재시도하는 방식은 유휴 상태가 된 GPU를 자신이 접근하기 전에 다른 누군가가 점유하는 여러 대기자 문제를 처리합니다.

부실 상태와 조정 패스

TTL이 있더라도 분산 상태에는 드리프트가 발생합니다. 작업자가 잠시 파티션된 상태에서 여전히 GPU를 점유하고 있다고 생각하는 동안 해당 임대가 삭제될 수 있습니다. 구독자가 이벤트 발생 순간에 재연결 중이었다면 만료 이벤트가 누락될 수 있습니다. 이벤트에만 의존하여 구축된 이벤트 기반 시스템은 시간이 지남에 따라 이러한 작은 불일치들이 누적됩니다. 해결책은 주기적인 조정입니다. 즉, 느린 간격으로 각 GPU 작업자에게 실제로 무엇을 하고 있는지 독립적으로 묻고, 현실에 맞게 Redis를 수정하는 것입니다. 이벤트는 일반적인 경우에 시스템을 빠르고 최신 상태로 유지합니다. 조정 패스는 이벤트가 놓친 드리프트를 포착하여 장기적으로 시스템의 정확성을 유지합니다. 이벤트로부터의 빠른 경로, 조정을 통한 정확성 안전장치는 모든 이벤트 기반 스케줄러에 좋은 기본 자세입니다.

더 간단한 설계로 충분할 때

이 메커니즘은 GPU가 희소하고 비싼 자원이며, 작업이 충분히 길어서 바쁜 GPU로 작업을 라우팅하는 것이 큰 비용을 초래하는 실수가 될 때 그 가치를 발휘합니다. 작업이 짧고 비용이 저렴하다면, 유휴 워커가 작업을 가져가는 단순한 큐 방식이 더 간단하고 그만큼 효과적입니다. 각 작업이 순식간에 끝나므로 약간 불완전하게 할당되더라도 그 비용은 무시할 수 있기 때문입니다. 배치를 대신 처리해주는 관리형 추론 서비스가 있다면, 그 서비스를 사용하고 이 모든 내용은 건너뛰십시오. 바로 이 문제를 해결하기 위해 비용을 지불하고 있는 것이기 때문입니다.

다음 세 가지가 모두 해당하는 경우에 임대 및 알림 스케줄러를 구축하십시오. GPU가 비싸서 유휴 시간이 실제 손실이 될 때, 작업이 길어서 잘못된 할당이 상당한 시간 동안 요청을 지연시킬 때, 그리고 관리형 엔드포인트를 대여하는 대신 직접 풀을 운영하고 있을 때입니다. 이러한 조건에서, 원자적 예약을 위한 Redis 임대와 즉각적인 핸드오프를 위한 키스페이스 알림은 폴링 없이 고가의 하드웨어를 계속 사용 중 상태로 유지하는 스케줄러를 제공합니다. 그리고 TTL과 조정(reconcile) 과정은 워커가 중단되거나 이벤트가 유실될 때 시스템을 정상적으로 유지해 줍니다.