エムスリーテックブログ

エムスリー(m3)のエンジニア・開発メンバーによる技術ブログです

AI・機械学習チームでのインターンでBigQueryのローカルテスト基盤を作った話

こんにちは、10月後半の2週間、エムスリーのAI・機械学習チームでインターンをしていた後藤です。

今回は私の行ったタスクと、インターン生からみてエムスリーはどのような会社なのかについて書いていこうと思います。

他の学生の方々のインターン選びやBigQueryのテスト手法に悩んでいるエンジニアの方々の参考になれば幸いです。

BigQueryのローカルテスト基盤を作った話

背景

AI・機械学習チームでは、BigQueryに日々蓄積されている大規模データから所望のデータを抽出しモデルの学習や分析に用いています。この抽出の過程で共通する処理をまとめたライブラリであるm3downloaderをチーム用に整備しており、m3downloaderを通じてBigQueryとやりとりすることで、コードを共通化することで可読性を高めたり作業の効率化に役立てています。

この便利なm3downloaderですがコアとなるSQL部分について、テストによって品質の保証がなされていないという問題がありました。そもそもBigQueryを含むロジックに対するテストが難しいというのが根本的な原因です。BigQueryはテスト用のエミュレータを提供していないため、テストを行うにはBigQueryのサービス自体にアクセスする必要があり、費用などを踏まえるとCIで定期的にテストを実行するのが現実的ではないからです。

BigQuery Emulatorの登場

この問題に真っ向から立ち向かったのがgoccyさんの開発するGo製のOSSであるBigQuery Emulatorです。2022年の6月に公開されたBigQuery Emulatorですが、その名の通りBigQueryの動作をエミュレートするプログラムであり、このソフトウェアを用いることでBigQueryとの通信を含むロジックをローカル環境でテストできるようになりました!

github.com

そこで、m3downloaderにおいてBigQueryとの通信を含むロジックをテストする手法をBigQuery Emulatorを用いて確立させるというのが、私のインターンでのタスクになりました。

bqemulatormanagerの作成

BigQuery EmulatorをPythonから扱えるようラップしたbqemulatormanagerというライブラリを開発し、OSSとして公開しました。まだまだ実用には遠い出来栄えですが同様のものを作成する際に何かの参考になればと思います。bqemulatormanagerにはBigQuery Emulatorとの通信部分に加え、テストを楽にするための機能が2つ搭載されています。

github.com

スキーマの自動取得

1つ目が、テーブルのスキーマの自動ダウンロードです。BigQueryにおいてはテーブルを作成する際に、各カラムに入る値の情報をスキーマとして指定してあげる必要があります。1つのテーブルにおいてカラム数が数十に上ることもあるため、この指定をテストコードを追加するたびに人力で行うのは得策とは言えません。そこで、テストしたいテーブルは本番環境に既に存在していることに着目し、対応するスキーマを自動取得してエミュレータに渡す処理が入っています。加えて、本番環境との通信回数を減らすため、取得したスキーマをyamlファイルに保存しており、ファイル内に対応するスキーマが存在する場合はそちらから情報を取得しています。

並列処理への対応

2つ目が並列処理への対応です。m3downloaderにおいてはCIを高速に終わらせるためにテストケースを並列で実行しています。そのため、何も考えずに実装すると立ち上がったエミュレータ同士がバッティングするという問題が発生する(した)ため、ポートが被らないよう回避する機能を追加しています。

テストコードの導入

上記のbqemulatormanagerの開発に加えて、作成したbqemulatormanagerを用いて実際にm3downloaderで使われている関数のテストコードも作成しました。最終的に2つの関数のテストに対応し、BigQueryテスト基盤作成への第一歩を踏み出すことができたと思います。実際に作成したテストの例を紹介しようと思います。なお、公開用にテーブル名などを修正しています。

まず、コアとなるSQLがこちらになります。

SELECT
  code,
  service_name
FROM (
  SELECT
    user_use_list.code,
    service.name AS service_name
  FROM
    `dataset1.user_use_list` user_use_list
  INNER JOIN
    `dataset1.service` service
  ON
    user_use_list.service_id = service.id )
