Reconcile trial assignments by comparing suggestion and trials being …

…executed (#1831) * Reconcile trial assignments by comparing suggestion and trials being executed * Add a unit test for ReconcileSuggestions * Change add count
kubeflow · Apr 17, 2022 · bc5b82b · bc5b82b
1 parent e031b5e
commit bc5b82b
Show file tree

Hide file tree

Showing 2 changed files with 71 additions and 4 deletions.
diff --git a/pkg/controller.v1beta1/experiment/experiment_controller.go b/pkg/controller.v1beta1/experiment/experiment_controller.go
@@ -346,9 +346,8 @@ func (r *ReconcileExperiment) ReconcileTrials(instance *experimentsv1beta1.Exper
 func (r *ReconcileExperiment) createTrials(instance *experimentsv1beta1.Experiment, trialList []trialsv1beta1.Trial, addCount int32) error {
 
 	logger := log.WithValues("Experiment", types.NamespacedName{Name: instance.GetName(), Namespace: instance.GetNamespace()})
-	currentCount := int32(len(trialList))
 	logger.Info("Reconcile Suggestion", "addCount", addCount)
-	trials, err := r.ReconcileSuggestions(instance, currentCount, addCount)
+	trials, err := r.ReconcileSuggestions(instance, trialList, addCount)
 	if err != nil {
 		logger.Error(err, "Get suggestions error")
 		return err
@@ -457,9 +456,15 @@ func (r *ReconcileExperiment) deleteTrials(instance *experimentsv1beta1.Experime
 }
 
 // ReconcileSuggestions gets or creates the suggestion if needed.
-func (r *ReconcileExperiment) ReconcileSuggestions(instance *experimentsv1beta1.Experiment, currentCount, addCount int32) ([]suggestionsv1beta1.TrialAssignment, error) {
+func (r *ReconcileExperiment) ReconcileSuggestions(instance *experimentsv1beta1.Experiment, trialList []trialsv1beta1.Trial, addCount int32) ([]suggestionsv1beta1.TrialAssignment, error) {
 	logger := log.WithValues("Experiment", types.NamespacedName{Name: instance.GetName(), Namespace: instance.GetNamespace()})
 	var assignments []suggestionsv1beta1.TrialAssignment
+	currentCount := int32(len(trialList))
+	trialNames := map[string]bool{}
+	for _, trial := range trialList {
+		trialNames[trial.Name] = true
+	}
+
 	suggestionRequestsCount := currentCount + addCount
 
 	logger.Info("GetOrCreateSuggestion", "name", instance.Name, "Suggestion Requests", suggestionRequestsCount)
@@ -476,7 +481,11 @@ func (r *ReconcileExperiment) ReconcileSuggestions(instance *experimentsv1beta1.
 				suggestion := original.DeepCopy()
 				if len(suggestion.Status.Suggestions) > int(currentCount) {
 					suggestions := suggestion.Status.Suggestions
-					assignments = suggestions[currentCount:]
+					for _, suggestion := range suggestions {
+						if !trialNames[suggestion.Name] {
+							assignments = append(assignments, suggestion)
+						}
+					}
 				}
 				if suggestion.Spec.Requests != suggestionRequestsCount {
 					suggestion.Spec.Requests = suggestionRequestsCount

diff --git a/pkg/controller.v1beta1/experiment/experiment_controller_test.go b/pkg/controller.v1beta1/experiment/experiment_controller_test.go
@@ -96,6 +96,64 @@ func TestAdd(t *testing.T) {
 	g.Expect(Add(mgr)).NotTo(gomega.HaveOccurred())
 }
 
+func TestReconcileSuggestions(t *testing.T) {
+	g := gomega.NewGomegaWithT(t)
+
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+	mockSuggestion := suggestionmock.NewMockSuggestion(mockCtrl)
+
+	mockCtrl2 := gomock.NewController(t)
+	defer mockCtrl2.Finish()
+	mockGenerator := manifestmock.NewMockGenerator(mockCtrl2)
+
+	// Setup the Manager and Controller. Wrap the Controller Reconcile function so it writes each request to a
+	// channel when it is finished.
+	mgr, err := manager.New(cfg, manager.Options{MetricsBindAddress: "0"})
+	g.Expect(err).NotTo(gomega.HaveOccurred())
+
+	r := &ReconcileExperiment{
+		Client:     mgr.GetClient(),
+		scheme:     mgr.GetScheme(),
+		Suggestion: mockSuggestion,
+		Generator:  mockGenerator,
+		collector:  experimentUtil.NewExpsCollector(mgr.GetCache(), prometheus.NewRegistry()),
+		recorder:   mgr.GetEventRecorderFor(ControllerName),
+	}
+
+	suggestionRestartNo := newFakeSuggestion()
+	mockSuggestion.EXPECT().GetOrCreateSuggestion(gomock.Any(), gomock.Any()).Return(
+		suggestionRestartNo, nil).AnyTimes()
+	mockSuggestion.EXPECT().UpdateSuggestion(gomock.Any()).Return(nil).AnyTimes()
+
+	instance := newFakeInstance()
+
+	// ReconcileSuggestions should return the missing trial assignments
+	assignments, err := r.ReconcileSuggestions(
+		instance,
+		[]trialsv1beta1.Trial{
+			{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      trialName + "-1",
+					Namespace: namespace,
+				},
+				Spec: trialsv1beta1.TrialSpec{},
+			},
+			{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      trialName + "-3",
+					Namespace: namespace,
+				},
+				Spec: trialsv1beta1.TrialSpec{},
+			},
+		},
+		1,
+	)
+	g.Expect(err).NotTo(gomega.HaveOccurred())
+	g.Expect(len(assignments)).To(gomega.Equal(1))
+	g.Expect(assignments[0].Name).To(gomega.Equal(trialName + "-2"))
+}
+
 func TestReconcile(t *testing.T) {
 	g := gomega.NewGomegaWithT(t)