[Feat] Consume StormService spec.mode in the controller update path and PodAutoscaler

docs/source/designs/aibrix-stormservice.rst

@@ -40,6 +40,8 @@ Stormservice supports two deployment modes: **Replica Mode** and **Pooled Mode**
     1. These two modes are mutually exclusive. The mode is declared through the `stormservice.spec.mode` field, which accepts `Replica` or `Pooled`.
     2. `spec.mode` is optional and is not defaulted. When it is omitted the mode is inferred for backward compatibility from `stormservice.spec.replicas`: replica mode when `replicas > 1`, otherwise pooled mode.
     3. When `spec.mode` is set to `Pooled`, `spec.replicas` must stay at `1`; roles are scaled through `spec.template.spec.roles[].replicas`.
+    4. A declared `spec.mode` drives the update path: `Replica` uses the rolling update path and `Pooled` uses the in-place update path, even when `spec.updateStrategy.type` holds the (possibly CRD-defaulted) `RollingUpdate` value. Declaring `mode: Replica` together with `updateStrategy.type: InPlaceUpdate` is rejected by the webhook. When `spec.mode` is omitted, `spec.updateStrategy.type` keeps selecting the update path as before.
+    5. A declared `spec.mode` is also the source of truth for PodAutoscaler role-level scaling. The `autoscaling.aibrix.ai/storm-service-mode` annotation is deprecated and only honored when the target StormService does not declare `spec.mode`.
 
 
 Replica Mode

docs/source/features/autoscaling/metric-based-autoscaling.rst

@@ -243,9 +243,9 @@ Example:
 StormService Role-Level Autoscaling
 ------------------------------------
 
-For StormService in pooled mode (``replicas=1``), different roles (e.g., prefill and decode) can be autoscaled independently. This enables fine-grained control where each role scales based on its specific metrics.
+For StormService in pooled mode (``spec.mode: Pooled``), different roles (e.g., prefill and decode) can be autoscaled independently. This enables fine-grained control where each role scales based on its specific metrics.
 
-Use the ``subTargetSelector`` field to target a specific role within a StormService. Additionally, add the annotation `autoscaling.aibrix.ai/storm-service-mode: "pool"` to the PodAutoscaler object. This helps the AIBrix autoscaler better distinguish ``replicas=1`` scenarios.
+Use the ``subTargetSelector`` field to target a specific role within a StormService, and declare ``spec.mode`` on the StormService (``Pooled`` to scale the targeted role, ``Replica`` to scale ``spec.replicas``). The autoscaler reads ``spec.mode`` to route role-level scaling; ``replicas=1`` alone cannot distinguish the two modes. The PodAutoscaler annotation ``autoscaling.aibrix.ai/storm-service-mode`` is deprecated and only honored as a compatibility fallback when the target StormService does not declare ``spec.mode``.
 
 **Key features:**
 

pkg/controller/podautoscaler/workload_scale.go

@@ -39,8 +39,34 @@ import (
 	"github.com/vllm-project/aibrix/pkg/controller/constants"
 )
 
+// AutoscalingStormServiceModeAnnotationKey is the PodAutoscaler annotation that
+// historically selected how role-level scaling applies to a StormService target:
+// "replica" scales StormService.spec.replicas, anything else scales the targeted
+// role's replicas.
+//
+// Deprecated: declare StormService.spec.mode instead. The annotation is only
+// honored as a compatibility fallback when the target StormService does not
+// declare spec.mode, and may be removed once spec.mode is broadly adopted.
 const AutoscalingStormServiceModeAnnotationKey = "autoscaling.aibrix.ai/storm-service-mode"
 
