Skip to content

Commit 32c047a

Browse files
committed
Update device plugin stub with new GetPreferredAllocation() call
1 parent 202c4f0 commit 32c047a

File tree

1 file changed

+53
-13
lines changed

1 file changed

+53
-13
lines changed

pkg/kubelet/cm/devicemanager/device_plugin_stub.go

Lines changed: 53 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -33,10 +33,11 @@ import (
3333

3434
// Stub implementation for DevicePlugin.
3535
type Stub struct {
36-
devs []*pluginapi.Device
37-
socket string
38-
resourceName string
39-
preStartContainerFlag bool
36+
devs []*pluginapi.Device
37+
socket string
38+
resourceName string
39+
preStartContainerFlag bool
40+
getPreferredAllocationFlag bool
4041

4142
stop chan interface{}
4243
wg sync.WaitGroup
@@ -47,12 +48,24 @@ type Stub struct {
4748
// allocFunc is used for handling allocation request
4849
allocFunc stubAllocFunc
4950

51+
// getPreferredAllocFunc is used for handling getPreferredAllocation request
52+
getPreferredAllocFunc stubGetPreferredAllocFunc
53+
5054
registrationStatus chan watcherapi.RegistrationStatus // for testing
5155
endpoint string // for testing
5256

5357
}
5458

55-
// stubAllocFunc is the function called when receive an allocation request from Kubelet
59+
// stubGetPreferredAllocFunc is the function called when a getPreferredAllocation request is received from Kubelet
60+
type stubGetPreferredAllocFunc func(r *pluginapi.PreferredAllocationRequest, devs map[string]pluginapi.Device) (*pluginapi.PreferredAllocationResponse, error)
61+
62+
func defaultGetPreferredAllocFunc(r *pluginapi.PreferredAllocationRequest, devs map[string]pluginapi.Device) (*pluginapi.PreferredAllocationResponse, error) {
63+
var response pluginapi.PreferredAllocationResponse
64+
65+
return &response, nil
66+
}
67+
68+
// stubAllocFunc is the function called when an allocation request is received from Kubelet
5669
type stubAllocFunc func(r *pluginapi.AllocateRequest, devs map[string]pluginapi.Device) (*pluginapi.AllocateResponse, error)
5770

5871
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
6275
}
6376

6477
// NewDevicePluginStub returns an initialized DevicePlugin Stub.
65-
func NewDevicePluginStub(devs []*pluginapi.Device, socket string, name string, preStartContainerFlag bool) *Stub {
78+
func NewDevicePluginStub(devs []*pluginapi.Device, socket string, name string, preStartContainerFlag bool, getPreferredAllocationFlag bool) *Stub {
6679
return &Stub{
67-
devs: devs,
68-
socket: socket,
69-
resourceName: name,
70-
preStartContainerFlag: preStartContainerFlag,
80+
devs: devs,
81+
socket: socket,
82+
resourceName: name,
83+
preStartContainerFlag: preStartContainerFlag,
84+
getPreferredAllocationFlag: getPreferredAllocationFlag,
7185

7286
stop: make(chan interface{}),
7387
update: make(chan []*pluginapi.Device),
7488

75-
allocFunc: defaultAllocFunc,
89+
allocFunc: defaultAllocFunc,
90+
getPreferredAllocFunc: defaultGetPreferredAllocFunc,
7691
}
7792
}
7893

94+
// SetGetPreferredAllocFunc sets allocFunc of the device plugin
95+
func (m *Stub) SetGetPreferredAllocFunc(f stubGetPreferredAllocFunc) {
96+
m.getPreferredAllocFunc = f
97+
}
98+
7999
// SetAllocFunc sets allocFunc of the device plugin
80100
func (m *Stub) SetAllocFunc(f stubAllocFunc) {
81101
m.allocFunc = f
@@ -174,7 +194,10 @@ func (m *Stub) Register(kubeletEndpoint, resourceName string, pluginSockDir stri
174194
Version: pluginapi.Version,
175195
Endpoint: path.Base(m.socket),
176196
ResourceName: resourceName,
177-
Options: &pluginapi.DevicePluginOptions{PreStartRequired: m.preStartContainerFlag},
197+
Options: &pluginapi.DevicePluginOptions{
198+
PreStartRequired: m.preStartContainerFlag,
199+
GetPreferredAllocationAvailable: m.getPreferredAllocationFlag,
200+
},
178201
}
179202

180203
_, err = client.Register(context.Background(), reqt)
@@ -186,7 +209,11 @@ func (m *Stub) Register(kubeletEndpoint, resourceName string, pluginSockDir stri
186209

187210
// GetDevicePluginOptions returns DevicePluginOptions settings for the device plugin.
188211
func (m *Stub) GetDevicePluginOptions(ctx context.Context, e *pluginapi.Empty) (*pluginapi.DevicePluginOptions, error) {
189-
return &pluginapi.DevicePluginOptions{PreStartRequired: m.preStartContainerFlag}, nil
212+
options := &pluginapi.DevicePluginOptions{
213+
PreStartRequired: m.preStartContainerFlag,
214+
GetPreferredAllocationAvailable: m.getPreferredAllocationFlag,
215+
}
216+
return options, nil
190217
}
191218

192219
// PreStartContainer resets the devices received
@@ -216,6 +243,19 @@ func (m *Stub) Update(devs []*pluginapi.Device) {
216243
m.update <- devs
217244
}
218245

246+
// GetPreferredAllocation gets the preferred allocation from a set of available devices
247+
func (m *Stub) GetPreferredAllocation(ctx context.Context, r *pluginapi.PreferredAllocationRequest) (*pluginapi.PreferredAllocationResponse, error) {
248+
klog.Infof("GetPreferredAllocation, %+v", r)
249+
250+
devs := make(map[string]pluginapi.Device)
251+
252+
for _, dev := range m.devs {
253+
devs[dev.ID] = *dev
254+
}
255+
256+
return m.getPreferredAllocFunc(r, devs)
257+
}
258+
219259
// Allocate does a mock allocation
220260
func (m *Stub) Allocate(ctx context.Context, r *pluginapi.AllocateRequest) (*pluginapi.AllocateResponse, error) {
221261
klog.Infof("Allocate, %+v", r)

0 commit comments

Comments
 (0)