LitServe 리뷰

최근 Lightning AI에서 AI 모델을 쉽게 서빙할 수 있도록 하는 프레임워크인 LitServe를 공개하였습니다.

https://lightning.ai/docs/litserve

Lightning AI | Turn ideas into AI, Lightning fast

 

LitServe는 FastAPI를 기반으로 구축된 AI 모델을 위한 사용하기 쉽고 유연한 서빙 엔진입니다. 배칭(batch), 스트리밍(streaming), GPU 오토스케일링(autoscaling)과 같은 기능을 제공하며, 모델별로 FastAPI 서버를 재구축할 필요를 없애줍니다. 특히, LitServe는 AI에 특화된 멀티 워커(multi-worker) 처리 덕분에 기본 FastAPI보다 최소 2배 빠릅니다.

 

Toy Example

import litserve as ls

class SimpleLitAPI(ls.LitAPI):
    def setup(self, device):
        self.model1 = lambda x: x**2
        self.model2 = lambda x: x**3

    def decode_request(self, request):
        return request["input"]

    def predict(self, x):
        squared = self.model1(x)
        cubed = self.model2(x)
        output = squared + cubed
        return {"output": output}

    def encode_response(self, output):
        return {"output": output}

if __name__ == "__main__":
    api = SimpleLitAPI()
    server = ls.LitServer(api, accelerator="gpu")
    server.run(port=8000)

 

 

LitAPI

LitAPI 클래스는 모델에 대한 요청을 처리하는 프로세스를 작성합니다. 기본적으로 setup, decode_request, predict, encode_response의 4가지 api를 기본적으로 작성합니다.

class SimpleLitAPI(ls.LitAPI):
    def setup(self, device): ...

    def decode_request(self, request): ...
    
    def predict(self, x): ...

    def encode_response(self, output): ...

 

서버가 배칭(batch)를 사용하지 않는 경우 한 번에 하나의 입력을 처리합니다. 

The sequence of LitAPI method calls for a single request.

서버가 시작될 때 처음 한 번 setup이 호출됩니다. setup을 사용하여 다음 작업을 수행할 수 있습니다:

• 데이터 로드
• 모델 초기화
• 데이터베이스 연결 설정
• 임베딩(embeddings) 로드
• 기타 등등...

서버가 초기화된 후, 각 요청에 대해 다음 메서드들이 이 순서로 호출됩니다:

• decode_request: 들어오는 요청 페이로드를 모델이 사용할 수 있는 입력으로 변환합니다.
• predict: decode_request의 출력을 사용하여 모델에서 추론을 수행합니다.
• encode_response: predict의 출력을 응답 페이로드로 변환합니다.

 

배치(batch)가 활성화되면 서버는 최대 max_batch_size까지 요청을 그룹화하고 이를 예측 메서드에 함께 보냅니다. decode_request의 출력은 배치 크기로 그룹화됩니다. 그런 다음 해당 그룹이 예측 메소드로 전송됩니다.

With batching, requests are batched and sent at once through this sequence of LitAPI methods.

 

서버가 시작될 때 처음 한 번 setup이 호출됩니다. 이는 위의 Not Batched와 같습니다.

서버가 초기화된 후, 배칭(batching)이 활성화된 경우, 서버는 LitServer(max_batch_size=n)에서 지정한 n개의 요청을 배치로 그룹화하고, 이 배치를 한 번에 LitAPI를 통해 다음 순서로 처리합니다:

• decode_request: 들어오는 요청 페이로드를 모델이 사용할 수 있는 입력으로 변환합니다.
• batch: 여러 요청을 어떻게 그룹화할지 정의할 수 있습니다. 지정하지 않으면, LitServe가 자동으로 이를 처리하는 휴리스틱(heuristics)을 사용합니다.
• predict: decode_request의 출력을 사용하여 모델에서 추론을 수행합니다.
• unbatch: 배치를 개별 응답으로 어떻게 매핑할지 정의할 수 있습니다. 지정하지 않으면, LitServe가 자동으로 이를 처리하는 휴리스틱을 사용합니다.
• encode_response: predict의 출력을 응답 페이로드로 변환합니다.

 

Batched 케이스는 Not Batched 케이스와는 다르게 batch, unbatch를 추가적으로 사용해야 합니다. batch 메서드는 decode_request에서 들어오는 개별 입력들의 List인 inputs를 인자로 받습니다. 이 inputs 인자는 하나의 배치가 됩니다. unbatch 메서드는 predict 메서드의 배치 처리된 결과값을 encode_response에 전달하기 전에 개별 결과로 분할합니다. 

class CustomBatchedAPI(ls.examples.SimpleBatchedAPI):
    def batch(self, inputs):
        return np.stack(inputs)

    def unbatch(self, output):
        return list(output)

 

 

