Adding HF token env var when using multi llm NIM without caching (#709)

Signed-off-by: Vishesh Tanksale <vtanksale@nvidia.com>

Author: visheshtanksale

api/apps/v1alpha1/nimservice_types.go

@@ -279,6 +279,17 @@ func (n *NIMService) GetStandardEnv() []corev1.EnvVar {
 			Name:  "NIM_CACHE_PATH",
 			Value: utils.DefaultModelStorePath,
 		},
+		{
+			Name: "NGC_API_KEY",
+			ValueFrom: &corev1.EnvVarSource{
+				SecretKeyRef: &corev1.SecretKeySelector{
+					LocalObjectReference: corev1.LocalObjectReference{
+						Name: n.Spec.AuthSecret,
+					},
+					Key: "NGC_API_KEY",
+				},
+			},
+		},
 		{
 			Name:  "OUTLINES_CACHE_DIR",
 			Value: "/tmp/outlines",
@@ -355,11 +366,6 @@ func (n *NIMService) getLWSCommonEnv() []corev1.EnvVar {
 	return env
 }
 
-// GetLWSLeaderEnvFrom returns the env from sources for the leader worker set.
-func (n *NIMService) GetLWSCommonEnvFrom() []corev1.EnvFromSource {
-	return n.GetEnvFrom()
-}
-
 func (n *NIMService) GetLWSLeaderEnv() []corev1.EnvVar {
 	env := n.getLWSCommonEnv()
 
@@ -564,23 +570,6 @@ func (n *NIMService) GetEnv() []corev1.EnvVar {
 	return envVarList
 }
 
-// GetEnvFrom returns merged slice of standard and user specified env from sources.
-func (n *NIMService) GetEnvFrom() []corev1.EnvFromSource {
-	if n.Spec.AuthSecret != "" {
-		return []corev1.EnvFromSource{
-			{
-				SecretRef: &corev1.SecretEnvSource{
-					LocalObjectReference: corev1.LocalObjectReference{
-						Name: n.Spec.AuthSecret,
-					},
-				},
-			},
-		}
-	}
-	// no secrets to source the env variables
-	return []corev1.EnvFromSource{}
-}
-
 // GetImage returns container image for the NIMService.
 func (n *NIMService) GetImage() string {
 	return fmt.Sprintf("%s:%s", n.Spec.Image.Repository, n.Spec.Image.Tag)
@@ -1065,7 +1054,6 @@ func (n *NIMService) GetDeploymentParams() *rendertypes.DeploymentParams {
 	// Set container spec
 	params.ContainerName = n.GetContainerName()
 	params.Env = n.GetEnv()
-	params.EnvFrom = n.GetEnvFrom()
 	params.Args = n.GetArgs()
 	params.Command = n.GetCommand()
 	params.Resources = n.GetResources()
@@ -1137,8 +1125,6 @@ func (n *NIMService) GetLWSParams() *rendertypes.LeaderWorkerSetParams {
 	params.Command = n.GetCommand()
 	params.LeaderEnvs = n.GetLWSLeaderEnv()
 	params.WorkerEnvs = n.GetLWSWorkerEnv()
-	params.LeaderEnvFrom = n.GetLWSCommonEnvFrom()
-	params.WorkerEnvFrom = n.GetLWSCommonEnvFrom()
 	params.UserID = n.GetUserID()
 	params.GroupID = n.GetGroupID()
 	params.Image = n.GetImage()

internal/controller/platform/kserve/nimservice.go

@@ -477,6 +477,28 @@ func (r *NIMServiceReconciler) renderAndSyncInferenceService(ctx context.Context
 			},
 		}, isvcParams.Env)
 	}
+	// If NIMCache is not provided, and the model is from Hugging Face, add the HF_TOKEN to the environment variables
+	if nimService.GetNIMCacheName() == "" {
+		env := utils.FindEnvByValue(isvcParams.Env, "NIM_MODEL_NAME")
+		if env != nil {
+			if strings.HasPrefix(env.Value, "hf://") {
+				isvcParams.Env = utils.RemoveEnvVar(isvcParams.Env, "NGC_API_KEY")
+				isvcParams.Env = utils.MergeEnvVars(isvcParams.Env, []corev1.EnvVar{
+					{
+						Name: "HF_TOKEN",
+						ValueFrom: &corev1.EnvVarSource{
+							SecretKeyRef: &corev1.SecretKeySelector{
+								LocalObjectReference: corev1.LocalObjectReference{
+									Name: nimService.Spec.AuthSecret,
+								},
+								Key: "HF_TOKEN",
+							},
+						},
+					},
+				})
+			}
+		}
+	}
 	// Setup volume mounts with model store
 	isvcParams.Volumes = nimService.GetVolumes(modelPVC)
 	isvcParams.VolumeMounts = nimService.GetVolumeMounts(modelPVC)

internal/controller/platform/standalone/nimservice.go

@@ -448,6 +448,28 @@ func (r *NIMServiceReconciler) reconcileNIMService(ctx context.Context, nimServi
 				Value: utils.DefaultModelStorePath,
 			}}, deploymentParams.Env)
 		}
