feat: support configuring init container image

Signed-off-by: rudeigerc <rudeigerc@gmail.com>

Author: rudeigerc

pkg/controller/inference/service_controller.go

@@ -116,10 +116,11 @@ func (r *ServiceReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ct
 		return ctrl.Result{}, err
 	}
 
-	workloadApplyConfiguration, err := buildWorkloadApplyConfiguration(service, models)
+	workloadApplyConfiguration, err := buildWorkloadApplyConfiguration(service, models, configs)
 	if err != nil {
 		return ctrl.Result{}, err
 	}
+
 	if err := setControllerReferenceForWorkload(service, workloadApplyConfiguration, r.Scheme); err != nil {
 		return ctrl.Result{}, err
 	}
@@ -162,7 +163,7 @@ func (r *ServiceReconciler) SetupWithManager(mgr ctrl.Manager) error {
 		Complete(r)
 }
 
-func buildWorkloadApplyConfiguration(service *inferenceapi.Service, models []*coreapi.OpenModel) (*applyconfigurationv1.LeaderWorkerSetApplyConfiguration, error) {
+func buildWorkloadApplyConfiguration(service *inferenceapi.Service, models []*coreapi.OpenModel, configs *helper.GlobalConfigs) (*applyconfigurationv1.LeaderWorkerSetApplyConfiguration, error) {
 	workload := applyconfigurationv1.LeaderWorkerSet(service.Name, service.Namespace)
 
 	leaderWorkerTemplate := applyconfigurationv1.LeaderWorkerTemplate()
@@ -193,7 +194,7 @@ func buildWorkloadApplyConfiguration(service *inferenceapi.Service, models []*co
 	leaderWorkerTemplate.WithWorkerTemplate(&podTemplateSpecApplyConfiguration)
 
 	// The core logic to inject additional configurations.
-	injectModelProperties(leaderWorkerTemplate, models, service)
+	injectModelProperties(leaderWorkerTemplate, models, service, configs)
 
 	spec := applyconfigurationv1.LeaderWorkerSetSpec()
 	spec.WithLeaderWorkerTemplate(leaderWorkerTemplate)
@@ -215,17 +216,17 @@ func buildWorkloadApplyConfiguration(service *inferenceapi.Service, models []*co
 	return workload, nil
 }
 
-func injectModelProperties(template *applyconfigurationv1.LeaderWorkerTemplateApplyConfiguration, models []*coreapi.OpenModel, service *inferenceapi.Service) {
+func injectModelProperties(template *applyconfigurationv1.LeaderWorkerTemplateApplyConfiguration, models []*coreapi.OpenModel, service *inferenceapi.Service, configs *helper.GlobalConfigs) {
 	isMultiNodesInference := template.LeaderTemplate != nil
 
 	for i, model := range models {
 		source := modelSource.NewModelSourceProvider(model)
 		// Skip model-loader initContainer if llmaz.io/skip-model-loader annotation is set.
 		if !helper.SkipModelLoader(service) {
 			if isMultiNodesInference {
-				source.InjectModelLoader(template.LeaderTemplate, i)
+				source.InjectModelLoader(template.LeaderTemplate, i, configs.InitContainerImage)
 			}
-			source.InjectModelLoader(template.WorkerTemplate, i)
+			source.InjectModelLoader(template.WorkerTemplate, i, configs.InitContainerImage)
 		} else {
 			if isMultiNodesInference {
 				source.InjectModelEnvVars(template.LeaderTemplate)

pkg/controller_helper/modelsource/modelhub.go

@@ -62,16 +62,20 @@ func (p *ModelHubProvider) ModelPath(skipModelLoader bool) string {
 	return CONTAINER_MODEL_PATH + "models--" + strings.ReplaceAll(p.modelID, "/", "--")
 }
 
-func (p *ModelHubProvider) InjectModelLoader(template *coreapplyv1.PodTemplateSpecApplyConfiguration, index int) {
+func (p *ModelHubProvider) InjectModelLoader(template *coreapplyv1.PodTemplateSpecApplyConfiguration, index int, initContainerImage string) {
 	initContainerName := MODEL_LOADER_CONTAINER_NAME
 	if index != 0 {
 		initContainerName += "-" + strconv.Itoa(index)
 	}
 
+	if initContainerImage == "" {
+		initContainerImage = pkg.LOADER_IMAGE
+	}
+
 	// Handle initContainer.
 	initContainer := coreapplyv1.Container().
 		WithName(initContainerName).
-		WithImage(pkg.LOADER_IMAGE).
+		WithImage(initContainerImage).
 		WithVolumeMounts(coreapplyv1.VolumeMount().WithName(MODEL_VOLUME_NAME).WithMountPath(CONTAINER_MODEL_PATH))
 
 	// We have exactly one container in the template.Spec.Containers.

pkg/controller_helper/modelsource/modelhub_test.go

@@ -38,10 +38,11 @@ func Test_ModelHubProvider_InjectModelLoader(t *testing.T) {
 	ignorePatterns := []string{"*.tmp"}
 
 	tests := []struct {
-		name            string
-		provider        *ModelHubProvider
-		index           int
-		expectMainModel bool
+		name               string
+		provider           *ModelHubProvider
+		index              int
+		expectMainModel    bool
+		initContainerImage string
 	}{
 		{
 			name: "inject full modelhub with fileName, revision, allow/ignore",
@@ -67,6 +68,17 @@ func Test_ModelHubProvider_InjectModelLoader(t *testing.T) {
 			index:           1,
 			expectMainModel: false,
 		},
+		{
+			name: "inject with custom initContainerImage",
+			provider: &ModelHubProvider{
+				modelName: "llama3",
+				modelID:   "meta/llama-3",
+				modelHub:  "Huggingface",
+			},
+			index:              0,
+			expectMainModel:    true,
+			initContainerImage: "custom-loader-image:latest",
+		},
 	}
 
 	envSortOpt := cmpopts.SortSlices(func(a, b corev1.EnvVar) bool {
@@ -83,7 +95,7 @@ func Test_ModelHubProvider_InjectModelLoader(t *testing.T) {
 					),
 				)
 
-			tt.provider.InjectModelLoader(template, tt.index)
+			tt.provider.InjectModelLoader(template, tt.index, tt.initContainerImage)
 
 			assert.Len(t, template.Spec.InitContainers, 1)
 			initContainer := template.Spec.InitContainers[0]
@@ -92,8 +104,12 @@ func Test_ModelHubProvider_InjectModelLoader(t *testing.T) {
 			if tt.index != 0 {
 				expectedName += "-" + strconv.Itoa(tt.index)
 			}
+			expectedImage := tt.initContainerImage
+			if expectedImage == "" {
+				expectedImage = pkg.LOADER_IMAGE
+			}
 			assert.Equal(t, expectedName, *initContainer.Name)
-			assert.Equal(t, pkg.LOADER_IMAGE, *initContainer.Image)
+			assert.Equal(t, expectedImage, *initContainer.Image)
 
 			wantEnv := buildExpectedEnv(tt.provider)
 			if diff := cmp.Diff(wantEnv, initContainer.Env, envSortOpt); diff != "" {

pkg/controller_helper/modelsource/modelsource.go

@@ -59,7 +59,7 @@ type ModelSourceProvider interface {
 	ModelPath(skipModelLoader bool) 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 *coreapplyv1.PodTemplateSpecApplyConfiguration, index int)
+	InjectModelLoader(spec *coreapplyv1.PodTemplateSpecApplyConfiguration, index int, initContainerImage string)
 	// InjectModelEnvVars will inject the model credentials env to the model-runner container.
 	// This is used when the model-loader initContainer is not injected, and the model loading is handled by the model-runner container.
 	InjectModelEnvVars(spec *coreapplyv1.PodTemplateSpecApplyConfiguration)

pkg/controller_helper/modelsource/modelsource_test.go

@@ -23,6 +23,7 @@ import (
 	coreapplyv1 "k8s.io/client-go/applyconfigurations/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"
 )
@@ -130,7 +131,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)
 		})

pkg/controller_helper/modelsource/uri.go

@@ -80,7 +80,7 @@ func (p *URIProvider) ModelPath(skipModelLoader bool) string {
 	return CONTAINER_MODEL_PATH + "models--" + splits[len(splits)-1]
 }
 
-func (p *URIProvider) InjectModelLoader(template *coreapplyv1.PodTemplateSpecApplyConfiguration, index int) {
+func (p *URIProvider) InjectModelLoader(template *coreapplyv1.PodTemplateSpecApplyConfiguration, index int, initContainerImage string) {
 	// We don't have additional operations for Ollama, just load in runtime.
 	if p.protocol == Ollama {
 		return
@@ -112,10 +112,16 @@ func (p *URIProvider) InjectModelLoader(template *coreapplyv1.PodTemplateSpecApp
 	if index != 0 {
 		initContainerName += "-" + strconv.Itoa(index)
 	}
+
+	// Handle the image of initContainer.
+	if initContainerImage == "" {
+		initContainerImage = pkg.LOADER_IMAGE
+	}
+
 	// Handle initContainer.
 	initContainer := coreapplyv1.Container().
 		WithName(initContainerName).
-		WithImage(pkg.LOADER_IMAGE).
+		WithImage(initContainerImage).
 		WithVolumeMounts(
 			coreapplyv1.VolumeMount().
 				WithName(MODEL_VOLUME_NAME).

test/util/validation/validate_service.go

@@ -72,6 +72,22 @@ 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 {
 			if helper.SkipModelLoader(service) {
 				if service.Spec.WorkloadTemplate.LeaderTemplate != nil {
@@ -85,11 +101,11 @@ func ValidateService(ctx context.Context, k8sClient client.Client, service *infe
 			} else {
 				// 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
 				}
 			}
@@ -114,15 +130,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 {
@@ -143,8 +159,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
@@ -437,25 +454,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