MLflow Deployments APIを用いたDatabricksへの外部LLMのデプロイメント

以前のMLflow AI Gatewayは非推奨になってMLflow Deployments APIに置き換えられたんですね。

https://mlflow.org/docs/latest/llms/gateway/index.html

# MLflow AI Gatewayのざっくりとした説明

大規模言語モデル(LLM)を含む様々なモデルをサービングする際、資格情報やレート制限などの設定とともに管理したいと言うニーズがあります。これを容易にするための仕組みが[MLflow AI Gateway](https://qiita.com/taka_yayoi/items/789120e35a2f213eec94)でした。

# MLflow Deployments APIのざっくりとした説明

モデルをサービングエンドポイントに配置することをデプロイメント(deployment)と呼びます。このデプロイメントを簡単に行えるようにしてくれるのが[MLflow Deployments API](https://mlflow.org/docs/latest/python_api/mlflow.deployments.html)です。上記MLflow AI Gatewayの機能も取り込まれています。

# MLflow Deployments APIを用いた外部LLMのデプロイメント

以下のチュートリアルをウォークスルーします。

https://docs.databricks.com/ja/generative-ai/external-models/external-models-tutorial.html

## 外部モデルをサポートするMLflowをインストールする

```py
%pip install mlflow[genai]==2.9.0
dbutils.library.restartPython()
```

## 外部モデルエンドポイントを作成して管理する

以下を実行する前に[シークレットにOpenAI APIキーを格納](https://qiita.com/taka_yayoi/items/d52a90d8b694d730d139#%E3%82%B7%E3%83%BC%E3%82%AF%E3%83%AC%E3%83%83%E3%83%88%E3%81%AE%E8%A8%AD%E5%AE%9A)しておきます。シークレットのスコープは`demo-token-takaaki.yayoi`シークレットのキー名は`openai_api_key`です。この情報を以下でエンドポイントを作成する際に設定として埋め込んでいます。

```py
import mlflow.deployments

client = mlflow.deployments.get_deploy_client("databricks")
client.create_endpoint(
    name="taka-openai-chat-endpoint",
    config={
        "served_entities": [{
            "external_model": {
                "name": "gpt-3.5-turbo",
                "provider": "openai",
                "task": "llm/v1/chat",
                "openai_config": {
                    "openai_api_key": "{{secrets/demo-token-takaaki.yayoi/openai_api_key}}"
                }
            }
        }]
    }
)
```

実行すると**サービング**の一覧にエンドポイントが追加されます。
![Screenshot 2024-01-11 at 18.14.21.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/0c5111e3-000f-2a13-adbc-280b4d7e2261.png)
![Screenshot 2024-01-11 at 18.14.55.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/0b7e40fe-d517-8f78-b343-ced1447a78fc.png)

これまでのモデルサービングではトレーニングした機械学習モデルしかデプロイできませんでしたが、外部のモデルもあたかも通常のモデルのようにサービングエンドポイントにデプロイできるようになったと言うことですね。

右上の**エンドポイントにクエリー**で動作確認することができます。
![Screenshot 2024-01-11 at 18.17.00.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/61e3b4fa-11dc-97d9-efb0-c28a0e5901dd.png)

## 外部モデルエンドポイントに要求を送信する

チュートリアルではexperimentalの`predict`メソッドを使っていますが、私の環境では動かなかったので従来の手法でアクセスします。

```py
import os

# 事前にCLIでシークレットにトークンを登録しておきます
token = dbutils.secrets.get("demo-token-takaaki.yayoi", "token")

os.environ["DATABRICKS_TOKEN"] = token
```

```py
import os
import requests
import numpy as np
import pandas as pd
  
# モデルを呼び出してスコアリングを実行
def call_llm(question):
  
  # エンドポイントのURL
  url = '<エンドポイントのURL>'
  headers = {'Authorization': f'Bearer {os.environ.get("DATABRICKS_TOKEN")}'}

  data = {
  "messages": [
    {
      "role": "user",
      "content": question
    }
  ],
  "max_tokens": 128
} 
  #print(data)
  
  response = requests.request(method='POST', headers=headers, url=url, json=data)
  if response.status_code != 200:
    raise Exception(f'Request failed with status {response.status_code}, {response.text}')
  return response.json()  
```

```py
response = call_llm("フランスの首都は？")
response
```
```
{'id': 'chatcmpl-8flJ3vDfiOPyfwNk3iFTPmBgjEtAM',
 'object': 'chat.completion',
 'created': 1704963985,
 'model': 'gpt-3.5-turbo-0613',
 'choices': [{'index': 0,
   'message': {'role': 'assistant', 'content': 'フランスの首都はパリです。'},
   'logprobs': None,
   'finish_reason': 'stop'}],
 'usage': {'prompt_tokens': 15, 'completion_tokens': 11, 'total_tokens': 26},
 'system_fingerprint': None}
 ```

これでDatabricksから外部モデルにアクセスする際には、このエンドポイント経由にすることができます。これによって、資格情報やレート制限の設定を集中管理できるということで、もともとAI Gatewayでやりたかったことを踏襲していると言うわけです。

なお、レート制限はエンドポイントの画面から設定できます。
![Screenshot 2024-01-11 at 18.20.41.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/3fbf7e69-9d44-a4be-e106-bf1ce2a599d9.png)
![Screenshot 2024-01-11 at 18.20.31.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/88a15687-e636-fc6c-fa5d-91014217b29a.png)

以下のマニュアルにあるように、OpenAI以外にもanthropic、cohere、aws-bedrockなどもサポートしています。

https://docs.databricks.com/ja/generative-ai/external-models/index.html


### Databricksクイックスタートガイド

[Databricksクイックスタートガイド](https://www.amazon.co.jp/dp/B09V1YXFVQ/)


### Databricks無料トライアル

[Databricks無料トライアル](https://databricks.com/jp/try-databricks)