機械学習基盤をGKE Autopilotに移行してコストを削減した

リブセンスで機械学習基盤の開発・運用をしている赤坂（yyyyskkk）です。

我々のチームでは今年の７月ごろにGKE Standard（以下Standardクラスタ）上に構築していた機械学習基盤をGKE Autopilot（以下Autopilotクラスタ）に移行しました。GKE Autopilotとはノードやポッドを自動で管理してくれるクラスタです（詳しくはGoogleのブログをご覧ください）。この記事ではなぜAutopilotクラスタに移行したのか、移行する上でどんな作業が必要だったかという話を書きます。

なぜAutopilotクラスタに移行したのか
検証
移行に必要だった作業
移行後の効果
これから

なぜAutopilotクラスタに移行したのか

高額なノードが複数立ち上がる問題

弊社では機械学習基盤をGKE上に構築しています。Argo Workflowsを利用してレコメンド生成などのバッチ処理を動かしている他、社内向けのWebサービスなどもGKE上で運用しています（Argo Workflowsについては導入時のブログ記事をご覧いただくと、おおよそどのようなものかが分かると思います）。

この機械学習基盤はタイトルにある通りStandardクラスタで運用されていましたが、本来1ノードだけで良いはずの時間帯で2ノードが動作し続けるという現象が起きていました。バッチ処理は夜間や早朝に動かすことが多く、昼間は必要なリソースが少なくなるのですが、この時間帯でオートスケールが適切に動作しなかったのです。いつからかは正確にはわからないのですが2020年夏頃には認識されており、チーム内で問題になっていました。

ワークフローは全て冪等に書かれており、同じ処理が日に二度実行されたとしてもサービスに影響があるわけではありません。しかし、機械学習用のノードはスペックが高く高額です。高スペックなノードの料金を倍支払うことになるのはもったいないわけです。

system podが原因？

調べたところsystem podがオートスケールを阻害するケースがあるらしい、ということがわかりました（参考）。しかしながら明確な対処法もなく、解決されそうな雰囲気もありません。FAQではPDBを付与せよと書かれていますが、system podに付与するのは管理が手間ですし、検証も大変です。そこで浮上したのがAutopilotクラスタへの移行という案でした。

StandardクラスタとAutopilotクラスタの違い

StandardクラスタとAutopilotクラスタを比較する上で、重要な違いの一つだったのが、課金体系です。従来のStandardクラスタではノード単位で課金が発生します。仮にほとんどリソースを消費していなくても、ノードを使用すると1ノード分丸々請求が来るということが起こり得ます。一方、Autopilotクラスタはリソース単位で課金されます。つまり実際のCPUやメモリの使用量に応じて請求が発生するわけです。したがってAutopilotクラスタに移行することで、日々実行しているジョブに対して過大であった使用料金が削減されることが期待できました。

またAutopilotクラスタのそもそもの利点も移行の強い動機でした。Autopilotクラスタであればノードプールを管理する必要もなくなり、podごとに必要なリソースを確保できます。運用コストを下げつつ、ノードに縛られない柔軟なシステム構成を実現できます。良いことづくめですね。

検証

とはいえ社内にAutopilotクラスタについての知見があるわけでもなく、普段動かしているArgo Workflowsのワークフローが動作するかといったことでさえ定かではなかったので、まずはsandbox用の環境でワークフローの動作確認と大まかな費用の推定を実施することにしました。

動作確認は動かすだけです。この過程で軽微ではあるものの修正が必要なことがわかったりしました（後述）。

費用の推定では、sandbox環境では本番で動かしているジョブのうち一部だけを動作させました。一部のジョブだけを動作させている状態で「もし全てのジョブを動作させたら費用がどれくらいになるか」ということを推定するために各ジョブにより消費されるリソースを元にジョブ実行にかかる費用を見積もりました。費用と消費するリソースを詳細に集計するためにCloudBillingの機能を使い、BigQueryへデータをエクスポートしています。BigQueryにエクスポートできれば、クエリで集計もできるので分析が容易になります。

