Add flyteadmin server config option to enable grpc histogram metrics (f…

…lyteorg#4231) Signed-off-by: squiishyy <[email protected]>
squiishyy · Oct 18, 2023 · 4928a44 · 4928a44
1 parent 2f3fe6c
commit 4928a44
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 3 deletions.
diff --git a/flyteadmin/pkg/config/config.go b/flyteadmin/pkg/config/config.go
@@ -47,9 +47,10 @@ type DataProxyUploadConfig struct {
 }
 
 type GrpcConfig struct {
-	Port                int  `json:"port" pflag:",On which grpc port to serve admin"`
-	ServerReflection    bool `json:"serverReflection" pflag:",Enable GRPC Server Reflection"`
-	MaxMessageSizeBytes int  `json:"maxMessageSizeBytes" pflag:",The max size in bytes for incoming gRPC messages"`
+	Port                 int  `json:"port" pflag:",On which grpc port to serve admin"`
+	ServerReflection     bool `json:"serverReflection" pflag:",Enable GRPC Server Reflection"`
+	MaxMessageSizeBytes  int  `json:"maxMessageSizeBytes" pflag:",The max size in bytes for incoming gRPC messages"`
+	EnableGrpcHistograms bool `json:"enableGrpcHistograms" pflag:",Enable grpc histograms"`
 }
 
 // KubeClientConfig contains the configuration used by flyteadmin to configure its internal Kubernetes Client.

diff --git a/flyteadmin/pkg/config/serverconfig_flags.go b/flyteadmin/pkg/config/serverconfig_flags.go
diff --git a/flyteadmin/pkg/config/serverconfig_flags_test.go b/flyteadmin/pkg/config/serverconfig_flags_test.go
diff --git a/flyteadmin/pkg/server/service.go b/flyteadmin/pkg/server/service.go
@@ -80,6 +80,11 @@ func newGRPCServer(ctx context.Context, pluginRegistry *plugins.Registry, cfg *c
 	pluginRegistry.RegisterDefault(plugins.PluginIDUnaryServiceMiddleware, grpcmiddleware.ChainUnaryServer(
 		RequestIDInterceptor, auth.BlanketAuthorization, auth.ExecutionUserIdentifierInterceptor))
 
+	if cfg.GrpcConfig.EnableGrpcHistograms {
+		logger.Debugf(ctx, "enabling grpc histogram metrics")
+		grpcprometheus.EnableHandlingTimeHistogram()
+	}
+
 	// Not yet implemented for streaming
 	var chainedUnaryInterceptors grpc.UnaryServerInterceptor
 	if cfg.Security.UseAuth {