From 7c5a18c987fa51b4dd97589c39fed27cb7ba6ea2 Mon Sep 17 00:00:00 2001
From: Jiaxin Shan <seedjeffwan@gmail.com>
Date: Mon, 20 Jan 2025 23:22:57 -0800
Subject: [PATCH] Check the HPA ownerReference in request enqueue (#582)

* Update runtime api with guidance on nightly install

* Update PodAutoscaler to only enqueue managed objects

---------

Signed-off-by: Jiaxin Shan <seedjeffwan@gmail.com>
---
 development/tutorials/podautoscaler/pa.yaml   | 13 ++++--
 docs/source/features/runtime.rst              |  4 ++
 .../podautoscaler/podautoscaler_controller.go | 40 ++++++++++++++++---
 3 files changed, 47 insertions(+), 10 deletions(-)

diff --git a/development/tutorials/podautoscaler/pa.yaml b/development/tutorials/podautoscaler/pa.yaml
index 47bb8f2e..ce1a1598 100644
--- a/development/tutorials/podautoscaler/pa.yaml
+++ b/development/tutorials/podautoscaler/pa.yaml
@@ -10,9 +10,14 @@ spec:
   scaleTargetRef:
     apiVersion: apps/v1
     kind: Deployment
-    name: llama2-70b
+    name: mock-llama2-7b
   minReplicas: 1
   maxReplicas: 10
-  targetMetric: "gpu_cache_usage_perc"
-  targetValue: "50"
-  scalingStrategy: "KPA"
+  metricsSources:
+  - metricSourceType: "pod"
+    protocolType: "http"
+    port: "8000"
+    path: "/metrics"
+    targetMetric: "gpu_cache_usage_perc"
+    targetValue: "40"
+  scalingStrategy: "HPA"
diff --git a/docs/source/features/runtime.rst b/docs/source/features/runtime.rst
index 2d30a5b7..8cc61899 100644
--- a/docs/source/features/runtime.rst
+++ b/docs/source/features/runtime.rst
@@ -13,6 +13,10 @@ The AI Runtime hides various implementation details on the inference engine side
 
     ``python3 -m pip install aibrix``
 
+    If you want to use nightly version, you can install from code.
+
+    ``python3 -m pip install -e .``
+
 
 Metric Standardization
 ----------------------
diff --git a/pkg/controller/podautoscaler/podautoscaler_controller.go b/pkg/controller/podautoscaler/podautoscaler_controller.go
index 2a9b145d..0362d458 100644
--- a/pkg/controller/podautoscaler/podautoscaler_controller.go
+++ b/pkg/controller/podautoscaler/podautoscaler_controller.go
@@ -21,13 +21,13 @@ import (
 	"fmt"
 	"time"
 
+	autoscalingv1alpha1 "github.com/aibrix/aibrix/api/autoscaling/v1alpha1"
 	"github.com/aibrix/aibrix/pkg/controller/podautoscaler/metrics"
 
 	"github.com/aibrix/aibrix/pkg/controller/podautoscaler/scaler"
 	podutil "github.com/aibrix/aibrix/pkg/utils"
-
-	autoscalingv1alpha1 "github.com/aibrix/aibrix/api/autoscaling/v1alpha1"
 	podutils "github.com/aibrix/aibrix/pkg/utils"
+
 	autoscalingv2 "k8s.io/api/autoscaling/v2"
 	corev1 "k8s.io/api/core/v1"
 	apiequality "k8s.io/apimachinery/pkg/api/equality"
@@ -80,6 +80,34 @@ func newReconciler(mgr manager.Manager) (reconcile.Reconciler, error) {
 	return reconciler, nil
 }
 
+// for hpa related changes, let's make sure we only enqueue related PodAutoscaler objects
+func filterHPAObject(ctx context.Context, object client.Object) []reconcile.Request {
+	hpa, ok := object.(*autoscalingv2.HorizontalPodAutoscaler)
+	if !ok {
+		klog.Warningf("unexpected object type: %T, %s:%s, HPA object is expected here.", object, object.GetNamespace(), object.GetName())
+		return nil
+	}
+
+	// Iterate through ownerReferences to find the PodAutoscaler.
+	// if found, enqueue the managing PodAutoscaler object
+	for _, ownerRef := range hpa.OwnerReferences {
+		if ownerRef.Kind == "PodAutoScaler" && ownerRef.Controller != nil && *ownerRef.Controller {
+			// Enqueue the managing PodAutoscaler object
+			return []reconcile.Request{
+				{
+					NamespacedName: types.NamespacedName{
+						Namespace: hpa.Namespace,
+						Name:      ownerRef.Name,
+					},
+				},
+			}
+		}
+	}
+
+	// no managed pod autoscaler found, no need to enqueue original object.
+	return []reconcile.Request{}
+}
+
 // add adds a new Controller to mgr with r as the reconcile.Reconciler
 func add(mgr manager.Manager, r reconcile.Reconciler) error {
 	// Build raw source for periodical requeue events from event channel
@@ -92,7 +120,7 @@ func add(mgr manager.Manager, r reconcile.Reconciler) error {
 	// and HorizontalPodAutoscaler objects.
 	err := ctrl.NewControllerManagedBy(mgr).
 		For(&autoscalingv1alpha1.PodAutoscaler{}).
-		Watches(&autoscalingv2.HorizontalPodAutoscaler{}, &handler.EnqueueRequestForObject{}).
+		Watches(&autoscalingv2.HorizontalPodAutoscaler{}, handler.EnqueueRequestsFromMapFunc(filterHPAObject)).
 		WatchesRawSource(src, &handler.EnqueueRequestForObject{}).
 		Complete(r)
 
@@ -124,7 +152,7 @@ type PodAutoscalerReconciler struct {
 	eventCh        chan event.GenericEvent
 }
 
-func (r *PodAutoscalerReconciler) deleteScaler(request types.NamespacedName) {
+func (r *PodAutoscalerReconciler) deleteStaleScalerInCache(request types.NamespacedName) {
 	// When deleting, we only have access to the Namespace and Name, not other attributes in pa_types.
 	// We should scan `AutoscalerMap` and remove the matched objects.
 	// Note that due to the OwnerRef, the created HPA object will automatically be removed when AIBrix-HPA is deleted.
@@ -162,7 +190,7 @@ func (r *PodAutoscalerReconciler) Reconcile(ctx context.Context, req ctrl.Reques
 	var pa autoscalingv1alpha1.PodAutoscaler
 	if err := r.Get(ctx, req.NamespacedName, &pa); err != nil {
 		if errors.IsNotFound(err) {
-			r.deleteScaler(req.NamespacedName)
+			r.deleteStaleScalerInCache(req.NamespacedName)
 			// Object might have been deleted after reconcile request, clean it and return.
 			klog.Infof("PodAutoscaler resource not found. Clean scaler object in memory since object %s must have been deleted", req.NamespacedName)
 			return ctrl.Result{}, nil
@@ -270,7 +298,7 @@ func (r *PodAutoscalerReconciler) reconcileHPA(ctx context.Context, pa autoscali
 		return ctrl.Result{}, err
 	} else {
 		// Update the existing HPA if it already exists.
-		klog.InfoS("Updating existing HPA", "HPA", hpaName)
+		klog.V(4).InfoS("Updating existing HPA to desired state", "HPA", hpaName)
 
 		err = r.Update(ctx, hpa)
 		if err != nil {