Merged master

Signed-off-by: Kevin Su <[email protected]>

.github/workflows/buf_publish.yaml

@@ -0,0 +1,16 @@
+name: Publish Buf Package
+
+on:
+  release:
+    types: [created]
+
+jobs:
+  buf:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - uses: bufbuild/buf-setup-action@v1
+      - uses: bufbuild/buf-push-action@v1
+        with:
+          buf_token: ${{ secrets.BUF_TOKEN }}
+          input: 'protos'

.gitignore

@@ -15,4 +15,5 @@ docs/build/
 .vscode/
 tmp/
 .python-version
-__pycache__
+__pycache__/
+

buf.work.yaml

@@ -0,0 +1,3 @@
+version: v1
+directories:
+  - protos

clients/go/admin/client.go

@@ -192,6 +192,10 @@ func initializeClients(ctx context.Context, cfg *Config, tokenCache pkce.TokenCa
 		opts = append(opts, authOpt)
 	}
 
+	if cfg.DefaultServiceConfig != "" {
+		opts = append(opts, grpc.WithDefaultServiceConfig(cfg.DefaultServiceConfig))
+	}
+
 	adminConnection, err := NewAdminConnection(ctx, cfg, opts...)
 	if err != nil {
 		logger.Panicf(ctx, "failed to initialize Admin connection. Err: %s", err.Error())

clients/go/admin/client_test.go

@@ -20,6 +20,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"golang.org/x/oauth2"
+	_ "google.golang.org/grpc/balancer/roundrobin" //nolint
 )
 
 func TestInitializeAndGetAdminClient(t *testing.T) {
@@ -262,6 +263,42 @@ func Test_getPkceAuthTokenSource(t *testing.T) {
 	})
 }
 
+func TestGetDefaultServiceConfig(t *testing.T) {
+	u, _ := url.Parse("localhost:8089")
+	adminServiceConfig := &Config{
+		Endpoint:             config.URL{URL: *u},
+		DefaultServiceConfig: `{"loadBalancingConfig": [{"round_robin":{}}]}`,
+	}
+
+	assert.NoError(t, SetConfig(adminServiceConfig))
+
+	ctx := context.Background()
+	t.Run("legal", func(t *testing.T) {
+		u, err := url.Parse("http://localhost:8089")
+		assert.NoError(t, err)
+		clientSet, err := ClientSetBuilder().WithConfig(&Config{Endpoint: config.URL{URL: *u}, DefaultServiceConfig: `{"loadBalancingConfig": [{"round_robin":{}}]}`}).Build(ctx)
+		assert.NoError(t, err)
+		assert.NotNil(t, clientSet)
+		assert.NotNil(t, clientSet.AdminClient())
+		assert.NotNil(t, clientSet.AuthMetadataClient())
+		assert.NotNil(t, clientSet.IdentityClient())
+		assert.NotNil(t, clientSet.HealthServiceClient())
+	})
+	t.Run("illegal default service config", func(t *testing.T) {
+		defer func() {
+			if r := recover(); r == nil {
+				t.Errorf("The code did not panic")
+			}
+		}()
+
+		u, err := url.Parse("http://localhost:8089")
+		assert.NoError(t, err)
+		clientSet, err := ClientSetBuilder().WithConfig(&Config{Endpoint: config.URL{URL: *u}, DefaultServiceConfig: `{"loadBalancingConfig": [{"foo":{}}]}`}).Build(ctx)
+		assert.Error(t, err)
+		assert.Nil(t, clientSet)
+	})
+}
+
 func ExampleClientSetBuilder() {
 	ctx := context.Background()
 	// Create a client set that initializes the connection with flyte admin and sets up Auth (if needed).

clients/go/admin/config.go

@@ -1,7 +1,6 @@
 // Initializes an Admin Client that exposes all implemented services by FlyteAdmin server. The library supports different
-// authentication flows (see AuthType). It initializes the grpc connection once and reuses it. The gRPC connection is
-// sticky (it hogs one server and keeps the connection alive). For better load balancing against the server, place a
-// proxy service in between instead.
+// authentication flows (see AuthType). It initializes the grpc connection once and reuses it. A grpc load balancing policy
+// can be configured as well.
 package admin
 
 import (
@@ -68,6 +67,12 @@ type Config struct {
 	PkceConfig pkce.Config `json:"pkceConfig" pflag:",Config for Pkce authentication flow."`
 
 	Command []string `json:"command" pflag:",Command for external authentication token generation"`
+
+	// Set the gRPC service config formatted as a json string https://github.com/grpc/grpc/blob/master/doc/service_config.md
+	// eg. {"loadBalancingConfig": [{"round_robin":{}}], "methodConfig": [{"name":[{"service": "foo", "method": "bar"}, {"service": "baz"}], "timeout": "1.000000001s"}]}
+	// find the full schema here https://github.com/grpc/grpc-proto/blob/master/grpc/service_config/service_config.proto#L625
+	// Note that required packages may need to be preloaded to support certain service config. For example "google.golang.org/grpc/balancer/roundrobin" should be preloaded to have round-robin policy supported.
+	DefaultServiceConfig string `json:"defaultServiceConfig" pdflag:",Set the default service config for the admin gRPC client"`
 }
 
 var (