Update unit tests for customizer with new config templates

Signed-off-by: Shiva Krishna, Merla <[email protected]>

Author: shivamerla

config/samples/nemo/latest/nemocustomizer_config.yaml

@@ -22,7 +22,7 @@ data:
     targets:
       meta/[email protected]:
         base_model: meta/llama-3.1-8b-instruct
-        enabled: false
+        enabled: true
         model_path: llama-3_1-8b-instruct_2_0
         model_uri: ngc://nvidia/nemo/llama-3_1-8b-instruct-nemo:2.0
         name: [email protected]
@@ -49,7 +49,7 @@ data:
         precision: bf16-mixed
       meta/[email protected]:
         base_model: meta/llama-3.2-1b-instruct
-        enabled: false
+        enabled: true
         model_path: llama32_1b-instruct_2_0
         model_uri: ngc://nvidia/nemo/llama-3_2-1b-instruct:2.0
         name: [email protected]

internal/controller/nemocustomizer_controller.go

@@ -569,12 +569,19 @@ func (r *NemoCustomizerReconciler) addModelDownloadJobsConfig(ctx context.Contex
 		"ngcAPISecretKey":         n.Spec.ModelDownloadJobs.NGCSecret.Key,
 		"ngcSecretName":           n.Spec.ModelDownloadJobs.NGCSecret.Name,
 		"ngcSecretKey":            n.Spec.ModelDownloadJobs.NGCSecret.Key,
-		"hfSecretName":            n.Spec.ModelDownloadJobs.HFSecret.Name,
-		"hfSecretKey":             n.Spec.ModelDownloadJobs.HFSecret.Key,
 		"securityContext":         n.Spec.ModelDownloadJobs.SecurityContext,
 		"ttlSecondsAfterFinished": n.Spec.ModelDownloadJobs.TTLSecondsAfterFinished,
 		"pollIntervalSeconds":     n.Spec.ModelDownloadJobs.PollIntervalSeconds,
 	}
+
+	// add HF secret only if present
+	if n.Spec.ModelDownloadJobs.HFSecret.Name != "" {
+		if modelDownloadJobs, ok := cfg["model_download_jobs"].(map[string]interface{}); ok {
+			modelDownloadJobs["hfSecretName"] = n.Spec.ModelDownloadJobs.HFSecret.Name
+			modelDownloadJobs["hfSecretKey"] = n.Spec.ModelDownloadJobs.HFSecret.Key
+		}
+	}
+
 	return nil
 }
 

internal/controller/nemocustomizer_controller_test.go