+// stormServiceScalingMode resolves the deployment mode used to route role-level
+// scaling for a StormService target. A declared StormService.spec.mode is the
+// source of truth. When spec.mode is unset, the deprecated
+// autoscaling.aibrix.ai/storm-service-mode annotation on the PodAutoscaler is
+// honored as a compatibility fallback ("replica" selects replica mode), and any
+// other value keeps the legacy pooled default. spec.replicas is intentionally
+// not used for inference here: replicas == 1 cannot distinguish a pooled
+// StormService from a scaled-down replica-mode one.
+func stormServiceScalingMode(pa *autoscalingv1alpha1.PodAutoscaler, ss *orchestrationv1alpha1.StormService) orchestrationv1alpha1.StormServiceMode {
+	if ss.Spec.Mode != "" {
+		return ss.Spec.Mode
+	}
+	if pa.Annotations[AutoscalingStormServiceModeAnnotationKey] == "replica" {
+		return orchestrationv1alpha1.StormServiceReplicaMode
+	}
+	return orchestrationv1alpha1.StormServicePooledMode
+}
+
 // WorkloadScale provides scaling operations for different workload types.
 // It provides the mechanism to get/set replica counts on workload resources,
 // while AutoScaler provides the intelligence to compute desired replica counts.
@@ -152,9 +178,11 @@ func (s *workloadScale) getCurrentReplicasForRole(ctx context.Context, pa *autos
 		return 0, err
 	}
 
