Add autoscaling

Author: adityapb

Dockerfile

@@ -7,6 +7,8 @@ ARG RELEASE_VERSION
 
 ADD . /go/src/github.com/kubeflow/mpi-operator
 WORKDIR /go/src/github.com/kubeflow/mpi-operator
+RUN apt update
+RUN apt install -y build-essential cmake zlib1g-dev
 RUN make RELEASE_VERSION=${RELEASE_VERSION} mpi-operator.$VERSION
 RUN ln -s mpi-operator.${VERSION} _output/cmd/bin/mpi-operator
 

Makefile

@@ -14,6 +14,7 @@
 
 BIN_DIR=_output/cmd/bin
 REPO_PATH="github.com/kubeflow/mpi-operator"
+CHARM_REPO="https://github.com/charmplusplus/charm.git"
 REL_OSARCH="linux/amd64"
 GitSHA=$(shell git rev-parse HEAD)
 Date=$(shell date "+%Y-%m-%d %H:%M:%S")
@@ -47,10 +48,10 @@ CRD_OPTIONS ?= "crd:generateEmbeddedObjectMeta=true"
 
 build: all
 
-all: ${BIN_DIR} fmt vet tidy lint test mpi-operator.v2
+all: ${BIN_DIR} fmt vet tidy lint test charm mpi-operator.v2
 
 .PHONY: mpi-operator.v2
-mpi-operator.v2:
+mpi-operator.v2: charm
 	go build -ldflags ${LD_FLAGS_V2} -o ${BIN_DIR}/mpi-operator.v2 ./cmd/mpi-operator/
 
 ${BIN_DIR}:
@@ -209,3 +210,11 @@ volcano-scheduler-crd: volcano-scheduler
 volcano-scheduler-deploy: volcano-scheduler-crd
 	mkdir -p $(PROJECT_DIR)/dep-manifests/volcano-scheduler/
 	cp -f /tmp/volcano.sh/volcano/installer/volcano-development.yaml $(PROJECT_DIR)/dep-manifests/volcano-scheduler/
+
+.PHONY: charm
+charm: ${BIN_DIR}
+	rm -rf $(PROJECT_DIR)/dep-libs/charm
+	mkdir -p $(PROJECT_DIR)/dep-libs/charm
+	git clone $(CHARM_REPO) $(PROJECT_DIR)/dep-libs/charm
+	cd $(PROJECT_DIR)/dep-libs/charm && git checkout shrinkexpand-fix && ./build charm++ netlrts-linux-x86_64 --enable-shrinkexpand -j8 --force
+	cd pkg/controller && $(PROJECT_DIR)/dep-libs/charm/bin/charmc -language c++ -seq -o ${PROJECT_DIR}/${BIN_DIR}/rescale_client rescale_client.C -lccs-client

cmd/mpi-operator/app/server.go

