support elastic jobset

pkg/controllers/elastic_jobset.go

@@ -0,0 +1,89 @@
+/*
+Copyright 2024 The Kubernetes Authors.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package controllers
+
+import (
+	"fmt"
+	"slices"
+	"strconv"
+
+	batchv1 "k8s.io/api/batch/v1"
+
+	jobset "sigs.k8s.io/jobset/api/jobset/v1alpha2"
+)
+
+func indexFunc(a, b batchv1.Job) int {
+	jobIndexA, errA := strconv.Atoi(a.Labels[jobset.JobIndexKey])
+	jobIndexB, errB := strconv.Atoi(b.Labels[jobset.JobIndexKey])
+	if errA != nil {
+		return 0
+	}
+	if errB != nil {
+		return 0
+	}
+	if jobIndexA > jobIndexB {
+		return 1
+	} else if jobIndexA < jobIndexB {
+		return -1
+	} else {
+		return 0
+	}
+}
+
+// jobsToDeleteForDownScale gathers the excess jobs during a downscale
+// and deletes the jobs
+func jobsToDeleteForDownScale(replicatedJobs []jobset.ReplicatedJob, replicatedJobStatuses []jobset.ReplicatedJobStatus, jobItems []batchv1.Job) ([]*batchv1.Job, error) {
+	jobsToDelete := []*batchv1.Job{}
+	type payload struct {
+		batchJobs []batchv1.Job
+		rjStatus  jobset.ReplicatedJobStatus
+		replicas  int32
+	}
+	replicatedJobToBatchJobMap := map[string]payload{}
+	for _, replicatedJob := range replicatedJobs {
+		status := findReplicatedJobStatus(replicatedJobStatuses, replicatedJob.Name)
+		newPayload := &payload{}
+		newPayload.rjStatus = status
+		newPayload.replicas = replicatedJob.Replicas
+		for _, val := range jobItems {
+			if val.Labels[jobset.ReplicatedJobNameKey] != replicatedJob.Name {
+				continue
+			}
+			newPayload.batchJobs = append(newPayload.batchJobs, val)
+		}
+		slices.SortFunc(newPayload.batchJobs, indexFunc)
+		replicatedJobToBatchJobMap[replicatedJob.Name] = *newPayload
+	}
+	for _, jobAndStatus := range replicatedJobToBatchJobMap {
+		countOfJobsToDelete := jobAndStatus.rjStatus.Ready - jobAndStatus.replicas
+		if countOfJobsToDelete > 0 {
+			jobsWeDeleted := 0
+			for i := len(jobAndStatus.batchJobs) - 1; i >= 0; i-- {
+
+				jobIndex, err := strconv.Atoi(jobAndStatus.batchJobs[i].Labels[jobset.JobIndexKey])
+				if err != nil {
+					return nil, fmt.Errorf("unable get integer from job index key")
+				}
+				if jobIndex >= int(countOfJobsToDelete) {
+					jobsWeDeleted = jobsWeDeleted + 1
+					jobsToDelete = append(jobsToDelete, &jobAndStatus.batchJobs[i])
+				}
+				if jobsWeDeleted == int(countOfJobsToDelete) {
+					break
+				}
+			}
+		}
+	}
+	return jobsToDelete, nil
+}

pkg/controllers/elastic_jobset_test.go

@@ -0,0 +1,252 @@
+/*
+Copyright 2023 The Kubernetes Authors.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package controllers
+
+import (
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	batchv1 "k8s.io/api/batch/v1"
+	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+
+	jobset "sigs.k8s.io/jobset/api/jobset/v1alpha2"
+)
+
+func TestJobsToDeleteDownScale(t *testing.T) {
+
+	tests := []struct {
+		name                        string
+		replicatedJobs              []jobset.ReplicatedJob
+		replicatedJobStatus         []jobset.ReplicatedJobStatus
+		jobs                        []batchv1.Job
+		expectedJobsThatWereDeleted []batchv1.Job
+		gotError                    error
+	}{
+		{
+			name: "no elastic downscale",
+			replicatedJobs: []jobset.ReplicatedJob{
+				{
+					Name:     "test",
+					Template: batchv1.JobTemplateSpec{},
+					Replicas: 2,
+				},
+			},
+			replicatedJobStatus: []jobset.ReplicatedJobStatus{
+				{
+					Name:  "test",
+					Ready: 1,
+				},
+			},
+			jobs: []batchv1.Job{
+				{
+					ObjectMeta: v1.ObjectMeta{
+						Labels: map[string]string{
+							jobset.ReplicatedJobNameKey: "test",
+							jobset.JobIndexKey:          "0",
+						},
+					},
+				},
+			},
+		},
+		{
+			name: "elastic upscale; do nothing",
+			replicatedJobs: []jobset.ReplicatedJob{
+				{
+					Name:     "test",
+					Template: batchv1.JobTemplateSpec{},
+					Replicas: 2,
+				},
+			},
+			replicatedJobStatus: []jobset.ReplicatedJobStatus{
+				{
+					Name:  "test",
+					Ready: 1,
+				},
+			},
+			jobs: []batchv1.Job{
+				{
+					ObjectMeta: v1.ObjectMeta{
+						Labels: map[string]string{
+							jobset.ReplicatedJobNameKey: "test",
+							jobset.JobIndexKey:          "0",
+						},
+					},
+				},
+			},
+		},
+		{
+			name: "elastic downscale is needed",
+			replicatedJobs: []jobset.ReplicatedJob{
+				{
+					Name:     "test",
+					Template: batchv1.JobTemplateSpec{},
+					Replicas: 1,
+				},
+			},
+			replicatedJobStatus: []jobset.ReplicatedJobStatus{
+				{
+					Name:  "test",
+					Ready: 2,
+				},
+			},
+			jobs: []batchv1.Job{
+				{
+					ObjectMeta: v1.ObjectMeta{
+						Labels: map[string]string{
+							jobset.ReplicatedJobNameKey: "test",
+							jobset.JobIndexKey:          "0",
+						},
+					},
+				},
+				{
+					ObjectMeta: v1.ObjectMeta{
+						Labels: map[string]string{
+							jobset.ReplicatedJobNameKey: "test",
+							jobset.JobIndexKey:          "1",
+						},
+					},
+				},
+			},
+			expectedJobsThatWereDeleted: []batchv1.Job{
+				{
+					ObjectMeta: v1.ObjectMeta{
+						Labels: map[string]string{
+							jobset.ReplicatedJobNameKey: "test",
+							jobset.JobIndexKey:          "1",
+						},
+					},
+				},
+			},
+		},
+		{
+			name: "elastic downscale is needed for second replicated job",
+			replicatedJobs: []jobset.ReplicatedJob{
+				{
+					Name:     "test",
+					Template: batchv1.JobTemplateSpec{},
+					Replicas: 2,
+				},
+				{
+					Name:     "test-2",
+					Template: batchv1.JobTemplateSpec{},
+					Replicas: 2,
+				},
+			},
+			replicatedJobStatus: []jobset.ReplicatedJobStatus{
+				{
+					Name:  "test",
+					Ready: 2,
+				},
+				{
+					Name:  "test-2",
+					Ready: 4,
+				},
+			},
+			jobs: []batchv1.Job{
+				{
+					ObjectMeta: v1.ObjectMeta{
+						Labels: map[string]string{
+							jobset.ReplicatedJobNameKey: "test",
+							jobset.JobIndexKey:          "0",
+						},
+					},
+				},
+				{
+					ObjectMeta: v1.ObjectMeta{
+						Labels: map[string]string{
+							jobset.ReplicatedJobNameKey: "test",
+							jobset.JobIndexKey:          "1",
+						},
+					},
+				},
+				{
+					ObjectMeta: v1.ObjectMeta{
+						Labels: map[string]string{
+							jobset.ReplicatedJobNameKey: "test-2",
+							jobset.JobIndexKey:          "2",
+						},
+					},
+				},
+				{
+					ObjectMeta: v1.ObjectMeta{
+						Labels: map[string]string{
+							jobset.ReplicatedJobNameKey: "test-2",
+							jobset.JobIndexKey:          "3",
+						},
+					},
+				},
+				{
+					ObjectMeta: v1.ObjectMeta{
+						Labels: map[string]string{
+							jobset.ReplicatedJobNameKey: "test-2",
+							jobset.JobIndexKey:          "0",
+						},
+					},
+				},
+				{
+					ObjectMeta: v1.ObjectMeta{
+						Labels: map[string]string{
+							jobset.ReplicatedJobNameKey: "test-2",
+							jobset.JobIndexKey:          "1",
+						},
+					},
+				},
+			},
+			expectedJobsThatWereDeleted: []batchv1.Job{
+				{
+					ObjectMeta: v1.ObjectMeta{
+						Labels: map[string]string{
+							jobset.ReplicatedJobNameKey: "test-2",
+							jobset.JobIndexKey:          "3",
+						},
+					},
+				},
+				{
+					ObjectMeta: v1.ObjectMeta{
+						Labels: map[string]string{
+							jobset.ReplicatedJobNameKey: "test-2",
+							jobset.JobIndexKey:          "2",
+						},
+					},
+				},
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			actual, err := jobsToDeleteForDownScale(tc.replicatedJobs, tc.replicatedJobStatus, tc.jobs)
+			if diff := cmp.Diff(tc.gotError, err); diff != "" {
+				t.Errorf("unexpected finished value (+got/-want): %s", diff)
+			}
+			if len(actual) != len(tc.expectedJobsThatWereDeleted) {
+				t.Errorf("unexpected length mismatch for deleted jobs: got: %d want: %d", len(actual), len(tc.expectedJobsThatWereDeleted))
+			}
+			if tc.expectedJobsThatWereDeleted != nil {
+				for i := range actual {
+					actualReplicatedJobName := actual[i].ObjectMeta.Labels[jobset.ReplicatedJobNameKey]
+					actualJobIndexKey := actual[i].ObjectMeta.Labels[jobset.JobIndexKey]
+					expectedReplicatedJobName := tc.expectedJobsThatWereDeleted[i].ObjectMeta.Labels[jobset.ReplicatedJobNameKey]
+					expectedJobIndexKey := tc.expectedJobsThatWereDeleted[i].ObjectMeta.Labels[jobset.JobIndexKey]
+					if diff := cmp.Diff(actualReplicatedJobName, expectedReplicatedJobName); diff != "" {
+						t.Errorf("unexpected replicated job name (+got/-want): %s", diff)
+					}
+					if diff := cmp.Diff(actualJobIndexKey, expectedJobIndexKey); diff != "" {
+						t.Errorf("unexpected job index (+got/-want): %s", diff)
+					}
+				}
+			}
+		})
+	}
+}

pkg/controllers/jobset_controller.go

@@ -203,6 +203,13 @@ func (r *JobSetReconciler) reconcile(ctx context.Context, js *jobset.JobSet, upd
 		return ctrl.Result{}, err
 	}
 
+	// ElasticJobSet can downscale and upscale Jobs of a given replicated job
+	// On downscale, we need to gather the jobs to delete.
+	if err := r.downscaleElasticJobs(ctx, js, rjobStatuses); err != nil {
+		log.Error(err, "unable to downscale elastic jobs")
+		return ctrl.Result{}, err
+	}
+
 	// Handle suspending a jobset or resuming a suspended jobset.
 	jobsetSuspended := jobSetSuspended(js)
 	if jobsetSuspended {
@@ -520,6 +527,22 @@ func (r *JobSetReconciler) reconcileReplicatedJobs(ctx context.Context, js *jobs
 	return nil
 }
 
+// We need to check if the replicas of a replicated job are mismatching
+// If they are, we should delete the extra jobs
+func (r *JobSetReconciler) downscaleElasticJobs(ctx context.Context, js *jobset.JobSet, replicatedJobStatus []jobset.ReplicatedJobStatus) error {
+	// Get all active jobs owned by JobSet.
+	var childJobList batchv1.JobList
+	if err := r.List(ctx, &childJobList, client.InNamespace(js.Namespace), client.MatchingFields{constants.JobOwnerKey: js.Name}); err != nil {
+		return err
+	}
+
+	jobsToDelete, err := jobsToDeleteForDownScale(js.Spec.ReplicatedJobs, replicatedJobStatus, childJobList.Items)
+	if err != nil {
+		return err
+	}
+	return r.deleteJobs(ctx, jobsToDelete)
+}
+
 func (r *JobSetReconciler) createJobs(ctx context.Context, js *jobset.JobSet, jobs []*batchv1.Job) error {
 	log := ctrl.LoggerFrom(ctx)