【kubeflow文档】Kubeflow Training Operator

What is Training Operator

        Training Operator是一个Kubernetes原生项目，用于对使用各种ML框架（如PyTorch、TensorFlow、XGBoost等）创建的机器学习（ML）模型进行微调和可扩展的分布式训练。

        用户可以将HuggingFace、DeepSpeed或Megatron LM等其他ML库与Training Operator集成，以协调他们在Kubernetes上的ML培训。

        Training Operator允许您使用Kubernetes工作负载，通过Kubernete自定义资源API或使用Training Operator Python SDK有效地训练您的大型模型。
        用户可以使用Training Operator和MPIJob运行高性能计算（HPC）任务，因为它支持在大量用于HPC的Kubernetes上运行消息传递接口（MPI）。

        Training Operator负责调度适当的Kubernetes工作负载，以针对不同的ML框架实现各种分布式训练策略。

---

Why Training Operator

        Training Operator解决了AI/ML生命周期中的模型训练和模型微调步骤，如图所示：

• Training Operator简化了运行分布式训练和微调的能力。

        用户可以使用training Operator提供的API和接口，轻松地将其模型训练从单机扩展到大规模分布式Kubernetes集群。

• Training Operator具有可扩展性和可移植性。

        用户可以在任何有Kubernetes集群的云上部署Training Operator，用户可以将自己用任何编程语言编写的ML框架与Training Operator集成。

• Training Operator与Kubernetes生态系统集成

        用户可以利用Kubernetes的高级调度技术，如Kueue、Volcano和YuniKorn与Training Operator，优化ML训练资源的成本耗费。

---

How Training Operator

        资源TFJob

apiVersion: "kubeflow.org/v1"
kind: TFJob
metadata:
  name: tfjob-simple
  namespace: kubeflow
spec:
  tfReplicaSpecs:
    Worker:
      replicas: 2
      restartPolicy: OnFailure
      template:
        spec:
          containers:
            - name: tensorflow
              image: kubeflow/tf-mnist-with-summaries:latest
              command:
                - "python"
                - "/var/tf_mnist/mnist_with_summaries.py"

        资源PyTorchJob 

apiVersion: "kubeflow.org/v1"
kind: PyTorchJob
metadata:
  name: pytorch-simple-001
  namespace: kubeflow
spec:
  pytorchReplicaSpecs:
    Master:
      replicas: 1
      restartPolicy: OnFailure
      template:
        spec:
          containers:
            - name: pytorch
              image: kubeflowkatib/pytorch-mnist:v1beta1-45c5727
              imagePullPolicy: Always
              command:
                - "python3"
                - "/opt/pytorch-mnist/mnist.py"
                - "--epochs=1"
    Worker:
      replicas: 2
      restartPolicy: OnFailure
      template:
        spec:
          containers:
            - name: pytorch
              image: kubeflowkatib/pytorch-mnist:v1beta1-45c5727
              imagePullPolicy: Always
              command:
                - "python3"
                - "/opt/pytorch-mnist/mnist.py"
                - "--epochs=1"

          将上述运行到k8s平台上，结果如下：

Training Operator源码实现

        Training Operator启动流程，这里直接使用了kubebuild创建，具体可以参考kubebuild创建项目（Introduction - The Kubebuilder Book），可以了解controller-runtime具体实现，manager controller等。

        这里分别调TFJobKind PyTorchJobKind MXJobKind XGBoostJobKind MPIJobKind PaddleJobKind 的SetupWithManager创建controller，将各种CRD的Reconcile添加到Controller中，这样当workqueue中存在事件时，会交由Reconcile处理。

        例如TFjob的Reconcile实现，最终要根据定义的spec设置成期待的状态，这里就是最核心的就是调用ReconcileJobs，再分别调用ReconcilePods，ReconcileServices。ReconcilePods这里会生成pod的spec，其中包括环境变量TF_CONFIG，例如：

        value: '{"cluster":{"worker":["tfjob-simple-worker-0.kubeflow.svc:2222","tfjob-simple-worker-1.kubeflow.svc:2222"]},"task":{"type":"worker","index":0},"environment":"cloud"}'

总结:

    Training Operator由kubebuiler创建项目，创建了TFJobKind PyTorchJobKind MXJobKind XGBoostJobKind MPIJobKind PaddleJobKind资源，调整成期待的状态，其实就是创建pod以及service。

参考：

  Overview | Kubeflow