feat: support configuring init container image

rudeigerc · rudeigerc · commit 621c32d6bff4 · 2025-06-07T00:58:17.000+08:00
Signed-off-by: rudeigerc &lt;rudeigerc@gmail.com&gt;
diff --git a/pkg/controller/inference/service_controller.go b/pkg/controller/inference/service_controller.go
@@ -44,6 +44,7 @@ import (
 
 	coreapi "github.com/inftyai/llmaz/api/core/v1alpha1"
 	inferenceapi "github.com/inftyai/llmaz/api/inference/v1alpha1"
+	"github.com/inftyai/llmaz/pkg"
 	helper "github.com/inftyai/llmaz/pkg/controller_helper"
 	modelSource "github.com/inftyai/llmaz/pkg/controller_helper/modelsource"
 	"github.com/inftyai/llmaz/pkg/util"
@@ -116,7 +117,12 @@ func (r *ServiceReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ct
 		return ctrl.Result{}, err
 	}
 
-	workloadApplyConfiguration := buildWorkloadApplyConfiguration(service, models)
+	initContainerImage := configs.InitContainerImage
+	if initContainerImage == "" {
+		initContainerImage = pkg.LOADER_IMAGE
+	}
+
+	workloadApplyConfiguration := buildWorkloadApplyConfiguration(service, models, initContainerImage)
 	if err := setControllerReferenceForWorkload(service, workloadApplyConfiguration, r.Scheme); err != nil {
 		return ctrl.Result{}, err
 	}
@@ -159,7 +165,7 @@ func (r *ServiceReconciler) SetupWithManager(mgr ctrl.Manager) error {
 		Complete(r)
 }
 
-func buildWorkloadApplyConfiguration(service *inferenceapi.Service, models []*coreapi.OpenModel) *applyconfigurationv1.LeaderWorkerSetApplyConfiguration {
+func buildWorkloadApplyConfiguration(service *inferenceapi.Service, models []*coreapi.OpenModel, initContainerImage string) *applyconfigurationv1.LeaderWorkerSetApplyConfiguration {
 	workload := applyconfigurationv1.LeaderWorkerSet(service.Name, service.Namespace)
 
 	leaderWorkerTemplate := applyconfigurationv1.LeaderWorkerTemplate()
@@ -169,7 +175,7 @@ func buildWorkloadApplyConfiguration(service *inferenceapi.Service, models []*co
 	leaderWorkerTemplate.WithWorkerTemplate(service.Spec.WorkloadTemplate.WorkerTemplate)
 
 	// The core logic to inject additional configurations.
-	injectModelProperties(leaderWorkerTemplate, models, service)
+	injectModelProperties(leaderWorkerTemplate, models, service, initContainerImage)
 
 	spec := applyconfigurationv1.LeaderWorkerSetSpec()
 	spec.WithLeaderWorkerTemplate(leaderWorkerTemplate)
@@ -191,15 +197,15 @@ func buildWorkloadApplyConfiguration(service *inferenceapi.Service, models []*co
 	return workload
 }
 
-func injectModelProperties(template *applyconfigurationv1.LeaderWorkerTemplateApplyConfiguration, models []*coreapi.OpenModel, service *inferenceapi.Service) {
+func injectModelProperties(template *applyconfigurationv1.LeaderWorkerTemplateApplyConfiguration, models []*coreapi.OpenModel, service *inferenceapi.Service, initContainerImage string) {
 	isMultiNodesInference := template.LeaderTemplate != nil
 
 	for i, model := range models {
 		source := modelSource.NewModelSourceProvider(model)
 		if isMultiNodesInference {
-			source.InjectModelLoader(template.LeaderTemplate, i)
+			source.InjectModelLoader(template.LeaderTemplate, i, initContainerImage)
 		}
-		source.InjectModelLoader(template.WorkerTemplate, i)
+		source.InjectModelLoader(template.WorkerTemplate, i, initContainerImage)
 	}
 
 	// We only consider the main model's requirements for now.
diff --git a/pkg/controller_helper/modelsource/modelhub.go b/pkg/controller_helper/modelsource/modelhub.go
@@ -22,8 +22,6 @@ import (
 
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/utils/ptr"
-
-	"github.com/inftyai/llmaz/pkg"
 )
 
 var _ ModelSourceProvider = &ModelHubProvider{}
@@ -57,7 +55,7 @@ func (p *ModelHubProvider) ModelPath() string {
 	return CONTAINER_MODEL_PATH + "models--" + strings.ReplaceAll(p.modelID, "/", "--")
 }
 
-func (p *ModelHubProvider) InjectModelLoader(template *corev1.PodTemplateSpec, index int) {
+func (p *ModelHubProvider) InjectModelLoader(template *corev1.PodTemplateSpec, index int, initContainerImage string) {
 	initContainerName := MODEL_LOADER_CONTAINER_NAME
 	if index != 0 {
 		initContainerName += "-" + strconv.Itoa(index)
@@ -66,7 +64,7 @@ func (p *ModelHubProvider) InjectModelLoader(template *corev1.PodTemplateSpec, i
 	// Handle initContainer.
 	initContainer := &corev1.Container{
 		Name:  initContainerName,
-		Image: pkg.LOADER_IMAGE,
+		Image: initContainerImage,
 		VolumeMounts: []corev1.VolumeMount{
 			{
 				Name:      MODEL_VOLUME_NAME,
diff --git a/pkg/controller_helper/modelsource/modelsource.go b/pkg/controller_helper/modelsource/modelsource.go
@@ -52,8 +52,9 @@ type ModelSourceProvider interface {
 	ModelName() string
 	ModelPath() string
 	// InjectModelLoader will inject the model loader to the spec,
-	// index refers to the suffix of the initContainer name, like model-loader, model-loader-1.
-	InjectModelLoader(spec *corev1.PodTemplateSpec, index int)
+	// index refers to the suffix of the initContainer name, like model-loader, model-loader-1,
+	// initContainerImage is the image used for the model loader.
+	InjectModelLoader(spec *corev1.PodTemplateSpec, index int, initContainerImage string)
 }
 
 func NewModelSourceProvider(model *coreapi.OpenModel) ModelSourceProvider {
diff --git a/pkg/controller_helper/modelsource/modelsource_test.go b/pkg/controller_helper/modelsource/modelsource_test.go
@@ -23,6 +23,7 @@ import (
 	corev1 "k8s.io/api/core/v1"
 
 	coreapi "github.com/inftyai/llmaz/api/core/v1alpha1"
+	"github.com/inftyai/llmaz/pkg"
 	"github.com/inftyai/llmaz/test/util"
 	"github.com/inftyai/llmaz/test/util/wrapper"
 )
@@ -123,7 +124,7 @@ func TestEnvInjectModelLoader(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			tt.provider.InjectModelLoader(tt.template, 0)
+			tt.provider.InjectModelLoader(tt.template, 0, pkg.LOADER_IMAGE)
 			initContainer := tt.template.Spec.InitContainers[0]
 			assert.Subset(t, initContainer.Env, tt.template.Spec.Containers[0].Env)
 		})
diff --git a/pkg/controller_helper/modelsource/uri.go b/pkg/controller_helper/modelsource/uri.go
@@ -22,8 +22,6 @@ import (
 
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/utils/ptr"
-
-	"github.com/inftyai/llmaz/pkg"
 )
 
 var _ ModelSourceProvider = &URIProvider{}
@@ -73,7 +71,7 @@ func (p *URIProvider) ModelPath() string {
 	return CONTAINER_MODEL_PATH + "models--" + splits[len(splits)-1]
 }
 
-func (p *URIProvider) InjectModelLoader(template *corev1.PodTemplateSpec, index int) {
+func (p *URIProvider) InjectModelLoader(template *corev1.PodTemplateSpec, index int, initContainerImage string) {
 	// We don't have additional operations for Ollama, just load in runtime.
 	if p.protocol == Ollama {
 		return
@@ -111,7 +109,7 @@ func (p *URIProvider) InjectModelLoader(template *corev1.PodTemplateSpec, index
 	// Handle initContainer.
 	initContainer := &corev1.Container{
 		Name:  initContainerName,
-		Image: pkg.LOADER_IMAGE,
+		Image: initContainerImage,
 		VolumeMounts: []corev1.VolumeMount{
 			{
 				Name:      MODEL_VOLUME_NAME,
diff --git a/test/config/others/global-configmap.yaml b/test/config/others/global-configmap.yaml
@@ -6,4 +6,4 @@ metadata:
 data:
   config.data: |
     scheduler-name: inftyai-scheduler
-    init-container-image: inftyai/model-loader:v0.0.10
+    init-container-image: docker.m.daocloud.io/inftyai/model-loader:v0.0.10
diff --git a/test/util/validation/validate_service.go b/test/util/validation/validate_service.go
@@ -72,14 +72,30 @@ func ValidateService(ctx context.Context, k8sClient client.Client, service *infe
 			models = append(models, model)
 		}
 
+		// Fetch global configuration.
+		cm := corev1.ConfigMap{}
+		if err := k8sClient.Get(ctx, types.NamespacedName{Name: "llmaz-global-config", Namespace: "llmaz-system"}, &cm); err != nil {
+			return err
+		}
+
+		data, err := helper.ParseGlobalConfigmap(&cm)
+		if err != nil {
+			return fmt.Errorf("failed to parse global configmap: %v", err)
+		}
+
+		initContainerImage := data.InitContainerImage
+		if initContainerImage == "" {
+			initContainerImage = pkg.LOADER_IMAGE
+		}
+
 		for index, model := range models {
 			// Validate injecting modelLoaders
 			if service.Spec.WorkloadTemplate.LeaderTemplate != nil {
-				if err := ValidateModelLoader(model, index, *workload.Spec.LeaderWorkerTemplate.LeaderTemplate, service); err != nil {
+				if err := ValidateModelLoader(model, index, *workload.Spec.LeaderWorkerTemplate.LeaderTemplate, service, initContainerImage); err != nil {
 					return err
 				}
 			}
-			if err := ValidateModelLoader(model, index, workload.Spec.LeaderWorkerTemplate.WorkerTemplate, service); err != nil {
+			if err := ValidateModelLoader(model, index, workload.Spec.LeaderWorkerTemplate.WorkerTemplate, service, initContainerImage); err != nil {
 				return err
 			}
 		}
@@ -103,15 +119,15 @@ func ValidateService(ctx context.Context, k8sClient client.Client, service *infe
 			return err
 		}
 
-		if err := ValidateConfigmap(ctx, k8sClient, service); err != nil {
+		if err := ValidateSchedulerName(data.SchedulerName, service); err != nil {
 			return err
 		}
 
 		return nil
 	}, util.IntegrationTimeout, util.Interval).Should(gomega.Succeed())
 }
 
-func ValidateModelLoader(model *coreapi.OpenModel, index int, template corev1.PodTemplateSpec, service *inferenceapi.Service) error {
+func ValidateModelLoader(model *coreapi.OpenModel, index int, template corev1.PodTemplateSpec, service *inferenceapi.Service, initContainerImage string) error {
 	if model.Spec.Source.URI != nil {
 		protocol, _, _ := pkgUtil.ParseURI(string(*model.Spec.Source.URI))
 		if protocol == modelSource.Ollama {
@@ -132,8 +148,9 @@ func ValidateModelLoader(model *coreapi.OpenModel, index int, template corev1.Po
 		if initContainer.Name != containerName {
 			return fmt.Errorf("unexpected initContainer name, want %s, got %s", modelSource.MODEL_LOADER_CONTAINER_NAME, initContainer.Name)
 		}
-		if initContainer.Image != pkg.LOADER_IMAGE {
-			return fmt.Errorf("unexpected initContainer image, want %s, got %s", pkg.LOADER_IMAGE, initContainer.Image)
+
+		if initContainer.Image != initContainerImage {
+			return fmt.Errorf("unexpected initContainer image, want %s, got %s", initContainerImage, initContainer.Image)
 		}
 
 		var envStrings []string
@@ -356,25 +373,15 @@ func CheckServiceAvaliable() error {
 	return nil
 }
 
-func ValidateConfigmap(ctx context.Context, k8sClient client.Client, service *inferenceapi.Service) error {
-	cm := corev1.ConfigMap{}
-	if err := k8sClient.Get(ctx, types.NamespacedName{Name: "llmaz-global-config", Namespace: "llmaz-system"}, &cm); err != nil {
-		return err
-	}
-
-	data, err := helper.ParseGlobalConfigmap(&cm)
-	if err != nil {
-		return fmt.Errorf("failed to parse global configmap: %v", err)
-	}
-
+func ValidateSchedulerName(schedulerName string, service *inferenceapi.Service) error {
 	if service.Spec.WorkloadTemplate.LeaderTemplate != nil {
-		if service.Spec.WorkloadTemplate.LeaderTemplate.Spec.SchedulerName != data.SchedulerName {
-			return fmt.Errorf("unexpected scheduler name %s, want %s", service.Spec.WorkloadTemplate.LeaderTemplate.Spec.SchedulerName, data.SchedulerName)
+		if service.Spec.WorkloadTemplate.LeaderTemplate.Spec.SchedulerName != schedulerName {
+			return fmt.Errorf("unexpected scheduler name %s, want %s", service.Spec.WorkloadTemplate.LeaderTemplate.Spec.SchedulerName, schedulerName)
 		}
 	}
 
-	if service.Spec.WorkloadTemplate.WorkerTemplate.Spec.SchedulerName != data.SchedulerName {
-		return fmt.Errorf("unexpected scheduler name %s, want %s", service.Spec.WorkloadTemplate.WorkerTemplate.Spec.SchedulerName, data.SchedulerName)
+	if service.Spec.WorkloadTemplate.WorkerTemplate.Spec.SchedulerName != schedulerName {
+		return fmt.Errorf("unexpected scheduler name %s, want %s", service.Spec.WorkloadTemplate.WorkerTemplate.Spec.SchedulerName, schedulerName)
 	}
 
 	return nil
