Replaced 30 seconds time buffer with polling logic to wait until rayjob exists

Author: abhijeet-dhumal

tests/odh/mnist_ray_test.go

@@ -129,50 +129,56 @@ func mnistRay(t *testing.T, numGpus int) {
 			),
 		)
 
-	time.Sleep(30 * time.Second)
-
 	// Fetch created raycluster
 	rayClusterName := "mnisttest"
+	// Wait until raycluster is up and running
 	rayCluster, err := test.Client().Ray().RayV1().RayClusters(namespace.Name).Get(test.Ctx(), rayClusterName, metav1.GetOptions{})
 	test.Expect(err).ToNot(HaveOccurred())
 
 	// Initialise raycluster client to interact with raycluster to get rayjob details using REST-API
 	dashboardUrl := GetDashboardUrl(test, namespace, rayCluster)
 	rayClusterClientConfig := RayClusterClientConfig{Address: dashboardUrl.String(), Client: nil, InsecureSkipVerify: true}
 	rayClient, err := NewRayClusterClient(rayClusterClientConfig, test.Config().BearerToken)
-	if err != nil {
-		test.T().Errorf("%s", err)
-	}
+	test.Expect(err).ToNot(HaveOccurred(), fmt.Sprintf("Failed to create new raycluster client: %s", err))
 
+	// wait until rayjob exists
+	test.Eventually(func() []RayJobDetailsResponse {
+		rayJobs, err := rayClient.GetJobs()
+		test.Expect(err).ToNot(HaveOccurred(), fmt.Sprintf("Failed to fetch ray-jobs : %s", err))
+		return *rayJobs
+	}, TestTimeoutMedium, 2*time.Second).Should(HaveLen(1), "Ray job not found")
+
+	// Get test job-id
 	jobID := GetTestJobId(test, rayClient, dashboardUrl.Host)
-	test.Expect(jobID).ToNot(Equal(nil))
+	test.Expect(jobID).ToNot(BeEmpty())
 
 	// Wait for the job to be succeeded or failed
 	var rayJobStatus string
-	fmt.Printf("Waiting for job to be Succeeded...\n")
+	test.T().Logf("Waiting for job to be Succeeded...\n")
 	test.Eventually(func() string {
 		resp, err := rayClient.GetJobDetails(jobID)
-		test.Expect(err).ToNot(HaveOccurred())
+		test.Expect(err).ToNot(HaveOccurred(), fmt.Sprintf("Failed to get job details :%s", err))
 		rayJobStatusVal := resp.Status
 		if rayJobStatusVal == "SUCCEEDED" || rayJobStatusVal == "FAILED" {
-			fmt.Printf("JobStatus : %s\n", rayJobStatusVal)
+			test.T().Logf("JobStatus - %s\n", rayJobStatusVal)
 			rayJobStatus = rayJobStatusVal
 			return rayJobStatus
 		}
 		if rayJobStatus != rayJobStatusVal && rayJobStatusVal != "SUCCEEDED" {
-			fmt.Printf("JobStatus : %s...\n", rayJobStatusVal)
+			test.T().Logf("JobStatus - %s...\n", rayJobStatusVal)
 			rayJobStatus = rayJobStatusVal
 		}
 		return rayJobStatus
 	}, TestTimeoutDouble, 3*time.Second).Should(Or(Equal("SUCCEEDED"), Equal("FAILED")), "Job did not complete within the expected time")
-	test.Expect(rayJobStatus).To(Equal("SUCCEEDED"), "RayJob failed !")
-
 	// Store job logs in output directory
 	WriteRayJobAPILogs(test, rayClient, jobID)
 
+	// Assert ray-job status after job execution
+	test.Expect(rayJobStatus).To(Equal("SUCCEEDED"), "RayJob failed !")
+
 	// Make sure the RayCluster finishes and is deleted
 	test.Eventually(RayClusters(test, namespace.Name), TestTimeoutLong).
-		Should(HaveLen(0))
+		Should(BeEmpty())
 }
 
 func readMnistScriptTemplate(test Test, filePath string) []byte {

tests/odh/mnist_raytune_hpo_test.go

@@ -129,48 +129,55 @@ func mnistRayTuneHpo(t *testing.T, numGpus int) {
 				ContainElement(WithTransform(KueueWorkloadAdmitted, BeTrueBecause("Workload failed to be admitted"))),
 			),
 		)
-	time.Sleep(30 * time.Second)
 
 	// Fetch created raycluster
 	rayClusterName := "mnisthpotest"
