Adding cel validations on trainjob crd

Signed-off-by: Akshay Chitneni <[email protected]>
kubeflow · Oct 18, 2024 · 54ee0b8 · 54ee0b8
1 parent 6965c1a
commit 54ee0b8
Show file tree

Hide file tree

Showing 3 changed files with 144 additions and 8 deletions.
diff --git a/manifests/v2/base/crds/kubeflow.org_trainjobs.yaml b/manifests/v2/base/crds/kubeflow.org_trainjobs.yaml
@@ -197,15 +197,14 @@ spec:
                   They will be merged with the TrainingRuntime values.
                 type: object
               managedBy:
-                description: |-
-                  ManagedBy is used to indicate the controller or entity that manages a TrainJob.
-                  The value must be either an empty, `kubeflow.org/trainjob-controller` or
-                  `kueue.x-k8s.io/multikueue`. The built-in TrainJob controller reconciles TrainJob which
-                  don't have this field at all or the field value is the reserved string
-                  `kubeflow.org/trainjob-controller`, but delegates reconciling TrainJobs
-                  with a 'kueue.x-k8s.io/multikueue' to the Kueue. The field is immutable.
-                  Defaults to `kubeflow.org/trainjob-controller`
+                default: kubeflow.org/trainjob-controller
                 type: string
+                x-kubernetes-validations:
+                - message: ManagedBy must be kubeflow.org/trainjob-controller or kueue.x-k8s.io/multikueue
+                    if set
+                  rule: self in ['kubeflow.org/trainjob-controller', 'kueue.x-k8s.io/multikueue']
+                - message: ManagedBy value is immutable
+                  rule: self == oldSelf
               modelConfig:
                 description: Configuration of the pre-trained and trained model.
                 properties:
@@ -2736,16 +2735,22 @@ spec:
                 description: Reference to the training runtime.
                 properties:
                   apiGroup:
+                    default: kubeflow.org
                     description: |-
                       APIGroup of the runtime being referenced.
                       Defaults to `kubeflow.org`.
                     type: string
                   kind:
+                    default: ClusterTrainingRuntime
                     description: |-
                       Kind of the runtime being referenced.
                       It must be one of TrainingRuntime or ClusterTrainingRuntime.
                       Defaults to ClusterTrainingRuntime.
                     type: string
+                    x-kubernetes-validations:
+                    - message: Kind must be ClusterTrainingRuntime or TrainingRuntime
+                        if set
+                      rule: self in ['ClusterTrainingRuntime', 'TrainingRuntime']
                   name:
                     description: |-
                       Name of the runtime being referenced.
@@ -2756,6 +2761,7 @@ spec:
                 - name
                 type: object
               suspend:
+                default: false
                 description: |-
                   Whether the controller should suspend the running TrainJob.
                   Defaults to false.

diff --git a/pkg/apis/kubeflow.org/v2alpha1/trainjob_types.go b/pkg/apis/kubeflow.org/v2alpha1/trainjob_types.go
@@ -19,6 +19,7 @@ package v2alpha1
 import (
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"reflect"
 	jobsetv1alpha2 "sigs.k8s.io/jobset/api/jobset/v1alpha2"
 )
 
