feat(elastic): Implement main feature

Signed-off-by: Ce Gao <ce.gao@outlook.com>

pkg/apis/pytorch/v1/defaults.go

@@ -57,10 +57,32 @@ func setDefaultPort(spec *v1.PodSpec) {
 	}
 }
 
+func setElasticPolicy(job *PyTorchJob) {
+	if job.Spec.ElasticPolicy != nil {
+		if job.Spec.ElasticPolicy.MaxReplicas != nil &&
+			job.Spec.ElasticPolicy.MinReplicas != nil {
+			return
+		} else if job.Spec.ElasticPolicy.MaxReplicas != nil {
+			// Set MinRepliacs to elasticPolicy.MaxReplicas.
+			job.Spec.ElasticPolicy.MinReplicas = job.Spec.ElasticPolicy.MaxReplicas
+		} else if job.Spec.ElasticPolicy.MinReplicas != nil {
+			job.Spec.ElasticPolicy.MaxReplicas = job.Spec.ElasticPolicy.MinReplicas
+		} else {
+			workerReplicas := job.Spec.PyTorchReplicaSpecs[PyTorchReplicaTypeWorker].Replicas
+			// Set Min and Max to worker.spec.Replicas.
+			job.Spec.ElasticPolicy.MaxReplicas = workerReplicas
+			job.Spec.ElasticPolicy.MinReplicas = workerReplicas
+		}
+	}
+}
+
 func setDefaultReplicas(spec *common.ReplicaSpec) {
 	if spec.Replicas == nil {
 		spec.Replicas = Int32(1)
 	}
+}
+
+func setDefaultRestartPolicy(spec *common.ReplicaSpec) {
 	if spec.RestartPolicy == "" {
 		spec.RestartPolicy = DefaultRestartPolicy
 	}
@@ -95,12 +117,11 @@ func SetDefaults_PyTorchJob(job *PyTorchJob) {
 	// Update the key of PyTorchReplicaSpecs to camel case.
 	setTypeNamesToCamelCase(job)
 
-	for rType, spec := range job.Spec.PyTorchReplicaSpecs {
-		// Set default replicas to 1.
+	for _, spec := range job.Spec.PyTorchReplicaSpecs {
 		setDefaultReplicas(spec)
-		if rType == PyTorchReplicaTypeMaster {
-			// Set default port to pytorch container of Master.
-			setDefaultPort(&spec.Template.Spec)
-		}
+		setDefaultRestartPolicy(spec)
+		setDefaultPort(&spec.Template.Spec)
 	}
+	// Set default elastic policy.
+	setElasticPolicy(job)
 }

pkg/apis/pytorch/v1/defaults_test.go

@@ -0,0 +1,160 @@
+package v1
+
+import (
+	"testing"
+
+	"github.com/onsi/ginkgo"
+	"github.com/onsi/gomega"
+
+	v1 "github.com/kubeflow/common/pkg/apis/common/v1"
+)
+
+func TestSetElasticPolicy(t *testing.T) {
+	gomega.RegisterFailHandler(ginkgo.Fail)
+
+	type args struct {
+		job *PyTorchJob
+	}
+	type result struct {
+		expectedMinReplicas *int32
+		expectedMaxReplicas *int32
+	}
+	tests := []struct {
+		name   string
+		args   args
+		result result
+	}{
+		{
+			name: "minReplicas and maxReplicas to null",
+			args: args{
+				job: &PyTorchJob{
+					Spec: PyTorchJobSpec{
+						ElasticPolicy: &ElasticPolicy{},
+						PyTorchReplicaSpecs: map[v1.ReplicaType]*v1.ReplicaSpec{
+							PyTorchReplicaTypeWorker: {
+								Replicas: int32Ptr(1),
+							},
+						},
+					},
+				},
+			},
+			result: result{
+				expectedMinReplicas: int32Ptr(1),
+				expectedMaxReplicas: int32Ptr(1),
+			},
+		},
+		{
+			name: "minReplicas and maxReplicas to 1",
+			args: args{
+				job: &PyTorchJob{
+					Spec: PyTorchJobSpec{
+						ElasticPolicy: &ElasticPolicy{
+							MaxReplicas: int32Ptr(1),
+							MinReplicas: int32Ptr(1),
+						},
+						PyTorchReplicaSpecs: map[v1.ReplicaType]*v1.ReplicaSpec{
+							PyTorchReplicaTypeWorker: {
+								Replicas: int32Ptr(1),
+							},
+						},
+					},
+				},
+			},
+			result: result{
+				expectedMinReplicas: int32Ptr(1),
+				expectedMaxReplicas: int32Ptr(1),
+			},
+		},
+		{
+			name: "minReplicas and maxReplicas to 1",
+			args: args{
+				job: &PyTorchJob{
+					Spec: PyTorchJobSpec{
+						ElasticPolicy: &ElasticPolicy{
+							MaxReplicas: int32Ptr(1),
+							MinReplicas: int32Ptr(1),
+						},
+						PyTorchReplicaSpecs: map[v1.ReplicaType]*v1.ReplicaSpec{
+							PyTorchReplicaTypeWorker: {
+								Replicas: int32Ptr(1),
+							},
+						},
+					},
+				},
+			},
+			result: result{
+				expectedMinReplicas: int32Ptr(1),
+				expectedMaxReplicas: int32Ptr(1),
+			},
+		},
+		{
+			name: "minReplicas to null, maxRepliacs to 1",
+			args: args{
+				job: &PyTorchJob{
+					Spec: PyTorchJobSpec{
+						ElasticPolicy: &ElasticPolicy{
+							MaxReplicas: int32Ptr(1),
+							MinReplicas: nil,
+						},
+						PyTorchReplicaSpecs: map[v1.ReplicaType]*v1.ReplicaSpec{
+							PyTorchReplicaTypeWorker: {
+								Replicas: int32Ptr(1),
+							},
+						},
+					},
+				},
+			},
+			result: result{
+				expectedMinReplicas: int32Ptr(1),
+				expectedMaxReplicas: int32Ptr(1),
+			},
+		},
+		{
+			name: "maxRepliacs to null, minReplicas to 1",
+			args: args{
+				job: &PyTorchJob{
+					Spec: PyTorchJobSpec{
+						ElasticPolicy: &ElasticPolicy{
+							MaxReplicas: nil,
+							MinReplicas: int32Ptr(1),
+						},
+						PyTorchReplicaSpecs: map[v1.ReplicaType]*v1.ReplicaSpec{
+							PyTorchReplicaTypeWorker: {
+								Replicas: int32Ptr(1),
+							},
+						},
+					},
+				},
+			},
+			result: result{
+				expectedMinReplicas: int32Ptr(1),
+				expectedMaxReplicas: int32Ptr(1),
+			},
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			setElasticPolicy(test.args.job)
+			if test.result.expectedMinReplicas != nil {
+				gomega.Expect(test.args.job.Spec.ElasticPolicy.MinReplicas).
+					To(gomega.Equal(test.result.expectedMinReplicas))
+			} else {
+				gomega.Expect(test.args.job.Spec.ElasticPolicy.MinReplicas).
+					To(gomega.BeNil())
+			}
+
+			if test.result.expectedMaxReplicas != nil {
+				gomega.Expect(test.args.job.Spec.ElasticPolicy.MaxReplicas).
+					To(gomega.Equal(test.result.expectedMaxReplicas))
+			} else {
+				gomega.Expect(test.args.job.Spec.ElasticPolicy.MaxReplicas).
+					To(gomega.BeNil())
+			}
+		})
+	}
+}
+
+func int32Ptr(n int) *int32 {
+	val := int32(n)
+	return &val
+}