From 4c92093ca7cda79be4ce2b32af416c441b5ff92e Mon Sep 17 00:00:00 2001 From: Ukri Niemimuukko Date: Mon, 28 Dec 2020 16:44:21 +0200 Subject: [PATCH] api: add PreferredAllocator interface This adds a PreferredAllocator interface so that plugins can optionally implement the API. Signed-off-by: Ukri Niemimuukko --- pkg/deviceplugin/api.go | 6 +++++ pkg/deviceplugin/manager.go | 17 ++++++++++---- pkg/deviceplugin/manager_test.go | 4 ++-- pkg/deviceplugin/server.go | 38 ++++++++++++++++++-------------- pkg/deviceplugin/server_test.go | 2 +- 5 files changed, 44 insertions(+), 23 deletions(-) diff --git a/pkg/deviceplugin/api.go b/pkg/deviceplugin/api.go index 07915de8..10f861e1 100644 --- a/pkg/deviceplugin/api.go +++ b/pkg/deviceplugin/api.go @@ -96,6 +96,12 @@ type PostAllocator interface { PostAllocate(*pluginapi.AllocateResponse) error } +// PreferredAllocator is an optional interface implemented by device plugins. +type PreferredAllocator interface { + // GetPreferredAllocation defines the list of devices preferred for allocating next. + GetPreferredAllocation(*pluginapi.PreferredAllocationRequest) (*pluginapi.PreferredAllocationResponse, error) +} + // ContainerPreStarter is an optional interface implemented by device plugins. type ContainerPreStarter interface { // PreStartContainer defines device initialization function before container is started. diff --git a/pkg/deviceplugin/manager.go b/pkg/deviceplugin/manager.go index 3414bbca..1afd9428 100644 --- a/pkg/deviceplugin/manager.go +++ b/pkg/deviceplugin/manager.go @@ -22,6 +22,10 @@ import ( pluginapi "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1" ) +type postAllocateFunc func(*pluginapi.AllocateResponse) error +type preStartContainerFunc func(*pluginapi.PreStartContainerRequest) error +type getPreferredAllocationFunc func(*pluginapi.PreferredAllocationRequest) (*pluginapi.PreferredAllocationResponse, error) + // updateInfo contains info for added, updated and deleted devices. type updateInfo struct { Added DeviceTree @@ -73,7 +77,7 @@ type Manager struct { devicePlugin Scanner namespace string servers map[string]devicePluginServer - createServer func(string, func(*pluginapi.AllocateResponse) error, func(*pluginapi.PreStartContainerRequest) error) devicePluginServer + createServer func(string, postAllocateFunc, preStartContainerFunc, getPreferredAllocationFunc) devicePluginServer } // NewManager creates a new instance of Manager. @@ -107,8 +111,9 @@ func (m *Manager) Run() { func (m *Manager) handleUpdate(update updateInfo) { klog.V(4).Info("Received dev updates:", update) for devType, devices := range update.Added { - var postAllocate func(*pluginapi.AllocateResponse) error - var preStartContainer func(*pluginapi.PreStartContainerRequest) error + var postAllocate postAllocateFunc + var preStartContainer preStartContainerFunc + var getPreferredAllocation getPreferredAllocationFunc if postAllocator, ok := m.devicePlugin.(PostAllocator); ok { postAllocate = postAllocator.PostAllocate @@ -118,7 +123,11 @@ func (m *Manager) handleUpdate(update updateInfo) { preStartContainer = containerPreStarter.PreStartContainer } - m.servers[devType] = m.createServer(devType, postAllocate, preStartContainer) + if preferredAllocator, ok := m.devicePlugin.(PreferredAllocator); ok { + getPreferredAllocation = preferredAllocator.GetPreferredAllocation + } + + m.servers[devType] = m.createServer(devType, postAllocate, preStartContainer, getPreferredAllocation) go func(dt string) { err := m.servers[dt].Serve(m.namespace) if err != nil { diff --git a/pkg/deviceplugin/manager_test.go b/pkg/deviceplugin/manager_test.go index 314de7f6..a31494a6 100644 --- a/pkg/deviceplugin/manager_test.go +++ b/pkg/deviceplugin/manager_test.go @@ -265,7 +265,7 @@ func TestHandleUpdate(t *testing.T) { mgr := Manager{ devicePlugin: &devicePluginStub{}, servers: tt.servers, - createServer: func(string, func(*pluginapi.AllocateResponse) error, func(*pluginapi.PreStartContainerRequest) error) devicePluginServer { + createServer: func(string, postAllocateFunc, preStartContainerFunc, getPreferredAllocationFunc) devicePluginServer { return &serverStub{} }, } @@ -279,7 +279,7 @@ func TestHandleUpdate(t *testing.T) { func TestRun(t *testing.T) { mgr := NewManager("testnamespace", &devicePluginStub{}) - mgr.createServer = func(string, func(*pluginapi.AllocateResponse) error, func(*pluginapi.PreStartContainerRequest) error) devicePluginServer { + mgr.createServer = func(string, postAllocateFunc, preStartContainerFunc, getPreferredAllocationFunc) devicePluginServer { return &serverStub{} } mgr.Run() diff --git a/pkg/deviceplugin/server.go b/pkg/deviceplugin/server.go index c3432f8b..70463e0b 100644 --- a/pkg/deviceplugin/server.go +++ b/pkg/deviceplugin/server.go @@ -51,27 +51,30 @@ type devicePluginServer interface { // server implements devicePluginServer and pluginapi.PluginInterfaceServer interfaces. type server struct { - devType string - grpcServer *grpc.Server - updatesCh chan map[string]DeviceInfo - devices map[string]DeviceInfo - postAllocate func(*pluginapi.AllocateResponse) error - preStartContainer func(*pluginapi.PreStartContainerRequest) error - state serverState - stateMutex sync.Mutex + devType string + grpcServer *grpc.Server + updatesCh chan map[string]DeviceInfo + devices map[string]DeviceInfo + postAllocate postAllocateFunc + preStartContainer preStartContainerFunc + getPreferredAllocation getPreferredAllocationFunc + state serverState + stateMutex sync.Mutex } // newServer creates a new server satisfying the devicePluginServer interface. func newServer(devType string, - postAllocate func(*pluginapi.AllocateResponse) error, - preStartContainer func(*pluginapi.PreStartContainerRequest) error) devicePluginServer { + postAllocate postAllocateFunc, + preStartContainer preStartContainerFunc, + getPreferredAllocation getPreferredAllocationFunc) devicePluginServer { return &server{ - devType: devType, - updatesCh: make(chan map[string]DeviceInfo, 1), // TODO: is 1 needed? - devices: make(map[string]DeviceInfo), - postAllocate: postAllocate, - preStartContainer: preStartContainer, - state: uninitialized, + devType: devType, + updatesCh: make(chan map[string]DeviceInfo, 1), // TODO: is 1 needed? + devices: make(map[string]DeviceInfo), + postAllocate: postAllocate, + preStartContainer: preStartContainer, + getPreferredAllocation: getPreferredAllocation, + state: uninitialized, } } @@ -162,6 +165,9 @@ func (srv *server) PreStartContainer(ctx context.Context, rqt *pluginapi.PreStar } func (srv *server) GetPreferredAllocation(ctx context.Context, rqt *pluginapi.PreferredAllocationRequest) (*pluginapi.PreferredAllocationResponse, error) { + if srv.getPreferredAllocation != nil { + return srv.getPreferredAllocation(rqt) + } return nil, errors.New("GetPreferredAllocation should not be called as this device plugin doesn't implement it") } diff --git a/pkg/deviceplugin/server_test.go b/pkg/deviceplugin/server_test.go index bc302c82..627b321f 100644 --- a/pkg/deviceplugin/server_test.go +++ b/pkg/deviceplugin/server_test.go @@ -530,7 +530,7 @@ func TestGetPreferredAllocation(t *testing.T) { } func TestNewServer(t *testing.T) { - _ = newServer("test", nil, nil) + _ = newServer("test", nil, nil, nil) } func TestUpdate(t *testing.T) {