diff --git a/pkg/workspace/controllers/workspace_controller.go b/pkg/workspace/controllers/workspace_controller.go index d5d9eee52..63173d5b2 100644 --- a/pkg/workspace/controllers/workspace_controller.go +++ b/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 } diff --git a/pkg/workspace/inference/preset-inferences.go b/pkg/workspace/inference/preset-inferences.go index 4bc182353..8961f5860 100644 --- a/pkg/workspace/inference/preset-inferences.go +++ b/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 } } diff --git a/pkg/workspace/inference/preset-inferences_test.go b/pkg/workspace/inference/preset-inferences_test.go index bebbab4e7..5a6651b2f 100644 --- a/pkg/workspace/inference/preset-inferences_test.go +++ b/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: diff --git a/test/e2e/inference_with_adapters_test.go b/test/e2e/inference_with_adapters_test.go index c2fd1ffe2..4567e1b4c 100644 --- a/test/e2e/inference_with_adapters_test.go +++ b/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,11 +92,6 @@ 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 @@ -101,6 +99,35 @@ func validateInitContainers(workspaceObj *kaitov1alpha1.Workspace, expectedInitC }) } +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) }) - }) diff --git a/test/e2e/utils/utils.go b/test/e2e/utils/utils.go index 5dc011f76..48b18e890 100644 --- a/test/e2e/utils/utils.go +++ b/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) +}