Kf operators use GetReplicaFunc (Error Handling)

flyteplugins/go/tasks/plugins/k8s/kfoperators/common/common_operator.go

@@ -266,27 +266,27 @@
 	return nil
 }
 
 func ToReplicaSpec(ctx context.Context, taskCtx pluginsCore.TaskExecutionContext, primaryContainerName string) (*commonOp.ReplicaSpec, error) {
 	podSpec, objectMeta, oldPrimaryContainerName, err := flytek8s.ToK8sPodSpec(ctx, taskCtx)
 	if err != nil {
 		return nil, flyteerr.Errorf(flyteerr.BadTaskSpecification, "Unable to create pod spec: [%v]", err.Error())
 	}
 
 	OverridePrimaryContainerName(podSpec, oldPrimaryContainerName, primaryContainerName)
 
 	cfg := config.GetK8sPluginConfig()
 	objectMeta.Annotations = utils.UnionMaps(cfg.DefaultAnnotations, objectMeta.Annotations, utils.CopyMap(taskCtx.TaskExecutionMetadata().GetAnnotations()))
 	objectMeta.Labels = utils.UnionMaps(cfg.DefaultLabels, objectMeta.Labels, utils.CopyMap(taskCtx.TaskExecutionMetadata().GetLabels()))
 
 	replicas := int32(0)
 	return &commonOp.ReplicaSpec{
 		Replicas: &replicas,
 		Template: v1.PodTemplateSpec{
 			ObjectMeta: *objectMeta,
 			Spec:       *podSpec,
 		},
 		RestartPolicy: commonOp.RestartPolicyNever,
 	}, nil
 }
 
 type kfDistributedReplicaSpec interface {
@@ -300,48 +300,56 @@
 	GetCommand() []string
 }
 
 func ToReplicaSpecWithOverrides(ctx context.Context, taskCtx pluginsCore.TaskExecutionContext, rs kfDistributedReplicaSpec, primaryContainerName string, isMaster bool) (*commonOp.ReplicaSpec, error) {
 	taskCtxOptions := []flytek8s.PluginTaskExecutionContextOption{}
 	if rs != nil && rs.GetResources() != nil {
 		resources, err := flytek8s.ToK8sResourceRequirements(rs.GetResources())
 		if err != nil {
 			return nil, flyteerr.Errorf(flyteerr.BadTaskSpecification, "invalid TaskSpecification on Resources [%v], Err: [%v]", resources, err.Error())
 		}
 		taskCtxOptions = append(taskCtxOptions, flytek8s.WithResources(resources))
 	}
 	newTaskCtx := flytek8s.NewPluginTaskExecutionContext(taskCtx, taskCtxOptions...)
 	replicaSpec, err := ToReplicaSpec(ctx, newTaskCtx, primaryContainerName)
 	if err != nil {
 		return nil, err
 	}
 
 	// Master should have a single replica
 	if isMaster {
 		replicas := int32(1)
 		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
 		}
 
 		replicaSpec.RestartPolicy = ParseRestartPolicy(rs.GetRestartPolicy())
 
 		if !isMaster {
 			replicas := rs.GetReplicas()
 			replicaSpec.Replicas = &replicas
 		}
 	}
 
 	return replicaSpec, nil
 }
+
+func GetReplicaCount(specs map[commonOp.ReplicaType]*commonOp.ReplicaSpec, replicaType commonOp.ReplicaType) *int32 {
+	if spec, ok := specs[replicaType]; ok && spec.Replicas != nil {
+		return spec.Replicas
+	}
+
+	return new(int32) // return 0 as default value
+}

flyteplugins/go/tasks/plugins/k8s/kfoperators/mpi/mpi.go

@@ -69,8 +69,8 @@
 
 		replicaSpec, err := common.ToReplicaSpec(ctx, taskCtx, kubeflowv1.MPIJobDefaultContainerName)
 		if err != nil {
 			return nil, flyteerr.Errorf(flyteerr.BadTaskSpecification, "Unable to create replica spec: [%v]", err.Error())
 		}
 		launcherReplicaSpec = replicaSpec.DeepCopy()
 		// TODO (jeev): Is this even a valid configuration. Can there be more than 1
 		// launcher? TaskTypeVersion 1 does not support overriding this value.
