MLflowのOpenAIフレーバーにおけるAPIキー管理方法の変更

こちらを読んでいたら以下のような記述が。

Specifying secrets for model serving with MLFLOW_OPENAI_SECRET_SCOPE is deprecated. Use secrets-based environment variables instead.

これまでは、OpenAIのAPIキーはDatabricksシークレットのスコープ名を環境変数MLFLOW_OPENAI_SECRET_SCOPEに指定する形で引き渡していましたが、これが非推奨になるとのこと。代わりにシークレットベースの環境変数を指定しろとのこと。なので、早速試してみます。

%pip install tiktoken
%pip install openai

import os
import mlflow
import openai

model_name = "takaakiyayoi_catalog.rag_chatbot.open_ai"

mlflow.set_registry_uri("databricks-uc")

with mlflow.start_run():
    model_info = mlflow.openai.log_model(
        model="gpt-3.5-turbo",
        task=openai.ChatCompletion,
        messages=[{"role": "system", "content": "You are an MLflow expert"}],
        artifact_path="model",
        registered_model_name = model_name,
    )

MLflowでトラッキングするとともに、Unity Catalogに登録します。

REST APIを使ってモデルサービングエンドポイントを作成しますが、この際にenvironment_varsにシークレットを参照する環境変数を設定します。以下の例ではスコープ名がdemo-token-takaaki.yayoiでキー名がopenaiのシークレットにOpenAIのAPIキーを格納しています。

import requests

secret_scope_name = "demo-token-takaaki.yayoi"
secret_key_name = "openai"
databricks_host = "<Databricksワークスペースのホスト名>"

endpoint_name = "my_serving_endpoint_taka"
model_name = model_name
model_version = 1
databricks_api_token = dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiToken().get()

# Create your endpoint
data = {
    "name": endpoint_name,
    "config": {
        "served_models": [
            {
                "model_name": model_name,
                "model_version": model_version,
                "workload_size": "Small",
                "scale_to_zero_enabled": True,
                "environment_vars": {
                    "OPENAI_API_KEY": f"{{{{secrets/{secret_scope_name}/{secret_key_name}}}}}"
                }
            }
        ]
    }
}

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

response = requests.post(
    url=f"https://{databricks_host}/api/2.0/serving-endpoints",
    json=data,
    headers=headers
)
print("Response status:", response.status_code)
print("Response text:", response.text, "\n")

しばらくするとエンドポイントが稼働し始めます。

動作確認のために(いつものように)streamlitで画面を作ります。

/Workspace/Users/[e-mail]/20231219_openai/streamlit.py

import streamlit as st 
import numpy as np 
from PIL import Image
import base64
import io

import os
import requests
import numpy as np
import pandas as pd

import json

from databricks.sdk.runtime import dbutils

st.header('OpenAI Chatbot on Databrikcs')

def score_model(prompt):
  # 1. パーソナルアクセストークンを設定してください
  token = dbutils.secrets.get("demo-token-takaaki.yayoi", "token")

  # 2. モデルエンドポイントのURLを設定してください
  url = '<モデルサービングエンドポイントのURL>'
  headers = {'Authorization': f'Bearer {token}'}
  
  data_json_str = f"""
  {{"inputs":[
{{"role": ["user"], "content": ["{prompt}"]}}
]
}}
"""

  data_json = json.loads(data_json_str)
  response = requests.request(method='POST', headers=headers, url=url, json=data_json)

  if response.status_code != 200:
    raise Exception(f'Request failed with status {response.status_code}, {response.text}')
  return response.json()

prompt = st.text_input("プロンプト")

if prompt != "":
  response = score_model(prompt)
  st.write(response['predictions'][0])

元のノートブックに戻ってstreamlitをインストールします。

%pip install streamlit watchdog
dbutils.library.restartPython()

from dbruntime.databricks_repl_context import get_context

def front_url(port):
    """
    フロントエンドを実行するための URL を返す

    Returns
    -------
    proxy_url : str
        フロントエンドのURL
    """
    ctx = get_context()
    proxy_url = f"https://{ctx.browserHostName}/driver-proxy/o/{ctx.workspaceId}/{ctx.clusterId}/{port}/"

    return proxy_url

PORT = 1501

# Driver ProxyのURLを表示
print(front_url(PORT))

# 利便性のためにリンクをHTML表示
displayHTML(f"<a href='{front_url(PORT)}' target='_blank' rel='noopener noreferrer'>別ウインドウで開く</a>")

streamlitを起動します。

streamlit_file = "/Workspace/Users/[e-mail]/20231219_openai/streamlit.py"
!streamlit run {streamlit_file} --server.port {PORT}

上のセルを実行した際に表示されるリンクをクリックするとGUIにアクセスできます。

ということで動作確認もできました！

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

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

Databricks無料トライアル

Databricks無料トライアル