Update device plugin stub with new GetPreferredAllocation() call

pkg/kubelet/cm/devicemanager/device_plugin_stub.go

@@ -37,6 +37,7 @@ type Stub struct {
 	socket                     string
 	resourceName               string
 	preStartContainerFlag      bool
+	getPreferredAllocationFlag bool
 
 	stop   chan interface{}
 	wg     sync.WaitGroup
@@ -47,12 +48,24 @@ type Stub struct {
 	// allocFunc is used for handling allocation request
 	allocFunc stubAllocFunc
 
+	// getPreferredAllocFunc is used for handling getPreferredAllocation request
+	getPreferredAllocFunc stubGetPreferredAllocFunc
+
 	registrationStatus chan watcherapi.RegistrationStatus // for testing
 	endpoint           string                             // for testing
 
 }
 
-// stubAllocFunc is the function called when receive an allocation request from Kubelet
+// stubGetPreferredAllocFunc is the function called when a getPreferredAllocation request is received from Kubelet
+type stubGetPreferredAllocFunc func(r *pluginapi.PreferredAllocationRequest, devs map[string]pluginapi.Device) (*pluginapi.PreferredAllocationResponse, error)
+
+func defaultGetPreferredAllocFunc(r *pluginapi.PreferredAllocationRequest, devs map[string]pluginapi.Device) (*pluginapi.PreferredAllocationResponse, error) {
+	var response pluginapi.PreferredAllocationResponse
+
+	return &response, nil
+}
+
+// stubAllocFunc is the function called when an allocation request is received from Kubelet
 type stubAllocFunc func(r *pluginapi.AllocateRequest, devs map[string]pluginapi.Device) (*pluginapi.AllocateResponse, error)
 
 func defaultAllocFunc(r *pluginapi.AllocateRequest, devs map[string]pluginapi.Device) (*pluginapi.AllocateResponse, error) {
@@ -62,20 +75,27 @@ func defaultAllocFunc(r *pluginapi.AllocateRequest, devs map[string]pluginapi.De
 }
 
 // NewDevicePluginStub returns an initialized DevicePlugin Stub.
-func NewDevicePluginStub(devs []*pluginapi.Device, socket string, name string, preStartContainerFlag bool) *Stub {
+func NewDevicePluginStub(devs []*pluginapi.Device, socket string, name string, preStartContainerFlag bool, getPreferredAllocationFlag bool) *Stub {
 	return &Stub{
 		devs:                       devs,
 		socket:                     socket,
 		resourceName:               name,
 		preStartContainerFlag:      preStartContainerFlag,
+		getPreferredAllocationFlag: getPreferredAllocationFlag,
 
 		stop:   make(chan interface{}),
 		update: make(chan []*pluginapi.Device),
 
 		allocFunc:             defaultAllocFunc,
+		getPreferredAllocFunc: defaultGetPreferredAllocFunc,
 	}
 }
 
+// SetGetPreferredAllocFunc sets allocFunc of the device plugin
+func (m *Stub) SetGetPreferredAllocFunc(f stubGetPreferredAllocFunc) {
+	m.getPreferredAllocFunc = f
+}
+
 // SetAllocFunc sets allocFunc of the device plugin
 func (m *Stub) SetAllocFunc(f stubAllocFunc) {
 	m.allocFunc = f
@@ -174,7 +194,10 @@ func (m *Stub) Register(kubeletEndpoint, resourceName string, pluginSockDir stri
 		Version:      pluginapi.Version,
 		Endpoint:     path.Base(m.socket),
 		ResourceName: resourceName,
-		Options:      &pluginapi.DevicePluginOptions{PreStartRequired: m.preStartContainerFlag},
+		Options: &pluginapi.DevicePluginOptions{
+			PreStartRequired:                m.preStartContainerFlag,
+			GetPreferredAllocationAvailable: m.getPreferredAllocationFlag,
+		},
 	}
 
 	_, err = client.Register(context.Background(), reqt)
@@ -186,7 +209,11 @@ func (m *Stub) Register(kubeletEndpoint, resourceName string, pluginSockDir stri
 
 // GetDevicePluginOptions returns DevicePluginOptions settings for the device plugin.
 func (m *Stub) GetDevicePluginOptions(ctx context.Context, e *pluginapi.Empty) (*pluginapi.DevicePluginOptions, error) {
-	return &pluginapi.DevicePluginOptions{PreStartRequired: m.preStartContainerFlag}, nil
+	options := &pluginapi.DevicePluginOptions{
+		PreStartRequired:                m.preStartContainerFlag,
+		GetPreferredAllocationAvailable: m.getPreferredAllocationFlag,
+	}
+	return options, nil
 }
 
 // PreStartContainer resets the devices received
@@ -216,6 +243,19 @@ func (m *Stub) Update(devs []*pluginapi.Device) {
 	m.update <- devs
 }
 
+// GetPreferredAllocation gets the preferred allocation from a set of available devices
+func (m *Stub) GetPreferredAllocation(ctx context.Context, r *pluginapi.PreferredAllocationRequest) (*pluginapi.PreferredAllocationResponse, error) {
+	klog.Infof("GetPreferredAllocation, %+v", r)
+
+	devs := make(map[string]pluginapi.Device)
+
+	for _, dev := range m.devs {
+		devs[dev.ID] = *dev
+	}
+
+	return m.getPreferredAllocFunc(r, devs)
+}
+
 // Allocate does a mock allocation
 func (m *Stub) Allocate(ctx context.Context, r *pluginapi.AllocateRequest) (*pluginapi.AllocateResponse, error) {
 	klog.Infof("Allocate, %+v", r)