Merge pull request #433 from klueska/imex-with-hostnames

Add support for using DNSNames instead of raw IPs for IMEX daemons

Author: klueska

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

@@ -91,4 +91,11 @@ type ComputeDomainNode struct {
 	Name      string `json:"name"`
 	IPAddress string `json:"ipAddress"`
 	CliqueID  string `json:"cliqueID"`
+	// The Index field is used to ensure a consistent IP-to-DNS name
+	// mapping across all machines within an IMEX domain. Each node's index
+	// directly determines its DNS name. It is marked as optional (but not
+	// omitempty) in order to support downgrades and avoid an API bump.
+	// +optional
+	// +kubebuilder:validation:Optional
+	Index int `json:"index"`
 }

cmd/compute-domain-controller/controller.go

@@ -38,6 +38,9 @@ type ManagerConfig struct {
 	// imageName is the full image name to use when rendering templates
 	imageName string
 
+	// maxNodesPerIMEXDomain is the maximum number of nodes per IMEX domain to allocate
+	maxNodesPerIMEXDomain int
+
 	// clientsets provides access to various Kubernetes API client interfaces
 	clientsets flags.ClientSets
 
@@ -67,12 +70,13 @@ func (c *Controller) Run(ctx context.Context) error {
 	workQueue := workqueue.New(workqueue.DefaultControllerRateLimiter())
 
 	managerConfig := &ManagerConfig{
-		driverName:           c.config.driverName,
-		driverNamespace:      c.config.flags.namespace,
-		additionalNamespaces: c.config.flags.additionalNamespaces.Value(),
-		imageName:            c.config.flags.imageName,
-		clientsets:           c.config.clientsets,
-		workQueue:            workQueue,
+		driverName:            c.config.driverName,
+		driverNamespace:       c.config.flags.namespace,
+		additionalNamespaces:  c.config.flags.additionalNamespaces.Value(),
+		imageName:             c.config.flags.imageName,
+		maxNodesPerIMEXDomain: c.config.flags.maxNodesPerIMEXDomain,
+		clientsets:            c.config.clientsets,
+		workQueue:             workQueue,
 	}
 
 	cdManager := NewComputeDomainManager(managerConfig)

cmd/compute-domain-controller/daemonset.go

@@ -50,6 +50,7 @@ type DaemonSetTemplateData struct {
 	ComputeDomainLabelValue   types.UID
 	ResourceClaimTemplateName string
 	ImageName                 string
+	MaxNodesPerIMEXDomain     int
 	FeatureGates              map[string]bool
 }
 
@@ -200,6 +201,7 @@ func (m *DaemonSetManager) Create(ctx context.Context, cd *nvapi.ComputeDomain)
 		ComputeDomainLabelValue:   cd.UID,
 		ResourceClaimTemplateName: rct.Name,
 		ImageName:                 m.config.imageName,
+		MaxNodesPerIMEXDomain:     m.config.maxNodesPerIMEXDomain,
 		FeatureGates:              featuregates.ToMap(),
 	}
 

cmd/compute-domain-controller/main.go

@@ -44,16 +44,23 @@ import (
 
 const (
 	DriverName = "compute-domain.nvidia.com"
+
+	// This constant provides a reasonable default for the maximum size of
+	// a given IMEX Domain. On GB200 and GB300 the limit is 18, so we pick
+	// this for now. It can be overridden as an environment variable or
+	// command line argument as required.
+	defaultMaxNodesPerIMEXDomain = 18
 )
 
 type Flags struct {
 	kubeClientConfig  flags.KubeClientConfig
 	loggingConfig     *flags.LoggingConfig
 	featureGateConfig *flags.FeatureGateConfig
 
-	podName   string
-	namespace string
-	imageName string
+	podName               string
+	namespace             string
+	imageName             string
+	maxNodesPerIMEXDomain int
 
 	httpEndpoint string
 	metricsPath  string
@@ -103,6 +110,13 @@ func newApp() *cli.App {
 			Destination: &flags.imageName,
 			EnvVars:     []string{"IMAGE_NAME"},
 		},
+		&cli.IntFlag{
+			Name:        "max-nodes-per-imex-domain",
+			Usage:       "The maximum number of possible nodes per IMEX domain",
+			Value:       defaultMaxNodesPerIMEXDomain,
+			EnvVars:     []string{"MAX_NODES_PER_IMEX_DOMAIN"},
+			Destination: &flags.maxNodesPerIMEXDomain,
+		},
 		&cli.StringFlag{
 			Category:    "HTTP server:",
 			Name:        "http-endpoint",

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{