Posted inAI

InfinityとDockerで構築する高パフォーマンスなEmbedding・Rerankingサーバーのセルフホスト：高速RAGのためのOpenAI互換API

AI tutorial - IT technology blog
AI tutorial - IT technology blog

本番環境のRAGパイプラインにおけるレイテンシのボトルネック

検索レイテンシは、検索拡張生成（RAG）システムにおいてユーザー体験を損なう最大の要因です。一般的なクエリのライフサイクルを考えてみましょう。バックエンドがリモートのEmbedding APIにテキストを送信し、ベクトルデータベースを検索し、上位のチャンクをクラウドのRerankerに転送して、最後にLLMへプロンプトを渡します。各ステップでインターネットの往復通信が発生するため、LLMが最初のトークンを生成し始めるまでに、ユーザーは1.5〜3秒ほど待たされることになります。

サードパーティのAPIは、ボトルネックをさらに悪化させます。OpenAIやCohereなどのエンドポイントを介してEmbeddingやRerankerを実行すると、ネットワークのジッター（揺らぎ）、厳格なレート制限、トークンごとの従量課金が発生します。

1日あたり100,000回の検索で、クエリごとに50件のドキュメントをリランキングすると、毎月のクラウド費用は容易に1,500ドルを超えてしまいます。さらに重要なのは、パブリックインターネット経由で社内ドキュメントや顧客のクエリを送信することにより、サードパーティの処理基盤にデータが露出するリスクがある点です。推論をローカル環境で実行すれば、これらのリスクを完全に排除できます。

そこで役立つのがInfinityです。Infinityは、Transformer Encoder向けにRustとPythonで構築された専用の推論エンジンです。ネイティブの動的バッチ処理、FP16およびINT8量子化、さらにEmbeddingとCross-Encoder Reranker向けにそのまま置き換え可能なOpenAI互換REST APIを提供します。

アプローチの比較：EmbeddingおよびReranking戦略

適切な推論エンジンの選択は、トラフィックパターンとハードウェア構成に直接左右されます。

アプローチ スループットとレイテンシ インフラオーバーヘッド OpenAI API互換性 Cross-Encoder Reranking
クラウドAPI（OpenAI, Cohere） p95レイテンシ 180〜450ms（ネットワーク依存） なし（フルマネージド） ネイティブ標準 Cohere、Voyage、Jinaなどの個別エンドポイントが必要
インプロセス SentenceTransformers シングルスレッドは高速（約25ms）だが、並行アクセス時にGILが飽和 アプリケーションとCPU・メモリを共有 なし（カスタムラッパーが必要） 手動のバッチループ処理。CPU枯渇リスクが高い
vLLM / TGI 高スループットだが、16GB〜24GB以上のVRAMが必要 過重（自己回帰型のDecoderモデル向け設計） 一部対応 Cross-Encoderの統合が不十分または非サポート
Infinity推論サーバー p95レイテンシ 8〜15ms。単一GPUで1,200文/秒以上 低（4 vCPUまたは手頃な8GB GPUで動作可能） /v1/embeddings および /v1/rerank に完全対応 Cross-Encoder向けのネイティブ継続的バッチ処理

セルフホストにおけるInfinityの評価

推論環境を社内へ移行することで圧倒的な高速化が図れますが、同時に運用の責任もチームが担うことになります。

メリット

  • 15ms未満の低レイテンシ：ローカルコンテナはUnixソケットやDockerのブリッジネットワークを介して通信するため、パブリックインターネットのハンドシェイクを完全に排除できます。
  • 予測可能なコンピュートコスト：ハードウェア費用は一定です。急な料金高騰やティアのアップグレードに悩まされることなく、5億トークン規模の埋め込みも処理できます。
  • データ漏洩リスクゼロ：機密データはプライベートサブネットやVPC内にとどまります。SOC 2、HIPAA、GDPRなどの厳格なデータ隔離要件も容易にクリアできます。
  • 容易なリプレイス：LangChain、LlamaIndex、あるいは素のHTTPクライアントにおいて、ベースURLとAPIキーの2行を変更するだけでエンドポイントを切り替えられます。
  • 複数モデルの同一ホスト運用：BAAI/bge-small-en-v1.5のようなDense Embedderと、BAAI/bge-reranker-baseのようなRerankerを、単一コンテナかつ4GB未満のVRAMで同時に実行できます。

トレードオフ

  • GPUのキャパシティプランニング：InfinityはONNX Runtime経由でCPU上でも動作しますが、高並行環境（50+ RPS）では最低8GBのVRAMを備えたGPUが必要です。
  • メンテナンスのオーバーヘッド：コンテナのヘルスチェック、VRAM消費量の監視、サーバー更新時の再起動対応などを管理する必要があります。
  • ディスクストレージ：EmbeddingおよびRerankingモデルのチェックポイント（重みとトークナイザー）を保存するために、2GB〜6GBの高速なNVMeストレージが必要です。

推奨アーキテクチャとハードウェア要件

