❄️

Snowpark Container ServiceでMLflowサーバーを立ち上げる

2024/01/13に公開

これは何?

Snowpark Container Serviceを使ってMLflowのサーバーを立ち上げたので、必要な手順や落とし穴の備忘録を残しておきます。
とりあえず立ち上げただけなので、まだ外部から機械学習モデルの情報を保存したりなどはできません。

そもそも

Snowpark Container Service(SPCS)って?

Snowflake上のリソースを使って、フルマネージドでコンテナを動かせるサービスです。
2024/1/11 現在パブリックプレビューの機能となっており、AWSの特定のリージョンとAzureで利用可能です。
詳しくは公式ドキュメントを参照ください。

公式のチュートリアルやハンズオンも充実しているため、手を動かしたい方はこちらから入るのも良さそうです。
僕は以下のチュートリアルを実行しました。
チュートリアル(Flaskで超簡易的なwebアプリを作成)
クイックスタート(jupyter notebookの立ち上げ)

以下の実践パートでも、多くのコマンドは上記のクイックスタートをかなり参考にさせてもらってます。

MLflowって?

実験、再現、デプロイ、モデルの保存を含む、機械学習のライフサイクルを管理するためのオープンソースプラットフォームです。(公式ドキュメントより)

簡単にいうと、MLflowはPythonで使用されるオープンソースのMLOpsライブラリです。

現在SnowflakeはMLOps周りの機能が弱く、pythonワークシートなどを用いて機械学習モデルを構築しても、モデルの管理は外部サービスを利用する必要がありました。こちらの弱点をコンテナで補えるのではないか、というのが本記事の動機になります。

動作環境

  • Snowflake環境
    • AWS - Asia Pacific (Mumbai)
    • (snowpark container servicesがパブリックプレビューになっている環境ならどこでも良さそう)
  • ローカル環境
    • MacBook Air (M1, 2020)

その1 Snowflake上でリソースの作成

基本的なリソースを作成

SnowflakeのSQLワークシートで以下のコマンドを実行します。

USE ROLE ACCOUNTADMIN;
CREATE ROLE CONTAINER_USER_ROLE;
GRANT CREATE DATABASE ON ACCOUNT TO ROLE CONTAINER_USER_ROLE;
GRANT CREATE WAREHOUSE ON ACCOUNT TO ROLE CONTAINER_USER_ROLE;
GRANT CREATE COMPUTE POOL ON ACCOUNT TO ROLE CONTAINER_USER_ROLE;
GRANT CREATE INTEGRATION ON ACCOUNT TO ROLE CONTAINER_USER_ROLE;
GRANT MONITOR USAGE ON ACCOUNT TO  ROLE  CONTAINER_USER_ROLE;
GRANT BIND SERVICE ENDPOINT ON ACCOUNT TO ROLE CONTAINER_USER_ROLE;
GRANT IMPORTED PRIVILEGES ON DATABASE snowflake TO ROLE CONTAINER_USER_ROLE;

grant role CONTAINER_USER_ROLE to role ACCOUNTADMIN;

// データベース、ウェアハウスを作成
USE ROLE CONTAINER_USER_ROLE;
CREATE OR REPLACE DATABASE CONTAINER_HOL_DB;

CREATE OR REPLACE WAREHOUSE CONTAINER_HOL_WH
  WAREHOUSE_SIZE = XSMALL
  AUTO_SUSPEND = 120
  AUTO_RESUME = TRUE;
  
// コンテナサービスの設定用ファイルを置く用の内部ステージを作成
CREATE STAGE IF NOT EXISTS specs
ENCRYPTION = (TYPE='SNOWFLAKE_SSE');

// コンテナがマウントするボリューム用の内部ステージを作成
CREATE STAGE IF NOT EXISTS volumes
ENCRYPTION = (TYPE='SNOWFLAKE_SSE')
DIRECTORY = (ENABLE = TRUE);

コンテナサービス特有のリソースを作成

作成するのは以下になります。

  • Security Integration
    • WebブラウザからUIベースでコンテナにログインするためのリソース
  • compute pool
    • コンテナサービスを動かすためのコンピューティングリソースの集合
    • Kubernetesベース
  • image Repository
    • コンテナサービスが使うdockerイメージを置いておくレポジトリ
Snowflake SQL worksheet
USE ROLE ACCOUNTADMIN;
CREATE SECURITY INTEGRATION IF NOT EXISTS snowservices_ingress_oauth
  TYPE=oauth
  OAUTH_CLIENT=snowservices_ingress
  ENABLED=true;

USE ROLE CONTAINER_USER_ROLE;
CREATE COMPUTE POOL IF NOT EXISTS CONTAINER_HOL_POOL
MIN_NODES = 1
MAX_NODES = 1
INSTANCE_FAMILY = standard_1;

CREATE IMAGE REPOSITORY CONTAINER_HOL_DB.PUBLIC.IMAGE_REPO;

その2 ローカルでコンテナイメージの動作確認

Snowpark Container Servicesを動かすためには、前節の最後で作ったimage repogitroyにコンテナイメージをプッシュする必要があります。 まずはローカルで動かしてみましょう。

mlflowの公式コンテナイメージをpullする

> docker pull ghcr.io/mlflow/mlflow:v2.9.2

ローカルで動作確認

まずはdocker-composeファイルを作成します。以下に、がちゃがちゃと実験を重ねながら作成したdocker-composeファイルを示します。

docker-compose.yaml
version: '3'
services:
  mlflow:
    image: ghcr.io/mlflow/mlflow:v2.9.2
    ports:
      - "8080:5000"
    command:
      mlflow server --host 0.0.0.0 --port 5000
    volumes:
      - ./mlruns:/mlflow/mlruns:rw
      - ./artifact_store:/mlartifacts:rw

docker composeコマンドを使用してコンテナを起動し、動作を確認してみましょう。

> cd {docker-compose.yamlが置いてあるディレクトリ}
> docker compose up
[+] Running 1/0
 ⠿ Container mlflow-mlflow-1  Created                                                                                                               0.0s
Attaching to mlflow-mlflow-1
mlflow-mlflow-1  | [2024-01-10 15:21:42 +0000] [15] [INFO] Starting gunicorn 21.2.0
mlflow-mlflow-1  | [2024-01-10 15:21:42 +0000] [15] [INFO] Listening at: http://0.0.0.0:5000 (15)
mlflow-mlflow-1  | [2024-01-10 15:21:42 +0000] [15] [INFO] Using worker: sync
mlflow-mlflow-1  | [2024-01-10 15:21:42 +0000] [16] [INFO] Booting worker with pid: 16
mlflow-mlflow-1  | [2024-01-10 15:21:42 +0000] [17] [INFO] Booting worker with pid: 17
mlflow-mlflow-1  | [2024-01-10 15:21:43 +0000] [18] [INFO] Booting worker with pid: 18
mlflow-mlflow-1  | [2024-01-10 15:21:43 +0000] [22] [INFO] Booting worker with pid: 22

立ち上がったようです。localhost:8080にアクセスすると、ちゃんとmlflowが立ち上がっていることがわかります。

ここでは本筋とずれるので詳細は省きますが、ちゃんとモデルも登録できます。
登録方法などは、MLflow公式のチュートリアルなどが参考になります。

その3 Snowflake上にコンテナイメージと設定ファイルをプッシュする

Snowpark Container Serviceを使うために必要なリソースを、ローカルからプッシュします。

コンテナイメージをImage Repogitoryにプッシュする


公式githubより

まずは、Image Repogitoryのエンドポイントを確認してメモします。
SnowflakeのSQLワークシートで以下のコマンドを実行します。

Snowflake SQL worksheet
SHOW IMAGE REPOSITORIES IN SCHEMA CONTAINER_HOL_DB.PUBLIC;