@@ -109,10 +109,12 @@ var _ = Describe("NemoCustomizer Controller", func() {
 		Expect(err).ToNot(HaveOccurred())
 		trainingCM := loadConfigMapFromFile(filepath.Join(testDataDir, "training_config.yaml"))
 		modelsCM := loadConfigMapFromFile(filepath.Join(testDataDir, "models_config.yaml"))
+		modelTargetsCM := loadConfigMapFromFile(filepath.Join(testDataDir, "models_config_targets.yaml"))
 
 		// Register the test ConfigMaps in the cluster
 		Expect(reconciler.GetClient().Create(ctx, trainingCM)).To(Succeed())
 		Expect(reconciler.GetClient().Create(ctx, modelsCM)).To(Succeed())
+		Expect(reconciler.GetClient().Create(ctx, modelTargetsCM)).To(Succeed())
 
 		nemoCustomizer = &appsv1alpha1.NemoCustomizer{
 			ObjectMeta: metav1.ObjectMeta{
@@ -365,6 +367,13 @@ var _ = Describe("NemoCustomizer Controller", func() {
 		if err := k8sClient.Get(ctx, modelsCMName, modelsCM); err == nil {
 			Expect(k8sClient.Delete(ctx, modelsCM)).To(Succeed())
 		}
+
+		// Delete the models targets config
+		modelTargetsCMName := types.NamespacedName{Name: "nemo-model-target-config", Namespace: "default"}
+		modelTargetsCM := &corev1.ConfigMap{}
+		if err := k8sClient.Get(ctx, modelTargetsCMName, modelTargetsCM); err == nil {
+			Expect(k8sClient.Delete(ctx, modelTargetsCM)).To(Succeed())
+		}
 	})
 
 	Describe("Reconcile", func() {
@@ -575,6 +584,74 @@ var _ = Describe("NemoCustomizer Controller", func() {
 			Expect(parsed).To(HaveKeyWithValue("mlflow_tracking_url", "http://mlflow-tracking.nemo.svc.cluster.local:80"))
 		})
 
+		It("should create customizer config with model targets and templates", func() {
+			namespacedName := types.NamespacedName{Name: nemoCustomizer.Name, Namespace: "default"}
+
+			// use models config map with targets and templates
+			nemoCustomizer.Spec.Models.Name = "nemo-model-config-targets"
+
+			err := client.Create(context.TODO(), nemoCustomizer)
+			Expect(err).NotTo(HaveOccurred())
+			err = client.Create(context.TODO(), secrets)
+			Expect(err).NotTo(HaveOccurred())
+
+			result, err := reconciler.Reconcile(context.TODO(), reconcile.Request{NamespacedName: namespacedName})
+			Expect(err).NotTo(HaveOccurred())
+			Expect(result).To(Equal(ctrl.Result{}))
+
+			err = client.Get(ctx, namespacedName, nemoCustomizer)
+			Expect(err).ToNot(HaveOccurred())
+			Expect(nemoCustomizer.Finalizers).To(ContainElement(NemoCustomizerFinalizer))
+
+			// Check that the customizer config map was created
+			configMap := &corev1.ConfigMap{}
+			err = client.Get(context.TODO(), types.NamespacedName{
+				Name:      nemoCustomizer.Name,
+				Namespace: nemoCustomizer.Namespace,
+			}, configMap)
+			Expect(err).NotTo(HaveOccurred())
+
+			// Verify metadata
+			Expect(configMap.Name).To(Equal(nemoCustomizer.GetName()))
+			Expect(configMap.Namespace).To(Equal(nemoCustomizer.Namespace))
+			// Verify key exists
+			Expect(configMap.Data).To(HaveKey("config.yaml"))
+			configData := configMap.Data["config.yaml"]
+			Expect(configData).NotTo(BeEmpty())
+
+			// Verify presence of top-level keys
+			Expect(configData).To(ContainSubstring("training:"))
+			Expect(configData).To(ContainSubstring("customizationTargets:"))
+			Expect(configData).To(ContainSubstring("customizationConfigTemplates:"))
+			Expect(configData).ToNot(ContainSubstring("models:")) // "models" is deprecated from NMR v25.06
+			Expect(configData).To(ContainSubstring("model_download_jobs:"))
+
+			// Unmarshal the full merged config
+			var parsed map[string]interface{}
+			err = yaml.Unmarshal([]byte(configData), &parsed)
+			Expect(err).NotTo(HaveOccurred())
+
+			// Validate model download jobs config with new secret params from NMR v25.06
+			training, ok := parsed["model_download_jobs"].(map[string]interface{})
+			Expect(ok).To(BeTrue(), "expected 'training' to be a map")
+			Expect(training).To(HaveKey("ngcSecretName"))
+			Expect(training).To(HaveKey("ngcSecretKey"))
+
+			// Validate model targets
+			modelTargets, ok := parsed["customizationTargets"].(map[string]interface{})
+			Expect(ok).To(BeTrue(), "expected 'modelTargets' to be a map")
+			Expect(modelTargets).To(HaveKey("targets"))
+
+			// Validate model config templates
+			configTemplates, ok := parsed["customizationConfigTemplates"].(map[string]interface{})
+			Expect(ok).To(BeTrue(), "expected 'configTemplates' to be a map")
+			Expect(configTemplates).To(HaveKey("templates"))
+
+			Expect(parsed).To(HaveKeyWithValue("entity_store_url", "http://nemoentitystore-sample.nemo.svc.cluster.local:8000"))
+			Expect(parsed).To(HaveKeyWithValue("nemo_data_store_url", "http://nemodatastore-sample.nemo.svc.cluster.local:8000"))
+			Expect(parsed).To(HaveKeyWithValue("mlflow_tracking_url", "http://mlflow-tracking.nemo.svc.cluster.local:80"))
+		})
+
 		It("should delete HPA when NemoCustomizer is updated", func() {
 			namespacedName := types.NamespacedName{Name: nemoCustomizer.Name, Namespace: "default"}
 			err := client.Create(context.TODO(), nemoCustomizer)