fix: secret update patch (#709)

**Reason for Change**:
Workspace Secret Update Patch

**Requirements**

- [ ] added unit tests and e2e tests (if applicable).

**Issue Fixed**:
<!-- If this PR fixes GitHub issue 4321, add "Fixes #4321" to the next
line. -->

**Notes for Reviewers**:

Signed-off-by: Bangqi Zhu <[email protected]>
Co-authored-by: Bangqi Zhu <[email protected]>
Co-authored-by: Fei Guo <[email protected]>

pkg/workspace/controllers/workspace_controller.go

@@ -769,6 +769,9 @@ func (c *WorkspaceReconciler) applyInference(ctx context.Context, wObj *kaitov1a
 						deployment.Annotations[kaitov1alpha1.WorkspaceRevisionAnnotation] = revisionStr
 						spec.Template.Spec.Volumes = volumes
 
+						_, imagePullSecrets := inference.GetInferenceImageInfo(ctx, wObj, inferenceParam)
+						deployment.Spec.Template.Spec.ImagePullSecrets = imagePullSecrets
+
 						if err := c.Update(ctx, deployment); err != nil {
 							return
 						}

pkg/workspace/inference/preset-inferences.go

@@ -99,6 +99,13 @@ func updateTorchParamsForDistributedInference(ctx context.Context, kubeClient cl
 func GetInferenceImageInfo(ctx context.Context, workspaceObj *kaitov1alpha1.Workspace, presetObj *model.PresetParam) (string, []corev1.LocalObjectReference) {
 	imagePullSecretRefs := []corev1.LocalObjectReference{}
 	// Check if the workspace preset's access mode is private
+	if len(workspaceObj.Inference.Adapters) > 0 {
+		for _, adapter := range workspaceObj.Inference.Adapters {
+			for _, secretName := range adapter.Source.ImagePullSecrets {
+				imagePullSecretRefs = append(imagePullSecretRefs, corev1.LocalObjectReference{Name: secretName})
+			}
+		}
+	}
 	if string(workspaceObj.Inference.Preset.AccessMode) == string(kaitov1alpha1.ModelImageAccessModePrivate) {
 		imageName := workspaceObj.Inference.Preset.PresetOptions.Image
 		for _, secretName := range workspaceObj.Inference.Preset.PresetOptions.ImagePullSecrets {
@@ -110,13 +117,7 @@ func GetInferenceImageInfo(ctx context.Context, workspaceObj *kaitov1alpha1.Work
 		imageTag := presetObj.Tag
 		registryName := os.Getenv("PRESET_REGISTRY_NAME")
 		imageName = fmt.Sprintf("%s/kaito-%s:%s", registryName, imageName, imageTag)
-		if len(workspaceObj.Inference.Adapters) > 0 {
-			for _, adapter := range workspaceObj.Inference.Adapters {
-				for _, secretName := range adapter.Source.ImagePullSecrets {
-					imagePullSecretRefs = append(imagePullSecretRefs, corev1.LocalObjectReference{Name: secretName})
-				}
-			}
-		}
+
 		return imageName, imagePullSecretRefs
 	}
 }

pkg/workspace/inference/preset-inferences_test.go

@@ -117,7 +117,7 @@ func TestCreatePresetInference(t *testing.T) {
 			}
 			mockClient.CreateOrUpdateObjectInMap(svc)
 
-			createdObject, _ := CreatePresetInference(context.TODO(), workspace, test.MockWorkspaceWithPresetHash, inferenceObj, useHeadlessSvc, mockClient)
+			createdObject, _ := CreatePresetInference(context.TODO(), workspace, "1", inferenceObj, useHeadlessSvc, mockClient)
 			createdWorkload := ""
 			switch createdObject.(type) {
 			case *appsv1.Deployment:

test/e2e/inference_with_adapters_test.go

@@ -31,7 +31,7 @@ var validAdapters1 = []kaitov1alpha1.AdapterSpec{
 			Name:  imageName1,
 			Image: fullImageName1,
 			ImagePullSecrets: []string{
-				aiModelsRegistrySecret,
+				utils.GetEnv("AI_MODELS_REGISTRY_SECRET"),
 			},
 		},
 		Strength: &DefaultStrength,
@@ -43,6 +43,9 @@ var validAdapters2 = []kaitov1alpha1.AdapterSpec{
 		Source: &kaitov1alpha1.DataSource{
 			Name:  imageName2,
 			Image: fullImageName2,
+			ImagePullSecrets: []string{
+				utils.GetEnv("E2E_ACR_REGISTRY_SECRET"),
+			},
 		},
 		Strength: &DefaultStrength,
 	},
@@ -89,18 +92,42 @@ func validateInitContainers(workspaceObj *kaitov1alpha1.Workspace, expectedInitC
 				return false
 			}
 			initContainer, expectedInitContainer := initContainers[0], expectedInitContainers[0]
-			if expectedInitContainer.Name == imageName1 { //only the first adapter need to check imagePullSecrets
-				if dep.Spec.Template.Spec.ImagePullSecrets == nil || len(dep.Spec.Template.Spec.ImagePullSecrets) == 0 {
-					return false
-				}
-			}
 
 			// GinkgoWriter.Printf("Resource '%s' not ready. Ready replicas: %d\n", workspaceObj.Name, readyReplicas)
 			return initContainer.Image == expectedInitContainer.Image && initContainer.Name == expectedInitContainer.Name
 		}, 20*time.Minute, utils.PollInterval).Should(BeTrue(), "Failed to wait for initContainers to be ready")
 	})
 }
 
+func validateImagePullSecrets(workspaceObj *kaitov1alpha1.Workspace, expectedImagePullSecrets []string) {
+	By("Checking the ImagePullSecrets", func() {
+		Eventually(func() bool {
+			var err error
+
+			dep := &appsv1.Deployment{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      workspaceObj.Name,
+					Namespace: workspaceObj.Namespace,
+				},
+			}
+			err = utils.TestingCluster.KubeClient.Get(ctx, client.ObjectKey{
+				Namespace: workspaceObj.Namespace,
+				Name:      workspaceObj.Name,
+			}, dep)
+
+			if err != nil {
+				GinkgoWriter.Printf("Error fetching resource: %v\n", err)
+				return false
+			}
+			if dep.Spec.Template.Spec.ImagePullSecrets == nil {
+				return false
+			}
+
+			return utils.CompareSecrets(dep.Spec.Template.Spec.ImagePullSecrets, expectedImagePullSecrets)
+		}, 5*time.Minute, utils.PollInterval).Should(BeTrue(), "Failed to wait for ImagePullSecrets to be ready")
+	})
+}
+
 func validateAdapterAdded(workspaceObj *kaitov1alpha1.Workspace, deploymentName string, adapterName string) {
 	By("Checking the Adapters", func() {
 		Eventually(func() bool {
@@ -134,7 +161,6 @@ func validateAdapterAdded(workspaceObj *kaitov1alpha1.Workspace, deploymentName
 var _ = Describe("Workspace Preset", func() {
 	BeforeEach(func() {
 		loadTestEnvVars()
-
 		loadModelVersions()
 	})
 
@@ -172,6 +198,7 @@ var _ = Describe("Workspace Preset", func() {
 		validateRevision(workspaceObj, "1")
 
 		validateInitContainers(workspaceObj, expectedInitContainers1)
+		validateImagePullSecrets(workspaceObj, validAdapters1[0].Source.ImagePullSecrets)
 		validateAdapterAdded(workspaceObj, workspaceObj.Name, imageName1)
 
 		workspaceObj = updateCustomWorkspaceWithAdapter(workspaceObj, validAdapters2)
@@ -187,7 +214,7 @@ var _ = Describe("Workspace Preset", func() {
 
 		validateRevision(workspaceObj, "2")
 		validateInitContainers(workspaceObj, expectedInitContainers2)
+		validateImagePullSecrets(workspaceObj, validAdapters2[0].Source.ImagePullSecrets)
 		validateAdapterAdded(workspaceObj, workspaceObj.Name, imageName2)
 	})
-
 })

test/e2e/utils/utils.go

@@ -12,6 +12,8 @@ import (
 	"math/rand"
 	"os"
 	"path/filepath"
+	"reflect"
+	"sort"
 	"strings"
 	"time"
 
@@ -401,3 +403,19 @@ func GeneratePodTemplate(name, namespace, image string, labels map[string]string
 		},
 	}
 }
+
+func CompareSecrets(refs []corev1.LocalObjectReference, secrets []string) bool {
+	if len(refs) != len(secrets) {
+		return false
+	}
+
+	refSecrets := make([]string, len(refs))
+	for i, ref := range refs {
+		refSecrets[i] = ref.Name
+	}
+
+	sort.Strings(refSecrets)
+	sort.Strings(secrets)
+
+	return reflect.DeepEqual(refSecrets, secrets)
+}