出力結果のREPOSITORY_URL (例:org-account.registry.snowflakecomputing.com/CONTAINER_hol_db/public/image_repo) をメモっておきます。イメージをpushするときに使います。

ではプッシュしていきます。Snowflake上のImage Repositoryにログインし、ローカルのイメージにタグをつけます。

# e.g. if repository_url = org-account.registry.snowflakecomputing.com/CONTAINER_hol_db/public/image_repo, snowflake_registry_hostname = org-account.registry.snowflakecomputing.com
docker login <snowflake_registry_hostname> -u <user_name>
  > prompt for password
docker tag ghcr.io/mlflow/mlflow:v2.9.2 <repository_url>/mlflow:dev

<snowflake_registry_hostname>については上のコードブロック内のコメントを参照してください。

タグをつけたイメージをプッシュします。

docker push <repository_url>/mlflow:dev

これでプッシュ完了です。SnowSightで以下コマンドを実行することで、きちんとプッシュできているかを確認することができます。

USE ROLE CONTAINER_USER_ROLE;
CALL SYSTEM$REGISTRY_LIST_IMAGES('/CONTAINER_HOL_DB/PUBLIC/IMAGE_REPO');

設定ファイルをステージにプッシュする

docker-compose.yamlファイルを元に設定ファイルを書きます。フォーマットはこちら
snowpark container services は裏でkubernetesが動いているので、設定ファイルの書き方もkubernetesの設定ファイルとパラレルな記法になっています。

mlflow.yaml
spec:
  containers:
  - name: mlflow
    image: <repogitory-uri>/container_hol_db/public/image_repo/mlflow:ver2
    command:
      - mlflow
      - server
      - --host
      - 0.0.0.0
      - --port
      - 5000
    volumeMounts:
    - name: mlruns-volume
      mountPath: /mlflow/mlruns
    - name: artifact-store-volume
      mountPath: /mlartifacts
  endpoints:
  - name: snowpark-mlflow
    port: 5000
    public: true
  volumes:
  - name: mlruns-volume
    source: "@volumes/mlruns"
  - name: artifact-store-volume
    source: "@volumes/artifacts_store"
  networkPolicyConfig:
    allowInternetEgress: true

こちらを、stage上にコピーします。SnowSightでドラッグ&ドロップすることでコピーすることもできますし、SnowCLIで実行することもできます。
後者で実行する場合は、クイックスタートのこちらのセクションを参考にconnection情報を作成した後、以下コマンドを実行することでコピーできます。

cd {mlflow.yamlが置いてあるディレクトリ}
snow object stage copy ./mlflow.yaml @specs --overwrite --connection CONTAINER_hol

上記手順が済んだら、SnowSightを使用して以下のコマンドを実行し、きちんとコピーできているかを確認しましょう。

USE ROLE CONTAINER_USER_ROLE;
LS @CONTAINER_HOL_DB.PUBLIC.SPECS;


以上でローカルでの作業は完了です。

その4 サービスを動かす

ではいよいよSnowflake上でコンテナサービスを作成して、動かしてみましょう。

SnowflakeのSQLワークシートで以下を実行します。

create service CONTAINER_HOL_DB.PUBLIC.mlflow
    in compute pool CONTAINER_HOL_POOL
    from @specs
    spec='mlflow.yaml';

コンテナサービスが立ち上がるまでしばらく時間がかかります。以下コマンドでコンテナの状態を確認できます。

CALL SYSTEM$GET_SERVICE_STATUS('CONTAINER_HOL_DB.PUBLIC.mlflow');


Statusが'Ready'になったら、コンテナのエンドポイント(アクセス用のURL)を表示させましょう。

SHOW ENDPOINTS IN SERVICE mlflow;

エンドポイントにアクセスすると、MLflowサーバーが立ち上がってることがわかります!やったー!

まとめ

Snowpark Container ServicesをつかってMLflow serverを立ち上げることができました!
こちらにモデルを保存できるようになったら、また記事を書きます。

Discussion