From 425ab64e933f2a57ddf3a23cf5ec830e77e95109 Mon Sep 17 00:00:00 2001
From: Sebastien Tardif <sebtardif@ncf.ca>
Date: Tue, 20 Oct 2026 10:35:05 -0700
Subject: [PATCH] Keep separate precompiled pools when kernel names match

getNodePools keyed each precompiled pool on the sanitized kernel.
That string drops x86_64, aarch64, and _64k, so an x86 node and an
arm node, or a 4k node and a 64k node, became one pool. The selector
and image tag then used only the first node's raw kernel. Key the
pool map on the raw kernel as well, and sort equal names by that
kernel. Nodes that share a raw kernel still share one pool.

Signed-off-by: Sebastien Tardif <sebtardif@ncf.ca>
---
 internal/state/nodepool.go      | 16 +++++++++--
 internal/state/nodepool_test.go | 51 +++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 3 deletions(-)

diff --git a/internal/state/nodepool.go b/internal/state/nodepool.go
index 21bde3f9d..ae48b34b6 100644
--- a/internal/state/nodepool.go
+++ b/internal/state/nodepool.go
@@ -139,9 +139,16 @@ func getNodePools(ctx context.Context, k8sClient client.Client, cr *nvidiav1alph
 			nodePool.name = rhcosVersion
 		}
 
-		if _, exists := nodePoolMap[nodePool.name]; !exists {
+		// Pool names drop arch and _64k. The node selector and the
+		// precompiled image tag keep the raw kernel, so x86_64 and
+		// aarch64 (or 4k and 64k) stay separate pools when those names match.
+		poolKey := nodePool.name
+		if nodePool.kernel != "" {
+			poolKey = nodePool.name + "/" + nodePool.kernel
+		}
+		if _, exists := nodePoolMap[poolKey]; !exists {
 			logger.Info("Detected new node pool", "NodePool", nodePool)
-			nodePoolMap[nodePool.name] = nodePool
+			nodePoolMap[poolKey] = nodePool
 		}
 	}
 
@@ -154,7 +161,10 @@ func sortedNodePools(nodePoolMap map[string]nodePool) []nodePool {
 		nodePools = append(nodePools, nodePool)
 	}
 	sort.Slice(nodePools, func(i, j int) bool {
-		return nodePools[i].name < nodePools[j].name
+		if nodePools[i].name != nodePools[j].name {
+			return nodePools[i].name < nodePools[j].name
+		}
+		return nodePools[i].kernel < nodePools[j].kernel
 	})
 	return nodePools
 }
diff --git a/internal/state/nodepool_test.go b/internal/state/nodepool_test.go
index e7e36955b..aea13853d 100644
--- a/internal/state/nodepool_test.go
+++ b/internal/state/nodepool_test.go
@@ -307,6 +307,57 @@ func TestGetNodePoolsPartitionsPrecompiledNodesByKernel(t *testing.T) {
 	require.Equal(t, "5.15.0-70-generic_x86_64", nodePools[0].nodeSelector[nfdKernelLabelKey])
 }
 
+func TestGetNodePoolsKeepsPrecompiledKernelsThatSanitizeTogether(t *testing.T) {
+	require.NoError(t, corev1.AddToScheme(scheme.Scheme))
+
+	// x86_64 and aarch64 sanitize to the same string. See TestGetSanitizedKernelVersion.
+	const (
+		x86Kernel = "5.14.0-427.37.1.el9_4.x86_64"
+		armKernel = "5.14.0-427.37.1.el9_4.aarch64"
+	)
+	node := func(name, kernel string) *corev1.Node {
+		return &corev1.Node{ObjectMeta: metav1.ObjectMeta{
+			Name: name,
+			Labels: map[string]string{
+				consts.GPUPresentLabel:        "true",
+				consts.NVIDIADriverOwnerLabel: "driver-a",
+				nfdOSReleaseIDLabelKey:        "rhel",
+				nfdOSVersionIDLabelKey:        "9.4",
+				nfdOSVersionIDMajorLabelKey:   "9",
+				nfdKernelLabelKey:             kernel,
+			},
+		}}
+	}
+	k8sClient := fake.NewClientBuilder().
+		WithScheme(scheme.Scheme).
+		WithObjects(
+			node("x86-node", x86Kernel),
+			node("arm-node", armKernel),
+			node("second-x86-node", x86Kernel),
+		).
+		Build()
+	driver := &nvidiav1alpha1.NVIDIADriver{
+		ObjectMeta: metav1.ObjectMeta{Name: "driver-a"},
+		Spec: nvidiav1alpha1.NVIDIADriverSpec{
+			UsePrecompiled: new(true),
+		},
+	}
+
+	nodePools, err := getNodePools(context.Background(), k8sClient, driver, false)
+
+	require.NoError(t, err)
+	require.Len(t, nodePools, 2)
+	got := map[string]string{}
+	for _, pool := range nodePools {
+		require.Equal(t, "rhel9-5.14.0-427.37.1.el9.4", pool.name)
+		got[pool.kernel] = pool.nodeSelector[nfdKernelLabelKey]
+	}
+	require.Equal(t, map[string]string{
+		x86Kernel: x86Kernel,
+		armKernel: armKernel,
+	}, got)
+}
+
 func TestGetNodePoolsPartitionsOpenShiftNodesByRHCOSVersion(t *testing.T) {
 	require.NoError(t, corev1.AddToScheme(scheme.Scheme))