-	// replica mode, return the replicas directly
-	// we can not easily use `*ss.Spec.Replicas > 1` as condition since 1 could be pool or replica both case under autoscaling scenarios
-	if pa.Annotations[AutoscalingStormServiceModeAnnotationKey] == "replica" {
+	// Replica mode scales the whole StormService, so report spec.replicas directly.
+	if stormServiceScalingMode(pa, ss) == orchestrationv1alpha1.StormServiceReplicaMode {
+		if ss.Spec.Replicas == nil {
+			return 0, nil
+		}
 		return *ss.Spec.Replicas, nil
 	}
 
@@ -246,9 +274,9 @@ func (s *workloadScale) setDesiredReplicasForRole(ctx context.Context, pa *autos
 			return err
 		}
 		upd := cur.DeepCopy()
-		// TODO: tricky part. it's hard to know replica=1 is pooling or replica mode, we can use autoscaling to limit it.
-		// we can extract the method and fallback to ss object annotation to check as well.
-		if pa.Annotations[AutoscalingStormServiceModeAnnotationKey] == "replica" {
+		// Replica mode scales the whole StormService through spec.replicas; pooled mode
+		// scales the targeted role. See stormServiceScalingMode for the resolution order.
+		if stormServiceScalingMode(pa, cur) == orchestrationv1alpha1.StormServiceReplicaMode {
 			upd.Spec.Replicas = ptr.To(replicas)
 			return s.client.Patch(ctx, upd, client.MergeFrom(cur))
 		}

pkg/controller/podautoscaler/workload_scale_test.go

@@ -30,10 +30,41 @@ import (
 	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/utils/ptr"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/client/fake"
 )
 
+func TestStormServiceScalingMode(t *testing.T) {
+	tests := map[string]struct {
+		specMode   orchestrationv1alpha1.StormServiceMode
+		annotation string
+		want       orchestrationv1alpha1.StormServiceMode
+	}{
+		"declared replica mode wins":                     {specMode: orchestrationv1alpha1.StormServiceReplicaMode, want: orchestrationv1alpha1.StormServiceReplicaMode},
+		"declared pooled mode wins over annotation":      {specMode: orchestrationv1alpha1.StormServicePooledMode, annotation: "replica", want: orchestrationv1alpha1.StormServicePooledMode},
+		"declared replica mode wins over annotation":     {specMode: orchestrationv1alpha1.StormServiceReplicaMode, annotation: "pool", want: orchestrationv1alpha1.StormServiceReplicaMode},
+		"no mode falls back to replica annotation":       {annotation: "replica", want: orchestrationv1alpha1.StormServiceReplicaMode},
+		"no mode with pool annotation defaults to pool":  {annotation: "pool", want: orchestrationv1alpha1.StormServicePooledMode},
+		"no mode and no annotation defaults to pool":     {want: orchestrationv1alpha1.StormServicePooledMode},
+		"no mode with unknown annotation stays pooled":   {annotation: "bogus", want: orchestrationv1alpha1.StormServicePooledMode},
+		"declared pooled mode with no annotation pooled": {specMode: orchestrationv1alpha1.StormServicePooledMode, want: orchestrationv1alpha1.StormServicePooledMode},
+	}
+
+	for name, tc := range tests {
+		t.Run(name, func(t *testing.T) {
+			pa := &autoscalingv1alpha1.PodAutoscaler{}
+			if tc.annotation != "" {
+				pa.Annotations = map[string]string{AutoscalingStormServiceModeAnnotationKey: tc.annotation}
+			}
+			ss := &orchestrationv1alpha1.StormService{
+				Spec: orchestrationv1alpha1.StormServiceSpec{Mode: tc.specMode},
+			}
+			assert.Equal(t, tc.want, stormServiceScalingMode(pa, ss))
+		})
+	}
+}
+
 func TestGetCurrentReplicasFromScale(t *testing.T) {
 	expectedReplicas := int32(2)
 
@@ -166,6 +197,74 @@ func TestGetCurrentReplicasFromScale(t *testing.T) {
 			},
 			scale: scaleStormService,
 		},
+		{
+			// spec.mode replaces the annotation as the mode signal.
+			name: "storm_service_with_declared_replica_mode",
+			pa: &autoscalingv1alpha1.PodAutoscaler{
+				ObjectMeta: v1.ObjectMeta{
+					Namespace: "default",
+				},
+				Spec: autoscalingv1alpha1.PodAutoscalerSpec{
+					SubTargetSelector: &autoscalingv1alpha1.SubTargetSelector{},
+					ScaleTargetRef: corev1.ObjectReference{
+						Kind: "StormService",
+						Name: "test-storm",
+					},
+				},
+			},
+			ss: &orchestrationv1alpha1.StormService{
+				ObjectMeta: v1.ObjectMeta{
+					Name:      "test-storm",
+					Namespace: "default",
+				},
+				Spec: orchestrationv1alpha1.StormServiceSpec{
+					Mode:     orchestrationv1alpha1.StormServiceReplicaMode,
+					Replicas: &expectedReplicas,
+				},
+			},
+			scale: scaleStormService,
+		},
+		{
+			// A declared pooled mode wins over a stale "replica" annotation:
+			// the role status is reported, not spec.replicas.
+			name: "storm_service_declared_pooled_mode_wins_over_annotation",
+			pa: &autoscalingv1alpha1.PodAutoscaler{
+				ObjectMeta: v1.ObjectMeta{
+					Namespace: "default",
+					Annotations: map[string]string{
+						AutoscalingStormServiceModeAnnotationKey: "replica",
+					},
+				},
+				Spec: autoscalingv1alpha1.PodAutoscalerSpec{
+					SubTargetSelector: &autoscalingv1alpha1.SubTargetSelector{
+						RoleName: "prefill",
+					},
+					ScaleTargetRef: corev1.ObjectReference{
+						Kind: "StormService",
+						Name: "test-storm",
+					},
+				},
+			},
+			ss: &orchestrationv1alpha1.StormService{
+				ObjectMeta: v1.ObjectMeta{
+					Name:      "test-storm",
+					Namespace: "default",
+				},
+				Spec: orchestrationv1alpha1.StormServiceSpec{
+					Mode:     orchestrationv1alpha1.StormServicePooledMode,
+					Replicas: ptr.To(int32(5)), // must not be reported in pooled mode
+				},
+				Status: orchestrationv1alpha1.StormServiceStatus{
+					RoleStatuses: []orchestrationv1alpha1.RoleStatus{
+						{
+							Name:     "prefill",
+							Replicas: expectedReplicas,
+						},
+					},
+				},
+			},
+			scale: scaleStormService,
+		},
 	}
 
 	for _, tt := range table {
@@ -455,6 +554,92 @@ func TestSetDesiredReplicas(t *testing.T) {
 				assert.Equal(t, expectedReplicas, *ss.Spec.Template.Spec.Roles[0].Replicas)
 			},
 		},
+		{
+			// spec.mode replaces the annotation as the mode signal: a declared
+			// replica mode updates spec.replicas without any annotation set.
+			name: "storm_service_with_declared_replica_mode",
+			pa: &autoscalingv1alpha1.PodAutoscaler{
+				ObjectMeta: v1.ObjectMeta{
+					Namespace: "default",
+				},
+				Spec: autoscalingv1alpha1.PodAutoscalerSpec{
+					SubTargetSelector: &autoscalingv1alpha1.SubTargetSelector{
+						RoleName: "",
+					},
+					ScaleTargetRef: corev1.ObjectReference{
+						Kind: "StormService",
+						Name: "test-storm",
+					},
+				},
+			},
+			deployment: &appsv1.Deployment{},
+			ss: &orchestrationv1alpha1.StormService{
+				ObjectMeta: v1.ObjectMeta{
+					Name:      "test-storm",
+					Namespace: "default",
+				},
+				Spec: orchestrationv1alpha1.StormServiceSpec{
+					Mode:     orchestrationv1alpha1.StormServiceReplicaMode,
+					Replicas: &currentReplicas,
+				},
+			},
+			assertReplicas: func(t *testing.T, fakeClient client.Client) {
+				ss := &orchestrationv1alpha1.StormService{}
+				err := fakeClient.Get(context.TODO(), client.ObjectKey{Namespace: "default", Name: "test-storm"}, ss)
+				assert.NoError(t, err)
+				assert.Equal(t, expectedReplicas, *ss.Spec.Replicas)
+			},
+		},
+		{
+			// A declared pooled mode wins over a stale "replica" annotation: the
+			// targeted role scales while spec.replicas stays untouched.
+			name: "storm_service_declared_pooled_mode_wins_over_annotation",
+			pa: &autoscalingv1alpha1.PodAutoscaler{
+				ObjectMeta: v1.ObjectMeta{
+					Namespace: "default",
+					Annotations: map[string]string{
+						AutoscalingStormServiceModeAnnotationKey: "replica",
+					},
+				},
+				Spec: autoscalingv1alpha1.PodAutoscalerSpec{
+					SubTargetSelector: &autoscalingv1alpha1.SubTargetSelector{
+						RoleName: "prefill",
+					},
+					ScaleTargetRef: corev1.ObjectReference{
+						Kind: "StormService",
+						Name: "test-storm",
+					},
+				},
+			},
+			deployment: &appsv1.Deployment{},
+			ss: &orchestrationv1alpha1.StormService{
+				ObjectMeta: v1.ObjectMeta{
+					Name:      "test-storm",
+					Namespace: "default",
+				},
+				Spec: orchestrationv1alpha1.StormServiceSpec{
+					Mode:     orchestrationv1alpha1.StormServicePooledMode,
+					Replicas: ptr.To(int32(1)),
+					Template: orchestrationv1alpha1.RoleSetTemplateSpec{
+						Spec: &orchestrationv1alpha1.RoleSetSpec{
+							Roles: []orchestrationv1alpha1.RoleSpec{
+								{
+									Name:     "prefill",
+									Replicas: &currentReplicas,
+								},
+							},
+						},
+					},
+				},
+			},
+			assertReplicas: func(t *testing.T, fakeClient client.Client) {
+				ss := &orchestrationv1alpha1.StormService{}
+				err := fakeClient.Get(context.TODO(), client.ObjectKey{Namespace: "default", Name: "test-storm"}, ss)
+				assert.NoError(t, err)
+				assert.Equal(t, expectedReplicas, *ss.Spec.Template.Spec.Roles[0].Replicas)
+				assert.Equal(t, int32(1), *ss.Spec.Replicas)
+			},
+		},
 	}
 
 	for _, tt := range table {

pkg/controller/stormservice/sync.go

@@ -272,17 +272,16 @@ func (r *StormServiceReconciler) rollout(ctx context.Context, stormService, curr
 	if len(updated) == int(expectReplica) {
 		return nil
 	}
-	switch stormService.Spec.UpdateStrategy.Type {
-	case "":
-		// By default use RollingUpdate strategy
-		fallthrough
-	case orchestrationv1alpha1.RollingUpdateStormServiceStrategyType:
-		return r.rollingUpdate(allRoleSets, stormService, current, currentCR, updateCR)
-	case orchestrationv1alpha1.InPlaceUpdateStormServiceStrategyType:
+	// The update path follows the declared spec.mode when it is set and falls back to
+	// the legacy updateStrategy.type selection otherwise, see EffectiveUpdateStrategyType.
+	strategyType, err := EffectiveUpdateStrategyType(stormService)
+	if err != nil {
+		return err
+	}
+	if strategyType == orchestrationv1alpha1.InPlaceUpdateStormServiceStrategyType {
 		return r.inPlaceUpdate(allRoleSets, stormService, current, currentCR, updateCR)
-	default:
-		return fmt.Errorf("unexpected stormService strategy type: %s", stormService.Spec.UpdateStrategy.Type)
 	}
+	return r.rollingUpdate(allRoleSets, stormService, current, currentCR, updateCR)
 }
 
 // rollingUpdate: rolling update logic for replica mode

pkg/controller/stormservice/utils.go

@@ -17,6 +17,7 @@ limitations under the License.
 package stormservice
 
 import (
+	"fmt"
 	"sort"
 
 	ctrlutil "github.com/vllm-project/aibrix/pkg/controller/util"
@@ -93,6 +94,51 @@ func IsRollingUpdate(stormService *orchestrationv1alpha1.StormService) bool {
 	return stormService.Spec.UpdateStrategy.Type == "" || stormService.Spec.UpdateStrategy.Type == orchestrationv1alpha1.RollingUpdateStormServiceStrategyType
 }
 
+// EffectiveUpdateStrategyType returns the update strategy type that drives the rollout path.
+//
+// A declared spec.mode is the source of truth: Replica mode replaces RoleSets through the
+// RollingUpdate path and Pooled mode updates its single RoleSet through the InPlaceUpdate
+// path. The webhook rejects a declared InPlaceUpdate strategy combined with Replica mode.
+// The reverse conflict (Pooled with RollingUpdate) cannot be rejected at admission because
+// the CRD defaults spec.updateStrategy.type to RollingUpdate whenever the updateStrategy
+// block is present, so a RollingUpdate value is indistinguishable from the default; the
+// controller resolves that conflict in favor of the declared mode and logs the override.
+//
+// When spec.mode is unset, the legacy updateStrategy.type selection is kept so existing
+// manifests behave as before. ResolvedMode is intentionally not consulted for the inferred
+// case: its replicas-based inference does not always match the declared strategy type
+// (for example replicas: 1 with RollingUpdate infers Pooled but must keep rolling).
+func EffectiveUpdateStrategyType(stormService *orchestrationv1alpha1.StormService) (orchestrationv1alpha1.StormServiceUpdateStrategyType, error) {
+	declaredType := stormService.Spec.UpdateStrategy.Type
+	if stormService.Spec.Mode != "" {
+		switch mode := stormService.Spec.ResolvedMode(); mode {
+		case orchestrationv1alpha1.StormServiceReplicaMode:
+			if declaredType == orchestrationv1alpha1.InPlaceUpdateStormServiceStrategyType {
+				klog.Infof("stormservice %s/%s declares mode %s, overriding updateStrategy.type %s with %s",
+					stormService.Namespace, stormService.Name, mode, declaredType, orchestrationv1alpha1.RollingUpdateStormServiceStrategyType)
+			}
+			return orchestrationv1alpha1.RollingUpdateStormServiceStrategyType, nil
+		case orchestrationv1alpha1.StormServicePooledMode:
+			if declaredType == orchestrationv1alpha1.RollingUpdateStormServiceStrategyType {
+				klog.Infof("stormservice %s/%s declares mode %s, overriding updateStrategy.type %s with %s",
+					stormService.Namespace, stormService.Name, mode, declaredType, orchestrationv1alpha1.InPlaceUpdateStormServiceStrategyType)
+			}
+			return orchestrationv1alpha1.InPlaceUpdateStormServiceStrategyType, nil
+		default:
+			return "", fmt.Errorf("unexpected stormService mode: %s", mode)
+		}
+	}
+	switch declaredType {
+	case "":
+		// By default use RollingUpdate strategy
+		return orchestrationv1alpha1.RollingUpdateStormServiceStrategyType, nil
+	case orchestrationv1alpha1.RollingUpdateStormServiceStrategyType, orchestrationv1alpha1.InPlaceUpdateStormServiceStrategyType:
+		return declaredType, nil
+	default:
+		return "", fmt.Errorf("unexpected stormService strategy type: %s", declaredType)
+	}
+}
+
 // ResolveFenceposts resolves both maxSurge and maxUnavailable. This needs to happen in one
 // step. For example:
 //