トップ > Machine Learning > 機械学習API基盤にregression test を追加する

機械学習API基盤にregression test を追加する

Machine Learning MLOps

こんにちは、キャディでMLOpsをやっている志水です。機械学習の推論基盤にregression testを追加したところ依存パッケージのアップデート等が楽になり開発者体験がすごくよくなったので、その詳細について書きます。

[toc]

推論基盤の運用

MLOpsチームでは機械学習モデルの推論API基盤を開発運用していています。こちらに関しての詳細は以前のTechブログ をご参照ください。

チームでGoogleのソフトウェアエンジニアリング本を読んだことをきっかけに、現在のプロダクトで改良できる部分を議論しました。

図1. 現状のデプロイフローと、起きえるエラーについて議論した図

現状のデプロイフローでは機械学習エンジニアが以下を手動で行っています。

  1. 実験時に作成したデータとモデルファイル(.ptファイル)を用意
  2. TorchServe でサーブするためにAPI用のDockerコンテナを作成
  3. dev環境へのデプロイと疎通確認
  4. stg環境へのデプロイと負荷試験
  5. prod環境へのデプロイ

ホワイトボードでデプロイフローと既存のテストを整理した結果、機械学習エンジニアが開発時に出した結果と最終的なAPIが同じ結果を出すことを保証するために、手動のテストでカバーしている範囲が多いことがわかりました。 （既存のCIのテストは前処理の違いによる不具合が起きたことで追加したUnit testが大半でした。）

機械学習エンジニアとも相談し、自動化するテストの構成や内容についていくつかのパターンを出しました。

  • 「これからのパッケージアップデートに対して現在と推論が変わらないことを保証る」を目的に、サンプル図面に対して推論を行い、過去の推論結果と比較するregression test
  • 「実験当時の推論結果と、デプロイされたものの精度の一致」を目的に、実験時のデータと推論結果が一致することを確認するテスト

実装の工数やそこから得られるbenefitを天秤にかけ、「これからのパッケージアップデートに対して現在と推論が変わらないことを保証する」ことを目的に、サンプル図面に対して推論を行い、そこで今までの推論結果と同じ結果を返していることを保証するregression test をCIに追加することにしました。

regression testの追加

regression test の構成は

  • サンプル図面とその推論結果を用意し docker-compose.yml で APIを立てて推論リクエストを送り（下記batchプロセス）
  • 返ってきた現在の推論結果が用意しておいた推論結果と一致することを確認する

ようにしました。

図2. サンプル図面

このようなサンプル図面に対して、以下のような推論結果を予め用意しています。

図面_id,pred,name
WA-20220616-ABC-01/RT-1,9.607873916625977,thickness

これをCI内で実行するための docker-compose.yml は以下のようになります。

# docker-compose.yml
services:
  api:
    build:
      context: .
      dockerfile: Dockerfile
    ports:
      - "7080:7080"
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:7080/ping"]
      interval: 30s
      timeout: 10s
      retries: 5

  batch:
    build:
      context: .
      dockerfile: tools/Dockerfile  # 中でテストを呼び出す
    depends_on:
      api:
        condition: service_healthy

    environment:
      API_PORT: 7080
      API_HOST: api

CI

CIでテストが走ってくれれば、エンジニアが人手で確認することなくデプロイされている推論APIが以前のバージョンと乖離していないことを保証でき、安心できます。 キャディでは複数の機械学習APIを運用しているため paths-filter を活用し、変更があったAPIだけテストが発火するようにしています。

github actions 内では以下のようにシンプルにテストを実行しています。

docker compose build
docker compose up api -d
docker compose run batch

この docker-compose.yml を使ったテストは1つのマシンで複数のプロセスが走るという、Googleのソフトウェアエンジニアリングに登場するmedium size のテストになります。関数ごとのテストはunit testなどのsmall testでカバーします。 small test がたくさんあり、medium(やlarge)のテストがそれより少ない状態がピラミッド型のバランスのいいtest suiteとされています。それを実現するためにsmall testはカバレッジ高くたくさん書いており、regression testのようなmedium sizeのテストはクリティカルな部分のみに書いています。

依存パッケージの更新

我々のチームではrenovate というツールによって依存パッケージのバージョンアップデートを管理しています

regression test を導入するまでは、例えばPyTorchにセキュリティパッチがあたりバージョンが上がった場合に、APIを作成した各機械学習エンジニアに想定した結果と精度が変わっていないかを確認してもらう必要があり、手間も時間もかかっていました。

regression test を導入してからはrenovateでパッケージバージョンが上がるPRが上がってきた場合、CI上のregression test通っていれば推論APIの挙動の一貫性についてある程度自信が持てるため、依存関係の更新が高速かつ安全になりました。

当初想定してた以上に開発者としての体験はよく、今後も積極的にテストを開発していくモメンタムがチームに生まれています。

終わりに

今回は推論API基盤に対するregression testの導入に至った経緯と、その具体的な方法について紹介しました。

キャディでは一緒に働く仲間を探しているので、募集要項 から気になる求人があればご気軽にご連絡いただければ幸いです。