fix: Deal with deleted experiments when restoring from cache

apis/mlops/scheduler/storage.proto

@@ -11,4 +11,12 @@ message PipelineSnapshot {
   uint32 lastVersion = 2;
   repeated PipelineWithState versions = 3;
   bool deleted = 4;
-}
+}
+
+message ExperimentSnapshot {
+  Experiment experiment = 1;
+  // to mark the experiment as deleted, this is currently required as we persist all
+  // experiments in the local scheduler state (badgerdb) so that events can be replayed
+  // on restart, which would guard against lost events in communication.
+  bool deleted = 2;
+}

operator/controllers/mlops/experiment_controller.go

@@ -51,7 +51,7 @@ func (r *ExperimentReconciler) handleFinalizer(ctx context.Context, logger logr.
 	} else { // experiment is being deleted
 		if utils.ContainsStr(experiment.ObjectMeta.Finalizers, constants.ExperimentFinalizerName) {
 			// Handle unload in scheduler
-			if err, retry := r.Scheduler.StopExperiment(ctx, experiment); err != nil {
+			if retry, err := r.Scheduler.StopExperiment(ctx, experiment, nil); err != nil {
 				if retry {
 					return true, err
 				} else {
@@ -102,7 +102,7 @@ func (r *ExperimentReconciler) Reconcile(ctx context.Context, req ctrl.Request)
 		return reconcile.Result{}, err
 	}
 
-	err, retry := r.Scheduler.StartExperiment(ctx, experiment)
+	retry, err := r.Scheduler.StartExperiment(ctx, experiment, nil)
 	if err != nil {
 		r.updateStatusFromError(ctx, logger, experiment, err)
 		if retry {

operator/controllers/mlops/pipeline_controller.go

@@ -112,7 +112,7 @@ func (r *PipelineReconciler) Reconcile(ctx context.Context, req ctrl.Request) (c
 		return reconcile.Result{}, err
 	}
 
-	err, retry := r.Scheduler.LoadPipeline(ctx, pipeline)
+	retry, err := r.Scheduler.LoadPipeline(ctx, pipeline, nil)
 	if err != nil {
 		r.updateStatusFromError(ctx, logger, pipeline, retry, err)
 		if retry {

operator/pkg/utils/testing/client.go

@@ -16,5 +16,5 @@ import (
 )
 
 func NewFakeClient(scheme *runtime.Scheme, objs ...client.Object) client.Client {
-	return fake.NewClientBuilder().WithScheme(scheme).WithObjects(objs...).Build()
+	return fake.NewClientBuilder().WithScheme(scheme).WithObjects(objs...).WithStatusSubresource(objs...).Build()
 }

operator/scheduler/client.go

@@ -253,7 +253,7 @@ func (s *SchedulerClient) checkErrorRetryable(resource string, resourceName stri
 }
 
 func retryFn(
-	fn func(context context.Context, conn *grpc.ClientConn, namespace string) error,
+	fn func(context context.Context, grcpClient scheduler.SchedulerClient, namespace string) error,
 	conn *grpc.ClientConn, namespace string, logger logr.Logger,
 ) error {
 	logger.Info("RetryFn", "namespace", namespace)
@@ -262,7 +262,8 @@ func retryFn(
 	}
 	backOffExp := backoff.NewExponentialBackOff()
 	fnWithArgs := func() error {
-		return fn(context.Background(), conn, namespace)
+		grcpClient := scheduler.NewSchedulerClient(conn)
+		return fn(context.Background(), grcpClient, namespace)
 	}
 	err := backoff.RetryNotify(fnWithArgs, backOffExp, logFailure)
 	if err != nil {

operator/scheduler/experiment.go

@@ -26,31 +26,38 @@ import (
 	"github.com/seldonio/seldon-core/operator/v2/pkg/utils"
 )
 
-func (s *SchedulerClient) StartExperiment(ctx context.Context, experiment *v1alpha1.Experiment) (error, bool) {
+func (s *SchedulerClient) StartExperiment(ctx context.Context, experiment *v1alpha1.Experiment, grpcClient scheduler.SchedulerClient) (bool, error) {
 	logger := s.logger.WithName("StartExperiment")
-	conn, err := s.getConnection(experiment.Namespace)
-	if err != nil {
-		return err, true
+	var err error
+	if grpcClient == nil {
+		conn, err := s.getConnection(experiment.Namespace)
+		if err != nil {
+			return true, err
+		}
+		grpcClient = scheduler.NewSchedulerClient(conn)
 	}
-	grcpClient := scheduler.NewSchedulerClient(conn)
+
 	req := &scheduler.StartExperimentRequest{
 		Experiment: experiment.AsSchedulerExperimentRequest(),
 	}
 	logger.Info("Start", "experiment name", experiment.Name)
-	_, err = grcpClient.StartExperiment(
+	_, err = grpcClient.StartExperiment(
 		ctx,
 		req,
 		grpc_retry.WithMax(SchedulerConnectMaxRetries),
 		grpc_retry.WithBackoff(grpc_retry.BackoffExponential(SchedulerConnectBackoffScalar)),
 	)
-	return err, s.checkErrorRetryable(experiment.Kind, experiment.Name, err)
+	return s.checkErrorRetryable(experiment.Kind, experiment.Name, err), err
 }
 
-func (s *SchedulerClient) StopExperiment(ctx context.Context, experiment *v1alpha1.Experiment) (error, bool) {
+func (s *SchedulerClient) StopExperiment(ctx context.Context, experiment *v1alpha1.Experiment, conn *grpc.ClientConn) (bool, error) {
 	logger := s.logger.WithName("StopExperiment")
-	conn, err := s.getConnection(experiment.Namespace)
-	if err != nil {
-		return err, true
+	var err error
+	if conn == nil {
+		conn, err = s.getConnection(experiment.Namespace)
+		if err != nil {
+			return true, err
+		}
 	}
 	grcpClient := scheduler.NewSchedulerClient(conn)
 	req := &scheduler.StopExperimentRequest{
@@ -63,18 +70,33 @@ func (s *SchedulerClient) StopExperiment(ctx context.Context, experiment *v1alph
 		grpc_retry.WithMax(SchedulerConnectMaxRetries),
 		grpc_retry.WithBackoff(grpc_retry.BackoffExponential(SchedulerConnectBackoffScalar)),
 	)
-	return err, s.checkErrorRetryable(experiment.Kind, experiment.Name, err)
+	return s.checkErrorRetryable(experiment.Kind, experiment.Name, err), err
 }
 
 // namespace is not used in this function
-func (s *SchedulerClient) SubscribeExperimentEvents(ctx context.Context, conn *grpc.ClientConn, namespace string) error {
+func (s *SchedulerClient) SubscribeExperimentEvents(ctx context.Context, grcpClient scheduler.SchedulerClient, namespace string) error {
 	logger := s.logger.WithName("SubscribeExperimentEvents")
-	grcpClient := scheduler.NewSchedulerClient(conn)
 
 	stream, err := grcpClient.SubscribeExperimentStatus(ctx, &scheduler.ExperimentSubscriptionRequest{SubscriberName: "seldon manager"}, grpc_retry.WithMax(1))
 	if err != nil {
 		return err
 	}
+
+	// get experiments from the scheduler
+	// if there are no experiments in the scheduler state then we need to create them
+	// this is likely because of a restart of the scheduler that migrated the state
+	// to v2 (where we delete the experiments from the scheduler state)
+	numExperimentsFromScheduler, err := getNumExperimentsFromScheduler(ctx, grcpClient)
+	if err != nil {
+		return err
+	}
+	// if there are no experiments in the scheduler state then we need to create them if they exist in k8s
+	// also remove finalizers from experiments that are being deleted
+	if numExperimentsFromScheduler == 0 {
+		handleLoadedExperiments(ctx, namespace, s, grcpClient)
+		handlePendingDeleteExperiments(ctx, namespace, s)
+	}
+
 	for {
 		event, err := stream.Recv()
 		if err != nil {