MLflow AI GatewayによるLLMの一元管理

code
Published

September 23, 2023

はじめに

近年、大規模言語モデル（LLM）の活用が進んでおり、テキストデータのセンチメント分析からリアルタイムのチャットボットの導入まで、多岐にわたるアプリケーションでその価値を発揮しています。しかしながら、多くの組織でLLMを効果的に管理と利用するためには、セキュリティ、コスト、ガバナンスなどの課題が存在します。

MLflow AI Gatewayはこれらの課題に対処するためのツールとして設計され、特定のモデルプロバイダ（例：OpenAI、Anthropicなど）に対して、一元化されたインターフェースを提供します。これにより、APIトークンの管理や利用コストのコントロール、ガバナンスの強化が実現します。

本記事では、MLflow AI Gatewayのセットアップ方法と、APIドキュメントへのアクセス方法、そしてGatewayを利用したレクエスト方法について解説します。ここでは、特にOpenAIを例に挙げて説明しますが、MLflow AI Gatewayは他の多くのプロバイダにも対応しています。 詳細は公式ドキュメントを参考にしてください。

Note

MLflow 2.5以降が必要です。

MLflow AI Gatewayのセットアップ

このセクションでは、MLflow AI Gatewayのセットアップについてに解説します。 環境の準備、必要なコンフィグレーションの設定、Gatewayの起動について説明します。 詳細は、GitHubレポジトリを参考にしてください。

インストール

pip install 'mlflow[gateway]'

OpenAI API Keyの設定

pip install 'mlflow[gateway]'

設定ファイルの用意

MLflow AI Gatewayの動作には、適切な設定ファイルが必要です。このファイルでは、ルート、プロバイダー、モデルなどの重要なパラメータを定義します。

ルート

routesセクションでは、各ルートの名称、タイプ、関連するモデルなどを定義します。ルートは、外部からのリクエストがどのモデルにルーティングされるかを制御します。

プロバイダーとモデル

modelセクションでは、使用するモデルのプロバイダーと名称を指定します。プロバイダーは、モデルが存在するクラウドサービスやプラットフォームを指します。また、configセクションでは、特定のプロバイダーに必要な設定を行います。例えば、APIキーの設定などです。

routes:
- name: completions
    route_type: llm/v1/completions
    model:
    provider: openai
    name: gpt-3.5-turbo
    config:
        openai_api_key: $OPENAI_API_KEY

- name: chat
    route_type: llm/v1/chat
    model:
    provider: openai
    name: gpt-3.5-turbo
    config:
        openai_api_key: $OPENAI_API_KEY

- name: embeddings
    route_type: llm/v1/embeddings
    model:
    provider: openai
    name: text-embedding-ada-002
    config:
        openai_api_key: $OPENAI_API_KEY

APIドキュメントの確認

上記のようにMLflow AI Gatewayサーバーを起動すると、Swaggerによるlocalhost:5000に自動的にAPIドキュメントが生成されます。

Fluent APIによるリクエストの実行

MLflow AI Gatewayがセットアップされたら、Fluent APIを利用してリクエストを実行できます。他にも、Client API, REST APIも利用可能です。

Code 
from mlflow.gateway import query, set_gateway_uri

set_gateway_uri(gateway_uri="http://localhost:5000")

response = query(
    "chat",
    {"messages": [{"role": "user", "content": "What is the best day of the week?"}]},
)

import pprint
pprint.pprint(response)
{'candidates': [{'message': {'content': 'The best day of the week is '
                                        'subjective and can vary depending on '
                                        'individual preferences and '
                                        'circumstances. Some people may prefer '
                                        'weekends (Saturday and Sunday) as '
                                        'they offer a break from work or '
                                        'school and allow for relaxation or '
                                        'leisure activities. Others may enjoy '
                                        'Fridays as they mark the end of the '
                                        'workweek and the beginning of the '
                                        'weekend. Ultimately, the best day of '
                                        'the week is a personal preference.',
                             'role': 'assistant'},
                 'metadata': {'finish_reason': 'stop'}}],
 'metadata': {'input_tokens': 16,
              'model': 'gpt-3.5-turbo-0613',
              'output_tokens': 78,
              'route_type': 'llm/v1/chat',
              'total_tokens': 94}}

まとめ

本記事では、MLflow AI Gatewayのセットアップ方法とその使用方法について詳細に説明しました。MLflow AI Gatewayを使用することで、大規模言語モデルの管理が効果的に行えるため、セキュリティ、コスト、ガバナンスの各観点での課題に対処できます。また、異なるプロバイダのモデルも一元化されたインターフェースを通じてアクセスできるため、開発の効率化が図れます。

参考

実装