検証の結果Autopilotクラスタに移行することで費用が削減できるだろうという推測が得られました。

移行に必要だった作業

基本的に標準クラスタで使用していた設定ファイルはそのまま流用できるという想定でしたが完全にそのままというわけには行きませんでした。実際にStandardクラスタからAutopilotクラスタに移行する上で必要だったことを説明します。

Argo WorkflowsのExecutorを変更した

前述の通り、バッチ処理にはArgo Workflowsを利用しています。Standardクラスタ時代はpnsというExecutorを利用していました。クラスタのバージョンをv1.19にするタイミングでArgo Workflowsで利用している機能との兼ね合いでpnsを指定する必要があったからです。Executorのドキュメントを見ると分かる通り、Autopilotクラスタで機能すると明記されているのはemissaryだけなので、移行のタイミングでpnsからemissaryへの変更を実施しました。

そもそも先のドキュメントを見れば分かる通り、emissary以外のExecutorはdeprecatedでありv3.4からは削除されるので、変更することは特に問題ではない、というか変更する方が望ましいでしょう。

pnsとemissaryの細かい違いとして、emissaryでは各ステップで必ずなんらかのコマンドを実行する必要がある、というものがあります。一部のワークフローではファイルをGCSのバケットに保存するためにArtifactsを使っています。ボリュームをマウントして、その中のファイルをArtifactsとしてアップロードするだけなので、なんのコマンドも指定していなかったのですが、emissaryでは次のようなエラーがでます。

when using the emissary executor you must either explicitly specify the command, or list the image's command in the index: https://argoproj.github.io/argo-workflows/workflow-executors/#emissary-emissary

ドキュメントを読むと workflow-controller-configmap の記述を変更することで、imageごとにデフォルトで実行するコマンドを指定できるようです。ただ、アップロードするだけのステップはそれほど多用しているわけではないので、代わりに echo するだけで何もしないコマンドを実行することにしました。

- name: upload
  outputs:
    artifacts:
    - name: workdir
      path: /workspace/hozon/suru/file
      archive: {none: {}}
      gcs:
        bucket: livesense.no.gcs-bucket
        key: /hozon/sakino/bucket/no/path
  container:
    image: debian:latest
    command: ["echo", "upload"] # この行を追加した
    volumeMounts:
    - name: workspace
      mountPath: /workspace

メモリ不足が発生したためresource requestの設定を追加した

ワークフローをAutopilotクラスタに移行中、一部のワークフロー（の一部のステップ）がOOMにより失敗するという事象が起こりました。これ自体は珍しいことではなく、（Standardクラスタの頃でも）データ量が急に増えると機械学習の重い処理を実行するステップでOOMになるということは度々経験しています。しかしこの時失敗したのは、これまで失敗したことのないもっと軽い処理でした。

Autopilotクラスタでは、pod仕様で定義されていないリソースはデフォルトの値が使われます（参考）。このデフォルト値がCPU 0.5vCPUとメモリ 2GiBなのですが、一部でメモリ2GiBでは足りない処理があったのです。機械学習の重い処理を実行するステップでは、もともと明示的にメモリを指定していたので引っかかることがありませんでした。

一部のステップにresource requestを明示的に指定することで対応しました。

- name: sql_wo_jikkou_suruzo
  inputs:
    parameters:
    - name: sqlfile
    - name: outfile
  container:
    # ↓2Giだと足りないことがあるので追加した
    resources:
      requests:
        cpu: 1
        memory: 4Gi
        ephemeral-storage: 1Gi
      limits:
        cpu: 1
        memory: 4Gi
        ephemeral-storage: 1Gi

ephemeral-storageの指定はなくても良いのですが、Deploymentなどで設定がなかった場合、リソースを追加・変更した際に「Warning: Autopilot increased resource requests for Deployment argo/workflow-controller to meet requirements. See http://g.co/gke/autopilot-resources.」というワーニングが出ることがあるので追記しています。

CronWorkflowでジョブが二重起動してしまう問題に対応した

