Repository: intelligent-machine-learning/dlrover (language: python)

Issue:
Support setting the image for dlrover master in a job yaml file.



Discussion / hints:

Fix patch:
diff --git a/dlrover/go/operator/config/default/manager_auth_proxy_patch.yaml b/dlrover/go/operator/config/default/manager_auth_proxy_patch.yaml
index dc4081d60..7204d829c 100644
--- a/dlrover/go/operator/config/default/manager_auth_proxy_patch.yaml
+++ b/dlrover/go/operator/config/default/manager_auth_proxy_patch.yaml
@@ -37,3 +37,4 @@ spec:
         - "--health-probe-bind-address=:8081"
         - "--metrics-bind-address=127.0.0.1:8080"
         - "--leader-elect"
+        - "--master-image=registry.cn-hangzhou.aliyuncs.com/intell-ai/dlrover:test"
diff --git a/dlrover/go/operator/main.go b/dlrover/go/operator/main.go
index 6ed0aab9d..f34a96d9e 100644
--- a/dlrover/go/operator/main.go
+++ b/dlrover/go/operator/main.go
@@ -53,11 +53,14 @@ func main() {
 	var metricsAddr string
 	var enableLeaderElection bool
 	var probeAddr string
+	var masterImage string
 	flag.StringVar(&metricsAddr, "metrics-bind-address", ":8080", "The address the metric endpoint binds to.")
 	flag.StringVar(&probeAddr, "health-probe-bind-address", ":8081", "The address the probe endpoint binds to.")
 	flag.BoolVar(&enableLeaderElection, "leader-elect", false,
 		"Enable leader election for controller manager. "+
 			"Enabling this will ensure there is only one active controller manager.")
+	flag.StringVar(&masterImage, "master-image", "registry.cn-hangzhou.aliyuncs.com/intell-ai/dlrover:latest",
+		"The image to launch a dlrover master Pod of an ElasticJob.")
 	opts := zap.Options{
 		Development: true,
 	}
@@ -90,7 +93,7 @@ func main() {
 		os.Exit(1)
 	}
 
-	if err = controllers.NewElasticJobReconciler(mgr).SetupWithManager(mgr); err != nil {
+	if err = controllers.NewElasticJobReconciler(mgr, masterImage).SetupWithManager(mgr); err != nil {
 		setupLog.Error(err, "unable to create controller", "controller", "ElasticJob")
 		os.Exit(1)
 	}
diff --git a/dlrover/go/operator/pkg/controllers/elasticjob_controller.go b/dlrover/go/operator/pkg/controllers/elasticjob_controller.go
index 88482a949..5acd3b183 100644
--- a/dlrover/go/operator/pkg/controllers/elasticjob_controller.go
+++ b/dlrover/go/operator/pkg/controllers/elasticjob_controller.go
@@ -46,18 +46,20 @@ const (
 // ElasticJobReconciler reconciles a ElasticJob object
 type ElasticJobReconciler struct {
 	client.Client
-	Scheme   *runtime.Scheme
-	Recorder record.EventRecorder
-	Log      logr.Logger
+	Scheme      *runtime.Scheme
+	Recorder    record.EventRecorder
+	Log         logr.Logger
+	masterImage string
 }
 
 // NewElasticJobReconciler creates a JobReconciler
-func NewElasticJobReconciler(mgr ctrl.Manager) *ElasticJobReconciler {
+func NewElasticJobReconciler(mgr ctrl.Manager, masterImage string) *ElasticJobReconciler {
 	r := &ElasticJobReconciler{
-		Client:   mgr.GetClient(),
-		Scheme:   mgr.GetScheme(),
-		Recorder: mgr.GetEventRecorderFor("elasticjob-controller"),
-		Log:      ctrl.Log.WithName("controllers").WithName("ElasticJob"),
+		Client:      mgr.GetClient(),
+		Scheme:      mgr.GetScheme(),
+		Recorder:    mgr.GetEventRecorderFor("elasticjob-controller"),
+		Log:         ctrl.Log.WithName("controllers").WithName("ElasticJob"),
+		masterImage: masterImage,
 	}
 	return r
 }
@@ -191,6 +193,7 @@ func (r *ElasticJobReconciler) stopRunningPods(job *elasticv1alpha1.ElasticJob)
 }
 
 func (r *ElasticJobReconciler) createEasydlMaster(job *elasticv1alpha1.ElasticJob) error {
+	master.NewMasterTemplateToJob(job, r.masterImage)
 	masterManager := common.ReplicaManagers[master.ReplicaTypeTrainerMaster]
 	err := masterManager.ReconcilePods(r.Client, job, nil)
 	if err != nil {
diff --git a/dlrover/go/operator/pkg/controllers/master/master.go b/dlrover/go/operator/pkg/controllers/master/master.go
index 78377358c..1b4582d4e 100644
--- a/dlrover/go/operator/pkg/controllers/master/master.go
+++ b/dlrover/go/operator/pkg/controllers/master/master.go
@@ -33,11 +33,10 @@ const (
 	initMasterContainerMemory  = "2Gi"
 	initMasterContainerStorage = "2Gi"
 	masterCommand              = "python -m dlrover.python.master.main"
-	masterImage                = "registry.cn-hangzhou.aliyuncs.com/intell-ai/dlrover:test"
 	masterServicePort          = 50001
 	initMasterIndex            = 0
-
-	envMasterAddrKey = "MASTER_ADDR"
+	defaultImagePullPolicy     = "IfNotPresent"
+	envMasterAddrKey           = "MASTER_ADDR"
 
 	// ReplicaTypeTrainerMaster is the type for DLRover Master replica.
 	ReplicaTypeTrainerMaster commonv1.ReplicaType = "dlrover-master"
@@ -47,46 +46,16 @@ const (
 type Manager struct{}
 
 func init() {
-	common.ReplicaManagers[ReplicaTypeTrainerMaster] = newManager()
-}
-
-func newManager() *Manager {
-	return &Manager{}
+	common.ReplicaManagers[ReplicaTypeTrainerMaster] = &Manager{}
 }
 
 func (m *Manager) newJobMaster(
 	job *elasticv1alpha1.ElasticJob, replicaIndex int,
 ) *corev1.Pod {
-	command := masterCommand + fmt.Sprintf(
-		" --namespace %s --job_name %s --port %d",
-		job.Namespace, job.Name, masterServicePort,
-	)
-	container := corev1.Container{
-		Name:            "main",
-		Image:           masterImage,
-		ImagePullPolicy: "IfNotPresent",
-		Command:         []string{"/bin/bash", "-c", command},
-		Resources: corev1.ResourceRequirements{
-			Requests: corev1.ResourceList{
-				corev1.ResourceCPU:              resource.MustParse(initMasterContainerCPU),
-				corev1.ResourceMemory:           resource.MustParse(initMasterContainerMemory),
-				corev1.ResourceEphemeralStorage: resource.MustParse(initMasterContainerStorage),
-			},
-			Limits: corev1.ResourceList{
-				corev1.ResourceCPU:              resource.MustParse(initMasterContainerCPU),
-				corev1.ResourceMemory:           resource.MustParse(initMasterContainerMemory),
-				corev1.ResourceEphemeralStorage: resource.MustParse(initMasterContainerStorage),
-			},
-		},
-	}
-	podTemplate := &corev1.PodTemplateSpec{
-		Spec: corev1.PodSpec{
-			Containers:    []corev1.Container{container},
-			RestartPolicy: corev1.RestartPolicyNever,
-		},
-	}
 	masterName := newJobMasterName(job.Name)
-	pod := common.NewPod(job, podTemplate, masterName)
+	pod := common.NewPod(
+		job, &job.Spec.ReplicaSpecs[ReplicaTypeTrainerMaster].Template, masterName,
+	)
 	pod.Labels[common.LabelReplicaTypeKey] = string(ReplicaTypeTrainerMaster)
 	pod.Labels[common.LabelReplicaIndexKey] = fmt.Sprintf("%d", replicaIndex)
 	return pod
@@ -240,3 +209,49 @@ func (m *Manager) StopRunningPods(
 	}
 	return nil
 }
+
+// NewMasterTemplateToJob sets configurations to the master template of a job.
+func NewMasterTemplateToJob(job *elasticv1alpha1.ElasticJob, masterImage string) {
+	command := masterCommand + fmt.Sprintf(
+		" --namespace %s --job_name %s --port %d",
+		job.Namespace, job.Name, masterServicePort,
+	)
+	container := corev1.Container{
+		Name:            "main",
+		Image:           masterImage,
+		ImagePullPolicy: defaultImagePullPolicy,
+		Command:         []string{"/bin/bash", "-c", command},
+		Resources: corev1.ResourceRequirements{
+			Requests: corev1.ResourceList{
+				corev1.ResourceCPU:              resource.MustParse(initMasterContainerCPU),
+				corev1.ResourceMemory:           resource.MustParse(initMasterContainerMemory),
+				corev1.ResourceEphemeralStorage: resource.MustParse(initMasterContainerStorage),
+			},
+			Limits: corev1.ResourceList{
+				corev1.ResourceCPU:              resource.MustParse(initMasterContainerCPU),
+				corev1.ResourceMemory:           resource.MustParse(initMasterContainerMemory),
+				corev1.ResourceEphemeralStorage: resource.MustParse(initMasterContainerStorage),
+			},
+		},
+	}
+	podTemplate := &corev1.PodTemplateSpec{
+		Spec: corev1.PodSpec{
+			Containers:    []corev1.Container{container},
+			RestartPolicy: corev1.RestartPolicyNever,
+		},
+	}
+	if _, ok := job.Spec.ReplicaSpecs[ReplicaTypeTrainerMaster]; ok {
+		mainContainer := job.Spec.ReplicaSpecs[ReplicaTypeTrainerMaster].ReplicaSpec.Template.Spec.Containers[0]
+		if mainContainer.Image != "" {
+			podTemplate.Spec.Containers[0].Image = mainContainer.Image
+		}
+		if mainContainer.ImagePullPolicy != "" {
+			podTemplate.Spec.Containers[0].ImagePullPolicy = mainContainer.ImagePullPolicy
+		}
+	}
+	job.Spec.ReplicaSpecs[ReplicaTypeTrainerMaster] = &elasticv1alpha1.ReplicaSpec{
+		ReplicaSpec: commonv1.ReplicaSpec{
+			Template: *podTemplate,
+		},
+	}
+}
diff --git a/scripts/build_wheel.sh b/scripts/build_wheel.sh
index 59908bdef..11dde03ce 100644
--- a/scripts/build_wheel.sh
+++ b/scripts/build_wheel.sh
@@ -22,4 +22,4 @@ make -f dlrover/Makefile
 # Create dlrover package
 echo "Building the wheel for dlrover."
 rm -rf ./build/lib
-python setup.py --quiet bdist_wheel --dist-dir ./build
+python setup.py --quiet bdist_wheel


Test patch:
diff --git a/dlrover/go/operator/pkg/controllers/master/master_test.go b/dlrover/go/operator/pkg/controllers/master/master_test.go
index db0a30bb5..5acf91e7e 100644
--- a/dlrover/go/operator/pkg/controllers/master/master_test.go
+++ b/dlrover/go/operator/pkg/controllers/master/master_test.go
@@ -15,7 +15,9 @@ package master
 
 import (
 	elasticv1alpha1 "github.com/intelligent-machine-learning/easydl/dlrover/go/operator/api/v1alpha1"
+	commonv1 "github.com/intelligent-machine-learning/easydl/dlrover/go/operator/pkg/common/api/v1"
 	"github.com/stretchr/testify/assert"
+	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"strings"
 	"testing"
@@ -30,11 +32,48 @@ func TestCreateMasterPod(t *testing.T) {
 			Labels:      map[string]string{},
 		},
 	}
-
-	manager := newManager()
+	job.Spec.ReplicaSpecs = make(map[commonv1.ReplicaType]*elasticv1alpha1.ReplicaSpec)
+	NewMasterTemplateToJob(job, "dlrover-master:test")
+	manager := &Manager{}
 	pod := manager.newJobMaster(job, initMasterIndex)
 	assert.Equal(t, pod.Name, "elasticjob-test-ps-dlrover-master")
+	assert.Equal(t, pod.Spec.Containers[0].Image, "dlrover-master:test")
+	assert.Equal(t, string(pod.Spec.Containers[0].ImagePullPolicy), "IfNotPresent")
 	assert.True(t, strings.Contains(pod.Spec.Containers[0].Command[2], "--namespace dlrover"))
 	assert.True(t, strings.Contains(pod.Spec.Containers[0].Command[2], "--job_name test-ps"))
 	assert.True(t, strings.Contains(pod.Spec.Containers[0].Command[2], "--port 50001"))
 }
+
+func TestCreateMasterPodWithImage(t *testing.T) {
+	job := &elasticv1alpha1.ElasticJob{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:        "test-ps",
+			Namespace:   "dlrover",
+			Annotations: map[string]string{},
+			Labels:      map[string]string{},
+		},
+	}
+
+	container := corev1.Container{
+		Image:           "dlrover-master:test-v0",
+		ImagePullPolicy: "Always",
+	}
+	job.Spec.ReplicaSpecs = make(map[commonv1.ReplicaType]*elasticv1alpha1.ReplicaSpec)
+	job.Spec.ReplicaSpecs[ReplicaTypeTrainerMaster] = &elasticv1alpha1.ReplicaSpec{
+		ReplicaSpec: commonv1.ReplicaSpec{
+			Template: corev1.PodTemplateSpec{
+				Spec: corev1.PodSpec{
+					Containers:    []corev1.Container{container},
+					RestartPolicy: corev1.RestartPolicyNever,
+				},
+			},
+		},
+	}
+
+	NewMasterTemplateToJob(job, "dlrover-master:test")
+	manager := &Manager{}
+	pod := manager.newJobMaster(job, initMasterIndex)
+	assert.Equal(t, pod.Name, "elasticjob-test-ps-dlrover-master")
+	assert.Equal(t, pod.Spec.Containers[0].Image, "dlrover-master:test-v0")
+	assert.Equal(t, string(pod.Spec.Containers[0].ImagePullPolicy), "Always")
+}