feat(replica-spec): Update corresponding golang files based on protob…

…uf changes Resolves: #4408 Signed-off-by: Chi-Sheng Liu <[email protected]>
flyteorg · May 18, 2024 · 81bceca · 81bceca
1 parent ed46572
commit 81bceca
Show file tree

Hide file tree

Showing 5 changed files with 1,002 additions and 576 deletions.
diff --git a/flyteplugins/go/tasks/plugins/k8s/kfoperators/common/common_operator.go b/flyteplugins/go/tasks/plugins/k8s/kfoperators/common/common_operator.go
@@ -294,16 +294,35 @@ type kfDistributedReplicaSpec interface {
 	GetImage() string
 	GetResources() *core.Resources
 	GetRestartPolicy() kfplugins.RestartPolicy
+	GetCommon() *kfplugins.CommonReplicaSpec
 }
 
 type allowsCommandOverride interface {
 	GetCommand() []string
 }
 
 func ToReplicaSpecWithOverrides(ctx context.Context, taskCtx pluginsCore.TaskExecutionContext, rs kfDistributedReplicaSpec, primaryContainerName string, isMaster bool) (*commonOp.ReplicaSpec, error) {
+	var replicas int32
+	var image string
+	var resources *core.Resources
+	var restartPolicy kfplugins.RestartPolicy
+
+	common := rs.GetCommon()
+	if common != nil {
+		replicas = common.GetReplicas()
+		image = common.GetImage()
+		resources = common.GetResources()
+		restartPolicy = common.GetRestartPolicy()
+	} else {
+		replicas = rs.GetReplicas()
+		image = rs.GetImage()
+		resources = rs.GetResources()
+		restartPolicy = rs.GetRestartPolicy()
+	}
+
 	taskCtxOptions := []flytek8s.PluginTaskExecutionContextOption{}
-	if rs != nil && rs.GetResources() != nil {
-		resources, err := flytek8s.ToK8sResourceRequirements(rs.GetResources())
+	if resources != nil {
+		resources, err := flytek8s.ToK8sResourceRequirements(resources)
 		if err != nil {
 			return nil, flyteerr.Errorf(flyteerr.BadTaskSpecification, "invalid TaskSpecification on Resources [%v], Err: [%v]", resources, err.Error())
 		}
@@ -321,26 +340,23 @@ func ToReplicaSpecWithOverrides(ctx context.Context, taskCtx pluginsCore.TaskExe
 		replicaSpec.Replicas = &replicas
 	}
 
-	if rs != nil {
-		var command []string
-		if v, ok := rs.(allowsCommandOverride); ok {
-			command = v.GetCommand()
-		}
-		if err := OverrideContainerSpec(
-			&replicaSpec.Template.Spec,
-			primaryContainerName,
-			rs.GetImage(),
-			command,
-		); err != nil {
-			return nil, err
-		}
+	var command []string
+	if v, ok := rs.(allowsCommandOverride); ok {
+		command = v.GetCommand()
+	}
+	if err := OverrideContainerSpec(
+		&replicaSpec.Template.Spec,
+		primaryContainerName,
+		image,
+		command,
+	); err != nil {
+		return nil, err
+	}
 
-		replicaSpec.RestartPolicy = ParseRestartPolicy(rs.GetRestartPolicy())
+	replicaSpec.RestartPolicy = ParseRestartPolicy(restartPolicy)
 
-		if !isMaster {
-			replicas := rs.GetReplicas()
-			replicaSpec.Replicas = &replicas
-		}
+	if !isMaster {
+		replicaSpec.Replicas = &replicas
 	}
 
 	return replicaSpec, nil