add num-devices parameter

Author: bart0sh

cmd/dra-example-kubeletplugin/discovery.go

@@ -28,8 +28,7 @@ import (
 	"github.com/google/uuid"
 )
 
-func enumerateAllPossibleDevices() (AllocatableDevices, error) {
-	numGPUs := 8
+func enumerateAllPossibleDevices(numGPUs int) (AllocatableDevices, error) {
 	seed := os.Getenv("NODE_NAME")
 	uuids := generateUUIDs(seed, numGPUs)
 

cmd/dra-example-kubeletplugin/main.go

@@ -44,8 +44,9 @@ type Flags struct {
 	kubeClientConfig flags.KubeClientConfig
 	loggingConfig    *flags.LoggingConfig
 
-	nodeName string
-	cdiRoot  string
+	nodeName   string
+	cdiRoot    string
+	numDevices int
 }
 
 type Config struct {
@@ -79,6 +80,13 @@ func newApp() *cli.App {
 			Destination: &flags.cdiRoot,
 			EnvVars:     []string{"CDI_ROOT"},
 		},
+		&cli.IntFlag{
+			Name:        "num-devices",
+			Usage:       "The number of devices to be generated.",
+			Value:       8,
+			Destination: &flags.numDevices,
+			EnvVars:     []string{"NUM_DEVICES"},
+		},
 	}
 	cliFlags = append(cliFlags, flags.kubeClientConfig.Flags()...)
 	cliFlags = append(cliFlags, flags.loggingConfig.Flags()...)

cmd/dra-example-kubeletplugin/state.go

@@ -63,7 +63,7 @@ type DeviceState struct {
 }
 
 func NewDeviceState(config *Config) (*DeviceState, error) {
-	allocatable, err := enumerateAllPossibleDevices()
+	allocatable, err := enumerateAllPossibleDevices(config.flags.numDevices)
 	if err != nil {
 		return nil, fmt.Errorf("error enumerating all possible devices: %v", err)
 	}

deployments/helm/dra-example-driver/templates/kubeletplugin.yaml

@@ -53,6 +53,8 @@ spec:
           valueFrom:
             fieldRef:
               fieldPath: metadata.namespace
+        - name: NUM_DEVICES
+          value: "8"
         volumeMounts:
         - name: plugins-registry
           mountPath: /var/lib/kubelet/plugins_registry