Pyppeteer(with headless Chromium) + GitHub Actionsでoptuna-dashboardの継続的E2Eテスト

以前 optuna-dashboard というWebツールを開発・公開しました。 もともと Goptuna のために実装したReact.js + TypeScript製のSPAのWebツールでしたが、Optunaでも使えるようにしたところ、周りでも使ってるよという声をいただくことが増えてきて、公式に利用が推奨されるようになりました。

Optuna v2.7.0 released, with new tutorials, examples, and code improvements! @c_bata_ has fully redesigned the dashboard. Try out the new version with `pip install optuna-dashboard` and then `optuna-dashboard $STORAGE_URL`!https://t.co/MynWjxFYBF pic.twitter.com/8mGCrXzu3N

— OptunaAutoML (@OptunaAutoML) 2021年4月5日

基本的に1人で雑に開発をしていたのですが、バグを出すといろんな方に迷惑をかけそうなので、もうすこしちゃんとテストまわりを整備することにしました。 Python (Bottle) 製のAPIサーバーのユニットテストはもともと結構ちゃんと書いていたのですが、問題はフロントエンドのコードのテストです。 自分がフロントエンドのテストテクニックに詳しくないというのもあるのですが、Plotly.jsによるグラフ描画が主な処理になるため、ユニットテストで保証できる振る舞いはそれほど多くありません。またグラフの描画に必要な情報もやや複雑で、OptunaのFrozenTrialという試行情報を各テストケースで用意するのは大変です。

そこで今回はPyppeteerを使ってVisual regression testingのようなアプローチ *1をとったのですが、結構便利でやってよかったなと思ったので記事に残しておきます。

Pyppeteerを使ってさまざまな目的関数に対するDashboardの表示を確認する

OptunaのFrozenTrialに相当する情報を手動で用意するのは面倒なので、実際にさまざまな目的関数をOptunaで評価して、その評価結果の入ったInMemoryStorageをもとにDashboardを起動、Pyppeteerでアクセスします。コードを簡略化するとざっくり↓のような感じです (全体のコードは こちら)。

import asyncio
import threading
import time
import optuna
import os

from optuna_dashboard.app import create_app
from pyppeteer import launch
from typing import List, Tuple
from wsgiref.simple_server import make_server

host = "127.0.0.1"
port = 8080
output_dir = "tmp"


def create_optuna_storage() -> optuna.storages.InMemoryStorage:
    storage = optuna.storages.InMemoryStorage()

    # Single-objective study
    study = optuna.create_study(study_name="single-objective", storage=storage)

    def objective_single(trial: optuna.Trial) -> float:
        x1 = trial.suggest_float("x1", 0, 10)
        x2 = trial.suggest_float("x2", 0, 10)
        return (x1 - 2) ** 2 + (x2 - 5) ** 2

    study.optimize(objective_single, n_trials=100)

    # Pruning with no intermediate values
    study = optuna.create_study(
        study_name="binh-korn-function-with-constraints", storage=storage
    )

    def objective_prune_with_no_trials(trial: optuna.Trial) -> float:
        x = trial.suggest_float("x", -15, 30)
        y = trial.suggest_float("y", -15, 30)
        v = x ** 2 + y ** 2
        if v > 100:
            raise optuna.TrialPruned()
        return v

    study.optimize(objective_prune_with_no_trials, n_trials=100)

    # No trials
    optuna.create_study(study_name="no trials", storage=storage)
    return storage


async def take_screenshots(study_ids: List[int]) -> None:
    browser = await launch()
    page = await browser.newPage()
    await page.setViewport({"width": 1200, "height": 3000})

    for study_id in study_ids:
        await page.goto(f"http://{host}:{port}/dashboard/studies/{study_id}")
        time.sleep(5)
        await page.screenshot(
            {"path": os.path.join(output_dir, f"study-{study_id}.png")}
        )
    await browser.close()


def main() -> None:
    os.makedirs(output_dir, exist_ok=True)

    storage = create_optuna_storage()
    app = create_app(storage)
    httpd = make_server(host, port, app)
    thread = threading.Thread(target=httpd.serve_forever)
    thread.start()

    study_ids = [s._study_id for s in storage.get_all_study_summaries()]
    loop = asyncio.get_event_loop()
    loop.run_until_complete(take_screenshots(study_ids))

    httpd.shutdown()
    httpd.server_close()
    thread.join()


if __name__ == "__main__":
    main()

このプログラムの流れは次のような感じです。

1. Optunaで実際にOptunaで様々な目的関数の最適化を回す (create_optuna_storage() 関数)。
2. InMemoryStorageを使ってoptuna-dashboardのAPIサーバーを別スレッドで立ち上げ
3. Pyppeteerを使ってheadless Chromiumで各ページにアクセスして、スクリーンショットを撮影
4. Pyppeteerの処理が一通り終われば、WSGIRefのサーバーを終了

