fix(operator): aggregate discovered GPU counts across matching nodes

components/src/dynamo/profiler/utils/dgdr_v1beta1_types.py

@@ -229,19 +229,19 @@ class HardwareSpec(BaseModel):
 
     gpuSku: Optional[GPUSKUType] = Field(
         default=None,
-        description="GPUSKU selects the GPU type to target. When omitted, auto-detected by selecting the GPU with the highest node count, then highest VRAM. In mixed-GPU clusters, set this to choose which GPU type to use. Discovery and totalGpus are then restricted to nodes matching this SKU. The legacy value gb200_sxm is deprecated; use gb200 instead.",
+        description="GPUSKU selects the GPU type to target. When omitted, auto-detected as the GPU type of the node with the most GPUs (ties broken by highest VRAM per GPU). In mixed-GPU clusters, set this to choose which GPU type to use. Discovery and totalGpus are then restricted to nodes matching this SKU. The legacy value gb200_sxm is deprecated; use gb200 instead.",
     )
     vramMb: Optional[float] = Field(
         default=None,
         description="VRAMMB is the VRAM per GPU in MiB. When omitted, auto-detected from cluster GPU nodes.",
     )
     totalGpus: Optional[int] = Field(
         default=None,
-        description="TotalGPUs is the GPU budget for profiling and deployment. The profiler uses this to determine parallelism and replica count. When omitted, computed by counting GPUs on discovered nodes (filtered by gpuSku when set), temporarily capped at 32 to limit profiler search space. This cap may be removed in a future release. Set this field explicitly to override.",
+        description="TotalGPUs is the GPU budget for profiling and deployment. The profiler uses this to determine parallelism and replica count. When omitted, computed by summing GPUs across discovered nodes of the selected GPU type (gpuSku, or the auto-detected type), temporarily capped at 32 to limit profiler search space. This cap may be removed in a future release. Set this field explicitly to override.",
     )
     numGpusPerNode: Optional[int] = Field(
         default=None,
-        description="NumGPUsPerNode is the number of GPUs per node. When omitted, auto-detected from cluster GPU nodes.",
+        description="NumGPUsPerNode is the number of GPUs per node. When omitted, auto-detected as the highest GPU count on any discovered node of the selected GPU type. In clusters whose nodes have different GPU counts, set this explicitly to the node size you intend to schedule on.",
     )
     interconnect: Optional[str] = Field(
         default=None,

deploy/operator/api/v1beta1/dynamographdeploymentrequest_types.go

@@ -407,8 +407,8 @@ type FeaturesSpec struct {
 // If all four fields are set, discovery is skipped.
 type HardwareSpec struct {
 	// GPUSKU selects the GPU type to target.
-	// When omitted, auto-detected by selecting the GPU with the highest
-	// node count, then highest VRAM. In mixed-GPU clusters, set this to
+	// When omitted, auto-detected as the GPU type of the node with the most GPUs
+	// (ties broken by highest VRAM per GPU). In mixed-GPU clusters, set this to
 	// choose which GPU type to use. Discovery and totalGpus are then
 	// restricted to nodes matching this SKU.
 	// The legacy value gb200_sxm is deprecated; use gb200 instead.
@@ -423,15 +423,17 @@ type HardwareSpec struct {
 
 	// TotalGPUs is the GPU budget for profiling and deployment.
 	// The profiler uses this to determine parallelism and replica count.
-	// When omitted, computed by counting GPUs on discovered nodes
-	// (filtered by gpuSku when set), temporarily capped at 32 to
+	// When omitted, computed by summing GPUs across discovered nodes of the
+	// selected GPU type (gpuSku, or the auto-detected type), temporarily capped at 32 to
 	// limit profiler search space. This cap may be removed in a future
 	// release. Set this field explicitly to override.
 	// +optional
 	TotalGPUs *int32 `json:"totalGpus,omitempty"`
 
 	// NumGPUsPerNode is the number of GPUs per node.
-	// When omitted, auto-detected from cluster GPU nodes.
+	// When omitted, auto-detected as the highest GPU count on any discovered node
+	// of the selected GPU type. In clusters whose nodes have different GPU counts,
+	// set this explicitly to the node size you intend to schedule on.
 	// +optional
 	NumGPUsPerNode *int32 `json:"numGpusPerNode,omitempty"`
 	// Interconnect describes the primary GPU-to-GPU interconnect *within a node*.

deploy/operator/config/crd/bases/nvidia.com_dynamographdeploymentrequests.yaml

@@ -663,8 +663,8 @@ spec:
                             - mi300
                       description: |-
                         GPUSKU selects the GPU type to target.
-                        When omitted, auto-detected by selecting the GPU with the highest
-                        node count, then highest VRAM. In mixed-GPU clusters, set this to
+                        When omitted, auto-detected as the GPU type of the node with the most GPUs
+                        (ties broken by highest VRAM per GPU). In mixed-GPU clusters, set this to
                         choose which GPU type to use. Discovery and totalGpus are then
                         restricted to nodes matching this SKU.
                         The legacy value gb200_sxm is deprecated; use gb200 instead.
@@ -693,7 +693,9 @@ spec:
                     numGpusPerNode:
                       description: |-
                         NumGPUsPerNode is the number of GPUs per node.
-                        When omitted, auto-detected from cluster GPU nodes.
+                        When omitted, auto-detected as the highest GPU count on any discovered node
+                        of the selected GPU type. In clusters whose nodes have different GPU counts,
+                        set this explicitly to the node size you intend to schedule on.
                       format: int32
                       type: integer
                     rdma:
@@ -721,8 +723,8 @@ spec:
                       description: |-
                         TotalGPUs is the GPU budget for profiling and deployment.
                         The profiler uses this to determine parallelism and replica count.
-                        When omitted, computed by counting GPUs on discovered nodes
-                        (filtered by gpuSku when set), temporarily capped at 32 to
+                        When omitted, computed by summing GPUs across discovered nodes of the
+                        selected GPU type (gpuSku, or the auto-detected type), temporarily capped at 32 to
                         limit profiler search space. This cap may be removed in a future
                         release. Set this field explicitly to override.
                       format: int32

deploy/operator/internal/controller/dynamographdeploymentrequest_controller.go

@@ -1900,7 +1900,7 @@ func (r *DynamoGraphDeploymentRequestReconciler) enrichHardwareFromDiscovery(ctx
 	logger.Info("GPU discovery completed successfully",
 		"gpusPerNode", gpuInfo.GPUsPerNode,
 		"nodesWithGPUs", gpuInfo.NodesWithGPUs,
-		"totalGpus", gpuInfo.GPUsPerNode*gpuInfo.NodesWithGPUs,
+		"totalGpus", gpuInfo.TotalGPUs,
 		"model", gpuInfo.Model,
 		"vramMiB", gpuInfo.VRAMPerGPU,
 		"system", gpuInfo.System,
@@ -1935,14 +1935,14 @@ func (r *DynamoGraphDeploymentRequestReconciler) enrichHardwareFromDiscovery(ctx
 	if hw.TotalGPUs == nil && gpuInfo != nil {
 		// TODO: This is a temporary limit to prevent the profiler from using too many GPUs.
 		// Will be removed once a fix is in the Profiler/AIC.
-		const defaultMaxAutoGPUs = int32(32)
-		total := int32(gpuInfo.GPUsPerNode * gpuInfo.NodesWithGPUs)
+		const defaultMaxAutoGPUs = 32
+		total := gpuInfo.TotalGPUs
 		if total > defaultMaxAutoGPUs {
 			logger.Info("Capping auto-discovered TotalGPUs at default limit; set hardware.totalGpus to override",
 				"discovered", total, "cap", defaultMaxAutoGPUs)
 			total = defaultMaxAutoGPUs
 		}
-		hw.TotalGPUs = &total
+		hw.TotalGPUs = new(int32(total))
 		changed = true
 	}
 	if hw.Interconnect == "" && gpuInfo != nil && gpuInfo.Interconnect != "" {

deploy/operator/internal/controller/dynamographdeploymentrequest_controller_envtest_test.go

@@ -23,6 +23,7 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
+	"fmt"
 	"slices"
 	"strings"
 	"time"
@@ -3007,13 +3008,13 @@ spec:
 			Expect(k8sClient.Create(ctx, dgdr)).Should(Succeed())
 			defer func() { _ = k8sClient.Delete(ctx, dgdr) }()
 
-			GinkgoT().Log("creating two eight-GPU nodes for node-label discovery")
+			GinkgoT().Log("creating eight-GPU and four-GPU nodes for node-label discovery")
 			gpuNodes := make([]*corev1.Node, 0, 2)
-			for _, nodeName := range []string{"gpu-worker-budget-1", "gpu-worker-budget-2"} {
+			for i, count := range []string{"8", "4"} {
 				node := &corev1.Node{
-					Name: nodeName,
+					Name: fmt.Sprintf("gpu-worker-budget-%d", i),
 					Labels: map[string]string{
-						gpu.LabelGPUCount:   "8",
+						gpu.LabelGPUCount:   count,
 						gpu.LabelGPUProduct: "H200-SXM5-141GB",
 						gpu.LabelGPUMemory:  "141312",
 					},
@@ -3042,7 +3043,7 @@ spec:
 			Expect(k8sClient.Get(ctx, request.NamespacedName, &persisted)).Should(Succeed())
 			Expect(persisted.Spec.Hardware).NotTo(BeNil())
 			Expect(persisted.Spec.Hardware.TotalGPUs).NotTo(BeNil())
-			Expect(*persisted.Spec.Hardware.TotalGPUs).To(Equal(int32(16)))
+			Expect(*persisted.Spec.Hardware.TotalGPUs).To(Equal(int32(12)))
 
 			GinkgoT().Log("creating the profiling Job from the persisted DGDR")
 			_, err = reconciler.Reconcile(ctx, request)
@@ -3110,6 +3111,7 @@ spec:
 				VRAMPerGPU:    81920,
 				System:        "h100_sxm",
 				NodesWithGPUs: 1,
+				TotalGPUs:     8,
 			}
 			cache := gpu.NewGPUDiscoveryCache()
 			cache.Set("", mockGPU, 10*time.Minute)
@@ -3224,6 +3226,7 @@ spec:
 				VRAMPerGPU:    81920,
 				System:        "h100_sxm",
 				NodesWithGPUs: 1,
+				TotalGPUs:     8,
 			}
 			cache := gpu.NewGPUDiscoveryCache()
 			cache.Set("", mockGPU, 10*time.Minute)
@@ -3339,6 +3342,7 @@ spec:
 				VRAMPerGPU:    81920,
 				System:        "h100_sxm",
 				NodesWithGPUs: 1,
+				TotalGPUs:     8,
 			}
 			cache := gpu.NewGPUDiscoveryCache()
 			cache.Set("", mockGPU, 10*time.Minute)
@@ -3961,6 +3965,7 @@ spec:
 				VRAMPerGPU:    81920,
 				System:        "a100_sxm",
 				NodesWithGPUs: 1,
+				TotalGPUs:     8,
 			}
 			cache := gpu.NewGPUDiscoveryCache()
 			cache.Set("", mockGPU, 10*time.Minute)

deploy/operator/internal/controller/enrich_hardware_test.go

@@ -273,6 +273,46 @@ func TestEnrichHardwareFromDiscovery(t *testing.T) {
 	}
 }
 
+func TestEnrichHardwareFromDiscovery_TotalGPUs(t *testing.T) {
+	tests := []struct {
+		name          string
+		counts        []int
+		explicitTotal *int32
+		wantTotal     int32
+	}{
+		{name: "heterogeneous nodes", counts: []int{8, 4}, wantTotal: 12},
+		{name: "aggregate below cap despite larger maximum times nodes", counts: []int{8, 8, 8, 4, 2}, wantTotal: 30},
+		{name: "aggregate above cap", counts: []int{8, 8, 8, 8, 4}, wantTotal: 32},
+		{name: "explicit total above cap", counts: []int{8, 4}, explicitTotal: new(int32(64)), wantTotal: 64},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Log("creating H100 nodes and an unrelated A100 node")
+			objects := []client.Object{gpuNode("a100-node", "A100-SXM4-40GB", 8, 40960)}
+			for i, count := range tt.counts {
+				objects = append(objects, gpuNode(fmt.Sprintf("h100-node-%d", i), "H100-SXM5-80GB", count, 81920))
+			}
+			r := newFakeReconciler(objects...)
+			dgdr := &nvidiacomv1beta1.DynamoGraphDeploymentRequest{
+				Spec: nvidiacomv1beta1.DynamoGraphDeploymentRequestSpec{
+					Hardware: &nvidiacomv1beta1.HardwareSpec{
+						GPUSKU:    nvidiacomv1beta1.GPUSKUTypeH100SXM,
+						TotalGPUs: tt.explicitTotal,
+					},
+				},
+			}
+
+			t.Log("enriching missing hardware fields while preserving an explicit budget")
+			changed, err := r.enrichHardwareFromDiscovery(t.Context(), dgdr)
+			require.NoError(t, err)
+			require.True(t, changed)
+			require.NotNil(t, dgdr.Spec.Hardware.TotalGPUs)
+			assert.Equal(t, tt.wantTotal, *dgdr.Spec.Hardware.TotalGPUs)
+			assert.Equal(t, int32(8), *dgdr.Spec.Hardware.NumGPUsPerNode)
+		})
+	}
+}
+
 func TestEnrichHardwareFromDiscovery_WritesOptionalHardwareMetadata(t *testing.T) {
 	r := newFakeReconciler()
 	r.GPUDiscovery = gpupkg.NewGPUDiscovery(nil)
@@ -281,6 +321,7 @@ func TestEnrichHardwareFromDiscovery_WritesOptionalHardwareMetadata(t *testing.T
 		NodeName:      "n1",
 		GPUsPerNode:   8,
 		NodesWithGPUs: 2,
+		TotalGPUs:     16,
 		Model:         "NVIDIA H100 80GB HBM3",
 		VRAMPerGPU:    81079,
 		System:        nvidiacomv1beta1.GPUSKUTypeH100PCIe,
@@ -313,6 +354,7 @@ func TestEnrichHardwareFromDiscovery_FillsOptionalMetadataWhenRequiredFieldsSet(
 		NodeName:      "n1",
 		GPUsPerNode:   8,
 		NodesWithGPUs: 4,
+		TotalGPUs:     32,
 		Model:         "NVIDIA H100 80GB HBM3",
 		VRAMPerGPU:    81079,
 		System:        nvidiacomv1beta1.GPUSKUTypeH100PCIe,
@@ -367,6 +409,7 @@ func TestCreateProfilingJobPersistsDiscoveredHardware(t *testing.T) {
 		NodeName:      "n1",
 		GPUsPerNode:   8,
 		NodesWithGPUs: 2,
+		TotalGPUs:     16,
 		Model:         "NVIDIA H100 80GB HBM3",
 		VRAMPerGPU:    81079,
 		System:        nvidiacomv1beta1.GPUSKUTypeH100PCIe,

deploy/operator/internal/gpu/discovery.go

@@ -165,8 +165,9 @@ var gpuRules = []gpuRule{
 // GPUInfo contains discovered GPU configuration from cluster nodes
 type GPUInfo struct {
 	NodeName         string                      // Name of the node with this GPU configuration
-	GPUsPerNode      int                         // Maximum GPUs per node found in the cluster
-	NodesWithGPUs    int                         // Number of nodes that have GPUs
+	GPUsPerNode      int                         // Maximum GPUs per node with the selected SKU
+	NodesWithGPUs    int                         // Number of discovered nodes with the selected SKU
+	TotalGPUs        int                         // Sum of GPUs across discovered nodes with the selected SKU, before any budget cap
 	Model            string                      // GPU product name (e.g., "H100-SXM5-80GB")
 	VRAMPerGPU       int                         // VRAM in MiB per GPU
 	System           nvidiacomv1beta1.GPUSKUType // AIC hardware system identifier (e.g., "h100_sxm", "h200_sxm"), empty if unknown
@@ -257,7 +258,8 @@ func (g *GPUDiscovery) DiscoverGPUsFromDCGM(ctx context.Context, k8sClient clien
 //
 // When filterSKU is non-empty, only nodes whose inferred SKU matches are
 // considered. When empty, the best node is selected first (highest GPU count,
-// then VRAM) and then only nodes with the same SKU are counted.
+// then VRAM). Each node with the selected SKU contributes once to NodesWithGPUs,
+// and its GPU count is summed into TotalGPUs.
 //
 // The function performs the following:
 //
@@ -269,16 +271,16 @@ func (g *GPUDiscovery) DiscoverGPUsFromDCGM(ctx context.Context, k8sClient clien
 //  6. Selects the "best" GPU node (filtered by SKU when set) based on:
 //     - Highest GPU count
 //     - Highest VRAM per GPU (tie-breaker)
-//  7. Counts only nodes matching the selected SKU for NodesWithGPUs.
-//  8. Caches the result per SKU for a short duration to avoid repeated scraping.
+//  7. Caches the result per SKU for a short duration to avoid repeated scraping.
 //
 // Behavior Notes:
 //
 //   - Scrapes pods directly instead of using a Service ClusterIP to avoid
 //     load-balancing ambiguity in multi-node clusters.
 //   - If at least one pod is successfully scraped, partial failures are tolerated.
 //   - If all pods fail to scrape, an aggregated error is returned.
-//   - Assumes DCGM exporter runs as a DaemonSet (one pod per GPU node).
+//   - Multiple exporters on one node use the report with the highest GPU count,
+//     then highest VRAM, for both node selection and aggregation.
 //
 // Returns:
 //   - *GPUInfo for the selected node
@@ -351,6 +353,7 @@ func (g *GPUDiscovery) discoverGPUsFromDCGMFilteredUncached(ctx context.Context,
 		nodeName string
 	}
 	allNodes := make([]nodeInfo, 0, len(dcgmPods))
+	nodeIndexes := make(map[string]int)
 	var scrapeErrors []error
 
 	for _, pod := range dcgmPods {
@@ -364,7 +367,18 @@ func (g *GPUDiscovery) discoverGPUsFromDCGMFilteredUncached(ctx context.Context,
 			continue
 		}
 
-		allNodes = append(allNodes, nodeInfo{info: info, sku: InferHardwareSystem(info.Model), nodeName: pod.Spec.NodeName})
+		// Running pods have a scheduled node name; keep its strongest report in first-seen node order.
+		n := nodeInfo{info: info, sku: InferHardwareSystem(info.Model), nodeName: pod.Spec.NodeName}
+		if index, seen := nodeIndexes[n.nodeName]; seen {
+			previous := allNodes[index].info
+			if info.GPUsPerNode > previous.GPUsPerNode ||
+				(info.GPUsPerNode == previous.GPUsPerNode && info.VRAMPerGPU > previous.VRAMPerGPU) {
+				allNodes[index] = n
+			}
+			continue
+		}
+		nodeIndexes[n.nodeName] = len(allNodes)
+		allNodes = append(allNodes, n)
 	}
 
 	if len(allNodes) == 0 {
@@ -397,18 +411,20 @@ func (g *GPUDiscovery) discoverGPUsFromDCGMFilteredUncached(ctx context.Context,
 		return nil, fmt.Errorf("no GPU metrics could be parsed from any DCGM pod")
 	}
 
-	// Count only nodes with the same SKU as the selected best node,
+	// Count nodes and GPUs with the same SKU as the selected best node,
 	// and detect RDMA on matching nodes only. On a cold cache and a no-RDMA
 	// cluster, this performs one Node read per matching node; that keeps a
 	// single negative node from masking RDMA on another node.
 	nodesWithGPUs := 0
+	totalGPUs := 0
 	var rdmaDetected bool
 	var rdmaType string
 	for _, n := range allNodes {
 		if n.sku != bestSKU {
 			continue
 		}
 		nodesWithGPUs++
+		totalGPUs += n.info.GPUsPerNode
 		if !rdmaDetected {
 			rdma, rType := detectRDMAFromNode(ctx, k8sClient, n.nodeName)
 			if rdma {
@@ -433,6 +449,7 @@ func (g *GPUDiscovery) discoverGPUsFromDCGMFilteredUncached(ctx context.Context,
 	bestNode.System = bestSKU
 	bestNode.CloudProvider = cloudProvider
 	bestNode.NodesWithGPUs = nodesWithGPUs
+	bestNode.TotalGPUs = totalGPUs
 	bestNode.RDMAEnabled = rdmaDetected
 	bestNode.RDMAType = rdmaType
 
@@ -854,16 +871,18 @@ func DiscoverGPUsFiltered(ctx context.Context, k8sClient client.Reader, filterSK
 			len(nodeList.Items), LabelGPUCount, LabelGPUProduct, LabelGPUMemory)
 	}
 
-	// Count only nodes with the same SKU as the selected best node,
+	// Count nodes and GPUs with the same SKU as the selected best node,
 	// and detect RDMA on matching nodes only. On a cold cache and a no-RDMA
 	// cluster, this performs one Node read per matching node; that keeps a
 	// single negative node from masking RDMA on another node.
 	nodesWithGPUs := 0
+	totalGPUs := 0
 	var rdmaDetected bool
 	var rdmaType string
 	for _, n := range allNodes {
 		if n.sku == bestSKU {
 			nodesWithGPUs++
+			totalGPUs += n.info.GPUsPerNode
 			if !rdmaDetected {
 				rdma, rType := detectRDMAFromNode(ctx, k8sClient, n.info.NodeName)
 				if rdma {
@@ -884,12 +903,13 @@ func DiscoverGPUsFiltered(ctx context.Context, k8sClient client.Reader, filterSK
 
 	bestNode.System = bestSKU
 	bestNode.NodesWithGPUs = nodesWithGPUs
+	bestNode.TotalGPUs = totalGPUs
 	bestNode.RDMAEnabled = rdmaDetected
 	bestNode.RDMAType = rdmaType
 	logger.Info("GPU discovery completed",
 		"gpusPerNode", bestNode.GPUsPerNode,
 		"nodesWithGPUs", bestNode.NodesWithGPUs,
-		"totalGpus", bestNode.GPUsPerNode*bestNode.NodesWithGPUs,
+		"totalGpus", bestNode.TotalGPUs,
 		"model", bestNode.Model,
 		"vram", bestNode.VRAMPerGPU,
 		"system", bestNode.System,