AWS Lambda Pythonでsqlite-vssによるベクトル検索を利用する

昨今LLMの台頭により、テキストをベクトル化して類似文書の検索に利用する手法が流行っています。今回はAWSでこの検索を実現するための一方法として、SQLiteのプラグインであるsqlite-vssをAWS Lambda上で使う方法をまとめます。

github.com

意外とハマりどころや特有の考慮事項が多いので、必見です！

アーキテクチャ

LambdaでSQLite？と思った方のため、このアーキテクチャの要点をまとめます。

このアーキテクチャのメリットは、完全なサーバーレスでベクトル検索を実行できる点です。OpenSearchやPostgres (pgvector)、Redisなどのインスタンスを管理する必要はありません。サーバーレスの利点はもはや言うまでもないでしょう。

また、SQLiteを使うため、ベクトルだけでなく他のリレーショナルなデータをあわせて格納できる点も便利でしょう。例えば検索対象のドキュメント群に関するメタデータとembeddingベクトルをあわせて保存するような使い方が想定できます。

一方SQLiteを使う明らかなデメリットは、分散環境下で書き込みを同期することが難しい点です。それぞれのLambdaはローカルのファイルシステムにある別々のデータベースファイルを参照しているため、書き込みは同期されませんし、そもそもLambdaのライフサイクルに応じて変更はリセットされます。これはS3やLitestreamといったソリューションによりなんとかできるかもしれませんが、いずれにせよ大変ではあるでしょう。では、どうすれば楽できるでしょうか？

CQRS的な考え方を使いましょう。データベースから検索するだけのアプリにとっては、そのデータベースは読み込み専用です。上記のアーキテクチャにおいては、LambdaにとってSQLiteデータベースは読み込みのみ・書き込まないという役割の分離ができます。これにより、Lambda側ではデータベースへの書き込み同期という問題を無視できます。

では書き込みはどうするのでしょうか？ドキュメント検索などの用途でベクトルデータベースを用いる場合、データベースは大量のドキュメントをクロールして作成される事が多いでしょう。このクローラーがデータベースを更新する役割です。

データベースが更新されたとき、どのようにLambdaに反映すればよいでしょうか？今回はLambdaのコードパッケージにSQLiteデータベースを同梱する想定なので、DBの更新はLambdaを再デプロイすることで行います。この方法は更新の頻度が高い場合はそれだけコールドスタートが頻発するなど、不都合があるかもしれません。しかし実用上は、この更新処理はインクリメンタルに実行するというよりは、定期的なバッチ処理で十分なことも多いはずです。例えば1日1回のデプロイで更新する程度であれば、大きな問題はないでしょう。

あるいはSQLiteデータベースをS3などに配置し、Lambdaから定期的にプルするという方法も考えられます。しかしこれは明らかに考慮事項が増えるので、一旦は上記の方法がシンプルで良いでしょう。

ということで、SQLiteによるベクトル検索 on Lambdaが実用できそうなことがわかりました。次は実際にデプロイする方法を見ていきます。

Lambdaでsqlite-vssを利用するためのコード

Lambdaでsqlite-vssを利用するためには、以下のポイントが重要です。

sqlite-vssをPythonから呼び出す

sqlite-vssは各言語から簡単に利用できるよう、ライブラリが用意されています。以下はPython向けパッケージをインストールするためのPoetryコマンドです:

poetry add sqlite-vss

このパッケージはSQLiteのエクステンションとしてロードするためのビルド済みバイナリが含まれます。このため一部のプラットフォームでは利用できず、現状はMac (x86_64/arm両方)とLinux (x86_64のみ)サポートされていることに注意してください。また、Ventura未満のバージョンでArm Macを使う場合は、2023/07現在 -allow-prereleases フラグが必要です (参考)。

パッケージをインストールできたら、あとはこちらのドキュメントに従えばPythonからsqlite-vssを利用できます: sqlite-vss with Python | sqlite-vss

Tips: MacでSQLiteのエクステンションをロードできるようにする

Macで使う場合、通常の手順でインストールしたPythonではSQLiteのエクステンションをロードできないことがあります。pyenvでPythonをインストールする際にフラグを指定することで、この問題を解消できます (参考):

brew install sqlite
export LDFLAGS="-L/opt/homebrew/opt/sqlite/lib"
export CPPFLAGS="-I/opt/homebrew/opt/sqlite/include"

# Pythonバージョンは任意で設定する (以下は3.10.9の例)
CONFIGURE_OPTS=--enable-loadable-sqlite-extensions pyenv install 3.10.9

Lambda向けのDockerfile

sqlite-vssを利用する場合、コンテナLambdaを利用するのがお勧めです。標準のPythonランタイムでは、依存するライブラリ (glibc, blasなど) が実行環境に存在しないため、sqlite-vssを利用するのが困難です。同様の理由で、Lambda向けのDockerイメージ amazon/aws-lambda-python もglibcが古いために利用が困難です (参考)。

踏まえると、Dockerfileは以下のようになります:

FROM --platform=linux/amd64 python:3.10

