機械学習手法を用いる研究向け基盤 on Kubernetes ~開発編~

DeepLearning Kubernetes k8s 情報インフラ 深層学習

はじめに

この記事では大学院修了にともなって、これまで構築してきた Kubernetes (K8s) を用いた ML 基盤のうち、開発したソフトウェアを中心に紹介していきます。 6 年間の振り返りは こちらの記事 を、運用に関する紹介は こちらの記事（執筆中）を参照してください。

システムの構成

構築・運用していたシステム構成は以下の図のようになっており、ユーザは専用のクライアントコマンド経由で Web IDE や Jupyter を起動してデータの分析ができるようになっています。

high_level_architecture
アーキテクチャ

クライアントコマンド

前述のクライアントコマンド（以下 cmd）の機能などについて軽く説明します。 cmd のソースコードは以下で公開していますが、cmd は様々な環境で動作させるために作ったわけではないため、ソースコードは参考程度にしてください。

gitlab.com

cmd は主に K8s API を叩くシンプルな物になっており、

  1. 永続ストレージの作成・削除・容量変更
  2. Web IDE や Jupyter の起動・削除・詳細の表示
  3. kubeconfig の入手

の 3 つの機能を提供しています。 永続ストレージの作成では、コードのこの部分 で示すように、

  1. CephFS 用 Persistent Volume Claim (PVC)
  2. NFS ストレージ用 PVC
  3. ストレージバックアップ用 CronJob
  4. 各種設定用 ConfigMap

を作成しており、WEB IDE や Jupyter の起動では コードのこの部分で示すように、後述する Custom Resource を作成しています。

また本システムはプライベートネットワークに対してのみ公開しているため、直接 kubeconfig を配布する形で認証を行なっています。cmd は コードのこの部分 で示すように、次節で紹介する authenticator に kubeconfig のリクエストを送り、受け取った kubeconfig をローカルに保存する仕組みになっています。

authenticator

前述した authenicator について紹介します。 authenticator は以下で公開しており、cmd と違い他の環境でも使用できるようになっていると思います。

gitlab.com

authenticator は、

  1. authenticator
  2. config-reloader

の 2 つのソフトウェアから構成されており、authenticator は

  1. 指定されたラベルを持つ、 ServiceAccount に対する kubeconfig の生成
  2. kubeconfig を配布するための HTTP サーバの起動

を行います。 1 の kubeconfig の生成では、 こちら に示すように、initContainer で kubeconfig を生成し、 2 のHTTP サーバの起動では、 こちらで示すように、container で HTTP サーバを起動します。

また config-reloader は authenticator Pod を再起動するだけの単純な物ですが、こちら に示すように Deployment の spec.template.metadata.annotations フィールドを編集することによって、 Pod を再起動するように実装しています。この実装方法は、以下に示すように本家 kubectl rollout restart コマンドでも同様の方法で実装されています。

github.com

...
    case *appsv1.Deployment:
        if obj.Spec.Paused {
            return nil, errors.New("can't restart paused deployment (run rollout resume first)")
        }
        if obj.Spec.Template.ObjectMeta.Annotations == nil {
            obj.Spec.Template.ObjectMeta.Annotations = make(map[string]string)
        }
        obj.Spec.Template.ObjectMeta.Annotations["kubectl.kubernetes.io/restartedAt"] = time.Now().Format(time.RFC3339)
        return runtime.Encode(scheme.Codecs.LegacyCodec(appsv1.SchemeGroupVersion), obj)
...

この config-reloader をこちら で示すように、 CronJob で定期的に実行することで、authenticator server Pod を再起動し、kubeconfig の再生成を定期的に行っています。

K8s Custom Controller

続いて K8s Custom Controller についてですが、構築・運用していたシステムでは、cks-operator と imperator の 2 種類の K8s Custom Controller を実装して運用の自動化を行なっていました。

cks-operator

cks-operator は以下で公開していますが、汎用的に使用できるように実装はされていないので、参考程度にしてください。

gitlab.com

cks-operator で使用する MachineLearning Custom Resource (ML CR) はこちら や以下で示すようになっています。

...
// MachineLearningSpec defines the desired state of MachineLearning
type MachineLearningSpec struct {

    // +kubebuilder:validation:Required
    MachineLearningApps []MachineLearningApps `json:"apps"`
}

// MachineLearningApps is apps for machine learning
type MachineLearningApps struct {

    // +kubebuilder:validation:Required
    AppName string `json:"name"`

    // +kubebuilder:validation:Required
    Namespace string `json:"namespace"`

    // +kubebuilder:validation:Required
    MLAppSpec MLAppSpec `json:"spec"`
}

// MLAppSpec is detail spec of apps for machine learning
type MLAppSpec struct {

    // +kubebuilder:validation:Required
    // +kubebuilder:validation:Enum=jupyterlab;codeserver
    Type ApplicationType `json:"type"`

    // +kubebuilder:validation:Required
    Machine MachineSpec `json:"machine"`

    // +kubebuilder:validation:Required
    AppImage string `json:"appImage"`

    // +optional
    ImagePullSecrets []string `json:"imagePullSecrets,omitempty"`
}

type ApplicationType string

