Retain for job and metricsCollector (#572)

kubeflow · May 23, 2019 · cc5f367 · cc5f367
1 parent 9df08fd
commit cc5f367
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 5 deletions.
diff --git a/pkg/api/operators/apis/trial/v1alpha2/trial_types.go b/pkg/api/operators/apis/trial/v1alpha2/trial_types.go
@@ -33,9 +33,13 @@ type TrialSpec struct {
 	// and let the corresponding resource controller (e.g. tf-operator) handle
 	// the rest.
 	RunSpec string `json:"runSpec,omitempty"`
+	// Whether to retain the trial run object after completed.
+	RetainRun bool `json:"retainRun,omitempty"`
 
 	// Raw text for the metrics collector spec. This must be a CronJob object.
 	MetricsCollectorSpec string `json:"metricsCollectorSpec,omitempty"`
+	// Whether to retain the trial metrics collector CronJob object after completed.
+	RetainMetricsCollector bool `json:"retainMetricsCollector,omitempty"`
 }
 
 type TrialStatus struct {

diff --git a/pkg/controller/v1alpha2/experiment/experiment_util.go b/pkg/controller/v1alpha2/experiment/experiment_util.go
@@ -73,6 +73,13 @@ func (r *ReconcileExperiment) createTrialInstance(expInstance *experimentsv1alph
 	}
 	trial.Spec.MetricsCollectorSpec = mcSpec
 
+	if expInstance.Spec.TrialTemplate != nil {
+		trial.Spec.RetainRun = expInstance.Spec.TrialTemplate.Retain
+	}
+	if expInstance.Spec.MetricsCollectorSpec != nil {
+		trial.Spec.RetainMetricsCollector = expInstance.Spec.MetricsCollectorSpec.Retain
+	}
+
 	if err := r.Create(context.TODO(), trial); err != nil {
 		logger.Error(err, "Trial create error", "Trial name", trial.Name)
 		return err

diff --git a/pkg/controller/v1alpha2/trial/trial_controller.go b/pkg/controller/v1alpha2/trial/trial_controller.go
@@ -144,11 +144,6 @@ func (r *ReconcileTrial) Reconcile(request reconcile.Request) (reconcile.Result,
 
 	instance := original.DeepCopy()
 
-	if instance.IsCompleted() {
-
-		return reconcile.Result{}, nil
-
-	}
 	if !instance.IsCreated() {
 		//Trial not created in DB
 		if instance.Status.StartTime == nil {
@@ -245,6 +240,9 @@ func (r *ReconcileTrial) reconcileJob(instance *trialsv1alpha2.Trial, desiredJob
 	err = r.Get(context.TODO(), types.NamespacedName{Name: desiredJob.GetName(), Namespace: desiredJob.GetNamespace()}, deployedJob)
 	if err != nil {
 		if errors.IsNotFound(err) {
+			if instance.IsCompleted()  {
+				return nil, nil
+			}
 			logger.Info("Creating Job", "kind", kind,
 				"name", desiredJob.GetName())
 			err = r.Create(context.TODO(), desiredJob)
@@ -256,6 +254,15 @@ func (r *ReconcileTrial) reconcileJob(instance *trialsv1alpha2.Trial, desiredJob
 			logger.Error(err, "Trial Get error")
 			return nil, err
 		}
+	} else {
+		if instance.IsCompleted() && !instance.Spec.RetainRun {
+			if err = r.Delete(context.TODO(), desiredJob); err != nil {
+				logger.Error(err, "Delete job error")
+				return nil, err
+			} else {
+				return nil, nil
+			}
+		}
 	}
 
 	msg := "Trial is running"
@@ -311,6 +318,9 @@ func (r *ReconcileTrial) reconcileMetricsCollector(instance *trialsv1alpha2.Tria
 	}, deployedMetricsCollector)
 	if err != nil {
 		if errors.IsNotFound(err) {
+			if instance.IsCompleted()  {
+				return nil, nil
+			}
 			logger.Info("Creating Metrics Collector",
 				"name", desiredMetricsCollector.GetName(),
 				"namespace", desiredMetricsCollector.GetNamespace())
@@ -323,6 +333,15 @@ func (r *ReconcileTrial) reconcileMetricsCollector(instance *trialsv1alpha2.Tria
 			logger.Error(err, "Metrics Collector Get error")
 			return nil, err
 		}
+	} else {
+		if instance.IsCompleted() && !instance.Spec.RetainMetricsCollector {
+			if err = r.Delete(context.TODO(), desiredMetricsCollector); err != nil {
+				logger.Error(err, "Delete Metrics Collector error")
+				return nil, err
+			} else {
+				return nil, nil
+			}
+		}
 	}
 
 	return deployedMetricsCollector, nil