Kubeflow--TFJob实现机制分析

前言：

虽然Kubeflow提供了一大堆组件，涵盖了机器学习的方方面面，但模型训练肯定是Kubeflow最重要的功能。

在深入探讨tfjob的实现机制 分析其源码前，强烈建议先阅读的奇虎大神的文章：炼丹师的工程修养之四： TensorFlow的分布式训练和K8S - 知乎

笔者也对该文章进行了转载，并解决了里面引用的一些github 仓库地址过时的问题：TensorFlow的分布式训练和K8S_chenxy02的博客-CSDN博客

TFJob的本质

KubeFlow针对各种各样的机器学习框架提供了训练的能力。方法是定义了各种各样的training-operator，这些Operator的本质，是K8S的CRD （CustomResourceDefinition）。本文主要以TFJob为例进行分析讲解。

Custom Resource 是什么？ 首先，Resource是K8S中的一个基础概念，一个Resource是多个Object的集合，Object又是一个基础概念，Pod，Controller，Job，Service都是Object。所以多个Pod集合就是Pods Resource。Custom Resource就是K8S提供的扩展API可以让用户定义自己Resource。

一句话，TFJob Operator就是开源社区基于k8s提供的扩展API，提供了TensorFlow的训练能力，从名字也能看出来，这个实现是类似Job控制器的一种方式。

为什么要有TFJob

首先，我们了解一下没有TFJob，我们怎么在K8S上做Tensorflow的分布式训练（还是建议先阅读“前言”所述文章）。

由于Tensorflow已经对于分布式训练做了很好的支持，开发人员只需要将 TF_CONFIG 通过环境变量传递给相应的容器即可。甚至通过使用Strategy，可以自动读取环境变量的TF_CONFIG信息并正确解析用于分布式训练集群的配置，不需要开发人员在代码里做任何显性的判断。

# TF_CONFIG 示例
{
  "cluster": {
    "ps": [
      "host1:2222"
    ],
    "worker": [
      "host2:3333",
      "host3:3333",
      "host3:4444"
    ]
  },
  "task": {
    "index": 0,   # 索引
    "type": "ps"  # 指定该容器承当的角色
  }
}

如何准确的运行多个容器

我们需要在k8s上创建多个Job，针对多个Job分别编写相应的配置文件，通过Service机制解决IP不确定的问题，每个Job的TF_CONFIG也要修改。对于MultiWorker架构可能仅仅只需要改index一个地方，如果是PS架构，配置起来可能还要更麻烦一些。

如何快速准确的运行多个容器

上述提到的Job的配置文件，结构一致，内容大部分是相同的。我们完全可以把变化的部分抽象出来作为参数，使用jinja之类的工具在运行前动态渲染生成出一堆配置文件，然后再扔给K8S就好了。这就是TFJob的精髓了。

TFJob Operator源码分析

TFJob的核心就是为了每一个Worker编写一个Job的配置文件，然后一个一个的提交。这些配置文件也完全可以通过代码来自动生成以简化操作。而KubeFlow或者说TFJob Operator做的更“K8S”一些。他用的是K8S提供的Costom Resource API来实现的。

上述我们提到TFJob的本质是CRD，对于编写CRD的知识可参考: 深入解析 Kubebuilder：让编写 CRD 变得更简单 - 知乎  粗略总结一下，就是Kubebuilder 作为脚手架工具已经做了很多，我们主要实现 Reconcile 方法即可。

TFJob的代码仓库地址为：https://github.com/kubeflow/training-operator  至于怎么集成进Kubeflow，怎么找到这个仓库地址，可参考：Kubeflow安装及代码架构解读_chenxy02的博客-CSDN博客

接下来分析一下TFJob Operator的源代码（以笔者项目中所用的v1.3版本为例），用到了大量的k8s的API，我们只把重要的地方摘出来：

