Adding support for DRA

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

Author: visheshtanksale

api/apps/v1alpha1/nimservice_types.go

@@ -223,16 +223,6 @@ func (n *NIMService) GetLWSName() string {
 	return fmt.Sprintf("%s-lws", n.GetName())
 }
 
-// GetMultiNodeGPUsPerPod returns the number of GPUs per pod for the multi-node NIMService.
-func (n *NIMService) GetMultiNodeGPUsPerPod() int {
-	gpuQuantity, ok := n.Spec.Resources.Requests["nvidia.com/gpu"]
-	if !ok {
-		// return 0 if no GPU limit is specified because auto determine base on tp*pp/(.spec.multiNode.size) is a TODO
-		return 0
-	}
-	return int(gpuQuantity.Value())
-}
-
 // GetPVCName returns the name to be used for the PVC based on the custom spec
 // Prefers pvc.Name if explicitly set by the user in the NIMService instance.
 func (n *NIMService) GetPVCName(pvc PersistentVolumeClaim) string {
@@ -324,7 +314,7 @@ func (n *NIMService) GetStandardEnv() []corev1.EnvVar {
 	return envVars
 }
 
-func (n *NIMService) getLWSCommonEnv() []corev1.EnvVar {
+func (n *NIMService) getLWSCommonEnv(multiNodeGPUsPerPod int) []corev1.EnvVar {
 	env := n.GetEnv()
 
 	env = utils.MergeEnvVars([]corev1.EnvVar{
@@ -342,7 +332,7 @@ func (n *NIMService) getLWSCommonEnv() []corev1.EnvVar {
 		},
 		{
 			Name:  "NIM_TENSOR_PARALLEL_SIZE",
-			Value: fmt.Sprintf("%d", n.GetMultiNodeGPUsPerPod()),
+			Value: fmt.Sprintf("%d", multiNodeGPUsPerPod),
 		},
 		{
 			Name:  "NIM_PIPELINE_PARALLEL_SIZE",
@@ -360,8 +350,8 @@ func (n *NIMService) getLWSCommonEnv() []corev1.EnvVar {
 	return env
 }
 
-func (n *NIMService) GetLWSLeaderEnv() []corev1.EnvVar {
-	env := n.getLWSCommonEnv()
+func (n *NIMService) GetLWSLeaderEnv(multiNodeGPUsPerPod int) []corev1.EnvVar {
+	env := n.getLWSCommonEnv(multiNodeGPUsPerPod)
 
 	mpiTimeout := DefaultMPITimeout
 	if n.Spec.MultiNode.MPI != nil && n.Spec.MultiNode.MPI.MPIStartTimeout != 0 {
@@ -383,7 +373,7 @@ func (n *NIMService) GetLWSLeaderEnv() []corev1.EnvVar {
 		},
 		{
 			Name:  "GPUS_PER_NODE",
-			Value: fmt.Sprintf("%d", n.GetMultiNodeGPUsPerPod()),
+			Value: fmt.Sprintf("%d", multiNodeGPUsPerPod),
 		},
 		{
 			Name:  "CLUSTER_START_TIMEOUT",
@@ -409,8 +399,8 @@ func (n *NIMService) GetLWSLeaderEnv() []corev1.EnvVar {
 	return env
 }
 
-func (n *NIMService) GetLWSWorkerEnv() []corev1.EnvVar {
-	env := n.getLWSCommonEnv()
+func (n *NIMService) GetLWSWorkerEnv(multiNodeGPUsPerPod int) []corev1.EnvVar {
+	env := n.getLWSCommonEnv(multiNodeGPUsPerPod)
 	env = utils.MergeEnvVars([]corev1.EnvVar{
 		{
 			Name:  "NIM_LEADER_ROLE",
@@ -1093,7 +1083,7 @@ func (n *NIMService) GetDeploymentParams() *rendertypes.DeploymentParams {
 	return params
 }
 
-func (n *NIMService) GetLWSParams() *rendertypes.LeaderWorkerSetParams {
+func (n *NIMService) GetLWSParams(multiNodeGPUsPerPod int) *rendertypes.LeaderWorkerSetParams {
 	params := &rendertypes.LeaderWorkerSetParams{}
 
 	// Set metadata
@@ -1110,8 +1100,8 @@ func (n *NIMService) GetLWSParams() *rendertypes.LeaderWorkerSetParams {
 	params.ContainerName = n.GetContainerName()
 	params.Args = n.GetArgs()
 	params.Command = n.GetCommand()
-	params.LeaderEnvs = n.GetLWSLeaderEnv()
-	params.WorkerEnvs = n.GetLWSWorkerEnv()
+	params.LeaderEnvs = n.GetLWSLeaderEnv(multiNodeGPUsPerPod)
+	params.WorkerEnvs = n.GetLWSWorkerEnv(multiNodeGPUsPerPod)
 	params.UserID = n.GetUserID()
 	params.GroupID = n.GetGroupID()
 	params.Image = n.GetImage()
@@ -1200,14 +1190,14 @@ func (n *NIMService) GetStatefulSetParams() *rendertypes.StatefulSetParams {
 	return params
 }
 
-func (n *NIMService) generateMPIConfigData() map[string]string {
+func (n *NIMService) generateMPIConfigData(multiNodeGPUsPerPod int) map[string]string {
 	// Construct ConfigMap data
 	data := make(map[string]string)
 	for i := 0; i < n.Spec.Replicas; i++ {
-		hostfile := fmt.Sprintf("localhost slots=%d\n", n.GetMultiNodeGPUsPerPod())
+		hostfile := fmt.Sprintf("localhost slots=%d\n", multiNodeGPUsPerPod)
 		for j := 1; j < n.Spec.MultiNode.Size; j++ {
 			workerHostname := fmt.Sprintf("%s-%d-%d.%s.%s.svc slots=%d",
-				n.GetLWSName(), i, j, n.GetLWSName(), n.GetNamespace(), n.GetMultiNodeGPUsPerPod())
+				n.GetLWSName(), i, j, n.GetLWSName(), n.GetNamespace(), multiNodeGPUsPerPod)
 			hostfile += workerHostname + "\n"
 		}
 		dataKey := fmt.Sprintf("hostfile-%d", i)
@@ -1216,14 +1206,14 @@ func (n *NIMService) generateMPIConfigData() map[string]string {
 	return data
 }
 
-func (n *NIMService) GetMPIConfigParams() *rendertypes.ConfigMapParams {
+func (n *NIMService) GetMPIConfigParams(multiNodeGPUsPerPod int) *rendertypes.ConfigMapParams {
 	if n.Spec.MultiNode == nil {
 		return nil
 	}
 	return &rendertypes.ConfigMapParams{
 		Name:          fmt.Sprintf("%s-mpi-config", n.GetName()),
 		Namespace:     n.GetNamespace(),
-		ConfigMapData: n.generateMPIConfigData(),
+		ConfigMapData: n.generateMPIConfigData(multiNodeGPUsPerPod),
 		Labels:        n.GetLabels(),
 		Annotations:   n.GetAnnotations(),
 	}

internal/controller/platform/kserve/nimservice.go

@@ -411,6 +411,12 @@ func (r *NIMServiceReconciler) renderAndSyncInferenceService(ctx context.Context
 	var conType, failedCon string
 	var renderObj client.Object
 
+	multiNodeGPUsPerPod, err := shared.GetMultiNodeGPUsPerPod(ctx, r.Client, nimService)
+	if err != nil {
+		logger.Error(err, "Failed to get multi-node GPUs per pod")
+		return err
+	}
+
 	// Setup env for explicit override profile is specified
 	if modelProfile != "" {
 		profileEnv = &[]corev1.EnvVar{
@@ -432,7 +438,7 @@ func (r *NIMServiceReconciler) renderAndSyncInferenceService(ctx context.Context
 
 			// Auto assign GPU resources in case of the optimized profile
 			if profile != nil {
-				gpuResources, err = r.addGPUResources(ctx, nimService, profile)
+				gpuResources, err = r.addGPUResources(ctx, nimService, profile, multiNodeGPUsPerPod)
 				if err != nil {
 					logger.Error(err, "Failed to get GPU resources")
 					return err
@@ -445,7 +451,7 @@ func (r *NIMServiceReconciler) renderAndSyncInferenceService(ctx context.Context
 
 	initContainers = nimService.GetInitContainers()
 	namedDraResources := shared.GenerateNamedDRAResources(nimService)
-	err := r.reconcileDRAResources(ctx, nimService, namedDraResources)
+	err = r.reconcileDRAResources(ctx, nimService, namedDraResources)
 	if err != nil {
 		logger.Error(err, "Failed to reconcile DRAResources")
 		return err
@@ -548,7 +554,7 @@ func (r *NIMServiceReconciler) getNIMCacheProfile(ctx context.Context, nimServic
 // If the TP value is not present, the function defaults to allocating 1 GPU.
 //
 // In case of multi-node NIMs, this function assigns the number of GPUs equal to .spec.multiNode.gpuPerWorker.
-func (r *NIMServiceReconciler) addGPUResources(ctx context.Context, nimService *appsv1alpha1.NIMService, profile *appsv1alpha1.NIMProfile) (*corev1.ResourceRequirements, error) {
+func (r *NIMServiceReconciler) addGPUResources(ctx context.Context, nimService *appsv1alpha1.NIMService, profile *appsv1alpha1.NIMProfile, multiNodeGPUsPerPod int) (*corev1.ResourceRequirements, error) {
 	logger := log.FromContext(ctx)
 
 	// TODO: Refine this to detect GPU claims and only assign GPU resources if no GPU claims are present.
@@ -577,7 +583,7 @@ func (r *NIMServiceReconciler) addGPUResources(ctx context.Context, nimService *
 	// if deployed as multi-node, use the GPU per worker value to assign GPU resources to each worker
 	// TODO auto determine base on tp*pp/(.spec.multiNode.size)
 	if nimService.Spec.MultiNode != nil {
-		gpuQuantity, err = apiResource.ParseQuantity(fmt.Sprintf("%d", nimService.GetMultiNodeGPUsPerPod()))
+		gpuQuantity, err = apiResource.ParseQuantity(fmt.Sprintf("%d", multiNodeGPUsPerPod))
 		if err != nil {
 			logger.Error(err, "Failed to parse GPU per worker value")
 			return nil, err

internal/controller/platform/kserve/nimservice_test.go

@@ -1691,7 +1691,7 @@ var _ = Describe("NIMServiceReconciler for a KServe platform", func() {
 				},
 			}
 
-			resources, err := reconciler.addGPUResources(context.TODO(), nimService, profile)
+			resources, err := reconciler.addGPUResources(context.TODO(), nimService, profile, 8)
 			Expect(err).ToNot(HaveOccurred())
 			Expect(resources).ToNot(BeNil())
 
@@ -1706,7 +1706,7 @@ var _ = Describe("NIMServiceReconciler for a KServe platform", func() {
 				Config: map[string]string{"tp": "4"},
 			}
 
-			resources, err := reconciler.addGPUResources(context.TODO(), nimService, profile)
+			resources, err := reconciler.addGPUResources(context.TODO(), nimService, profile, 0)
 			Expect(err).ToNot(HaveOccurred())
 			Expect(resources).ToNot(BeNil())
 
@@ -1720,7 +1720,7 @@ var _ = Describe("NIMServiceReconciler for a KServe platform", func() {
 				Config: map[string]string{},
 			}
 
-			resources, err := reconciler.addGPUResources(context.TODO(), nimService, profile)
+			resources, err := reconciler.addGPUResources(context.TODO(), nimService, profile, 8)
 			Expect(err).ToNot(HaveOccurred())
 			Expect(resources).ToNot(BeNil())
 
@@ -1736,7 +1736,7 @@ var _ = Describe("NIMServiceReconciler for a KServe platform", func() {
 				Config: map[string]string{},
 			}
 
-			resources, err := reconciler.addGPUResources(context.TODO(), nimService, profile)
+			resources, err := reconciler.addGPUResources(context.TODO(), nimService, profile, 0)
 			Expect(err).ToNot(HaveOccurred())
 			Expect(resources).ToNot(BeNil())
 
@@ -1750,7 +1750,7 @@ var _ = Describe("NIMServiceReconciler for a KServe platform", func() {
 				Config: map[string]string{"tp": "invalid"},
 			}
 
-			_, err := reconciler.addGPUResources(context.TODO(), nimService, profile)
+			_, err := reconciler.addGPUResources(context.TODO(), nimService, profile, 0)
 			Expect(err).To(HaveOccurred())
 		})
 	})

internal/controller/platform/standalone/nimservice.go

@@ -349,6 +349,12 @@ func (r *NIMServiceReconciler) reconcileNIMService(ctx context.Context, nimServi
 	var conType, failedCon string
 	var renderObj client.Object
 
+	multiNodeGPUsPerPod, err := shared.GetMultiNodeGPUsPerPod(ctx, r.GetClient(), nimService)
+	if err != nil {
+		logger.Error(err, "failed to get GPU count for DRA resources")
+		return ctrl.Result{}, err
+	}
+
 	if modelProfile != "" {
 		profileEnv = &[]corev1.EnvVar{{
 			Name:  "NIM_MODEL_PROFILE",
@@ -366,7 +372,7 @@ func (r *NIMServiceReconciler) reconcileNIMService(ctx context.Context, nimServi
 
 			// Auto assign GPU resources in case of the optimized profile
 			if profile != nil {
-				gpuResources, err = r.addGPUResources(ctx, nimService, profile)
+				gpuResources, err = r.addGPUResources(ctx, nimService, profile, multiNodeGPUsPerPod)
 				if err != nil {
 					logger.Error(err, "Failed to get GPU resources")
 					return ctrl.Result{}, err
@@ -386,7 +392,7 @@ func (r *NIMServiceReconciler) reconcileNIMService(ctx context.Context, nimServi
 	}
 
 	if nimService.Spec.MultiNode != nil && nimService.Spec.MultiNode.BackendType == appsv1alpha1.NIMBackendTypeLWS {
-		lwsParams := nimService.GetLWSParams()
+		lwsParams := nimService.GetLWSParams(multiNodeGPUsPerPod)
 		lwsParams.PodResourceClaims = shared.GetPodResourceClaims(namedDraResources)
 		lwsParams.OrchestratorType = string(r.GetOrchestratorType())
 		lwsParams.LeaderVolumes = nimService.GetLeaderVolumes(*modelPVC)
@@ -428,7 +434,7 @@ func (r *NIMServiceReconciler) reconcileNIMService(ctx context.Context, nimServi
 		renderObj = &lws.LeaderWorkerSet{}
 
 		// Create configmap for MPI
-		err = r.createMultiNodeVolumeObjects(ctx, nimService)
+		err = r.createMultiNodeVolumeObjects(ctx, nimService, multiNodeGPUsPerPod)
 		if err != nil {
 			return ctrl.Result{}, fmt.Errorf("failed to create multi-node volumes: %v", err)
 		}
@@ -523,8 +529,8 @@ func (r *NIMServiceReconciler) reconcileNIMService(ctx context.Context, nimServi
 	return ctrl.Result{}, nil
 }
 
-func (r *NIMServiceReconciler) createMultiNodeVolumeObjects(ctx context.Context, nimService *appsv1alpha1.NIMService) error {
-	if err := r.createMultiNodeConfigMap(ctx, nimService, nimService.GetMPIConfigParams()); err != nil {
+func (r *NIMServiceReconciler) createMultiNodeVolumeObjects(ctx context.Context, nimService *appsv1alpha1.NIMService, multiNodeGPUsPerPod int) error {
+	if err := r.createMultiNodeConfigMap(ctx, nimService, nimService.GetMPIConfigParams(multiNodeGPUsPerPod)); err != nil {
 		return fmt.Errorf("failed to create MPI configmap for %s: %v", nimService.Name, err)
 	}
 
@@ -975,7 +981,7 @@ func (r *NIMServiceReconciler) getNIMCacheProfile(ctx context.Context, nimServic
 // If the TP value is not present, the function defaults to allocating 1 GPU.
 //
 // In case of multi-node NIMs, this function assigns the number of GPUs equal to .spec.multiNode.gpuPerWorker.
-func (r *NIMServiceReconciler) addGPUResources(ctx context.Context, nimService *appsv1alpha1.NIMService, profile *appsv1alpha1.NIMProfile) (*corev1.ResourceRequirements, error) {
+func (r *NIMServiceReconciler) addGPUResources(ctx context.Context, nimService *appsv1alpha1.NIMService, profile *appsv1alpha1.NIMProfile, multiNodeGPUsPerPod int) (*corev1.ResourceRequirements, error) {
 	logger := log.FromContext(ctx)
 
 	// TODO: Refine this to detect GPU claims and only assign GPU resources if no GPU claims are present.
@@ -1004,7 +1010,7 @@ func (r *NIMServiceReconciler) addGPUResources(ctx context.Context, nimService *
 	// if deployed as multi-node, use the GPU per worker value to assign GPU resources to each worker
 	// TODO auto determine base on tp*pp/(.spec.multiNode.size)
 	if nimService.Spec.MultiNode != nil {
-		gpuQuantity, err = apiResource.ParseQuantity(fmt.Sprintf("%d", nimService.GetMultiNodeGPUsPerPod()))
+		gpuQuantity, err = apiResource.ParseQuantity(fmt.Sprintf("%d", multiNodeGPUsPerPod))
 		if err != nil {
 			logger.Error(err, "Failed to parse GPU per worker value")
 			return nil, err
@@ -1083,5 +1089,6 @@ func (r *NIMServiceReconciler) reconcileDRAResources(ctx context.Context, nimSer
 			return fmt.Errorf("failed to reconcile DRAResource %s: %w", namedDraResource.ResourceName, err)
 		}
 	}
+
 	return nil
 }

internal/controller/platform/standalone/nimservice_test.go

@@ -1038,8 +1038,8 @@ var _ = Describe("NIMServiceReconciler for a standalone platform", func() {
 				},
 			}
 
-			leaderEnv := utils.SortKeys(nimService.GetLWSLeaderEnv())
-			workerEnv := utils.SortKeys(nimService.GetLWSWorkerEnv())
+			leaderEnv := utils.SortKeys(nimService.GetLWSLeaderEnv(8))
+			workerEnv := utils.SortKeys(nimService.GetLWSWorkerEnv(8))
 
 			Expect(reflect.DeepEqual(leaderEnv, []corev1.EnvVar{
 				{
@@ -2317,7 +2317,7 @@ var _ = Describe("NIMServiceReconciler for a standalone platform", func() {
 				},
 			}
 
-			resources, err := reconciler.addGPUResources(context.TODO(), nimService, profile)
+			resources, err := reconciler.addGPUResources(context.TODO(), nimService, profile, 8)
 			Expect(err).ToNot(HaveOccurred())
 			Expect(resources).ToNot(BeNil())
 
@@ -2332,7 +2332,7 @@ var _ = Describe("NIMServiceReconciler for a standalone platform", func() {
 				Config: map[string]string{"tp": "4"},
 			}
 
-			resources, err := reconciler.addGPUResources(context.TODO(), nimService, profile)
+			resources, err := reconciler.addGPUResources(context.TODO(), nimService, profile, 0)
 			Expect(err).ToNot(HaveOccurred())
 			Expect(resources).ToNot(BeNil())
 
@@ -2346,7 +2346,7 @@ var _ = Describe("NIMServiceReconciler for a standalone platform", func() {
 				Config: map[string]string{},
 			}
 
-			resources, err := reconciler.addGPUResources(context.TODO(), nimService, profile)
+			resources, err := reconciler.addGPUResources(context.TODO(), nimService, profile, 0)
 			Expect(err).ToNot(HaveOccurred())
 			Expect(resources).ToNot(BeNil())
 
@@ -2362,7 +2362,7 @@ var _ = Describe("NIMServiceReconciler for a standalone platform", func() {
 				Config: map[string]string{},
 			}
 
-			resources, err := reconciler.addGPUResources(context.TODO(), nimService, profile)
+			resources, err := reconciler.addGPUResources(context.TODO(), nimService, profile, 0)
 			Expect(err).ToNot(HaveOccurred())
 			Expect(resources).ToNot(BeNil())
 
@@ -2376,7 +2376,7 @@ var _ = Describe("NIMServiceReconciler for a standalone platform", func() {
 				Config: map[string]string{"tp": "invalid"},
 			}
 
-			_, err := reconciler.addGPUResources(context.TODO(), nimService, profile)
+			_, err := reconciler.addGPUResources(context.TODO(), nimService, profile, 0)
 			Expect(err).To(HaveOccurred())
 		})
 	})

internal/k8sutil/resourceclaimsutil.go

@@ -74,3 +74,11 @@ func GetResourceClaimState(claim *resourcev1beta2.ResourceClaim) string {
 	}
 	return strings.Join(states, ",")
 }
+
+func GetResourceClaimTemplate(ctx context.Context, k8sclient client.Client, name string, namespace string) (*resourcev1beta2.ResourceClaimTemplate, error) {
+	claimTemplate := &resourcev1beta2.ResourceClaimTemplate{}
+	if err := k8sclient.Get(ctx, client.ObjectKey{Name: name, Namespace: namespace}, claimTemplate); err != nil {
+		return nil, err
+	}
+	return claimTemplate, nil
+}