Extract Trainjob from the Kubeflow directory (#7081)

Author: kaisoz

charts/kueue/templates/webhook/manifests.yaml

@@ -242,26 +242,6 @@ webhooks:
           - xgboostjobs
     sideEffects: None
     reinvocationPolicy: '{{ .Values.mutatingWebhook.reinvocationPolicy }}'
-  - admissionReviewVersions:
-      - v1
-    clientConfig:
-      service:
-        name: '{{ include "kueue.fullname" . }}-webhook-service'
-        namespace: '{{ .Release.Namespace }}'
-        path: /mutate-trainer-kubeflow-org-v1alpha1-trainjob
-    failurePolicy: Fail
-    name: mtrainjob.kb.io
-    rules:
-      - apiGroups:
-          - trainer.kubeflow.org
-        apiVersions:
-          - v1alpha1
-        operations:
-          - CREATE
-        resources:
-          - trainjobs
-    sideEffects: None
-    reinvocationPolicy: '{{ .Values.mutatingWebhook.reinvocationPolicy }}'
   - admissionReviewVersions:
       - v1
     clientConfig:
@@ -430,6 +410,26 @@ webhooks:
           - statefulsets
     sideEffects: None
     reinvocationPolicy: '{{ .Values.mutatingWebhook.reinvocationPolicy }}'
+  - admissionReviewVersions:
+      - v1
+    clientConfig:
+      service:
+        name: '{{ include "kueue.fullname" . }}-webhook-service'
+        namespace: '{{ .Release.Namespace }}'
+        path: /mutate-trainer-kubeflow-org-v1alpha1-trainjob
+    failurePolicy: Fail
+    name: mtrainjob.kb.io
+    rules:
+      - apiGroups:
+          - trainer.kubeflow.org
+        apiVersions:
+          - v1alpha1
+        operations:
+          - CREATE
+        resources:
+          - trainjobs
+    sideEffects: None
+    reinvocationPolicy: '{{ .Values.mutatingWebhook.reinvocationPolicy }}'
   - admissionReviewVersions:
       - v1
     clientConfig:
@@ -730,26 +730,6 @@ webhooks:
         resources:
           - xgboostjobs
     sideEffects: None
-  - admissionReviewVersions:
-      - v1
-    clientConfig:
-      service:
-        name: '{{ include "kueue.fullname" . }}-webhook-service'
-        namespace: '{{ .Release.Namespace }}'
-        path: /validate-trainer-kubeflow-org-v1alpha1-trainjob
-    failurePolicy: Fail
-    name: vtrainjob.kb.io
-    rules:
-      - apiGroups:
-          - trainer.kubeflow.org
-        apiVersions:
-          - v1alpha1
-        operations:
-          - CREATE
-          - UPDATE
-        resources:
-          - trainjobs
-    sideEffects: None
   - admissionReviewVersions:
       - v1
     clientConfig:
@@ -916,6 +896,26 @@ webhooks:
         resources:
           - statefulsets
     sideEffects: None
+  - admissionReviewVersions:
+      - v1
+    clientConfig:
+      service:
+        name: '{{ include "kueue.fullname" . }}-webhook-service'
+        namespace: '{{ .Release.Namespace }}'
+        path: /validate-trainer-kubeflow-org-v1alpha1-trainjob
+    failurePolicy: Fail
+    name: vtrainjob.kb.io
+    rules:
+      - apiGroups:
+          - trainer.kubeflow.org
+        apiVersions:
+          - v1alpha1
+        operations:
+          - CREATE
+          - UPDATE
+        resources:
+          - trainjobs
+    sideEffects: None
   - admissionReviewVersions:
       - v1
     clientConfig:

config/components/webhook/manifests.yaml

@@ -176,25 +176,6 @@ webhooks:
     resources:
     - xgboostjobs
   sideEffects: None
-- admissionReviewVersions:
-  - v1
-  clientConfig:
-    service:
-      name: webhook-service
-      namespace: system
-      path: /mutate-trainer-kubeflow-org-v1alpha1-trainjob
-  failurePolicy: Fail
-  name: mtrainjob.kb.io
-  rules:
-  - apiGroups:
-    - trainer.kubeflow.org
-    apiVersions:
-    - v1alpha1
-    operations:
-    - CREATE
-    resources:
-    - trainjobs
-  sideEffects: None
 - admissionReviewVersions:
   - v1
   clientConfig:
@@ -311,6 +292,25 @@ webhooks:
     resources:
     - statefulsets
   sideEffects: None
+- admissionReviewVersions:
+  - v1
+  clientConfig:
+    service:
+      name: webhook-service
+      namespace: system
+      path: /mutate-trainer-kubeflow-org-v1alpha1-trainjob
+  failurePolicy: Fail
+  name: mtrainjob.kb.io
+  rules:
+  - apiGroups:
+    - trainer.kubeflow.org
+    apiVersions:
+    - v1alpha1
+    operations:
+    - CREATE
+    resources:
+    - trainjobs
+  sideEffects: None
 - admissionReviewVersions:
   - v1
   clientConfig:
@@ -554,26 +554,6 @@ webhooks:
     resources:
     - xgboostjobs
   sideEffects: None
-- admissionReviewVersions:
-  - v1
-  clientConfig:
-    service:
-      name: webhook-service
-      namespace: system
-      path: /validate-trainer-kubeflow-org-v1alpha1-trainjob
-  failurePolicy: Fail
-  name: vtrainjob.kb.io
-  rules:
-  - apiGroups:
-    - trainer.kubeflow.org
-    apiVersions:
-    - v1alpha1
-    operations:
-    - CREATE
-    - UPDATE
-    resources:
-    - trainjobs
-  sideEffects: None
 - admissionReviewVersions:
   - v1
   clientConfig:
@@ -694,6 +674,26 @@ webhooks:
     resources:
     - statefulsets
   sideEffects: None
+- admissionReviewVersions:
+  - v1
+  clientConfig:
+    service:
+      name: webhook-service
+      namespace: system
+      path: /validate-trainer-kubeflow-org-v1alpha1-trainjob
+  failurePolicy: Fail
+  name: vtrainjob.kb.io
+  rules:
+  - apiGroups:
+    - trainer.kubeflow.org
+    apiVersions:
+    - v1alpha1
+    operations:
+    - CREATE
+    - UPDATE
+    resources:
+    - trainjobs
+  sideEffects: None
 - admissionReviewVersions:
   - v1
   clientConfig:

pkg/controller/jobs/jobs.go

@@ -23,11 +23,11 @@ import (
 	_ "sigs.k8s.io/kueue/pkg/controller/jobs/job"
 	_ "sigs.k8s.io/kueue/pkg/controller/jobs/jobset"
 	_ "sigs.k8s.io/kueue/pkg/controller/jobs/kubeflow/jobs"
-	_ "sigs.k8s.io/kueue/pkg/controller/jobs/kubeflow/trainjob"
 	_ "sigs.k8s.io/kueue/pkg/controller/jobs/leaderworkerset"
 	_ "sigs.k8s.io/kueue/pkg/controller/jobs/mpijob"
 	_ "sigs.k8s.io/kueue/pkg/controller/jobs/pod"
 	_ "sigs.k8s.io/kueue/pkg/controller/jobs/raycluster"
 	_ "sigs.k8s.io/kueue/pkg/controller/jobs/rayjob"
 	_ "sigs.k8s.io/kueue/pkg/controller/jobs/statefulset"
+	_ "sigs.k8s.io/kueue/pkg/controller/jobs/trainjob"
 )

…kubeflow/trainjob/trainjob_controller.go …ler/jobs/trainjob/trainjob_controller.gopkg/controller/jobs/kubeflow/trainjob/trainjob_controller.go renamed to pkg/controller/jobs/trainjob/trainjob_controller.go pkg/controller/jobs/kubeflow/trainjob/trainjob_controller.go renamed to pkg/controller/jobs/trainjob/trainjob_controller.go

@@ -23,7 +23,7 @@ import (
 	"fmt"
 	"strconv"
 
-	kftrainerapi "github.com/kubeflow/trainer/v2/pkg/apis/trainer/v1alpha1"
+	kftrainer "github.com/kubeflow/trainer/v2/pkg/apis/trainer/v1alpha1"
 	kftrainerruntime "github.com/kubeflow/trainer/v2/pkg/runtime"
 	kftrainerruntimecore "github.com/kubeflow/trainer/v2/pkg/runtime/core"
 	kftrainerjobset "github.com/kubeflow/trainer/v2/pkg/runtime/framework/plugins/jobset"
@@ -42,15 +42,15 @@ import (
 
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
 	"sigs.k8s.io/kueue/pkg/controller/jobframework"
-	kJobset "sigs.k8s.io/kueue/pkg/controller/jobs/jobset"
+	workloadjobset "sigs.k8s.io/kueue/pkg/controller/jobs/jobset"
 	"sigs.k8s.io/kueue/pkg/features"
 	"sigs.k8s.io/kueue/pkg/podset"
 	clientutil "sigs.k8s.io/kueue/pkg/util/client"
 	"sigs.k8s.io/kueue/pkg/util/slices"
 )
 
 var (
-	gvk                    = kftrainerapi.GroupVersion.WithKind("TrainJob")
+	gvk                    = kftrainer.GroupVersion.WithKind("TrainJob")
 	FrameworkName          = "trainer.kubeflow.org/trainjob"
 	TrainJobControllerName = "trainer.kubeflow.org/trainjob-controller"
 )
@@ -66,8 +66,8 @@ func init() {
 		NewJob:            NewJob,
 		NewReconciler:     NewReconciler,
 		SetupWebhook:      SetupTrainJobWebhook,
-		JobType:           &kftrainerapi.TrainJob{},
-		AddToScheme:       kftrainerapi.AddToScheme,
+		JobType:           &kftrainer.TrainJob{},
+		AddToScheme:       kftrainer.AddToScheme,
 		MultiKueueAdapter: &multiKueueAdapter{},
 	}))
 }
