Revert "Removing NGC Home path and HF_HOME env (#202)"

This reverts commit 425f91d. Signed-off-by: Vishesh Tanksale <[email protected]>
NVIDIA · Nov 5, 2024 · a7f4823 · a7f4823
1 parent 425f91d
commit a7f4823
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 7 deletions.
diff --git a/internal/controller/nimcache_controller.go b/internal/controller/nimcache_controller.go
@@ -1075,13 +1075,21 @@ func (r *NIMCacheReconciler) constructJob(ctx context.Context, nimCache *appsv1a
 			{
 				Name:    NIMCacheContainerName,
 				Image:   nimCache.Spec.Source.NGC.ModelPuller,
-				Args:    []string{"download-to-cache"},
+				Command: []string{"download-to-cache"},
 				EnvFrom: nimCache.Spec.Source.EnvFromSecrets(),
 				Env: []corev1.EnvVar{
+					{
+						Name:  "HF_HOME",
+						Value: "/model-store/huggingface",
+					},
 					{
 						Name:  "NIM_CACHE_PATH",
 						Value: "/model-store",
 					},
+					{
+						Name:  "NGC_HOME",
+						Value: "/model-store/ngc",
+					},
 				},
 				VolumeMounts: []corev1.VolumeMount{
 					{
@@ -1118,7 +1126,6 @@ func (r *NIMCacheReconciler) constructJob(ctx context.Context, nimCache *appsv1a
 				Name: nimCache.Spec.Source.NGC.PullSecret,
 			},
 		}
-
 		// Pass specific profiles to download based on user selection or auto-selection
 		selectedProfiles, err := getSelectedProfiles(nimCache)
 		if err != nil {
@@ -1132,9 +1139,9 @@ func (r *NIMCacheReconciler) constructJob(ctx context.Context, nimCache *appsv1a
 
 		if len(selectedProfiles) > 0 {
 			if utils.ContainsElement(selectedProfiles, AllProfiles) {
-				job.Spec.Template.Spec.Containers[0].Args = append(job.Spec.Template.Spec.Containers[0].Args, "--all")
+				job.Spec.Template.Spec.Containers[0].Args = []string{"--all"}
 			} else {
-				job.Spec.Template.Spec.Containers[0].Args = append(job.Spec.Template.Spec.Containers[0].Args, "--profiles")
+				job.Spec.Template.Spec.Containers[0].Args = []string{"--profiles"}
 				job.Spec.Template.Spec.Containers[0].Args = append(job.Spec.Template.Spec.Containers[0].Args, selectedProfiles...)
 			}
 		}

diff --git a/internal/controller/nimcache_controller_test.go b/internal/controller/nimcache_controller_test.go
@@ -393,7 +393,8 @@ var _ = Describe("NIMCache Controller", func() {
 			Expect(job.Name).To(Equal(getJobName(nimCache)))
 			Expect(job.Spec.Template.Spec.Containers[0].Image).To(Equal("nvcr.io/nim:test"))
 			Expect(job.Spec.Template.Spec.ImagePullSecrets[0].Name).To(Equal("my-secret"))
-			Expect(job.Spec.Template.Spec.Containers[0].Args).To(ContainElements("download-to-cache", "--profiles", "36fc1fa4fc35c1d54da115a39323080b08d7937dceb8ba47be44f4da0ec720ff"))
+			Expect(job.Spec.Template.Spec.Containers[0].Command).To(ContainElements("download-to-cache"))
+			Expect(job.Spec.Template.Spec.Containers[0].Args).To(ContainElements("--profiles", "36fc1fa4fc35c1d54da115a39323080b08d7937dceb8ba47be44f4da0ec720ff"))
 			Expect(*job.Spec.Template.Spec.SecurityContext.RunAsUser).To(Equal(int64(1000)))
 			Expect(*job.Spec.Template.Spec.SecurityContext.FSGroup).To(Equal(int64(2000)))
 			Expect(*job.Spec.Template.Spec.SecurityContext.RunAsNonRoot).To(Equal(true))
@@ -423,7 +424,8 @@ var _ = Describe("NIMCache Controller", func() {
 			Expect(job.Name).To(Equal(getJobName(nimCache)))
 			Expect(job.Spec.Template.Spec.Containers[0].Image).To(Equal("nvcr.io/nim:test"))
 			Expect(job.Spec.Template.Spec.ImagePullSecrets[0].Name).To(Equal("my-secret"))
-			Expect(job.Spec.Template.Spec.Containers[0].Args).To(ContainElements("download-to-cache", "--profiles", "36fc1fa4fc35c1d54da115a39323080b08d7937dceb8ba47be44f4da0ec720ff", "04fdb4d11f01be10c31b00e7c0540e2835e89a0079b483ad2dd3c25c8cc12345"))
+			Expect(job.Spec.Template.Spec.Containers[0].Command).To(ContainElements("download-to-cache"))
+			Expect(job.Spec.Template.Spec.Containers[0].Args).To(ContainElements("--profiles", "36fc1fa4fc35c1d54da115a39323080b08d7937dceb8ba47be44f4da0ec720ff", "04fdb4d11f01be10c31b00e7c0540e2835e89a0079b483ad2dd3c25c8cc12345"))
 			Expect(*job.Spec.Template.Spec.SecurityContext.RunAsUser).To(Equal(int64(1000)))
 			Expect(*job.Spec.Template.Spec.SecurityContext.FSGroup).To(Equal(int64(2000)))
 			Expect(*job.Spec.Template.Spec.SecurityContext.RunAsNonRoot).To(Equal(true))
@@ -449,7 +451,8 @@ var _ = Describe("NIMCache Controller", func() {
 			Expect(job.Name).To(Equal(getJobName(nimCache)))
 			Expect(job.Spec.Template.Spec.Containers[0].Image).To(Equal("nvcr.io/nim:test"))
 			Expect(job.Spec.Template.Spec.ImagePullSecrets[0].Name).To(Equal("my-secret"))
-			Expect(job.Spec.Template.Spec.Containers[0].Args).To(ContainElements("download-to-cache", "--all"))
+			Expect(job.Spec.Template.Spec.Containers[0].Command).To(ContainElements("download-to-cache"))
+			Expect(job.Spec.Template.Spec.Containers[0].Args).To(ContainElements("--all"))
 		})
 
 		It("should create a job with the correct specifications", func() {