+		// If NIMCache is not provided, and the model is from Hugging Face, add the HF_TOKEN to the environment variables
+		if nimService.GetNIMCacheName() == "" {
+			env := utils.FindEnvByValue(deploymentParams.Env, "NIM_MODEL_NAME")
+			if env != nil {
+				if strings.HasPrefix(env.Value, "hf://") {
+					deploymentParams.Env = utils.RemoveEnvVar(deploymentParams.Env, "NGC_API_KEY")
+					deploymentParams.Env = utils.MergeEnvVars(deploymentParams.Env, []corev1.EnvVar{
+						{
+							Name: "HF_TOKEN",
+							ValueFrom: &corev1.EnvVarSource{
+								SecretKeyRef: &corev1.SecretKeySelector{
+									LocalObjectReference: corev1.LocalObjectReference{
+										Name: nimService.Spec.AuthSecret,
+									},
+									Key: "HF_TOKEN",
+								},
+							},
+						},
+					})
+				}
+			}
+		}
 		// Setup volume mounts with model store
 		deploymentParams.Volumes = nimService.GetVolumes(modelPVC)
 		deploymentParams.VolumeMounts = nimService.GetVolumeMounts(modelPVC)

internal/controller/platform/standalone/nimservice_test.go

@@ -1123,6 +1123,17 @@ var _ = Describe("NIMServiceReconciler for a standalone platform", func() {
 					Name:  "NIM_CACHE_PATH",
 					Value: "/model-store",
 				},
+				{
+					Name: "NGC_API_KEY",
+					ValueFrom: &corev1.EnvVarSource{
+						SecretKeyRef: &corev1.SecretKeySelector{
+							LocalObjectReference: corev1.LocalObjectReference{
+								Name: "",
+							},
+							Key: "NGC_API_KEY",
+						},
+					},
+				},
 				{
 					Name:  "OUTLINES_CACHE_DIR",
 					Value: "/tmp/outlines",
@@ -1214,6 +1225,17 @@ var _ = Describe("NIMServiceReconciler for a standalone platform", func() {
 					Name:  "NIM_CACHE_PATH",
 					Value: "/model-store",
 				},
+				{
+					Name: "NGC_API_KEY",
+					ValueFrom: &corev1.EnvVarSource{
+						SecretKeyRef: &corev1.SecretKeySelector{
+							LocalObjectReference: corev1.LocalObjectReference{
+								Name: "",
+							},
+							Key: "NGC_API_KEY",
+						},
+					},
+				},
 				{
 					Name:  "OUTLINES_CACHE_DIR",
 					Value: "/tmp/outlines",

internal/render/types/types.go

@@ -70,7 +70,6 @@ type DeploymentParams struct {
 	Volumes            []corev1.Volume
 	VolumeMounts       []corev1.VolumeMount
 	Env                []corev1.EnvVar
-	EnvFrom            []corev1.EnvFromSource
 	Resources          *corev1.ResourceRequirements
 	NodeSelector       map[string]string
 	Tolerations        []corev1.Toleration
@@ -112,8 +111,6 @@ type LeaderWorkerSetParams struct {
 	LeaderVolumeMounts []corev1.VolumeMount
 	WorkerEnvs         []corev1.EnvVar
 	LeaderEnvs         []corev1.EnvVar
-	WorkerEnvFrom      []corev1.EnvFromSource
-	LeaderEnvFrom      []corev1.EnvFromSource
 	Resources          *corev1.ResourceRequirements
 	NodeSelector       map[string]string
 	Tolerations        []corev1.Toleration

internal/utils/utils.go

@@ -485,3 +485,22 @@ func IsVersionGreaterThanOrEqual(version string, minVersion string) bool {
 	}
 	return cv.AtLeast(mv)
 }
+
+func RemoveEnvVar(envs []corev1.EnvVar, name string) []corev1.EnvVar {
+	result := make([]corev1.EnvVar, 0, len(envs))
+	for _, e := range envs {
+		if e.Name != name {
+			result = append(result, e)
+		}
+	}
+	return result
+}
+
+func FindEnvByValue(envs []corev1.EnvVar, key string) *corev1.EnvVar {
+	for i := range envs {
+		if envs[i].Name == key {
+			return &envs[i]
+		}
+	}
+	return nil
+}

manifests/deployment.yaml

@@ -98,10 +98,6 @@ spec:
         {{- if .Env }}
           {{- .Env | yaml | nindent 10 }}
         {{- end }}
-        envFrom:
-        {{- if .EnvFrom }}
-          {{- .EnvFrom | yaml | nindent 10 }}
-        {{- end }}
         {{- with .Resources }}
         resources:
           {{ . | yaml | nindent 10 }}