Add a field SchedulerName to TFJob for specifying a scheduler (#408)

This commit adds a new field SchedulerName to the definition of TFJob.
The purpose of the field is specifying the scheduler name of the pods
created by tf-operator and let the scheduler (which wouldn't be the
default scheduler) handle them. It would be convenient for letting
kube-batchd (a component of kube-arbitrator) handle the pods.

pkg/apis/tensorflow/v1alpha1/types.go

@@ -56,6 +56,9 @@ type TFJobSpec struct {
 
 	// TerminationPolicy specifies the condition that the tfjob should be considered finished.
 	TerminationPolicy *TerminationPolicySpec `json:"terminationPolicy,omitempty"`
+
+	// SchedulerName specifies the name of scheduler which should handle the TFJob
+	SchedulerName string `json:"schedulerName,omitempty"`
 }
 
 type TerminationPolicySpec struct {

pkg/trainer/replicas.go

@@ -171,6 +171,8 @@ func (s *TFReplicaSet) CreatePodWithIndex(index int32) (*v1.Pod, error) {
 		Spec: *s.Spec.Template.Spec.DeepCopy(),
 	}
 
+	pod.Spec.SchedulerName = s.Job.SchedulerName()
+
 	// Configure the TFCONFIG environment variable.
 	tfConfig := TFConfig{
 		Cluster: s.Job.ClusterSpec(),

pkg/trainer/replicas_test.go

@@ -18,6 +18,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"reflect"
+	"strings"
 	"testing"
 	"time"
 
@@ -44,6 +45,8 @@ var (
 func TestTFReplicaSet(t *testing.T) {
 	clientSet := fake.NewSimpleClientset()
 
+	testSchedulerName := "test-scheduler"
+
 	jobSpec := &tfv1alpha1.TFJob{
 		ObjectMeta: meta_v1.ObjectMeta{
 			Name: "some-job",
@@ -67,6 +70,7 @@ func TestTFReplicaSet(t *testing.T) {
 					TFReplicaType: tfv1alpha1.PS,
 				},
 			},
+			SchedulerName: testSchedulerName,
 		},
 	}
 
@@ -169,6 +173,10 @@ func TestTFReplicaSet(t *testing.T) {
 			t.Fatalf("Expected 1 environment variable got %v", len(c.Env))
 		}
 
+		if strings.Compare(p.Spec.SchedulerName, testSchedulerName) != 0 {
+			t.Fatalf("p.Spec.Template.Spec.SchedulerName; Got %v; want %v", p.Spec.SchedulerName, testSchedulerName)
+		}
+
 		actualTFConfig := &TFConfig{}
 		if err := json.Unmarshal([]byte(c.Env[0].Value), actualTFConfig); err != nil {
 			t.Fatalf("Could not unmarshal TFConfig %v", err)

pkg/trainer/training.go

@@ -412,3 +412,7 @@ func (j *TrainingJob) name() string {
 func (j *TrainingJob) fullname() string {
 	return j.job.ObjectMeta.GetNamespace() + ":" + j.job.ObjectMeta.GetName()
 }
+
+func (j *TrainingJob) SchedulerName() string {
+	return j.job.Spec.SchedulerName
+}