【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-simplenamespace: kubeflow
spec:tfReplicaSpecs:Worker:replicas: 2restartPolicy: OnFailuretemplate:spec:containers:- name: tensorflowimage: kubeflow/tf-mnist-with-summaries:latestcommand:- "python"- "/var/tf_mnist/mnist_with_summaries.py"

资源PyTorchJob

apiVersion: "kubeflow.org/v1"
kind: PyTorchJob
metadata:name: pytorch-simple-001namespace: kubeflow
spec:pytorchReplicaSpecs:Master:replicas: 1restartPolicy: OnFailuretemplate:spec:containers:- name: pytorchimage: kubeflowkatib/pytorch-mnist:v1beta1-45c5727imagePullPolicy: Alwayscommand:- "python3"- "/opt/pytorch-mnist/mnist.py"- "--epochs=1"Worker:replicas: 2restartPolicy: OnFailuretemplate:spec:containers:- name: pytorchimage: kubeflowkatib/pytorch-mnist:v1beta1-45c5727imagePullPolicy: Alwayscommand:- "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"}'