Feature/cron scheduling rayjob 2426 (#3836)

* Add support for cron sheduling for ray jobs

---------

Signed-off-by: Kenny Han <[email protected]>

Author: DW-Han

docs/reference/api.md

@@ -248,6 +248,7 @@ _Appears in:_
 | `jobId` _string_ | If jobId is not set, a new jobId will be auto-generated. |  |  |
 | `submissionMode` _[JobSubmissionMode](#jobsubmissionmode)_ | SubmissionMode specifies how RayJob submits the Ray job to the RayCluster.<br />In "K8sJobMode", the KubeRay operator creates a submitter Kubernetes Job to submit the Ray job.<br />In "HTTPMode", the KubeRay operator sends a request to the RayCluster to create a Ray job.<br />In "InteractiveMode", the KubeRay operator waits for a user to submit a job to the Ray cluster. | K8sJobMode |  |
 | `entrypointResources` _string_ | EntrypointResources specifies the custom resources and quantities to reserve for the<br />entrypoint command. |  |  |
+| `schedule` _string_ | Schedule specifies a cron like string for scheduling Ray jobs.<br />When shutdownAfterJobFinishes is set to true, a new cluster is provisioned<br />per scheduled job, otherwise the job is scheduled on an existing cluster. |  |  |
 | `entrypointNumCpus` _float_ | EntrypointNumCpus specifies the number of cpus to reserve for the entrypoint command. |  |  |
 | `entrypointNumGpus` _float_ | EntrypointNumGpus specifies the number of gpus to reserve for the entrypoint command. |  |  |
 | `ttlSecondsAfterFinished` _integer_ | TTLSecondsAfterFinished is the TTL to clean up RayCluster.<br />It's only working when ShutdownAfterJobFinishes set to true. | 0 |  |

go.mod

@@ -87,6 +87,7 @@ require (
 	github.com/prometheus/client_model v0.6.1 // indirect
 	github.com/prometheus/common v0.62.0 // indirect
 	github.com/prometheus/procfs v0.15.1 // indirect
+	github.com/robfig/cron/v3 v3.0.1 // indirect
 	github.com/russross/blackfriday/v2 v2.1.0 // indirect
 	github.com/x448/float16 v0.8.4 // indirect
 	github.com/xlab/treeprint v1.2.0 // indirect

go.sum

@@ -54,6 +54,8 @@ const (
 	JobDeploymentStatusSuspended    JobDeploymentStatus = "Suspended"
 	JobDeploymentStatusRetrying     JobDeploymentStatus = "Retrying"
 	JobDeploymentStatusWaiting      JobDeploymentStatus = "Waiting"
+	JobDeploymentStatusScheduling   JobDeploymentStatus = "Scheduling"
+	JobDeploymentStatusScheduled    JobDeploymentStatus = "Scheduled"
 )
 
 // IsJobDeploymentTerminal returns true if the given JobDeploymentStatus
@@ -181,6 +183,11 @@ type RayJobSpec struct {
 	// entrypoint command.
 	// +optional
 	EntrypointResources string `json:"entrypointResources,omitempty"`
+	// Schedule specifies a cron like string for scheduling Ray jobs.
+	// When shutdownAfterJobFinishes is set to true, a new cluster is provisioned
+	// per scheduled job, otherwise the job is scheduled on an existing cluster.
+	// +optional
+	Schedule string `json:"schedule,omitempty"`
 	// EntrypointNumCpus specifies the number of cpus to reserve for the entrypoint command.
 	// +optional
 	EntrypointNumCpus float32 `json:"entrypointNumCpus,omitempty"`
@@ -233,6 +240,9 @@ type RayJobStatus struct {
 	// or the submitter Job has failed.
 	// +optional
 	EndTime *metav1.Time `json:"endTime,omitempty"`
+	// lastScheduledTime is the last time the job was successfully scheduled.
+	// +optional
+	LastScheduleTime *metav1.Time `json:"lastScheduleTime,omitempty"`
 	// Succeeded is the number of times this job succeeded.
 	// +kubebuilder:default:=0
 	// +optional

helm-chart/kuberay-operator/crds/ray.io_rayjobs.yaml

@@ -0,0 +1,121 @@
+apiVersion: ray.io/v1
+kind: RayJob
+metadata:
+  name: rayjob-schedule
+spec:
+  # schedule specifires a cron scheduling string telling the rayjob when to start schedule and run new jobs
+  # Here it runs at every 5 minutes of every hour of every day of every week of every year
+  schedule: "*/5 * * * *"
+
+  entrypoint: python /home/ray/samples/sample_code.py
+
+  # shutdownAfterJobFinishes specifies whether the RayCluster should be deleted after the RayJob finishes. Default is false.
+  # NOTE that the expected behavior with schedule is that the cluster will be deleted and recreated at each schedule if set to true, and it will keep using the same cluster otherwise
+  shutdownAfterJobFinishes: true
+
+  runtimeEnvYAML: |
+    pip:
+      - requests==2.26.0
+      - pendulum==2.1.2
+    env_vars:
+      counter_name: "test_counter"
+
+
+  rayClusterSpec:
+    rayVersion: '2.46.0'
+    headGroupSpec:
+      rayStartParams: {}
+      template:
+        spec:
+          containers:
+          - name: ray-head
+            image: rayproject/ray:2.46.0
+            ports:
+            - containerPort: 6379
+              name: gcs-server
+            - containerPort: 8265
+              name: dashboard
+            - containerPort: 10001
+              name: client
+            resources:
+              limits:
+                cpu: "1"
+              requests:
+                cpu: "200m"
+            volumeMounts:
+            - mountPath: /home/ray/samples
+              name: code-sample
+          volumes:
+          - name: code-sample
+            configMap:
+              name: ray-job-code-sample
+              items:
+              - key: sample_code.py
+                path: sample_code.py
+    workerGroupSpecs:
+    - replicas: 1
+      minReplicas: 1
+      maxReplicas: 5
+      groupName: small-group
+      rayStartParams: {}
+      template:
+        spec:
+          containers:
+          - name: ray-worker
+            image: rayproject/ray:2.46.0
+            resources:
+              limits:
+                cpu: "1"
+              requests:
+                cpu: "200m"
+                # SubmitterPodTemplate is the template for the pod that will run the `ray job submit` command against the RayCluster.
+                # If SubmitterPodTemplate is specified, the first container is assumed to be the submitter container.
+                # submitterPodTemplate:
+                #   spec:
+                #     restartPolicy: Never
+                #     containers:
+                #       - name: my-custom-rayjob-submitter-pod
+                #         image: rayproject/ray:2.46.0
+                #         # If Command is not specified, the correct command will be supplied at runtime using the RayJob spec `entrypoint` field.
+                #         # Specifying Command is not recommended.
+                #         # command: ["sh", "-c", "ray job submit --address=http://$RAY_DASHBOARD_ADDRESS --submission-id=$RAY_JOB_SUBMISSION_ID -- echo hello world"]
+
+
+######################Ray code sample#################################
+# this sample is from https://docs.ray.io/en/latest/cluster/job-submission.html#quick-start-example
+# it is mounted into the container and executed to show the Ray job at work
+---
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: ray-job-code-sample
+data:
+  sample_code.py: |
+    import ray
+    import os
+    import requests
+
+    ray.init()
+
+    @ray.remote
+    class Counter:
+        def __init__(self):
+            # Used to verify runtimeEnv
+            self.name = os.getenv("counter_name")
+            assert self.name == "test_counter"
+            self.counter = 0
+
+        def inc(self):
+            self.counter += 1
+
+        def get_counter(self):
+            return "{} got {}".format(self.name, self.counter)
+
+    counter = Counter.remote()
+
+    for _ in range(5):
+        ray.get(counter.inc.remote())
+        print(ray.get(counter.get_counter.remote()))
+
+    # Verify that the correct runtime env was used for the job.
+    assert requests.__version__ == "2.26.0"

ray-operator/apis/ray/v1/rayjob_types.go

@@ -9,6 +9,7 @@ import (
 	"time"
 
 	"github.com/go-logr/logr"
+	"github.com/robfig/cron/v3"
 	batchv1 "k8s.io/api/batch/v1"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/errors"
@@ -34,6 +35,8 @@ const (
 	RayJobDefaultRequeueDuration    = 3 * time.Second
 	RayJobDefaultClusterSelectorKey = "ray.io/cluster"
 	PythonUnbufferedEnvVarName      = "PYTHONUNBUFFERED"
+	// The buffer period in which a scheduled rajob can run since the last cron tick
+	ScheduleBuffer = 100 * time.Millisecond
 )
 
 // RayJobReconciler reconciles a RayJob object
@@ -168,6 +171,11 @@ func (r *RayJobReconciler) Reconcile(ctx context.Context, request ctrl.Request)
 				return ctrl.Result{RequeueAfter: RayJobDefaultRequeueDuration}, err
 			}
 		}
+		// We check the LastScheduleTime to know if its the first job
+		if rayJobInstance.Spec.Schedule != "" && rayJobInstance.Status.LastScheduleTime == nil {
+			rayJobInstance.Status.JobDeploymentStatus = rayv1.JobDeploymentStatusScheduled
+			break
+		}
 		// Set `Status.JobDeploymentStatus` to `JobDeploymentStatusInitializing`, and initialize `Status.JobId`
 		// and `Status.RayClusterName` prior to avoid duplicate job submissions and cluster creations.
 		logger.Info("JobDeploymentStatusNew")
@@ -449,9 +457,58 @@ func (r *RayJobReconciler) Reconcile(ctx context.Context, request ctrl.Request)
 				return ctrl.Result{RequeueAfter: RayJobDefaultRequeueDuration}, err
 			}
 		}
+		if rayJobInstance.Spec.Schedule != "" {
+			logger.Info("Rescheduling RayJob")
+			rayJobInstance.Status.JobDeploymentStatus = rayv1.JobDeploymentStatusScheduling
+			break
+		}
 
 		// If the RayJob is completed, we should not requeue it.
 		return ctrl.Result{}, nil
+	case rayv1.JobDeploymentStatusScheduling:
+		isJobDeleted, err := r.deleteSubmitterJob(ctx, rayJobInstance)
+		if err != nil {
+			return ctrl.Result{RequeueAfter: RayJobDefaultRequeueDuration}, err
+		}
+
+		if !isJobDeleted {
+			logger.Info("The release of the compute resources has not been completed yet. " +
+				"Wait for the resources to be deleted before the status transitions to avoid a resource leak.")
+			return ctrl.Result{RequeueAfter: RayJobDefaultRequeueDuration}, nil
+		}
+
+		if rayJobInstance.Spec.ShutdownAfterJobFinishes {
+			rayJobInstance.Status.RayClusterStatus = rayv1.RayClusterStatus{}
+			rayJobInstance.Status.RayClusterName = ""
+
+		}
+		rayJobInstance.Status.DashboardURL = ""
+		rayJobInstance.Status.JobId = ""
+		rayJobInstance.Status.Message = ""
+		rayJobInstance.Status.Reason = ""
+		rayJobInstance.Status.RayJobStatusInfo = rayv1.RayJobStatusInfo{}
+
+		rayJobInstance.Status.JobStatus = rayv1.JobStatusNew
+		rayJobInstance.Status.JobDeploymentStatus = rayv1.JobDeploymentStatusScheduled
+	case rayv1.JobDeploymentStatusScheduled:
+		// We get the time from the current time to the previous and next cron schedule times
+		// We pass in time.Now() as a parameter so easier unit testing and consistency
+		t1, t2, err := r.getNextAndPreviousScheduleDistance(ctx, time.Now(), rayJobInstance)
+		if err != nil {
+			logger.Error(err, "Could not get the previous and next distances for a cron schedule")
+			return ctrl.Result{}, err
+		}
+		// Checking if we are currently within a buffer to the previous cron schedule time
+		if t2 <= ScheduleBuffer {
+			logger.Info("The current time is within the buffer window of a cron tick", "NextScheduleTimeDuration", t1, "LastScheduleTimeDuration", t2, "Previous LastScheduleTime", rayJobInstance.Status.LastScheduleTime)
+			rayJobInstance.Status.LastScheduleTime = &metav1.Time{Time: time.Now()}
+			rayJobInstance.Status.JobStatus = rayv1.JobStatusNew
+			rayJobInstance.Status.JobDeploymentStatus = rayv1.JobDeploymentStatusNew
+		} else {
+			logger.Info("Waiting until the next reconcile to determine schedule", "nextScheduleDuration", t1, "currentTime", time.Now(), "lastScheduleTimeDuration", t2)
+			return ctrl.Result{RequeueAfter: t1}, nil
+		}
+
 	default:
 		logger.Info("Unknown JobDeploymentStatus", "JobDeploymentStatus", rayJobInstance.Status.JobDeploymentStatus)
 		return ctrl.Result{RequeueAfter: RayJobDefaultRequeueDuration}, nil
@@ -894,6 +951,23 @@ func (r *RayJobReconciler) constructRayClusterForRayJob(rayJobInstance *rayv1.Ra
 	return rayCluster, nil
 }
 
+func (r *RayJobReconciler) getNextAndPreviousScheduleDistance(ctx context.Context, currentTime time.Time, rayJobInstance *rayv1.RayJob) (time.Duration, time.Duration, error) {
+	logger := ctrl.LoggerFrom(ctx)
+	formatedCron := utils.FormatSchedule(rayJobInstance, r.Recorder)
+	cronSchedule, err := cron.ParseStandard(formatedCron)
+	if err != nil {
+		// this is likely a user error in defining the spec value
+		// we should log the error and not reconcile this cronjob until an update to spec
+		r.Recorder.Eventf(rayJobInstance, corev1.EventTypeWarning, "UnparseableSchedule", "unparseable schedule: %q : %s", rayJobInstance.Spec.Schedule, err)
+		return 0, 0, fmt.Errorf("the cron schedule provided is unparseable: %w", err)
+	}
+
+	t1 := utils.NextScheduleTimeDuration(logger, rayJobInstance, currentTime, cronSchedule)
+	t2 := utils.LastScheduleTimeDuration(logger, rayJobInstance, currentTime, cronSchedule)
+
+	return t1, t2, nil
+}
+
 func updateStatusToSuspendingIfNeeded(ctx context.Context, rayJob *rayv1.RayJob) bool {
 	logger := ctrl.LoggerFrom(ctx)
 	if !rayJob.Spec.Suspend {