Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rishabhsinha17 wants to merge 3 commits into
base: main
Choose a base branch
from
+734 −32
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c320d81
[Feat] Consume StormService spec.mode in controller update path and P…
rishabhsinha17 Aug 26, 2026
7e15a22
Merge branch 'main' into fix/storm-service-resolved-mode
varungup90 Aug 26, 2026
d419b4b
[Fix] Route IsRollingUpdate through EffectiveUpdateStrategyType so po…
rishabhsinha17 Aug 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions docs/source/designs/aibrix-stormservice.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 2 additions & 2 deletions docs/source/features/autoscaling/metric-based-autoscaling.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:**

Expand Down
40 changes: 34 additions & 6 deletions pkg/controller/podautoscaler/workload_scale.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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
}
Comment on lines +183 to +185

@gemini-code-assist gemini-code-assist Bot Aug 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Since spec.replicas defaults to 1 when omitted (as documented in stormservice_types.go), returning 0 here when ss.Spec.Replicas is nil can lead to incorrect autoscaling calculations (e.g., division by zero or incorrect scaling ratios). It should return 1 instead to reflect the default replica count.

Suggested change
if ss.Spec.Replicas == nil {
return 0, nil
}
if ss.Spec.Replicas == nil {
return 1, nil
}

@googs1025 googs1025 Aug 27, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you please double-check this nil branch? Since spec.replicas is optional and documented as defaulting to 1, returning 0 here may make the autoscaler interpret an omitted replica count as zero replicas. Could you verify whether this field is defaulted before the autoscaler reads it, or whether this should resolve nil to the documented default of 1?

return *ss.Spec.Replicas, nil
}

Expand Down Expand Up @@ -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))
}
Expand Down
185 changes: 185 additions & 0 deletions pkg/controller/podautoscaler/workload_scale_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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 {
Expand Down
17 changes: 8 additions & 9 deletions pkg/controller/stormservice/sync.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

@gemini-code-assist gemini-code-assist Bot Aug 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

While reviewing the update path changes, we noticed a critical issue regarding how spec.replicas is handled when it is nil (omitted). Since spec.replicas is an optional field (*int32) and has no defaulting webhook or CRD default value, it can be nil in the cluster.

In updateStatus (which is called at line 64 of sync.go during reconciliation), the controller directly dereferences stormService.Spec.Replicas at lines 399 and 400:

stormService.Status.UpdatedReplicas == *stormService.Spec.Replicas
stormService.Status.Replicas == *stormService.Spec.Replicas

This will cause a critical nil pointer dereference panic and crash the controller whenever a StormService is created with replicas omitted.

Additionally, in scaling (line 148) and rollout (line 268), expectReplica defaults to 0 when Replicas is nil, which contradicts the documented default of 1 and will prevent the controller from creating any RoleSets or performing rollouts.

Please ensure that Replicas is safely checked for nil and defaulted to 1 throughout the controller reconcile loop.

@googs1025 googs1025 Aug 27, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you please double-check the spec.replicas == nil handling around the reconcile path? spec.replicas is optional, and the generated CRD does not seem to set a real default: 1. If omitted replicas can reach reconcile, scaling() / rollout() appear to derive expectReplica as 0, while updateStatus() still compares against *stormService.Spec.Replicas. Could you verify whether replicas is actually defaulted before reconcile in this codebase, or whether these paths should resolve nil to the documented default of 1?

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
Expand Down
Loading
Loading