UNION DISTINCT (
  SELECT
    users.code,
    service_name
  FROM (
    SELECT
      package_use_list.ID,
      service.name AS service_name
    FROM
      `dataset1.package_use_list` package_use_list
    INNER JOIN
      `dataset1.service` service
    ON
      package_use_list.SERVICE_ID = service.id ) lst
  INNER JOIN
    `dataset1.users` users
  ON
    users.ID = CAST(lst.ID AS STRING) )

サービスを利用しているユーザーのidとそのサービス名を取得するものになります。複数のテーブルに情報が分散しているのでUNIONを使ってまとめています。

次にテスト対象となる関数です。

def download_user_list(downloader: BigQueryDownloader) -> UserList:
    sql = read_sql(file_name='hogehoge/user_list.sql')
    user_list = downloader.download(sql)
    user_list = user_list.astype(dict(code=int))
    return UserList(user_list)

m3downloaderにおいてはBigQueryDownloaderの引数を変えてインスタンスを生成することで通信先などの振る舞いを変更できるようにしています。テスト時にはエミュレータと通信するように設定したdownloaderを渡します。

最後にテスト用コードです。

def test_result_download_user_list(self):

    # テスト用データの用意
service_df = pd.DataFrame([
        {
            'id': 0,
            'name': 'service-A'
        },
        {
            'id': 1,
            'name': 'service-B',
        },
        {
            'id': 2,
            'name': 'service-C',
        },
    ])

    use_df = pd.DataFrame([
        {
            'code': '0111',
            'service_id': 0
        },
        {
            'code': '0111',
            'service_id': 1
        },
        {
            'code': '0222',
            'service_id': 0
        },
    ])

    package_df = pd.DataFrame([
        {
            'SERVICE_ID': 0,
            'ID': 11,
        },
        {
            'SERVICE_ID': 2,
            'ID': 33,
        },
    ])

    user_df = pd.DataFrame([
        {
            'ID': '11',
            'code': '0111'
        },
        {
            'ID': '22',
            'code': '0222'
        },
        {
            'ID': '33',
            'code': '0333'
        },
        {
            'ID': '44',
            'code': '0444'
        },
        {
            'ID': '55',
            'code': '0555'
        },
    ])

    expected_df = pd.DataFrame([
        {
            'code': 111,
            'service_name': 'service-A'
        },
        {
            'code': 111,
            'service_name': 'service-B'
        },
        {
            'code': 222,
            'service_name': 'service-A'
        },
        {
            'code': 333,
            'service_name': 'service-C'
        },
    ])

    # エミュレータとやりとりするためのマネージャーを作成
    manager = bqemulatormanager.Manager(project='project_a')
    with manager:

        # データの挿入
        manager.load(use_df, 'dataset1.user_use_list')
        manager.load(service_df, 'dataset1.service')
        manager.load(package_df, 'dataset1.package_use_list')
        manager.load(user_df, 'dataset1.users')

        # m3downloaderで使用するための設定
        downloader = BigQueryDownloader(project_id='project_a', mode=BigQueryDownloaderMode.EMULATOR)
        downloader.client.set_manager(manager)

        # クエリの実行
        result = download_user_list(downloader=downloader)

    # 結果が想定通りかチェック
    assert_frame_equal(expected_df, result.data)

テストデータを作成しエミュレータに注入したのちに、テスト対象の関数を実行し結果が想定通りか確認しています。

ちなみにこちらのテストとほぼ同じものをbqemulatormanagerのexampleにも追加しているのでよければ色々試してみてください!

ここまで、作ったものを紹介してきました。複雑なクエリに対してもBigQueryの動作をうまく模倣できているというのがBigQuery Emulatorの素晴らしい点の1つかと思います。

ただ、一方でBigQuery Emulatorを使用する際に気をつけるべきポイントがあるというのもこの2週間でわかってきました。最後にその点について共有したいと思います。

BigQuery Emulatorを使用する際に気をつけるべきポイント

私が確認した限りでは次の場合のSQL文について上手く動作しません。

  • テーブル名にワイルドカードを含む場合
  • DISTINCT の中のサブクエリでDISTINCTを用いる場合
  • UDFとテーブル名を併用する場合

