feat(backend): add gRPC metrics to api-server (RPS/latency), optimize execution spec reporting (#12010)

* - add standard grpc metrics to api-server
- add report gap histogram
- optimize create or update tasks query

Signed-off-by: ntny <[email protected]>
Signed-off-by: arpechenin <[email protected]>

* - move to the latest version of grpc-prometheus

Signed-off-by: arpechenin <[email protected]>

* - register metrics

Signed-off-by: arpechenin <[email protected]>

* - bump metrics lib

Signed-off-by: arpechenin <[email protected]>

* - fix metrics injection

Signed-off-by: arpechenin <[email protected]>

* merge

Signed-off-by: arpechenin <[email protected]>

---------

Signed-off-by: ntny <[email protected]>
Signed-off-by: arpechenin <[email protected]>
Co-authored-by: arpechenin <[email protected]>

Author: ntny

backend/src/apiserver/api-server-metrics.md

@@ -0,0 +1,32 @@
+## Metrics
+
+Standard gRPC metrics are available via https://github.com/grpc-ecosystem/go-grpc-middleware/tree/main/providers/prometheus
+
+### Sample
+
+gRPC RPMs by grpc_method and grpc_code for the ml-pipeline server:
+```yaml
+sum by (grpc_service, grpc_method, grpc_code) (
+rate(grpc_server_handled_total{app=~"ml-pipeline", kubernetes_namespace="kubeflow"}[1m])
+) * 60
+```
+
+95th percentile gRPC latency by grpc_service and grpc_method for the ml-pipeline server:
+```yaml
+histogram_quantile(
+  0.95,
+  sum by (grpc_service, grpc_method, le) (
+    rate(grpc_server_handling_seconds_bucket{
+      app="ml-pipeline",
+      kubernetes_namespace="kubeflow"
+    }[1m])
+  )
+)
+```
+
+Gap in seconds between creating an execution spec (Argo or other backend) for a recurring run and reporting it via the persistence agent
+```yaml
+histogram_quantile(0.95, sum(rate(resource_manager_recurring_run_report_gap_bucket[1h])) by (le))
+```
+
+

backend/src/apiserver/main.go

@@ -28,6 +28,7 @@ import (
 	"strings"
 	"sync"
 
+	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
 	"google.golang.org/grpc/credentials"
 
 	"github.com/fsnotify/fsnotify"
@@ -247,18 +248,32 @@ func grpcCustomMatcher(key string) (string, bool) {
 
 func startRPCServer(resourceManager *resource.ResourceManager, tlsCfg *tls.Config) {
 	var s *grpc.Server
+
+	grpc_prometheus.EnableHandlingTimeHistogram(
+		grpc_prometheus.WithHistogramBuckets([]float64{
+			0.005, 0.01, 0.03, 0.1, 0.3, 1, 3, 10, 15, 30, 60, 120, 300, // 5 ms -> 5 min
+		}),
+	)
+
 	if tlsCfg != nil {
 		glog.Info("Starting RPC server (TLS enabled)")
 		tlsCredentials := credentials.NewTLS(tlsCfg)
 		s = grpc.NewServer(
 			grpc.Creds(tlsCredentials),
-			grpc.UnaryInterceptor(apiServerInterceptor),
+			grpc.ChainUnaryInterceptor(
+				grpc_prometheus.UnaryServerInterceptor,
+				apiServerInterceptor,
+			),
 			grpc.MaxRecvMsgSize(math.MaxInt32),
 		)
 	} else {
 		glog.Info("Starting RPC server")
-		s = grpc.NewServer(grpc.UnaryInterceptor(apiServerInterceptor), grpc.MaxRecvMsgSize(math.MaxInt32))
+		s = grpc.NewServer(grpc.ChainUnaryInterceptor(
+			grpc_prometheus.UnaryServerInterceptor,
+			apiServerInterceptor,
+		), grpc.MaxRecvMsgSize(math.MaxInt32))
 	}
+
 	listener, err := net.Listen("tcp", *rpcPortFlag)
 	if err != nil {
 		glog.Fatalf("Failed to start RPC server: %v", err)

backend/src/apiserver/resource/resource_manager.go

@@ -22,6 +22,7 @@ import (
 	"net"
 	"reflect"
 	"strconv"
+	"time"
 
 	apiv2beta1 "github.com/kubeflow/pipelines/backend/api/v2beta1/go_client"
 	scheduledworkflow "github.com/kubeflow/pipelines/backend/src/crd/pkg/apis/scheduledworkflow/v1beta1"
@@ -79,6 +80,13 @@ var (
 		Help: "The current number of failed workflow runs",
 	}, extraLabels)
 
+	// Gap in seconds between creating an execution spec (Argo or other backend) for a recurring run and reporting it via the persistence agent.
+	recurringPipelineRunReportGap = promauto.NewHistogram(prometheus.HistogramOpts{
+		Name:    "resource_manager_recurring_run_report_gap",
+		Help:    "Recurring Run Report Delay",
+		Buckets: prometheus.ExponentialBuckets(0.5, 2, 10), // 0.5s -> 4min
+	})
+
 	// Map API enum values to Kubernetes DeletionPropagation values
 	propagationPolicyMap = map[apiv2beta1.DeletePropagationPolicy]v1.DeletionPropagation{
 		apiv2beta1.DeletePropagationPolicy_FOREGROUND: v1.DeletePropagationForeground,
@@ -1286,8 +1294,8 @@ func (r *ResourceManager) DeleteJob(ctx context.Context, jobID string, propagati
 
 // Creates new tasks or updates existing ones.
 // This is not a part of internal API exposed to persistence agent only.
-func (r *ResourceManager) CreateOrUpdateTasks(t []*model.Task) ([]*model.Task, error) {
-	tasks, err := r.taskStore.CreateOrUpdateTasks(t)
+func (r *ResourceManager) CreateOrUpdateTasks(t []*model.Task, runID string) ([]*model.Task, error) {
+	tasks, err := r.taskStore.CreateOrUpdateTasks(t, runID)
 	if err != nil {
 		return nil, util.Wrap(err, "Failed to create or update tasks")
 	}
@@ -1436,6 +1444,10 @@ func (r *ResourceManager) ReportWorkflowResource(ctx context.Context, execSpec u
 			},
 		}
 		run, err = r.runStore.CreateRun(run)
+		if r.options.CollectMetrics && !execStatus.StartedAtTime().Time.IsZero() {
+			reportGap := time.Since(execStatus.StartedAtTime().Time).Seconds()
+			recurringPipelineRunReportGap.Observe(reportGap)
+		}
 		if err != nil {
 			return nil, util.Wrapf(err, "Failed to report a workflow due to error creating run %s", runId)
 		} else {

backend/src/apiserver/server/report_server.go

@@ -57,7 +57,7 @@ func (s *BaseReportServer) reportTasksFromExecution(execSpec util.ExecutionSpec,
 	if err != nil {
 		return nil, util.Wrap(err, "Failed to report tasks of an execution")
 	}
-	return s.resourceManager.CreateOrUpdateTasks(tasks)
+	return s.resourceManager.CreateOrUpdateTasks(tasks, runId)
 }
 
 // Reports a workflow.

backend/src/apiserver/storage/task_store.go

@@ -61,7 +61,7 @@ type TaskStoreInterface interface {
 	ListTasks(filterContext *model.FilterContext, opts *list.Options) ([]*model.Task, int, string, error)
 
 	// Creates new tasks or updates the existing ones.
-	CreateOrUpdateTasks(tasks []*model.Task) ([]*model.Task, error)
+	CreateOrUpdateTasks(tasks []*model.Task, runID string) ([]*model.Task, error)
 }
 
 type TaskStore struct {
@@ -324,15 +324,15 @@ func (s *TaskStore) GetTask(id string) (*model.Task, error) {
 }
 
 // Updates missing fields with existing data entries.
-func (s *TaskStore) patchWithExistingTasks(tasks []*model.Task) error {
+func (s *TaskStore) patchWithExistingTasks(tasks []*model.Task, runID string) error {
 	var podNames []string
 	for _, task := range tasks {
 		podNames = append(podNames, task.PodName)
 	}
 	sql, args, err := sq.
 		Select(taskColumns...).
 		From("tasks").
-		Where(sq.Eq{"PodName": podNames}).
+		Where(sq.Eq{"PodName": podNames, "RunUUID": runID}).
 		ToSql()
 	if err != nil {
 		return util.NewInternalServerError(err, "Failed to create query to check existing tasks")
@@ -359,7 +359,7 @@ func (s *TaskStore) patchWithExistingTasks(tasks []*model.Task) error {
 }
 
 // Creates new entries or updates existing ones.
-func (s *TaskStore) CreateOrUpdateTasks(tasks []*model.Task) ([]*model.Task, error) {
+func (s *TaskStore) CreateOrUpdateTasks(tasks []*model.Task, runID string) ([]*model.Task, error) {
 	buildQuery := func(ts []*model.Task) (string, []interface{}, error) {
 		sqlInsert := sq.Insert("tasks").Columns(taskColumnsWithPayload...)
 		for _, t := range ts {
@@ -405,7 +405,7 @@ func (s *TaskStore) CreateOrUpdateTasks(tasks []*model.Task) ([]*model.Task, err
 
 	// Check for existing tasks and fill empty field with existing data.
 	// Assumes that PodName column is a unique key.
-	if err := s.patchWithExistingTasks(tasks); err != nil {
+	if err := s.patchWithExistingTasks(tasks, runID); err != nil {
 		return nil, util.NewInternalServerError(err, "Failed to check for existing tasks")
 	}
 	for _, task := range tasks {

backend/src/apiserver/storage/task_store_test.go

@@ -504,7 +504,7 @@ func TestTaskStore_patchWithExistingTasks(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			err := taskStore.patchWithExistingTasks(tt.tasks)
+			err := taskStore.patchWithExistingTasks(tt.tasks, defaultFakeRunIdTwo)
 			if tt.wantErr {
 				assert.NotNil(t, err)
 				assert.Contains(t, err.Error(), tt.errMsg)
@@ -611,7 +611,7 @@ func TestTaskStore_UpdateOrCreateTasks(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			got, err := taskStore.CreateOrUpdateTasks(tt.tasks)
+			got, err := taskStore.CreateOrUpdateTasks(tt.tasks, defaultFakeRunIdTwo)
 			assert.Nil(t, err)
 			assert.Equal(t, tt.want, got)
 		})

go.mod

@@ -74,6 +74,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/credentials v1.17.67
 	github.com/aws/aws-sdk-go-v2/service/s3 v1.58.3
 	github.com/aws/smithy-go v1.22.3
+	github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0
 	gorm.io/driver/mysql v1.6.0
 	gorm.io/driver/postgres v1.6.0
 	gorm.io/driver/sqlite v1.6.0