Repository: kubeflow/trainer
Pull request #2555: Implement CustomValidation UT for MPI plugin
Author: login: tenzen-y
Labels: name: size/L
name: lgtm
name: approved
Created: 2025-03-21T12:09:21Z | Merged: 2025-03-27T12:11:24Z

<!--  Thanks for sending a pull request! Here are some tips for you:
1. If this is your first time, check our contributor guidelines: https://www.kubeflow.org/docs/about/contributing
2. To know more about Kubeflow Trainer, check the developer guide:
    https://github.com/kubeflow/trainer/blob/master/CONTRIBUTING.md
3. If you want *faster* PR reviews, check how: https://git.k8s.io/community/contributors/guide/pull-requests.md#best-practices-for-faster-reviews
-->

**What this PR does / why we need it**:

Part of https://github.com/kubeflow/trainer/issues/2556

**Which issue(s) this PR fixes** _(optional, in `Fixes #<issue number>, #<issue number>, ...` format, will close the issue(s) when PR gets merged)_:
Fixes #

**Checklist:**

- [ ] [Docs](https://www.kubeflow.org/docs/components/trainer/) included if any changes are user facing



Diff:
diff --git a/pkg/runtime/framework/plugins/mpi/mpi.go b/pkg/runtime/framework/plugins/mpi/mpi.go
index c20175203b..b516e565de 100644
--- a/pkg/runtime/framework/plugins/mpi/mpi.go
+++ b/pkg/runtime/framework/plugins/mpi/mpi.go
@@ -47,6 +47,10 @@ import (
 	"github.com/kubeflow/trainer/pkg/runtime/framework"
 )
 
+var (
+	numProcPerNodePath = field.NewPath("spec").Child("trainer").Child("numProcPerNode")
+)
+
 // TODO : Support MPICH and IntelMPI implementations.
 
 type MPI struct {
@@ -85,12 +89,9 @@ func (m *MPI) Validate(_ client.Object, runtimeInfo *runtime.Info, _, newJobObj
 		return nil, allErrs
 	}
 
-	specPath := field.NewPath("spec")
-	if newJobObj.Spec.Trainer != nil && newJobObj.Spec.Trainer.NumProcPerNode != nil {
-		numProcPerNodePath := specPath.Child("trainer").Child("numProcPerNode")
-		numProcPerNode := *newJobObj.Spec.Trainer.NumProcPerNode
-		if numProcPerNode.Type != intstr.Int {
-			allErrs = append(allErrs, field.Invalid(numProcPerNodePath, newJobObj.Spec.Trainer.NumProcPerNode, "must have an int value"))
+	if trainJobTrainer := newJobObj.Spec.Trainer; trainJobTrainer != nil && trainJobTrainer.NumProcPerNode != nil {
+		if trainJobTrainer.NumProcPerNode.Type != intstr.Int {
+			allErrs = append(allErrs, field.Invalid(numProcPerNodePath, *trainJobTrainer.NumProcPerNode, "must have an int value for MPI TrainJob"))
 		}
 	}
 	return nil, allErrs
diff --git a/pkg/runtime/framework/plugins/mpi/mpi_test.go b/pkg/runtime/framework/plugins/mpi/mpi_test.go
index 4103c81f3a..66cac48845 100644
--- a/pkg/runtime/framework/plugins/mpi/mpi_test.go
+++ b/pkg/runtime/framework/plugins/mpi/mpi_test.go
@@ -29,11 +29,13 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	apiruntime "k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/util/intstr"
+	"k8s.io/apimachinery/pkg/util/validation/field"
 	corev1ac "k8s.io/client-go/applyconfigurations/core/v1"
 	"k8s.io/klog/v2/ktesting"
 	"k8s.io/utils/ptr"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/client/interceptor"
+	"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
 
 	trainer "github.com/kubeflow/trainer/pkg/apis/trainer/v1alpha1"
 	"github.com/kubeflow/trainer/pkg/constants"
@@ -739,3 +741,90 @@ trainJob-node-1-0.trainJob slots=1
 		})
 	}
 }