実際にOptunaで最適化を回し、そのStorage情報を使ってJSON APIサーバーを立ち上げているので、FrozenTrialのfixtureを手動で頑張って用意する必要はありません。実行すると↓の画像のようにスクリーンショットが生成されるため、様々な目的関数に対する表示結果を簡単に確認できるようになりました。

実際結構便利で、表示中にクラッシュしていたりすると真っ白のスクリーンショットが生成されるのですぐに気が付きますし、質的変数とかの表示がちゃんとできてるかみたいなユニットテストでの確認が難しい問題も手軽に目視で確認できます。

GitHub Actionsで動かす。

せっかくなのでGitHub Actionsで動かしてみます。 撮影したスクリーンショットをGitHubのコメント欄に貼り付けるようにしようかと思ったのですが、結構いろんな目的関数で評価していることもあり画像の枚数が多く全部コメント欄に貼り付けるのはやや大変です。 将来的には画像をGIFとかにまとめて一気に見れるようにするのもいいかなと思いましたが、今回はもう少し簡易的なチェックのみを実行するようにします。

async def contains_study_name(page: Page, study_name: str) -> bool:
    h6_elements = await page.querySelectorAll("h6")
    for element in h6_elements:
        title = await page.evaluate("(element) => element.textContent", element)
        if title == study_name:
            return True
    return False


async def take_screenshots(storage: optuna.storages.BaseStorage) -> List[str]:
    validation_errors: List[str] = []

    browser = await launch()
    page = await browser.newPage()
    await page.setViewport({"width": args.width, "height": args.height})

    study_ids = {s._study_id: s.study_name for s in storage.get_all_study_summaries()}
    for study_id, study_name in study_ids.items():
        await page.goto(f"http://{args.host}:{args.port}/dashboard/studies/{study_id}")
        time.sleep(args.sleep)

        if not args.skip_screenshot:
            await page.screenshot(
                {"path": os.path.join(args.output_dir, f"study-{study_name}.png")}
            )

        is_crashed = not await contains_study_name(page, study_name)
        if is_crashed:
            validation_errors.append(
                f"Page is crashed at study_name='{study_name}' (id={study_id})"
            )

    await browser.close()
    return validation_errors

チェックしているのは、study_nameを含むh6タグが存在するかどうかです。 もしアプリケーションがクラッシュして真っ白になっていた場合、このチェックに引っかかるので最低限クラッシュしていないということはチェックできるようになりました。これをGitHub Actionsで動かします。

name: integration-tests
on:
  pull_request:
    branches:
      - main
    paths:
      - '.github/workflows/integration-tests.yml'
      - '**.py'
      - '**.ts'
      - '**.tsx'
      - 'package.json'
      - 'package-lock.json'
      - 'tsconfig.json'
jobs:
  test:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v2

      - name: Set up Node v14
        uses: actions/setup-node@v2-beta
        with:
          node-version: '14'
      - run: npm install
      - run: npm run build:dev

      - name: Set up Python
        uses: actions/setup-python@v2
        with:
          python-version: '3.x'
          architecture: x64
      - name: Install dependencies
        run: |
          python -m pip install --progress-bar off --upgrade pip setuptools
          pip install --progress-bar off .
          pip install --progress-bar off pyppeteer
      - name: Cache headless chromium
        id: cache-chromium
        uses: actions/cache@v2
        with:
          path: ./local-chromium
          key: chromium
      - run: pyppeteer-install
        if: steps.cache-chromium.outputs.cache-hit != 'true'
        env:
          PYPPETEER_HOME: ./local-chromium

      - run: python visual_regression_test.py --skip-screenshot --sleep 3
        env:
          PYPPETEER_HOME: ./local-chromium

GitHub Actionsの設定はこんな感じで、思ったよりもスッキリかけました。特にハマりどころもなくスッと動きました。

1. Nodeや依存ライブラリをインストールして、JSのコードをビルド
2. Pythonや依存ライブラリをインストールして、optuna-dashboardを動かせるようにする。
3. actions/cache を使ってheadless Chromiumのバイナリ (~150MB) をキャッシュ (PYPPETEER_HOME 環境変数でキャッシュ先のディレクトリを変更できます)。
4. Visual regression testingを実行。スクリーンショットの撮影は必要ないので、スキップしています。またページ遷移してからchromiumの描画が完了するまで少し待つ必要があるので、各ページ3秒待ってからチェックを実行しています。

まとめ

Pyppeteerを使って手軽にVisual regression testが走らせられるようになり、動作確認がだいぶ楽になりました。 Google Summer of Codeのプロジェクトテーマになっていることもあり、先にこの辺を整備しておけてよかったかなと思います。 1つ気になるのはPyppeteerの開発が活発ではなさそうな点ですが、かなり便利なツールだったのでまた動かなくなることが増えてきたら自分も開発手伝いたいなと思います。

*1:開発のなかで見た目に変化が入ることは多いため、なにか画像のdiffをとるようなSnapshot testing的なことはやっていません。あまりこの辺の用語の厳密な定義がよくわからなかったので「のような」とぼかしました