Refactor

Signed-off-by: win5923 <[email protected]>

Author: win5923

kubectl-plugin/pkg/cmd/job/job_submit.go

@@ -14,7 +14,6 @@ import (
 	"time"
 
 	"k8s.io/apimachinery/pkg/api/meta"
-	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/cli-runtime/pkg/genericclioptions"
 	"k8s.io/cli-runtime/pkg/genericiooptions"
@@ -34,9 +33,10 @@ import (
 )
 
 const (
-	dashboardAddr      = "http://localhost:8265"
-	clusterTimeout     = 120.0
-	portforwardtimeout = 60.0
+	dashboardAddr         = "http://localhost:8265"
+	clusterTimeout        = 120.0
+	portforwardtimeout    = 60.0
+	rayjobDeletionTimeout = 30.0
 )
 
 type SubmitJobOptions struct {
@@ -336,44 +336,18 @@ func (options *SubmitJobOptions) Run(ctx context.Context, factory cmdutil.Factor
 	}
 	fmt.Printf("Submitted RayJob %s.\n", options.RayJob.GetName())
 
-	// Continuously checks for Kubernetes events related to the RayJobDeletionPolicy.
-	// If an event indicates that the RayJobDeletionPolicy feature gate must be enabled, throw an error and delete the RayJob.
-	go func() {
-		ticker := time.NewTicker(5 * time.Second)
-		defer ticker.Stop()
-
-		for {
-			select {
-			case <-time.After(clusterTimeout * time.Second):
-				return
-			case <-ticker.C:
-				eventList, err := k8sClients.KubernetesClient().CoreV1().Events(*options.configFlags.Namespace).List(ctx, metav1.ListOptions{
-					FieldSelector: fmt.Sprintf("involvedObject.name=%s", options.RayJob.GetName()),
-				})
-				if err != nil {
-					fmt.Printf("Error listing events: %v\n", err)
-					return
-				}
-
-				// Check for error events related to RayJobDeletionPolicy feature gate
-				for _, event := range eventList.Items {
-					if strings.Contains(event.Message, "RayJobDeletionPolicy feature gate must be enabled to use the DeletionPolicy feature") {
-						if event.FirstTimestamp.Time.After(startTime) || event.LastTimestamp.Time.After(startTime) {
-							fmt.Printf("Deleting RayJob...\n")
-							err = k8sClients.RayClient().RayV1().RayJobs(*options.configFlags.Namespace).Delete(ctx, options.RayJob.GetName(), v1.DeleteOptions{})
-							if err != nil {
-								fmt.Printf("Failed to clean up Ray job: %v\n", err)
-							} else {
-								fmt.Printf("Cleaned Up RayJob: %s\n", options.RayJob.GetName())
-							}
-							log.Fatalf("%s", event.Message)
-						}
-					}
-				}
-				return
+	if options.deletionPolicy != "" {
+		err = k8sClients.WaitRayJobDeletionPolicyEnabled(ctx, *options.configFlags.Namespace, options.RayJob.Name, startTime, rayjobDeletionTimeout)
+		if err != nil {
+			fmt.Printf("Deleting RayJob...\n")
+			deleteErr := k8sClients.RayClient().RayV1().RayJobs(*options.configFlags.Namespace).Delete(ctx, options.RayJob.GetName(), v1.DeleteOptions{})
+			if deleteErr != nil {
+				return fmt.Errorf("Failed to clean up Ray job after time out.: %w", deleteErr)
 			}
+			fmt.Printf("Cleaned Up RayJob: %s\n", options.RayJob.GetName())
+			return fmt.Errorf("%w", err)
 		}
-	}()
+	}
 
 	if len(options.RayJob.GetName()) > 0 {
 		// Add timeout?

kubectl-plugin/pkg/cmd/version/version_test.go

@@ -110,6 +110,10 @@ func (c fakeClient) RayClient() rayclient.Interface {
 	return nil
 }
 
+func (c fakeClient) WaitRayJobDeletionPolicyEnabled(_ context.Context, _, _ string, _ time.Time, _ time.Duration) error {
+	return nil
+}
+
 // Tests the Run() step of the command and checks the output.
 func TestRayVersionRun(t *testing.T) {
 	testContext := "test-context"

kubectl-plugin/pkg/util/client/client.go

@@ -8,6 +8,7 @@ import (
 
 	dockerparser "github.com/novln/docker-parser"
 	"github.com/ray-project/kuberay/kubectl-plugin/pkg/util"
+	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/meta"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/fields"
@@ -25,6 +26,7 @@ type Client interface {
 	// GetRayHeadSvcName retrieves the name of RayHead service for the given RayCluster, RayJob, or RayService.
 	GetRayHeadSvcName(ctx context.Context, namespace string, resourceType util.ResourceType, name string) (string, error)
 	GetKubeRayOperatorVersion(ctx context.Context) (string, error)
+	WaitRayJobDeletionPolicyEnabled(ctx context.Context, namespace, name string, startTime time.Time, timeout time.Duration) error
 	WaitRayClusterProvisioned(ctx context.Context, namespace, name string, timeout time.Duration) error
 }
 
@@ -185,3 +187,40 @@ func (c *k8sClient) getRayHeadSvcNameByRayService(ctx context.Context, namespace
 	svcName := rayService.Status.ActiveServiceStatus.RayClusterStatus.Head.ServiceName
 	return svcName, nil
 }
+
+// WaitRayJobDeletionPolicyEnabled blocks until an event indicates that the RayJobDeletionPolicy feature gate must be enabled.
+func (c *k8sClient) WaitRayJobDeletionPolicyEnabled(ctx context.Context, namespace, name string, startTime time.Time, timeout time.Duration) error {
+	timeoutCtx, cancel := context.WithTimeout(ctx, timeout*time.Second)
+	defer cancel()
+
+	watcher, err := c.KubernetesClient().CoreV1().Events(namespace).Watch(ctx, metav1.ListOptions{
+		FieldSelector: fmt.Sprintf("involvedObject.name=%s", name),
+	})
+	if err != nil {
+		return fmt.Errorf("failed to watch events for RayJob %s in namespace %s: %w", name, namespace, err)
+	}
+	defer watcher.Stop()
+
+	for {
+		select {
+		case <-timeoutCtx.Done():
+			// If the RayJobDeletionPolicy feature event is not received within the timeout period, it is considered enabled.
+			return nil
+		case event := <-watcher.ResultChan():
+			if event.Type == watch.Error {
+				return fmt.Errorf("error watching events: %v", event.Object)
+			}
+
+			e, ok := event.Object.(*corev1.Event)
+			if !ok {
+				continue
+			}
+
+			if strings.Contains(e.Message, "RayJobDeletionPolicy feature gate must be enabled to use the DeletionPolicy feature") {
+				if e.FirstTimestamp.Time.After(startTime) || e.LastTimestamp.Time.After(startTime) {
+					return fmt.Errorf("%s", e.Message)
+				}
+			}
+		}
+	}
+}