秒間100〜250リクエスト程度を処理する環境の場合、以下の最小スペックを推奨します：

  • GPU環境（推奨）：8GB以上のVRAMを搭載したNVIDIA GPU（RTX 3060、RTX 4060 Ti、またはクラウドのT4/L4）×1、4 CPUコア、16GBシステムメモリ。
  • CPU構成（代替案）：INT8量子化を利用した8 vCPUコア（AVX-512拡張対応のIntel Xeon Ice LakeまたはAMD EPYC）、16GB RAM。
  • ホストOS：Docker EngineおよびDocker Compose v2が導入されたUbuntu 22.04 LTSまたは24.04 LTS。
  • ドライバ構成：NVIDIA Driver 535以上およびNVIDIA Container Toolkit。
  • 推奨モデル構成：
    • Embedding：超高速なベクトル検索には BAAI/bge-small-en-v1.5（384次元、重み約133MB）。多言語対応が必須の場合は BAAI/bge-m3（1024次元）を採用。
    • Reranker：BAAI/bge-reranker-base（重み約1.1GB）または mixedbread-ai/mxbai-rerank-base-v1。

Docker Composeを使ったデプロイ手順

ステップ1：NVIDIA Container Toolkitのセットアップ

CPUのみで実行する場合はこの手順をスキップしてください。GPUアクセラレーションを使用する場合は、NVIDIAリポジトリを登録し、Dockerのデフォルトランタイムを設定します：

# NVIDIA Container Toolkitリポジトリの追加
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \
  && curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
    sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
    sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list

sudo apt-get update
sudo apt-get install -y nvidia-container-toolkit
sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker

ステップ2：Docker Compose構成の定義

Infinity用の作業ディレクトリを作成し、新しいCompose設定ファイルを開きます：

mkdir -p ~/infinity-stack && cd ~/infinity-stack
nano docker-compose.yml

以下のサービス定義を貼り付けます。これにより、EmbeddingモデルとCross-Encoderモデルの両方が、半精度（FP16）を使用して共有エンジンに読み込まれます：

services:
  infinity:
    image: michaelf34/infinity:latest
    container_name: infinity-inference
    restart: unless-stopped
    ports:
      - "7997:7997"
    environment:
      - INFINITY_MODEL_ID=BAAI/bge-small-en-v1.5;BAAI/bge-reranker-base
      - INFINITY_BATCH_SIZE=64
      - INFINITY_API_KEY=sk-inf-local-secret-key-9923
      - INFINITY_DTYPE=float16
    volumes:
      - ./cache:/root/.cache/huggingface
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]

CPUで動作させる場合は、deploy ブロック全体を削除し、INFINITY_DTYPE を int8 または float32 に変更してください。

ステップ3：サービスの起動とヘルスチェック

バックグラウンドでコンテナを起動し、ダウンロードの進行状況をモニタリングします：

docker compose up -d
docker compose logs -f infinity

初回の実行時、InfinityはHugging Faceからモデルの重みを ./cache にダウンロードします。モデルのコンパイルとウォームアップが完了すると、ログに準備完了のメッセージが出力されます：

INFO: Started server process on http://0.0.0.0:7997
INFO: Application startup complete. Ready for requests.

ステップ4：ローカルエンドポイントへの問い合わせ

まず、標準的なOpenAIのペイロード構造を使用して、ベクトル生成をテストします：

curl http://localhost:7997/v1/embeddings \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer sk-inf-local-secret-key-9923" \
  -d '{
    "model": "BAAI/bge-small-en-v1.5",
    "input": ["セルフホストされたEmbeddingにより、ネットワークレイテンシはほぼゼロになります。"]
  }'

次に、取得した候補ドキュメントのバッチをRerankerに渡します。エンドポイントが関連性スコアを算出し、ソートされたインデックスを返します：

curl http://localhost:7997/v1/rerank \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer sk-inf-local-secret-key-9923" \
  -d '{
    "model": "BAAI/bge-reranker-base",
    "query": "ベクトル検索の応答時間を短縮するにはどうすればよいですか？",
    "documents": [
      "Infinity内部の継続的な動的バッチ処理により、GPU使用率が最適化されます。",
      "水出しコーヒーは室温で16時間抽出する必要があります。",
      "HNSWインデックスは近似最近傍探索を高速化します。"
    ],
    "top_n": 2
  }'

ステップ5：Pythonアプリケーションとの接続

InfinityはOpenAIのスキーマに完全準拠しているため、標準の openai パッケージをそのまま使用できます。カスタムのアダプターコードを書く必要はありません。

from openai import OpenAI

# クライアントの接続先をローカルのInfinityサーバーに設定
client = OpenAI(
    base_url="http://localhost:7997/v1",
    api_key="sk-inf-local-secret-key-9923"
)

def embed_documents(texts: list[str]) -> list[list[float]]:
    response = client.embeddings.create(
        model="BAAI/bge-small-en-v1.5",
        input=texts
    )
    return [record.embedding for record in response.data]

sample_chunks = [
    "セルフホスト型の推論により、p95検索時間を12msまで短縮できます。",
    "Cross-Encoderは、誤検出されたベクトルの一致を正確に除外します。"
]

vectors = embed_documents(sample_chunks)
print(f"{len(vectors)} 件のEmbeddingを生成 | ベクトルの次元数: {len(vectors[0])}")

データベースの近接環境でInfinityを実行することにより、検索時間を10〜15msまで短縮できます。サードパーティAPIへの依存をなくし、機密データを保護しながら、RAGのパフォーマンスを完全にコントロールすることが可能になります。

Share:
Tags: