KEP-2170: Expose the TrainingRuntime and ClusterTrainingRuntime Kind

Signed-off-by: Yuki Iwai <[email protected]>
kubeflow · Oct 10, 2024 · f146a7d · f146a7d
1 parent 23d7b44
commit f146a7d
Show file tree

Hide file tree

Showing 8 changed files with 19 additions and 11 deletions.
diff --git a/pkg/apis/kubeflow.org/v2alpha1/trainingruntime_types.go b/pkg/apis/kubeflow.org/v2alpha1/trainingruntime_types.go
@@ -22,6 +22,13 @@ import (
 	jobsetv1alpha2 "sigs.k8s.io/jobset/api/jobset/v1alpha2"
 )
 
+const (
+	// TrainingRuntimeKind is the GroupVersionKind Kind name for the TrainingRuntime.
+	TrainingRuntimeKind string = "TrainingRuntime"
+	// ClusterTrainingRuntimeKind is the GroupVersionKind Kind name for the ClusterTrainingRuntime.
+	ClusterTrainingRuntimeKind string = "ClusterTrainingRuntime"
+)
+
 // +kubebuilder:object:root=true
 // +kubebuilder:storageversion
 // +kubebuilder:resource:scope=Cluster

diff --git a/pkg/runtime.v2/core/clustertrainingruntime.go b/pkg/runtime.v2/core/clustertrainingruntime.go
@@ -42,7 +42,7 @@ var _ runtime.Runtime = (*ClusterTrainingRuntime)(nil)
 
 var ClusterTrainingRuntimeGroupKind = schema.GroupKind{
 	Group: kubeflowv2.GroupVersion.Group,
-	Kind:  "ClusterTrainingRuntime",
+	Kind:  kubeflowv2.ClusterTrainingRuntimeKind,
 }.String()
 
 func NewClusterTrainingRuntime(context.Context, client.Client, client.FieldIndexer) (runtime.Runtime, error) {

diff --git a/pkg/runtime.v2/core/clustertrainingruntime_test.go b/pkg/runtime.v2/core/clustertrainingruntime_test.go
@@ -47,7 +47,7 @@ func TestClusterTrainingRuntimeNewObjects(t *testing.T) {
 		"succeeded to build JobSet and PodGroup": {
 			trainJob: testingutil.MakeTrainJobWrapper(t, metav1.NamespaceDefault, "test-job").
 				UID("uid").
-				TrainingRuntimeRef(kubeflowv2.SchemeGroupVersion.WithKind("ClusterTrainingRuntime"), "test-runtime").
+				TrainingRuntimeRef(kubeflowv2.SchemeGroupVersion.WithKind(kubeflowv2.ClusterTrainingRuntimeKind), "test-runtime").
 				Trainer(
 					testingutil.MakeTrainJobTrainerWrapper(t).
 						ContainerImage("test:trainjob").
@@ -93,7 +93,7 @@ func TestClusterTrainingRuntimeNewObjects(t *testing.T) {
 		"missing trainingRuntime resource": {
 			trainJob: testingutil.MakeTrainJobWrapper(t, metav1.NamespaceDefault, "test-job").
 				UID("uid").
-				TrainingRuntimeRef(kubeflowv2.SchemeGroupVersion.WithKind("ClusterTrainingRuntime"), "test-runtime").
+				TrainingRuntimeRef(kubeflowv2.SchemeGroupVersion.WithKind(kubeflowv2.ClusterTrainingRuntimeKind), "test-runtime").
 				Trainer(
 					testingutil.MakeTrainJobTrainerWrapper(t).
 						ContainerImage("test:trainjob").

diff --git a/pkg/runtime.v2/core/trainingruntime.go b/pkg/runtime.v2/core/trainingruntime.go
@@ -49,7 +49,7 @@ type TrainingRuntime struct {
 
 var TrainingRuntimeGroupKind = schema.GroupKind{
 	Group: kubeflowv2.GroupVersion.Group,
-	Kind:  "TrainingRuntime",
+	Kind:  kubeflowv2.TrainingRuntimeKind,
 }.String()
 
 var _ runtime.Runtime = (*TrainingRuntime)(nil)

diff --git a/pkg/runtime.v2/core/trainingruntime_test.go b/pkg/runtime.v2/core/trainingruntime_test.go
@@ -47,7 +47,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 		"succeeded to build JobSet and PodGroup": {
 			trainJob: testingutil.MakeTrainJobWrapper(t, metav1.NamespaceDefault, "test-job").
 				UID("uid").
-				TrainingRuntimeRef(kubeflowv2.SchemeGroupVersion.WithKind("TrainingRuntime"), "test-runtime").
+				TrainingRuntimeRef(kubeflowv2.SchemeGroupVersion.WithKind(kubeflowv2.TrainingRuntimeKind), "test-runtime").
 				Trainer(
 					testingutil.MakeTrainJobTrainerWrapper(t).
 						ContainerImage("test:trainjob").
@@ -93,7 +93,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 		"missing trainingRuntime resource": {
 			trainJob: testingutil.MakeTrainJobWrapper(t, metav1.NamespaceDefault, "test-job").
 				UID("uid").
-				TrainingRuntimeRef(kubeflowv2.SchemeGroupVersion.WithKind("TrainingRuntime"), "test-runtime").
+				TrainingRuntimeRef(kubeflowv2.SchemeGroupVersion.WithKind(kubeflowv2.TrainingRuntimeKind), "test-runtime").
 				Trainer(
 					testingutil.MakeTrainJobTrainerWrapper(t).
 						ContainerImage("test:trainjob").

diff --git a/pkg/runtime.v2/indexer/indexer.go b/pkg/runtime.v2/indexer/indexer.go
@@ -38,7 +38,7 @@ func IndexTrainJobTrainingRuntimes(obj client.Object) []string {
 		Kind:  ptr.Deref(trainJob.Spec.TrainingRuntimeRef.Kind, ""),
 	}
 	if runtimeRefGroupKind.Group == kubeflowv2.GroupVersion.Group &&
-		(runtimeRefGroupKind.Kind == "TrainingRuntime" || runtimeRefGroupKind.Kind == "ClusterTrainingRuntime") {
+		(runtimeRefGroupKind.Kind == kubeflowv2.TrainingRuntimeKind || runtimeRefGroupKind.Kind == kubeflowv2.ClusterTrainingRuntimeKind) {
 		return []string{trainJob.Spec.TrainingRuntimeRef.Name}
 	}
 	return nil

diff --git a/pkg/util.v2/testing/wrapper.go b/pkg/util.v2/testing/wrapper.go
@@ -238,7 +238,7 @@ func MakeTrainingRuntimeWrapper(t *testing.T, namespace, name string) *TrainingR
 		TrainingRuntime: kubeflowv2.TrainingRuntime{
 			TypeMeta: metav1.TypeMeta{
 				APIVersion: kubeflowv2.SchemeGroupVersion.String(),
-				Kind:       "TrainingRuntime",
+				Kind:       kubeflowv2.TrainingRuntimeKind,
 			},
 			ObjectMeta: metav1.ObjectMeta{
 				Namespace: namespace,
@@ -309,7 +309,7 @@ func MakeClusterTrainingRuntimeWrapper(t *testing.T, name string) *ClusterTraini
 		ClusterTrainingRuntime: kubeflowv2.ClusterTrainingRuntime{
 			TypeMeta: metav1.TypeMeta{
 				APIVersion: kubeflowv2.SchemeGroupVersion.String(),
-				Kind:       "ClusterTrainingRuntime",
+				Kind:       kubeflowv2.ClusterTrainingRuntimeKind,
 			},
 			ObjectMeta: metav1.ObjectMeta{
 				Name: name,

diff --git a/pkg/webhook.v2/setup.go b/pkg/webhook.v2/setup.go
@@ -19,15 +19,16 @@ package webhookv2
 import (
 	ctrl "sigs.k8s.io/controller-runtime"
 
+	kubeflowv2 "github.com/kubeflow/training-operator/pkg/apis/kubeflow.org/v2alpha1"
 	runtime "github.com/kubeflow/training-operator/pkg/runtime.v2"
 )
 
 func Setup(mgr ctrl.Manager, runtimes map[string]runtime.Runtime) (string, error) {
 	if err := setupWebhookForClusterTrainingRuntime(mgr, runtimes); err != nil {
-		return "ClusterTrainingRuntime", err
+		return kubeflowv2.ClusterTrainingRuntimeKind, err
 	}
 	if err := setupWebhookForTrainingRuntime(mgr, runtimes); err != nil {
-		return "TrainingRuntime", err
+		return kubeflowv2.TrainingRuntimeKind, err
 	}
 	if err := setupWebhookForTrainJob(mgr, runtimes); err != nil {
 		return "TrainJob", err