Make maxNodesPerIMEXDomain configurable (default at 18)

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

Author: klueska

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/dnsnames.go

@@ -30,7 +30,6 @@ import (
 )
 
 const (
-	maxDNSNames   = 18
 	hostsFilePath = "/etc/hosts"
 	dnsNameFormat = "compute-domain-daemon-%d"
 )
@@ -41,17 +40,19 @@ type IPToDNSNameMap map[string]string
 // DNSNameManager manages the allocation of static DNS names to IP addresses.
 type DNSNameManager struct {
 	sync.Mutex
-	ipToDNSName     IPToDNSNameMap
-	cliqueID        string
-	nodesConfigPath string
+	ipToDNSName           IPToDNSNameMap
+	cliqueID              string
+	maxNodesPerIMEXDomain int
+	nodesConfigPath       string
 }
 
 // NewDNSNameManager creates a new DNS name manager.
-func NewDNSNameManager(cliqueID string, nodesConfigPath string) *DNSNameManager {
+func NewDNSNameManager(cliqueID string, maxNodesPerIMEXDomain int, nodesConfigPath string) *DNSNameManager {
 	return &DNSNameManager{
-		ipToDNSName:     make(IPToDNSNameMap),
-		cliqueID:        cliqueID,
-		nodesConfigPath: nodesConfigPath,
+		ipToDNSName:           make(IPToDNSNameMap),
+		cliqueID:              cliqueID,
+		maxNodesPerIMEXDomain: maxNodesPerIMEXDomain,
+		nodesConfigPath:       nodesConfigPath,
 	}
 }
 
@@ -134,7 +135,7 @@ func (m *DNSNameManager) allocateDNSName(ip string) (string, error) {
 	}
 
 	// Find the next available DNS name
-	for i := 0; i < maxDNSNames; i++ {
+	for i := 0; i < m.maxNodesPerIMEXDomain; i++ {
 		dnsName := fmt.Sprintf(dnsNameFormat, i)
 		// Check if this DNS name is already in use
 		inUse := false
@@ -151,7 +152,7 @@ func (m *DNSNameManager) allocateDNSName(ip string) (string, error) {
 	}
 
 	// If all DNS names are used, return an error
-	return "", fmt.Errorf("no DNS names available (max: %d)", maxDNSNames)
+	return "", fmt.Errorf("no DNS names available (max: %d)", m.maxNodesPerIMEXDomain)
 }
 
 // updateHostsFile updates the /etc/hosts file with current IP to DNS name mappings.
@@ -221,14 +222,14 @@ func (m *DNSNameManager) WriteNodesConfig() error {
 	defer f.Close()
 
 	// Write static DNS names
-	for i := 0; i < maxDNSNames; i++ {
+	for i := 0; i < m.maxNodesPerIMEXDomain; i++ {
 		dnsName := fmt.Sprintf(dnsNameFormat, i)
 		if _, err := fmt.Fprintf(f, "%s\n", dnsName); err != nil {
 			return fmt.Errorf("failed to write to nodes config file: %w", err)
 		}
 	}
 
-	klog.Infof("Created static nodes config file with %d DNS names using format %s", maxDNSNames, dnsNameFormat)
+	klog.Infof("Created static nodes config file with %d DNS names using format %s", m.maxNodesPerIMEXDomain, dnsNameFormat)
 
 	return nil
 }

cmd/compute-domain-daemon/main.go

@@ -51,6 +51,7 @@ type Flags struct {
 	computeDomainNamespace string
 	nodeName               string
 	podIP                  string
+	maxNodesPerIMEXDomain  int
 	loggingConfig          *flags.LoggingConfig
 	featureGateConfig      *flags.FeatureGateConfig
 }
@@ -124,6 +125,12 @@ func newApp() *cli.App {
 			EnvVars:     []string{"POD_IP"},
 			Destination: &flags.podIP,
 		},
+		&cli.IntFlag{
+			Name:        "max-nodes-per-imex-domain",
+			Usage:       "The maximum number of possible nodes per IMEX domain",
+			EnvVars:     []string{"MAX_NODES_PER_IMEX_DOMAIN"},
+			Destination: &flags.maxNodesPerIMEXDomain,
+		},
 	}
 	cliFlags = append(cliFlags, flags.featureGateConfig.Flags()...)
 	cliFlags = append(cliFlags, flags.loggingConfig.Flags()...)
@@ -181,7 +188,7 @@ func run(ctx context.Context, cancel context.CancelFunc, flags *Flags) error {
 	var dnsNameManager *DNSNameManager
 	if featuregates.Enabled(featuregates.IMEXDaemonsWithDNSNames) {
 		// Prepare DNS name manager
-		dnsNameManager = NewDNSNameManager(flags.cliqueID, nodesConfigPath)
+		dnsNameManager = NewDNSNameManager(flags.cliqueID, flags.maxNodesPerIMEXDomain, nodesConfigPath)
 
 		// Create static nodes config file with DNS names
 		if err := dnsNameManager.WriteNodesConfig(); err != nil {

templates/compute-domain-daemon.tmpl.yaml

@@ -26,6 +26,8 @@ spec:
         image: {{ .ImageName }}
         command: ["compute-domain-daemon", "-v", "6", "run"]
         env:
+        - name: MAX_NODES_PER_IMEX_DOMAIN
+          value: "{{ .MaxNodesPerIMEXDomain }}"
         - name: NODE_NAME
           valueFrom:
             fieldRef: