【TECH BLOG】Vertex AI Pipelinesによる機械学習ワークフローの自動化

はじめに

こんにちは。検索基盤部の倉澤です。

私たちは、ZOZOTOWNの検索機能の改善に取り組んでいます。ZOZOTOWNのおすすめ順検索ではランキング学習を用いた検索機能の改善に取り組んでおり、A/Bテストにて効果を測定しています。

ランキング学習やElasticsearch Learning to Rankプラグインについては過去の記事で紹介していますので、併せてご覧ください。

私たちは、機械学習モデルの開発からデプロイまでの一連の処理を実行するワークフローの構築にGoogle Cloud Platform（GCP）のVertex AI Pipelinesを利用しています。

本記事では、Vertex AI Pipelines採用前の運用とその課題点について説明し、次にVertex AI Pipelinesで構築したワークフロー概要とその運用について紹介します。

Vertex AI Pipelines採用の背景

Vertex AI Pipelines採用に至った背景として、従来の運用と抱えていた課題点を紹介します。

従来の運用

Vertex AI Pipelinesを採用する以前は、GitHubで管理されているスクリプトを、開発者が各自GCPに立てたインスタンスのJupyter Notebook上で順に実行していました。

機械学習モデルの学習期間や特徴量などのパラメータは設定ファイルで管理しており、papermillでNotebookを自動実行して機械学習モデルを生成します。そして、ElasticsearchのLearning to Rankプラグインで指定された形式にモデルを変換し、手動でデプロイを行っていました。

papermillとは

Jupyter Notebookを実行するライブラリとして記載したpapermillについて簡単に説明します。

papermillは、Jupyter Notebookに定義された各セルを実行するPythonライブラリです。 実行時にパラメータを渡すことで予めセルに定義されたデフォルトのパラメータの上書きが可能です。実行されたJupyter Notebookは別名のNotebookに保存できます。

私たちは、papermillをCLIで実行していました。

papermill input.ipynb output.ipynb -f parameter.yaml

設定ファイルは以下のようにYAMLファイルとして定義できます。

# parameter.yaml
train_start_date: 20220101
train_end_date: 20220102
valid_start_date: 20220103
valid_end_date: 20220104
test_start_date: 20220105
test_end_date: 20220106
features:
  - feature1
  - feature2
  - feature3

抱えていた課題と解決策

従来の運用フローでは、モデルの数だけ同様の作業を手動で繰り返しており、以下の点を課題に感じていました。

• 各タスクの実施作業と実施完了の確認作業の工数が多い
• 実行前に設定ファイルの変更に対するレビューが無いので、誤りがあった場合は機械学習モデルを再度生成し直す必要がある
• 機械学習モデル生成の一連のタスクが途中で失敗した際に、一から再実行する必要がある

これらの課題を解決するために、各タスクを依存関係通りに実行でき、さらに再実行時にはキャッシュが利用できるワークフローエンジンの導入を検討しました。

候補となるワークフローエンジンはいくつかありましたが、弊社MLOpsブロックがVertex AI Pipelinesの実行環境の整備を進めていることもあり、より導入コストが低いVertex AI Pipelinesを選びました。

Vertex AI Pipelinesの実行環境については過去の記事で紹介していますので、併せてご覧ください。

Vertex AI Pipelinesによるワークフローの構築

私たちが構成した機械学習モデルの開発からデプロイまでのワークフローの概要を紹介します。以下の図は、実際に運用しているVertex AI Pipelinesのコンソール画面から確認できるワークフローの全体像です。

本ワークフローではおおよそ以下のことを行っております。

1. 学習データセット生成に必要な期間のデータが揃っているかの確認
2. 学習データセットの生成
3. ハイパーパラメータチューニング及び最適なパラメータによる学習
4. 評価及びオフライン評価結果の描画
5. デプロイ判定

それぞれ順に説明します。

1. 学習データセット生成に必要な期間のデータが揃っているかの確認

学習に必要となる期間のデータが、対象のBigQueryテーブルに存在しているかの欠損チェックを行います。以下のようなAssertionクエリをコンポーネントから実行し、指定期間のデータが存在しているかを確認します。

-- check_bq_table.sql

DECLARE target_dates ARRAY<DATE>;
DECLARE x INT64 DEFAULT 1;

SET target_dates =  (
SELECT ARRAY_AGG(date ORDER BY date) FROM UNNEST(GENERATE_DATE_ARRAY('{{ start_date }}', '{{ end_date }}', INTERVAL 1 DAY) ) AS date
);

WHILE x <= ARRAY_LENGTH(target_dates) DO
  ASSERT EXISTS(
    SELECT {{ period_column }}
    FROM `{{ full_table_id }}`
    WHERE DATE({{ period_column }}) = target_dates [ORDINAL(x)]
  ) AS 'target date does not exist in this table';

  SET x = x+ 1;
END WHILE;

2. 学習データセットの生成

学習・検証・テストのデータセットの生成をします。 以下のYAMLファイルはコンポーネントの入出力を定義し、学習・検証・テストのデータセットを出力しています。

続きはこちら