2022年10月28日現在、これらの場合については上手くエミュレータが動かないので何らかの回避策を考える必要があります。「OSSは動いても動かなくても嬉しい」というのはインターン中に飛び出したメンターの名言ですが、これは非常によく言ったもので、つまり今あげた3つがOSSへのコミットチャンスと言えるでしょう。実際私も、メンターに発見して頂いた軽微なバグを修正するためのPRを出すことができました。先ほど述べた3つの未対応クエリを動くようにするのは現状の私では太刀打ちできない修正ですが、もっと強くなってこれらに対応するPRを出せるようになりたいと思える良い機会でした。

最後はポエムっぽくなってしまいましたが、以上がインターンで行ったタスクのまとめになります。続いて、インターンで得られた事やチームの雰囲気などを紹介していこうと思います。

インターンの話

進め方について

これは私の場合の一例なので、時期や人によって異なる可能性がある事に注意してください。

まず、形式ですが私の場合は関西在住のためフルリモートで9:00-17:45定時の7時間45分勤務でした。リモートワークに不安もありましたが、設定済みのPCを発送していただいたことに始まり、活発に動いているSlackや毎日の朝会・夕会・もくもく会を通して十分なコミニュケーションを取ることができ、特に困った事態も発生せず最終日を迎えることができています。さらに毎日10分ほどランダムに決定した話題で雑談する時間があったり6日目には歓迎会を開いていただいたりとフルリモートでも人と話せなくて辛いと言ったことは起こりませんでした。なお雑談時間についてはこちらで記事になっているので気になる人は確認してみてください。

次にタスクの決め方について説明します。私の場合は初日にチームリーダーの方々やメンターと面談をし、機械学習から開発まで幅広いタスクの中から自分の興味あるものを選択できました。私はML側の選考経由で入った気がしますが、それでも私の希望したBigQueryのテスト基盤作成という開発メインのタスクを任せて頂きました。

タスクの決定後は基本的に一人でその仕事に取り組むことになります。適宜メンターやチームリーダに相談する機会があったり、詰まった時に適当にSlackに流せば誰かが反応してくれたりとここでも楽しくのびのびと開発を行うことができたと思います。

その後、5日が経ったタイミングで中間発表、10日目に最終発表をチームの前で行います。私の場合はその後全社的なエンジニアの勉強会であるTech Talkで発表予定です。発表といっても研究発表のような固い感じではなく(人によってイメージは違うと思いますが)、カジュアルな雰囲気でやったことを報告して議論するという感じです。チャットもワイワイしています。

発表の雰囲気についてはこちらのチャンネルのTech Talkの動画が参考になるかもしれません。 www.youtube.com

AI・機械学習チームについて

最後にAI・機械学習チームについて私の感じたことを紹介して行きます。

事前にエムスリーのエンジニアは強い人が多いイメージがありましたが、想定通り(以上?)でした。エンジニアリング、サイエンス両面での知識の幅と深さには圧倒されました。また知識だけでなく、問題があるときにその原因や対処法を切り分けて整理する力も凄くインターン期間中は何度も助けられました。

個人的にAI・機械学習チームで一番魅力的なのは開発から機械学習までを担当するという仕事の幅広さです。私は学部生の頃はアルバイトなどでweb開発、研究室に配属されてからは機械学習を主にやってきましたが、就職するにあたってどちらか一本に道を絞らなけばいけないと漠然と考えていました。しかし、両面の仕事を担当し強い専門性を発揮されているチームの方々を見て、こういうキャリアプランもあるんだと感動すると共に自分もその方向性を目指したいと思いました。開発から機械学習まで扱うAI・機械学習チームですが、その広さを象徴しているのがチームでの勉強会です。私が参加したある回では前半にPythonでのプロジェクトのモノレポ管理について深い議論を交わした後にみんなでPU学習の理論について勉強していました。このギャップが私に強く刺さりました。

終わりに

2週間という短い間でしたが、多くのことを経験でき自身の成長に繋がった良いインターンだったと思います。メンターの北川さんを始め、チームの皆さん、また、BigQuery Emulatorの開発者であるgoccyさんにはとても感謝しています。ありがとうございました。

もし、これを読んでいる学生で学部時代にweb開発、研究室では機械学習をやっているような人がいればAI・機械学習チームの仕事は非常に魅力的に感じると思うので是非インターンに参加してみてください。(もちろんそうじゃない人も!)

jobs.m3.com