@@ -108,12 +108,12 @@
 
 		launcherReplicaSpec, err = common.ToReplicaSpecWithOverrides(ctx, taskCtx, kfMPITaskExtraArgs.GetLauncherReplicas(), kubeflowv1.MPIJobDefaultContainerName, true)
 		if err != nil {
 			return nil, flyteerr.Errorf(flyteerr.BadTaskSpecification, "Unable to create launcher replica spec: [%v]", err.Error())
 		}
 
 		workerReplicaSpec, err = common.ToReplicaSpecWithOverrides(ctx, taskCtx, kfMPITaskExtraArgs.GetWorkerReplicas(), kubeflowv1.MPIJobDefaultContainerName, false)
 		if err != nil {
 			return nil, flyteerr.Errorf(flyteerr.BadTaskSpecification, "Unable to create worker replica spec: [%v]", err.Error())
 		}
 
 		if kfMPITaskExtraArgs.GetRunPolicy() != nil {
@@ -129,7 +129,7 @@
 		return nil, fmt.Errorf("number of workers must be greater than 0")
 	}
 	if *launcherReplicaSpec.Replicas <= 0 {
 		return nil, fmt.Errorf("number of launchers must be greater than 0")
 	}
 
 	jobSpec := kubeflowv1.MPIJobSpec{
@@ -162,8 +162,8 @@
 		return pluginsCore.PhaseInfoUndefined, fmt.Errorf("failed to convert resource data type")
 	}
 
-	numWorkers = app.Spec.MPIReplicaSpecs[kubeflowv1.MPIJobReplicaTypeWorker].Replicas
-	numLauncherReplicas = app.Spec.MPIReplicaSpecs[kubeflowv1.MPIJobReplicaTypeLauncher].Replicas
+	numWorkers = common.GetReplicaCount(app.Spec.MPIReplicaSpecs, kubeflowv1.MPIJobReplicaTypeWorker)
+	numLauncherReplicas = common.GetReplicaCount(app.Spec.MPIReplicaSpecs, kubeflowv1.MPIJobReplicaTypeLauncher)
 
 	taskLogs, err := common.GetLogs(pluginContext, common.MPITaskType, app.ObjectMeta, false,
 		*numWorkers, *numLauncherReplicas, 0, 0)

flyteplugins/go/tasks/plugins/k8s/kfoperators/mpi/mpi_test.go

@@ -776,3 +776,17 @@ func TestBuildResourceMPIV1ResourceTolerations(t *testing.T) {
 	assert.NotContains(t, mpiJob.Spec.MPIReplicaSpecs[kubeflowv1.MPIJobReplicaTypeLauncher].Template.Spec.Tolerations, gpuToleration)
 	assert.Contains(t, mpiJob.Spec.MPIReplicaSpecs[kubeflowv1.MPIJobReplicaTypeWorker].Template.Spec.Tolerations, gpuToleration)
 }
+
+func TestGetReplicaCount(t *testing.T) {
+	mpiResourceHandler := mpiOperatorResourceHandler{}
+	tfObj := dummyMPICustomObj(1, 1, 0)
+	taskTemplate := dummyMPITaskTemplate("the job", tfObj)
+	resource, err := mpiResourceHandler.BuildResource(context.TODO(), dummyMPITaskContext(taskTemplate, resourceRequirements, nil))
+	assert.NoError(t, err)
+	assert.NotNil(t, resource)
+	MPIJob, ok := resource.(*kubeflowv1.MPIJob)
+	assert.True(t, ok)
+
+	assert.NotNil(t, common.GetReplicaCount(MPIJob.Spec.MPIReplicaSpecs, kubeflowv1.MPIJobReplicaTypeWorker))
+	assert.NotNil(t, common.GetReplicaCount(MPIJob.Spec.MPIReplicaSpecs, kubeflowv1.MPIJobReplicaTypeLauncher))
+}

flyteplugins/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch.go

@@ -67,8 +67,8 @@
 
 		replicaSpec, err := common.ToReplicaSpec(ctx, taskCtx, kubeflowv1.PytorchJobDefaultContainerName)
 		if err != nil {
 			return nil, flyteerr.Errorf(flyteerr.BadTaskSpecification, "Unable to create replica spec: [%v]", err.Error())
 		}
 		masterReplicaSpec = replicaSpec.DeepCopy()
 		masterReplicas := int32(1)
 		masterReplicaSpec.Replicas = &masterReplicas