長時間動作し続けるpodでは、動作中にノードがスケールダウンして、podが別のノードに移される（evict）ということがよくあります。evictによりpodは別ノードで再作成されるのですが、CronJobは再作成されたpodを管理できず、再度podを作成してしまいます。結果的に同じ処理が二重に動くという事象が発生することがありました。これ自体はStandardクラスタでも発生することがあり、対策（後述）を講じていたのですがAutopilotクラスタでは使えなくなったので別の方法で対処しました。

なおStandardクラスタでも同様の問題が起きると書きましたが、深刻度ではAutopilotクラスタの方が上です。ジョブを二つ動かせば使用するリソースは２倍になり、使用料金も２倍になります。また機械学習系のジョブでは実行時間が10時間を超えることもあります。実行が長時間にわたると、再作成されたpodがさらにevictされ、3重、4重、と多重起動することもありました。費用削減のために移行しようとしているのに、これでは意味がありません。

さてこの問題に対処するためにStandardクラスタで運用していた時は各podに"cluster-autoscaler.kubernetes.io/safe-to-evict": "false"というアノテーションを付与し、podが別ノードに移らないようにすることで二重起動を防いでいました。しかし、Autopilotクラスタではこのアノテーションは使用できません（v1.21から）。

検討の結果、CronJobをCronWorkflowsで置き換えることにしました。

CronWorkflowsを採用した経緯ですが、私たちのチームではArgo WorkflowsをキックするCronJobを定期的に実行することで、日次、週次のバッチ処理を実現していました。CronJobを使っていたのはArgo WorkflowsにCronWorkflowsが実装される前からArgo Workflowsを利用しているからです。

ジョブはCronitorで監視している都合上argo submitに-–waitフラグをつけて実行しています。--waitフラグを外せば、argo submitはすぐに終了するのでCronJobもすぐに終了し、多重起動の問題は起こらなくなることは予想できました。しかし、そのためにはワークフロー側でCronitorの通知をキックするように修正を入れる必要があり、少々手間です。

そこで、「Argo Workflowsを定期的に実行する」という点で同等の機能をもつCronWorkflowsを試してみたというわけです。

もともとのCronJobは以下のように定義されていました。

# CronJobの例
# 具体的な処理はWorkflowTemplateに定義してある。
apiVersion: batch/v1beta1
kind: CronJob
metadata:
  name: mainichi-jikkou-suru-yatu
spec:
  schedule: "00 19 * * *"
  startingDeadlineSeconds: 600
  jobTemplate:
    spec:
      backoffLimit: 0
      template:
        spec:
          restartPolicy: Never
          containers:
          - name: main
            image: argo-wo-jikkou-dekiru-image
            command: ["cronitor", "exec", "Cronitor-no-code", "sh", "-euc"]
            args:
            - >-
              argo submit --wait
              --from workflowtemplates/mainichi-jikkou-suru-template
              --name workflow-`date +%Y-%m-%d`

CronWorkflowだと次のようになります。

apiVersion: argoproj.io/v1alpha1
kind: CronWorkflow
metadata:
  name: mainichi-jikkou-suru-yatu
spec:
  schedule: "00 19 * * *"
  timezone: "Asia/Tokyo"
  concurrencyPolicy: Replace
  startingDeadlineSeconds: 600
  failedJobsHistoryLimit: 10
  workflowSpec:
    entrypoint: main
    templates:
    - name: main
      container:
        image: argo-wo-jikkou-dekiru-image
        command: ["cronitor", "exec", "Cronitor-no-code", "sh", "-euc"]
        args:
        - >-
          argo submit --wait
          --from workflowtemplates/mainichi-jikkou-suru-template
          --name workflow-`date +%Y-%m-%d`

細かい違いはありますが、この程度であれば書き換えの負担はさほど大きくなく、既存の設定ファイルを流用できます。CronJobをCronWorkflowsに置き換えた結果、多い時で週に数回は起きていた多重起動が、まったく発生しなくなりました。本来のWorklfowとWorkflowを起動するためのCronWorkflowでワークフローは２倍になりますが、一方で今まで使っていたCronJobはなくなるので大きなデメリットではないと考えています。