From 338cceadfb9f0eb3be4476b38323f1dbf94b6544 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fabio=20Gr=C3=A4tz?= <fabiogratz@googlemail.com>
Date: Mon, 10 Apr 2023 13:03:49 +0200
Subject: [PATCH 1/4] Set elastic config in pytorchjob
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Fabio Grätz <fabiogratz@googlemail.com>
---
 .../k8s/kfoperators/pytorch/pytorch.go        | 63 ++++++++++++++-----
 1 file changed, 48 insertions(+), 15 deletions(-)

diff --git a/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch.go b/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch.go
index 287590ed5..06758e1b6 100644
--- a/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch.go
+++ b/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch.go
@@ -73,26 +73,59 @@ func (pytorchOperatorResourceHandler) BuildResource(ctx context.Context, taskCtx
 		return nil, fmt.Errorf("number of worker should be more then 0")
 	}
 
-	jobSpec := kubeflowv1.PyTorchJobSpec{
-		PyTorchReplicaSpecs: map[commonOp.ReplicaType]*commonOp.ReplicaSpec{
-			kubeflowv1.PyTorchJobReplicaTypeMaster: {
-				Template: v1.PodTemplateSpec{
-					ObjectMeta: *objectMeta,
-					Spec:       *podSpec,
+	var jobSpec kubeflowv1.PyTorchJobSpec
+
+	elasticConfig := pytorchTaskExtraArgs.GetElasticConfig()
+
+	if elasticConfig != nil {
+		minReplicas := elasticConfig.GetMinReplicas()
+		maxReplicas := elasticConfig.GetMaxReplicas()
+		nProcPerNode := elasticConfig.GetNprocPerNode()
+		maxRestarts := elasticConfig.GetMaxRestarts()
+		rdzvBackend := kubeflowv1.RDZVBackend(elasticConfig.GetRdzvBackend())
+
+		jobSpec = kubeflowv1.PyTorchJobSpec{
+			ElasticPolicy: &kubeflowv1.ElasticPolicy{
+				MinReplicas:  &minReplicas,
+				MaxReplicas:  &maxReplicas,
+				RDZVBackend:  &rdzvBackend,
+				NProcPerNode: &nProcPerNode,
+				MaxRestarts:  &maxRestarts,
+			},
+			PyTorchReplicaSpecs: map[commonOp.ReplicaType]*commonOp.ReplicaSpec{
+				kubeflowv1.PyTorchJobReplicaTypeWorker: {
+					Replicas: &workers,
+					Template: v1.PodTemplateSpec{
+						ObjectMeta: *objectMeta,
+						Spec:       *podSpec,
+					},
+					RestartPolicy: commonOp.RestartPolicyNever,
 				},
-				RestartPolicy: commonOp.RestartPolicyNever,
 			},
-			kubeflowv1.PyTorchJobReplicaTypeWorker: {
-				Replicas: &workers,
-				Template: v1.PodTemplateSpec{
-					ObjectMeta: *objectMeta,
-					Spec:       *podSpec,
+		}
+
+	} else {
+
+		jobSpec = kubeflowv1.PyTorchJobSpec{
+			PyTorchReplicaSpecs: map[commonOp.ReplicaType]*commonOp.ReplicaSpec{
+				kubeflowv1.PyTorchJobReplicaTypeMaster: {
+					Template: v1.PodTemplateSpec{
+						ObjectMeta: *objectMeta,
+						Spec:       *podSpec,
+					},
+					RestartPolicy: commonOp.RestartPolicyNever,
+				},
+				kubeflowv1.PyTorchJobReplicaTypeWorker: {
+					Replicas: &workers,
+					Template: v1.PodTemplateSpec{
+						ObjectMeta: *objectMeta,
+						Spec:       *podSpec,
+					},
+					RestartPolicy: commonOp.RestartPolicyNever,
 				},
-				RestartPolicy: commonOp.RestartPolicyNever,
 			},
-		},
+		}
 	}
-
 	job := &kubeflowv1.PyTorchJob{
 		TypeMeta: metav1.TypeMeta{
 			Kind:       kubeflowv1.PytorchJobKind,

From b644427ef3af6ccb9b6771cdfa6191b2a6209e35 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fabio=20Gr=C3=A4tz?= <fabiogratz@googlemail.com>
Date: Sat, 22 Apr 2023 20:02:31 +0200
Subject: [PATCH 2/4] Add test for pytorch elastic task
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Fabio Grätz <fabiogratz@googlemail.com>
---
 .../k8s/kfoperators/pytorch/pytorch_test.go   | 41 +++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch_test.go b/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch_test.go
index f1979dbf2..3e0c88413 100644
--- a/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch_test.go
+++ b/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch_test.go
@@ -69,6 +69,13 @@ func dummyPytorchCustomObj(workers int32) *plugins.DistributedPyTorchTrainingTas
 	}
 }
 
+func dummyElasticPytorchCustomObj(workers int32, elasticConfig plugins.ElasticConfig) *plugins.DistributedPyTorchTrainingTask {
+	return &plugins.DistributedPyTorchTrainingTask{
+		Workers:       workers,
+		ElasticConfig: &elasticConfig,
+	}
+}
+
 func dummyPytorchTaskTemplate(id string, pytorchCustomObj *plugins.DistributedPyTorchTrainingTask) *core.TaskTemplate {
 
 	ptObjJSON, err := utils.MarshalToString(pytorchCustomObj)
@@ -282,6 +289,39 @@ func dummyPytorchJobResource(pytorchResourceHandler pytorchOperatorResourceHandl
 	}
 }
 
+func TestBuildResourcePytorchElastic(t *testing.T) {
+	pytorchResourceHandler := pytorchOperatorResourceHandler{}
+
+	ptObj := dummyElasticPytorchCustomObj(2, plugins.ElasticConfig{MinReplicas: 1, MaxReplicas: 2, NprocPerNode: 4, RdzvBackend: "c10d"})
+	taskTemplate := dummyPytorchTaskTemplate("the job", ptObj)
+
+	resource, err := pytorchResourceHandler.BuildResource(context.TODO(), dummyPytorchTaskContext(taskTemplate))
+	assert.NoError(t, err)
+	assert.NotNil(t, resource)
+
+	pytorchJob, ok := resource.(*kubeflowv1.PyTorchJob)
+	assert.True(t, ok)
+	assert.Equal(t, int32(2), *pytorchJob.Spec.PyTorchReplicaSpecs[kubeflowv1.PyTorchJobReplicaTypeWorker].Replicas)
+	assert.NotNil(t, pytorchJob.Spec.ElasticPolicy)
+	assert.Equal(t, int32(1), *pytorchJob.Spec.ElasticPolicy.MinReplicas)
+	assert.Equal(t, int32(2), *pytorchJob.Spec.ElasticPolicy.MaxReplicas)
+	assert.Equal(t, int32(4), *pytorchJob.Spec.ElasticPolicy.NProcPerNode)
+	assert.Equal(t, kubeflowv1.RDZVBackend("c10d"), *pytorchJob.Spec.ElasticPolicy.RDZVBackend)
+
+	assert.Equal(t, 1, len(pytorchJob.Spec.PyTorchReplicaSpecs))
+	assert.Contains(t, pytorchJob.Spec.PyTorchReplicaSpecs, kubeflowv1.PyTorchJobReplicaTypeWorker)
+
+	var hasContainerWithDefaultPytorchName = false
+
+	for _, container := range pytorchJob.Spec.PyTorchReplicaSpecs[kubeflowv1.PyTorchJobReplicaTypeWorker].Template.Spec.Containers {
+		if container.Name == kubeflowv1.PytorchJobDefaultContainerName {
+			hasContainerWithDefaultPytorchName = true
+		}
+	}
+
+	assert.True(t, hasContainerWithDefaultPytorchName)
+}
+
 func TestBuildResourcePytorch(t *testing.T) {
 	pytorchResourceHandler := pytorchOperatorResourceHandler{}
 
@@ -295,6 +335,7 @@ func TestBuildResourcePytorch(t *testing.T) {
 	pytorchJob, ok := resource.(*kubeflowv1.PyTorchJob)
 	assert.True(t, ok)
 	assert.Equal(t, int32(100), *pytorchJob.Spec.PyTorchReplicaSpecs[kubeflowv1.PyTorchJobReplicaTypeWorker].Replicas)
+	assert.Nil(t, pytorchJob.Spec.ElasticPolicy)
 
 	for _, replicaSpec := range pytorchJob.Spec.PyTorchReplicaSpecs {
 		var hasContainerWithDefaultPytorchName = false

From e9f1501ee5c818f071b3e8bd9c3d8b238b540777 Mon Sep 17 00:00:00 2001
From: Ketan Umare <ketan.umare@gmail.com>
Date: Sun, 23 Apr 2023 21:11:50 -0700
Subject: [PATCH 3/4] updated flyteidl

Signed-off-by: Ketan Umare <ketan.umare@gmail.com>
---
 go.mod | 2 +-
 go.sum | 2 ++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/go.mod b/go.mod
index 4d54cbdea..78a49a78b 100644
--- a/go.mod
+++ b/go.mod
@@ -12,7 +12,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/athena v1.0.0
 	github.com/bstadlbauer/dask-k8s-operator-go-client v0.1.0
 	github.com/coocood/freecache v1.1.1
-	github.com/flyteorg/flyteidl v1.3.14
+	github.com/flyteorg/flyteidl v1.3.19
 	github.com/flyteorg/flytestdlib v1.0.15
 	github.com/go-test/deep v1.0.7
 	github.com/golang/protobuf v1.5.2
diff --git a/go.sum b/go.sum
index 179728dec..f81a39c3b 100644
--- a/go.sum
+++ b/go.sum
@@ -234,6 +234,8 @@ github.com/fatih/color v1.13.0 h1:8LOYc1KYPPmyKMuN8QV2DNRWNbLo6LZ0iLs8+mlH53w=
 github.com/fatih/color v1.13.0/go.mod h1:kLAiJbzzSOZDVNGyDpeOxJ47H46qBXwg5ILebYFFOfk=
 github.com/flyteorg/flyteidl v1.3.14 h1:o5M0g/r6pXTPu5PEurbYxbQmuOu3hqqsaI2M6uvK0N8=
 github.com/flyteorg/flyteidl v1.3.14/go.mod h1:Pkt2skI1LiHs/2ZoekBnyPhuGOFMiuul6HHcKGZBsbM=
+github.com/flyteorg/flyteidl v1.3.19 h1:i79Dh7UoP8Z4LEJ2ox6jlfZVJtFZ+r4g84CJj1gh22Y=
+github.com/flyteorg/flyteidl v1.3.19/go.mod h1:Pkt2skI1LiHs/2ZoekBnyPhuGOFMiuul6HHcKGZBsbM=
 github.com/flyteorg/flytestdlib v1.0.15 h1:kv9jDQmytbE84caY+pkZN8trJU2ouSAmESzpTEhfTt0=
 github.com/flyteorg/flytestdlib v1.0.15/go.mod h1:ghw/cjY0sEWIIbyCtcJnL/Gt7ZS7gf9SUi0CCPhbz3s=
 github.com/flyteorg/stow v0.3.6 h1:jt50ciM14qhKBaIrB+ppXXY+SXB59FNREFgTJqCyqIk=

From c8de6e2e9850b9ce1ef1fa51f9d1146ad1cff8e5 Mon Sep 17 00:00:00 2001
From: Ketan Umare <ketan.umare@gmail.com>
Date: Sun, 23 Apr 2023 21:57:22 -0700
Subject: [PATCH 4/4] fixing lint issues

Signed-off-by: Ketan Umare <ketan.umare@gmail.com>
---
 .../k8s/kfoperators/pytorch/pytorch_test.go   | 20 +++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch_test.go b/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch_test.go
index 3e0c88413..150bdb59a 100644
--- a/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch_test.go
+++ b/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch_test.go
@@ -267,7 +267,7 @@ func dummyPytorchJobResource(pytorchResourceHandler pytorchOperatorResourceHandl
 	}
 
 	ptObj := dummyPytorchCustomObj(workers)
-	taskTemplate := dummyPytorchTaskTemplate("the job", ptObj)
+	taskTemplate := dummyPytorchTaskTemplate("job1", ptObj)
 	resource, err := pytorchResourceHandler.BuildResource(context.TODO(), dummyPytorchTaskContext(taskTemplate))
 	if err != nil {
 		panic(err)
@@ -293,7 +293,7 @@ func TestBuildResourcePytorchElastic(t *testing.T) {
 	pytorchResourceHandler := pytorchOperatorResourceHandler{}
 
 	ptObj := dummyElasticPytorchCustomObj(2, plugins.ElasticConfig{MinReplicas: 1, MaxReplicas: 2, NprocPerNode: 4, RdzvBackend: "c10d"})
-	taskTemplate := dummyPytorchTaskTemplate("the job", ptObj)
+	taskTemplate := dummyPytorchTaskTemplate("job2", ptObj)
 
 	resource, err := pytorchResourceHandler.BuildResource(context.TODO(), dummyPytorchTaskContext(taskTemplate))
 	assert.NoError(t, err)
@@ -326,13 +326,13 @@ func TestBuildResourcePytorch(t *testing.T) {
 	pytorchResourceHandler := pytorchOperatorResourceHandler{}
 
 	ptObj := dummyPytorchCustomObj(100)
-	taskTemplate := dummyPytorchTaskTemplate("the job", ptObj)
+	taskTemplate := dummyPytorchTaskTemplate("job3", ptObj)
 
-	resource, err := pytorchResourceHandler.BuildResource(context.TODO(), dummyPytorchTaskContext(taskTemplate))
+	res, err := pytorchResourceHandler.BuildResource(context.TODO(), dummyPytorchTaskContext(taskTemplate))
 	assert.NoError(t, err)
-	assert.NotNil(t, resource)
+	assert.NotNil(t, res)
 
-	pytorchJob, ok := resource.(*kubeflowv1.PyTorchJob)
+	pytorchJob, ok := res.(*kubeflowv1.PyTorchJob)
 	assert.True(t, ok)
 	assert.Equal(t, int32(100), *pytorchJob.Spec.PyTorchReplicaSpecs[kubeflowv1.PyTorchJobReplicaTypeWorker].Replicas)
 	assert.Nil(t, pytorchJob.Spec.ElasticPolicy)
@@ -433,17 +433,17 @@ func TestReplicaCounts(t *testing.T) {
 			ptObj := dummyPytorchCustomObj(test.workerReplicaCount)
 			taskTemplate := dummyPytorchTaskTemplate("the job", ptObj)
 
-			resource, err := pytorchResourceHandler.BuildResource(context.TODO(), dummyPytorchTaskContext(taskTemplate))
+			res, err := pytorchResourceHandler.BuildResource(context.TODO(), dummyPytorchTaskContext(taskTemplate))
 			if test.expectError {
 				assert.Error(t, err)
-				assert.Nil(t, resource)
+				assert.Nil(t, res)
 				return
 			}
 
 			assert.NoError(t, err)
-			assert.NotNil(t, resource)
+			assert.NotNil(t, res)
 
-			job, ok := resource.(*kubeflowv1.PyTorchJob)
+			job, ok := res.(*kubeflowv1.PyTorchJob)
 			assert.True(t, ok)
 
 			assert.Len(t, job.Spec.PyTorchReplicaSpecs, len(test.contains))
