DRA kubelet: list supported gRPC services during registration

Listing supported gRPC services (e.g. drav1alpha3.Node, drav1beta1.DRAPlugin)
during registration enables the kubelet to determine in advance which methods
it can call.

Versioning by Kubernetes release makes less sense because it doesn't say
anything about which gRPC service is supported. New ones might get added and
obsolete ones removed. Some services might be optional.

In the past, this versioning support wasn't really used. At least one version
had to be provided and kubelet tried to use the plugin with the highest
version. This version comparison gets dropped. In the unlikely situation
that different plugins register under the same name, the most recent one is
used.

Because advertising gRPC services is a new convention, plugins only reporting
some version are treated as providing the old alpha gRPC service.

Author: pohly

pkg/kubelet/cm/dra/manager_test.go

@@ -581,7 +581,7 @@ func TestPrepareResources(t *testing.T) {
 			defer draServerInfo.teardownFn()
 
 			plg := plugin.NewRegistrationHandler(nil, getFakeNode)
-			if err := plg.RegisterPlugin(test.driverName, draServerInfo.socketName, []string{"1.27"}, pluginClientTimeout); err != nil {
+			if err := plg.RegisterPlugin(test.driverName, draServerInfo.socketName, []string{drapb.DRAPluginService}, pluginClientTimeout); err != nil {
 				t.Fatalf("failed to register plugin %s, err: %v", test.driverName, err)
 			}
 			defer plg.DeRegisterPlugin(test.driverName) // for sake of next tests
@@ -718,7 +718,7 @@ func TestUnprepareResources(t *testing.T) {
 			defer draServerInfo.teardownFn()
 
 			plg := plugin.NewRegistrationHandler(nil, getFakeNode)
-			if err := plg.RegisterPlugin(test.driverName, draServerInfo.socketName, []string{"1.27"}, pluginClientTimeout); err != nil {
+			if err := plg.RegisterPlugin(test.driverName, draServerInfo.socketName, []string{drapb.DRAPluginService}, pluginClientTimeout); err != nil {
 				t.Fatalf("failed to register plugin %s, err: %v", test.driverName, err)
 			}
 			defer plg.DeRegisterPlugin(test.driverName) // for sake of next tests
@@ -888,7 +888,7 @@ func TestParallelPrepareUnprepareResources(t *testing.T) {
 	defer draServerInfo.teardownFn()
 
 	plg := plugin.NewRegistrationHandler(nil, getFakeNode)
-	if err := plg.RegisterPlugin(driverName, draServerInfo.socketName, []string{"1.27"}, nil); err != nil {
+	if err := plg.RegisterPlugin(driverName, draServerInfo.socketName, []string{drapb.DRAPluginService}, nil); err != nil {
 		t.Fatalf("failed to register plugin %s, err: %v", driverName, err)
 	}
 	defer plg.DeRegisterPlugin(driverName)

pkg/kubelet/cm/dra/plugin/plugin.go

@@ -25,12 +25,10 @@ import (
 	"time"
 
 	"google.golang.org/grpc"
-	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/status"
 
-	utilversion "k8s.io/apimachinery/pkg/util/version"
 	"k8s.io/klog/v2"
 	drapbv1alpha4 "k8s.io/kubelet/pkg/apis/dra/v1alpha4"
 	drapbv1beta1 "k8s.io/kubelet/pkg/apis/dra/v1beta1"
@@ -59,27 +57,19 @@ type Plugin struct {
 	backgroundCtx context.Context
 	cancel        func(cause error)
 
-	mutex                   sync.Mutex
-	conn                    *grpc.ClientConn
-	supportedAPI            apiVersion
-	endpoint                string
-	highestSupportedVersion *utilversion.Version
-	clientCallTimeout       time.Duration
+	mutex             sync.Mutex
+	conn              *grpc.ClientConn
+	endpoint          string
+	chosenService     string // e.g. drapbv1beta1.DRAPluginService
+	clientCallTimeout time.Duration
 }
 
-type apiVersion string
-
-const (
-	apiV1alpha4 = apiVersion("v1alpha4")
-	apiV1beta1  = apiVersion("v1beta1")
-)
-
-func (p *Plugin) getOrCreateGRPCConn() (*grpc.ClientConn, apiVersion, error) {
+func (p *Plugin) getOrCreateGRPCConn() (*grpc.ClientConn, error) {
 	p.mutex.Lock()
 	defer p.mutex.Unlock()
 
 	if p.conn != nil {
-		return p.conn, p.supportedAPI, nil
+		return p.conn, nil
 	}
 
 	ctx := p.backgroundCtx
@@ -101,18 +91,18 @@ func (p *Plugin) getOrCreateGRPCConn() (*grpc.ClientConn, apiVersion, error) {
 		grpc.WithChainUnaryInterceptor(newMetricsInterceptor(p.name)),
 	)
 	if err != nil {
-		return nil, "", err
+		return nil, err
 	}
 
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
 	defer cancel()
 
 	if ok := conn.WaitForStateChange(ctx, connectivity.Connecting); !ok {
-		return nil, "", errors.New("timed out waiting for gRPC connection to be ready")
+		return nil, errors.New("timed out waiting for gRPC connection to be ready")
 	}
 
 	p.conn = conn
-	return p.conn, "", nil
+	return p.conn, nil
 }
 
 func (p *Plugin) NodePrepareResources(
@@ -123,7 +113,7 @@ func (p *Plugin) NodePrepareResources(
 	logger := klog.FromContext(ctx)
 	logger.V(4).Info("Calling NodePrepareResources rpc", "request", req)
 
-	conn, supportedAPI, err := p.getOrCreateGRPCConn()
+	conn, err := p.getOrCreateGRPCConn()
 	if err != nil {
 		return nil, err
 	}
@@ -132,29 +122,17 @@ func (p *Plugin) NodePrepareResources(
 	defer cancel()
 
 	var response *drapbv1beta1.NodePrepareResourcesResponse
-	switch supportedAPI {
-	case apiV1beta1:
+	switch p.chosenService {
+	case drapbv1beta1.DRAPluginService:
 		nodeClient := drapbv1beta1.NewDRAPluginClient(conn)
 		response, err = nodeClient.NodePrepareResources(ctx, req)
-	case apiV1alpha4:
+	case drapbv1alpha4.NodeService:
 		nodeClient := drapbv1alpha4.NewNodeClient(conn)
 		response, err = nodeClient.NodePrepareResources(ctx, req)
 	default:
-		// Try it, fall back if necessary.
-		supportedAPI = apiV1beta1
-		nodeClient := drapbv1beta1.NewDRAPluginClient(conn)
-		response, err = nodeClient.NodePrepareResources(ctx, req)
-		if err != nil && status.Convert(err).Code() == codes.Unimplemented {
-			supportedAPI = apiV1alpha4
-			nodeClient := drapbv1alpha4.NewNodeClient(conn)
-			response, err = nodeClient.NodePrepareResources(ctx, req)
-		}
-		if err == nil || status.Convert(err).Code() != codes.Unimplemented {
-			// Store discovered version for future use.
-			p.mutex.Lock()
-			p.supportedAPI = supportedAPI
-			p.mutex.Unlock()
-		}
+		// Shouldn't happen, validateSupportedServices should only
+		// return services we support here.
+		return nil, fmt.Errorf("internal error: unsupported chosen service: %q", p.chosenService)
 	}
 	logger.V(4).Info("Done calling NodePrepareResources rpc", "response", response, "err", err)
 	return response, err
@@ -168,7 +146,7 @@ func (p *Plugin) NodeUnprepareResources(
 	logger := klog.FromContext(ctx)
 	logger.V(4).Info("Calling NodeUnprepareResource rpc", "request", req)
 
-	conn, supportedAPI, err := p.getOrCreateGRPCConn()
+	conn, err := p.getOrCreateGRPCConn()
 	if err != nil {
 		return nil, err
 	}
@@ -177,29 +155,17 @@ func (p *Plugin) NodeUnprepareResources(
 	defer cancel()
 
 	var response *drapbv1beta1.NodeUnprepareResourcesResponse
-	switch supportedAPI {
-	case apiV1beta1:
+	switch p.chosenService {
+	case drapbv1beta1.DRAPluginService:
 		nodeClient := drapbv1beta1.NewDRAPluginClient(conn)
 		response, err = nodeClient.NodeUnprepareResources(ctx, req)
-	case apiV1alpha4:
+	case drapbv1alpha4.NodeService:
 		nodeClient := drapbv1alpha4.NewNodeClient(conn)
 		response, err = nodeClient.NodeUnprepareResources(ctx, req)
 	default:
-		// Try it, fall back if necessary.
-		supportedAPI = apiV1beta1
-		nodeClient := drapbv1beta1.NewDRAPluginClient(conn)
-		response, err = nodeClient.NodeUnprepareResources(ctx, req)
-		if err != nil && status.Convert(err).Code() == codes.Unimplemented {
-			supportedAPI = apiV1alpha4
-			nodeClient := drapbv1alpha4.NewNodeClient(conn)
-			response, err = nodeClient.NodeUnprepareResources(ctx, req)
-		}
-		if err == nil || status.Convert(err).Code() != codes.Unimplemented {
-			// Store discovered version for future use.
-			p.mutex.Lock()
-			p.supportedAPI = supportedAPI
-			p.mutex.Unlock()
-		}
+		// Shouldn't happen, validateSupportedServices should only
+		// return services we support here.
+		return nil, fmt.Errorf("internal error: unsupported chosen service: %q", p.chosenService)
 	}
 	logger.V(4).Info("Done calling NodeUnprepareResources rpc", "response", response, "err", err)
 	return response, err

pkg/kubelet/cm/dra/plugin/plugin_test.go

@@ -27,23 +27,20 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"google.golang.org/grpc"
-	drapb "k8s.io/kubelet/pkg/apis/dra/v1beta1"
+	drapbv1alpha4 "k8s.io/kubelet/pkg/apis/dra/v1alpha4"
+	drapbv1beta1 "k8s.io/kubelet/pkg/apis/dra/v1beta1"
 	"k8s.io/kubernetes/test/utils/ktesting"
 )
 
-const (
-	v1alpha4Version = "v1alpha4"
-)
-
-type fakeV1alpha4GRPCServer struct {
-	drapb.UnimplementedDRAPluginServer
+type fakeGRPCServer struct {
+	drapbv1beta1.UnimplementedDRAPluginServer
 }
 
-var _ drapb.DRAPluginServer = &fakeV1alpha4GRPCServer{}
+var _ drapbv1beta1.DRAPluginServer = &fakeGRPCServer{}
 
-func (f *fakeV1alpha4GRPCServer) NodePrepareResources(ctx context.Context, in *drapb.NodePrepareResourcesRequest) (*drapb.NodePrepareResourcesResponse, error) {
-	return &drapb.NodePrepareResourcesResponse{Claims: map[string]*drapb.NodePrepareResourceResponse{"claim-uid": {
-		Devices: []*drapb.Device{
+func (f *fakeGRPCServer) NodePrepareResources(ctx context.Context, in *drapbv1beta1.NodePrepareResourcesRequest) (*drapbv1beta1.NodePrepareResourcesResponse, error) {
+	return &drapbv1beta1.NodePrepareResourcesResponse{Claims: map[string]*drapbv1beta1.NodePrepareResourceResponse{"claim-uid": {
+		Devices: []*drapbv1beta1.Device{
 			{
 				RequestNames: []string{"test-request"},
 				CDIDeviceIDs: []string{"test-cdi-id"},
@@ -52,14 +49,14 @@ func (f *fakeV1alpha4GRPCServer) NodePrepareResources(ctx context.Context, in *d
 	}}}, nil
 }
 
-func (f *fakeV1alpha4GRPCServer) NodeUnprepareResources(ctx context.Context, in *drapb.NodeUnprepareResourcesRequest) (*drapb.NodeUnprepareResourcesResponse, error) {
+func (f *fakeGRPCServer) NodeUnprepareResources(ctx context.Context, in *drapbv1beta1.NodeUnprepareResourcesRequest) (*drapbv1beta1.NodeUnprepareResourcesResponse, error) {
 
-	return &drapb.NodeUnprepareResourcesResponse{}, nil
+	return &drapbv1beta1.NodeUnprepareResourcesResponse{}, nil
 }
 
 type tearDown func()
 
-func setupFakeGRPCServer(version string) (string, tearDown, error) {
+func setupFakeGRPCServer(service string) (string, tearDown, error) {
 	p, err := os.MkdirTemp("", "dra_plugin")
 	if err != nil {
 		return "", nil, err
@@ -81,12 +78,14 @@ func setupFakeGRPCServer(version string) (string, tearDown, error) {
 	}
 
 	s := grpc.NewServer()
-	switch version {
-	case v1alpha4Version:
-		fakeGRPCServer := &fakeV1alpha4GRPCServer{}
-		drapb.RegisterDRAPluginServer(s, fakeGRPCServer)
+	fakeGRPCServer := &fakeGRPCServer{}
+	switch service {
+	case drapbv1beta1.DRAPluginService:
+		drapbv1beta1.RegisterDRAPluginServer(s, fakeGRPCServer)
+	case drapbv1alpha4.NodeService:
+		drapbv1alpha4.RegisterNodeServer(s, fakeGRPCServer)
 	default:
-		return "", nil, fmt.Errorf("unsupported version: %s", version)
+		return "", nil, fmt.Errorf("unsupported gRPC service: %s", service)
 	}
 
 	go func() {
@@ -104,7 +103,8 @@ func setupFakeGRPCServer(version string) (string, tearDown, error) {
 
 func TestGRPCConnIsReused(t *testing.T) {
 	tCtx := ktesting.Init(t)
-	addr, teardown, err := setupFakeGRPCServer(v1alpha4Version)
+	service := drapbv1beta1.DRAPluginService
+	addr, teardown, err := setupFakeGRPCServer(service)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -119,10 +119,11 @@ func TestGRPCConnIsReused(t *testing.T) {
 		name:              pluginName,
 		backgroundCtx:     tCtx,
 		endpoint:          addr,
+		chosenService:     service,
 		clientCallTimeout: defaultClientCallTimeout,
 	}
 
-	conn, _, err := p.getOrCreateGRPCConn()
+	conn, err := p.getOrCreateGRPCConn()
 	defer func() {
 		err := conn.Close()
 		if err != nil {
@@ -148,8 +149,8 @@ func TestGRPCConnIsReused(t *testing.T) {
 				return
 			}
 
-			req := &drapb.NodePrepareResourcesRequest{
-				Claims: []*drapb.Claim{
+			req := &drapbv1beta1.NodePrepareResourcesRequest{
+				Claims: []*drapbv1beta1.Claim{
 					{
 						Namespace: "dummy-namespace",
 						UID:       "dummy-uid",
@@ -233,21 +234,27 @@ func TestNewDRAPluginClient(t *testing.T) {
 
 func TestNodeUnprepareResources(t *testing.T) {
 	for _, test := range []struct {
-		description   string
-		serverSetup   func(string) (string, tearDown, error)
-		serverVersion string
-		request       *drapb.NodeUnprepareResourcesRequest
+		description string
+		serverSetup func(string) (string, tearDown, error)
+		service     string
+		request     *drapbv1beta1.NodeUnprepareResourcesRequest
 	}{
 		{
-			description:   "server supports v1alpha4",
-			serverSetup:   setupFakeGRPCServer,
-			serverVersion: v1alpha4Version,
-			request:       &drapb.NodeUnprepareResourcesRequest{},
+			description: "server supports v1alpha4",
+			serverSetup: setupFakeGRPCServer,
+			service:     drapbv1alpha4.NodeService,
+			request:     &drapbv1beta1.NodeUnprepareResourcesRequest{},
+		},
+		{
+			description: "server supports v1beta1",
+			serverSetup: setupFakeGRPCServer,
+			service:     drapbv1beta1.DRAPluginService,
+			request:     &drapbv1beta1.NodeUnprepareResourcesRequest{},
 		},
 	} {
 		t.Run(test.description, func(t *testing.T) {
 			tCtx := ktesting.Init(t)
-			addr, teardown, err := setupFakeGRPCServer(test.serverVersion)
+			addr, teardown, err := setupFakeGRPCServer(test.service)
 			if err != nil {
 				t.Fatal(err)
 			}
@@ -258,10 +265,11 @@ func TestNodeUnprepareResources(t *testing.T) {
 				name:              pluginName,
 				backgroundCtx:     tCtx,
 				endpoint:          addr,
+				chosenService:     test.service,
 				clientCallTimeout: defaultClientCallTimeout,
 			}
 
-			conn, _, err := p.getOrCreateGRPCConn()
+			conn, err := p.getOrCreateGRPCConn()
 			defer func() {
 				err := conn.Close()
 				if err != nil {