@@ -87,6 +88,7 @@ type TrainJobSpec struct {
 
 	// Whether the controller should suspend the running TrainJob.
 	// Defaults to false.
+	// +kubebuilder:default=false
 	Suspend *bool `json:"suspend,omitempty"`
 
 	// ManagedBy is used to indicate the controller or entity that manages a TrainJob.
@@ -96,6 +98,9 @@ type TrainJobSpec struct {
 	// `kubeflow.org/trainjob-controller`, but delegates reconciling TrainJobs
 	// with a 'kueue.x-k8s.io/multikueue' to the Kueue. The field is immutable.
 	// Defaults to `kubeflow.org/trainjob-controller`
+	// +kubebuilder:default="kubeflow.org/trainjob-controller"
+	// +kubebuilder:validation:XValidation:rule="self in ['kubeflow.org/trainjob-controller', 'kueue.x-k8s.io/multikueue']", message="ManagedBy must be kubeflow.org/trainjob-controller or kueue.x-k8s.io/multikueue if set"
+	// +kubebuilder:validation:XValidation:rule="self == oldSelf", message="ManagedBy value is immutable"
 	ManagedBy *string `json:"managedBy,omitempty"`
 }
 
@@ -108,11 +113,14 @@ type RuntimeRef struct {
 
 	// APIGroup of the runtime being referenced.
 	// Defaults to `kubeflow.org`.
+	// +kubebuilder:default="kubeflow.org"
 	APIGroup *string `json:"apiGroup,omitempty"`
 
 	// Kind of the runtime being referenced.
 	// It must be one of TrainingRuntime or ClusterTrainingRuntime.
 	// Defaults to ClusterTrainingRuntime.
+	// +kubebuilder:default="ClusterTrainingRuntime"
+	// +kubebuilder:validation:XValidation:rule="self in ['ClusterTrainingRuntime', 'TrainingRuntime']", message="Kind must be ClusterTrainingRuntime or TrainingRuntime if set"
 	Kind *string `json:"kind,omitempty"`
 }
 
@@ -256,6 +264,10 @@ type TrainJobStatus struct {
 	ReplicatedJobsStatus []jobsetv1alpha2.ReplicatedJobStatus `json:"replicatedJobsStatus,omitempty"`
 }
 
+var (
+	TrainJobKind = reflect.TypeOf(TrainJob{}).Name()
+)
+
 func init() {
 	SchemeBuilder.Register(&TrainJob{}, &TrainJobList{})
 }
diff --git a/test/integration/controller.v2/trainjob_controller_test.go b/test/integration/controller.v2/trainjob_controller_test.go
@@ -29,6 +29,8 @@ import (
 
 var _ = ginkgo.Describe("TrainJob controller", ginkgo.Ordered, func() {
 	var ns *corev1.Namespace
+	apiGroup := kubeflowv2.GroupVersion.Group
+	trainingRuntimeKind := kubeflowv2.TrainingRuntimeKind
 
 	ginkgo.BeforeAll(func() {
 		fwk = &framework.Framework{}
@@ -71,4 +73,120 @@ var _ = ginkgo.Describe("TrainJob controller", ginkgo.Ordered, func() {
 			gomega.Expect(k8sClient.Create(ctx, trainJob)).Should(gomega.Succeed())
 		})
 	})
+
+	ginkgo.When("TrainJob CR Validation", func() {
+		ginkgo.AfterEach(func() {
+			gomega.Expect(k8sClient.DeleteAllOf(ctx, &kubeflowv2.TrainJob{}, client.InNamespace(ns.Name))).Should(
+				gomega.Succeed())
+		})
+
+		ginkgo.It("Should succeed in creating TrainJob", func() {
+
+			managedBy := "kubeflow.org/trainjob-controller"
+
+			trainingRuntimeRef := kubeflowv2.RuntimeRef{
+				Name:     "InvalidRuntimeRef",
+				APIGroup: &apiGroup,
+				Kind:     &trainingRuntimeKind,
+			}
+			jobSpec := kubeflowv2.TrainJobSpec{
+				RuntimeRef: trainingRuntimeRef,
+				ManagedBy:  &managedBy,
+			}
+			trainJob := &kubeflowv2.TrainJob{
+				TypeMeta: metav1.TypeMeta{
+					APIVersion: kubeflowv2.SchemeGroupVersion.String(),
+					Kind:       kubeflowv2.TrainJobKind,
+				},
+				ObjectMeta: metav1.ObjectMeta{
+					GenerateName: "valid-trainjob-",
+					Namespace:    ns.Name,
+				},
+				Spec: jobSpec,
+			}
+
+			err := k8sClient.Create(ctx, trainJob)
+			gomega.Expect(err).Should(gomega.Succeed())
+		})
+
+		ginkgo.It("Should fail in creating TrainJob with invalid spec.trainingRuntimeRef", func() {
+
+			kind := "InvalidRuntime"
+
+			trainingRuntimeRef := kubeflowv2.RuntimeRef{
+				Name:     "InvalidRuntimeRef",
+				APIGroup: &apiGroup,
+				Kind:     &kind,
+			}
+			jobSpec := kubeflowv2.TrainJobSpec{
+				RuntimeRef: trainingRuntimeRef,
+			}
+			trainJob := &kubeflowv2.TrainJob{
+				TypeMeta: metav1.TypeMeta{
+					APIVersion: kubeflowv2.SchemeGroupVersion.String(),
+					Kind:       kubeflowv2.TrainJobKind,
+				},
+				ObjectMeta: metav1.ObjectMeta{
+					GenerateName: "invalid-trainjob-",
+					Namespace:    ns.Name,
+				},
+				Spec: jobSpec,
+			}
+			gomega.Expect(k8sClient.Create(ctx, trainJob)).To(gomega.MatchError(
+				gomega.ContainSubstring("spec.runtimeRef.kind: Invalid value")))
+		})
+
+		ginkgo.It("Should fail in creating TrainJob with invalid spec.managedBy", func() {
+			managedBy := "invalidManagedBy"
+			jobSpec := kubeflowv2.TrainJobSpec{
+				ManagedBy: &managedBy,
+			}
+			trainJob := &kubeflowv2.TrainJob{
+				TypeMeta: metav1.TypeMeta{
+					APIVersion: kubeflowv2.SchemeGroupVersion.String(),
+					Kind:       kubeflowv2.TrainJobKind,
+				},
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "invalid-trainjob",
+					Namespace: ns.Name,
+				},
+				Spec: jobSpec,
+			}
+			gomega.Expect(k8sClient.Create(ctx, trainJob)).To(gomega.MatchError(
+				gomega.ContainSubstring("spec.managedBy: Invalid value")))
+		})
+
+		ginkgo.It("Should fail in updating spec.managedBy", func() {
+
+			managedBy := "kubeflow.org/trainjob-controller"
+
+			trainingRuntimeRef := kubeflowv2.RuntimeRef{
+				Name:     "InvalidRuntimeRef",
+				APIGroup: &apiGroup,
+				Kind:     &trainingRuntimeKind,
+			}
+			jobSpec := kubeflowv2.TrainJobSpec{
+				RuntimeRef: trainingRuntimeRef,
+				ManagedBy:  &managedBy,
+			}
+			trainJob := &kubeflowv2.TrainJob{
+				TypeMeta: metav1.TypeMeta{
+					APIVersion: kubeflowv2.SchemeGroupVersion.String(),
+					Kind:       kubeflowv2.TrainJobKind,
+				},
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "job-with-failed-update",
+					Namespace: ns.Name,
+				},
+				Spec: jobSpec,
+			}
+
+			gomega.Expect(k8sClient.Create(ctx, trainJob)).Should(gomega.Succeed())
+			updatedManagedBy := "kueue.x-k8s.io/multikueue"
+			jobSpec.ManagedBy = &updatedManagedBy
+			trainJob.Spec = jobSpec
+			gomega.Expect(k8sClient.Update(ctx, trainJob)).To(gomega.MatchError(
+				gomega.ContainSubstring("ManagedBy value is immutable")))
+		})
+	})
 })