@@ -33,10 +33,11 @@ import (
33
33
34
34
// Stub implementation for DevicePlugin.
35
35
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
40
41
41
42
stop chan interface {}
42
43
wg sync.WaitGroup
@@ -47,12 +48,24 @@ type Stub struct {
47
48
// allocFunc is used for handling allocation request
48
49
allocFunc stubAllocFunc
49
50
51
+ // getPreferredAllocFunc is used for handling getPreferredAllocation request
52
+ getPreferredAllocFunc stubGetPreferredAllocFunc
53
+
50
54
registrationStatus chan watcherapi.RegistrationStatus // for testing
51
55
endpoint string // for testing
52
56
53
57
}
54
58
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
56
69
type stubAllocFunc func (r * pluginapi.AllocateRequest , devs map [string ]pluginapi.Device ) (* pluginapi.AllocateResponse , error )
57
70
58
71
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
62
75
}
63
76
64
77
// 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 {
66
79
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 ,
71
85
72
86
stop : make (chan interface {}),
73
87
update : make (chan []* pluginapi.Device ),
74
88
75
- allocFunc : defaultAllocFunc ,
89
+ allocFunc : defaultAllocFunc ,
90
+ getPreferredAllocFunc : defaultGetPreferredAllocFunc ,
76
91
}
77
92
}
78
93
94
+ // SetGetPreferredAllocFunc sets allocFunc of the device plugin
95
+ func (m * Stub ) SetGetPreferredAllocFunc (f stubGetPreferredAllocFunc ) {
96
+ m .getPreferredAllocFunc = f
97
+ }
98
+
79
99
// SetAllocFunc sets allocFunc of the device plugin
80
100
func (m * Stub ) SetAllocFunc (f stubAllocFunc ) {
81
101
m .allocFunc = f
@@ -174,7 +194,10 @@ func (m *Stub) Register(kubeletEndpoint, resourceName string, pluginSockDir stri
174
194
Version : pluginapi .Version ,
175
195
Endpoint : path .Base (m .socket ),
176
196
ResourceName : resourceName ,
177
- Options : & pluginapi.DevicePluginOptions {PreStartRequired : m .preStartContainerFlag },
197
+ Options : & pluginapi.DevicePluginOptions {
198
+ PreStartRequired : m .preStartContainerFlag ,
199
+ GetPreferredAllocationAvailable : m .getPreferredAllocationFlag ,
200
+ },
178
201
}
179
202
180
203
_ , err = client .Register (context .Background (), reqt )
@@ -186,7 +209,11 @@ func (m *Stub) Register(kubeletEndpoint, resourceName string, pluginSockDir stri
186
209
187
210
// GetDevicePluginOptions returns DevicePluginOptions settings for the device plugin.
188
211
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
190
217
}
191
218
192
219
// PreStartContainer resets the devices received
@@ -216,6 +243,19 @@ func (m *Stub) Update(devs []*pluginapi.Device) {
216
243
m .update <- devs
217
244
}
218
245
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
+
219
259
// Allocate does a mock allocation
220
260
func (m * Stub ) Allocate (ctx context.Context , r * pluginapi.AllocateRequest ) (* pluginapi.AllocateResponse , error ) {
221
261
klog .Infof ("Allocate, %+v" , r )
0 commit comments