fix(scheduler): Scheduler wait on server connect

apis/mlops/scheduler/scheduler.proto

@@ -186,6 +186,7 @@ message ModelStatusRequest {
 
 message ServerNotifyRequest {
   repeated ServerNotify servers = 1;
+  bool isFirstSync = 2;
 }
 
 message ServerNotify {

k8s/helm-charts/seldon-core-v2-setup/templates/seldon-v2-components.yaml

@@ -903,6 +903,7 @@ spec:
         - --db-path=/mnt/scheduler/db
         - --allow-plaintxt=$(ALLOW_PLAINTXT)
         - --kafka-config-path=/mnt/kafka/kafka.json
+        - --scheduler-ready-timeout-seconds=$(SCHEDULER_READY_TIMEOUT_SECONDS)
         command:
         - /bin/scheduler
         env:
@@ -949,6 +950,8 @@ spec:
           value: '{{ .Values.security.envoy.ssl.downstream.server.caPath }}'
         - name: ENVOY_DOWNSTREAM_CLIENT_TLS_CA_LOCATION
           value: '{{ .Values.security.envoy.ssl.downstream.server.clientCaPath }}'
+        - name: SCHEDULER_READY_TIMEOUT_SECONDS
+          value: '{{ .Values.scheduler.schedulerReadyTimeoutSeconds }}'
         - name: ALLOW_PLAINTXT
           value: "true"
         - name: POD_NAMESPACE

k8s/helm-charts/seldon-core-v2-setup/values.yaml

@@ -213,6 +213,7 @@ scheduler:
     runAsUser: 1000
     runAsGroup: 1000
     runAsNonRoot: true
+  schedulerReadyTimeoutSeconds: 600
 
 serverConfig:
   terminationGracePeriodSeconds: 120

k8s/helm-charts/seldon-core-v2-setup/values.yaml.template

@@ -213,6 +213,7 @@ scheduler:
     runAsUser: 1000
     runAsGroup: 1000
     runAsNonRoot: true
+  schedulerReadyTimeoutSeconds: 600
 
 serverConfig:
   terminationGracePeriodSeconds: 120

k8s/kustomize/helm-components-sc/patch_scheduler.yaml

@@ -59,6 +59,8 @@ spec:
             value: '{{ .Values.security.envoy.ssl.downstream.server.caPath }}'
           - name: ENVOY_DOWNSTREAM_CLIENT_TLS_CA_LOCATION
             value: '{{ .Values.security.envoy.ssl.downstream.server.clientCaPath }}'
+          - name: SCHEDULER_READY_TIMEOUT_SECONDS
+            value: '{{ .Values.scheduler.schedulerReadyTimeoutSeconds }}'
     volumeClaimTemplates:
     - name: scheduler-state
       spec:

k8s/yaml/components.yaml

@@ -547,6 +547,7 @@ spec:
         - --db-path=/mnt/scheduler/db
         - --allow-plaintxt=$(ALLOW_PLAINTXT)
         - --kafka-config-path=/mnt/kafka/kafka.json
+        - --scheduler-ready-timeout-seconds=$(SCHEDULER_READY_TIMEOUT_SECONDS)
         command:
         - /bin/scheduler
         env:
@@ -590,6 +591,8 @@ spec:
           value: '/tmp/certs/dds/ca.crt'
         - name: ENVOY_DOWNSTREAM_CLIENT_TLS_CA_LOCATION
           value: '/tmp/certs/ddc/ca.crt'
+        - name: SCHEDULER_READY_TIMEOUT_SECONDS
+          value: '600'
         - name: ALLOW_PLAINTXT
           value: "true"
         - name: POD_NAMESPACE

operator/config/seldonconfigs/default.yaml

@@ -214,11 +214,14 @@ spec:
         - --db-path=/mnt/scheduler/db
         - --allow-plaintxt=$(ALLOW_PLAINTXT)
         - --kafka-config-path=/mnt/kafka/kafka.json
+        - --scheduler-ready-timeout-seconds=$(SCHEDULER_READY_TIMEOUT_SECONDS)
         command:
         - /bin/scheduler
         env:
         - name: ALLOW_PLAINTXT
           value: "true"
+        - name: SCHEDULER_READY_TIMEOUT_SECONDS
+          value: 600
         - name: POD_NAMESPACE
           valueFrom:
             fieldRef:

operator/controllers/mlops/server_controller.go

@@ -90,7 +90,7 @@ func (r *ServerReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctr
 		return reconcile.Result{}, nil
 	}
 
-	err := r.Scheduler.ServerNotify(ctx, nil, []v1alpha1.Server{*server})
+	err := r.Scheduler.ServerNotify(ctx, nil, []v1alpha1.Server{*server}, false)
 	if err != nil {
 		r.updateStatusFromError(ctx, logger, server, err)
 		return reconcile.Result{}, err

operator/scheduler/server.go

@@ -23,7 +23,7 @@ import (
 	"github.com/seldonio/seldon-core/operator/v2/apis/mlops/v1alpha1"
 )
 
-func (s *SchedulerClient) ServerNotify(ctx context.Context, grpcClient scheduler.SchedulerClient, servers []v1alpha1.Server) error {
+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 {
 		return nil
@@ -61,7 +61,8 @@ func (s *SchedulerClient) ServerNotify(ctx context.Context, grpcClient scheduler
 		})
 	}
 	request := &scheduler.ServerNotifyRequest{
-		Servers: requests,
+		Servers:     requests,
+		IsFirstSync: isFirstSync,
 	}
 	_, err := grpcClient.ServerNotify(
 		ctx,

operator/scheduler/server_test.go

@@ -161,7 +161,7 @@ func TestServerNotify(t *testing.T) {
 				requests_servers: []*scheduler.ServerNotify{},
 			}
 			controller := newMockControllerClient()
-			err := controller.ServerNotify(context.Background(), &grpcClient, test.servers)
+			err := controller.ServerNotify(context.Background(), &grpcClient, test.servers, false)
 			g.Expect(err).To(BeNil())
 
 			if len(test.servers) != 0 {

operator/scheduler/utils.go

@@ -129,7 +129,7 @@ func handleRegisteredServers(
 		return
 	}
 
-	if err := s.ServerNotify(ctx, grpcClient, serverList.Items); err != nil {
+	if err := s.ServerNotify(ctx, grpcClient, serverList.Items, true); err != nil {
 		s.logger.Error(err, "Failed to notify servers", "servers", serverList.Items)
 	}
 }

scheduler/all-host-network.yaml

@@ -167,6 +167,8 @@ services:
       - --disable-autoscaling
       - "--kafka-config-path"
       - "/mnt/config/kafka-host.json"
+      - "--scheduler-ready-timeout-seconds"
+      - ${SCHEDULER_READY_TIMEOUT_SECONDS}
     volumes:
       - type: bind
         source: ./config

scheduler/all-internal.yaml

@@ -206,6 +206,8 @@ services:
       - --disable-autoscaling
       - "--kafka-config-path"
       - "/mnt/config/kafka-internal.json"
+      - "--scheduler-ready-timeout-seconds"
+      - ${SCHEDULER_READY_TIMEOUT_SECONDS}
     volumes:
       - type: bind
         source: ./config

scheduler/cmd/scheduler/main.go

@@ -35,28 +35,34 @@ import (
 	"github.com/seldonio/seldon-core/scheduler/v2/pkg/store"
 	"github.com/seldonio/seldon-core/scheduler/v2/pkg/store/experiment"
 	"github.com/seldonio/seldon-core/scheduler/v2/pkg/store/pipeline"
+	"github.com/seldonio/seldon-core/scheduler/v2/pkg/synchroniser"
 	"github.com/seldonio/seldon-core/scheduler/v2/pkg/tracing"
 	"github.com/seldonio/seldon-core/scheduler/v2/pkg/util"
 )
 
 var (
-	envoyPort               uint
-	agentPort               uint
-	agentMtlsPort           uint
-	schedulerPort           uint
-	schedulerMtlsPort       uint
-	chainerPort             uint
-	namespace               string
-	pipelineGatewayHost     string
-	pipelineGatewayHttpPort int
-	pipelineGatewayGrpcPort int
-	logLevel                string
-	tracingConfigPath       string
-	dbPath                  string
-	nodeID                  string
-	allowPlaintxt           bool //scheduler server
-	autoscalingDisabled     bool
-	kafkaConfigPath         string
+	envoyPort                    uint
+	agentPort                    uint
+	agentMtlsPort                uint
+	schedulerPort                uint
+	schedulerMtlsPort            uint
+	chainerPort                  uint
+	namespace                    string
+	pipelineGatewayHost          string
+	pipelineGatewayHttpPort      int
+	pipelineGatewayGrpcPort      int
+	logLevel                     string
+	tracingConfigPath            string
+	dbPath                       string
+	nodeID                       string
+	allowPlaintxt                bool //scheduler server
+	autoscalingDisabled          bool
+	kafkaConfigPath              string
+	schedulerReadyTimeoutSeconds uint
+)
+
+const (
+	xDSWaitTimeout = time.Duration(10 * time.Second)
 )
 
 func init() {
@@ -98,12 +104,17 @@ func init() {
 
 	// Whether to enable autoscaling, default is true
 	flag.BoolVar(&autoscalingDisabled, "disable-autoscaling", false, "Disable autoscaling feature")
+
+	// Kafka config path
 	flag.StringVar(
 		&kafkaConfigPath,
 		"kafka-config-path",
 		"/mnt/config/kafka.json",
 		"Path to kafka configuration file",
 	)
+
+	// Timeout for scheduler to be ready
+	flag.UintVar(&schedulerReadyTimeoutSeconds, "scheduler-ready-timeout-seconds", 300, "Timeout for scheduler to be ready")
 }
 
 func getNamespace() string {
@@ -137,6 +148,8 @@ func main() {
 	logger.Infof("Setting log level to %s", logLevel)
 	logger.SetLevel(logIntLevel)
 
+	logger.Debugf("Scheduler ready timeout is set to %d seconds", schedulerReadyTimeoutSeconds)
+
 	done := make(chan bool, 1)
 
 	namespace = getNamespace()
@@ -149,16 +162,8 @@ func main() {
 	defer eventHub.Close()
 	go makeSignalHandler(logger, done)
 
-	// Start xDS server
 	// Create a cache
 	xdsCache := cache.NewSnapshotCache(false, cache.IDHash{}, logger)
-	ctx := context.Background()
-	srv := envoyServerControlPlaneV3.NewServer(ctx, xdsCache, nil)
-	xdsServer := envoyServer.NewXDSServer(srv, logger)
-	err = xdsServer.StartXDSServer(envoyPort)
-	if err != nil {
-		log.WithError(err).Fatalf("Failed to start envoy xDS server")
-	}
 
 	tracer, err := tracing.NewTraceProvider("seldon-scheduler", &tracingConfigPath, logger)
 	if err != nil {
@@ -167,6 +172,7 @@ func main() {
 		defer tracer.Stop()
 	}
 
+	// Create stores
 	ss := store.NewMemoryStore(logger, store.NewLocalSchedulerStore(), eventHub)
 	ps := pipeline.NewPipelineStore(logger, eventHub, ss)
 	es := experiment.NewExperimentServer(logger, eventHub, ss, ps)
@@ -178,19 +184,13 @@ func main() {
 		GrpcPort: pipelineGatewayGrpcPort,
 	}
 
+	// Create envoy incremental processor
 	_, err = processor.NewIncrementalProcessor(xdsCache, nodeID, logger, ss, es, ps, eventHub, &pipelineGatewayDetails, cleaner)
 	if err != nil {
 		log.WithError(err).Fatalf("Failed to create incremental processor")
 	}
 
-	sched := scheduler.NewSimpleScheduler(
-		logger,
-		ss,
-		scheduler.DefaultSchedulerConfig(ss),
-	)
-	logger.Infof("Autoscaling service is set to %t", !autoscalingDisabled)
-	as := agent.NewAgentServer(logger, ss, sched, eventHub, !autoscalingDisabled)
-
+	// scheduler <-> dataflow grpc
 	dataFlowLoadBalancer := util.NewRingLoadBalancer(1)
 	kafkaConfigMap, err := config.NewKafkaConfig(kafkaConfigPath)
 	if err != nil {
@@ -223,17 +223,57 @@ func main() {
 		log.Warn("Not running with scheduler local DB")
 	}
 
-	s := schedulerServer.NewSchedulerServer(logger, ss, es, ps, sched, eventHub)
+	// Setup synchroniser
+	var sync synchroniser.Synchroniser
+
+	if namespace == "" {
+		// running outside k8s
+		sync = synchroniser.NewSimpleSynchroniser(time.Duration(schedulerReadyTimeoutSeconds) * time.Second)
+	} else {
+		sync = synchroniser.NewServerBasedSynchroniser(eventHub, logger, time.Duration(schedulerReadyTimeoutSeconds)*time.Second)
+	}
+
+	// scheduler scheduling models service
+	sched := scheduler.NewSimpleScheduler(
+		logger,
+		ss,
+		scheduler.DefaultSchedulerConfig(ss),
+		sync,
+	)
+
+	// scheduler <-> controller grpc
+	s := schedulerServer.NewSchedulerServer(logger, ss, es, ps, sched, eventHub, sync)
 	err = s.StartGrpcServers(allowPlaintxt, schedulerPort, schedulerMtlsPort)
 	if err != nil {
 		log.WithError(err).Fatalf("Failed to start server gRPC servers")
 	}
 
+	// scheduler <-> agent  grpc
+	logger.Infof("Autoscaling service is set to %t", !autoscalingDisabled)
+	as := agent.NewAgentServer(logger, ss, sched, eventHub, !autoscalingDisabled)
 	err = as.StartGrpcServer(allowPlaintxt, agentPort, agentMtlsPort)
 	if err != nil {
 		log.WithError(err).Fatalf("Failed to start agent gRPC server")
 	}
 
+	// wait for model servers to be ready
+	sync.WaitReady()
+
+	// extra wait to allow routes state to get created
+	time.Sleep(xDSWaitTimeout)
+
+	// Start envoy xDS server, this is done after the scheduler is ready
+	// so that the xDS server can start sending valid updates to envoy.
+	ctx := context.Background()
+	srv := envoyServerControlPlaneV3.NewServer(ctx, xdsCache, nil)
+	xdsServer := envoyServer.NewXDSServer(srv, logger)
+	err = xdsServer.StartXDSServer(envoyPort)
+	if err != nil {
+		log.WithError(err).Fatalf("Failed to start envoy xDS server")
+	}
+
+	log.Info("Scheduler is ready")
+
 	// Wait for completion
 	<-done
 

scheduler/env.all

@@ -62,3 +62,5 @@ HODOMETER_RECEIVER_LOG_LEVEL=info
 KAFKA_MESSAGE_MAX_BYTES=1000000000
 
 OTEL_EXPORTER_OTLP_PROTOCOL=grpc
+
+SCHEDULER_READY_TIMEOUT_SECONDS=30

scheduler/pkg/agent/server.go

@@ -390,7 +390,12 @@ func (s *Server) Subscribe(request *pb.AgentSubscribeRequest, stream pb.AgentSer
 	}
 	s.mutex.Unlock()
 
-	err := s.syncMessage(request)
+	s.logger.Debugf("Add Server Replica %+v with config %+v", request, request.ReplicaConfig)
+	err := s.store.AddServerReplica(request)
+	if err != nil {
+		return err
+	}
+	err = s.scheduleModelsFromRequest(request)
 	if err != nil {
 		return err
 	}
@@ -421,13 +426,7 @@ func (s *Server) StopAgentStreams() {
 	}
 }
 
-func (s *Server) syncMessage(request *pb.AgentSubscribeRequest) error {
-	s.logger.Debugf("Add Server Replica %+v with config %+v", request, request.ReplicaConfig)
-	err := s.store.AddServerReplica(request)
-	if err != nil {
-		return err
-	}
-
+func (s *Server) scheduleModelsFromRequest(request *pb.AgentSubscribeRequest) error {
 	// we have to reschedule models that are loaded on the incoming agent
 	// this is because we can have a network glitch that causes the communication between the agent and the scheduler
 	// to drop and the scheduler loading the models on other servers.
@@ -439,7 +438,7 @@ func (s *Server) syncMessage(request *pb.AgentSubscribeRequest) error {
 		}
 	}
 
-	_, err = s.scheduler.ScheduleFailedModels()
+	_, err := s.scheduler.ScheduleFailedModels()
 	if err != nil {
 		return err
 	}