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

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

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

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

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

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

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

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

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

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

以下を実行する前にシークレットにOpenAI APIキーを格納しておきます。シークレットのスコープはdemo-token-takaaki.yayoiシークレットのキー名はopenai_api_keyです。この情報を以下でエンドポイントを作成する際に設定として埋め込んでいます。

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}}"
                }
            }
        }]
    }
)

実行するとサービングの一覧にエンドポイントが追加されます。

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

右上のエンドポイントにクエリーで動作確認することができます。

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

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

import os

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

os.environ["DATABRICKS_TOKEN"] = token

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()

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でやりたかったことを踏襲していると言うわけです。

なお、レート制限はエンドポイントの画面から設定できます。

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

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