Merge pull request NVIDIA#186 from ArangoGutierrez/107/131

Clean up resourceSlices during exit

Author: ArangoGutierrez

Makefile

@@ -90,7 +90,14 @@ vendor:
 	go mod vendor
 	go mod verify
 
-check-vendor: vendor
+# TODO: Remove once the upstream change is released
+vendor-patch:
+	@echo "Applying patch to "k8s.io/dynamic-resource-allocation/resourceslice""
+	git apply -p3 --directory=vendor/ patches/resourceslice.patch || { \
+		echo "Failed to apply patch"; exit 1; \
+	}
+
+check-vendor: vendor vendor-patch
 	git diff --quiet HEAD -- go.mod go.sum vendor
 
 COVERAGE_FILE := coverage.out

cmd/nvidia-dra-controller/imex.go

@@ -19,6 +19,8 @@ package main
 import (
 	"context"
 	"fmt"
+	"strings"
+	"sync"
 	"time"
 
 	v1 "k8s.io/api/core/v1"
@@ -35,30 +37,34 @@ import (
 )
 
 const (
-	DriverName       = "gpu.nvidia.com"
-	ImexDomainLabel  = "nvidia.com/gpu.imex-domain"
-	ImexChannelLimit = 128
+	DriverName         = "gpu.nvidia.com"
+	ImexDomainLabel    = "nvidia.com/gpu.imex-domain"
+	ImexChannelLimit   = 128
+	DriverChannelLimit = 2048
 )
 
-type DriverResources resourceslice.DriverResources
+type ImexManager struct {
+	waitGroup sync.WaitGroup
+	clientset kubernetes.Interface
+}
 
-func StartIMEXManager(ctx context.Context, config *Config) error {
+func StartIMEXManager(ctx context.Context, config *Config) (*ImexManager, error) {
 	// Build a client set config
 	csconfig, err := config.flags.kubeClientConfig.NewClientSetConfig()
 	if err != nil {
-		return fmt.Errorf("error creating client set config: %w", err)
+		return nil, fmt.Errorf("error creating client set config: %w", err)
 	}
 
 	// Create a new clientset
 	clientset, err := kubernetes.NewForConfig(csconfig)
 	if err != nil {
-		return fmt.Errorf("error creating dynamic client: %w", err)
+		return nil, fmt.Errorf("error creating dynamic client: %w", err)
 	}
 
 	// Fetch the current Pod object
 	pod, err := clientset.CoreV1().Pods(config.flags.namespace).Get(ctx, config.flags.podName, metav1.GetOptions{})
 	if err != nil {
-		return fmt.Errorf("error fetching pod: %w", err)
+		return nil, fmt.Errorf("error fetching pod: %w", err)
 	}
 
 	// Set the owner of the ResourceSlices we will create
@@ -69,46 +75,61 @@ func StartIMEXManager(ctx context.Context, config *Config) error {
 		UID:        pod.UID,
 	}
 
+	// Create the manager itself
+	m := &ImexManager{
+		clientset: clientset,
+	}
+
 	// Stream added/removed IMEX domains from nodes over time
 	klog.Info("Start streaming IMEX domains from nodes...")
-	addedDomainsCh, removedDomainsCh, err := streamImexDomains(ctx, clientset)
+	addedDomainsCh, removedDomainsCh, err := m.streamImexDomains(ctx)
 	if err != nil {
-		return fmt.Errorf("error streaming IMEX domains: %w", err)
+		return nil, fmt.Errorf("error streaming IMEX domains: %w", err)
 	}
 
 	// Add/Remove resource slices from IMEX domains as they come and go
 	klog.Info("Start publishing IMEX channels to ResourceSlices...")
-	err = manageResourceSlices(ctx, clientset, owner, addedDomainsCh, removedDomainsCh)
+	err = m.manageResourceSlices(ctx, owner, addedDomainsCh, removedDomainsCh)
 	if err != nil {
-		return fmt.Errorf("error managing resource slices: %w", err)
+		return nil, fmt.Errorf("error managing resource slices: %w", err)
 	}
 
-	return nil
+	return m, nil
 }
 
 // manageResourceSlices reacts to added and removed IMEX domains and triggers the creation / removal of resource slices accordingly.
-func manageResourceSlices(ctx context.Context, clientset kubernetes.Interface, owner resourceslice.Owner, addedDomainsCh <-chan string, removedDomainsCh <-chan string) error {
-	driverResources := resourceslice.DriverResources{}
-	controller, err := resourceslice.StartController(ctx, clientset, DriverName, owner, &driverResources)
+func (m *ImexManager) manageResourceSlices(ctx context.Context, owner resourceslice.Owner, addedDomainsCh <-chan string, removedDomainsCh <-chan string) error {
+	driverResources := resourceslice.DriverResources{
+		Pools: make(map[string]resourceslice.Pool),
+	}
+	controller, err := resourceslice.StartController(ctx, m.clientset, DriverName, owner, &driverResources)
 	if err != nil {
 		return fmt.Errorf("error starting resource slice controller: %w", err)
 	}
 
+	imexChannelOffsets := make(map[string]map[string]int)
+
+	m.waitGroup.Add(1)
 	go func() {
+		defer m.waitGroup.Done()
 		for {
 			select {
 			case addedDomain := <-addedDomainsCh:
 				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
+				if err := calculateImexChannelLimit(addedDomain, &imexChannelOffsets); err != nil {
+					klog.Error(err)
+					return
+				}
+				newDriverResources := driverResources.DeepCopy()
+				newDriverResources.Pools[addedDomain] = generateImexChannelPool(addedDomain, imexChannelOffsets)
+				controller.Update(newDriverResources)
+				driverResources = *newDriverResources
 			case removedDomain := <-removedDomainsCh:
 				klog.Infof("Removing channels for removed IMEX domain: %v", removedDomain)
-				newDriverResources := DriverResources(driverResources).DeepCopy()
+				newDriverResources := driverResources.DeepCopy()
 				delete(newDriverResources.Pools, removedDomain)
-				controller.Update(&newDriverResources)
-				driverResources = newDriverResources
+				controller.Update(newDriverResources)
+				driverResources = *newDriverResources
 			case <-ctx.Done():
 				return
 			}
@@ -118,19 +139,23 @@ func manageResourceSlices(ctx context.Context, clientset kubernetes.Interface, o
 	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),
+// Stop stops a running ImexManager.
+func (m *ImexManager) Stop() error {
+	if m == nil {
+		return nil
 	}
-	for p := range d.Pools {
-		driverResources.Pools[p] = generateImexChannelPool(p, ImexChannelLimit)
+
+	m.waitGroup.Wait()
+	klog.Info("Cleaning up all resourceSlices")
+	if err := m.cleanupResourceSlices(); err != nil {
+		return fmt.Errorf("error cleaning up resource slices: %w", err)
 	}
-	return driverResources
+
+	return nil
 }
 
 // streamImexDomains returns two channels that streams imexDomans that are added and removed from nodes over time.
-func streamImexDomains(ctx context.Context, clientset kubernetes.Interface) (<-chan string, <-chan string, error) {
+func (m *ImexManager) streamImexDomains(ctx context.Context) (<-chan string, <-chan string, error) {
 	// Create channels to stream IMEX domain ids that are added / removed
 	addedDomainCh := make(chan string)
 	removedDomainCh := make(chan string)
@@ -147,7 +172,7 @@ func streamImexDomains(ctx context.Context, clientset kubernetes.Interface) (<-c
 
 	// Create a shared informer factory for nodes
 	informerFactory := informers.NewSharedInformerFactoryWithOptions(
-		clientset,
+		m.clientset,
 		time.Minute*10, // Resync period
 		informers.WithTweakListOptions(func(options *metav1.ListOptions) {
 			options.LabelSelector = labelSelector
@@ -206,7 +231,11 @@ func streamImexDomains(ctx context.Context, clientset kubernetes.Interface) (<-c
 	}
 
 	// Start the informer and wait for it to sync
-	go informerFactory.Start(ctx.Done())
+	m.waitGroup.Add(1)
+	go func() {
+		defer m.waitGroup.Done()
+		informerFactory.Start(ctx.Done())
+	}()
 
 	// Wait for the informer caches to sync
 	if !cache.WaitForCacheSync(ctx.Done(), nodeInformer.HasSynced) {
@@ -217,10 +246,13 @@ func streamImexDomains(ctx context.Context, clientset kubernetes.Interface) (<-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, imexChannelOffsets map[string]map[string]int) resourceslice.Pool {
+	id := strings.Split(imexDomain, ".")
+	numChannels := imexChannelOffsets[id[0]][id[1]] + ImexChannelLimit
+
+	// Generate dchannels from offset to offset+ImexChannelLimit
 	var devices []resourceapi.Device
-	for i := 0; i < numChannels; i++ {
+	for i := imexChannelOffsets[id[0]][id[1]]; i < numChannels; i++ {
 		d := resourceapi.Device{
 			Name: fmt.Sprintf("imex-channel-%d", i),
 			Basic: &resourceapi.BasicDevice{
@@ -259,3 +291,46 @@ func generateImexChannelPool(imexDomain string, numChannels int) resourceslice.P
 
 	return pool
 }
+
+// cleanupResourceSlices removes all resource slices created by the IMEX manager.
+func (m *ImexManager) cleanupResourceSlices() error {
+	// Delete all resource slices created by the IMEX manager
+	ops := metav1.ListOptions{
+		FieldSelector: fmt.Sprintf("%s=%s", resourceapi.ResourceSliceSelectorDriver, DriverName),
+	}
+	l, err := m.clientset.ResourceV1alpha3().ResourceSlices().List(context.Background(), ops)
+	if err != nil {
+		return fmt.Errorf("error listing resource slices: %w", err)
+	}
+
+	for _, rs := range l.Items {
+		err := m.clientset.ResourceV1alpha3().ResourceSlices().Delete(context.Background(), rs.Name, metav1.DeleteOptions{})
+		if err != nil {
+			return fmt.Errorf("error deleting resource slice %s: %w", rs.Name, err)
+		}
+	}
+
+	return nil
+}
+
+// calculateImexChannelLimit calculates the number of IMEX channels that can be allocated to a given IMEX domain.
+// if the limit is reached, it will return an error.
+func calculateImexChannelLimit(addedDomainsCh string, imexChannelOffsets *map[string]map[string]int) error {
+	id := strings.Split(addedDomainsCh, ".")
+
+	// Check if the IMEX domain is already in the map
+	if _, ok := (*imexChannelOffsets)[id[0]]; !ok {
+		(*imexChannelOffsets)[id[0]] = make(map[string]int)
+	}
+
+	// Check if the clique is already in the map
+	if _, ok := (*imexChannelOffsets)[id[0]][id[1]]; !ok {
+		offset := (len((*imexChannelOffsets)[id[0]]) * ImexChannelLimit)
+		if offset >= DriverChannelLimit {
+			return fmt.Errorf("error adding IMEX domain %s: channel limit reached", id[0])
+		}
+		(*imexChannelOffsets)[id[0]][id[1]] = offset
+	}
+
+	return nil
+}

cmd/nvidia-dra-controller/main.go

@@ -17,12 +17,15 @@
 package main
 
 import (
+	"context"
 	"fmt"
 	"net"
 	"net/http"
 	"net/http/pprof"
 	"os"
+	"os/signal"
 	"path"
+	"syscall"
 
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
@@ -132,7 +135,6 @@ func newApp() *cli.App {
 			return flags.loggingConfig.Apply()
 		},
 		Action: func(c *cli.Context) error {
-			ctx := c.Context
 			mux := http.NewServeMux()
 			flags.deviceClasses = sets.New[string](c.StringSlice("device-classes")...)
 
@@ -154,14 +156,26 @@ func newApp() *cli.App {
 				}
 			}
 
+			sigs := make(chan os.Signal, 1)
+			signal.Notify(sigs, syscall.SIGTERM, syscall.SIGINT)
+
+			var imexManager *ImexManager
+			ctx, cancel := context.WithCancel(c.Context)
+			defer func() {
+				cancel()
+				if err := imexManager.Stop(); err != nil {
+					klog.Errorf("Error stopping IMEX manager: %v", err)
+				}
+			}()
+
 			if flags.deviceClasses.Has(ImexChannelType) {
-				err = StartIMEXManager(ctx, config)
+				imexManager, err = StartIMEXManager(ctx, config)
 				if err != nil {
 					return fmt.Errorf("start IMEX manager: %w", err)
 				}
 			}
 
-			<-ctx.Done()
+			<-sigs
 
 			return nil
 		},