# lapack is required for sqlite-vss
RUN apt-get update && apt-get install -y liblapack-dev \ 
    && apt-get clean && rm -rf /var/lib/apt/lists/*
RUN pip install "poetry==1.5.1"

# Use Lambda web adapter (optional)
COPY --from=public.ecr.aws/awsguru/aws-lambda-adapter:0.7.0 /lambda-adapter /opt/extensions/lambda-adapter
ENV PORT=8000

WORKDIR /app
COPY poetry.lock pyproject.toml ./

RUN poetry config virtualenvs.create false \
  && poetry install --only main --no-interaction --no-ansi

// ここでSQLiteデータベースのファイルもコピーする
COPY . ./

CMD ["python", "backend/api.py"]

基本的には普通のPython (w/ poetry)のDockerfileです。何点かポイントがあります:

ベースイメージはDocker公式のpythonを利用しています。新しめのイメージなら何でも良いと思います。
liblapack-devをインストールしていることに注意してください。これはsqlite-vssの依存関係の1つですが、今回のベースイメージには不足しているため、追加でインストールします。ちなみにyum系なら yum install -y lapack で同様にインストール可能です。
コンテナイメージ内にSQLiteデータベースのファイルを含めます。冒頭に説明した通り、データベースを更新する際はイメージをビルドしてLambdaを再デプロイします。
Lambda Web Adapterを利用していますが、これは必須ではありません。FastAPIをサクッとコンテナLambdaで動かすには非常に便利なので、ぜひお試しください。

aws.amazon.com

Tips: SQLiteのバージョンを上げる

sqlite-vssではSQLiteのバージョンにより若干使い方が変わります。具体的には、SQLiteが3.41.0未満の場合は、クエリする際のパラメータ指定方法がやや冗長になります (参考)。できればより便利な記法を使える3.41.0以上のSQLiteを使いたいものです。

Pythonで import sqlite3 した際のSQLiteのバージョンはDockerのベースイメージにより決まるため、python:3.10 などのイメージを使う場合は不可抗力的に3.41.0未満になってしまうことがあります。バージョンを変えるためには、 pysqlite3-binary パッケージを導入するのが楽です。これはパッケージ内にSQLiteバイナリを同梱することで、任意のバージョンのSQLiteを使えるようにしたパッケージです。pysqlite3-binaryの0.5.1から、SQLiteが3.41.0以上(3.42.0)になっています。

pysqlite3-binary は現状Linuxでのみ利用できるため、Poetryもplatform指定でaddしましょう:

poetry add pysqlite3-binary --platform linux

pysqlite3-binaryはsqlite3と同じAPIを持つため、以下のようなコードを書けば、Macではsqlite3・Linuxではpysqlite3-binaryを利用するようにできます:

import sqlite_vss
try:
    import pysqlite3 as sqlite3
except ModuleNotFoundError:
    import sqlite3

db = sqlite3.connect("test.db", check_same_thread=False)
db.enable_load_extension(True)
sqlite_vss.load(db)
db.enable_load_extension(False)

ちなみに check_same_thread=False を指定しているのは、FastAPIで簡単にSQLiteを使うための措置です。こうしないとFastAPIのスレッド間で db を共有できず不便でした。Lambdaではリクエストは同時に1件しか処理されないため、特に問題ないでしょう (そもそも読み込みしかしないので、データレースは起きない)。

これにより、任意のDockerイメージで最新のSQLiteを使えるようになりました。

コンテナLambdaのデプロイ

あとは通常の手順でAWSにデプロイするだけです。例えばAWS CDKを使う場合、コードは以下の通りです。Amazon API Gateway HTTP APIから叩けるようにしています。

import * as cdk from 'aws-cdk-lib';
import { Construct } from 'constructs';
import { DockerImageCode, DockerImageFunction } from 'aws-cdk-lib/aws-lambda';
import { HttpApi } from '@aws-cdk/aws-apigatewayv2-alpha';
import { HttpLambdaIntegration } from '@aws-cdk/aws-apigatewayv2-integrations-alpha';
import { Platform } from 'aws-cdk-lib/aws-ecr-assets';

export class SqliteVssOnLambdaStack extends cdk.Stack {
  constructor(scope: Construct, id: string, props?: cdk.StackProps) {
    super(scope, id, props);

    const handler = new DockerImageFunction(this, 'Handler', {
      code: DockerImageCode.fromImageAsset('backend', { platform: Platform.LINUX_AMD64 }),
    });

    const api = new HttpApi(this, 'Api');
    const integration = new HttpLambdaIntegration('Integration', handler);
    api.addRoutes({
      path: '/{proxy+}',
      integration,
    });
    new cdk.CfnOutput(this, 'ApiUrl', { value: api.url! });
  }
}

実際のアプリケーションのコードはほぼ含めていませんが、公式ドキュメントを見ればわかるので省いています。

気付き・注意事項

上記のアーキテクチャを構築・運用する上でいくつかの考慮事項があるので、以下に気づいた範囲でまとめます:

Faissについて

sqlite-vssにおいてはベクトル検索のためにFaissというライブラリが使われています。Faissはsqlite-vssを利用する上で完全に抽象化されているわけではないので、理解しておく方が良さそうです。概要はこちらを見るのが分かりやすいです。

Introduction to Facebook AI Similarity Search (Faiss) | Pinecone

ベクトル検索は類似度が高いベクトル同士を見つけることが目的です。単純な方法は全てのベクトルと類似度を計算すれば良いでしょう。この方法は厳密な解が得られ正確であるものの、計算量がベクトルの個数に比例するため、大規模なデータでは使いづらいです。厳密解を近似することで、解の精度が下がることをしつつ、計算量を下げることができる近似アルゴリズム (Approximate Nearest Neighbor, ANN) が存在します。

sqlite-vssでのデフォルト設定では全探査のアルゴリズムのため、厳密解が得られるものの計算量は非効率です (とはいえ10万とか小規模なデータセットなら十分実用的なようです (後述))。 sqlite-vssが利用しているFaissでは、いくつかの代表的なANNアルゴリズムをサポートしています。アルゴリズムの選定法はFaiss公式Wikiが詳しいです。

Guidelines to choose an index · facebookresearch/faiss Wiki · GitHub

一部の近似アルゴリズムを使う場合は、学習の手順が追加で必要になります。sqlite-vssでは、SQLiteのインターフェースから学習処理を実行できます: operation='training' 利用するデータセットのベクトルの一部をランダムにサンプルし、教師データとして与えるようです。

このスライドも、ベクトル数とアルゴリズムの対応に関して参考になります。少し上で見積もった数字と、概ね一致しています。

また、デフォルトではベクトルをインデックスに入れる際圧縮しないため、インデックスのサイズが大きくなるようです。これはProduct Quantizationという手法で節約できます。これも検索精度とRAMサイズのトレードオフになります。

とはいえ次で考える通り、実用上Lambdaではあまり大規模なデータセットを扱えなさそうなので、検索アルゴリズムはデフォルトのままで良いかもしれません。

ベクトル数の限界について

このアーキテクチャだとどの程度の件数のベクトルまでデータベースに保持できるでしょうか？これを考えるには、1. 検索処理時間 と、2. RAM消費量 を考える必要があります。

検索処理時間は、デフォルトのアルゴリズムだと計算能力に反比例、ベクトル数に比例 (後述)します。1000件程度のベクトル数 (次元数1536) で試したところ、RAM=128MBのLambdaだと約5msで検索できました。仮に50ms以内にレスポンスを得るとした場合、RAM=10GB (最大) のLambdaでは、

10240/128 * 50/5 * 1000 = 0.8e6

ということで80万件程度であれば処理できるようです。Lambdaでは計算能力がRAMサイズに比例します。ちなみにRAM=1GB、500msにしても同じく80万件です。

次にRAM消費量は、デフォルトのアルゴリズムでは1ベクトル当たり 4x(次元数) Bytes 消費します (参考)。次元数=1536、RAM=10GBとすれば、

10240*1e6 / (4*1536) = 1.67e6

ということで167万件程度が限界になります。

上記は非常にざっくりとした見積もりですが、先程のスライドとも概ね一致している数字です。いずれにせよ数千万以上の規模は難しいですし、1~10万程度であれば十分実用できそうなことが分かります。

コスト効率について

上記の通り、ベクトル件数が多い場合はLambdaに高いRAMサイズ (1~10GB) を割り当てる必要が生じるでしょう。

Lambdaに高いRAMサイズを割り当てる場合、気になるのはI/O待機時間の課金です。例えば外部のLLMに文章のembeddingベクトルを計算させているときは、Lambdaは何も仕事ができません。Lambdaは同じインスタンスで1リクエストしか処理できないため、I/O多重化によるコスト節約が難しいのです。

これを回避するには、ベクトル検索用のLambdaは別のAPIとして切り出し、I/Oが発生しないようにするのがベターです。こうすれば、RAMサイズの大きなLambdaの利用率を高め、コスト効率を高めることができるでしょう。

Lambdaのコールドスタートについて

より深刻かもしれない問題もあります。sqlite-vssは初回クエリ時にFaissのインデックスを初期化・メモリ上に展開するため、Lambdaのコールドスタート後1発目のクエリ応答時間が長くなります。

例えば上記で試した1000件のベクトル・128MBのRAMの場合、1発目のクエリのみ2秒程度 (通常は約5ms) かかりました。この初期化時間もベクトル数に応じて大きくなるはずなので、より大きなデータセットでは深刻になりえます。

これはLambdaのSnapStartが本命の解決策 (はよ…！) ですが、一旦はタイムアウトしたらクライアント側でリトライするしかないかもしれません。

ベクトル検索を自分で管理することについて

正直なところ、ベクトル検索を自前で管理するのは、必要な知識や経験が多く難しそうです。特にアルゴリズムの部分はパラメータも多く、この分野に詳しい専門家が不在の状況では最適化は困難に思えます。

これらを考慮すべきなのはOpenSearchなどを使っても同様のようです。(以下の記事は各アルゴリズムの紹介に加えてベンチマークもあり、非常に分かりやすいのでおすすめです。) 最適化の余地があるのは好ましいことでもあり、いい感じにやってくれれば良いんだけどなという思いも捨てきれません。