Databricks Model ServingでSGLangを使ったLLM Chatエンドポイントを作る

# 導入

以下の記事でSGLangのスループットがかなり高いことを取り上げました。

https://qiita.com/isanakamishiro2/items/45b7c6489321027a3cbe

高いスループットといえば、Databricksには[プロビジョニングされたスループットの基盤モデルAPI](https://docs.databricks.com/ja/machine-learning/foundation-models/deploy-prov-throughput-foundation-model-apis.html)があり、本番運用を想定した高スループットエンドポイントを構築できます。

ただ、これに対応しているLLMはLlamaやMistral、DBRXなど一部であり、Google GemmaやQwenなどは未サポートです。

ということは、対応モデルは素直にプロビジョニングされたスループットの基盤モデルAPIを使用し、対応していないものはDatabricksモデルサービングの中でSGLangを使えばいいんじゃない？と思ったので、実際にやってみました。

流れは以前の下記記事とだいたい同じです。

https://qiita.com/isanakamishiro2/items/69ad78a157380d5182f9

構築環境はDatabricks on AWS、
DBRは15.4ML LTS(Beta)、インスタンスタイプはg5.xlarge(GPU)です。

# Step1. パッケージインストール

SGlangに必要なパッケージをインストール。また、```mlflow```も最新化しておきます。

```python
# transformers 最新化
%pip install -U transformers

# vLLM
%pip install vllm

# flashinfer
%pip install https://github.com/flashinfer-ai/flashinfer/releases/download/v0.1.3/flashinfer-0.1.3+cu121torch2.3-cp311-cp311-linux_x86_64.whl

# sglang
%pip install "sglang[srt]"

# mlflow
%pip install "mlflow-skinny[databricks]>=2.15.1"

dbutils.library.restartPython()
```

# Step2. カスタムPyFuncモデルの定義

MLflowに登録するためのカスタムPyFuncモデルクラスを定義します。

詳細は以下の公式Docを参照ください。

https://mlflow.org/docs/latest/llms/transformers/tutorials/conversational/pyfunc-chat-model.html#Customizing-the-model

```python
from typing import List, Union, Iterator
import uuid

import mlflow
from mlflow.types.llm import ChatResponse, ChatMessage, ChatParams
from pprint import pprint

import sglang as sgl
from sglang.lang.chat_template import get_chat_template


@sgl.function
def inference(s, messages:List[ChatMessage]):
    """ シンプルなMessageの推論処理 """
    regex = None

    for m in messages:
        if m.role == "system":
            s += sgl.system(m.content)
        elif m.role == "user":
            s += sgl.user(m.content)
        elif m.role == "assistant":
            s += sgl.assistant(m.content)
        elif m.role == "regex":
            regex = m.content

    s += sgl.assistant(sgl.gen("answer", regex=regex))

# SGLangを使ったカスタムChatモデル
class SGLangChatModel(mlflow.pyfunc.ChatModel):
    def __init__(
        self,
        model_name: str,
        chat_template_name: str,
        mem_fraction_static: float = 0.80,
        disable_cuda_graph: bool = False,
    ):
        self.model_name = model_name
        self.chat_template_name = chat_template_name
        self.mem_fraction_static = mem_fraction_static
        self.disable_cuda_graph = disable_cuda_graph
        self.runtime = None

    def load_context(self, context):
        model_directory = context.artifacts["llm-model"]

        runtime = sgl.Runtime(
            model_path=model_directory,
            mem_fraction_static=self.mem_fraction_static,
            disable_cuda_graph=self.disable_cuda_graph,
        )
        if self.chat_template_name:
            runtime.endpoint.chat_template = get_chat_template(self.chat_template_name)

        sgl.set_default_backend(runtime)

        self.runtime = runtime

    def predict(
        self,
        context,
        messages: List[mlflow.types.llm.ChatMessage],
        params: mlflow.types.llm.ChatParams,
    ):
        # 推論設定の構築
        _params = self._build_inference_settings(params)
        state = inference.run(
            messages=messages,
            temperature=_params.temperature,
            top_k=_params.top_k,
            top_p=_params.top_p,
            max_new_tokens=_params.max_tokens,
            frequency_penalty=_params.frequency_penalty,
            presence_penalty=_params.presence_penalty,
            stop=_params.stop,
            stream=False,
        )

        id = str(uuid.uuid4())
        full_text = state.get_var("answer")
        meta_info = state.get_meta_info("answer")
        return self._build_final_response(id, full_text, meta_info)

    def _build_inference_settings(self, params):
        # コピーを作って内容を補正
        _params = ChatParams(**params.to_dict())

        _params.max_tokens = _params.max_tokens or 50
        _params.stop = _params.stop or ()
        _params.top_k = _params.top_k or -1
        _params.top_p = _params.top_p or 1.0
        _params.frequency_penalty = _params.frequency_penalty or 0.0
        _params.presence_penalty = _params.presence_penalty or 0.0

        return _params

    def _build_final_response(self, id, final_text, meta_info):
        """最終レスポンスを生成する"""

        usage = {
            "prompt_tokens": meta_info["prompt_tokens"],
            "completion_tokens": meta_info["completion_tokens"],
            "total_tokens": meta_info["prompt_tokens"] + meta_info["completion_tokens"],
        }

        finish_reason = meta_info["finish_reason"]

        response = {
            "id": id,
            "model": self.model_name,
            "choices": [
                {
                    "index": 0,
                    "message": {
                        "role": "assistant",
                        "content": final_text,
                        "name": "answer",
                    },
                    "finish_reason": finish_reason,
                }
            ],
            "usage": usage,
        }

        return ChatResponse(**response)

    def __getstate__(self):
        state = self.__dict__.copy()
        del state["runtime"]
        return state
    
    def __del__(self):
        if self.runtime:
            self.runtime.shutdown()
```

# Step3. モデルの登録

作成したカスタムクラスを利用して、Unity Catalog上にMLflowモデルを登録します。

まずはモデルをロギングする際に必要な各種パラメータの準備。

モデルのパスを変数にいれていますが、今回は以下のモデルをあらかじめUnity Catalog Volumes上にダウンロードしています。
Gemma2 2Bを日本語ファインチューニングしたモデルです。パラメータサイズ比でかなり高い性能を示しています。

https://huggingface.co/HODACHI/EZO-Common-T2-2B-gemma-2-it


```python
import mlflow
import os

mlflow.set_registry_uri("databricks-uc")

model_path = "/Volumes/training/llm/model_snapshots/models--HODACHI--EZO-Common-T2-2B-gemma-2-it/"
model_name = "EZO_Common_T2_2B_gemma_2_it_sglang_chatmodel"
registered_model_name = f"training.llm.{model_name}"

extra_pip_requirements = [
    "torch==2.4.0 --index-url https://download.pytorch.org/whl/cu121",
    "vllm==0.5.4",
    "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.1.3/flashinfer-0.1.3+cu121torch2.3-cp311-cp311-linux_x86_64.whl",
    "sglang[srt]",
]

pip_requirements = mlflow.pyfunc.get_default_pip_requirements() + extra_pip_requirements

input_example = {
    "messages": [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Hello!"},
    ],
    "max_tokens": 100,
    "temperature": 0.0,
}
```

カスタムクラスのインスタンスを作成して、ロギングを実行します。

```python

model = SGLangChatModel(
    model_name=model_name,
    chat_template_name="gemma-it",
    mem_fraction_static=0.7,
)

with mlflow.start_run() as run:
    _ = mlflow.pyfunc.log_model(
        artifact_path="model",
        python_model=model,
        artifacts={
            "llm-model": model_path,
        },
        pip_requirements=pip_requirements,
        input_example=input_example,
        example_no_conversion=True,
        await_registration_for=1200,  # モデルサイズが大きいので長めの待ち時間にします
        registered_model_name=registered_model_name,  # 登録モデル名 in Unity Catalog
    )
```

# Step4. エンドポイントの作成

モデルの登録が無事できたら、モデルサービング機能を使ってエンドポイントを作成します。

```python
import requests
import json
from mlflow import MlflowClient

# MLflowモデルの最新バージョンを取得
client=MlflowClient()
versions = [mv.version for mv in client.search_model_versions(f"name='{registered_model_name}'")]
model_version = versions[0]

# 現在のノートブックのコンテキストからAPIエンドポイントとトークンを取得
API_ROOT = dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiUrl().get() 
API_TOKEN = dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiToken().get()

# MLflowエンドポイント名を設定
endpoint_name = "EZO-Common-T2-2B-gemma-2-it-sglang-endpoint"

data = {
    "name": endpoint_name,
    "config":{
        "served_entities": [
        {
            "entity_name": registered_model_name,
            "entity_version": model_version,
            "workload_type": "GPU_MEDIUM",
            "workload_size": "Small",
            "scale_to_zero_enabled": True
        }]
    },
}

headers = {"Context-Type": "text/json", "Authorization": f"Bearer {API_TOKEN}"}

response = requests.post(url=f"{API_ROOT}/api/2.0/serving-endpoints", json=data, headers=headers)

print(json.dumps(response.json(), indent=4))
```

処理が正しく実行されれば、エンドポイントの作成が開始されます。
私の環境ではだいたい20分程度でデプロイが完了しました。

# Step5. クエリの実行

無事にエンドポイントが出来たので、試しに実行してみます。
今回はMLflow Deployments APIを使ってエンドポイントにクエリを渡しました。

```python
from IPython.display import display, Markdown
import mlflow.deployments

client = mlflow.deployments.get_deploy_client("databricks")
endpoint = "EZO-Common-T2-2B-gemma-2-it-sglang-endpoint"

chat_response = client.predict(
    endpoint=endpoint,
    inputs={
        "messages": [
            {
              "role": "system",
              "content": "あなたは優秀なAIアシスタントです。ハルシネーションを起こさずに回答してください。"
            },
            {
              "role": "user",
              "content": "Databricksとは何ですか？"
            }
        ],
        "temperature": 0.0,
        "max_tokens": 500,
        "top_p":0.6,
        "top_k":5,        
    }
)

content = chat_response["choices"][0]["message"]["content"]
display(Markdown(content))

usage = chat_response["usage"]
print(f"Usage: {usage}")
```

結果は以下のようになりました。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/d394e2a0-5f0c-fbe7-8c0b-9a5f1d396ce0.png)

