set annotation automatically when EnableGangScheduling is set to true…

…[1032](kubeflow/training-operator#1032) (#181)
kubeflow · Jun 27, 2019 · 6aa39a4 · 6aa39a4
1 parent 213df9c
commit 6aa39a4
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/pkg/controller.v1/pytorch/pod.go b/pkg/controller.v1/pytorch/pod.go
@@ -222,6 +222,12 @@ func (pc *PyTorchController) createNewPod(job *pyv1.PyTorchJob, rtype pyv1.PyTor
 		} else {
 			podTemplate.Spec.SchedulerName = gangSchedulerName
 		}
+
+		if podTemplate.Annotations == nil {
+			podTemplate.Annotations = map[string]string{}
+		}
+		// we create the podGroup with the same name as the pyTorch job
+		podTemplate.Annotations["scheduling.k8s.io/group-name"] = job.Name
 	}
 
 	err = pc.PodControl.CreatePodsWithControllerRef(job.Namespace, podTemplate, job, controllerRef)