文件：cmd/training-operator.v1/main.go (入口文件)

    // 接收传入的参数
	flag.Parse()

	ctrl.SetLogger(zap.New(zap.UseFlagOptions(&opts)))

	// 初始化一个Manager,Kubebuilder的核心组件
    // 负责运行所有的 Controllers,初始化共享caches(包含 listAndWatch 功能),初始化 clients 用于与 Api Server 通信。
	mgr, err := ctrl.NewManager(ctrl.GetConfigOrDie(), ctrl.Options{
		Scheme:                 scheme,
        ......
		Namespace:              namespace,
	})
	if err != nil {
		setupLog.Error(err, "unable to start manager")
		os.Exit(1)
	}

	// TODO: We need a general manager. all rest reconciler addsToManager
	// Based on the user configuration, we start different controllers
	if enabledSchemes.Empty() {
		enabledSchemes.FillAll()
	}
	for _, s := range enabledSchemes {
		// 得到相应的 SetupWithManager方法
		setupFunc, supported := controllerv1.SupportedSchemeReconciler[s]
		if !supported {
			setupLog.Error(fmt.Errorf("cannot find %s in supportedSchemeReconciler", s),
				"scheme not supported", "scheme", s)
			os.Exit(1)
		}
		// 调用 SetupWithManager 方法传入 Manager
		if err = setupFunc(mgr, enableGangScheduling); err != nil {
			setupLog.Error(err, "unable to create controller", "controller", s)
			os.Exit(1)
		}
	}

（重点关注笔者加中文注释的地方，以下同）

文件：pkg/controller.v1/register_controller.go

type ReconcilerSetupFunc func(manager manager.Manager, enableGangScheduling bool) error

var SupportedSchemeReconciler = map[string]ReconcilerSetupFunc{
	tensorflowv1.Kind: func(mgr manager.Manager, enableGangScheduling bool) error {
		return tensorflowcontroller.NewReconciler(mgr, enableGangScheduling).SetupWithManager(mgr)
	},
	pytorchv1.Kind: func(mgr manager.Manager, enableGangScheduling bool) error {
		return pytorchcontroller.NewReconciler(mgr, enableGangScheduling).SetupWithManager(mgr)
	},
	mxnetv1.Kind: func(mgr manager.Manager, enableGangScheduling bool) error {
		return mxnetcontroller.NewReconciler(mgr, enableGangScheduling).SetupWithManager(mgr)
	},
	xgboostv1.Kind: func(mgr manager.Manager, enableGangScheduling bool) error {
		return xgboostcontroller.NewReconciler(mgr, enableGangScheduling).SetupWithManager(mgr)
	},
	mpiv1.Kind: func(mgr manager.Manager, enableGangScheduling bool) error {
		return mpicontroller.NewReconciler(mgr, enableGangScheduling).SetupWithManager(mgr)
	},
}

在这个文件，我们可以看到目前training-operator所支持的分布式训练框架。

文件地址：pkg/controller.v1/tensorflow/tfjob_controller.go

// 完成 Client等组件的创建，返回TFJobReconciler
func NewReconciler(mgr manager.Manager, enableGangScheduling bool) *TFJobReconciler {
	r := &TFJobReconciler{
		Client:   mgr.GetClient(),
		Scheme:   mgr.GetScheme(),
		recorder: mgr.GetEventRecorderFor(controllerName),
		Log:      log.Log,
	}

	cfg := mgr.GetConfig()
	kubeClientSet := kubeclientset.NewForConfigOrDie(cfg)
	volcanoClientSet := volcanoclient.NewForConfigOrDie(cfg)
	sharedInformers := informers.NewSharedInformerFactory(kubeClientSet, 0)
	priorityClassInformer := sharedInformers.Scheduling().V1beta1().PriorityClasses()

	r.JobController = common.JobController{
		Controller:                  r,
		Expectations:                expectation.NewControllerExpectations(),
		Config:                      common.JobControllerConfiguration{EnableGangScheduling: enableGangScheduling},
		WorkQueue:                   &util.FakeWorkQueue{},
		Recorder:                    r.recorder,
		KubeClientSet:               kubeClientSet,
		VolcanoClientSet:            volcanoClientSet,
        ……
	}

	return r
}

// TFJobReconciler reconciles a TFJob object
type TFJobReconciler struct {
	common.JobController
	client.Client
	Scheme   *runtime.Scheme
	recorder record.EventRecorder
	Log      logr.Logger
}

