fix(operator): incorrect expected replicas notification (#6890)

domsolutions · web-flow · commit 24950ebe91ad · 2025-10-20T11:39:03.000+01:00
* wip

* wip

* stateful sets working

* only notify scheduler of replicas if deployment successful
diff --git a/operator/controllers/mlops/server_controller.go b/operator/controllers/mlops/server_controller.go
@@ -30,7 +30,6 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/controller-runtime/pkg/reconcile"
 
-	"github.com/seldonio/seldon-core/operator/v2/apis/mlops/v1alpha1"
 	mlopsv1alpha1 "github.com/seldonio/seldon-core/operator/v2/apis/mlops/v1alpha1"
 	"github.com/seldonio/seldon-core/operator/v2/controllers/reconcilers/common"
 	serverreconcile "github.com/seldonio/seldon-core/operator/v2/controllers/reconcilers/server"
@@ -106,12 +105,6 @@ func (r *ServerReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctr
 		return reconcile.Result{}, reconcile.TerminalError(err)
 	}
 
-	if err := r.Scheduler.ServerNotify(ctx, nil, []v1alpha1.Server{*server}, false); err != nil {
-		logger.Error(err, "Failed calling ServerNotify", "name", req.Name, "namespace", req.Namespace, "spec", server.Spec)
-		r.updateStatusFromError(ctx, logger, server, err)
-		return reconcile.Result{}, err
-	}
-
 	server.Spec.Replicas = ptr.To(int32(scalingSpec.Replicas))
 
 	var sr common.Reconciler
@@ -142,13 +135,23 @@ func (r *ServerReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctr
 		return reconcile.Result{}, err
 	}
 
+	// attempt to deploy
 	err = sr.Reconcile()
 	if err != nil {
 		logger.Error(err, "Failed reconciling", "name", req.Name, "namespace", req.Namespace, "spec", server.Spec)
 		r.updateStatusFromError(ctx, logger, server, err)
 		return reconcile.Result{}, err
 	}
 
+	// only after successful deployment of the pods, we notify the scheduler of the number of Servers it can expect that
+	// will try and connect. It's important we do this after we deploy, otherwise if deployment fails, the scheduler
+	// may expect the wrong number, and end up waiting for the 10-minute timeout until loading models onto the Servers.
+	if err := r.Scheduler.ServerNotify(ctx, nil, []mlopsv1alpha1.Server{*server}, false); err != nil {
+		logger.Error(err, "Failed calling ServerNotify", "name", req.Name, "namespace", req.Namespace, "spec", server.Spec)
+		r.updateStatusFromError(ctx, logger, server, err)
+		return reconcile.Result{}, err
+	}
+
 	conditions := sr.GetConditions()
 	for _, condition := range conditions {
 		server.Status.SetCondition(condition)
diff --git a/operator/controllers/mlops/suite_test.go b/operator/controllers/mlops/suite_test.go
@@ -467,14 +467,15 @@ var _ = Describe("Controller", func() {
 	})
 
 	Context("When creating a Server spec", func() {
-		It("Accepts a valid Spec with minReplicas = 2 replicas not set, therefore replicas should take minReplicas value of 1", func() {
+		It("Accepts a valid Spec with minReplicas = 2, replicas not set, therefore replicas should take minReplicas value of 2", func() {
 			testID := "test-mlserver-2"
 			minReplicas := ptr.To(int32(2))
 
 			serverNotifyCalled := &atomic.Bool{}
 			addExpectedServerNotifyCall(testID, mlopsv1alpha1.ScalingSpec{
 				MinReplicas: minReplicas,
 				MaxReplicas: ptr.To(int32(2)),
+				Replicas:    minReplicas,
 			}, serverNotifyCalled)
 
 			ctx := context.Background()
diff --git a/operator/main.go b/operator/main.go
@@ -159,7 +159,7 @@ func main() {
 	schedulerClient := scheduler.NewSchedulerClient(logger,
 		mgr.GetClient(),
 		mgr.GetEventRecorderFor("scheduler-client"),
-		*tlsOptions)
+		*tlsOptions, useDeploymentsForServers)
 
 	if err = (&mlopscontrollers.ModelReconciler{
 		Client:    mgr.GetClient(),
diff --git a/operator/scheduler/client.go b/operator/scheduler/client.go
@@ -50,31 +50,33 @@ const (
 
 type SchedulerClient struct {
 	client.Client
-	logger         logr.Logger
-	callOptions    []grpc.CallOption
-	recorder       record.EventRecorder
-	seldonRuntimes map[string]*grpc.ClientConn // map of namespace to grpc connection
-	tlsOptions     tls.TLSOptions
-	mu             sync.Mutex
+	logger                   logr.Logger
+	callOptions              []grpc.CallOption
+	recorder                 record.EventRecorder
+	seldonRuntimes           map[string]*grpc.ClientConn // map of namespace to grpc connection
+	tlsOptions               tls.TLSOptions
+	useDeploymentsForServers bool
+	mu                       sync.Mutex
 }
 
 //	connect on demand by add getConnection(namespace) which if not existing calls connect to scheduler.
 //
 // For this will need to know ports (hardwire for now to 9004 and 9044 - ssl comes fom envvar - so always
 // the same for all schedulers
-func NewSchedulerClient(logger logr.Logger, client client.Client, recorder record.EventRecorder, tlsOptions tls.TLSOptions) *SchedulerClient {
+func NewSchedulerClient(logger logr.Logger, client client.Client, recorder record.EventRecorder, tlsOptions tls.TLSOptions, useDeploymentsForServers bool) *SchedulerClient {
 	opts := []grpc.CallOption{
 		grpc.MaxCallSendMsgSize(math.MaxInt32),
 		grpc.MaxCallRecvMsgSize(math.MaxInt32),
 	}
 
 	return &SchedulerClient{
-		Client:         client,
-		logger:         logger.WithName("SchedulerClient"),
-		callOptions:    opts,
-		recorder:       recorder,
-		seldonRuntimes: make(map[string]*grpc.ClientConn),
-		tlsOptions:     tlsOptions,
+		Client:                   client,
+		logger:                   logger.WithName("SchedulerClient"),
+		callOptions:              opts,
+		recorder:                 recorder,
+		seldonRuntimes:           make(map[string]*grpc.ClientConn),
+		tlsOptions:               tlsOptions,
+		useDeploymentsForServers: useDeploymentsForServers,
 	}
 }
 
@@ -161,7 +163,6 @@ func (s *SchedulerClient) handleControlPlaneEvent(ctx context.Context, grpcClien
 			s.logger.Error(err, "Failed to send registered server to scheduler")
 			return err
 		}
-
 		return nil
 	case scheduler.ControlPlaneResponse_SEND_EXPERIMENTS:
 		err := s.handleExperiments(ctx, grpcClient, namespace)
diff --git a/operator/scheduler/control_plane_test.go b/operator/scheduler/control_plane_test.go
@@ -16,7 +16,9 @@ import (
 	"testing"
 	"time"
 
+	"github.com/gotidy/ptr"
 	. "github.com/onsi/gomega"
+	v1 "k8s.io/api/apps/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 
@@ -161,6 +163,17 @@ func TestControlPlaneEvents(t *testing.T) {
 					},
 					Spec: mlopsv1alpha1.ServerSpec{},
 				},
+				&v1.StatefulSet{
+					TypeMeta: metav1.TypeMeta{},
+					ObjectMeta: metav1.ObjectMeta{
+						Name:       "foo",
+						Namespace:  "default",
+						Generation: 1,
+					},
+					Spec: v1.StatefulSetSpec{
+						Replicas: ptr.Int32(1),
+					},
+				},
 			},
 			expected_requests_pipelines: []*scheduler.LoadPipelineRequest{
 				{
@@ -233,7 +246,7 @@ func TestControlPlaneEvents(t *testing.T) {
 		t.Run(test.name, func(t *testing.T) {
 			grpcClient := mockSchedulerGrpcClient{}
 
-			controller := newMockControllerClient(test.existing_resources...)
+			controller := newMockControllerClient(false, test.existing_resources...)
 
 			err := controller.SubscribeControlPlaneEvents(context.Background(), &grpcClient, "")
 			g.Expect(err).To(BeNil())
diff --git a/operator/scheduler/experiment_test.go b/operator/scheduler/experiment_test.go
@@ -325,7 +325,7 @@ func TestSubscribeExperimentsEvents(t *testing.T) {
 					responses_subscribe_experiments: test.results,
 				}
 			}
-			controller := newMockControllerClient(test.existing_resources...)
+			controller := newMockControllerClient(false, test.existing_resources...)
 			err := controller.handleExperiments(context.Background(), &grpcClient, "")
 			g.Expect(err).To(BeNil())
 			err = controller.SubscribeExperimentEvents(context.Background(), &grpcClient, "")
diff --git a/operator/scheduler/model_test.go b/operator/scheduler/model_test.go
@@ -189,7 +189,7 @@ func TestSubscribeModelEvents(t *testing.T) {
 					responses_subscribe_models: test.results,
 				}
 			}
-			controller := newMockControllerClient(test.existing_resources...)
+			controller := newMockControllerClient(false, test.existing_resources...)
 			err := controller.handleModels(context.Background(), &grpcClient, "")
 			g.Expect(err).To(BeNil())
 			err = controller.SubscribeModelEvents(context.Background(), &grpcClient, "")
diff --git a/operator/scheduler/pipeline_test.go b/operator/scheduler/pipeline_test.go
@@ -419,7 +419,7 @@ func TestSubscribePipelineEvents(t *testing.T) {
 					responses_subscribe_pipelines: test.results,
 				}
 			}
-			controller := newMockControllerClient(test.existing_resources...)
+			controller := newMockControllerClient(false, test.existing_resources...)
 			err := controller.handlePipelines(context.Background(), &grpcClient, "")
 			g.Expect(err).To(BeNil())
 			err = controller.SubscribePipelineEvents(context.Background(), &grpcClient, "")
diff --git a/operator/scheduler/server.go b/operator/scheduler/server.go
@@ -29,11 +29,12 @@ import (
 )
 
 func (s *SchedulerClient) ServerNotify(ctx context.Context, grpcClient scheduler.SchedulerClient, servers []v1alpha1.Server, isFirstSync bool) error {
+	logger := s.logger.WithName("NotifyServer")
 	if len(servers) == 0 {
+		logger.Info("No servers to notify scheduler, skipping")
 		return nil
 	}
 
-	logger := s.logger.WithName("NotifyServer")
 	if grpcClient == nil {
 		// we assume that all servers are in the same namespace
 		namespace := servers[0].Namespace
diff --git a/operator/scheduler/server_test.go b/operator/scheduler/server_test.go
@@ -244,7 +244,7 @@ func TestServerNotify(t *testing.T) {
 			grpcClient := mockSchedulerGrpcClient{
 				requests_servers: []*scheduler.ServerNotify{},
 			}
-			controller := newMockControllerClient()
+			controller := newMockControllerClient(false)
 			err := controller.ServerNotify(context.Background(), &grpcClient, test.servers, false)
 			g.Expect(err).To(BeNil())
 
@@ -470,7 +470,7 @@ func TestSubscribeServerEvents(t *testing.T) {
 			}
 
 			existing_resources := []client.Object{&test.existingServer}
-			controller := newMockControllerClient(existing_resources...)
+			controller := newMockControllerClient(false, existing_resources...)
 			err := controller.SubscribeServerEvents(context.Background(), &grpcClient, "")
 			g.Expect(err).To(BeNil())
 
diff --git a/operator/scheduler/utils.go b/operator/scheduler/utils.go
@@ -14,6 +14,9 @@ import (
 	"fmt"
 	"io"
 
+	v1 "k8s.io/api/apps/v1"
+	"k8s.io/apimachinery/pkg/api/errors"
+	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/client-go/util/retry"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 
@@ -136,6 +139,54 @@ func (s *SchedulerClient) handleLoadedModels(
 	return nil
 }
 
+func (s *SchedulerClient) handleServerStatefulSet(ctx context.Context, namespace string, servers *[]v1alpha1.Server, server *v1alpha1.Server) error {
+	found := &v1.StatefulSet{}
+	err := s.Client.Get(ctx, types.NamespacedName{
+		Name:      server.Name,
+		Namespace: namespace,
+	}, found)
+	if err != nil {
+		if errors.IsNotFound(err) {
+			// this may not have been deployed due to it being invalid i.e. invalid scaling spec
+			s.logger.V(1).Info("Statefulset not found, skipping", "server", server.Name)
+			return nil
+		}
+		return fmt.Errorf("failed to get statefulset: %w", err)
+	}
+
+	server.Spec.Replicas = found.Spec.Replicas
+	*servers = append(*servers, *server)
+	return nil
+}
+
+func (s *SchedulerClient) handleServerDeployment(ctx context.Context, namespace string, servers *[]v1alpha1.Server, server *v1alpha1.Server) error {
+	found := &v1.Deployment{}
+	err := s.Client.Get(ctx, types.NamespacedName{
+		Name:      server.Name,
+		Namespace: namespace,
+	}, found)
+	if err != nil {
+		if errors.IsNotFound(err) {
+			// this may not have been deployed due to it being invalid i.e. invalid scaling spec
+			s.logger.V(1).Info("Deployment not found, skipping", "server", server.Name)
+			return nil
+		}
+		return fmt.Errorf("failed to get deployment: %w", err)
+	}
+
+	server.Spec.Replicas = found.Spec.Replicas
+	*servers = append(*servers, *server)
+	return nil
+}
+
+// handleRegisteredServers we need to inform the scheduler of the Servers deployed, in particular their:
+// - replicas
+// - max/min replicas
+// It is possible for a customer to change their Server CR to an invalid state which can not be deployed i.e. invalid
+// scaling spec, ServerConfig CR field is changed such as VPC size, which can not be changed when a StatefulSet is already
+// deployed. In which case, we can't rely on the Server CR as the source of truth of what's deployed. Instead, we
+// get each Server CR and then check how many replicas are actually deployed. This mitigates the scheduler waiting
+// the 10-minute timeout window for Servers to connect which haven't been deployed.
 func (s *SchedulerClient) handleRegisteredServers(
 	ctx context.Context, grpcClient scheduler.SchedulerClient, namespace string) error {
 	serverList := &v1alpha1.ServerList{}
@@ -149,10 +200,30 @@ func (s *SchedulerClient) handleRegisteredServers(
 		return fmt.Errorf("failed to list servers: %w", err)
 	}
 
-	if err := s.ServerNotify(ctx, grpcClient, serverList.Items, true); err != nil {
-		return err
+	servers := make([]v1alpha1.Server, 0)
+
+	for _, server := range serverList.Items {
+		if !server.ObjectMeta.DeletionTimestamp.IsZero() {
+			// ServerNotify will handle setting replicas = 0 before sending to scheduler
+			servers = append(servers, server)
+			continue
+		}
+
+		if s.useDeploymentsForServers {
+			if err := s.handleServerDeployment(ctx, server.Namespace, &servers, &server); err != nil {
+				return fmt.Errorf("failed to handle server deployment: %w", err)
+			}
+			continue
+		}
+
+		if err := s.handleServerStatefulSet(ctx, server.Namespace, &servers, &server); err != nil {
+			return fmt.Errorf("failed to handle server statefulset: %w", err)
+		}
 	}
 
+	if err := s.ServerNotify(ctx, grpcClient, servers, true); err != nil {
+		return err
+	}
 	return nil
 }
 
diff --git a/operator/scheduler/utils_test.go b/operator/scheduler/utils_test.go

Original file line number	Diff line number	Diff line change
`@@ -325,7 +325,7 @@ func TestSubscribeExperimentsEvents(t *testing.T) {`
`325`	`325`	`responses_subscribe_experiments: test.results,`
`326`	`326`	`}`
`327`	`327`	`}`
`328`		`- controller := newMockControllerClient(test.existing_resources...)`
	`328`	`+ controller := newMockControllerClient(false, test.existing_resources...)`
`329`	`329`	`err := controller.handleExperiments(context.Background(), &grpcClient, "")`
`330`	`330`	`g.Expect(err).To(BeNil())`
`331`	`331`	`err = controller.SubscribeExperimentEvents(context.Background(), &grpcClient, "")`
Original file line number	Diff line number	Diff line change
`@@ -189,7 +189,7 @@ func TestSubscribeModelEvents(t *testing.T) {`
`189`	`189`	`responses_subscribe_models: test.results,`
`190`	`190`	`}`
`191`	`191`	`}`
`192`		`- controller := newMockControllerClient(test.existing_resources...)`
	`192`	`+ controller := newMockControllerClient(false, test.existing_resources...)`
`193`	`193`	`err := controller.handleModels(context.Background(), &grpcClient, "")`
`194`	`194`	`g.Expect(err).To(BeNil())`
`195`	`195`	`err = controller.SubscribeModelEvents(context.Background(), &grpcClient, "")`
Original file line number	Diff line number	Diff line change
`@@ -419,7 +419,7 @@ func TestSubscribePipelineEvents(t *testing.T) {`
`419`	`419`	`responses_subscribe_pipelines: test.results,`
`420`	`420`	`}`
`421`	`421`	`}`
`422`		`- controller := newMockControllerClient(test.existing_resources...)`
	`422`	`+ controller := newMockControllerClient(false, test.existing_resources...)`
`423`	`423`	`err := controller.handlePipelines(context.Background(), &grpcClient, "")`
`424`	`424`	`g.Expect(err).To(BeNil())`
`425`	`425`	`err = controller.SubscribePipelineEvents(context.Background(), &grpcClient, "")`
Original file line number	Diff line number	Diff line change
`@@ -29,11 +29,12 @@ import (`
`29`	`29`	`)`
`30`	`30`
`31`	`31`	`func (s *SchedulerClient) ServerNotify(ctx context.Context, grpcClient scheduler.SchedulerClient, servers []v1alpha1.Server, isFirstSync bool) error {`
	`32`	`+ logger := s.logger.WithName("NotifyServer")`
`32`	`33`	`if len(servers) == 0 {`
	`34`	`+ logger.Info("No servers to notify scheduler, skipping")`
`33`	`35`	`return nil`
`34`	`36`	`}`
`35`	`37`
`36`		`- logger := s.logger.WithName("NotifyServer")`
`37`	`38`	`if grpcClient == nil {`
`38`	`39`	`// we assume that all servers are in the same namespace`
`39`	`40`	`namespace := servers[0].Namespace`
Original file line number	Diff line number	Diff line change
`@@ -244,7 +244,7 @@ func TestServerNotify(t *testing.T) {`
`244`	`244`	`grpcClient := mockSchedulerGrpcClient{`
`245`	`245`	`requests_servers: []*scheduler.ServerNotify{},`
`246`	`246`	`}`
`247`		`- controller := newMockControllerClient()`
	`247`	`+ controller := newMockControllerClient(false)`
`248`	`248`	`err := controller.ServerNotify(context.Background(), &grpcClient, test.servers, false)`
`249`	`249`	`g.Expect(err).To(BeNil())`
`250`	`250`
`@@ -470,7 +470,7 @@ func TestSubscribeServerEvents(t *testing.T) {`
`470`	`470`	`}`
`471`	`471`
`472`	`472`	`existing_resources := []client.Object{&test.existingServer}`
`473`		`- controller := newMockControllerClient(existing_resources...)`
	`473`	`+ controller := newMockControllerClient(false, existing_resources...)`
`474`	`474`	`err := controller.SubscribeServerEvents(context.Background(), &grpcClient, "")`
`475`	`475`	`g.Expect(err).To(BeNil())`
`476`	`476`