Ensure consistent dnsname --> IMEXDaemonIP mapping on all nodes

Without this the IMEX daemons were getting confused if the same DNS name
was used to by different nodes to point to differernt IMEX daemons in
the ensemble.

Signed-off-by: Kevin Klues <[email protected]>

Author: klueska

api/nvidia.com/resource/v1beta1/computedomain.go

@@ -91,4 +91,5 @@ type ComputeDomainNode struct {
 	Name      string `json:"name"`
 	IPAddress string `json:"ipAddress"`
 	CliqueID  string `json:"cliqueID"`
+	Index     int    `json:"index"`
 }

cmd/compute-domain-daemon/computedomain.go

@@ -218,9 +218,16 @@ func (m *ComputeDomainManager) UpdateComputeDomainNodeInfo(ctx context.Context,
 
 	// If there isn't one, create one and append it to the list
 	if nodeInfo == nil {
+		// Get the next available index for this new node
+		nextIndex, err := getNextAvailableIndex(newCD.Status.Nodes, m.config.maxNodesPerIMEXDomain)
+		if err != nil {
+			return fmt.Errorf("error getting next available index: %w", err)
+		}
+
 		nodeInfo = &nvapi.ComputeDomainNode{
 			Name:     m.config.nodeName,
 			CliqueID: m.config.cliqueID,
+			Index:    nextIndex,
 		}
 		newCD.Status.Nodes = append(newCD.Status.Nodes, nodeInfo)
 	}
@@ -243,6 +250,46 @@ func (m *ComputeDomainManager) UpdateComputeDomainNodeInfo(ctx context.Context,
 	return nil
 }
 
+// The Index field in the Nodes section of the ComputeDomain status ensures a
+// consistent IP-to-DNS name mapping across all machines within a given IMEX
+// domain. Each node's index directly determines its DNS name using the format
+// "compute-domain-daemon-{index}".
+//
+// getNextAvailableIndex finds the next available index for the current node by
+// seeing which ones are already taken by other nodes in the ComputeDomain
+// status. It fills in gaps where it can, and returns an error if no index is
+// available within maxNodesPerIMEXDomain.
+//
+// By filling gaps in the index sequence (rather than always appending), we
+// maintain stable DNS names for existing nodes even when intermediate nodes
+// are removed from the compute domain and new ones are added.
+func getNextAvailableIndex(nodes []*nvapi.ComputeDomainNode, maxNodesPerIMEXDomain int) (int, error) {
+	if len(nodes) >= maxNodesPerIMEXDomain {
+		return -1, fmt.Errorf("cannot add more nodes, already at maximum (%d)", maxNodesPerIMEXDomain)
+	}
+
+	// Create a map to track used indices
+	usedIndices := make(map[int]bool)
+
+	// Collect all currently used indices
+	for _, node := range nodes {
+		usedIndices[node.Index] = true
+	}
+
+	// Find the next available index, starting from 0 and filling gaps
+	nextIndex := 0
+	for usedIndices[nextIndex] {
+		nextIndex++
+	}
+
+	// Ensure nextIndex is within the range 0..maxNodesPerIMEXDomain
+	if nextIndex < 0 || nextIndex >= maxNodesPerIMEXDomain {
+		return -1, fmt.Errorf("no available indices within maxNodesPerIMEXDomain (%d)", maxNodesPerIMEXDomain)
+	}
+
+	return nextIndex, nil
+}
+
 // If we've reached the expected number of nodes and if there was actually a
 // change compared to the previously known set of nodes: pass info to IMEX
 // daemon controller.

cmd/compute-domain-daemon/controller.go

@@ -35,6 +35,7 @@ type ManagerConfig struct {
 	computeDomainNamespace string
 	cliqueID               string
 	podIP                  string
+	maxNodesPerIMEXDomain  int
 }
 
 // ControllerConfig holds the configuration for the controller.
@@ -45,6 +46,7 @@ type ControllerConfig struct {
 	computeDomainNamespace string
 	cliqueID               string
 	podIP                  string
+	maxNodesPerIMEXDomain  int
 }
 
 // Controller manages the lifecycle of compute domain operations.
@@ -73,6 +75,7 @@ func NewController(config *ControllerConfig) (*Controller, error) {
 		computeDomainNamespace: config.computeDomainNamespace,
 		cliqueID:               config.cliqueID,
 		podIP:                  config.podIP,
+		maxNodesPerIMEXDomain:  config.maxNodesPerIMEXDomain,
 	}
 
 	controller := &Controller{

cmd/compute-domain-daemon/dnsnames.go

@@ -83,14 +83,15 @@ func (m *DNSNameManager) UpdateDNSNameMappings(nodes []*nvapi.ComputeDomainNode)
 		}
 	}
 