const (
    CodeServer ApplicationType = "codeserver"
    JupyterLab ApplicationType = "jupyterlab"
)

func (t ApplicationType) String() string {
    return string(t)
}

type MachineSpec struct {

    // +kubebuilder:validation:Required
    Type string `json:"type"`

    // +kubebuilder:validation:Required
    Group string `json:"group"`
}

// MachineLearningStatus defines the observed state of MachineLearning
type MachineLearningStatus struct {
    Conditions []metav1.Condition `json:"conditions,omitempty"`
}
...

ML CR が作成されると、 cks-operator 内の machinelearning-controller は ML CR に定義された

  1. Pod を起動するための Namespace
  2. 後述する Imperator で定義された MachineType （計算機スペック）
  3. 使用するアプリケーションの種類 (codeserver もしくは jupyterlab)
  4. 使用するコンテナイメージ
  5. etc...

をもとに

  1. Ingress の設定
  2. NFS ストレージ・Rook/Ceph から事前に払い出された PV の Pod へのマウント
  3. Deployment の作成

を行う仕組みになっており、Web IDE や Jupyter を起動するために必要な種々の作業を自動化しています。

また前述のクライアントコマンドでは、この ML CR を操作することで Web IDE や Jupyter の起動や削除を行なっています。

imperator

imperator （インペラトル） は以下で公開しており、汎用的に使用できるように実装されているので、ぜひ使用してみてください。

github.com

f:id:tenzen_hgst:20220318221920p:plain
imperator-architecture

運用していたシステムでは、種々の事情で

  1. 様々な物理スペックのサーバを使用
  2. 研究室メンバー（ユーザ）全員が同時に使用できるほどの計算機スペックがない

という問題点があり、ユーザが同時の使用できる計算機リソースを制限した上で、できるだけサーバの物理スペックを使い切れるように Pod をスケジューリングし、どの計算リソース、とりわけどの GPU 種別がどれくらい空いているかをユーザに通知する必要がありました。

これらの課題を解決するために Imperator を開発しました。Imperator の思想はとても単純で、リソース制限がかかった何もしない Pod によって K8s 上のリソースをロックし、本物の Pod がスケジュールされた時に、偽物の Pod を削除して本物の Pod に計算機資源を譲る物となっています。特性上、GPU や FPGA などのデバイスをロックすることが主目的であるため、そのようなデバイスを使用しない環境などではあまり効果はないと思われます。

似たような仕組みとして Public Cloud 環境などで、Cluster Autoscaler を使用する際、負荷のスパイクにも対応するため、ダミーの Pod を配置することがあると思いますが、そこから着想を得ました。

仕組み

Imperator は以下のような Machine CR が定義されると動作するようになっています。

apiVersion: imperator.tenzen-y.io/v1alpha1
kind: Machine
metadata:
  name: machine-with-gpus
  labels:
    imperator.tenzen-y.io/machine-group: gpu-machines
spec:
  nodePool:
    - name: kind-control-plane
      mode: ready
      taint: false
      machineType:
        - name: compute-xmedium
  machineTypes:
    - name: compute-xmedium
      spec:
        cpu: 4000m
        memory: 12Gi
        gpu:
          type: nvidia.com/gpu
          num: 1
          family: ampere
      available: 1

Imperator は主に

  1. Machine Controller
  2. MachineNodePool Controller

の 2 種の Custom Controller および、Admission Mutating Webhook である Pod Resource Injector から構成されており、管理者、ユーザそれぞれに対する Imperator の動作は以下のようになっています。

f:id:tenzen_hgst:20220318222133p:plain
管理者が Machine CR を作成してからの動き

f:id:tenzen_hgst:20220318222203p:plain
ユーザが Pod を作成してからの動き

まず Machine Controller は、 Machine CR の spec.machineTypes で定義された計算機スペック(MachineType) を満たす偽物の Pod (Reservation Pod) と本物の Pod (Guest Pod) の状態や個数を管理することを主な役割とし、どの Node がどの machineType で使用可能かを管理する MachineNodePool Controller 用の MachineNodePool CR の管理も行います。

次に MachineNodePool Controller は、MachineNodePool CR の定義に従ってどの Node にどの MachineType をスケジュール可能にするかを決定し、Node に Annotation や Label、Taint を付与する役割をもった Custom Controller です。

最後に Pod Resource Injector ですが、これは本物の Pod (Guest Pod) の作成要求が来た際に、Guest Pod を正しい Node へスケジューリングできるように Guest Pod の resource.request や resource.limit、NodeAffinity などにパラメータを注入する役割を持っています。

本記事ではこれ以上詳細な説明はしませんが、より詳しい情報は以下に記述してあるので、参照してください。

github.com

まとめ

今回は構築したシステムで使用するために開発したソフトウェアを中心に紹介しました。 本記事では全体的な紹介にとどめたため各ソフトウェアを詳しく掘り下げたりしていませんが、掘り下げて欲しい物などがあれば、Twitter で言っていただけると紹介記事を書くかもしれません。

次回は ”運用編” として、使用していたミドルウェアをはじめ、運用で使用していた Ansible PlayBook、CI でのマニフェストへのテストなどについて紹介したいと思います。

最後まで読んでいただきありがとうございました。