Add AKS-owned labels to Flex Nodes

pkg/config/adapter.go

@@ -20,6 +20,13 @@ const (
 
 	// aksAADServerID is the Azure AD server application ID for AKS.
 	aksAADServerID = "6dae42f8-4368-4678-94ff-3960e28e3630"
+
+	managedNodeLabel      = "kubernetes.azure.com/managed"
+	agentPoolNodeLabel    = "kubernetes.azure.com/agentpool"
+	modeNodeLabel         = "kubernetes.azure.com/mode"
+	nodePoolTypeNodeLabel = "kubernetes.azure.com/nodepool-type"
+	flexNodePoolType      = "FlexNodes"
+	userNodeMode          = "user"
 )
 
 // ToAgentConfig converts a FlexNode Config to the shared agent library's
@@ -44,7 +51,7 @@ func ToAgentConfig(cfg *Config, machineName string) *agentconfig.AgentConfig {
 		Kubelet: agentconfig.AgentKubeletConfig{
 			ApiServer:          cfg.APIServerURL(),
 			NodeIP:             cfg.Node.Kubelet.NodeIP,
-			Labels:             cfg.Node.Labels,
+			Labels:             kubeletNodeLabels(cfg),
 			RegisterWithTaints: cfg.Node.Taints,
 			Configuration:      kubeletConfig,
 		},
@@ -162,6 +169,20 @@ func kubeReservedOrDefault(cfg *Config, maxPods int) map[string]string {
 	return defaultKubeReserved(runtime.NumCPU(), hostTotalMemoryMi(), maxPods)
 }
 
+func kubeletNodeLabels(cfg *Config) map[string]string {
+	labels := maps.Clone(cfg.Node.Labels)
+	if labels == nil {
+		labels = make(map[string]string)
+	}
+
+	// Keep AKS-owned labels out of the custom labels sent in the Machine goal.
+	labels[managedNodeLabel] = "false"
+	labels[agentPoolNodeLabel] = cfg.Azure.TargetAgentPoolName
+	labels[modeNodeLabel] = userNodeMode
+	labels[nodePoolTypeNodeLabel] = flexNodePoolType
+	return labels
+}
+
 // ResolveMachineGoalState converts FlexNode config to the shared agent config
 // and resolves the nspawn machine goal state. Bootstrap and preflight both use
 // this helper so preflight validates the same sources that bootstrap consumes.

pkg/config/adapter_test.go

@@ -10,12 +10,59 @@ import (
 	"github.com/Azure/unbounded/pkg/agent/goalstates"
 )
 
+func TestToAgentConfigKubeletLabels(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name   string
+		labels map[string]string
+	}{
+		{
+			name:   "custom labels",
+			labels: map[string]string{"workload": "edge"},
+		},
+		{
+			name: "no custom labels",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			customLabels := maps.Clone(tt.labels)
+			cfg := &Config{
+				Azure: AzureConfig{TargetAgentPoolName: "flexnode-edge"},
+				Node:  NodeConfig{Labels: customLabels},
+			}
+			wantLabels := maps.Clone(tt.labels)
+			if wantLabels == nil {
+				wantLabels = make(map[string]string)
+			}
+			wantLabels[managedNodeLabel] = "false"
+			wantLabels[agentPoolNodeLabel] = "flexnode-edge"
+			wantLabels[modeNodeLabel] = userNodeMode
+			wantLabels[nodePoolTypeNodeLabel] = flexNodePoolType
+
+			agentCfg := ToAgentConfig(cfg, "kube1")
+
+			if !maps.Equal(agentCfg.Kubelet.Labels, wantLabels) {
+				t.Errorf("Kubelet.Labels = %#v, want %#v", agentCfg.Kubelet.Labels, wantLabels)
+			}
+			if !maps.Equal(cfg.Node.Labels, tt.labels) {
+				t.Errorf("Node.Labels mutated to %#v; Machine custom labels must remain %#v", cfg.Node.Labels, tt.labels)
+			}
+		})
+	}
+}
+
 func TestToAgentConfig_BootstrapToken(t *testing.T) {
 	t.Parallel()
 
 	cfg := &Config{
 		Azure: AzureConfig{
-			BootstrapToken: &BootstrapTokenConfig{Token: "abcdef.0123456789abcdef"},
+			BootstrapToken:      &BootstrapTokenConfig{Token: "abcdef.0123456789abcdef"},
+			TargetAgentPoolName: "flexnode-edge",
 		},
 		Components: ComponentsConfig{Kubernetes: "1.30.0"},
 		Networking: NetworkingConfig{DNSServiceIP: "10.0.0.10"},
@@ -68,8 +115,15 @@ func TestToAgentConfig_BootstrapToken(t *testing.T) {
 	if ac.Kubelet.Auth.ExecCredential != nil {
 		t.Fatalf("Kubelet.Auth.ExecCredential should be nil for bootstrap token auth")
 	}
-	if len(ac.Kubelet.Labels) != 1 || ac.Kubelet.Labels["env"] != "test" {
-		t.Fatalf("Kubelet.Labels=%v, want map[env:test]", ac.Kubelet.Labels)
+	wantLabels := map[string]string{
+		"env":                 "test",
+		managedNodeLabel:      "false",
+		agentPoolNodeLabel:    "flexnode-edge",
+		modeNodeLabel:         userNodeMode,
+		nodePoolTypeNodeLabel: flexNodePoolType,
+	}
+	if !maps.Equal(ac.Kubelet.Labels, wantLabels) {
+		t.Fatalf("Kubelet.Labels=%v, want %v", ac.Kubelet.Labels, wantLabels)
 	}
 	if len(ac.Kubelet.RegisterWithTaints) != 1 || ac.Kubelet.RegisterWithTaints[0] != "dedicated=infra:NoSchedule" {
 		t.Fatalf("Kubelet.RegisterWithTaints=%v, want [dedicated=infra:NoSchedule]", ac.Kubelet.RegisterWithTaints)

pkg/config/config.go

@@ -456,9 +456,6 @@ func (c *Config) setNodeDefaults() {
 	if c.Node.Labels == nil {
 		c.Node.Labels = make(map[string]string)
 	}
-	// Mark node as unmanaged by cloud controller manager by default, otherwise ccm will delete this node if node is not ready
-	// doc: https://cloud-provider-azure.sigs.k8s.io/topics/cross-resource-group-nodes/#unmanaged-nodes
-	c.Node.Labels["kubernetes.azure.com/managed"] = "false"
 
 	// Set default kubelet configuration if not provided
 	if c.Node.Kubelet.Verbosity == 0 {
@@ -932,6 +929,11 @@ func (c *NodeConfig) validate() error {
 	if c.MaxPods < 0 || int64(c.MaxPods) > math.MaxInt32 {
 		return fmt.Errorf("node.maxPods must be between 0 and %d, inclusive", math.MaxInt32)
 	}
+	for _, label := range []string{managedNodeLabel, agentPoolNodeLabel, modeNodeLabel, nodePoolTypeNodeLabel} {
+		if _, configured := c.Labels[label]; configured {
+			return fmt.Errorf("node.labels cannot configure AKS-owned label %q", label)
+		}
+	}
 	return c.Kubelet.validate()
 }
 

pkg/config/kubelet_config_test.go

@@ -188,3 +188,45 @@ func TestNodeConfigValidateMaxPods(t *testing.T) {
 		})
 	}
 }
+
+func TestNodeConfigValidateAKSOwnedLabels(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name    string
+		labels  map[string]string
+		wantErr bool
+	}{
+		{name: "custom label", labels: map[string]string{"workload": "edge"}},
+		{name: "managed", labels: map[string]string{managedNodeLabel: "false"}, wantErr: true},
+		{name: "agent pool", labels: map[string]string{agentPoolNodeLabel: "pool"}, wantErr: true},
+		{name: "mode", labels: map[string]string{modeNodeLabel: "user"}, wantErr: true},
+		{name: "node pool type", labels: map[string]string{nodePoolTypeNodeLabel: "FlexNodes"}, wantErr: true},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			cfg := NodeConfig{
+				MaxPods: 110,
+				Labels:  tt.labels,
+				Kubelet: KubeletConfig{
+					Verbosity:            2,
+					ImageGCHighThreshold: 85,
+					ImageGCLowThreshold:  80,
+				},
+			}
+			err := cfg.validate()
+			if tt.wantErr {
+				if err == nil || !strings.Contains(err.Error(), "AKS-owned label") {
+					t.Fatalf("validate error = %v, want AKS-owned label error", err)
+				}
+				return
+			}
+			if err != nil {
+				t.Fatalf("validate: %v", err)
+			}
+		})
+	}
+}