feat: add trainer service to rpc package (#2209)
Signed-off-by: Gaius <gaius.qi@gmail.com>
This commit is contained in:
Gaius
parent
baf581ea63
commit
f8934e701b
|
|
@ -53,7 +53,6 @@ func GetV2(ctx context.Context, dynconfig config.Dynconfig, opts ...grpc.DialOpt
|
|||
append([]grpc.DialOption{
|
||||
grpc.WithDefaultServiceConfig(pkgbalancer.BalancerServiceConfig),
|
||||
grpc.WithUnaryInterceptor(grpc_middleware.ChainUnaryClient(
|
||||
rpc.ConvertErrorUnaryClientInterceptor,
|
||||
otelgrpc.UnaryClientInterceptor(),
|
||||
grpc_prometheus.UnaryClientInterceptor,
|
||||
grpc_zap.UnaryClientInterceptor(logger.GrpcLogger.Desugar()),
|
||||
|
|
@ -64,7 +63,6 @@ func GetV2(ctx context.Context, dynconfig config.Dynconfig, opts ...grpc.DialOpt
|
|||
rpc.RefresherUnaryClientInterceptor(dynconfig),
|
||||
)),
|
||||
grpc.WithStreamInterceptor(grpc_middleware.ChainStreamClient(
|
||||
rpc.ConvertErrorStreamClientInterceptor,
|
||||
otelgrpc.StreamClientInterceptor(),
|
||||
grpc_prometheus.StreamClientInterceptor,
|
||||
grpc_zap.StreamClientInterceptor(logger.GrpcLogger.Desugar()),
|
||||
|
|
@ -93,7 +91,6 @@ func GetV2ByAddr(ctx context.Context, target string, opts ...grpc.DialOption) (V
|
|||
append([]grpc.DialOption{
|
||||
grpc.WithDefaultServiceConfig(pkgbalancer.BalancerServiceConfig),
|
||||
grpc.WithUnaryInterceptor(grpc_middleware.ChainUnaryClient(
|
||||
rpc.ConvertErrorUnaryClientInterceptor,
|
||||
otelgrpc.UnaryClientInterceptor(),
|
||||
grpc_prometheus.UnaryClientInterceptor,
|
||||
grpc_zap.UnaryClientInterceptor(logger.GrpcLogger.Desugar()),
|
||||
|
|
@ -103,7 +100,6 @@ func GetV2ByAddr(ctx context.Context, target string, opts ...grpc.DialOption) (V
|
|||
),
|
||||
)),
|
||||
grpc.WithStreamInterceptor(grpc_middleware.ChainStreamClient(
|
||||
rpc.ConvertErrorStreamClientInterceptor,
|
||||
otelgrpc.StreamClientInterceptor(),
|
||||
grpc_prometheus.StreamClientInterceptor,
|
||||
grpc_zap.StreamClientInterceptor(logger.GrpcLogger.Desugar()),
|
||||
|
|
|
|||
|
|
@ -0,0 +1,96 @@
|
|||
/*
|
||||
* Copyright 2023 The Dragonfly Authors
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
//go:generate mockgen -destination mocks/client_v1_mock.go -source client_v1.go -package mocks
|
||||
|
||||
package client
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
|
||||
grpc_zap "github.com/grpc-ecosystem/go-grpc-middleware/logging/zap"
|
||||
grpc_retry "github.com/grpc-ecosystem/go-grpc-middleware/retry"
|
||||
grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
|
||||
"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
|
||||
"google.golang.org/grpc"
|
||||
|
||||
trainerv1 "d7y.io/api/pkg/apis/trainer/v1"
|
||||
|
||||
logger "d7y.io/dragonfly/v2/internal/dflog"
|
||||
)
|
||||
|
||||
const (
|
||||
// maxRetries is maximum number of retries.
|
||||
maxRetries = 3
|
||||
|
||||
// backoffWaitBetween is waiting for a fixed period of
|
||||
// time between calls in backoff linear.
|
||||
backoffWaitBetween = 500 * time.Millisecond
|
||||
)
|
||||
|
||||
// GetV1ByAddr returns v1 version of the trainer client by address.
|
||||
func GetV1ByAddr(ctx context.Context, target string, opts ...grpc.DialOption) (V1, error) {
|
||||
conn, err := grpc.DialContext(
|
||||
ctx,
|
||||
target,
|
||||
append([]grpc.DialOption{
|
||||
grpc.WithUnaryInterceptor(grpc_middleware.ChainUnaryClient(
|
||||
otelgrpc.UnaryClientInterceptor(),
|
||||
grpc_prometheus.UnaryClientInterceptor,
|
||||
grpc_zap.UnaryClientInterceptor(logger.GrpcLogger.Desugar()),
|
||||
grpc_retry.UnaryClientInterceptor(
|
||||
grpc_retry.WithMax(maxRetries),
|
||||
grpc_retry.WithBackoff(grpc_retry.BackoffLinear(backoffWaitBetween)),
|
||||
),
|
||||
)),
|
||||
grpc.WithStreamInterceptor(grpc_middleware.ChainStreamClient(
|
||||
otelgrpc.StreamClientInterceptor(),
|
||||
grpc_prometheus.StreamClientInterceptor,
|
||||
grpc_zap.StreamClientInterceptor(logger.GrpcLogger.Desugar()),
|
||||
)),
|
||||
}, opts...)...,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &v1{
|
||||
TrainerClient: trainerv1.NewTrainerClient(conn),
|
||||
ClientConn: conn,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// V1 is the interface for v1 version of the grpc client.
|
||||
type V1 interface {
|
||||
// Train models of scheduler using dataset.
|
||||
Train(context.Context, ...grpc.CallOption) (trainerv1.Trainer_TrainClient, error)
|
||||
|
||||
// Close tears down the ClientConn and all underlying connections.
|
||||
Close() error
|
||||
}
|
||||
|
||||
// v1 provides v1 version of the trainer grpc function.
|
||||
type v1 struct {
|
||||
trainerv1.TrainerClient
|
||||
*grpc.ClientConn
|
||||
}
|
||||
|
||||
// Train models of scheduler using dataset.
|
||||
func (v *v1) Train(ctx context.Context, opts ...grpc.CallOption) (trainerv1.Trainer_TrainClient, error) {
|
||||
return v.TrainerClient.Train(ctx, opts...)
|
||||
}
|
||||
|
|
@ -0,0 +1,71 @@
|
|||
// Code generated by MockGen. DO NOT EDIT.
|
||||
// Source: client_v1.go
|
||||
|
||||
// Package mocks is a generated GoMock package.
|
||||
package mocks
|
||||
|
||||
import (
|
||||
context "context"
|
||||
reflect "reflect"
|
||||
|
||||
trainer "d7y.io/api/pkg/apis/trainer/v1"
|
||||
gomock "github.com/golang/mock/gomock"
|
||||
grpc "google.golang.org/grpc"
|
||||
)
|
||||
|
||||
// MockV1 is a mock of V1 interface.
|
||||
type MockV1 struct {
|
||||
ctrl *gomock.Controller
|
||||
recorder *MockV1MockRecorder
|
||||
}
|
||||
|
||||
// MockV1MockRecorder is the mock recorder for MockV1.
|
||||
type MockV1MockRecorder struct {
|
||||
mock *MockV1
|
||||
}
|
||||
|
||||
// NewMockV1 creates a new mock instance.
|
||||
func NewMockV1(ctrl *gomock.Controller) *MockV1 {
|
||||
mock := &MockV1{ctrl: ctrl}
|
||||
mock.recorder = &MockV1MockRecorder{mock}
|
||||
return mock
|
||||
}
|
||||
|
||||
// EXPECT returns an object that allows the caller to indicate expected use.
|
||||
func (m *MockV1) EXPECT() *MockV1MockRecorder {
|
||||
return m.recorder
|
||||
}
|
||||
|
||||
// Close mocks base method.
|
||||
func (m *MockV1) Close() error {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "Close")
|
||||
ret0, _ := ret[0].(error)
|
||||
return ret0
|
||||
}
|
||||
|
||||
// Close indicates an expected call of Close.
|
||||
func (mr *MockV1MockRecorder) Close() *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Close", reflect.TypeOf((*MockV1)(nil).Close))
|
||||
}
|
||||
|
||||
// Train mocks base method.
|
||||
func (m *MockV1) Train(arg0 context.Context, arg1 ...grpc.CallOption) (trainer.Trainer_TrainClient, error) {
|
||||
m.ctrl.T.Helper()
|
||||
varargs := []interface{}{arg0}
|
||||
for _, a := range arg1 {
|
||||
varargs = append(varargs, a)
|
||||
}
|
||||
ret := m.ctrl.Call(m, "Train", varargs...)
|
||||
ret0, _ := ret[0].(trainer.Trainer_TrainClient)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// Train indicates an expected call of Train.
|
||||
func (mr *MockV1MockRecorder) Train(arg0 interface{}, arg1 ...interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
varargs := append([]interface{}{arg0}, arg1...)
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Train", reflect.TypeOf((*MockV1)(nil).Train), varargs...)
|
||||
}
|
||||
|
|
@ -0,0 +1,95 @@
|
|||
/*
|
||||
* Copyright 2023 The Dragonfly Authors
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package server
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
|
||||
grpc_zap "github.com/grpc-ecosystem/go-grpc-middleware/logging/zap"
|
||||
grpc_ratelimit "github.com/grpc-ecosystem/go-grpc-middleware/ratelimit"
|
||||
grpc_recovery "github.com/grpc-ecosystem/go-grpc-middleware/recovery"
|
||||
grpc_validator "github.com/grpc-ecosystem/go-grpc-middleware/validator"
|
||||
grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
|
||||
"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
|
||||
"google.golang.org/grpc"
|
||||
"google.golang.org/grpc/health"
|
||||
healthpb "google.golang.org/grpc/health/grpc_health_v1"
|
||||
"google.golang.org/grpc/keepalive"
|
||||
"google.golang.org/grpc/reflection"
|
||||
|
||||
trainerv1 "d7y.io/api/pkg/apis/trainer/v1"
|
||||
|
||||
logger "d7y.io/dragonfly/v2/internal/dflog"
|
||||
"d7y.io/dragonfly/v2/pkg/rpc"
|
||||
)
|
||||
|
||||
const (
|
||||
// DefaultQPS is default qps of grpc server.
|
||||
DefaultQPS = 10 * 1000
|
||||
|
||||
// DefaultBurst is default burst of grpc server.
|
||||
DefaultBurst = 20 * 1000
|
||||
|
||||
// DefaultMaxConnectionIdle is default max connection idle of grpc keepalive.
|
||||
DefaultMaxConnectionIdle = 10 * time.Minute
|
||||
|
||||
// DefaultMaxConnectionAge is default max connection age of grpc keepalive.
|
||||
DefaultMaxConnectionAge = 12 * time.Hour
|
||||
|
||||
// DefaultMaxConnectionAgeGrace is default max connection age grace of grpc keepalive.
|
||||
DefaultMaxConnectionAgeGrace = 5 * time.Minute
|
||||
)
|
||||
|
||||
// New returns grpc server instance and register service on grpc server.
|
||||
func New(trainerServerV1 trainerv1.TrainerServer, opts ...grpc.ServerOption) *grpc.Server {
|
||||
limiter := rpc.NewRateLimiterInterceptor(DefaultQPS, DefaultBurst)
|
||||
|
||||
grpcServer := grpc.NewServer(append([]grpc.ServerOption{
|
||||
grpc.KeepaliveParams(keepalive.ServerParameters{
|
||||
MaxConnectionIdle: DefaultMaxConnectionIdle,
|
||||
MaxConnectionAge: DefaultMaxConnectionAge,
|
||||
MaxConnectionAgeGrace: DefaultMaxConnectionAgeGrace,
|
||||
}),
|
||||
grpc.UnaryInterceptor(grpc_middleware.ChainUnaryServer(
|
||||
grpc_ratelimit.UnaryServerInterceptor(limiter),
|
||||
otelgrpc.UnaryServerInterceptor(),
|
||||
grpc_prometheus.UnaryServerInterceptor,
|
||||
grpc_zap.UnaryServerInterceptor(logger.GrpcLogger.Desugar()),
|
||||
grpc_validator.UnaryServerInterceptor(),
|
||||
grpc_recovery.UnaryServerInterceptor(),
|
||||
)),
|
||||
grpc.StreamInterceptor(grpc_middleware.ChainStreamServer(
|
||||
grpc_ratelimit.StreamServerInterceptor(limiter),
|
||||
otelgrpc.StreamServerInterceptor(),
|
||||
grpc_prometheus.StreamServerInterceptor,
|
||||
grpc_zap.StreamServerInterceptor(logger.GrpcLogger.Desugar()),
|
||||
grpc_validator.StreamServerInterceptor(),
|
||||
grpc_recovery.StreamServerInterceptor(),
|
||||
)),
|
||||
}, opts...)...)
|
||||
|
||||
// Register servers on v1 version of the grpc server.
|
||||
trainerv1.RegisterTrainerServer(grpcServer, trainerServerV1)
|
||||
|
||||
// Register health on grpc server.
|
||||
healthpb.RegisterHealthServer(grpcServer, health.NewServer())
|
||||
|
||||
// Register reflection on grpc server.
|
||||
reflection.Register(grpcServer)
|
||||
return grpcServer
|
||||
}
|
||||
Loading…
Reference in New Issue