|
| 1 | +// /* |
| 2 | +// Copyright 2025 The Grove Authors. |
| 3 | +// |
| 4 | +// Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | +// you may not use this file except in compliance with the License. |
| 6 | +// You may obtain a copy of the License at |
| 7 | +// |
| 8 | +// http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | +// |
| 10 | +// Unless required by applicable law or agreed to in writing, software |
| 11 | +// distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | +// See the License for the specific language governing permissions and |
| 14 | +// limitations under the License. |
| 15 | +// */ |
| 16 | + |
| 17 | +package setup |
| 18 | + |
| 19 | +import ( |
| 20 | + "context" |
| 21 | + "fmt" |
| 22 | + "sort" |
| 23 | + |
| 24 | + "github.com/ai-dynamo/grove/operator/e2e/utils" |
| 25 | + v1 "k8s.io/api/core/v1" |
| 26 | + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" |
| 27 | + k8stypes "k8s.io/apimachinery/pkg/types" |
| 28 | + "k8s.io/client-go/kubernetes" |
| 29 | + "k8s.io/client-go/rest" |
| 30 | +) |
| 31 | + |
| 32 | +const ( |
| 33 | + // WorkerNodeLabelKey is the label key used to identify worker nodes in e2e tests. |
| 34 | + // This can be changed if infrastructure changes. |
| 35 | + WorkerNodeLabelKey = "node_role.e2e.grove.nvidia.com" |
| 36 | + // WorkerNodeLabelValue is the label value for worker node identification in e2e tests. |
| 37 | + WorkerNodeLabelValue = "agent" |
| 38 | + |
| 39 | + // TopologyLabelZone is the Kubernetes label key for zone topology domain. |
| 40 | + TopologyLabelZone = "kubernetes.io/zone" |
| 41 | + // TopologyLabelBlock is the Kubernetes label key for the block topology domain. |
| 42 | + TopologyLabelBlock = "kubernetes.io/block" |
| 43 | + // TopologyLabelRack is the Kubernetes label key for the rack topology domain. |
| 44 | + TopologyLabelRack = "kubernetes.io/rack" |
| 45 | + // TopologyLabelHostname is the Kubernetes label key for the hostname topology domain. |
| 46 | + TopologyLabelHostname = "kubernetes.io/hostname" |
| 47 | + |
| 48 | + // NodesPerZone is the number of nodes per zone. |
| 49 | + NodesPerZone = 28 |
| 50 | + // NodesPerBlock is the number of nodes per block (28 / 2 blocks). |
| 51 | + NodesPerBlock = 14 |
| 52 | + // NodesPerRack is the number of nodes per rack (28 / 4 racks). |
| 53 | + NodesPerRack = 7 |
| 54 | +) |
| 55 | + |
| 56 | +// GetZoneForNodeIndex returns the zone label for a given node index. |
| 57 | +// Both the index parameter and the returned zone number are 0-based. |
| 58 | +// e.g., nodes 0-27 → zone-0, nodes 28-55 → zone-1, etc. |
| 59 | +func GetZoneForNodeIndex(idx int) string { |
| 60 | + zoneNum := idx / NodesPerZone |
| 61 | + return fmt.Sprintf("zone-%d", zoneNum) |
| 62 | +} |
| 63 | + |
| 64 | +// GetBlockForNodeIndex returns the block label for a given node index. |
| 65 | +// Both the index parameter and the returned block number are 0-based. |
| 66 | +// e.g., nodes 0-13 → block-0, nodes 14-27 → block-1 |
| 67 | +func GetBlockForNodeIndex(idx int) string { |
| 68 | + blockNum := idx / NodesPerBlock |
| 69 | + return fmt.Sprintf("block-%d", blockNum) |
| 70 | +} |
| 71 | + |
| 72 | +// GetRackForNodeIndex returns the rack label for a given node index. |
| 73 | +// Both the index parameter and the returned rack number are 0-based. |
| 74 | +// e.g., nodes 0-6 → rack-0, nodes 7-13 → rack-1, etc. |
| 75 | +func GetRackForNodeIndex(idx int) string { |
| 76 | + rackNum := idx / NodesPerRack |
| 77 | + return fmt.Sprintf("rack-%d", rackNum) |
| 78 | +} |
| 79 | + |
| 80 | +// GetWorkerNodeLabelSelector returns the label selector for worker nodes in e2e tests. |
| 81 | +// Returns a formatted string "key=value" for use with Kubernetes label selectors. |
| 82 | +func GetWorkerNodeLabelSelector() string { |
| 83 | + return fmt.Sprintf("%s=%s", WorkerNodeLabelKey, WorkerNodeLabelValue) |
| 84 | +} |
| 85 | + |
| 86 | +// applyTopologyLabels applies hierarchical topology labels to worker nodes in the k3d cluster. |
| 87 | +// Creates a 4-level topology hierarchy: zone -> block -> rack -> host (kubernetes.io/hostname already exists) |
| 88 | +// Distribution strategy for 28 worker nodes: |
| 89 | +// - Zone: all nodes in "zone-0" |
| 90 | +// - Block: nodes 0-13 in "block-0", nodes 14-27 in "block-1" |
| 91 | +// - Rack: 4 racks total (2 per block), 7 hosts per rack |
| 92 | +func applyTopologyLabels(ctx context.Context, restConfig *rest.Config, logger *utils.Logger) error { |
| 93 | + logger.Info("🏷️ Applying hierarchical topology labels to worker nodes...") |
| 94 | + |
| 95 | + // Create clientset |
| 96 | + clientset, err := kubernetes.NewForConfig(restConfig) |
| 97 | + if err != nil { |
| 98 | + return fmt.Errorf("failed to create clientset: %w", err) |
| 99 | + } |
| 100 | + |
| 101 | + // Get all worker nodes (filter by label set during cluster creation) |
| 102 | + workerLabelSelector := GetWorkerNodeLabelSelector() |
| 103 | + nodes, err := clientset.CoreV1().Nodes().List(ctx, metav1.ListOptions{ |
| 104 | + LabelSelector: workerLabelSelector, |
| 105 | + }) |
| 106 | + if err != nil { |
| 107 | + return fmt.Errorf("failed to list worker nodes: %w", err) |
| 108 | + } |
| 109 | + |
| 110 | + if len(nodes.Items) == 0 { |
| 111 | + logger.Warn("⚠️ No worker nodes found for topology labeling") |
| 112 | + return nil |
| 113 | + } |
| 114 | + |
| 115 | + sortedNodes := make([]v1.Node, len(nodes.Items)) |
| 116 | + copy(sortedNodes, nodes.Items) |
| 117 | + sort.Slice(sortedNodes, func(i, j int) bool { return sortedNodes[i].Name < sortedNodes[j].Name }) |
| 118 | + |
| 119 | + for idx, node := range sortedNodes { |
| 120 | + topologyLabels := fmt.Sprintf(`{"metadata":{"labels":{"%s":"%s","%s":"%s","%s":"%s"}}}`, |
| 121 | + TopologyLabelZone, GetZoneForNodeIndex(idx), TopologyLabelBlock, GetBlockForNodeIndex(idx), TopologyLabelRack, GetRackForNodeIndex(idx)) |
| 122 | + |
| 123 | + _, err := clientset.CoreV1().Nodes().Patch( |
| 124 | + ctx, |
| 125 | + node.Name, |
| 126 | + k8stypes.StrategicMergePatchType, |
| 127 | + []byte(topologyLabels), |
| 128 | + metav1.PatchOptions{}, |
| 129 | + ) |
| 130 | + if err != nil { |
| 131 | + return fmt.Errorf("failed to patch node %s with topology labels: %w", node.Name, err) |
| 132 | + } |
| 133 | + } |
| 134 | + logger.Infof("✅ Applied topology labels to %d worker nodes", len(sortedNodes)) |
| 135 | + return nil |
| 136 | +} |
0 commit comments