diff --git a/pkg/kubelet/cm/devicemanager/device_plugin_stub.go b/pkg/kubelet/cm/devicemanager/device_plugin_stub.go index fb5ec5c0c42..bf9e96a9c1b 100644 --- a/pkg/kubelet/cm/devicemanager/device_plugin_stub.go +++ b/pkg/kubelet/cm/devicemanager/device_plugin_stub.go @@ -33,10 +33,11 @@ import ( // Stub implementation for DevicePlugin. type Stub struct { - devs []*pluginapi.Device - socket string - resourceName string - preStartContainerFlag bool + devs []*pluginapi.Device + 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, + devs: devs, + socket: socket, + resourceName: name, + preStartContainerFlag: preStartContainerFlag, + getPreferredAllocationFlag: getPreferredAllocationFlag, stop: make(chan interface{}), update: make(chan []*pluginapi.Device), - allocFunc: defaultAllocFunc, + 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)