fix(ci): Merge changes from v2 for release 2.8.5 (4)

apis/mlops/scheduler/scheduler.proto

@@ -390,7 +390,12 @@ message ControlPlaneSubscriptionRequest {
 }
 
 message ControlPlaneResponse {
-
+  enum Event {
+    UNKNOWN_EVENT = 0;
+    SEND_SERVERS = 1;  // initial sync for the servers
+    SEND_RESOURCES = 2;  // send models / pipelines / experiments
+  }
+  Event event = 1;
 }
 
 // [END Messages]

operator/scheduler/client.go

@@ -142,34 +142,37 @@ func (s *SchedulerClient) startEventHanders(namespace string, conn *grpc.ClientC
 	}()
 }
 
-func (s *SchedulerClient) handleStateOnReconnect(context context.Context, grpcClient scheduler.SchedulerClient, namespace string) error {
-	// on new reconnects we send a list of servers to the schedule
-	err := s.handleRegisteredServers(context, grpcClient, namespace)
-	if err != nil {
-		s.logger.Error(err, "Failed to send registered server to scheduler")
-	}
-
-	if err == nil {
-		err = s.handleExperiments(context, grpcClient, namespace)
+func (s *SchedulerClient) handleStateOnReconnect(context context.Context, grpcClient scheduler.SchedulerClient, namespace string, operation scheduler.ControlPlaneResponse_Event) error {
+	switch operation {
+	case scheduler.ControlPlaneResponse_SEND_SERVERS:
+		// on new reconnects we send a list of servers to the schedule
+		err := s.handleRegisteredServers(context, grpcClient, namespace)
 		if err != nil {
-			s.logger.Error(err, "Failed to send experiments to scheduler")
+			s.logger.Error(err, "Failed to send registered server to scheduler")
 		}
-	}
-
-	if err == nil {
-		err = s.handlePipelines(context, grpcClient, namespace)
+		return err
+	case scheduler.ControlPlaneResponse_SEND_RESOURCES:
+		err := s.handleExperiments(context, grpcClient, namespace)
 		if err != nil {
-			s.logger.Error(err, "Failed to send pipelines to scheduler")
+			s.logger.Error(err, "Failed to send experiments to scheduler")
 		}
-	}
-
-	if err == nil {
-		err = s.handleModels(context, grpcClient, namespace)
-		if err != nil {
-			s.logger.Error(err, "Failed to send models to scheduler")
+		if err == nil {
+			err = s.handlePipelines(context, grpcClient, namespace)
+			if err != nil {
+				s.logger.Error(err, "Failed to send pipelines to scheduler")
+			}
 		}
+		if err == nil {
+			err = s.handleModels(context, grpcClient, namespace)
+			if err != nil {
+				s.logger.Error(err, "Failed to send models to scheduler")
+			}
+		}
+		return err
+	default:
+		s.logger.Info("Unknown operation", "operation", operation)
+		return fmt.Errorf("Unknown operation %v", operation)
 	}
-	return err
 }
 
 func (s *SchedulerClient) RemoveConnection(namespace string) {

operator/scheduler/control_plane.go

@@ -50,7 +50,7 @@ func (s *SchedulerClient) SubscribeControlPlaneEvents(ctx context.Context, grpcC
 		logger.Info("Received event to handle state", "event", event)
 
 		fn := func() error {
-			return s.handleStateOnReconnect(ctx, grpcClient, namespace)
+			return s.handleStateOnReconnect(ctx, grpcClient, namespace, event.GetEvent())
 		}
 		_, err = execWithTimeout(fn, execTimeOut)
 		if err != nil {

operator/scheduler/utils_test.go

@@ -109,7 +109,7 @@ func (s *mockSchedulerServerSubscribeGrpcClient) Recv() (*scheduler.ServerStatus
 // Control Plane subscribe mock grpc client
 
 type mockControlPlaneSubscribeGrpcClient struct {
-	sent bool
+	sent int
 	grpc.ClientStream
 }
 
@@ -120,9 +120,13 @@ func newMockControlPlaneSubscribeGrpcClient() *mockControlPlaneSubscribeGrpcClie
 }
 
 func (s *mockControlPlaneSubscribeGrpcClient) Recv() (*scheduler.ControlPlaneResponse, error) {
-	if !s.sent {
-		s.sent = true
-		return &scheduler.ControlPlaneResponse{}, nil
+	if s.sent == 0 {
+		s.sent++
+		return &scheduler.ControlPlaneResponse{Event: scheduler.ControlPlaneResponse_SEND_SERVERS}, nil
+	}
+	if s.sent == 1 {
+		s.sent++
+		return &scheduler.ControlPlaneResponse{Event: scheduler.ControlPlaneResponse_SEND_RESOURCES}, nil
 	}
 	return nil, io.EOF
 }

scheduler/cmd/scheduler/main.go

@@ -55,10 +55,11 @@ var (
 	tracingConfigPath            string
 	dbPath                       string
 	nodeID                       string
-	allowPlaintxt                bool //scheduler server
+	allowPlaintxt                bool // scheduler server
 	autoscalingDisabled          bool
 	kafkaConfigPath              string
 	schedulerReadyTimeoutSeconds uint
+	deletedResourceTTLSeconds    uint
 )
 
 const (
@@ -115,6 +116,9 @@ func init() {
 
 	// Timeout for scheduler to be ready
 	flag.UintVar(&schedulerReadyTimeoutSeconds, "scheduler-ready-timeout-seconds", 300, "Timeout for scheduler to be ready")
+
+	// This TTL is set in badger DB
+	flag.UintVar(&deletedResourceTTLSeconds, "deleted-resource-ttl-seconds", 86400, "TTL for deleted experiments and pipelines (in seconds)")
 }
 
 func getNamespace() string {
@@ -211,11 +215,11 @@ func main() {
 	// Do here after other services created so eventHub events will be handled on pipeline/experiment load
 	// If we start earlier events will be sent but not received by services that start listening "late" to eventHub
 	if dbPath != "" {
-		err := ps.InitialiseOrRestoreDB(dbPath)
+		err := ps.InitialiseOrRestoreDB(dbPath, deletedResourceTTLSeconds)
 		if err != nil {
 			log.WithError(err).Fatalf("Failed to initialise pipeline db at %s", dbPath)
 		}
-		err = es.InitialiseOrRestoreDB(dbPath)
+		err = es.InitialiseOrRestoreDB(dbPath, deletedResourceTTLSeconds)
 		if err != nil {
 			log.WithError(err).Fatalf("Failed to initialise experiment db at %s", dbPath)
 		}

scheduler/pkg/server/control_plane.go

@@ -23,8 +23,15 @@ func (s *SchedulerServer) SubscribeControlPlane(req *pb.ControlPlaneSubscription
 		return err
 	}
 
-	fin := make(chan bool)
+	s.synchroniser.WaitReady()
+
+	err = s.sendResourcesMarker(stream)
+	if err != nil {
+		logger.WithError(err).Errorf("Failed to send resources marker to %s", req.GetSubscriberName())
+		return err
+	}
 
+	fin := make(chan bool)
 	s.controlPlaneStream.mu.Lock()
 	s.controlPlaneStream.streams[stream] = &ControlPlaneSubsription{
 		name:   req.GetSubscriberName(),
@@ -61,11 +68,20 @@ func (s *SchedulerServer) StopSendControlPlaneEvents() {
 // this is to mark the initial start of a new stream (at application level)
 // as otherwise the other side sometimes doesnt know if the scheduler has established a new stream explicitly
 func (s *SchedulerServer) sendStartServerStreamMarker(stream pb.Scheduler_SubscribeControlPlaneServer) error {
-	ssr := &pb.ControlPlaneResponse{}
+	ssr := &pb.ControlPlaneResponse{Event: pb.ControlPlaneResponse_SEND_SERVERS}
 	_, err := sendWithTimeout(func() error { return stream.Send(ssr) }, s.timeout)
 	if err != nil {
 		return err
 	}
+	return nil
+}
 
+// this is to mark a stage to send resources (models, pipelines, experiments) from the controller
+func (s *SchedulerServer) sendResourcesMarker(stream pb.Scheduler_SubscribeControlPlaneServer) error {
+	ssr := &pb.ControlPlaneResponse{Event: pb.ControlPlaneResponse_SEND_RESOURCES}
+	_, err := sendWithTimeout(func() error { return stream.Send(ssr) }, s.timeout)
+	if err != nil {
+		return err
+	}
 	return nil
 }

scheduler/pkg/server/control_plane_test.go

@@ -10,15 +10,22 @@ the Change License after the Change Date as each is defined in accordance with t
 package server
 
 import (
+	"context"
+	"fmt"
 	"testing"
 	"time"
 
 	. "github.com/onsi/gomega"
 	log "github.com/sirupsen/logrus"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
 
 	pb "github.com/seldonio/seldon-core/apis/go/v2/mlops/scheduler"
 
+	"github.com/seldonio/seldon-core/scheduler/v2/pkg/coordinator"
+	"github.com/seldonio/seldon-core/scheduler/v2/pkg/internal/testing_utils"
 	"github.com/seldonio/seldon-core/scheduler/v2/pkg/store"
+	"github.com/seldonio/seldon-core/scheduler/v2/pkg/synchroniser"
 )
 
 func TestStartServerStream(t *testing.T) {
@@ -67,7 +74,91 @@ func TestStartServerStream(t *testing.T) {
 				}
 
 				g.Expect(msr).ToNot(BeNil())
+				g.Expect(msr.Event).To(Equal(pb.ControlPlaneResponse_SEND_SERVERS))
 			}
+
+			err = test.server.sendResourcesMarker(stream)
+			if test.err {
+				g.Expect(err).ToNot(BeNil())
+			} else {
+				g.Expect(err).To(BeNil())
+
+				var msr *pb.ControlPlaneResponse
+				select {
+				case next := <-stream.msgs:
+					msr = next
+				default:
+					t.Fail()
+				}
+
+				g.Expect(msr).ToNot(BeNil())
+				g.Expect(msr.Event).To(Equal(pb.ControlPlaneResponse_SEND_RESOURCES))
+			}
+		})
+	}
+}
+
+func TestSubscribeControlPlane(t *testing.T) {
+	log.SetLevel(log.DebugLevel)
+	g := NewGomegaWithT(t)
+
+	type test struct {
+		name string
+	}
+	tests := []test{
+		{
+			name: "simple",
+		},
+	}
+
+	getStream := func(context context.Context, port int) (*grpc.ClientConn, pb.Scheduler_SubscribeControlPlaneClient) {
+		conn, _ := grpc.NewClient(fmt.Sprintf(":%d", port), grpc.WithTransportCredentials(insecure.NewCredentials()))
+		grpcClient := pb.NewSchedulerClient(conn)
+		client, _ := grpcClient.SubscribeControlPlane(
+			context,
+			&pb.ControlPlaneSubscriptionRequest{SubscriberName: "dummy"},
+		)
+		return conn, client
+	}
+
+	createTestScheduler := func() *SchedulerServer {
+		logger := log.New()
+		logger.SetLevel(log.WarnLevel)
+
+		eventHub, err := coordinator.NewEventHub(logger)
+		g.Expect(err).To(BeNil())
+
+		sync := synchroniser.NewSimpleSynchroniser(time.Duration(10 * time.Millisecond))
+
+		s := NewSchedulerServer(logger, nil, nil, nil, nil, eventHub, sync)
+		sync.Signals(1)
+
+		return s
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			server := createTestScheduler()
+			port, err := testing_utils.GetFreePortForTest()
+			if err != nil {
+				t.Fatal(err)
+			}
+			err = server.startServer(uint(port), false)
+			if err != nil {
+				t.Fatal(err)
+			}
+			time.Sleep(100 * time.Millisecond)
+
+			conn, client := getStream(context.Background(), port)
+
+			msg, _ := client.Recv()
+			g.Expect(msg.GetEvent()).To(Equal(pb.ControlPlaneResponse_SEND_SERVERS))
+
+			msg, _ = client.Recv()
+			g.Expect(msg.Event).To(Equal(pb.ControlPlaneResponse_SEND_RESOURCES))
+
+			conn.Close()
+			server.StopSendControlPlaneEvents()
 		})
 	}
 }