-	// Add new IPs to map (filling in holes where others were removed)
+	// Add new IPs to map
 	for _, node := range cliqueNodes {
 		// If IP already has a DNS name, skip it
 		if _, exists := ipToDNSName[node.IPAddress]; exists {
 			continue
 		}
 
-		dnsName, err := m.allocateDNSName(node.IPAddress)
+		// Construct the DNS name from the node index
+		dnsName, err := m.constructDNSName(node)
 		if err != nil {
 			return fmt.Errorf("failed to allocate DNS name for IP %s: %w", node.IPAddress, err)
 		}
@@ -127,32 +128,17 @@ func (m *DNSNameManager) LogDNSNameMappings() {
 	}
 }
 
-// allocateDNSName allocates a DNS name for an IP address, reusing existing DNS names if possible.
-func (m *DNSNameManager) allocateDNSName(ip string) (string, error) {
-	// If IP already has a DNS name, return it
-	if dnsName, exists := m.ipToDNSName[ip]; exists {
-		return dnsName, nil
+// contructDNSName constructs a DNS name for a node based on its index field.
+// Returns an error if the index is invalid or exceeds maxNodesPerIMEXDomain.
+func (m *DNSNameManager) constructDNSName(node *nvapi.ComputeDomainNode) (string, error) {
+	if node.Index < 0 {
+		return "", fmt.Errorf("node %s has invalid index %d", node.Name, node.Index)
 	}
-
-	// Find the next available DNS name
-	for i := 0; i < m.maxNodesPerIMEXDomain; i++ {
-		dnsName := fmt.Sprintf(dnsNameFormat, i)
-		// Check if this DNS name is already in use
-		inUse := false
-		for _, existingDNSName := range m.ipToDNSName {
-			if existingDNSName == dnsName {
-				inUse = true
-				break
-			}
-		}
-		if !inUse {
-			m.ipToDNSName[ip] = dnsName
-			return dnsName, nil
-		}
+	if node.Index >= m.maxNodesPerIMEXDomain {
+		return "", fmt.Errorf("node %s has invalid index %d, must be less than %d", node.Name, node.Index, m.maxNodesPerIMEXDomain)
 	}
-
-	// If all DNS names are used, return an error
-	return "", fmt.Errorf("no DNS names available (max: %d)", m.maxNodesPerIMEXDomain)
+	dnsName := fmt.Sprintf(dnsNameFormat, node.Index)
+	return dnsName, nil
 }
 
 // updateHostsFile updates the /etc/hosts file with current IP to DNS name mappings.

cmd/compute-domain-daemon/main.go

@@ -182,6 +182,7 @@ func run(ctx context.Context, cancel context.CancelFunc, flags *Flags) error {
 		computeDomainNamespace: flags.computeDomainNamespace,
 		nodeName:               flags.nodeName,
 		podIP:                  flags.podIP,
+		maxNodesPerIMEXDomain:  flags.maxNodesPerIMEXDomain,
 	}
 	klog.Infof("config: %v", config)
 

deployments/helm/nvidia-dra-driver-gpu/crds/resource.nvidia.com_computedomains.yaml

@@ -85,12 +85,15 @@ spec:
                   properties:
                     cliqueID:
                       type: string
+                    index:
+                      type: integer
                     ipAddress:
                       type: string
                     name:
                       type: string
                   required:
                   - cliqueID
+                  - index
                   - ipAddress
                   - name
                   type: object