LitServer

LitServer 객체는 모델 서빙의 요청 스케일링, 입력 배칭(batching), 응답 스트리밍(streaming), 그리고 GPU 자원 관리 등의 인프라 측면을 처리합니다. 모델 로직이 LitAPI로 설정되면, LitServer는 서빙 프로세스를 자동으로 최적화합니다. 이 시스템은 단순성을 염두에 두고 설계되어 초보자도 모델을 효율적으로 배포할 수 있으며, 전문가에게는 성능을 미세 조정할 수 있는 고급 제어 기능을 제공합니다.

LitServer lifecycle.

 

Parameters

from litserve import LitServer

LitServer(
    lit_api: LitAPI,
    accelerator: str = "auto",
    devices: Union[str, int] = "auto",
    workers_per_device: int = 1,
    timeout: Union[float, bool] = 30,
    max_batch_size: int = 1,
    batch_timeout: float = 0.0,
    api_path: str = "/predict",
    stream: bool = False,
    spec: Optional[LitSpec] = None,
)

LitServer 생성자에서의 인자를 보면 기존의 다른 서빙 프레임워크의 설정값들과 크게 다르지 않은 것을 볼 수 있습니다. 비교하자면 다른 프레임워크에 비해 설정값들이 약간은 간소화되어 있다는 느낌을 받을 수 있는데요, 이러한 부분들은 LitServe의 목표인 쉽게 배포하자는 기조에 맞다는 느낌을 줍니다.

 

 

LitServe Features

다음으로 LitServe에서 제공하는 기능들에 대해서 알아보겠습니다. 제가 생각했을 때 중요하다고 생각하는 부분에는 🔥 표시를 해두었으니, 내용을 빠르게 읽고 싶으신 분들은 해당 표시들을 확인해주세요.

 

Authentication

LIT_SERVER_API_KEY 환경변수를 설정해서 Authentication 기능을 쉽게 사용가능하도록 합니다. 간단하게 서버 실행시 LIT_SERVER_API_KEY를 지정해주면 됩니다.

LIT_SERVER_API_KEY=A_SECRET_KEY python server.py

 

client에서는 headers에 X-API-Key로 해당 키값을 전달하여 인증합니다.

import requests

# Define the API key in the request headers
headers = {"X-API-Key": "A_SECRET_KEY"}

response = requests.post("http://127.0.0.1:8000/predict", json={"input": 4.0}, headers=headers)
print(response.json())

 

Autoscaling 🔥

LitServe는 여러 장치(GPU, TPU, CPU 등) 또는 여러 시스템에서 동일한 모델의 여러 복사본을 병렬로 제공하여 높은 트래픽을 처리할 수 있습니다.

LitServe handles high throughput serving within and across machines.

devices 인자를 설정하여 머신에서 사용 가능한 최대 GPU 수를 지정할 수 있고, workers_per_device 인자를 설정하여 각 GPU 또는 CPU 디바이스에서 모델의 여러 복사본을 서빙하도록 LitServe를 확장할 수 있습니다. 또한 num_api_servers 인자를 설정하여 API Server를 여러 프로세스에서 시작할 수 있습니다.

개인적으로는 이러한 기능을 Autoscaling이라고 부를 수 있는지에 대해서는 의문입니다. Autoscaling이라 하면 현재 수요에 따라 활성 서버 또는 리소스 수를 조정하는 과정을 의미합니다. LitServe 코드를 살펴봤을 때 수요에 따라 서버 수나 리소스를 자동으로 조정하는 기능은 없어 보입니다.

 

Batching 🔥

Batching 처리는 위에서 설명했듯이 요청을 grouping하여 처리한 뒤, 개별 응답으로 보내주는 프로세스를 거칩니다. 이 과정에서 max_batch_size와 batch_timeout을 설정할 수 있습니다. 일반적인 다른 서빙 프레임워크와 마찬가지로 dynamic batching을 기본으로 제공하고 있습니다.

 

OpenAI Spec 🔥

현대 AI 언어 모델 배포에서는 Open AI spec이라는 서버 입력/출력 표준이 정착되었습니다. 이 표준은 서버로 들어가는 입력의 형식과 토큰 수 및 AI 모델에서 일반적으로 사용하는 기타 파라미터와 같은 서버에서 반환되는 정보의 형식을 보장합니다. LitServe API는 LitServe OpenAI spec을 사용하여 Open AI 표준을 준수할 수 있습니다.

import litserve as ls

class SimpleLitAPI(ls.LitAPI):
    def setup(self, device):
        self.model = None

    def predict(self, prompt):
        # `prompt` is a list of dictionary containing role and content
        # example: [{'role': 'user', 'content': 'How can I help you today?'}]
        yield "This is a sample generated output"

