Repository: numaproj/numaflow
Pull request #123 (2022-07-28 23:18:14)

feat: autoscaling with back pressure factor
Signed-off-by: Derek Wang <whynowy@gmail.com>

Consider back pressure factor into autoscaling.

1. Back pressure calculation is based on `avg-pending/available-buffer-length > threshold`, threshold defaults to 90%, available-buffer-length = buffer-length * buffer-limit;
2. When calculating desired replicas, if it's more than current replicas:
  a. If there's back pressure from directly connected vertices, instead of increasing the replicas, decrease by 1;
  b. If the back pressure is from downstream not directly connected vertices, no scaling up or down will happen.

<!--

Before you push your changes:

* Run `make pre-push -B` to fix codegen and lint problems.

Then, you MUST:

* Sign-off your commit (otherwise the DCO check will fail).
* Use [a conventional commit message](https://www.conventionalcommits.org/en/v1.0.0/) (otherwise the commit message check will fail).

If you did not do this, reset all your commit and replace them with a single commit:

```
git reset HEAD~1 ;# change 1 to how many commits you made
git commit --signoff -m 'feat: my feat. Fixes #1234'
```

When creating your PR: 

* Make sure that "Fixes #" or "Closes #" is in both the PR title (for release notes) and description (to automatically link and close the issue).
* Say how you tested your changes. If you changed the UI, attach screenshots.
* Set your PR as a draft initially.
* Your PR needs to pass the required checks before it can be approved. 
* Once required tests have passed, mark your PR "Ready for review".

If changes were requested, once you've made them, you MUST dismiss the review to get it reviewed again.

-->



Patch:
diff --git a/controllers/vertex/controller.go b/controllers/vertex/controller.go
index c4039abf..1b9642c0 100644
--- a/controllers/vertex/controller.go
+++ b/controllers/vertex/controller.go
@@ -69,9 +69,10 @@ func (r *vertexReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctr
 // reconcile does the real logic
 func (r *vertexReconciler) reconcile(ctx context.Context, vertex *dfv1.Vertex) (ctrl.Result, error) {
 	log := logging.FromContext(ctx)
+	vertexKey := scaling.KeyOfVertex(*vertex)
 	if !vertex.DeletionTimestamp.IsZero() {
 		log.Info("Deleting vertex")
-		r.scaler.StopWatching(fmt.Sprintf("%s/%s", vertex.Namespace, vertex.Name))
+		r.scaler.StopWatching(vertexKey)
 		return ctrl.Result{}, nil
 	}
 
@@ -98,7 +99,7 @@ func (r *vertexReconciler) reconcile(ctx context.Context, vertex *dfv1.Vertex) (
 	}
 
 	if vertex.Scalable() { // Add to autoscaling watcher
-		r.scaler.StartWatching(fmt.Sprintf("%s/%s", vertex.Namespace, vertex.Name))
+		r.scaler.StartWatching(vertexKey)
 	}
 
 	desiredReplicas := vertex.Spec.GetReplicas()
diff --git a/controllers/vertex/scaling/options.go b/controllers/vertex/scaling/options.go
index 205b8e60..38011ce0 100644
--- a/controllers/vertex/scaling/options.go
+++ b/controllers/vertex/scaling/options.go
@@ -1,18 +1,21 @@
 package scaling
 
 type options struct {
-	// Number of workers working on autoscaling
+	// Number of workers working on autoscaling.
 	workers int
 	// Time in milliseconds, each element in the work queue will be picked up in an interval of this period of time.
 	taskInterval int
+	// Threshold of considering there's back pressure, a float value less than 1.
+	backPressureThreshold float64
 }
 
 type Option func(*options)
 
 func defaultOptions() *options {
 	return &options{
-		workers:      20,
-		taskInterval: 30000,
+		workers:               20,
+		taskInterval:          30000,
+		backPressureThreshold: 0.9,
 	}
 }
 
@@ -21,3 +24,15 @@ func WithWorkers(n int) Option {
 		o.workers = n
 	}
 }
+
+func WithTaskInterval(n int) Option {
+	return func(o *options) {
+		o.taskInterval = n
+	}
+}
+
+func WithBackPressureThreshold(n float64) Option {
+	return func(o *options) {
+		o.backPressureThreshold = n
+	}
+}
diff --git a/controllers/vertex/scaling/scaling.go b/controllers/vertex/scaling/scaling.go
index 90dcfc6b..b5a34703 100644
--- a/controllers/vertex/scaling/scaling.go
+++ b/controllers/vertex/scaling/scaling.go
@@ -10,6 +10,7 @@ import (
 	"sync"
 	"time"
 
+	lru "github.com/hashicorp/golang-lru"
 	"go.uber.org/zap"
 	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/apimachinery/pkg/types"
@@ -29,8 +30,11 @@ type Scaler struct {
 	vertexList *list.List
 	lock       *sync.RWMutex
 	options    *options
+	// Cache to store the vertex metrics such as pending message number
+	vertexMetricsCache *lru.Cache
 }
 
+// NewScaler returns a Scaler instance.
 func NewScaler(client client.Client, opts ...Option) *Scaler {
 	s := &Scaler{
 		client:     client,
@@ -39,6 +43,8 @@ func NewScaler(client client.Client, opts ...Option) *Scaler {
 		vertexList: list.New(),
 		lock:       new(sync.RWMutex),
 	}
+	vertexMetricsCache, _ := lru.New(10000)
+	s.vertexMetricsCache = vertexMetricsCache
 	for _, opt := range opts {
 		if opt != nil {
 			opt(s.options)
@@ -107,6 +113,11 @@ func (s *Scaler) scale(ctx context.Context, id int, keyCh <-chan string) {
 // For UDF and sinks which have the read buffer information
 //    singleReplicaContribution = (totalAvailableBufferLength - pending) / currentReplicas
 //    desiredReplicas = targetAvailableBufferLength / singleReplicaContribution
+//
+// Back pressure factor
+// When desiredReplicas > currentReplics:
+// If there's back pressure in the directly connectied vertices, desiredReplicas = currentReplicas-1;
+// If there's back pressure in the downstream vertices (not connected), desiredReplicas remains the same.
 func (s *Scaler) scaleOneVertex(ctx context.Context, key string, worker int) error {
 	log := logging.FromContext(ctx).With("worker", fmt.Sprint(worker)).With("vertexKey", key)
 	log.Debugf("Working on key: %s", key)
@@ -190,6 +201,8 @@ func (s *Scaler) scaleOneVertex(ctx context.Context, key string, worker int) err
 		log.Debugf("Vertex %s has no pending messages information, skip scaling", vertex.Name)
 		return nil
 	}
+	// Add to cache for back pressure calculation
+	_ = s.vertexMetricsCache.Add(key+"/pending", pending)
 	totalBufferLength := int64(0)
 	targetAvailableBufferLength := int64(0)
 	if !vertex.IsASource() { // Only non-source vertex has buffer to read
@@ -202,12 +215,13 @@ func (s *Scaler) scaleOneVertex(ctx context.Context, key string, worker int) err
 			}
 			totalBufferLength = int64(float64(*bInfo.BufferLength) * *bInfo.BufferUsageLimit)
 			targetAvailableBufferLength = int64(float64(*bInfo.BufferLength) * float64(vertex.Spec.Scale.GetTargetBufferUsage()) / 100)
+			// Add to cache for back pressure calculation
+			_ = s.vertexMetricsCache.Add(*bInfo.BufferName+"/length", totalBufferLength)
 		}
 	}
-
+	current := int32(vertex.Spec.GetReplicas())
 	desired := s.desiredReplicas(ctx, vertex, rate, pending, totalBufferLength, targetAvailableBufferLength)
 	log.Debugf("Calculated desired replica number of vertex %q is: %v", vertex.Name, desired)
-	current := int32(vertex.Spec.GetReplicas())
 	max := vertex.Spec.Scale.GetMaxReplicas()
 	min := vertex.Spec.Scale.GetMinReplicas()
 	if desired > max {
@@ -228,6 +242,20 @@ func (s *Scaler) scaleOneVertex(ctx context.Context, key string, worker int) err
 		return s.patchVertexReplicas(ctx, vertex, current-diff) // We scale down gradually
 	}
 	if desired > current {
+		// When scaling up, need to check back pressure
+		directPressure, downstreamPressure := s.hasBackPressure(*pl, *vertex)
+		if directPressure {
+			if current > 1 {
+				log.Debugf("Vertex %s has direct back pressure from connected vertices, decreasing one replica")
+				return s.patchVertexReplicas(ctx, vertex, current-1)
+			} else {
+				log.Debugf("Vertex %s has direct back pressure from connected vertices, skip scaling")
+				return nil
+			}
+		} else if downstreamPressure {
+			log.Debugf("Vertex %s has back pressure in downstream vertices, skip scaling")
+			return nil
+		}
 		diff := desired - current
 		if diff > maxAllowed {
 			diff = maxAllowed
@@ -323,6 +351,37 @@ func (s *Scaler) Start(ctx context.Context) {
 	}
 }
 
+// hasBackPressure checks if there's back pressure in the downstream buffers
+// It returns if 2 bool values, which represent:
+// 1. If there's back pressure in the connected vertices;
+// 2. If there's back pressure in any of the downstream vertices.
+func (s *Scaler) hasBackPressure(pl dfv1.Pipeline, vertex dfv1.Vertex) (bool, bool) {
+	downstreamEdges := pl.GetDownstreamEdges(vertex.Spec.Name)
+	directPressure, downstreamPressure := false, false
+	for _, e := range downstreamEdges {
+		vertexKey := pl.Namespace + "/" + pl.Name + "-" + e.To
+		bufferName := dfv1.GenerateEdgeBufferName(pl.Namespace, pl.Name, e.From, e.To)
+		pendingVal, ok := s.vertexMetricsCache.Get(vertexKey + "/pending")
+		if !ok { // Vertex key has not been cached, skip it.
+			continue
+		}
+		pending := pendingVal.(int64)
+		bufferLengthVal, ok := s.vertexMetricsCache.Get(bufferName + "/length")
+		if !ok { // Buffer length has not been cached, skip it.
+			continue
+		}
+		length := bufferLengthVal.(int64)
+		if float64(pending)/float64(length) >= s.options.backPressureThreshold {
+			downstreamPressure = true
+			if e.From == vertex.Spec.Name {
+				directPressure = true
+				break
+			}
+		}
+	}
+	return directPressure, downstreamPressure
+}
+
 func (s *Scaler) patchVertexReplicas(ctx context.Context, vertex *dfv1.Vertex, desiredReplicas int32) error {
 	log := logging.FromContext(ctx)
 	origin := vertex.Spec.Replicas
@@ -337,3 +396,8 @@ func (s *Scaler) patchVertexReplicas(ctx context.Context, vertex *dfv1.Vertex, d
 	log.Infow("Auto scaling - vertex replicas changed", zap.Int32p("from", origin), zap.Int32("to", desiredReplicas), zap.String("pipeline", vertex.Spec.PipelineName), zap.String("vertex", vertex.Spec.Name))
 	return nil
 }
+
+// KeyOfVertex returns the unique key of a vertex
+func KeyOfVertex(vertex dfv1.Vertex) string {
+	return fmt.Sprintf("%s/%s", vertex.Namespace, vertex.Name)
+}
diff --git a/docs/quick-start.md b/docs/quick-start.md
index 247768a1..7c03195a 100644
--- a/docs/quick-start.md
+++ b/docs/quick-start.md
@@ -73,4 +73,4 @@ kubectl delete -f https://raw.githubusercontent.com/numaproj/numaflow/stable/exa
 
 ## Coming Up
 
-Next, let's take a peek at [an advance pipeline](./advanced-start.md), to learn some powerful features of Numaflow.
+Next, let's take a peek at [an advanced pipeline](./advanced-start.md), to learn some powerful features of Numaflow.
diff --git a/go.mod b/go.mod
index 43b0d36d..f2af46f5 100644
--- a/go.mod
+++ b/go.mod
@@ -20,6 +20,7 @@ require (
 	github.com/grpc-ecosystem/go-grpc-middleware v1.3.0
 	github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0
 	github.com/grpc-ecosystem/grpc-gateway v1.16.0
+	github.com/hashicorp/golang-lru v0.5.4
 	github.com/montanaflynn/stats v0.0.0-20171201202039-1bf9dbcd8cbe
 	github.com/nats-io/nats.go v1.16.0
 	github.com/numaproj/numaflow-go v0.0.1
diff --git a/go.sum b/go.sum
index 9c16c6da..3aa08240 100644
--- a/go.sum
+++ b/go.sum
@@ -520,6 +520,8 @@ github.com/hashicorp/go-uuid v1.0.2/go.mod h1:6SBZvOh/SIDV7/2o3Jml5SYk/TvGqwFJ/b
 github.com/hashicorp/go.net v0.0.1/go.mod h1:hjKkEWcCURg++eb33jQU7oqQcI9XDCnUzHA0oac0k90=
 github.com/hashicorp/golang-lru v0.5.0/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
 github.com/hashicorp/golang-lru v0.5.1/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
+github.com/hashicorp/golang-lru v0.5.4 h1:YDjusn29QI/Das2iO9M0BHnIbxPeyuCHsjMW+lJfyTc=
+github.com/hashicorp/golang-lru v0.5.4/go.mod h1:iADmTwqILo4mZ8BN3D2Q6+9jd8WM5uGBxy+E8yxSoD4=
 github.com/hashicorp/hcl v1.0.0 h1:0Anlzjpi4vEasTeNFn2mLJgTSwt0+6sfsiTG8qcWGx4=
 github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ=
 github.com/hashicorp/logutils v1.0.0/go.mod h1:QIAnNjmIWmVIIkWDTG1z5v++HQmx9WQRO+LraFDTW64=
diff --git a/pkg/apis/numaflow/v1alpha1/pipeline_types.go b/pkg/apis/numaflow/v1alpha1/pipeline_types.go
index f4e88725..5b8c52f2 100644
--- a/pkg/apis/numaflow/v1alpha1/pipeline_types.go
+++ b/pkg/apis/numaflow/v1alpha1/pipeline_types.go
@@ -122,6 +122,22 @@ func (p Pipeline) GetAllBuffers() []Buffer {
 	return r
 }
 
+// GetDownstreamEdges returns all the downstream edges of a vertex
+func (p Pipeline) GetDownstreamEdges(vertexName string) []Edge {
+	var f func(vertexName string, edges *[]Edge)
+	f = func(vertexName string, edges *[]Edge) {
+		for _, b := range p.Spec.Edges {
+			if b.From == vertexName {
+				*edges = append(*edges, b)
+				f(b.To, edges)
+			}
+		}
+	}
+	result := []Edge{}
+	f(vertexName, &result)
+	return result
+}
+
 func (p Pipeline) GetDaemonServiceName() string {
 	return fmt.Sprintf("%s-daemon-svc", p.Name)
 }
diff --git a/pkg/daemon/server/service/pipeline_watermark_query.go b/pkg/daemon/server/service/pipeline_watermark_query.go
index 6c0158b8..fd6f42b6 100644
--- a/pkg/daemon/server/service/pipeline_watermark_query.go
+++ b/pkg/daemon/server/service/pipeline_watermark_query.go
@@ -4,6 +4,8 @@ package service
 import (
 	"context"
 	"fmt"
+	"time"
+
 	"github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1"
 	"github.com/numaproj/numaflow/pkg/apis/proto/daemon"
 	"github.com/numaproj/numaflow/pkg/isbsvc"
@@ -13,7 +15,6 @@ import (
 	"github.com/numaproj/numaflow/pkg/watermark/generic"
 	"github.com/numaproj/numaflow/pkg/watermark/store/jetstream"
 	"go.uber.org/zap"
-	"time"
 )
 
 // watermarkFetchers used to store watermark metadata for propagation


Test patch:
diff --git a/pkg/apis/numaflow/v1alpha1/pipeline_types_test.go b/pkg/apis/numaflow/v1alpha1/pipeline_types_test.go
index 3aa04830..6b91aef0 100644
--- a/pkg/apis/numaflow/v1alpha1/pipeline_types_test.go
+++ b/pkg/apis/numaflow/v1alpha1/pipeline_types_test.go
@@ -188,3 +188,33 @@ func Test_PipelineMarkPhases(t *testing.T) {
 	s.MarkPhaseRunning()
 	assert.Equal(t, PipelinePhaseRunning, s.Phase)
 }
+
+func Test_GetDownstreamEdges(t *testing.T) {
+	pl := Pipeline{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "test-pl",
+			Namespace: "test-ns",
+		},
+		Spec: PipelineSpec{
+			Edges: []Edge{
+				{From: "input", To: "p1"},
+				{From: "p1", To: "p11"},
+				{From: "p1", To: "p2"},
+				{From: "p2", To: "output"},
+			},
+		},
+	}
+	edges := pl.GetDownstreamEdges("input")
+	assert.Equal(t, 4, len(edges))
+	assert.Equal(t, edges, pl.Spec.Edges)
+	assert.Equal(t, edges[2], Edge{From: "p1", To: "p2"})
+
+	edges = pl.GetDownstreamEdges("p1")
+	assert.Equal(t, 3, len(edges))
+
+	edges = pl.GetDownstreamEdges("output")
+	assert.Equal(t, 0, len(edges))
+
+	edges = pl.GetDownstreamEdges("notexisting")
+	assert.Equal(t, 0, len(edges))
+}
diff --git a/test/e2e/functional_test.go b/test/e2e/functional_test.go
index 7b8334ba..e7ec3d3c 100644
--- a/test/e2e/functional_test.go
+++ b/test/e2e/functional_test.go
@@ -80,7 +80,7 @@ func (s *FunctionalSuite) TestCreateSimplePipeline() {
 	assert.Equal(s.T(), "input", *bufferInfo.FromVertex)
 	m, err := client.GetVertexMetrics(context.Background(), pipelineName, "p1")
 	assert.NoError(s.T(), err)
-	assert.Equal(s.T(), pipelineName, m.Pipeline)
+	assert.Equal(s.T(), pipelineName, *m.Pipeline)
 }
 
 func (s *FunctionalSuite) TestFiltering() {