@@ -108,11 +108,11 @@ func (r *trainJobReconciler) Reconcile(ctx context.Context, req ctrl.Request) (c
 
 func (r *trainJobReconciler) SetupWithManager(mgr ctrl.Manager) error {
 	b := ctrl.NewControllerManagedBy(mgr).
-		For(&kftrainerapi.TrainJob{}).Owns(&kueue.Workload{}).Owns(&jobsetapi.JobSet{})
+		For(&kftrainer.TrainJob{}).Owns(&kueue.Workload{}).Owns(&jobsetapi.JobSet{})
 	return b.Complete(r)
 }
 
-type TrainJob kftrainerapi.TrainJob
+type TrainJob kftrainer.TrainJob
 
 var _ jobframework.GenericJob = (*TrainJob)(nil)
 var _ jobframework.JobWithCustomStop = (*TrainJob)(nil)
@@ -124,11 +124,11 @@ func NewJob() jobframework.GenericJob {
 }
 
 func fromObject(obj runtime.Object) *TrainJob {
-	return (*TrainJob)(obj.(*kftrainerapi.TrainJob))
+	return (*TrainJob)(obj.(*kftrainer.TrainJob))
 }
 
 func (t *TrainJob) Object() client.Object {
-	return (*kftrainerapi.TrainJob)(t)
+	return (*kftrainer.TrainJob)(t)
 }
 
 func (t *TrainJob) IsSuspended() bool {
@@ -167,7 +167,7 @@ func getChildJobSet(t *TrainJob) (*jobsetapi.JobSet, error) {
 		return nil, fmt.Errorf("unsupported runtime: %s", runtimeRefGK)
 	}
 
-	trainJob := (*kftrainerapi.TrainJob)(t)
+	trainJob := (*kftrainer.TrainJob)(t)
 	trSpec, err := getRuntimeSpec(trainJob)
 	if err != nil {
 		return nil, fmt.Errorf("runtime '%s' not found", trainJob.Spec.RuntimeRef.Name)
@@ -212,16 +212,16 @@ func jobsetApplyToJobset(jobsetApply *jobsetapplyapi.JobSetApplyConfiguration) (
 	return jobset, nil
 }
 
-func getRuntimeSpec(trainJob *kftrainerapi.TrainJob) (*kftrainerapi.TrainingRuntimeSpec, error) {
-	if *trainJob.Spec.RuntimeRef.Kind == kftrainerapi.ClusterTrainingRuntimeKind {
-		var ctr kftrainerapi.ClusterTrainingRuntime
+func getRuntimeSpec(trainJob *kftrainer.TrainJob) (*kftrainer.TrainingRuntimeSpec, error) {
+	if *trainJob.Spec.RuntimeRef.Kind == kftrainer.ClusterTrainingRuntimeKind {
+		var ctr kftrainer.ClusterTrainingRuntime
 		err := reconciler.client.Get(reconciler.ctx, client.ObjectKey{Name: trainJob.Spec.RuntimeRef.Name}, &ctr)
 		if err != nil {
 			return nil, err
 		}
 		return &ctr.Spec, nil
 	} else {
-		var tr kftrainerapi.TrainingRuntime
+		var tr kftrainer.TrainingRuntime
 		err := reconciler.client.Get(reconciler.ctx, client.ObjectKey{Namespace: trainJob.Namespace, Name: trainJob.Spec.RuntimeRef.Name}, &tr)
 		if err != nil {
 			return nil, err
@@ -235,7 +235,7 @@ func (t *TrainJob) PodSets() ([]kueue.PodSet, error) {
 	if err != nil {
 		return nil, err
 	}
-	return (*kJobset.JobSet)(jobset).PodSets()
+	return (*workloadjobset.JobSet)(jobset).PodSets()
 }
 
 func (t *TrainJob) RunWithPodSetsInfo(podSetsInfo []podset.PodSetInfo) error {
@@ -249,16 +249,16 @@ func (t *TrainJob) RunWithPodSetsInfo(podSetsInfo []podset.PodSetInfo) error {
 	}
 
 	if t.Spec.PodSpecOverrides == nil {
-		t.Spec.PodSpecOverrides = []kftrainerapi.PodSpecOverride{}
+		t.Spec.PodSpecOverrides = []kftrainer.PodSpecOverride{}
 	}
 	if t.Annotations == nil {
 		t.Annotations = map[string]string{}
 	}
 	t.Annotations[firstOverrideIdx] = strconv.Itoa(len(t.Spec.PodSpecOverrides))
 	for _, info := range podSetsInfo {
 		// The trainjob controller merges each podSpecOverride sequentially, so any existing user provided override will be processed first
-		t.Spec.PodSpecOverrides = append(t.Spec.PodSpecOverrides, kftrainerapi.PodSpecOverride{
-			TargetJobs: []kftrainerapi.PodSpecOverrideTargetJob{
+		t.Spec.PodSpecOverrides = append(t.Spec.PodSpecOverrides, kftrainer.PodSpecOverride{
+			TargetJobs: []kftrainer.PodSpecOverrideTargetJob{
 				{Name: string(info.Name)},
 			},
 			// TODO: Set the labels/annotations when supported. See https://github.com/kubeflow/trainer/pull/2785
@@ -320,10 +320,10 @@ func (t *TrainJob) RestorePodSetsInfo(_ []podset.PodSetInfo) bool {
 }
 
 func (t *TrainJob) Finished() (message string, success, finished bool) {
-	if c := apimeta.FindStatusCondition(t.Status.Conditions, kftrainerapi.TrainJobComplete); c != nil && c.Status == metav1.ConditionTrue {
+	if c := apimeta.FindStatusCondition(t.Status.Conditions, kftrainer.TrainJobComplete); c != nil && c.Status == metav1.ConditionTrue {
 		return c.Message, true, true
 	}
-	if c := apimeta.FindStatusCondition(t.Status.Conditions, kftrainerapi.TrainJobFailed); c != nil && c.Status == metav1.ConditionTrue {
+	if c := apimeta.FindStatusCondition(t.Status.Conditions, kftrainer.TrainJobFailed); c != nil && c.Status == metav1.ConditionTrue {
 		return c.Message, false, true
 	}
 	return message, success, false
@@ -356,15 +356,15 @@ func (t *TrainJob) ReclaimablePods() ([]kueue.ReclaimablePod, error) {
 	}
 
 	ret := make([]kueue.ReclaimablePod, 0, len(jobset.Spec.ReplicatedJobs))
-	statuses := slices.ToRefMap(t.Status.JobsStatus, func(js *kftrainerapi.JobStatus) string { return js.Name })
+	statuses := slices.ToRefMap(t.Status.JobsStatus, func(js *kftrainer.JobStatus) string { return js.Name })
 
 	for i := range jobset.Spec.ReplicatedJobs {
 		spec := &jobset.Spec.ReplicatedJobs[i]
 		if status, found := statuses[spec.Name]; found && status.Succeeded > 0 {
 			if status.Succeeded > 0 && status.Succeeded <= spec.Replicas {
 				ret = append(ret, kueue.ReclaimablePod{
 					Name:  kueue.NewPodSetReference(spec.Name),
-					Count: status.Succeeded * kJobset.PodsCountPerReplica(spec),
+					Count: status.Succeeded * workloadjobset.PodsCountPerReplica(spec),
 				})
 			}
 		}

…low/trainjob/trainjob_controller_test.go …obs/trainjob/trainjob_controller_test.gopkg/controller/jobs/kubeflow/trainjob/trainjob_controller_test.go renamed to pkg/controller/jobs/trainjob/trainjob_controller_test.go pkg/controller/jobs/kubeflow/trainjob/trainjob_controller_test.go renamed to pkg/controller/jobs/trainjob/trainjob_controller_test.go

@@ -46,12 +46,12 @@ import (
 	workloadjob "sigs.k8s.io/kueue/pkg/controller/jobs/job"
 	workloadjobset "sigs.k8s.io/kueue/pkg/controller/jobs/jobset"
 	workloadpytorchjob "sigs.k8s.io/kueue/pkg/controller/jobs/kubeflow/jobs/pytorchjob"
-	workloadtrainjob "sigs.k8s.io/kueue/pkg/controller/jobs/kubeflow/trainjob"
 	workloadmpijob "sigs.k8s.io/kueue/pkg/controller/jobs/mpijob"
 	workloadpod "sigs.k8s.io/kueue/pkg/controller/jobs/pod"
 	podconstants "sigs.k8s.io/kueue/pkg/controller/jobs/pod/constants"
 	workloadraycluster "sigs.k8s.io/kueue/pkg/controller/jobs/raycluster"
 	workloadrayjob "sigs.k8s.io/kueue/pkg/controller/jobs/rayjob"
+	workloadtrainjob "sigs.k8s.io/kueue/pkg/controller/jobs/trainjob"
 	"sigs.k8s.io/kueue/pkg/util/admissioncheck"
 	utilpod "sigs.k8s.io/kueue/pkg/util/pod"
 	utiltesting "sigs.k8s.io/kueue/pkg/util/testing"