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

api/orchestration/v1alpha1/stormservice_types.go

@@ -100,6 +100,18 @@ func (s *StormServiceSpec) ResolvedMode() StormServiceMode {
 	return StormServicePooledMode
 }
 
+// ResolvedReplicas returns the effective number of desired RoleSets. spec.replicas is
+// optional and documented to default to 1 when not specified, but neither the CRD schema
+// nor the mutating webhook materializes that default, so an omitted value reaches the
+// controllers as nil and is resolved here. An explicit 0 is preserved: the field is a
+// pointer precisely to tell explicit zero apart from not specified.
+func (s *StormServiceSpec) ResolvedReplicas() int32 {
+	if s.Replicas == nil {
+		return 1
+	}
+	return *s.Replicas
+}
+
 type RoleSetTemplateSpec struct {
 	// +optional
 	metav1.ObjectMeta `json:"metadata,omitempty" protobuf:"bytes,1,opt,name=metadata"`

api/orchestration/v1alpha1/stormservice_types_test.go

@@ -42,3 +42,21 @@ func TestStormServiceSpecResolvedMode(t *testing.T) {
 		})
 	}
 }
+
+func TestStormServiceSpecResolvedReplicas(t *testing.T) {
+	tests := map[string]struct {
+		replicas *int32
+		want     int32
+	}{
+		"unset resolves to the documented default of 1": {replicas: nil, want: 1},
+		"explicit zero is preserved":                    {replicas: ptr.To[int32](0), want: 0},
+		"explicit count is returned as is":              {replicas: ptr.To[int32](3), want: 3},
+	}
+
+	for name, tc := range tests {
+		t.Run(name, func(t *testing.T) {
+			spec := &StormServiceSpec{Replicas: tc.replicas}
+			assert.Equal(t, tc.want, spec.ResolvedReplicas())
+		})
+	}
+}

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,10 +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" {
-		return *ss.Spec.Replicas, nil
+	// Replica mode scales the whole StormService, so report spec.replicas. An omitted
+	// spec.replicas resolves to the documented default of 1 instead of reading as a
+	// scaled-to-zero workload, matching how the stormservice controller reconciles it.
+	if stormServiceScalingMode(pa, ss) == orchestrationv1alpha1.StormServiceReplicaMode {
+		return ss.Spec.ResolvedReplicas(), nil
 	}
 
 	roleName := pa.Spec.SubTargetSelector.RoleName
@@ -246,9 +273,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 {
@@ -189,6 +288,55 @@ func TestGetCurrentReplicasFromScale(t *testing.T) {
 	}
 }
 
+// TestGetCurrentReplicasReplicaModeNilReplicas covers the review scenario for an
+// omitted spec.replicas in replica mode: the field is optional with a documented
+// default of 1 that neither the CRD nor a webhook materializes, so the autoscaler
+// must resolve nil to 1 rather than report a scaled-to-zero workload.
+func TestGetCurrentReplicasReplicaModeNilReplicas(t *testing.T) {
+	scale := &unstructured.Unstructured{}
+	scale.SetAPIVersion("orchestration.aibrix.ai/v1alpha1")
+	scale.SetKind("StormService")
+
+	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 intentionally omitted.
+		},
+	}
+
+	scheme := runtime.NewScheme()
+	_ = autoscalingv1alpha1.AddToScheme(scheme)
+	_ = orchestrationv1alpha1.AddToScheme(scheme)
+
+	fakeClient := fake.NewClientBuilder().
+		WithScheme(scheme).
+		WithObjects(pa, ss).
+		Build()
+
+	workloadScale := NewWorkloadScale(fakeClient, nil)
+
+	currentReplicas, err := workloadScale.GetCurrentReplicasFromScale(context.TODO(), pa, scale)
+
+	assert.NoError(t, err)
+	assert.Equal(t, int32(1), currentReplicas)
+}
+
 func TestGetPodSelectorFromScale(t *testing.T) {
 	t.Run("llm_model", func(t *testing.T) {
 		scheme := runtime.NewScheme()
@@ -455,6 +603,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 {