@@ -91,12 +91,12 @@
 
 		masterReplicaSpec, err = common.ToReplicaSpecWithOverrides(ctx, taskCtx, kfPytorchTaskExtraArgs.GetMasterReplicas(), kubeflowv1.PytorchJobDefaultContainerName, true)
 		if err != nil {
 			return nil, flyteerr.Errorf(flyteerr.BadTaskSpecification, "Unable to create master replica spec: [%v]", err.Error())
 		}
 
 		workerReplicaSpec, err = common.ToReplicaSpecWithOverrides(ctx, taskCtx, kfPytorchTaskExtraArgs.GetWorkerReplicas(), kubeflowv1.PytorchJobDefaultContainerName, false)
 		if err != nil {
 			return nil, flyteerr.Errorf(flyteerr.BadTaskSpecification, "Unable to create worker replica spec: [%v]", err.Error())
 		}
 
 		if kfPytorchTaskExtraArgs.GetRunPolicy() != nil {
@@ -171,15 +171,18 @@
 // any operations that might take a long time (limits are configured system-wide) should be offloaded to the
 // background.
 func (pytorchOperatorResourceHandler) GetTaskPhase(_ context.Context, pluginContext k8s.PluginContext, resource client.Object) (pluginsCore.PhaseInfo, error) {
-	app := resource.(*kubeflowv1.PyTorchJob)
+	app, ok := resource.(*kubeflowv1.PyTorchJob)
+	if !ok {
+		return pluginsCore.PhaseInfoUndefined, fmt.Errorf("failed to convert resource data type")
+	}
 
 	// Elastic PytorchJobs don't use master replicas
 	hasMaster := false
 	if _, ok := app.Spec.PyTorchReplicaSpecs[kubeflowv1.PyTorchJobReplicaTypeMaster]; ok {
 		hasMaster = true
 	}
 
-	workersCount := app.Spec.PyTorchReplicaSpecs[kubeflowv1.PyTorchJobReplicaTypeWorker].Replicas
+	workersCount := common.GetReplicaCount(app.Spec.PyTorchReplicaSpecs, kubeflowv1.PyTorchJobReplicaTypeWorker)
 
 	taskLogs, err := common.GetLogs(pluginContext, common.PytorchTaskType, app.ObjectMeta, hasMaster, *workersCount, 0, 0, 0)
 	if err != nil {

flyteplugins/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch_test.go

@@ -959,3 +959,16 @@ func TestParseElasticConfig(t *testing.T) {
 	assert.Equal(t, int32(4), *elasticPolicy.NProcPerNode)
 	assert.Equal(t, kubeflowv1.RDZVBackend("c10d"), *elasticPolicy.RDZVBackend)
 }
+
+func TestGetReplicaCount(t *testing.T) {
+	pytorchResourceHandler := pytorchOperatorResourceHandler{}
+	tfObj := dummyPytorchCustomObj(1)
+	taskTemplate := dummyPytorchTaskTemplate("the job", tfObj)
+	resource, err := pytorchResourceHandler.BuildResource(context.TODO(), dummyPytorchTaskContext(taskTemplate, resourceRequirements, nil))
+	assert.NoError(t, err)
+	assert.NotNil(t, resource)
+	PytorchJob, ok := resource.(*kubeflowv1.PyTorchJob)
+	assert.True(t, ok)
+
+	assert.NotNil(t, common.GetReplicaCount(PytorchJob.Spec.PyTorchReplicaSpecs, kubeflowv1.PyTorchJobReplicaTypeWorker))
+}

flyteplugins/go/tasks/plugins/k8s/kfoperators/tensorflow/tensorflow.go

@@ -65,8 +65,8 @@
 
 		replicaSpec, err := common.ToReplicaSpec(ctx, taskCtx, kubeflowv1.TFJobDefaultContainerName)
 		if err != nil {
 			return nil, flyteerr.Errorf(flyteerr.BadTaskSpecification, "Unable to create replica spec: [%v]", err.Error())
 		}
 
 		replicaNumMap := map[commonOp.ReplicaType]int32{
 			kubeflowv1.TFJobReplicaTypeChief:  tensorflowTaskExtraArgs.GetChiefReplicas(),
@@ -105,7 +105,7 @@
 			}
 			rs, err := common.ToReplicaSpecWithOverrides(ctx, taskCtx, cfg, kubeflowv1.TFJobDefaultContainerName, false)
 			if err != nil {
 				return nil, flyteerr.Errorf(flyteerr.BadTaskSpecification, "Unable to create replica spec: [%v]", err.Error())
 			}
 			replicaSpecMap[t] = rs
 		}
@@ -139,24 +139,19 @@
 	return job, nil
 }
 
-func getReplicaCount(specs map[commonOp.ReplicaType]*commonOp.ReplicaSpec, replicaType commonOp.ReplicaType) *int32 {
-	if spec, ok := specs[replicaType]; ok && spec.Replicas != nil {
-		return spec.Replicas
-	}
-
-	return new(int32) // return 0 as default value
-}
-
 // Analyses the k8s resource and reports the status as TaskPhase. This call is expected to be relatively fast,
 // any operations that might take a long time (limits are configured system-wide) should be offloaded to the
 // background.
 func (tensorflowOperatorResourceHandler) GetTaskPhase(_ context.Context, pluginContext k8s.PluginContext, resource client.Object) (pluginsCore.PhaseInfo, error) {
-	app := resource.(*kubeflowv1.TFJob)
+	app, ok := resource.(*kubeflowv1.TFJob)
+	if !ok {
+		return pluginsCore.PhaseInfoUndefined, fmt.Errorf("failed to convert resource data type")
+	}
 
-	workersCount := getReplicaCount(app.Spec.TFReplicaSpecs, kubeflowv1.TFJobReplicaTypeWorker)
-	psReplicasCount := getReplicaCount(app.Spec.TFReplicaSpecs, kubeflowv1.TFJobReplicaTypePS)
-	chiefCount := getReplicaCount(app.Spec.TFReplicaSpecs, kubeflowv1.TFJobReplicaTypeChief)
-	evaluatorReplicasCount := getReplicaCount(app.Spec.TFReplicaSpecs, kubeflowv1.TFJobReplicaTypeEval)
+	workersCount := common.GetReplicaCount(app.Spec.TFReplicaSpecs, kubeflowv1.TFJobReplicaTypeWorker)
+	psReplicasCount := common.GetReplicaCount(app.Spec.TFReplicaSpecs, kubeflowv1.TFJobReplicaTypePS)
+	chiefCount := common.GetReplicaCount(app.Spec.TFReplicaSpecs, kubeflowv1.TFJobReplicaTypeChief)
+	evaluatorReplicasCount := common.GetReplicaCount(app.Spec.TFReplicaSpecs, kubeflowv1.TFJobReplicaTypeEval)
 
 	taskLogs, err := common.GetLogs(pluginContext, common.TensorflowTaskType, app.ObjectMeta, false,
 		*workersCount, *psReplicasCount, *chiefCount, *evaluatorReplicasCount)

flyteplugins/go/tasks/plugins/k8s/kfoperators/tensorflow/tensorflow_test.go

@@ -305,10 +305,10 @@ func TestGetReplicaCount(t *testing.T) {
 	tensorflowJob, ok := resource.(*kubeflowv1.TFJob)
 	assert.True(t, ok)
 
-	assert.NotNil(t, getReplicaCount(tensorflowJob.Spec.TFReplicaSpecs, kubeflowv1.TFJobReplicaTypeWorker))
-	assert.NotNil(t, getReplicaCount(tensorflowJob.Spec.TFReplicaSpecs, kubeflowv1.TFJobReplicaTypePS))
-	assert.NotNil(t, getReplicaCount(tensorflowJob.Spec.TFReplicaSpecs, kubeflowv1.TFJobReplicaTypeChief))
-	assert.NotNil(t, getReplicaCount(tensorflowJob.Spec.TFReplicaSpecs, kubeflowv1.TFJobReplicaTypeEval))
+	assert.NotNil(t, common.GetReplicaCount(tensorflowJob.Spec.TFReplicaSpecs, kubeflowv1.TFJobReplicaTypeWorker))
+	assert.NotNil(t, common.GetReplicaCount(tensorflowJob.Spec.TFReplicaSpecs, kubeflowv1.TFJobReplicaTypePS))
+	assert.NotNil(t, common.GetReplicaCount(tensorflowJob.Spec.TFReplicaSpecs, kubeflowv1.TFJobReplicaTypeChief))
+	assert.NotNil(t, common.GetReplicaCount(tensorflowJob.Spec.TFReplicaSpecs, kubeflowv1.TFJobReplicaTypeEval))
 }
 
 func TestBuildResourceTensorFlow(t *testing.T) {