+
+func TestValidate(t *testing.T) {
+	cases := map[string]struct {
+		info         *runtime.Info
+		oldObj       *trainer.TrainJob
+		newObj       *trainer.TrainJob
+		wantError    field.ErrorList
+		wantWarnings admission.Warnings
+	}{
+		"no action when info is nil ": {},
+		"no action when info does not have MLPolicySource": {
+			info: runtime.NewInfo(),
+			oldObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
+				Obj(),
+			newObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
+				Obj(),
+		},
+		"info does not have MPIPolicySource": {
+			info: runtime.NewInfo(
+				runtime.WithMLPolicySource(utiltesting.MakeMLPolicyWrapper().Obj()),
+			),
+			oldObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
+				Obj(),
+			newObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
+				Obj(),
+		},
+		"numProcPerNode typed is string": {
+			info: runtime.NewInfo(
+				runtime.WithMLPolicySource(utiltesting.MakeMLPolicyWrapper().
+					WithMLPolicySource(*utiltesting.MakeMLPolicySourceWrapper().
+						MPIPolicy(ptr.To[int32](1), ptr.To(trainer.MPIImplementationOpenMPI), nil, nil).
+						Obj(),
+					).
+					Obj()),
+			),
+			newObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
+				Trainer(utiltesting.MakeTrainJobTrainerWrapper().
+					NumProcPerNode(intstr.FromString("invalid")).
+					Obj(),
+				).
+				Obj(),
+			wantError: field.ErrorList{
+				field.Invalid(
+					field.NewPath("spec").Child("trainer").Child("numProcPerNode"),
+					intstr.FromString("invalid"),
+					"must have an int value for MPI TrainJob",
+				),
+			},
+		},
+		"numProcPerNode typed is int": {
+			info: runtime.NewInfo(
+				runtime.WithMLPolicySource(utiltesting.MakeMLPolicyWrapper().
+					WithMLPolicySource(*utiltesting.MakeMLPolicySourceWrapper().
+						MPIPolicy(ptr.To[int32](1), ptr.To(trainer.MPIImplementationOpenMPI), nil, nil).
+						Obj(),
+					).
+					Obj(),
+				),
+			),
+			newObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
+				Trainer(utiltesting.MakeTrainJobTrainerWrapper().
+					NumProcPerNode(intstr.FromInt32(8)).
+					Obj(),
+				).
+				Obj(),
+		},
+	}
+	for name, tc := range cases {
+		t.Run(name, func(t *testing.T) {
+			_, ctx := ktesting.NewTestContext(t)
+			var cancel func()
+			ctx, cancel = context.WithCancel(ctx)
+			t.Cleanup(cancel)
+			p, err := New(ctx, utiltesting.NewClientBuilder().Build(), nil)
+			if err != nil {
+				t.Fatalf("Failed to initialize MPI plugin: %v", err)
+			}
+			warnings, errs := p.(framework.CustomValidationPlugin).Validate(nil, tc.info, tc.oldObj, tc.newObj)
+			if diff := gocmp.Diff(tc.wantError, errs); len(diff) != 0 {
+				t.Errorf("Unexpected error from Validate (-want, +got): %s", diff)
+			}
+			if diff := gocmp.Diff(tc.wantWarnings, warnings); len(diff) != 0 {
+				t.Errorf("Unexpected warnings from Validate (-want, +got): %s", diff)
+			}
+		})
+	}
+}


Reviews:
user: login: astefanutti
body: LGTM, thanks.
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T13:26:34Z
user: login: tenzen-y
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-21T13:30:33Z
user: login: tenzen-y
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-21T13:30:40Z
user: login: astefanutti
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T13:44:01Z
user: login: andreyvelich
body: Thank you for the update @tenzen-y!
/lgtm
/approve

state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-27T12:10:56Z

Review comments:
diff_hunk: @@ -85,12 +89,9 @@ func (m *MPI) Validate(_ client.Object, runtimeInfo *runtime.Info, _, newJobObj
 		return nil, allErrs
 	}
 
