From fe3e5a007d407cd63395b1ffd73ffa3a5a5832d2 Mon Sep 17 00:00:00 2001
From: kunwuluan <kunwuluan@gmail.com>
Date: Fri, 21 Apr 2023 17:26:53 +0800
Subject: [PATCH] delete all pods and services when job is suspended

---
 pkg/controller.v1/pytorch/controller.go | 16 ++++++++--------
 pkg/controller.v1/pytorch/job.go        |  6 ++++--
 pkg/controller.v1/pytorch/status.go     | 13 +++++++++++--
 3 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/pkg/controller.v1/pytorch/controller.go b/pkg/controller.v1/pytorch/controller.go
index 048ecd0e..898ec5b2 100644
--- a/pkg/controller.v1/pytorch/controller.go
+++ b/pkg/controller.v1/pytorch/controller.go
@@ -264,13 +264,11 @@ func (pc *PyTorchController) processNextWorkItem() bool {
 	}
 
 	// Wait until queuing annotation is removed
-	if pytorchJob.Annotations != nil {
-		if suspend, exist := pytorchJob.Annotations[suspendInQueue]; exist && suspend == "true" {
-			infoMsg := fmt.Sprintf("Annotation %s is found, operator will not process until removed", suspendInQueue)
-			logger.Info(infoMsg)
-			return true
-		}
-	}
+	// if isSuspend(pytorchJob) {
+	// 	infoMsg := fmt.Sprintf("Annotation %s is found, operator will not process until removed", suspendInQueue)
+	// 	logger.Info(infoMsg)
+	// 	return true
+	// }
 
 	// Sync PyTorchJob to mapch the actual state to this desired state.
 	forget, err := pc.syncHandler(key)
@@ -373,7 +371,9 @@ func (pc *PyTorchController) reconcilePyTorchJobs(job *pyv1.PyTorchJob) error {
 	}
 
 	// If the PyTorchJob is terminated, delete all pods and services.
-	if isSucceeded(job.Status) || isFailed(job.Status) {
+	if isSucceeded(job.Status) || isFailed(job.Status) || isSuspend(job) {
+		logger.Infof("job %v is inactive(maybe in succeeed, failed or suspend status)", job.Name)
+
 		if err := pc.deletePodsAndServices(job, pods, services); err != nil {
 			return err
 		}
diff --git a/pkg/controller.v1/pytorch/job.go b/pkg/controller.v1/pytorch/job.go
index 306f9991..116a7012 100644
--- a/pkg/controller.v1/pytorch/job.go
+++ b/pkg/controller.v1/pytorch/job.go
@@ -2,11 +2,12 @@ package pytorch
 
 import (
 	"fmt"
-	"github.com/kubeflow/pytorch-operator/pkg/util"
 	"reflect"
 	"strings"
 	"time"
 
+	"github.com/kubeflow/pytorch-operator/pkg/util"
+
 	log "github.com/sirupsen/logrus"
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -159,8 +160,9 @@ func (pc *PyTorchController) deletePodsAndServices(job *pyv1.PyTorchJob, pods []
 		return nil
 	}
 
+	isSuspend := isSuspend(job)
 	// Delete nothing when the cleanPodPolicy is None.
-	if *job.Spec.CleanPodPolicy == common.CleanPodPolicyNone {
+	if *job.Spec.CleanPodPolicy == common.CleanPodPolicyNone && !isSuspend {
 		return nil
 	}
 
diff --git a/pkg/controller.v1/pytorch/status.go b/pkg/controller.v1/pytorch/status.go
index 7259ab0a..5c038cd5 100644
--- a/pkg/controller.v1/pytorch/status.go
+++ b/pkg/controller.v1/pytorch/status.go
@@ -20,12 +20,12 @@ import (
 	"fmt"
 	"time"
 
-	"k8s.io/api/core/v1"
+	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
 
-	pyv1 "github.com/kubeflow/pytorch-operator/pkg/apis/pytorch/v1"
 	common "github.com/kubeflow/common/job_controller/api/v1"
+	pyv1 "github.com/kubeflow/pytorch-operator/pkg/apis/pytorch/v1"
 	pylogger "github.com/kubeflow/tf-operator/pkg/logger"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
@@ -212,6 +212,15 @@ func hasCondition(status common.JobStatus, condType common.JobConditionType) boo
 	return false
 }
 
+func isSuspend(pytorchJob *pyv1.PyTorchJob) bool {
+	if pytorchJob.Annotations != nil {
+		if suspend, exist := pytorchJob.Annotations[suspendInQueue]; exist && suspend == "true" {
+			return true
+		}
+	}
+	return false
+}
+
 func isSucceeded(status common.JobStatus) bool {
 	return hasCondition(status, common.JobSucceeded)
 }