fix(scheduler): Controller to update the list of servers to scheduler on reconnect

apis/mlops/scheduler/scheduler.proto

@@ -185,14 +185,17 @@ message ModelStatusRequest {
 }
 
 message ServerNotifyRequest {
+  repeated ServerNotify servers = 1;
+}
+
+message ServerNotify {
   string name = 1;
   int32 expectedReplicas = 2;
   bool shared = 3;
   optional KubernetesMeta kubernetesMeta = 4;
 }
 
 message ServerNotifyResponse {
-
 }
 
 message ServerSubscriptionRequest {

operator/controllers/mlops/server_controller.go

@@ -29,6 +29,7 @@ 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"
@@ -89,7 +90,7 @@ func (r *ServerReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctr
 		return reconcile.Result{}, nil
 	}
 
-	err := r.Scheduler.ServerNotify(ctx, server)
+	err := r.Scheduler.ServerNotify(ctx, nil, []v1alpha1.Server{*server})
 	if err != nil {
 		r.updateStatusFromError(ctx, logger, server, err)
 		return reconcile.Result{}, err

operator/scheduler/pipeline_test.go

@@ -363,7 +363,7 @@ func TestSubscribePipelineEvents(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			// note that responses_pipelines is nill -> scheduler state is not existing
+			// note that if responses_pipelines is nil -> scheduler state is not existing
 			var grpcClient mockSchedulerGrpcClient
 			if !test.noSchedulerState {
 				grpcClient = mockSchedulerGrpcClient{
@@ -388,7 +388,7 @@ func TestSubscribePipelineEvents(t *testing.T) {
 				}
 			}
 
-			// check that we have reloaded the correct resources if the stata of the scheduler is not correct
+			// check that we have reloaded the correct resources if the state of the scheduler is not correct
 			if test.noSchedulerState {
 				activeResources := 0
 				for idx, req := range test.existing_resources {

operator/scheduler/server.go

@@ -23,40 +23,54 @@ import (
 	"github.com/seldonio/seldon-core/operator/v2/apis/mlops/v1alpha1"
 )
 
-func (s *SchedulerClient) ServerNotify(ctx context.Context, server *v1alpha1.Server) error {
+func (s *SchedulerClient) ServerNotify(ctx context.Context, grpcClient scheduler.SchedulerClient, servers []v1alpha1.Server) error {
 	logger := s.logger.WithName("NotifyServer")
-	conn, err := s.getConnection(server.Namespace)
-	if err != nil {
-		return err
+	if len(servers) == 0 {
+		return nil
 	}
-	grpcClient := scheduler.NewSchedulerClient(conn)
-
-	var replicas int32
-	if !server.ObjectMeta.DeletionTimestamp.IsZero() {
-		replicas = 0
-	} else if server.Spec.Replicas != nil {
-		replicas = *server.Spec.Replicas
-	} else {
-		replicas = 1
+
+	if grpcClient == nil {
+		// we assume that all servers are in the same namespace
+		namespace := servers[0].Namespace
+		conn, err := s.getConnection(namespace)
+		if err != nil {
+			return err
+		}
+		grpcClient = scheduler.NewSchedulerClient(conn)
 	}
 
+	var requests []*scheduler.ServerNotify
+	for _, server := range servers {
+		var replicas int32
+		if !server.ObjectMeta.DeletionTimestamp.IsZero() {
+			replicas = 0
+		} else if server.Spec.Replicas != nil {
+			replicas = *server.Spec.Replicas
+		} else {
+			replicas = 1
+		}
+
+		logger.Info("Notify server", "name", server.GetName(), "namespace", server.GetNamespace(), "replicas", replicas)
+		requests = append(requests, &scheduler.ServerNotify{
+			Name:             server.GetName(),
+			ExpectedReplicas: replicas,
+			KubernetesMeta: &scheduler.KubernetesMeta{
+				Namespace:  server.GetNamespace(),
+				Generation: server.GetGeneration(),
+			},
+		})
+	}
 	request := &scheduler.ServerNotifyRequest{
-		Name:             server.GetName(),
-		ExpectedReplicas: replicas,
-		KubernetesMeta: &scheduler.KubernetesMeta{
-			Namespace:  server.GetNamespace(),
-			Generation: server.GetGeneration(),
-		},
+		Servers: requests,
 	}
-	logger.Info("Notify server", "name", server.GetName(), "namespace", server.GetNamespace(), "replicas", replicas)
-	_, err = grpcClient.ServerNotify(
+	_, err := grpcClient.ServerNotify(
 		ctx,
 		request,
 		grpc_retry.WithMax(SchedulerConnectMaxRetries),
 		grpc_retry.WithBackoff(grpc_retry.BackoffExponential(SchedulerConnectBackoffScalar)),
 	)
 	if err != nil {
-		logger.Error(err, "Failed to send notify server to scheduler", "name", server.GetName(), "namespace", server.GetNamespace())
+		logger.Error(err, "Failed to send notify server to scheduler")
 		return err
 	}
 	return nil
@@ -75,6 +89,10 @@ func (s *SchedulerClient) SubscribeServerEvents(ctx context.Context, grpcClient
 	if err != nil {
 		return err
 	}
+
+	// on new reconnects we send a list of servers to the schedule
+	go handleRegisteredServers(ctx, namespace, s, grpcClient)
+
 	for {
 		event, err := stream.Recv()
 		if err != nil {