if __name__ == "__main__":
    # Enable the OpenAISpec in LitServer
    api = SimpleLitAPI()
    server = ls.LitServer(api, spec=ls.OpenAISpec())
    server.run(port=8000)

 

Request, response format

LitServe는 요청과 응답을 검증하고 형식을 지정하여 오류를 방지하고 안정적인 시스템을 보장할 수 있습니다. 이미지, 텍스트, 오디오, JSON, 파일 등 다양한 유형의 요청과 응답을 처리할 수 있습니다.

 

Serve LLMs 🔥

LitServe는 모든 AI 모델을 제공하도록 설계된 범용 프레임워크입니다. LitServe 자체에는 LLM에 특화된 기능들(kv-caching, batching, kernels, 등)에 대한 고려는 되어있지 않습니다. LLM 관련 솔루션의 경우 LitServe를 기반으로 구축된 LitGPT를 사용도록 가이드하고 있습니다. 물론 LitServe와 다른 LLM 모델 추론 프레임워크를 integration하여 사용할 수도 있습니다.

Use a LitGPT LLM

import litgpt
import litserve as ls

class SimpleLitAPI(ls.LitAPI):
    def setup(self, device):
        self.llm = litgpt.LLM.load("checkpoints/meta-llama/Meta-Llama-3-8B-Instruct")

    def decode_request(self, request):
        return request["prompt"]

    def predict(self, prompt):
        return self.llm.generate(prompt, max_new_tokens=200)

    def encode_response(self, output):
        return {"output": output}

if __name__ == "__main__":
    api = SimpleLitAPI()
    server = ls.LitServer(api)
    server.run(port=8000)

 

Use a Hugging Face LLM

from transformers import pipeline
import litserve as ls

class HFAPI(ls.LitAPI):
    def setup(self, device):
        self.pipe = pipeline("text-generation", model="meta-llama/Meta-Llama-3-8B-Instruct", trust_remote_code=True)

    def decode_request(self, request: dict):
        return request["prompt"]

    def predict(self, prompt):
        messages = [{"role": "user", "content": prompt}]
        return self.pipe(messages)

    def encode_response(self, output):
        return {"output": output[0]["generated_text"][-1]["content"]}

if __name__ == "__main__":
    api = HFAPI()
    server = ls.LitServer(api)
    server.run(port=8000)

 

Use a vLLM model

from vllm import LLM, SamplingParams
import litserve as ls

class VLLMLlamaAPI(ls.LitAPI):
    def setup(self, device):
        self.llm = LLM(model="meta-llama/Meta-Llama-3-8B-Instruct", quantization="fp8")

    def decode_request(self, request):
        return request["prompt"]

    def predict(self, prompt):
        sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=100)
        outputs =  self.llm.generate(prompt, sampling_params)
        return outputs[0].outputs[0].text

    def encode_response(self, output):
        return {"output": output}

if __name__ == "__main__":
    api = VLLMLlamaAPI()
    server = ls.LitServer(api)
    server.run(port=8000)

 

 

리뷰

범용적인 AI 모델 서빙 프레임워크로써의 LitServe는 굉장히 쉽고 직관적인 interface로 이루어져 있어서, 본인이 지향하는 'AI 모델을 위한 사용하기 쉽고 유연한 서빙 엔진'이라는 목표에 부합하는 것 같습니다. 제가 경험해봤던 TorchServe, Triton Inference Server와 같은 범용 모델 서빙 프레임워크와 비교하였을 때, 확실하게 유저 친화적이며 사용하기 쉽고 유연하게 디자인되어 있다는 인상을 받았습니다.

하지만 프로젝트가 초창기이다 보니 기존 프레임워크들에 누적되어 있는 노하우들을 완벽히 옮겨오지는 못한 것 같습니다. 예를 들어 모델을 배포하고 운영하다보면 observability에 대해서도 고민하게 되는데, 이러한 부분들에 대한 기능이 현재로써는 없어 보여서 아쉽습니다.

개인적으로는 Lightning AI에서 기존에 PyTorch Lightning과 Fabric 등으로 보여줬던, 깔끔하게 정돈된 코드 스타일을 갖춘 사용하기 쉬운 추상화된 인터페이스를 제공하겠다는 철학에 굉장히 깊게 공감하고 있으며 응원하고 있습니다. LitServe도 이러한 철학에 맞춰서 작업되고 있고, 앞으로 모델 배포의 여러 노하우들에 대한 추가 기능들이 구현된다면 사용하기 쉽고 강력한 모델 서빙 엔진이 되지 않을까 기대하고 있습니다 ☺️.