// SetupWithManager sets up the controller with the Manager.
func (r *TFJobReconciler) SetupWithManager(mgr ctrl.Manager) error {
	// 传入Reconciler，初始化controller
	c, err := controller.New(r.ControllerName(), mgr, controller.Options{
		Reconciler: r,
	})
    ……

	// Controller不断地查询队列，如果有变更消息则触发到我们自定义的
    // Reconcile逻辑(包括创建TFjob,pod,service等资源)
	// using onOwnerCreateFunc is easier to set defaults
	if err = c.Watch(&source.Kind{Type: &tfv1.TFJob{}}, &handler.EnqueueRequestForObject{},
		predicate.Funcs{CreateFunc: r.onOwnerCreateFunc()},
	); err != nil {
		return err
	}

	// inject watching for job related pod
	if err = c.Watch(&source.Kind{Type: &corev1.Pod{}}, &handler.EnqueueRequestForOwner{
		IsController: true,
		OwnerType:    &tfv1.TFJob{},
	}, predicate.Funcs{
		CreateFunc: util.OnDependentCreateFunc(r.Expectations),
		UpdateFunc: util.OnDependentUpdateFunc(&r.JobController),
		DeleteFunc: util.OnDependentDeleteFunc(r.Expectations),
	}); err != nil {
		return err
	}

	// inject watching for job related service
	if err = c.Watch(&source.Kind{Type: &corev1.Service{}}, &handler.EnqueueRequestForOwner{
		IsController: true,
		OwnerType:    &tfv1.TFJob{},
	}, predicate.Funcs{
		CreateFunc: util.OnDependentCreateFunc(r.Expectations),
		UpdateFunc: util.OnDependentUpdateFunc(&r.JobController),
		DeleteFunc: util.OnDependentDeleteFunc(r.Expectations),
	}); err != nil {
		return err
	}

	return nil
}

 以上所摘的两个方法，其实做的最重要的事情就是：初始化Controller，启动goroutime不断查询队列，如果有变更消息则触发到我们自定义的 Reconcile 逻辑。Reconcile方法如下：

// 自定义的 Reconcile 逻辑
// Reconcile is part of the main kubernetes reconciliation loop which aims to
// move the current state of the cluster closer to the desired state.
func (r *TFJobReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
	_ = log.FromContext(ctx)
	logger := r.Log.WithValues(tensorflowv1.Singular, req.NamespacedName)

	tfjob := &tensorflowv1.TFJob{}
	err := r.Get(ctx, req.NamespacedName, tfjob)
	if err != nil {
		logger.Info(err.Error(), "unable to fetch TFJob", req.NamespacedName.String())
		return ctrl.Result{}, client.IgnoreNotFound(err)
	}

	if err = validation.ValidateV1TFJobSpec(&tfjob.Spec); err != nil {
		logger.Info(err.Error(), "TFJob failed validation", req.NamespacedName.String())
	}

	// Check if reconciliation is needed
	jobKey, err := common.KeyFunc(tfjob)
	if err != nil {
		utilruntime.HandleError(fmt.Errorf("couldn't get jobKey for job object %#v: %v", tfjob, err))
	}

	// 获取Spec信息, 通过tfjob类型的实例传入
	replicaTypes := util.GetReplicaTypes(tfjob.Spec.TFReplicaSpecs)
	needReconcile := util.SatisfiedExpectations(r.Expectations, jobKey, replicaTypes)

	if !needReconcile || tfjob.GetDeletionTimestamp() != nil {
		logger.Info("reconcile cancelled, job does not need to do reconcile or has been deleted",
			"sync", needReconcile, "deleted", tfjob.GetDeletionTimestamp() != nil)
		return ctrl.Result{}, nil
	}

	// Set default priorities to tfjob
	r.Scheme.Default(tfjob)

	// Use common to reconcile the job related pod and service
	// 执行Reconcile逻辑！！！
	err = r.ReconcileJobs(tfjob, tfjob.Spec.TFReplicaSpecs, tfjob.Status, &tfjob.Spec.RunPolicy)
	if err != nil {
		logrus.Warnf("Reconcile Tensorflow Job error %v", err)
		return ctrl.Result{}, err
	}

	return ctrl.Result{}, nil
}

