Ensure each imex domain.cliqueId has a unique set of channel numbers

Signed-off-by: Carlos Eduardo Arango Gutierrez <[email protected]>

Author: ArangoGutierrez

cmd/nvidia-dra-controller/imex.go

@@ -19,6 +19,7 @@ package main
 import (
 	"context"
 	"fmt"
+	"strings"
 	"sync"
 	"time"
 
@@ -36,17 +37,20 @@ import (
 )
 
 const (
-	DriverName       = "gpu.nvidia.com"
-	ImexDomainLabel  = "nvidia.com/gpu.imex-domain"
-	ImexChannelLimit = 128
+	DriverName                    = "gpu.nvidia.com"
+	ImexDomainLabel               = "nvidia.com/gpu.imex-domain"
+	ResourceSliceImexChannelLimit = 128
+	DriverImexChannelLimit        = 2048
 )
 
 type ImexManager struct {
 	waitGroup sync.WaitGroup
 	clientset kubernetes.Interface
 }
 
-type DriverResources resourceslice.DriverResources
+// imexDomainOffsets represents the offset for assigning IMEX channels
+// to ResourceSlices for each <imex-domain, cliqueid> combination.
+type imexDomainOffsets map[string]map[string]int
 
 func StartIMEXManager(ctx context.Context, config *Config) (*ImexManager, error) {
 	// Build a client set config
@@ -99,29 +103,34 @@ func StartIMEXManager(ctx context.Context, config *Config) (*ImexManager, error)
 
 // manageResourceSlices reacts to added and removed IMEX domains and triggers the creation / removal of resource slices accordingly.
 func (m *ImexManager) manageResourceSlices(ctx context.Context, owner resourceslice.Owner, addedDomainsCh <-chan string, removedDomainsCh <-chan string) error {
-	driverResources := resourceslice.DriverResources{}
-	controller, err := resourceslice.StartController(ctx, m.clientset, DriverName, owner, &driverResources)
+	driverResources := &resourceslice.DriverResources{}
+	controller, err := resourceslice.StartController(ctx, m.clientset, DriverName, owner, driverResources)
 	if err != nil {
 		return fmt.Errorf("error starting resource slice controller: %w", err)
 	}
 
+	imexDomainOffsets := new(imexDomainOffsets)
 	m.waitGroup.Add(1)
 	go func() {
 		defer m.waitGroup.Done()
 		for {
 			select {
 			case addedDomain := <-addedDomainsCh:
+				offset, err := imexDomainOffsets.add(addedDomain, ResourceSliceImexChannelLimit, DriverImexChannelLimit)
+				if err != nil {
+					klog.Errorf("Error calculating channel offset for IMEX domain %s: %v", addedDomain, err)
+					return
+				}
 				klog.Infof("Adding channels for new IMEX domain: %v", addedDomain)
-				newDriverResources := DriverResources(driverResources).DeepCopy()
-				newDriverResources.Pools[addedDomain] = generateImexChannelPool(addedDomain, ImexChannelLimit)
-				controller.Update(&newDriverResources)
-				driverResources = newDriverResources
+				driverResources := driverResources.DeepCopy()
+				driverResources.Pools[addedDomain] = generateImexChannelPool(addedDomain, offset, ResourceSliceImexChannelLimit)
+				controller.Update(driverResources)
 			case removedDomain := <-removedDomainsCh:
 				klog.Infof("Removing channels for removed IMEX domain: %v", removedDomain)
-				newDriverResources := DriverResources(driverResources).DeepCopy()
-				delete(newDriverResources.Pools, removedDomain)
-				controller.Update(&newDriverResources)
-				driverResources = newDriverResources
+				driverResources := driverResources.DeepCopy()
+				delete(driverResources.Pools, removedDomain)
+				imexDomainOffsets.remove(removedDomain)
+				controller.Update(driverResources)
 			case <-ctx.Done():
 				return
 			}
@@ -146,17 +155,6 @@ func (m *ImexManager) Stop() error {
 	return nil
 }
 
-// DeepCopy will perform a deep copy of the provided DriverResources.
-func (d DriverResources) DeepCopy() resourceslice.DriverResources {
-	driverResources := resourceslice.DriverResources{
-		Pools: make(map[string]resourceslice.Pool),
-	}
-	for p := range d.Pools {
-		driverResources.Pools[p] = generateImexChannelPool(p, ImexChannelLimit)
-	}
-	return driverResources
-}
-
 // streamImexDomains returns two channels that streams imexDomans that are added and removed from nodes over time.
 func (m *ImexManager) streamImexDomains(ctx context.Context) (<-chan string, <-chan string, error) {
 	// Create channels to stream IMEX domain ids that are added / removed
@@ -249,10 +247,10 @@ func (m *ImexManager) streamImexDomains(ctx context.Context) (<-chan string, <-c
 }
 
 // generateImexChannelPool generates the contents of a ResourceSlice pool for a given IMEX domain.
-func generateImexChannelPool(imexDomain string, numChannels int) resourceslice.Pool {
-	// Generate dchannels from 0 to numChannels
+func generateImexChannelPool(imexDomain string, startChannel int, numChannels int) resourceslice.Pool {
+	// Generate channels from startChannel to offset+numChannels
 	var devices []resourceapi.Device
-	for i := 0; i < numChannels; i++ {
+	for i := startChannel; i < (startChannel + numChannels); i++ {
 		d := resourceapi.Device{
 			Name: fmt.Sprintf("imex-channel-%d", i),
 			Basic: &resourceapi.BasicDevice{
@@ -312,3 +310,60 @@ func (m *ImexManager) cleanupResourceSlices() error {
 
 	return nil
 }
+
+// add sets the offset where an IMEX domain's channels should start counting from.
+func (offsets imexDomainOffsets) add(imexDomain string, resourceSliceImexChannelLimit, driverImexChannelLimit int) (int, error) {
+	// Split the incoming imexDomain to split off its cliqueID
+	id := strings.SplitN(imexDomain, ".", 2)
+	if len(id) != 2 {
+		return -1, fmt.Errorf("error adding IMEX domain %s: invalid format", imexDomain)
+	}
+	imexDomain = id[0]
+	cliqueID := id[1]
+
+	// Check if the IMEX domain is already in the map
+	if _, ok := offsets[imexDomain]; !ok {
+		offsets[imexDomain] = make(map[string]int)
+	}
+
+	// Return early if the clique is already in the map
+	if offset, exists := offsets[imexDomain][cliqueID]; exists {
+		return offset, nil
+	}
+
+	// Track used offsets for the current imexDomain
+	usedOffsets := make(map[int]struct{})
+	for _, v := range offsets[imexDomain] {
+		usedOffsets[v] = struct{}{}
+	}
+
+	// Look for the first unused offset, stepping by resourceSliceImexChannelLimit
+	var offset int
+	for offset = 0; offset < driverImexChannelLimit; offset += resourceSliceImexChannelLimit {
+		if _, exists := usedOffsets[offset]; !exists {
+			break
+		}
+	}
+
+	// If we reach the limit, return an error
+	if offset == driverImexChannelLimit {
+		return -1, fmt.Errorf("error adding IMEX domain %s: channel limit reached", imexDomain)
+	}
+	offsets[imexDomain][cliqueID] = offset
+
+	return offset, nil
+}
+
+func (offsets imexDomainOffsets) remove(imexDomain string) {
+	id := strings.SplitN(imexDomain, ".", 2)
+	if len(id) != 2 {
+		return
+	}
+	imexDomain = id[0]
+	cliqueID := id[1]
+
+	delete(offsets[imexDomain], cliqueID)
+	if len(offsets[imexDomain]) == 0 {
+		delete(offsets, imexDomain)
+	}
+}