fix(connector): add gRPC keepalive and default UNAVAILABLE retry policy

flyteplugins/go/tasks/plugins/webapi/connector/client.go

@@ -11,6 +11,7 @@ import (
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/grpclog"
+	"google.golang.org/grpc/keepalive"
 	"google.golang.org/grpc/status"
 
 	"github.com/flyteorg/flyte/v2/flytestdlib/config"
@@ -20,6 +21,27 @@ import (
 
 const defaultTaskTypeVersion = 0
 
+// DefaultGRPCServiceConfig is the gRPC service config applied to a connector
+// connection when its Deployment does not set DefaultServiceConfig. It enables
+// round-robin load balancing across connector replicas and retries transient
+// UNAVAILABLE failures so a single dropped connection does not fail the task.
+// It is exported so callers that register connectors programmatically (e.g. the
+// operator app controller) can set it explicitly.
+const DefaultGRPCServiceConfig = `{
+  "loadBalancingConfig": [{"round_robin":{}}],
+  "methodConfig": [{
+    "name": [{"service": "flyteidl2.connector.AsyncConnectorService"}],
+    "retryPolicy": {
+      "maxAttempts": 4,
+      "initialBackoff": "0.2s",
+      "maxBackoff": "3s",
+      "backoffMultiplier": 2.0,
+      "retryableStatusCodes": ["UNAVAILABLE"]
+    }
+  }],
+  "retryThrottling": {"maxTokens": 10, "tokenRatio": 0.1}
+}`
+
 type Connector struct {
 	// ConnectorDeployment is the connector deployment where this connector is running.
 	ConnectorDeployment *Deployment
@@ -50,8 +72,22 @@ func getGrpcConnection(ctx context.Context, connector *Deployment) (*grpc.Client
 		opts = append(opts, grpc.WithTransportCredentials(creds))
 	}
 
-	if len(connector.DefaultServiceConfig) != 0 {
-		opts = append(opts, grpc.WithDefaultServiceConfig(connector.DefaultServiceConfig))
+	serviceConfig := connector.DefaultServiceConfig
+	if len(serviceConfig) == 0 {
+		serviceConfig = DefaultGRPCServiceConfig
+	}
+	opts = append(opts, grpc.WithDefaultServiceConfig(serviceConfig))
+
+	// Keepalive is opt-in per deployment. It is only safe for connectors fronted
+	// by an L7 gateway (e.g. Knative/Kourier Envoy) that reaps idle connections;
+	// a directly-reached grpc-go or grpcio (C-core) server rejects frequent idle
+	// pings with GOAWAY "too_many_pings".
+	if connector.Keepalive != nil {
+		opts = append(opts, grpc.WithKeepaliveParams(keepalive.ClientParameters{
+			Time:                connector.Keepalive.Time.Duration,
+			Timeout:             connector.Keepalive.Timeout.Duration,
+			PermitWithoutStream: connector.Keepalive.PermitWithoutStream,
+		}))
 	}
 
 	var err error

flyteplugins/go/tasks/plugins/webapi/connector/client_test.go

@@ -2,9 +2,15 @@ package connector
 
 import (
 	"context"
+	"encoding/json"
+	"strings"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/flyteorg/flyte/v2/flytestdlib/config"
 )
 
 func TestInitializeClients(t *testing.T) {
@@ -26,3 +32,72 @@ func TestInitializeClients(t *testing.T) {
 	_, ok = cs.asyncConnectorClients["x"]
 	assert.True(t, ok)
 }
+
+func TestDefaultGRPCServiceConfig(t *testing.T) {
+	// Must be valid JSON — grpc.WithDefaultServiceConfig silently ignores a
+	// malformed config, so a typo here would disable LB/retry without any error.
+	assert.True(t, json.Valid([]byte(DefaultGRPCServiceConfig)), "DefaultGRPCServiceConfig must be valid JSON")
+
+	var parsed struct {
+		LoadBalancingConfig []map[string]any `json:"loadBalancingConfig"`
+		MethodConfig        []struct {
+			Name []struct {
+				Service string `json:"service"`
+			} `json:"name"`
+			RetryPolicy struct {
+				MaxAttempts          int      `json:"maxAttempts"`
+				RetryableStatusCodes []string `json:"retryableStatusCodes"`
+			} `json:"retryPolicy"`
+		} `json:"methodConfig"`
+		RetryThrottling map[string]any `json:"retryThrottling"`
+	}
+	require.NoError(t, json.Unmarshal([]byte(DefaultGRPCServiceConfig), &parsed))
+
+	require.Len(t, parsed.LoadBalancingConfig, 1)
+	_, hasRoundRobin := parsed.LoadBalancingConfig[0]["round_robin"]
+	assert.True(t, hasRoundRobin, "expected round_robin load balancing")
+
+	require.Len(t, parsed.MethodConfig, 1)
+	mc := parsed.MethodConfig[0]
+	require.Len(t, mc.Name, 1)
+	assert.Equal(t, "flyteidl2.connector.AsyncConnectorService", mc.Name[0].Service)
+	assert.Greater(t, mc.RetryPolicy.MaxAttempts, 1)
+	assert.Equal(t, []string{"UNAVAILABLE"}, mc.RetryPolicy.RetryableStatusCodes)
+	assert.NotEmpty(t, parsed.RetryThrottling, "expected retryThrottling to bound retry storms")
+}
+
+func TestGetGrpcConnection(t *testing.T) {
+	ctx := context.Background()
+
+	// Empty DefaultServiceConfig must fall back to DefaultGRPCServiceConfig
+	// (round-robin + retry) rather than no service config at all.
+	conn, err := getGrpcConnection(ctx, &Deployment{Endpoint: "x", Insecure: true})
+	require.NoError(t, err)
+	require.NotNil(t, conn)
+	assert.NoError(t, conn.Close())
+
+	// A deployment-specific DefaultServiceConfig still takes precedence.
+	custom := `{"loadBalancingConfig": [{"round_robin":{}}]}`
+	conn, err = getGrpcConnection(ctx, &Deployment{Endpoint: "y", Insecure: true, DefaultServiceConfig: custom})
+	require.NoError(t, err)
+	require.NotNil(t, conn)
+	assert.NoError(t, conn.Close())
+
+	// Keepalive is opt-in: a deployment may enable it (gateway-fronted connectors).
+	conn, err = getGrpcConnection(ctx, &Deployment{
+		Endpoint: "z", Insecure: true,
+		Keepalive: &KeepaliveConfig{
+			Time:                config.Duration{Duration: 30 * time.Second},
+			Timeout:             config.Duration{Duration: 10 * time.Second},
+			PermitWithoutStream: true,
+		},
+	})
+	require.NoError(t, err)
+	require.NotNil(t, conn)
+	assert.NoError(t, conn.Close())
+}
+
+// ensure the constant referenced from config.go and the LB choice stay in sync
+func TestDefaultGRPCServiceConfigMentionsRoundRobin(t *testing.T) {
+	assert.True(t, strings.Contains(DefaultGRPCServiceConfig, "round_robin"))
+}

flyteplugins/go/tasks/plugins/webapi/connector/config.go

@@ -41,10 +41,13 @@ var (
 			},
 		},
 		DefaultConnector: Deployment{
-			Endpoint:             "",
-			Insecure:             true,
-			DefaultTimeout:       config.Duration{Duration: 10 * time.Second},
-			DefaultServiceConfig: `{"loadBalancingConfig": [{"round_robin":{}}]}`,
+			Endpoint:       "",
+			Insecure:       true,
+			DefaultTimeout: config.Duration{Duration: 10 * time.Second},
+			// DefaultServiceConfig is left empty so getGrpcConnection falls back to
+			// DefaultGRPCServiceConfig (round-robin LB + UNAVAILABLE retry). Set this
+			// per-deployment only to override that default.
+			DefaultServiceConfig: "",
 		},
 		ConnectorDeployments:  map[string]*Deployment{},
 		ConnectorForTaskTypes: map[string]string{},
@@ -95,6 +98,14 @@ type Deployment struct {
 	// DefaultServiceConfig sets default gRPC service config; check https://github.com/grpc/grpc/blob/master/doc/service_config.md for more details
 	DefaultServiceConfig string `json:"defaultServiceConfig,omitempty" yaml:"defaultServiceConfig,omitempty"`
 
+	// Keepalive configures gRPC client keepalive pings for this connector. Leave
+	// unset for connectors reached directly: a grpc-go or grpcio (C-core) server
+	// rejects frequent idle pings with GOAWAY "too_many_pings". Set it for
+	// connectors fronted by an L7 gateway (e.g. Knative/Kourier Envoy) that reaps
+	// idle connections, so the connection is kept warm and dead connections are
+	// detected proactively instead of on the next RPC.
+	Keepalive *KeepaliveConfig `json:"keepalive,omitempty" yaml:"keepalive,omitempty"`
+
 	// Timeouts defines various RPC timeout values for different plugin operations: CreateTask, GetTask, DeleteTask; if not configured, defaults to DefaultTimeout
 	Timeouts map[string]config.Duration `json:"timeouts,omitempty" yaml:"timeouts,omitempty"`
 
@@ -108,6 +119,23 @@ type Deployment struct {
 	Domain string `json:"domain,omitempty" yaml:"domain,omitempty"`
 }
 
+// KeepaliveConfig configures gRPC client keepalive pings for a connector
+// connection. See https://github.com/grpc/grpc/blob/master/doc/keepalive.md.
+type KeepaliveConfig struct {
+	// Time is the interval between keepalive pings on an otherwise idle
+	// connection. Keep it above the server/gateway enforcement threshold to
+	// avoid GOAWAY "too_many_pings".
+	Time config.Duration `json:"time,omitempty" yaml:"time,omitempty"`
+
+	// Timeout is how long the client waits for a keepalive ping ack before
+	// considering the connection dead and closing it.
+	Timeout config.Duration `json:"timeout,omitempty" yaml:"timeout,omitempty"`
+
+	// PermitWithoutStream sends keepalive pings even when there are no active
+	// RPCs. Required to keep an idle connection warm through a gateway.
+	PermitWithoutStream bool `json:"permitWithoutStream,omitempty" yaml:"permitWithoutStream,omitempty"`
+}
+
 func GetConfig() *Config {
 	return configSection.GetConfig().(*Config)
 }

flyteplugins/go/tasks/plugins/webapi/connector/config_test.go

@@ -53,7 +53,9 @@ func TestDefaultAgentConfig(t *testing.T) {
 	assert.Equal(t, "", cfg.DefaultConnector.Endpoint)
 	assert.True(t, cfg.DefaultConnector.Insecure)
 	assert.Equal(t, 10*time.Second, cfg.DefaultConnector.DefaultTimeout.Duration)
-	assert.Equal(t, `{"loadBalancingConfig": [{"round_robin":{}}]}`, cfg.DefaultConnector.DefaultServiceConfig)
+	// DefaultServiceConfig defaults to empty; getGrpcConnection falls back to
+	// DefaultGRPCServiceConfig when a deployment does not override it.
+	assert.Empty(t, cfg.DefaultConnector.DefaultServiceConfig)
 
 	assert.Empty(t, cfg.DefaultConnector.Timeouts)