最后，我们看一下上述方法引用的 "TFJob"类型是怎么定义

文件：pkg/apis/tensorflow/v1/types.go

// TFJob represents a TFJob resource.
type TFJob struct {
	// Standard Kubernetes type metadata.
	metav1.TypeMeta `json:",inline"`

	// +optional
	metav1.ObjectMeta `json:"metadata,omitempty"`

	// Specification of the desired state of the TFJob.
	// +optional
	Spec TFJobSpec `json:"spec,omitempty"`

	// Most recently observed status of the TFJob.
	Status commonv1.JobStatus `json:"status,omitempty"`
}

// TFJobSpec is a desired state description of the TFJob.
type TFJobSpec struct {
	// RunPolicy encapsulates various runtime policies of the distributed training
	// job, for example how to clean up resources and how long the job can stay
	// active.
	//+kubebuilder:validation:Optional
	RunPolicy commonv1.RunPolicy `json:"runPolicy"`

	// SuccessPolicy defines the policy to mark the TFJob as succeeded.
	// Default to "", using the default rules.
	// +optional
	SuccessPolicy *SuccessPolicy `json:"successPolicy,omitempty"`

	// A map of TFReplicaType (type) to ReplicaSpec (value). Specifies the TF cluster configuration.
	// For example,
	//   {
	//     "PS": ReplicaSpec,
	//     "Worker": ReplicaSpec,
	//   }
	TFReplicaSpecs map[commonv1.ReplicaType]*commonv1.ReplicaSpec `json:"tfReplicaSpecs"`

	// A switch to enable dynamic worker
	EnableDynamicWorker bool `json:"enableDynamicWorker,omitempty"`
}

// TFReplicaType is the type for TFReplica. Can be one of: "Chief"/"Master" (semantically equivalent),
// "Worker", "PS", or "Evaluator".

const (
	// TFReplicaTypePS is the type for parameter servers of distributed TensorFlow.
	TFReplicaTypePS commonv1.ReplicaType = "PS"

	// TFReplicaTypeWorker is the type for workers of distributed TensorFlow.
	// This is also used for non-distributed TensorFlow.
	TFReplicaTypeWorker commonv1.ReplicaType = "Worker"

	// TFReplicaTypeChief is the type for chief worker of distributed TensorFlow.
	// If there is "chief" replica type, it's the "chief worker".
	// Else, worker:0 is the chief worker.
	TFReplicaTypeChief commonv1.ReplicaType = "Chief"

	// TFReplicaTypeMaster is the type for master worker of distributed TensorFlow.
	// This is similar to chief, and kept just for backwards compatibility.
	TFReplicaTypeMaster commonv1.ReplicaType = "Master"

	// TFReplicaTypeEval is the type for evaluation replica in TensorFlow.
	TFReplicaTypeEval commonv1.ReplicaType = "Evaluator"
)

Tensorflow已经对于分布式训练做了很好的支持，开发人员只需要将 TF_CONFIG 通过环境变量传递给相应的容器。TF_CONFIG 就是通过TFjob类型的实例传入 最终K8S创建Pod时所使用的配置文件 。

有兴趣的同学，可以看一下PyTorchJob的定义（pkg/apis/pytorch/v1/types.go）。会发现里面的字段是大一样的。那是因为PyTorch与Tensorflow对于分布式训练有着不同的支持机制。

总结：

TFJob的核心就是帮助AI开发者能快速准确的创建多个资源，运行多个容器。开发者通过声明式API的方式，即可以快速的配置及创建分布式训练集群。

想要理解 Kubeflow中 training-operator的实现机制，需要了解CRM的实现方式、Kubernetes -APIs 、Kubebuilder的使用，同时对于深度学习框架对分布式支持有一定的认识。

 参考网址：

https://github.com/kubeflow/training-operator

TensorFlow Training (TFJob) | Kubeflow  

（目前网上 找到源码分析，所用版本都比较老，与现版本代码对不上，仅供参考)

炼丹师的工程修养之五：KubeFlow介绍和源码分析 - 知乎

Kubeflow@TFJob源码分析 - 知乎