機械学習バッチ実行環境を整備した話

はじめまして。

機械学習チームにてレコメンドの改善を行っているgumigumi4fです。

本記事では、機械学習チームの取り組みの一環として機械学習等のバッチを含むバッチ実行環境を整備した話について話したいと思います。

今までのバッチ実行環境

機械学習チームではピクシブ会社全体のサービスにおけるアイテムのレコメンド等を主に取り組んでおり、そのロジックは多岐に渡ります。

matrix factorizationを用いた手法
item間共起頻度に基づくシンプルなアイテムの推薦手法
ニューラルネットを用いた手法

これらの計算を行うためのバッチは実装者の好きな手法で実装されており、バッチを行う環境に関しても下記にようにバラバラになっていました。

オンプレの機械学習用マシンに各ユーザー毎に設定したcronでバッチが実行
gitlab-ciのスケジューリングによってバッチが実行
BigQueryで完結するような処理はJenkinsやAirflowにて実行

特にここで問題になるのが、オンプレにcronで実装されたバッチです。オンプレにて設定されているバッチには以下のような問題がありました。

ユーザー毎にcronを設定することからいつバッチが実行されるのかわからずリソースの競合が起きる可能性がある。
マシン自体が冗長化されていないためマシンが故障した際に全てのバッチが動作不能になる
バッチが正しく実行されたのかどうかの監視の実装が各々に任されるため適切な監視が行えない

このような問題を解決するためにシンプルなバッチを実行する環境を整備することにしました。

検討したバッチ実行環境

バッチ実行環境にはArgo WorkflowとDigdagの2種類について検討を行いました。

Argo Workflow

Argo Workflowはピクシブ内でも他チームで採用しているチームがあり機械学習チームでの利用も検討しましたが、問題となったのがKubernetesに対する知識の問題です。

Argo Workflowではワークフローの記述を当然yamlで行う必要があり、そのyamlはKubernetesに対する理解がないと簡単に扱えるようなものではありません。しかしながら、バッチ実行環境を検討している当時はKubernetesに対する理解がある人間がチーム内には自分しかおらず、バッチ実行環境を整えてもチーム内で利用を推進できない可能性があったため今回は導入を見送りました。

Digdag

Digdagでは同じArgo Workflowと同様にyamlを用いたワークフローの記述をするものの、記述はより単純でわかりやすいためとっつきやすいという点で良いという決断をし、チーム内での導入をすることにしました。

Digdag内からkubectlを叩いてjobを作成し、バッチを実行するようにすることで適切にリソースを確保した上でバッチが実行されるようになり、バッチが実行されたかどうかも共通の形でslackなどに通知する仕組みを作成することができるようになります。

実際のワークフローの一部 (例)

timezone: Asia/Tokyo

( 略 ) 

+run:
  +prepare:
    call>: k8s/job/cpu
    _export:
      job:
        name: prepare
        image: "us-central1-docker.pkg.dev/hoge/fuga/foo:bar"
        command: "hoge prepare --input ${input_file} --output ${output_file}"
        resource_cpu: 8000m
        resource_memory: 48G

  +training:
    call>: k8s/job/gpu
    _export:
      job:
        name: training
        image: "us-central1-docker.pkg.dev/hoge/fuga/foo:bar"
        command: "hoge training --dataset ${output_file} --model ${output_model}"
        resource_cpu: 12000m
        resource_memory: 64G

完成したバッチ実行環境

Digdagを用いたバッチの実行環境はGCP上に構築しました。

Digdagから学習や推論を行うジョブをk8s経由で走らせ、その結果をGCSやBQに格納することで活用するような仕組みとなっています。

最後に

今回、既存の機械学習バッチの実行基盤を整備しました。すでにdigdag上で大量のワークフローが運用されています。

機械学習チームでは自然言語処理、画像処理、テーブルデータなどの様々なデータとタスクがあり、インターン・アルバイト・新卒採用・中途採用問わず、採用を行っています。

もしご興味のある方は下記をご参照ください