Continuous Batching 알고리즘

class ContinuousBatchingScheduler:
    def __init__(self, max_batch_size=4):
        self.max_batch_size = max_batch_size
        self.waiting_queue = []  # 대기 중인 새 요청
        self.active_batch = []   # 현재 GPU에서 연산 중인 요청 슬롯

    def step(self):
        """매 토큰 생성(Forward Pass) 직후 호출되는 스케줄링 로직"""
        
        # 1. 종료된 요청 방출 (Eviction)
        retained_batch = []
        for req in self.active_batch:
            if req.is_finished(): # EOS 토큰이 나왔거나 최대 길이에 도달했는지 확인
                print(f"[{req.id}] 생성 완료. 배치에서 제거 및 결과 반환.")
            else:
                retained_batch.append(req)
        self.active_batch = retained_batch

        # 2. 빈자리만큼 새 요청 투입 (Injection)
        available_slots = self.max_batch_size - len(self.active_batch)
        while available_slots > 0 and self.waiting_queue:
            new_req = self.waiting_queue.pop(0)
            self.active_batch.append(new_req)
            print(f"[{new_req.id}] 빈 슬롯에 새 요청 즉시 투입.")
            available_slots -= 1

        # 3. 구성된 새로운 배치로 다음 1토큰 동시 추론 진행 (GPU Forward Pass)
        if self.active_batch:
            # model.forward(self.active_batch)
            pass

from vllm.engine.arg_utils import AsyncEngineArgs
from vllm.engine.async_llm_engine import AsyncLLMEngine

# Continuous Batching 스케줄러가 최적의 결정을 내릴 수 있도록 제어값을 부여
engine_args = AsyncEngineArgs(
    model="meta-llama/Llama-3-8B",
    
    # [핵심 파라미터 1] max_num_seqs
    # Continuous Batching 큐에서 '동시에 GPU 연산 슬롯에 올릴 수 있는 최대 문장(요청) 개수'
    # 이 값이 너무 크면 VRAM OOM이 발생하고, 너무 작으면 배치에 빈자리가 생겨 Throughput이 떨어짐
    max_num_seqs=256,
    
    # [핵심 파라미터 2] max_num_batched_tokens
    # 한 번의 스텝(Iteration)에서 처리할 수 있는 총 토큰 수의 합 (Prompt 토큰 + 생성 토큰)
    # 새 요청(프롬프트)이 중간에 Inject 될 때, 이 토큰 수 한계를 넘지 않는 선에서만 투입을 허용함
    max_num_batched_tokens=8192,
    
    # 앞서 학습한 PagedAttention과 결합되어 VRAM을 효율적으로 재활용
    gpu_memory_utilization=0.9
)

# 비동기 엔진 기동
engine = AsyncLLMEngine.from_engine_args(engine_args)

from fastapi import FastAPI
import asyncio
from prometheus_client import Gauge
import logging

logger = logging.getLogger(__name__)
app = FastAPI()

# 1. Prometheus 메트릭 정의 (그라파나 대시보드 연동용)
# 대기 중인 요청 수, 현재 처리 중인 요청 수, GPU 캐시 사용률을 추적합니다.
METRIC_PENDING_REQUESTS = Gauge('vllm_pending_requests', 'Number of requests waiting in queue')
METRIC_RUNNING_REQUESTS = Gauge('vllm_running_requests', 'Number of requests currently in continuous batch')
METRIC_KV_CACHE_USAGE = Gauge('vllm_kv_cache_usage_percent', 'GPU KV Cache usage percentage')

async def log_continuous_batching_stats(engine):
    """엔진 내부의 스케줄러 통계를 주기적으로 빼내어 메트릭으로 노출하는 백그라운드 태스크"""
    while True:
        try:
            # vLLM 엔진 내부의 현재 스케줄링 상태(Stats) 강제 조회
            stats = await engine.get_decoding_stats()
            
            # 메트릭 게이지 업데이트
            METRIC_PENDING_REQUESTS.set(stats.num_requests_waiting)
            METRIC_RUNNING_REQUESTS.set(stats.num_requests_running)
            METRIC_KV_CACHE_USAGE.set(stats.gpu_cache_usage)
            
            # 특정 임계치 초과 시 경고 로깅 (Auto-scaling 트리거 조건으로 활용 가능)
            if stats.num_requests_waiting > 100:
                logger.warning(f"🚨 Continuous Batching Queue 포화 상태! 대기 중: {stats.num_requests_waiting}건")
                
        except Exception as e:
            logger.error(f"메트릭 수집 중 오류: {e}")
            
        await asyncio.sleep(2.0) # 2초마다 수집

@app.on_event("startup")
async def startup_event():
    # 애플리케이션 시작 시 백그라운드 모니터링 루프 실행
    asyncio.create_task(log_continuous_batching_stats(engine))