support queued before scheduled (#708)

kubeflow · Nov 15, 2021 · 5a47651 · 5a47651
1 parent eb2933b
commit 5a47651
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 0 deletions.
diff --git a/pkg/apis/types/submit.go b/pkg/apis/types/submit.go
@@ -73,6 +73,9 @@ type CommonSubmitArgs struct {
 	// EnableRDMA enable rdma or not,match option --rdma
 	EnableRDMA bool `yaml:"enableRDMA"`
 
+	// EnableQueue enables the feature to queue jobs after they are scheduled.
+	EnableQueue bool `yaml:"enableQueue"`
+
 	// UseENI defines using eni or not
 	UseENI bool `yaml:"useENI"`
 

diff --git a/pkg/apis/types/training.go b/pkg/apis/types/training.go
@@ -131,6 +131,8 @@ type TrainingJobInfo struct {
 type TrainingJobStatus string
 
 const (
+	// TrainingJobQueuing means the job is queuing
+	TrainingJobQueuing TrainingJobStatus = "QUEUING"
 	// TrainingJobPending means the job is pending
 	TrainingJobPending TrainingJobStatus = "PENDING"
 	// TrainingJobRunning means the job is running

diff --git a/pkg/argsbuilder/const.go b/pkg/argsbuilder/const.go
@@ -7,4 +7,6 @@ const (
 	gangSchdName = "kube-batch"
 
 	aliyunENIAnnotation = "k8s.aliyun.com/eni"
+
+	jobSuspend = "scheduling.x-k8s.io/suspend"
 )
diff --git a/pkg/argsbuilder/submit.go b/pkg/argsbuilder/submit.go
@@ -120,6 +120,8 @@ func (s *SubmitArgsBuilder) AddCommandFlags(command *cobra.Command) {
 	command.Flags().BoolVar(&s.args.Coscheduling, "gang", false, "enable gang scheduling")
 	// use priority
 	command.Flags().StringVarP(&s.args.PriorityClassName, "priority", "p", "", "priority class name")
+	// enable Queue
+	command.Flags().BoolVar(&s.args.EnableQueue, "queue", false, "enables the feature to queue jobs after they are scheduled (Kube-queue needs to be pre-installed https://github.com/kube-queue/kube-queue)")
 	// add option --toleration,its' value will be get from viper
 	command.Flags().StringArrayVar(&tolerations, "toleration", []string{}, `tolerate some k8s nodes with taints,usage: "--toleration taint-key" or "--toleration all" `)
 	// add option --selector,its' value will be get from viper
@@ -174,6 +176,9 @@ func (s *SubmitArgsBuilder) Build() error {
 	if err := s.setAnnotations(); err != nil {
 		return err
 	}
+	if err := s.setQueue(); err != nil {
+		return err
+	}
 	if err := s.setLabels(); err != nil {
 		return err
 	}
@@ -323,6 +328,17 @@ func (s *SubmitArgsBuilder) setAnnotations() error {
 	return nil
 }
 
+// setQueue is used to add annotation for suspend status
+func (s *SubmitArgsBuilder) setQueue() error {
+	if s.args.EnableQueue {
+		if s.args.Annotations == nil {
+			s.args.Annotations = map[string]string{}
+		}
+		s.args.Annotations[jobSuspend] = "true"
+	}
+	return nil
+}
+
 // setAnnotations is used to handle option --annotation
 func (s *SubmitArgsBuilder) setLabels() error {
 	if s.args.Labels == nil {