Improving the condition for adding HF_TOKEN on NIM Service (#711)

* Improving the condition for adding HF_TOKEN on NIM Service

Signed-off-by: Vishesh Tanksale <[email protected]>

* Addressing review comments

Signed-off-by: Vishesh Tanksale <[email protected]>

* Addressing review comments

Signed-off-by: Vishesh Tanksale <[email protected]>

* Addressing review comments

Signed-off-by: Vishesh Tanksale <[email protected]>

---------

Signed-off-by: Vishesh Tanksale <[email protected]>

Author: visheshtanksale

api/apps/v1alpha1/common_types.go

@@ -36,6 +36,10 @@ const (
 	DefaultNamedPortGRPC = "grpc"
 	// DefaultNamedPortMetrics is the default name for metrics port.
 	DefaultNamedPortMetrics = "metrics"
+	// NGCAPIKey is the environment variable name for NGC API key.
+	NGCAPIKey = "NGC_API_KEY"
+	// HFToken is the environment variable name for Hugging Face token.
+	HFToken = "HF_TOKEN"
 )
 
 // Expose defines attributes to expose the service.

api/apps/v1alpha1/nimcache_types.go

@@ -331,6 +331,11 @@ func (n *NIMCache) IsUniversalNIM() bool {
 	return false
 }
 
+// IsHFModel returns true if the NIMCache is for a Hugging Face model.
+func (n *NIMCache) IsHFModel() bool {
+	return n.Spec.Source.HF != nil || n.Spec.Source.DataStore != nil
+}
+
 // IsOptimizedNIM returns true if the NIMCache is for an optimized NIM.
 func (n *NIMCache) IsOptimizedNIM() bool {
 	// Universal NIM is when the modelEndpoint is set in the NGCSource.

api/apps/v1alpha1/nimservice_types.go

@@ -25,6 +25,7 @@ import (
 	"fmt"
 	"maps"
 	"os"
+	"strings"
 
 	kserveconstants "github.com/kserve/kserve/pkg/constants"
 	monitoringv1 "github.com/prometheus-operator/prometheus-operator/pkg/apis/monitoring/v1"
@@ -283,13 +284,13 @@ func (n *NIMService) GetStandardEnv() []corev1.EnvVar {
 			Value: utils.DefaultModelStorePath,
 		},
 		{
-			Name: "NGC_API_KEY",
+			Name: NGCAPIKey,
 			ValueFrom: &corev1.EnvVarSource{
 				SecretKeyRef: &corev1.SecretKeySelector{
 					LocalObjectReference: corev1.LocalObjectReference{
 						Name: n.Spec.AuthSecret,
 					},
-					Key: "NGC_API_KEY",
+					Key: NGCAPIKey,
 				},
 			},
 		},
@@ -983,6 +984,17 @@ func (n *NIMService) IsHTTPRouteEnabled() bool {
 	return n.Spec.Router.Gateway != nil && n.Spec.Router.Gateway.HTTPRoutesEnabled
 }
 
+// IsHFModel returns true if the NIMService is using an HF model.
+func (n *NIMService) IsHFModel() bool {
+	env := utils.FindEnvByValue(n.GetEnv(), "NIM_MODEL_NAME")
+	if env != nil {
+		if strings.HasPrefix(env.Value, "hf://") {
+			return true
+		}
+	}
+	return false
+}
+
 func (n *NIMService) GetHTTPRouteSpec() gatewayv1.HTTPRouteSpec {
 	return n.Spec.Router.GenerateGatewayHTTPRouteSpec(n.GetNamespace(), n.GetName(), n.GetServicePort())
 }

internal/controller/nimbuild_controller.go

@@ -642,13 +642,13 @@ func (r *NIMBuildReconciler) constructEngineBuildPod(nimBuild *appsv1alpha1.NIMB
 					Value: inputNimProfile.Name,
 				},
 				{
-					Name: "NGC_API_KEY",
+					Name: appsv1alpha1.NGCAPIKey,
 					ValueFrom: &corev1.EnvVarSource{
 						SecretKeyRef: &corev1.SecretKeySelector{
 							LocalObjectReference: corev1.LocalObjectReference{
 								Name: nimCache.Spec.Source.NGC.AuthSecret,
 							},
-							Key: "NGC_API_KEY",
+							Key: appsv1alpha1.NGCAPIKey,
 						},
 					},
 				},

internal/controller/platform/kserve/nimservice.go

@@ -477,28 +477,24 @@ 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",
-							},
+	// If NIMCache or NIMService is a Hugging Face Multi-LLM NIM, add the HF_TOKEN to the environment variables
+	if nimCache.IsHFModel() || nimService.IsHFModel() {
+		isvcParams.Env = utils.RemoveEnvVar(isvcParams.Env, appsv1alpha1.NGCAPIKey)
+		isvcParams.Env = utils.MergeEnvVars(isvcParams.Env, []corev1.EnvVar{
+			{
+				Name: appsv1alpha1.HFToken,
+				ValueFrom: &corev1.EnvVarSource{
+					SecretKeyRef: &corev1.SecretKeySelector{
+						LocalObjectReference: corev1.LocalObjectReference{
+							Name: nimService.Spec.AuthSecret,
 						},
+						Key: appsv1alpha1.HFToken,
 					},
-				})
-			}
-		}
+				},
+			},
+		})
 	}
+
 	// Setup volume mounts with model store
 	isvcParams.Volumes = nimService.GetVolumes(modelPVC)
 	isvcParams.VolumeMounts = nimService.GetVolumeMounts(modelPVC)