Fix bug that can't distinguish standalone job (#828)

* Fix bug that can't distinguish standalone job

* Fix bug that update priority for single job failed

* filter pod with labels owner

Author: D0m021ng

pkg/common/schema/job.go

@@ -133,10 +133,11 @@ const (
 	PriorityClassHigh     = "high"
 	PriorityClassVeryHigh = "very-high"
 
-	JobOwnerLabel = "owner"
-	JobOwnerValue = "paddleflow"
-	JobIDLabel    = "paddleflow-job-id"
-	JobTTLSeconds = "padleflow/job-ttl-seconds"
+	JobOwnerLabel     = "owner"
+	JobOwnerValue     = "paddleflow"
+	JobIDLabel        = "paddleflow-job-id"
+	JobTTLSeconds     = "padleflow/job-ttl-seconds"
+	JobLabelFramework = "paddleflow-job-framework"
 
 	VolcanoJobNameLabel  = "volcano.sh/job-name"
 	QueueLabelKey        = "volcano.sh/queue-name"

pkg/job/runtime/kubernetes/controller/job_sync.go

@@ -359,12 +359,18 @@ func (j *JobSync) syncTaskStatus(taskSyncInfo *TaskSyncInfo) error {
 
 func responsibleForJob(obj interface{}) bool {
 	job := obj.(*unstructured.Unstructured)
+	jobName := job.GetName()
 	labels := job.GetLabels()
+	gvk := job.GroupVersionKind()
 	if labels != nil && labels[commonschema.JobOwnerLabel] == commonschema.JobOwnerValue {
-		log.Debugf("responsible for handle job. jobName:[%s]", job.GetName())
+		if gvk == k8s.PodGVK && labels[commonschema.JobLabelFramework] != string(commonschema.FrameworkStandalone) {
+			log.Debugf("job %s is not single job", jobName)
+			return false
+		}
+		log.Debugf("responsible for handle job. jobName:[%s]", jobName)
 		return true
 	}
-	log.Debugf("responsible for skip job. jobName:[%s]", job.GetName())
+	log.Debugf("responsible for skip job. jobName:[%s]", jobName)
 	return false
 }
 

pkg/job/runtime/kubernetes/controller/job_sync_handler.go

@@ -335,26 +335,26 @@ func getJobByTask(obj *unstructured.Unstructured) string {
 		log.Errorf("get job by task failed, obj is nil")
 		return ""
 	}
-	name := obj.GetName()
-	namespace := obj.GetNamespace()
 	labels := obj.GetLabels()
 	ownerReferences := obj.GetOwnerReferences()
 
-	if len(ownerReferences) == 0 {
-		// get job name for single job
-		if labels != nil && labels[schema.JobOwnerLabel] == schema.JobOwnerValue {
-			return name
-		} else {
-			log.Debugf("pod %s/%s not belong to paddlefow job, skip it.", namespace, name)
-			return ""
+	if labels == nil || labels[schema.JobOwnerLabel] != schema.JobOwnerValue {
+		return ""
+	}
+	// 1. get job name from ownerReferences, including workflow, PaddleJob
+	if len(ownerReferences) > 0 {
+		// get job name for distributed job
+		ownerReference := ownerReferences[0]
+		gvk := k8sschema.FromAPIVersionAndKind(ownerReference.APIVersion, ownerReference.Kind)
+		_, find := k8s.GVKJobStatusMap[gvk]
+		if find {
+			return ownerReference.Name
 		}
 	}
-	// get job name for distributed job
-	ownerReference := ownerReferences[0]
-	gvk := k8sschema.FromAPIVersionAndKind(ownerReference.APIVersion, ownerReference.Kind)
-	_, find := k8s.GVKJobStatusMap[gvk]
-	if !find {
-		return ""
+	// 2. get job name from pod labels
+	jobName, find := labels[schema.JobIDLabel]
+	if find {
+		return jobName
 	}
-	return ownerReference.Name
+	return ""
 }

pkg/job/runtime/kubernetes/executor/kube_builtin_job.go

@@ -63,7 +63,7 @@ func (j *KubeJob) setPodTemplateSpec(podSpec *corev1.PodTemplateSpec, task *sche
 	if podSpec == nil || task == nil {
 		return fmt.Errorf("podTemplateSpec or task is nil")
 	}
-	// TODO: set pod metadata
+	j.patchTaskMetadata(&podSpec.ObjectMeta, *task)
 	// set SchedulerName
 	podSpec.Spec.SchedulerName = config.GlobalServerConfig.Job.SchedulerName
 	// TODO: remove hard coded schedulerName when upstream package is fixed

pkg/job/runtime/kubernetes/executor/kubernetes_job.go

@@ -553,6 +553,7 @@ func (j *KubeJob) patchTaskMetadata(metadata *metav1.ObjectMeta, member schema.M
 	metadata.Annotations = j.appendAnnotationsIfAbsent(metadata.Annotations, member.Annotations)
 	metadata.Labels = j.appendLabelsIfAbsent(metadata.Labels, member.Labels)
 	metadata.Labels[schema.JobIDLabel] = j.ID
+	metadata.Labels[schema.JobOwnerLabel] = schema.JobOwnerValue
 }
 
 func (j *KubeJob) fillPodTemplateSpec(pod *corev1.PodTemplateSpec, member schema.Member) error {
@@ -647,6 +648,9 @@ func GetPodGroupName(jobID string) string {
 		if anno != nil {
 			pgName = anno[schedulingv1beta1.KubeGroupNameAnnotationKey]
 		}
+		if pgName == "" {
+			pgName = fmt.Sprintf("podgroup-%s", jobObj.GetUID())
+		}
 	default:
 		log.Warningf("the framework[%s] of job is not supported", job.Framework)
 		pgName = jobID

pkg/job/runtime/kubernetes/executor/pytorch_job.go

@@ -108,12 +108,25 @@ func (pj *PyTorchJob) validateCustomYaml(torchJobSpec *pytorchv1.PyTorchJobSpec)
 
 // customPyTorchJobSpec set custom PyTorchJob Spec
 func (pj *PyTorchJob) customPyTorchJobSpec(torchJobSpec *pytorchv1.PyTorchJobSpec) error {
+	if torchJobSpec == nil || torchJobSpec.PyTorchReplicaSpecs == nil {
+		err := fmt.Errorf("build custom %s failed, PyTorchJobSpec or PyTorchReplicaSpecs is nil", pj.String())
+		log.Errorf("%v", err)
+		return err
+	}
 	log.Debugf("patch %s spec:%#v", pj.String(), torchJobSpec)
 	err := pj.validateCustomYaml(torchJobSpec)
 	if err != nil {
 		return err
 	}
-	// TODO: patch pytorch job from user
+	// patch metadata
+	ps, find := torchJobSpec.PyTorchReplicaSpecs[pytorchv1.PyTorchReplicaTypeMaster]
+	if find && ps != nil {
+		pj.patchTaskMetadata(&ps.Template.ObjectMeta, schema.Member{})
+	}
+	worker, find := torchJobSpec.PyTorchReplicaSpecs[pytorchv1.PyTorchReplicaTypeWorker]
+	if find && worker != nil {
+		pj.patchTaskMetadata(&worker.Template.ObjectMeta, schema.Member{})
+	}
 	// check RunPolicy
 	return nil
 }

pkg/job/runtime/kubernetes/executor/ray_job.go

@@ -240,11 +240,21 @@ func (j *RayJob) validateCustomYaml(rayJobSpec *rayV1alpha1.RayJobSpec) error {
 
 // customPyTorchJobSpec set custom PyTorchJob Spec
 func (j *RayJob) customRayJobSpec(rayJobSpec *rayV1alpha1.RayJobSpec) error {
+	if rayJobSpec == nil {
+		err := fmt.Errorf("build custom %s failed, RayJobSpec is nil", j.String())
+		log.Errorf("%v", err)
+		return err
+	}
 	log.Debugf("patch %s spec:%#v", j.String(), rayJobSpec)
 	err := j.validateCustomYaml(rayJobSpec)
 	if err != nil {
 		return err
 	}
+	// patch metadata
+	j.patchTaskMetadata(&rayJobSpec.RayClusterSpec.HeadGroupSpec.Template.ObjectMeta, schema.Member{})
+	for i := range rayJobSpec.RayClusterSpec.WorkerGroupSpecs {
+		j.patchTaskMetadata(&rayJobSpec.RayClusterSpec.WorkerGroupSpecs[i].Template.ObjectMeta, schema.Member{})
+	}
 	// TODO: patch ray job from user
 	// check RunPolicy
 	return nil

pkg/job/runtime/kubernetes/executor/singlejob.go

@@ -80,7 +80,12 @@ func (sp *SingleJob) patchSinglePodVariable(pod *v1.Pod, jobID string) error {
 		priorityClass := sp.getPriorityClass()
 		pod.Spec.PriorityClassName = priorityClass
 	}
-	sp.fillPodSpec(&pod.Spec, nil)
+	// set single job mark
+	pod.Labels[schema.JobLabelFramework] = string(schema.FrameworkStandalone)
+	if err := sp.fillPodSpec(&pod.Spec, nil); err != nil {
+		log.Errorf("single job fillPodSpec failed, err: %v", err)
+		return err
+	}
 
 	// file container
 	if err := sp.fillContainersInPod(pod); err != nil {

pkg/job/runtime/kubernetes/executor/tensorflow_job.go

@@ -114,11 +114,24 @@ func (j *TFJob) validateCustomYaml(tfSpec *tfv1.TFJobSpec) error {
 // customPyTorchJobSpec set custom PyTorchJob Spec
 func (j *TFJob) customTFJobSpec(tfJobSpec *tfv1.TFJobSpec) error {
 	log.Debugf("patch %s spec:%#v", j.String(), tfJobSpec)
+	if tfJobSpec == nil || tfJobSpec.TFReplicaSpecs == nil {
+		err := fmt.Errorf("build custom %s failed, TFJobSpec or TFReplicaSpecs is nil", j.String())
+		log.Errorf("%v", err)
+		return err
+	}
 	err := j.validateCustomYaml(tfJobSpec)
 	if err != nil {
 		return err
 	}
-	// TODO: patch tf job from user
+	// patch metadata
+	ps, find := tfJobSpec.TFReplicaSpecs[tfv1.TFReplicaTypePS]
+	if find && ps != nil {
+		j.patchTaskMetadata(&ps.Template.ObjectMeta, schema.Member{})
+	}
+	worker, find := tfJobSpec.TFReplicaSpecs[tfv1.TFReplicaTypeWorker]
+	if find && worker != nil {
+		j.patchTaskMetadata(&worker.Template.ObjectMeta, schema.Member{})
+	}
 	// check RunPolicy
 	return nil
 }