最近の砂場活動その22: Argo WorkflowsでGKE上の機械学習バッチのワークフローを管理する

なぜ

  • 機械学習の典型的なワークフローはこんな感じになりがち
    • データの取得
    • 前処理
    • モデルの学習
    • 新しいデータに対して予測を行なう
  • ごくまれにコケる。こけたらretryなどをして欲しいし、その辺はワークフロエンジンで面倒を見て欲しい
  • 推薦タスクの機械学習ジョブをGKE上で動かしているので、GCPで動くワークフローエンジンを試してみたい
  • Kubernetesとあれこれ連携できると便利なので、Argo Workflowsを使ってみることにする
  • ...というのは前置きで、仕事でArgoを使うので気軽に遊べる砂場を作りたかったのであった

インストール

簡単にインストールできる。

kubectl create ns argo
kubectl apply -n argo -f https://raw.githubusercontent.com/argoproj/argo-workflows/stable/manifests/quick-start-postgres.yaml

helm経由で入れるときはこちら。

% kubectl create namespace argo
namespace/argo created
% helm install argo argo/argo --namespace argo
manifest_sorter.go:192: info: skipping unknown hook: "crd-install"
manifest_sorter.go:192: info: skipping unknown hook: "crd-install"
manifest_sorter.go:192: info: skipping unknown hook: "crd-install"
manifest_sorter.go:192: info: skipping unknown hook: "crd-install"
NAME: argo
LAST DEPLOYED: Mon May  3 16:31:18 2021
NAMESPACE: argo
STATUS: deployed
REVISION: 1
TEST SUITE: None
NOTES:
1. Get Argo Server external IP/domain by running:

kubectl --namespace argo get services -o wide | grep argo-server

2. Submit the hello-world workflow by running:

argo submit https://raw.githubusercontent.com/argoproj/argo/master/examples/hello-world.yaml --watch

テンプレートの具体例

こういう感じ。name: recommendというのをテンプレート的に宣言しておいて、変数化したいものをargumentsでバチバチ当てはめていく。下の例だと、前処理である{prep1, prep2}が並列で動いて、その後学習の{train1, train2}が走って、最後にpredictする、という感じ。predictのところでこけたら、Webコンソールからこけたところからretryする、といったこともできる。exit handler的なものも実装できるので、{成功, 失敗}したらSlackに通知というのも簡単にできる。

apiVersion: argoproj.io/v1alpha1
kind: CronWorkflow
metadata:
  name: {{ .Values.app.name }}
  namespace: {{ .Release.Namespace }}
spec:
  schedule: "0 */1 * * *"
  concurrencyPolicy: "Replace"
  startingDeadlineSeconds: 0
  workflowSpec:
    entrypoint: main
    serviceAccountName: {{ .Values.app.name }}-sa
    templates:
    - name: main
      steps:
        - - name: prep1
            template: recommend
            arguments:
              parameters:
                - name: script_name
                  value: "prep1.py"
          - name: prep2
            template: recommend
            arguments:
              parameters:
                - name: script_name
                  value: "prep2.py"
        - - name: train1
            template: recommend
            arguments:
              parameters:
                - name: script_name
                  value: "train1.py"
          - name: train2
            template: recommend
            arguments:
              parameters:
                - name: script_name
                  value: "train2.py"
        - - name: predict
            template: recommend
            arguments:
              parameters:
                - name: script_name
                  value: "predict.py"
    - name: recommend
      inputs:
        parameters:
          - name: script_name
      container:
        image: gcr.io/ml-news/{{ .Values.app.name | replace "-" "_"}}:latest
        command:
          - "python"
          - "{{ `{{inputs.parameters.script_name}}` }}"
        env:
          - name: OPENBLAS_NUM_THREADS
            value: "1"
          - name: PYTHONUNBUFFERED
            value: "1"
      resources:
        request:
          cpu: 200m
          memory: 1.5Gi

この設定の反映は前回やったhelmで行なうと自動化できて便利。

www.yasuhisay.info

Webコンソールを手元から見る

Ingressなどのサービスを準備するのは面倒(個人でやる場合、ロードバランサーのお金をケチりたい...)ので、もっと手軽な方法を試す。これが多分一番簡単。

argo server -n argo
open http://localhost:2746/

ワークフローの一覧を見る

argoはWebコンソールも見やすいが、cliも充実している。ワークフローのstatusなどもcliからさっと取れる。

% argo list
NAME                STATUS      AGE   DURATION   PRIORITY
hello-world-4xhtw   Succeeded   3m    10s        0
hello-world-vzgpx   Succeeded   4m    27s        0

ジョブを投げる

試行錯誤するときには手元からジョブを投げれると便利。こんな感じでさっと投げれる。

argo submit --watch https://raw.githubusercontent.com/argoproj/argo/master/examples/hello-world.yaml

ログもtail -f的な感じで見ることができて便利。

% argo logs @latest -n argo --follow
hello-world-4xhtw:  _____________
hello-world-4xhtw: < hello world >
hello-world-4xhtw:  -------------
hello-world-4xhtw:     \
hello-world-4xhtw:      \
hello-world-4xhtw:       \
hello-world-4xhtw:                     ##        .
hello-world-4xhtw:               ## ## ##       ==
hello-world-4xhtw:            ## ## ## ##      ===
hello-world-4xhtw:        /""""""""""""""""___/ ===
hello-world-4xhtw:   ~~~ {~~ ~~~~ ~~~ ~~~~ ~~ ~ /  ===- ~~~
hello-world-4xhtw:        \______ o          __/
hello-world-4xhtw:         \    \        __/
hello-world-4xhtw:           \____\______/

この辺のコマンドはよく使うので、Makefileで包んでいつも使っています。

RELEASE=post-list-members-to-bigquery
NAMESPACE=post-list-members-to-bigquery-ns

submit:
  argo submit --from cronwf/$(RELEASE) --namespace $(NAMESPACE) --watch

logs:
  argo logs @latest -n $(NAMESPACE) --follow