-	specPath := field.NewPath("spec")
-	if newJobObj.Spec.Trainer != nil && newJobObj.Spec.Trainer.NumProcPerNode != nil {
-		numProcPerNodePath := specPath.Child("trainer").Child("numProcPerNode")
-		numProcPerNode := *newJobObj.Spec.Trainer.NumProcPerNode
-		if numProcPerNode.Type != intstr.Int {
-			allErrs = append(allErrs, field.Invalid(numProcPerNodePath, newJobObj.Spec.Trainer.NumProcPerNode, "must have an int value"))
+	if trainJobTrainer := newJobObj.Spec.Trainer; trainJobTrainer != nil && trainJobTrainer.NumProcPerNode != nil {
+		if trainJobTrainer.NumProcPerNode.Type != intstr.Int {
+			allErrs = append(allErrs, field.Invalid(numProcPerNodePath, *trainJobTrainer.NumProcPerNode, "must have an int value"))
path: pkg/runtime/framework/plugins/mpi/mpi.go
user: login: astefanutti
body: Not strictly in the scope of that PR, but while we are at it, maybe the error message could be more specific, e.g. "must have an int value for MPI job".
created_at: 2025-03-21T13:26:22Z
author_association: CONTRIBUTOR
diff_hunk: @@ -85,12 +89,9 @@ func (m *MPI) Validate(_ client.Object, runtimeInfo *runtime.Info, _, newJobObj
 		return nil, allErrs
 	}
 
-	specPath := field.NewPath("spec")
-	if newJobObj.Spec.Trainer != nil && newJobObj.Spec.Trainer.NumProcPerNode != nil {
-		numProcPerNodePath := specPath.Child("trainer").Child("numProcPerNode")
-		numProcPerNode := *newJobObj.Spec.Trainer.NumProcPerNode
-		if numProcPerNode.Type != intstr.Int {
-			allErrs = append(allErrs, field.Invalid(numProcPerNodePath, newJobObj.Spec.Trainer.NumProcPerNode, "must have an int value"))
+	if trainJobTrainer := newJobObj.Spec.Trainer; trainJobTrainer != nil && trainJobTrainer.NumProcPerNode != nil {
+		if trainJobTrainer.NumProcPerNode.Type != intstr.Int {
+			allErrs = append(allErrs, field.Invalid(numProcPerNodePath, *trainJobTrainer.NumProcPerNode, "must have an int value"))
path: pkg/runtime/framework/plugins/mpi/mpi.go
user: login: tenzen-y
body: Good point, let me contain your request in this PR.

created_at: 2025-03-21T13:30:32Z
author_association: MEMBER
diff_hunk: @@ -85,12 +89,9 @@ func (m *MPI) Validate(_ client.Object, runtimeInfo *runtime.Info, _, newJobObj
 		return nil, allErrs
 	}
 
-	specPath := field.NewPath("spec")
-	if newJobObj.Spec.Trainer != nil && newJobObj.Spec.Trainer.NumProcPerNode != nil {
-		numProcPerNodePath := specPath.Child("trainer").Child("numProcPerNode")
-		numProcPerNode := *newJobObj.Spec.Trainer.NumProcPerNode
-		if numProcPerNode.Type != intstr.Int {
-			allErrs = append(allErrs, field.Invalid(numProcPerNodePath, newJobObj.Spec.Trainer.NumProcPerNode, "must have an int value"))
+	if trainJobTrainer := newJobObj.Spec.Trainer; trainJobTrainer != nil && trainJobTrainer.NumProcPerNode != nil {
+		if trainJobTrainer.NumProcPerNode.Type != intstr.Int {
+			allErrs = append(allErrs, field.Invalid(numProcPerNodePath, *trainJobTrainer.NumProcPerNode, "must have an int value"))
path: pkg/runtime/framework/plugins/mpi/mpi.go
user: login: tenzen-y
body: Done.
created_at: 2025-03-21T13:30:40Z
author_association: MEMBER
diff_hunk: @@ -85,12 +89,9 @@ func (m *MPI) Validate(_ client.Object, runtimeInfo *runtime.Info, _, newJobObj
 		return nil, allErrs
 	}
 
-	specPath := field.NewPath("spec")
-	if newJobObj.Spec.Trainer != nil && newJobObj.Spec.Trainer.NumProcPerNode != nil {
-		numProcPerNodePath := specPath.Child("trainer").Child("numProcPerNode")
-		numProcPerNode := *newJobObj.Spec.Trainer.NumProcPerNode
-		if numProcPerNode.Type != intstr.Int {
-			allErrs = append(allErrs, field.Invalid(numProcPerNodePath, newJobObj.Spec.Trainer.NumProcPerNode, "must have an int value"))
+	if trainJobTrainer := newJobObj.Spec.Trainer; trainJobTrainer != nil && trainJobTrainer.NumProcPerNode != nil {
+		if trainJobTrainer.NumProcPerNode.Type != intstr.Int {
+			allErrs = append(allErrs, field.Invalid(numProcPerNodePath, *trainJobTrainer.NumProcPerNode, "must have an int value"))
path: pkg/runtime/framework/plugins/mpi/mpi.go
user: login: astefanutti
body: Thanks, LGTM.
created_at: 2025-03-21T13:44:01Z
author_association: CONTRIBUTOR

PR comments:
user: login: coveralls
created_at: 2025-03-21T12:13:36Z
body: ## Pull Request Test Coverage Report for [Build 13992886013](https://coveralls.io/builds/72856501)


### Details

* **3** of **3**   **(100.0%)**  changed or added relevant lines in **1** file are covered.
* No unchanged relevant lines lost coverage.
* Overall coverage increased (+**0.5%**) to **64.859%**

---



|  Totals | [![Coverage Status](https://coveralls.io/builds/72856501/badge)](https://coveralls.io/builds/72856501) |
| :-- | --: |
| Change from base [Build 13991199064](https://coveralls.io/builds/72854979): |  0.5% |
| Covered Lines: | 1698 |
| Relevant Lines: | 2618 |

---
##### 💛  - [Coveralls](https://coveralls.io)

author_association: NONE
user: login: tenzen-y
created_at: 2025-03-21T12:32:35Z
body: /assign @kubeflow/wg-training-leads @astefanutti 
author_association: MEMBER
user: login: tenzen-y
created_at: 2025-03-27T11:53:26Z
body: @kubeflow/wg-training-leads @astefanutti @Electronic-Waste Can anyone add lgtm tag?

author_association: MEMBER
user: login: google-oss-prow[bot]
created_at: 2025-03-27T12:11:03Z
body: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/kubeflow/trainer/pull/2555#pullrequestreview-2721161098" title="Approved">andreyvelich</a>*, *<a href="https://github.com/kubeflow/trainer/pull/2555#pullrequestreview-2705943055" title="Approved">astefanutti</a>*

The full list of commands accepted by this bot can be found [here](https://oss.gprow.dev/command-help?repo=kubeflow%2Ftrainer).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/kubeflow/trainer/blob/master/OWNERS)~~ [andreyvelich]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->
author_association: NONE