+	// Wait until raycluster is up and running
 	rayCluster, err := test.Client().Ray().RayV1().RayClusters(namespace.Name).Get(test.Ctx(), rayClusterName, metav1.GetOptions{})
 	test.Expect(err).ToNot(HaveOccurred())
 
 	// Initialise raycluster client to interact with raycluster to get rayjob details using REST-API
 	dashboardUrl := GetDashboardUrl(test, namespace, rayCluster)
 	rayClusterClientConfig := RayClusterClientConfig{Address: dashboardUrl.String(), Client: nil, InsecureSkipVerify: true}
 	rayClient, err := NewRayClusterClient(rayClusterClientConfig, test.Config().BearerToken)
-	if err != nil {
-		test.T().Errorf("%s", err)
-	}
+	test.Expect(err).ToNot(HaveOccurred(), fmt.Sprintf("Failed to create new raycluster client: %s", err))
+
+	// Wait until the rayjob is created and running
+	test.Eventually(func() []RayJobDetailsResponse {
+		rayJobs, err := rayClient.GetJobs()
+		test.Expect(err).ToNot(HaveOccurred(), fmt.Sprintf("Failed to fetch ray-jobs : %s", err))
+		return *rayJobs
+	}, TestTimeoutMedium, 2*time.Second).Should(HaveLen(1), "Ray job not found")
 
+	// Get rayjob-ID
 	jobID := GetTestJobId(test, rayClient, dashboardUrl.Host)
-	test.Expect(jobID).ToNot(Equal(nil))
+	test.Expect(jobID).ToNot(BeEmpty())
 
-	// Wait for the job to be succeeded or failed
+	// Wait for the job to either succeed or fail
 	var rayJobStatus string
-	fmt.Printf("Waiting for job to be Succeeded...\n")
+	test.T().Logf("Waiting for job to be Succeeded...\n")
 	test.Eventually(func() string {
 		resp, err := rayClient.GetJobDetails(jobID)
-		test.Expect(err).ToNot(HaveOccurred())
+		test.Expect(err).ToNot(HaveOccurred(), fmt.Sprintf("Failed to get job details :%s", err))
 		rayJobStatusVal := resp.Status
 		if rayJobStatusVal == "SUCCEEDED" || rayJobStatusVal == "FAILED" {
-			fmt.Printf("JobStatus : %s\n", rayJobStatusVal)
+			test.T().Logf("JobStatus - %s\n", rayJobStatusVal)
 			rayJobStatus = rayJobStatusVal
 			return rayJobStatus
 		}
 		if rayJobStatus != rayJobStatusVal && rayJobStatusVal != "SUCCEEDED" {
-			fmt.Printf("JobStatus : %s...\n", rayJobStatusVal)
+			test.T().Logf("JobStatus - %s...\n", rayJobStatusVal)
 			rayJobStatus = rayJobStatusVal
 		}
 		return rayJobStatus
 	}, TestTimeoutDouble, 3*time.Second).Should(Or(Equal("SUCCEEDED"), Equal("FAILED")), "Job did not complete within the expected time")
-	test.Expect(rayJobStatus).To(Equal("SUCCEEDED"), "RayJob failed !")
-
 	// Store job logs in output directory
 	WriteRayJobAPILogs(test, rayClient, jobID)
 
+	// Assert ray-job status after job execution
+	test.Expect(rayJobStatus).To(Equal("SUCCEEDED"), "RayJob failed !")
+
 	// Make sure the RayCluster finishes and is deleted
 	test.Eventually(RayClusters(test, namespace.Name), TestTimeoutLong).
-		Should(HaveLen(0))
+		Should(BeEmpty())
 }

tests/odh/support.go

@@ -18,7 +18,6 @@ package odh
 
 import (
 	"embed"
-	"net/http"
 	"net/url"
 
 	. "github.com/onsi/gomega"
@@ -41,7 +40,6 @@ func ReadFile(t support.Test, fileName string) []byte {
 
 func GetDashboardUrl(test support.Test, namespace *v1.Namespace, rayCluster *rayv1.RayCluster) *url.URL {
 	dashboardName := "ray-dashboard-" + rayCluster.Name
-	test.T().Logf("Raycluster created : %s\n", rayCluster.Name)
 	route := GetRoute(test, namespace.Name, dashboardName)
 	hostname := route.Status.Ingress[0].Host
 	dashboardUrl, _ := url.Parse("https://" + hostname)
@@ -51,12 +49,6 @@ func GetDashboardUrl(test support.Test, namespace *v1.Namespace, rayCluster *ray
 }
 
 func GetTestJobId(test Test, rayClient RayClusterClient, hostName string) string {
-	listJobsReq, err := http.NewRequest("GET", "https://"+hostName+"/api/jobs/", nil)
-	if err != nil {
-		test.T().Errorf("failed to do get request: %s\n", err)
-	}
-	listJobsReq.Header.Add("Authorization", "Bearer "+test.Config().BearerToken)
-
 	allJobsData, err := rayClient.GetJobs()
 	test.Expect(err).ToNot(HaveOccurred())