Skip to content

Commit efbe53a

Browse files
shimibshimib
authored
feat: add queue depth, in-flight, and broker backlog metrics (llm-d#244)
* feat: add queue depth, in-flight, and broker backlog metrics Signed-off-by: Shimi Bandiel <shimib@google.com> * align PubSub metric labels and harden backlog reporting Signed-off-by: Shimi Bandiel <shimib@google.com> * reset metric on error Signed-off-by: Shimi Bandiel <shimib@google.com> --------- Signed-off-by: Shimi Bandiel <shimib@google.com>
1 parent 9b81d08 commit efbe53a

14 files changed

Lines changed: 558 additions & 87 deletions

File tree

‎cmd/main.go‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ func main() {
4949
var redisTracing bool
5050

5151
var drainTimeout time.Duration
52+
var backlogPollInterval time.Duration
5253

5354
var tlsCACert string
5455
var tlsCert string
@@ -64,6 +65,7 @@ func main() {
6465
flag.IntVar(&concurrency, "concurrency", 8, "number of concurrent workers")
6566
flag.DurationVar(&requestTimeout, "request-timeout", 5*time.Minute, "timeout for individual inference requests")
6667
flag.DurationVar(&drainTimeout, "drain-timeout", 2*time.Minute, "maximum time to wait for in-flight requests to complete after SIGTERM")
68+
flag.DurationVar(&backlogPollInterval, "metrics-backlog-poll-interval", 15*time.Second, "interval to poll the broker for queue backlog metrics (0 disables); only applies to flows that support it (redis-sortedset, gcp-pubsub)")
6769

6870
flag.StringVar(&requestMergePolicy, "request-merge-policy", "random-robin", "The request merge policy to use. Supported policies: random-robin")
6971
flag.StringVar(&messageQueueImpl, "message-queue-impl", "redis-pubsub", "The message queue implementation to use. Supported implementations: redis-pubsub, redis-sortedset, gcp-pubsub, gcp-pubsub-gated")
@@ -233,6 +235,12 @@ func main() {
233235
impl.Start(signalCtx)
234236
healthServer.SetReady()
235237

238+
if reporter, ok := impl.(pipeline.BacklogReporter); ok && backlogPollInterval > 0 {
239+
go pollBacklog(signalCtx, reporter, backlogPollInterval)
240+
} else if !ok {
241+
setupLog.Info("Selected flow does not support broker backlog metrics", "message-queue-impl", messageQueueImpl)
242+
}
243+
236244
<-signalCtx.Done()
237245
healthServer.SetNotReady()
238246

@@ -301,6 +309,34 @@ func buildTLSConfig(caCertPath, certPath, keyPath string, insecureSkipVerify boo
301309
return tlsConfig, nil
302310
}
303311

312+
// pollBacklog periodically queries the broker for queue backlog and publishes
313+
// it as the async_broker_backlog gauge until ctx is cancelled.
314+
func pollBacklog(ctx context.Context, reporter pipeline.BacklogReporter, interval time.Duration) {
315+
logger := ctrl.Log.WithName("backlog-poller")
316+
ticker := time.NewTicker(interval)
317+
defer ticker.Stop()
318+
319+
poll := func() {
320+
stats, err := reporter.QueueBacklog(ctx)
321+
if err != nil {
322+
logger.V(logging.DEFAULT).Error(err, "Failed to poll broker backlog")
323+
}
324+
for _, s := range stats {
325+
metrics.SetBrokerBacklog(s.QueueID, s.QueueName, float64(s.Depth))
326+
}
327+
}
328+
329+
poll()
330+
for {
331+
select {
332+
case <-ctx.Done():
333+
return
334+
case <-ticker.C:
335+
poll()
336+
}
337+
}
338+
}
339+
304340
func printAllFlags(setupLog logr.Logger) {
305341
flags := make(map[string]any)
306342
flag.VisitAll(func(f *flag.Flag) {

‎go.mod‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ replace github.com/llm-d-incubation/llm-d-async/api => ./api
1111
replace github.com/llm-d-incubation/llm-d-async/pipeline => ./pipeline
1212

1313
require (
14+
cloud.google.com/go/monitoring v1.29.0
1415
cloud.google.com/go/pubsub/v2 v2.6.0
1516
github.com/alicebob/miniredis/v2 v2.38.0
1617
github.com/go-logr/logr v1.4.3
@@ -28,6 +29,8 @@ require (
2829
go.opentelemetry.io/otel/trace v1.44.0
2930
go.uber.org/zap v1.28.0
3031
golang.org/x/oauth2 v0.36.0
32+
google.golang.org/api v0.274.0
33+
google.golang.org/protobuf v1.36.11
3134
k8s.io/client-go v0.34.8
3235
sigs.k8s.io/controller-runtime v0.22.5
3336
sigs.k8s.io/gateway-api-inference-extension v1.2.1
@@ -117,12 +120,10 @@ require (
117120
golang.org/x/time v0.15.0 // indirect
118121
golang.org/x/tools v0.44.0 // indirect
119122
gomodules.xyz/jsonpatch/v2 v2.4.0 // indirect
120-
google.golang.org/api v0.274.0 // indirect
121123
google.golang.org/genproto v0.0.0-20260319201613-d00831a3d3e7 // indirect
122124
google.golang.org/genproto/googleapis/api v0.0.0-20260526163538-3dc84a4a5aaa // indirect
123125
google.golang.org/genproto/googleapis/rpc v0.0.0-20260526163538-3dc84a4a5aaa // indirect
124126
google.golang.org/grpc v1.81.1 // indirect
125-
google.golang.org/protobuf v1.36.11 // indirect
126127
gopkg.in/evanphx/json-patch.v4 v4.13.0 // indirect
127128
gopkg.in/inf.v0 v0.9.1 // indirect
128129
gopkg.in/yaml.v3 v3.0.1 // indirect

‎go.sum‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@ cloud.google.com/go/compute/metadata v0.9.0 h1:pDUj4QMoPejqq20dK0Pg2N4yG9zIkYGdB
1111
cloud.google.com/go/compute/metadata v0.9.0/go.mod h1:E0bWwX5wTnLPedCKqk3pJmVgCBSM6qQI1yTBdEb3C10=
1212
cloud.google.com/go/iam v1.7.0 h1:JD3zh0C6LHl16aCn5Akff0+GELdp1+4hmh6ndoFLl8U=
1313
cloud.google.com/go/iam v1.7.0/go.mod h1:tetWZW1PD/m6vcuY2Zj/aU0eCHNPuxedbnbRTyKXvdY=
14+
cloud.google.com/go/monitoring v1.29.0 h1:AHhDsFaSax1/4k+qlIDX/SDGe6hggnfXJ9dkgD9qBPY=
15+
cloud.google.com/go/monitoring v1.29.0/go.mod h1:72NOVjJXHY/HBfoLT0+qlCZBT059+9VXLeAnL2PeeVM=
1416
cloud.google.com/go/pubsub/v2 v2.6.0 h1:8pjR0id+GTB+krKx5G6AGJoYrHog58w2Q89PCOrfM64=
1517
cloud.google.com/go/pubsub/v2 v2.6.0/go.mod h1:4anqvV/w8Pcgu2tO0qr2XgsF3GXHowzryfQ5gOnVmWY=
1618
github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=

‎pipeline/pipeline.go‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,21 @@ type RequestMergePolicy interface {
7878
MergeRequestChannels(channels []RequestChannel) EmbelishedRequestChannel
7979
}
8080

81+
// QueueBacklogStat reports the broker-side backlog for a single queue.
82+
type QueueBacklogStat struct {
83+
QueueID string
84+
QueueName string
85+
Depth int64
86+
}
87+
88+
// BacklogReporter is an optional capability for flows backed by a broker that
89+
// exposes a queryable backlog (e.g. Redis sorted sets, GCP PubSub). Flows whose
90+
// transport has no persisted queue (e.g. Redis Pub/Sub) do not implement it.
91+
type BacklogReporter interface {
92+
// QueueBacklog returns the current backlog for each configured queue.
93+
QueueBacklog(ctx context.Context) ([]QueueBacklogStat, error)
94+
}
95+
8196
type RequestChannel struct {
8297
Channel chan *api.InternalRequest
8398
IGWBaseURL string

‎pkg/async/random_robin_policy.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66

77
"github.com/llm-d-incubation/llm-d-async/api"
88
"github.com/llm-d-incubation/llm-d-async/pipeline"
9+
"github.com/llm-d-incubation/llm-d-async/pkg/metrics"
910
)
1011

1112
func NewRandomRobinPolicy() pipeline.RequestMergePolicy {
@@ -67,6 +68,7 @@ func (r *RandomRobinPolicy) MergeRequestChannels(channels []pipeline.RequestChan
6768
HttpHeaders: headers,
6869
RequestURL: requestURL,
6970
}
71+
metrics.IncQueueDepth(ir.QueueID, ir.RequestQueueName)
7072
mergedChannel <- erm
7173
}
7274

‎pkg/asyncworker/worker.go‎

Lines changed: 84 additions & 75 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ func Worker(consumeCtx, requestCtx context.Context, characteristics pipeline.Cha
4646
if msg.InternalRequest == nil || msg.PublicRequest == nil {
4747
continue
4848
}
49+
metrics.DecQueueDepth(msg.QueueID, msg.RequestQueueName)
4950
retryMsg := pipeline.RetryMessage{
5051
EmbelishedRequestMessage: msg,
5152
BackoffDurationSeconds: 0,
@@ -67,97 +68,105 @@ func Worker(consumeCtx, requestCtx context.Context, characteristics pipeline.Cha
6768
}
6869
queueID := msg.QueueID
6970
queueName := msg.RequestQueueName
70-
if msg.RetryCount == 0 {
71-
metrics.RecordAsyncReq(queueID, queueName)
72-
}
73-
payloadBytes := validateAndMarshal(requestCtx, resultChannel, msg)
74-
if payloadBytes == nil {
75-
continue
76-
}
71+
metrics.DecQueueDepth(queueID, queueName)
7772

78-
sendInferenceRequest := func() {
79-
reqCtx := requestCtx
80-
if md := msg.PublicRequest.ReqMetadata(); len(md) > 0 {
81-
reqCtx = otel.GetTextMapPropagator().Extract(reqCtx, propagation.MapCarrier(md))
82-
}
73+
processMessage := func() {
74+
metrics.IncInflight(queueID, queueName)
75+
defer metrics.DecInflight(queueID, queueName)
8376

84-
spanAttrs := []attribute.KeyValue{
85-
attribute.String(uotel.AttrRequestID, msg.PublicRequest.ReqID()),
86-
attribute.Int(uotel.AttrRetryCount, msg.RetryCount),
77+
if msg.RetryCount == 0 {
78+
metrics.RecordAsyncReq(queueID, queueName)
8779
}
88-
if queueID != "" {
89-
spanAttrs = append(spanAttrs, attribute.String(uotel.AttrQueueID, queueID))
90-
}
91-
if queueName != "" {
92-
spanAttrs = append(spanAttrs, attribute.String(uotel.AttrQueueName, queueName))
80+
payloadBytes := validateAndMarshal(requestCtx, resultChannel, msg)
81+
if payloadBytes == nil {
82+
return
9383
}
94-
reqCtx, span := uotel.StartSpan(reqCtx, "process-request",
95-
trace.WithAttributes(spanAttrs...),
96-
)
97-
defer span.End()
9884

99-
reqDeadline := time.Now().Add(requestTimeout)
100-
if dline := msg.PublicRequest.ReqDeadline(); dline > 0 {
101-
if msgDeadline := time.Unix(dline, 0); msgDeadline.Before(reqDeadline) {
102-
reqDeadline = msgDeadline
85+
sendInferenceRequest := func() {
86+
reqCtx := requestCtx
87+
if md := msg.PublicRequest.ReqMetadata(); len(md) > 0 {
88+
reqCtx = otel.GetTextMapPropagator().Extract(reqCtx, propagation.MapCarrier(md))
10389
}
104-
}
105-
reqCtx, cancel := context.WithDeadline(reqCtx, reqDeadline)
106-
defer cancel()
10790

108-
logger.V(logutil.DEBUG).Info("Sending inference request", "url", msg.RequestURL)
109-
responseBody, err := client.SendRequest(reqCtx, msg.RequestURL, msg.HttpHeaders, payloadBytes)
91+
spanAttrs := []attribute.KeyValue{
92+
attribute.String(uotel.AttrRequestID, msg.PublicRequest.ReqID()),
93+
attribute.Int(uotel.AttrRetryCount, msg.RetryCount),
94+
}
95+
if queueID != "" {
96+
spanAttrs = append(spanAttrs, attribute.String(uotel.AttrQueueID, queueID))
97+
}
98+
if queueName != "" {
99+
spanAttrs = append(spanAttrs, attribute.String(uotel.AttrQueueName, queueName))
100+
}
101+
reqCtx, span := uotel.StartSpan(reqCtx, "process-request",
102+
trace.WithAttributes(spanAttrs...),
103+
)
104+
defer span.End()
110105

111-
if err == nil {
112-
metrics.RecordSuccessfulReq(queueID, queueName)
113-
select {
114-
case resultChannel <- asyncapi.ResultMessage{
115-
ID: msg.PublicRequest.ReqID(),
116-
Payload: string(responseBody),
117-
Routing: msg.InternalRouting,
118-
Metadata: msg.PublicRequest.ReqMetadata(),
119-
}:
120-
case <-requestCtx.Done():
106+
reqDeadline := time.Now().Add(requestTimeout)
107+
if dline := msg.PublicRequest.ReqDeadline(); dline > 0 {
108+
if msgDeadline := time.Unix(dline, 0); msgDeadline.Before(reqDeadline) {
109+
reqDeadline = msgDeadline
110+
}
121111
}
122-
return
123-
}
112+
reqCtx, cancel := context.WithDeadline(reqCtx, reqDeadline)
113+
defer cancel()
124114

125-
if requestCtx.Err() != nil && (errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded)) {
126-
_, bgSpan := uotel.DetachedContext(reqCtx, "re-enqueue")
127-
bgSpan.SetAttributes(attribute.String(uotel.AttrRequestID, msg.PublicRequest.ReqID()))
128-
defer bgSpan.End()
129-
retryChannel <- pipeline.RetryMessage{
130-
EmbelishedRequestMessage: msg,
131-
BackoffDurationSeconds: 0,
115+
logger.V(logutil.DEBUG).Info("Sending inference request", "url", msg.RequestURL)
116+
responseBody, err := client.SendRequest(reqCtx, msg.RequestURL, msg.HttpHeaders, payloadBytes)
117+
118+
if err == nil {
119+
metrics.RecordSuccessfulReq(queueID, queueName)
120+
select {
121+
case resultChannel <- asyncapi.ResultMessage{
122+
ID: msg.PublicRequest.ReqID(),
123+
Payload: string(responseBody),
124+
Routing: msg.InternalRouting,
125+
Metadata: msg.PublicRequest.ReqMetadata(),
126+
}:
127+
case <-requestCtx.Done():
128+
}
129+
return
132130
}
133-
return
134-
}
135131

136-
var inferenceErr asyncapi.InferenceError
137-
if !errors.As(err, &inferenceErr) || inferenceErr.Category().Fatal() {
138-
span.RecordError(err)
139-
span.SetStatus(codes.Error, "inference request failed")
140-
span.SetAttributes(attribute.String(uotel.AttrErrorCategory, inferenceErrorCategory(err)))
141-
metrics.RecordFailedReq(queueID, queueName)
142-
select {
143-
case resultChannel <- CreateErrorResultMessage(msg.PublicRequest, msg.InternalRouting, fmt.Sprintf("Failed to send request to inference: %s", err.Error())):
144-
case <-requestCtx.Done():
132+
if requestCtx.Err() != nil && (errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded)) {
133+
_, bgSpan := uotel.DetachedContext(reqCtx, "re-enqueue")
134+
bgSpan.SetAttributes(attribute.String(uotel.AttrRequestID, msg.PublicRequest.ReqID()))
135+
defer bgSpan.End()
136+
retryChannel <- pipeline.RetryMessage{
137+
EmbelishedRequestMessage: msg,
138+
BackoffDurationSeconds: 0,
139+
}
140+
return
145141
}
146-
return
147-
}
148142

149-
if inferenceErr.Category().Sheddable() {
150-
metrics.RecordSheddedReq(queueID, queueName)
151-
}
152-
span.SetAttributes(attribute.String(uotel.AttrErrorCategory, string(inferenceErr.Category())))
153-
var retryAfter time.Duration
154-
var clientErr *asyncapi.ClientError
155-
if errors.As(err, &clientErr) {
156-
retryAfter = clientErr.RetryAfter
143+
var inferenceErr asyncapi.InferenceError
144+
if !errors.As(err, &inferenceErr) || inferenceErr.Category().Fatal() {
145+
span.RecordError(err)
146+
span.SetStatus(codes.Error, "inference request failed")
147+
span.SetAttributes(attribute.String(uotel.AttrErrorCategory, inferenceErrorCategory(err)))
148+
metrics.RecordFailedReq(queueID, queueName)
149+
select {
150+
case resultChannel <- CreateErrorResultMessage(msg.PublicRequest, msg.InternalRouting, fmt.Sprintf("Failed to send request to inference: %s", err.Error())):
151+
case <-requestCtx.Done():
152+
}
153+
return
154+
}
155+
156+
if inferenceErr.Category().Sheddable() {
157+
metrics.RecordSheddedReq(queueID, queueName)
158+
}
159+
span.SetAttributes(attribute.String(uotel.AttrErrorCategory, string(inferenceErr.Category())))
160+
var retryAfter time.Duration
161+
var clientErr *asyncapi.ClientError
162+
if errors.As(err, &clientErr) {
163+
retryAfter = clientErr.RetryAfter
164+
}
165+
retryMessage(requestCtx, msg, retryChannel, resultChannel, retryAfter)
157166
}
158-
retryMessage(requestCtx, msg, retryChannel, resultChannel, retryAfter)
167+
sendInferenceRequest()
159168
}
160-
sendInferenceRequest()
169+
processMessage()
161170
}
162171
}
163172
}

0 commit comments

Comments
 (0)