feat(roles): add support for role dependencies to enforce startup order

Signed-off-by: CYJiang <googs1025@gmail.com>

Author: googs1025

api/orchestration/v1alpha1/roleset_types.go

@@ -150,6 +150,22 @@ type RoleSpec struct {
 	// +optional
 	Replicas *int32 `json:"replicas,omitempty"`
 
+	// Dependencies specifies the list of role names that must be ready before this role starts.
+	// This is used to define a dependency graph between different roles, ensuring that certain
+	// roles only start once their dependencies have successfully started and are ready.
+	// Each element in the slice should correspond to the name of another role defined within
+	// the same StormService specification.
+	//
+	// For example, if RoleA depends on RoleB and RoleC, then RoleA's Dependencies would look like:
+	// Dependencies: []string{"RoleB", "RoleC"}
+	//
+	// The system will ensure that RoleA does not begin its startup process until both RoleB and
+	// RoleC have been marked as ready. If any specified dependency fails to start or becomes not
+	// ready, RoleA's startup (or restart) will be delayed until all dependencies are satisfied.
+	//
+	// +optional
+	Dependencies []string `json:"dependencies,omitempty"`
+
 	// UpgradeOrder specifies the order in which this role should be upgraded.
 	// Lower values are upgraded first. If not specified, roles upgrade after all explicitly ordered roles.
 	// +optional

api/orchestration/v1alpha1/zz_generated.deepcopy.go

@@ -38,6 +38,10 @@ spec:
               roles:
                 items:
                   properties:
+                    dependencies:
+                      items:
+                        type: string
+                      type: array
                     disruptionTolerance:
                       properties:
                         maxUnavailable:

config/crd/orchestration/orchestration.aibrix.ai_rolesets.yaml

@@ -113,6 +113,10 @@ spec:
                       roles:
                         items:
                           properties:
+                            dependencies:
+                              items:
+                                type: string
+                              type: array
                             disruptionTolerance:
                               properties:
                                 maxUnavailable:

config/crd/orchestration/orchestration.aibrix.ai_stormservices.yaml

@@ -38,6 +38,10 @@ spec:
               roles:
                 items:
                   properties:
+                    dependencies:
+                      items:
+                        type: string
+                      type: array
                     disruptionTolerance:
                       properties:
                         maxUnavailable:

dist/chart/crds/orchestration.aibrix.ai_rolesets.yaml

@@ -113,6 +113,10 @@ spec:
                       roles:
                         items:
                           properties:
+                            dependencies:
+                              items:
+                                type: string
+                              type: array
                             disruptionTolerance:
                               properties:
                                 maxUnavailable:

dist/chart/crds/orchestration.aibrix.ai_stormservices.yaml

@@ -22,11 +22,10 @@ import (
 	"math"
 	"strings"
 
-	orchestrationv1alpha1 "github.com/vllm-project/aibrix/api/orchestration/v1alpha1"
-
 	"k8s.io/klog/v2"
-
 	"sigs.k8s.io/controller-runtime/pkg/client"
+
+	orchestrationv1alpha1 "github.com/vllm-project/aibrix/api/orchestration/v1alpha1"
 )
 
 type RollingManager interface {
@@ -38,22 +37,35 @@ type RollingManagerSequential struct {
 }
 
 func (m *RollingManagerSequential) Next(ctx context.Context, roleSet *orchestrationv1alpha1.RoleSet) (err error) {
-	// 1. ensure pod replica meet expectations
+	// 1. Find the furthest role whose dependencies are ready
+	progressIndex := -1
+	for i, role := range roleSet.Spec.Roles {
+		if !isRoleDependenciesReady(roleSet, &role) {
+			break
+		}
+		progressIndex = i
+	}
+
+	if progressIndex == -1 {
+		return nil // even A is blocked (shouldn't happen)
+	}
+
+	// 2. Scale all roles from 0 to progressIndex
 	var scaling bool
-	for _, role := range roleSet.Spec.Roles {
-		klog.Infof("[RollingManagerSequential.Next] start to scale roleset %s/%s role %s", roleSet.Namespace, roleSet.Name, role.Name)
-		s, err := GetRoleSyncer(m.cli, &role).Scale(ctx, roleSet, &role)
+	for i := 0; i <= progressIndex; i++ {
+		role := &roleSet.Spec.Roles[i]
+		s, err := GetRoleSyncer(m.cli, role).Scale(ctx, roleSet, role)
 		if err != nil {
 			return err
 		}
 		scaling = scaling || s
 	}
 	if scaling {
-		klog.Infof("[RollingManagerSequential.Next] waiting for roleset %s/%s to be scaled", roleSet.Namespace, roleSet.Name)
+		// Wait for Pods to be created and status updated
 		return nil
 	}
 
-	// 2. Sort roles by upgrade order
+	// 3. Sort roles by upgrade order
 	klog.Infof("[RollingManagerSequential.Next] sorting roleset roles by UpgradeOrder")
 	sortedRoles := sortRolesByUpgradeOrder(roleSet.Spec.Roles)
 	var sequenceLines []string
@@ -72,7 +84,7 @@ func (m *RollingManagerSequential) Next(ctx context.Context, roleSet *orchestrat
 		roleSet.Name,
 		strings.Join(sequenceLines, "\n"))
 
-	// 3. do the rollout process for each role by order
+	// 4. do the rollout process for each role by order
 	for _, role := range sortedRoles {
 		klog.Infof("[RollingManagerSequential.Next] start to rollout roleset %s/%s role %s", roleSet.Namespace, roleSet.Name, role.Name)
 		err := GetRoleSyncer(m.cli, &role).Rollout(ctx, roleSet, &role)
@@ -95,29 +107,43 @@ type RollingManagerParallel struct {
 }
 
 func (m *RollingManagerParallel) Next(ctx context.Context, roleSet *orchestrationv1alpha1.RoleSet) (err error) {
-	// 1. ensure pod replica meet expectations
+	// 1. Find the furthest role whose dependencies are ready
+	progressIndex := -1
+	for i, role := range roleSet.Spec.Roles {
+		if !isRoleDependenciesReady(roleSet, &role) {
+			break
+		}
+		progressIndex = i
+	}
+
+	if progressIndex == -1 {
+		return nil // even A is blocked (shouldn't happen)
+	}
+
+	// 2: Scale all roles from 0 to progressIndex
 	var scaling bool
-	for _, role := range roleSet.Spec.Roles {
-		klog.Infof("[RollingManagerParallel.Next] start to scale roleset %s/%s role %s", roleSet.Namespace, roleSet.Name, role.Name)
-		s, err := GetRoleSyncer(m.cli, &role).Scale(ctx, roleSet, &role)
+	for i := 0; i <= progressIndex; i++ {
+		role := &roleSet.Spec.Roles[i]
+		s, err := GetRoleSyncer(m.cli, role).Scale(ctx, roleSet, role)
 		if err != nil {
 			return err
 		}
 		scaling = scaling || s
 	}
 	if scaling {
-		klog.Infof("[RollingManagerParallel.Next] waiting for roleset %s/%s to be scaled", roleSet.Namespace, roleSet.Name)
-		return
+		// Wait for Pods to be created and status updated
+		return nil
 	}
-	// 2. do the rollout process for each role
+
+	// 3. do the rollout process for each role
 	for _, role := range roleSet.Spec.Roles {
 		klog.Infof("[RollingManagerParallel.Next] start to rollout roleset %s/%s role %s", roleSet.Namespace, roleSet.Name, role.Name)
 		err := GetRoleSyncer(m.cli, &role).Rollout(ctx, roleSet, &role)
 		if err != nil {
 			return err
 		}
 	}
-	return
+	return nil
 }
 
 type RollingManagerInterleave struct {
@@ -127,21 +153,34 @@ type RollingManagerInterleave struct {
 // Interleaved rollout: update roles in alternating steps,
 // using (maxSurge + maxUnavailable) as the step size for all roles
 func (m *RollingManagerInterleave) Next(ctx context.Context, roleSet *orchestrationv1alpha1.RoleSet) (err error) {
-	// 1. ensure pod replica meet expectations
+	// 1. Find the furthest role whose dependencies are ready
+	progressIndex := -1
+	for i, role := range roleSet.Spec.Roles {
+		if !isRoleDependenciesReady(roleSet, &role) {
+			break
+		}
+		progressIndex = i
+	}
+
+	if progressIndex == -1 {
+		return nil // even A is blocked (shouldn't happen)
+	}
+
+	// 2. Scale all roles from 0 to progressIndex
 	var scaling bool
-	for _, role := range roleSet.Spec.Roles {
-		klog.Infof("[RollingManagerInterleave.Next] start to scale roleset %s/%s role %s", roleSet.Namespace, roleSet.Name, role.Name)
-		s, err := GetRoleSyncer(m.cli, &role).Scale(ctx, roleSet, &role)
+	for i := 0; i <= progressIndex; i++ {
+		role := &roleSet.Spec.Roles[i]
+		s, err := GetRoleSyncer(m.cli, role).Scale(ctx, roleSet, role)
 		if err != nil {
 			return err
 		}
 		scaling = scaling || s
 	}
 	if scaling {
-		klog.Infof("[RollingManagerInterleave.Next] waiting for roleset %s/%s to be scaled", roleSet.Namespace, roleSet.Name)
+		// Wait for Pods to be created and status updated
 		return nil
 	}
-	// 2. do the rollout process for each role
+	// 3. do the rollout process for each role
 	roleSteps := make(map[string]int32)
 	roleReadyStatus := make(map[string]bool)
 	currentStep := int32(math.MaxInt32)

pkg/client/applyconfiguration/orchestration/v1alpha1/rolespec.go

@@ -436,3 +436,35 @@ func sortRolesByUpgradeOrder(roles []orchestrationv1alpha1.RoleSpec) []orchestra
 	})
 	return sortedRoles
 }
+
+func isRoleDependenciesReady(roleSet *orchestrationv1alpha1.RoleSet, role *orchestrationv1alpha1.RoleSpec) bool {
+	if len(role.Dependencies) == 0 {
+		return true
+	}
+
+	// Initialize all roles to 0
+	roleStatusMap := make(map[string]int32)
+	for _, r := range roleSet.Spec.Roles {
+		roleStatusMap[r.Name] = 0
+	}
+	for _, rs := range roleSet.Status.Roles {
+		roleStatusMap[rs.Name] = rs.ReadyReplicas
+	}
+
+	for _, depName := range role.Dependencies {
+		depReady := roleStatusMap[depName]
+		expected := int32(1)
+		for _, r := range roleSet.Spec.Roles {
+			if r.Name == depName && r.Replicas != nil {
+				expected = *r.Replicas
+				break
+			}
+		}
+		if depReady < expected {
+			klog.V(4).Infof("Role %s depends on %s, "+
+				"but only %d/%d ready", role.Name, depName, depReady, expected)
+			return false
+		}
+	}
+	return true
+}