From ce655c0bb01ed2b9d70d7a4044bcaf346ac9fec0 Mon Sep 17 00:00:00 2001
From: gshaibi <39763067+gshaibi@users.noreply.github.com>
Date: Thu, 28 Nov 2024 11:19:30 +0200
Subject: [PATCH] RUN-20790 Don't run device plugin DaemonSet on fake nodes
 (#102)

---
 internal/status-updater/handlers/node/common.go        | 10 ++++++++++
 .../handlers/node/fake_node_deployments.go             |  4 ----
 internal/status-updater/handlers/node/labels.go        | 10 +++++++---
 3 files changed, 17 insertions(+), 7 deletions(-)
 create mode 100644 internal/status-updater/handlers/node/common.go

diff --git a/internal/status-updater/handlers/node/common.go b/internal/status-updater/handlers/node/common.go
new file mode 100644
index 0000000..0949783
--- /dev/null
+++ b/internal/status-updater/handlers/node/common.go
@@ -0,0 +1,10 @@
+package node
+
+import (
+	"github.com/run-ai/fake-gpu-operator/internal/common/constants"
+	v1 "k8s.io/api/core/v1"
+)
+
+func isFakeNode(node *v1.Node) bool {
+	return node != nil && node.Annotations[constants.AnnotationKwokNode] == "fake"
+}
diff --git a/internal/status-updater/handlers/node/fake_node_deployments.go b/internal/status-updater/handlers/node/fake_node_deployments.go
index f75ab9e..5d73313 100644
--- a/internal/status-updater/handlers/node/fake_node_deployments.go
+++ b/internal/status-updater/handlers/node/fake_node_deployments.go
@@ -119,7 +119,3 @@ func generateFakeNodeDeploymentFromTemplate(template *appsv1.Deployment, node *v
 
 	return deployment
 }
-
-func isFakeNode(node *v1.Node) bool {
-	return node != nil && node.Annotations[constants.AnnotationKwokNode] == "fake"
-}
diff --git a/internal/status-updater/handlers/node/labels.go b/internal/status-updater/handlers/node/labels.go
index 9b297eb..d148ccb 100644
--- a/internal/status-updater/handlers/node/labels.go
+++ b/internal/status-updater/handlers/node/labels.go
@@ -18,10 +18,14 @@ const (
 
 // labelNode labels the node with required labels for the fake-gpu-operator to function.
 func (p *NodeHandler) labelNode(node *v1.Node) error {
-	err := p.patchNodeLabels(node, map[string]interface{}{
+	labels := map[string]interface{}{
 		dcgmExporterLabelKey: "true",
-		devicePluginLabelKey: "true",
-	})
+	}
+	if !isFakeNode(node) {
+		labels[devicePluginLabelKey] = "true"
+	}
+
+	err := p.patchNodeLabels(node, labels)
 	if err != nil {
 		return fmt.Errorf("failed to label node %s: %w", node.Name, err)
 	}