fix: Deal with deleted experiments when restoring from cache (SeldonIO#5726)

* remove dead code path

* skip restoring an experiment if there is an error.

* add a note that we do not validate pipelines when we restore them

* deal with deleted experiments on restore

* use a call back for deleted experiments

* add test for multiple experiments in db

* update store to mark deleted experiments

* add experiment get (for testing)

* Add active field in experiment protos

* add deleted instead of active

* make deleted field not optional

* handle deleted in controller for experiments

* fix restoring of experiments

* add compare for the entire proto

* add pipeline get from db helper (for testing)

* add test for db check after adding pipeline

* add testing coverage

* revert changes to operator as they are not required anymore

* add experiment db migration helper

* reinstate delete helper for DBss

* simplify get from DB

* add testing for delete from db

* add scafolding to get the version from the (experiment) db

* use `dropall` helper to clear db

* optimize how to migrate to the new version

* refactor common code to utils

* add version to pipelinedb

* add helper to get the number of experiments

* add helper to count the number of expriments from the scheduler

* handle load experiments on startup of controller.

* remove finalizers for experiments if there are no experiments from scheduler

* simplify removing finalizers for experiments

* add tests for experiments utils

* refactor model handlers and add tests

* add pipeline handlers and tests

* add helper to get pipeline status from scheduler

* Add status subresoruce to fake client

* pass grpc client instead of conn to subscriptions

* add test for pipeline subscription

* add a test for pipeline termination

* add experiment tests

* add test case for pipelines

* check pipelineready status

* add a test case when pipeline is removed

* add note about expected state

* add 2 extra test cases to cover when the resource doesn exist in k8s

* deal with errors

* update copyright

* revert back upgrade to protoc

* use grpc client in StopExperiment instead of the underlying connection

* fix mis-spelled grpc vars in controller

* use Be[True/False]Because pattern in experiments and pipelines tests

* rename function for current db migration

* fix mispelling pipeline->experiment

Author: Sherif Akoush

apis/go/mlops/scheduler/storage.pb.go

@@ -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;
+}

apis/mlops/scheduler/storage.proto

@@ -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/experiment_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/controllers/mlops/pipeline_controller.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/pkg/utils/testing/client.go

@@ -132,9 +132,9 @@ func (s *SchedulerClient) RemoveConnection(namespace string) {
 
 // A smoke test allows us to quickly check if we actually have a functional grpc connection to the scheduler
 func (s *SchedulerClient) smokeTestConnection(conn *grpc.ClientConn) error {
-	grcpClient := scheduler.NewSchedulerClient(conn)
+	grpcClient := scheduler.NewSchedulerClient(conn)
 
-	stream, err := grcpClient.SubscribeModelStatus(context.TODO(), &scheduler.ModelSubscriptionRequest{SubscriberName: "seldon manager"}, grpc_retry.WithMax(1))
+	stream, err := grpcClient.SubscribeModelStatus(context.TODO(), &scheduler.ModelSubscriptionRequest{SubscriberName: "seldon manager"}, grpc_retry.WithMax(1))
 	if err != nil {
 		return err
 	}
@@ -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, grpcClient 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)
+		grpcClient := scheduler.NewSchedulerClient(conn)
+		return fn(context.Background(), grpcClient, namespace)
 	}
 	err := backoff.RetryNotify(fnWithArgs, backOffExp, logFailure)
 	if err != nil {

operator/scheduler/client.go

@@ -14,7 +14,6 @@ import (
 	"io"
 
 	grpc_retry "github.com/grpc-ecosystem/go-grpc-middleware/retry"
-	"google.golang.org/grpc"
 	v1 "k8s.io/api/core/v1"
 	"k8s.io/client-go/util/retry"
 	"sigs.k8s.io/controller-runtime/pkg/client"
@@ -26,55 +25,77 @@ 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, grpcClient scheduler.SchedulerClient) (bool, error) {
 	logger := s.logger.WithName("StopExperiment")
-	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.StopExperimentRequest{
 		Name: experiment.Name,
 	}
 	logger.Info("Stop", "experiment name", experiment.Name)
-	_, err = grcpClient.StopExperiment(
+	_, err = grpcClient.StopExperiment(
 		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
 }
 
 // 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, grpcClient 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))
+	stream, err := grpcClient.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, grpcClient)
+	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, grpcClient)
+		handlePendingDeleteExperiments(ctx, namespace, s)
+	}
+
 	for {
 		event, err := stream.Recv()
 		if err != nil {