このセルの実行にかかった時間は5.29秒でした。
API呼び出しのオーバーヘッド等はあると思いますが、単純に計算すると約58 tokens/secですね。
※　このテストだと全体のスループットの大きさはわからないので、今度測ってみようと思います。

また、少し小細工してStructured Outputもできるようにしたので、そこも試しておきます。

messageの中に```role=regex```を含めると、指定したJSONスキーマで出力するようになります。

```python
from IPython.display import display, Markdown
import mlflow.deployments

client = mlflow.deployments.get_deploy_client("databricks")

chat_response = client.predict(
    endpoint=endpoint,
    inputs={
        "messages": [
            {
              "role": "system",
              "content": "あなたは優秀なAIアシスタントです。ハルシネーションを起こさずに回答してください。"
            },
            {
              "role": "regex",
              "content": r"""\{ "weather": "(Sunny|Rain|Cloudy)" \}"""
            },
            {
              "role": "user",
              "content": "東京の明日の天気は？"
            },
        ],
        "temperature": 0.0,
        "max_tokens": 500,
        "top_p":0.6,
        "top_k":5,        
    }
)

content = chat_response["choices"][0]["message"]["content"]
display(Markdown(content))

usage = chat_response["usage"]
print(f"Usage: {usage}")
```

結果は以下の通り。
想定した形のJSONで出力されていますね。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/0cf3df08-6828-dc7e-a0ab-83c3968eecea.png)

# まとめ

SGLangを推論エンジンに使ってDatabricksモデルサービングのエンドポイントを立ててみました。
SGLang自体、バックエンドでOpenAI互換のAPIサーバを立てるので若干二度手間感がありますが、Databricksモデルサービングを使う方が管理上便利なことも多いのでこれもアリかと思っています。

十分なスループットが出そうですし、[プロビジョニングされたスループットの基盤モデルAPI](https://docs.databricks.com/ja/machine-learning/foundation-models/deploy-prov-throughput-foundation-model-apis.html)を使わない場合のLLM APIサーバはこの形で割といいのではないかと思っています。

とはいえ、推論エンジンは次々に新しいものも出てくるので、ユースケースにあったものなどをちゃんと考える必要はありそうですが。