Merge pull request #65 from bart0sh/PR003-add-num-devices-parameter

add num-devices parameter
kubernetes-sigs · Sep 30, 2024 · 2713efd · 2713efd
2 parents 303c1d2 + 009b41f
commit 2713efd
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 5 deletions.
diff --git a/cmd/dra-example-kubeletplugin/discovery.go b/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)
 

diff --git a/cmd/dra-example-kubeletplugin/main.go b/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()...)

diff --git a/cmd/dra-example-kubeletplugin/state.go b/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)
 	}

diff --git a/deployments/helm/dra-example-driver/templates/kubeletplugin.yaml b/deployments/helm/dra-example-driver/templates/kubeletplugin.yaml
@@ -53,6 +53,9 @@ spec:
           valueFrom:
             fieldRef:
               fieldPath: metadata.namespace
+        # Simulated number of devices the example driver will pretend to have.
+        - name: NUM_DEVICES
+          value: "8"
         volumeMounts:
         - name: plugins-registry
           mountPath: /var/lib/kubelet/plugins_registry