@@ -39,6 +39,7 @@ import (
 	"k8s.io/client-go/tools/leaderelection/resourcelock"
 	"k8s.io/client-go/tools/record"
 	"k8s.io/klog"
+	metrics "k8s.io/metrics/pkg/client/clientset/versioned"
 	schedclientset "sigs.k8s.io/scheduler-plugins/pkg/generated/clientset/versioned"
 	volcanoclient "volcano.sh/apis/pkg/client/clientset/versioned"
 
@@ -114,7 +115,7 @@ func Run(opt *options.ServerOption) error {
 	cfg.Burst = opt.Burst
 
 	// Create clients.
-	kubeClient, leaderElectionClientSet, mpiJobClientSet, volcanoClientSet, schedClientSet, err := createClientSets(cfg, opt.GangSchedulingName)
+	kubeClient, metricsClient, leaderElectionClientSet, mpiJobClientSet, volcanoClientSet, schedClientSet, err := createClientSets(cfg, opt.GangSchedulingName)
 	if err != nil {
 		return err
 	}
@@ -143,6 +144,7 @@ func Run(opt *options.ServerOption) error {
 
 		controller, err := controllersv1.NewMPIJobController(
 			kubeClient,
+			metricsClient,
 			mpiJobClientSet,
 			volcanoClientSet,
 			schedClientSet,
@@ -259,6 +261,7 @@ func createClientSets(
 	gangSchedulingName string,
 ) (
 	kubeclientset.Interface,
+	metrics.Interface,
 	kubeclientset.Interface,
 	mpijobclientset.Interface,
 	volcanoclient.Interface,
@@ -268,17 +271,22 @@ func createClientSets(
 
 	kubeClientSet, err := kubeclientset.NewForConfig(restclientset.AddUserAgent(config, "mpi-operator"))
 	if err != nil {
-		return nil, nil, nil, nil, nil, err
+		return nil, nil, nil, nil, nil, nil, err
+	}
+
+	metricsClientSet, err := metrics.NewForConfig(config)
+	if err != nil {
+		return nil, nil, nil, nil, nil, nil, err
 	}
 
 	leaderElectionClientSet, err := kubeclientset.NewForConfig(restclientset.AddUserAgent(config, "leader-election"))
 	if err != nil {
-		return nil, nil, nil, nil, nil, err
+		return nil, nil, nil, nil, nil, nil, err
 	}
 
 	mpiJobClientSet, err := mpijobclientset.NewForConfig(config)
 	if err != nil {
-		return nil, nil, nil, nil, nil, err
+		return nil, nil, nil, nil, nil, nil, err
 	}
 
 	var (
@@ -287,15 +295,15 @@ func createClientSets(
 	)
 	if gangSchedulingName == options.GangSchedulerVolcano {
 		if volcanoClientSet, err = volcanoclient.NewForConfig(restclientset.AddUserAgent(config, "volcano")); err != nil {
-			return nil, nil, nil, nil, nil, err
+			return nil, nil, nil, nil, nil, nil, err
 		}
 	} else if len(gangSchedulingName) != 0 {
 		if schedClientSet, err = schedclientset.NewForConfig(restclientset.AddUserAgent(config, "scheduler-plugins")); err != nil {
-			return nil, nil, nil, nil, nil, err
+			return nil, nil, nil, nil, nil, nil, err
 		}
 	}
 
-	return kubeClientSet, leaderElectionClientSet, mpiJobClientSet, volcanoClientSet, schedClientSet, nil
+	return kubeClientSet, metricsClientSet, leaderElectionClientSet, mpiJobClientSet, volcanoClientSet, schedClientSet, nil
 }
 
 func checkCRDExists(clientset mpijobclientset.Interface, namespace string) bool {

deploy/v2beta1/mpi-operator.yaml

@@ -80,9 +80,15 @@ spec:
                 additionalProperties:
                   description: ReplicaSpec is a description of the replica
                   properties:
-                    replicas:
+                    maxReplicas:
                       description: |-
-                        Replicas is the desired number of replicas of the given template.
+                        MaxReplicas is the desired number of replicas of the given template.
+                        If unspecified, defaults to 1.
+                      format: int32
+                      type: integer
+                    minReplicas:
+                      description: |-
+                        MaxReplicas is the desired number of replicas of the given template.
                         If unspecified, defaults to 1.
                       format: int32
                       type: integer
@@ -7839,6 +7845,12 @@ spec:
                   Defaults to 1.
                 format: int32
                 type: integer
+              priority:
+                default: 0
+                description: |-
+                  Specifies the job priority
+                format: int32
+                type: integer
               sshAuthMountPath:
                 default: /root/.ssh
                 description: |-
@@ -8090,8 +8102,10 @@ rules:
   - update
 - apiGroups:
   - ""
+  - metrics.k8s.io
   resources:
   - pods
+  - nodes
   verbs:
   - create
   - get
@@ -8162,12 +8176,14 @@ rules:
   - '*'
 - apiGroups:
   - scheduling.x-k8s.io
+  - metrics.k8s.io
   resources:
   - podgroups
   verbs:
   - '*'
 - apiGroups:
   - scheduling.k8s.io
+  - metrics.k8s.io
   resources:
   - priorityclasses
   verbs:
@@ -8225,6 +8241,7 @@ spec:
       - args:
         - -alsologtostderr
         - --lock-namespace=mpi-operator
-        image: mpioperator/mpi-operator:master
+        image: adityapb/mpi-operator:autoscale
+        imagePullPolicy: Always
         name: mpi-operator
       serviceAccountName: mpi-operator

go.mod

@@ -12,6 +12,7 @@ require (
 	k8s.io/apimachinery v0.29.4
 	k8s.io/apiserver v0.29.4
 	k8s.io/client-go v0.29.4
+	k8s.io/metrics v0.29.4
 	k8s.io/code-generator v0.29.4
 	k8s.io/klog v1.0.0
 	k8s.io/kube-openapi v0.0.0-20231010175941-2dd684a91f00

go.sum

@@ -361,6 +361,8 @@ k8s.io/klog/v2 v2.110.1 h1:U/Af64HJf7FcwMcXyKm2RPM22WZzyR7OSpYj5tg3cL0=
 k8s.io/klog/v2 v2.110.1/go.mod h1:YGtd1984u+GgbuZ7e08/yBuAfKLSO0+uR1Fhi6ExXjo=
 k8s.io/kube-openapi v0.0.0-20231010175941-2dd684a91f00 h1:aVUu9fTY98ivBPKR9Y5w/AuzbMm96cd3YHRTU83I780=
 k8s.io/kube-openapi v0.0.0-20231010175941-2dd684a91f00/go.mod h1:AsvuZPBlUDVuCdzJ87iajxtXuR9oktsTctW/R9wwouA=
+k8s.io/metrics v0.29.4 h1:06sZ63/Kt9HEb5GP/1y6xbHDz6XkxnHpu949UdXfoXQ=
+k8s.io/metrics v0.29.4/go.mod h1:ZN9peB0nLTqPZuwQna8ZUrPFJQ0i8QNH4pqRJopS+9c=
 k8s.io/utils v0.0.0-20240102154912-e7106e64919e h1:eQ/4ljkx21sObifjzXwlPKpdGLrCfRziVtos3ofG/sQ=
 k8s.io/utils v0.0.0-20240102154912-e7106e64919e/go.mod h1:OLgZIPagt7ERELqWJFomSt595RzquPNLL48iOWgYOg0=
 sigs.k8s.io/apiserver-network-proxy/konnectivity-client v0.28.0 h1:TgtAeesdhpm2SGwkQasmbeqDo8th5wOBA5h/AjTKA4I=

pkg/apis/kubeflow/v2beta1/default.go

@@ -31,8 +31,11 @@ func setDefaultsTypeLauncher(spec *ReplicaSpec) {
 	if spec.RestartPolicy == "" {
 		spec.RestartPolicy = DefaultLauncherRestartPolicy
 	}
-	if spec.Replicas == nil {
-		spec.Replicas = ptr.To[int32](1)
+	if spec.MaxReplicas == nil {
+		spec.MaxReplicas = ptr.To[int32](1)
+	}
+	if spec.MinReplicas == nil {
+		spec.MinReplicas = ptr.To[int32](1)
 	}
 }
 
@@ -44,8 +47,11 @@ func setDefaultsTypeWorker(spec *ReplicaSpec) {
 	if spec.RestartPolicy == "" {
 		spec.RestartPolicy = DefaultRestartPolicy
 	}
-	if spec.Replicas == nil {
-		spec.Replicas = ptr.To[int32](0)
+	if spec.MaxReplicas == nil {
+		spec.MaxReplicas = ptr.To[int32](0)
+	}
+	if spec.MinReplicas == nil {
+		spec.MinReplicas = ptr.To[int32](1)
 	}
 }
 

pkg/apis/kubeflow/v2beta1/types.go

@@ -154,6 +154,10 @@ type MPIJobSpec struct {
 	// +kubebuilder:default:=1
 	SlotsPerWorker *int32 `json:"slotsPerWorker,omitempty"`
 
+	// +optional
+	// +kubebuilder:default:=0
+	Priority *int32 `json:"priority,omitempty"`
+
 	// RunLauncherAsWorker indicates whether to run worker process in launcher
 	// Defaults to false.
 	// +optional
@@ -174,7 +178,7 @@ type MPIJobSpec struct {
 
 	// launcherCreationPolicy if WaitForWorkersReady, the launcher is created only after all workers are in Ready state. Defaults to AtStartup.
 	// +kubebuilder:validation:Enum:AtStartup;WaitForWorkersReady
-	// +kubebuilder:default:=AtStartup
+	// +kubebuilder:default:=WaitForWorkersReady
 	LauncherCreationPolicy LauncherCreationPolicy `json:"launcherCreationPolicy,omitempty"`
 
 	// MPIImplementation is the MPI implementation.
@@ -329,7 +333,9 @@ const (
 type ReplicaSpec struct {
 	// Replicas is the desired number of replicas of the given template.
 	// If unspecified, defaults to 1.
-	Replicas *int32 `json:"replicas,omitempty"`
+	MaxReplicas *int32 `json:"maxReplicas,omitempty"`
+
+	MinReplicas *int32 `json:"minReplicas,omitempty"`
 
 	// Template is the object that describes the pod that
 	// will be created for this replica. RestartPolicy in PodTemplateSpec

pkg/apis/kubeflow/v2beta1/zz_generated.deepcopy.go

@@ -53,8 +53,8 @@ func validateMPIJobName(job *kubeflow.MPIJob) field.ErrorList {
 	var allErrs field.ErrorList
 	var replicas int32 = 1
 	if workerSpec := job.Spec.MPIReplicaSpecs[kubeflow.MPIReplicaTypeWorker]; workerSpec != nil {
-		if workerSpec.Replicas != nil && *workerSpec.Replicas > 0 {
-			replicas = *workerSpec.Replicas
+		if workerSpec.MaxReplicas != nil && *workerSpec.MaxReplicas > 0 {
+			replicas = *workerSpec.MaxReplicas
 		}
 	}
 	maximumPodHostname := fmt.Sprintf("%s-worker-%d", job.Name, replicas-1)
@@ -119,8 +119,8 @@ func validateLauncherReplicaSpec(spec *kubeflow.ReplicaSpec, path *field.Path) f
 		return errs
 	}
 	errs = append(errs, validateReplicaSpec(spec, path)...)
-	if spec.Replicas != nil && *spec.Replicas != 1 {
-		errs = append(errs, field.Invalid(path.Child("replicas"), *spec.Replicas, "must be 1"))
+	if spec.MaxReplicas != nil && *spec.MaxReplicas != 1 {
+		errs = append(errs, field.Invalid(path.Child("replicas"), *spec.MaxReplicas, "must be 1"))
 	}
 	return errs
 }
@@ -131,15 +131,15 @@ func validateWorkerReplicaSpec(spec *kubeflow.ReplicaSpec, path *field.Path) fie
 		return errs
 	}
 	errs = append(errs, validateReplicaSpec(spec, path)...)
-	if spec.Replicas != nil && *spec.Replicas <= 0 {
-		errs = append(errs, field.Invalid(path.Child("replicas"), *spec.Replicas, "must be greater than or equal to 1"))
+	if spec.MaxReplicas != nil && *spec.MaxReplicas <= 0 {
+		errs = append(errs, field.Invalid(path.Child("replicas"), *spec.MaxReplicas, "must be greater than or equal to 1"))
 	}
 	return errs
 }
 
 func validateReplicaSpec(spec *kubeflow.ReplicaSpec, path *field.Path) field.ErrorList {
 	var errs field.ErrorList
-	if spec.Replicas == nil {
+	if spec.MaxReplicas == nil {
 		errs = append(errs, field.Required(path.Child("replicas"), "must define number of replicas"))
 	}
 	if !validRestartPolicies.Has(string(spec.RestartPolicy)) {