Add task type annotation for pods when EnableGangScheduling is true. (#…

…1268) Signed-off-by: jiangkaihua <[email protected]>
kubeflow · Apr 2, 2021 · 3d9e7c8 · 3d9e7c8
1 parent 8d179f7
commit 3d9e7c8
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/pkg/controller.v1/tensorflow/controller.go b/pkg/controller.v1/tensorflow/controller.go
@@ -55,6 +55,8 @@ const (
 	labelGroupName      = "group-name"
 	// Deprecated label for backwards compatibility. Has to be removed
 	labelTFJobName = "tf-job-name"
+	// volcanoTaskSpecKey task spec key used in pod annotation when EnableGangScheduling is true
+	volcanoTaskSpecKey = "volcano.sh/task-spec"
 )
 
 var (

diff --git a/pkg/controller.v1/tensorflow/pod.go b/pkg/controller.v1/tensorflow/pod.go
@@ -228,6 +228,7 @@ func (tc *TFController) createNewPod(tfjob *tfv1.TFJob, rt, index string, spec *
 			podTemplate.Annotations = map[string]string{}
 		}
 		podTemplate.Annotations[gangSchedulingPodGroupAnnotation] = tfjob.GetName()
+		podTemplate.Annotations[volcanoTaskSpecKey] = rt
 	}
 
 	err = tc.PodControl.CreatePodsWithControllerRef(tfjob.Namespace, podTemplate, tfjob, controllerRef)