Add option to keep Ray clusters after a job has finished

gateway/api/management/commands/free_resources.py

@@ -6,6 +6,7 @@
 
 from api.models import ComputeResource, Job
 from api.ray import kill_ray_cluster
+from main import settings as config
 
 
 logger = logging.getLogger("commands")
@@ -27,6 +28,13 @@ def handle(self, *args, **options):
 
             # only kill cluster if not in local mode and no jobs are running there
             if len(alive_jobs) == 0 and not settings.RAY_CLUSTER_MODE.get("local"):
+                if config.RAY_CLUSTER_NO_DELETE_ON_COMPLETE:
+                    logger.debug(
+                        "RAY_CLUSTER_NO_DELETE_ON_COMPLETE is enabled, "
+                        + "so cluster [%s] will not be removed",
+                        compute_resource.title,
+                    )
+                    return
                 kill_ray_cluster(compute_resource.title)
                 compute_resource.delete()
                 counter += 1

gateway/main/settings.py

@@ -295,3 +295,7 @@
 )
 
 RAY_SETUP_MAX_RETRIES = int(os.environ.get("RAY_SETUP_MAX_RETRIES", 30))
+
+RAY_CLUSTER_NO_DELETE_ON_COMPLETE = bool(
+    os.environ.get("RAY_CLUSTER_NO_DELETE_ON_COMPLETE", False)
+)

infrastructure/helm/quantumserverless/charts/gateway/templates/deployment.yaml

@@ -187,6 +187,10 @@ spec:
               value: {{ .Values.application.limits.maxJobsPerUser | quote }}
             - name: LIMITS_MAX_CLUSTERS
               value: {{ .Values.application.limits.maxComputeResources | quote }}
+{{- if .Values.application.limits.keepClusterOnComplete }}
+            - name: RAY_CLUSTER_NO_DELETE_ON_COMPLETE
+              value: "True"
+{{- end }}       
       {{- with .Values.nodeSelector }}
       nodeSelector:
         {{- toYaml . | nindent 8 }}

infrastructure/helm/quantumserverless/charts/gateway/templates/rayclustertemplate.yaml

@@ -9,6 +9,10 @@ data:
     metadata:
       name: {{`{{ cluster_name }}`}}
       namespace: {{ .Release.Namespace }}
+{{- if .Values.application.limits.keepClusterOnComplete }}
+      labels:
+        nodelete: "true"
+{{- end }}
     spec:
 {{- if .Values.application.ray.scrapeWithPrometheus }}      
       headServiceAnnotations:

infrastructure/helm/quantumserverless/charts/gateway/values.yaml

@@ -30,6 +30,7 @@ application:
   limits:
     maxJobsPerUser: 2
     maxComputeResources: 4
+    keepClusterOnComplete: False
 
 database:
   host: postgresql