🧠

Snowflake の E2E ML を試してみる

2022/12/14に公開

Python

📜 この記事は

Snowflake Advent Calendar 2022 Calendar 2 Day 14 の記事になります。Snowflake で Sales Engineer をしております、@tmasuo がお届けします。
この記事で触れる Snowflake の各機能のリリース状態は 2022 年 12 月 1 日時点に基づきます。

🐍 Snowpark for Python - GA

11 月に行われた Snowday 2022 にて Snowpark for Python の GA がついに発表されました🎉
昨年の 11 月に Private Preview、今年の 6 月（Snowflake Summit 2022）に Public Preview が開始され、待望の GA となりました。

思わず口角の上がるクマ太郎
個人的には、私が Snowflake に入社したのが昨年 10 月末で、当時 Snowpark は Scala/Java しか対応していなかったため、「この会社は Scala かけないと生き残れないのか、、Python じゃダメなのか、、」と思っていたためかなり感慨深かったりします。

この GA にて対応された Snowpark for Python は以下のコンポーネントから成り立っています。

クライアント API

Jupyter Notebook などから Snowflake のデータを DataFrame 形式で操作する際の API となります。DataFrame の処理は特定のアクションが実行された際に SQL 変換されて Snowflake 上で実行されます（遅延評価）。そのためデータ処理を Snowflake のウェアハウスで実行できる点がメリットとなっています。
Pandas DataFrame みたいなもんかな？と思った方、イメージあってます。ただ、クライアント側のメモリに展開して処理する方式ではなく遅延評価されるので差分があります。Pandas DataFrame との比較は @masato_takada がまとめてくれた記事が参考になります。

ユーザ定義関数（UDF）

Snowflake 上で実行される Python UDF になります。スカラー値を返す関数と、テーブルを返す関数どちらにも対応しています。
前述のクライアント API は SQL 変換されて Snowflake で実行されますが、UDF は Snowflake 上の実行環境で Python コードが実行されます。Python コードが実行できる利点としては豊富な Python エコシステムのライブラリを活用できる点にあると思います。Snowflake で実行できる Python パッケージは内蔵された Anaconda が管理しており、パッケージのリストはこちらから確認できます。

ストアドプロシージャ

Snowflake 上で実行される Python ストアドプロシージャです。UDF 同様 Python コードが Snowflake 上で実行されます。

❄️ 使いどこ

主に 3 つあると思います。

データエンジニアリング: データ変換やカスタムロジックの実行
データサイエンス & 機械学習: 特徴量エンジニアリングやモデルトレーニングおよび推論
アプリケーション開発: Python で記述されたアプリでのデータ処理

それぞれとても重要な要素がありますが、こちらの記事ではタイトルの通り、 データサイエンス & 機械学習にフォーカスしてゆきます。

データサイエンス & 機械学習

Snowpark for Python により、Snowflake での E2E 機械学習が実現できるようになりました。以下のダイアグラムのイメージです。

クライアント API により特徴量エンジニアリングが実行され、ストアドプロシージャによりモデルトレーニングを実施し、UDF で推論を行っています。
ここでストアドプロシージャが実行されるウェアハウスが Snowpark 用に最適化されたウェアハウスになります。2022 年 6 月（Snowflake Summit）時点ではハイメモリインスタンスなどと呼ばれていたものです。Public Preview 中でリージョン制約がありますが、AWS 東京や Azure 東日本リージョンでもお試しいただけます。

🧠 ML meets Snowflake を動かしてみよう

ということで、この一連の流れを Snowflake 上で動かしてみたいと思います。
コンテンツとしては以下の Quickstart を動かしてみます。

Step by Step で解説すると記事が間延びしてしまうので、こちらではいくつかポイントを絞って解説したいと思います。

1. DataFrame 比較

こちらの Notebook で実行している Pandas DataFrame と Snowpark DataFrame の比較です。
利用しているメモリ空間について、Pandas DataFrame は read_csv した結果、Snowpark DataFrame は session.table した結果で比較していますが、後者は全く利用されていないです。つまり遅延評価の評価タイミングが来ていないためクライアントサイドのメモリにデータが展開されていない状態を示しています。

また、各種ライブラリとの互換性を考えると、Snowpark DataFrame ではなく、Pandas DataFrame で取り回したいケースもあるかと思います。その場合、toPandas() というメソッドが利用できます。このメソッドを実行すると、クライアント側のメモリに DataFrame の内容が展開されるため、対象を絞ったり集計処理をかけてから toPandas() することが徹底されている印象があります。

limit してから toPandas

集計してから toPandas

このように DataFrame 操作を Snowflake にプッシュダウンすることで、Snowflake の柔軟なエンジンの活用とクライアント側のリソースの有効活用を同時に実現しています。

2. 推論のパフォーマンス最適化

推論は UDF で実行していますが、通常の実行と最適化された実行でパフォーマンスに違いがあります。ここでは、

ベクトル化された UDF
cachetools の活用

の 2 つを活用しています。
これらの動作を確認するため、UDF を呼び出している CTAS 内の SELECT 文を抜き出してパフォーマンスを計測してみました。3 回の計測値と平均値（単位は秒、小数点 2 位四捨五入）は以下の通りでした。

パターン	1 回目計測値	2 回目計測値	3 回目計測値	平均計測値
通常の実行	9.9	7.6	7.6	8.4
最適化 1 適応	2.4	2.3	2.3	2.3
最適化 2 適応	5.8	5.5	5.6	5.6
最適化 1+2 適応	2.3	2.3	2.2	2.3

それぞれの最適化について動作を確認してみます。

ベクトル化された UDF

Snowflake では UDF の実行をある程度まとめて実行するベクトル化された UDF をサポートしています。ここで利用する API がバッチ API と呼ばれ、パフォーマンス向上に寄与しています。具体的にクエリプロファイルを確認すると以下の通りです。

通常の実行: 行ごとに UDF が実行されている（処理行数 == UDF 呼び出し回数）
ベクトル化された実行: まとまった単位で UDF が実行（処理行数 != UDF 呼び出し回数）

常に処理が早くなるとは限りませんが、計測結果からも分かる通り今回は処理の効率化に大きく貢献しています。

cachetools

cachetools は Python でキャッシュを取り扱うライブラリです。UDF での推論は、ステージからモデルの読み込みを行うため、IO の観点で効率化を行うことができます。こちらもクエリプロファイルを確認すると以下の通りです。

通常の実行: ハンドラごとの実行時間 4.07ms
キャッシュを利用した実行: ハンドラごとの実行時間 2.72ms

IO の最適化は重要ですね。

✅ さいごに

Snowpark for Python の GA により、Snowflake の柔軟なエンジンを活用した E2E ML が実現できるようになりました。
また、ここでは触れてませんが Tag-based Masking（これも最近 GA しました！）などガバナンスとデータ活用のバランスを取ることができる点も、Snowflake で ML をする大きなメリットだとおもいます。ぜひお試しいただけると幸いです！