add priorities and schedulingSpec to SDK

Author: MichaelClifford

src/codeflare_sdk/cluster/cluster.py

@@ -89,8 +89,12 @@ def create_app_wrapper(self):
         instascale = self.config.instascale
         instance_types = self.config.machine_types
         env = self.config.envs
+<<<<<<< HEAD
         local_interactive = self.config.local_interactive
         image_pull_secrets = self.config.image_pull_secrets
+=======
+        priority = self.config.priority
+>>>>>>> 7e7a311 ( add priorities and schedulingSpec to SDK)
         return generate_appwrapper(
             name=name,
             namespace=namespace,
@@ -105,8 +109,12 @@ def create_app_wrapper(self):
             instascale=instascale,
             instance_types=instance_types,
             env=env,
+<<<<<<< HEAD
             local_interactive=local_interactive,
             image_pull_secrets=image_pull_secrets,
+=======
+            priority=priority,
+>>>>>>> 7e7a311 ( add priorities and schedulingSpec to SDK)
         )
 
     # creates a new cluster with the provided or default spec

src/codeflare_sdk/cluster/config.py

@@ -47,3 +47,4 @@ class ClusterConfiguration:
     image: str = "quay.io/project-codeflare/ray:2.5.0-py38-cu116"
     local_interactive: bool = False
     image_pull_secrets: list = field(default_factory=list)
+

src/codeflare_sdk/utils/generate_yaml.py

@@ -89,6 +89,27 @@ def update_labels(yaml, instascale, instance_types):
         metadata.pop("labels")
 
 
+def update_priority(yaml, item, priority):
+    if priority not in ["low", "default", "high"]:
+        sys.exit("Priority must be 'low', 'default', or 'high'")
+
+    priority_levels = {
+        "low": (1, "low-priority"),
+        "default": (5, "default-priority"),
+        "high": (10, "high-priority"),
+    }
+
+    priority_level = priority_levels[priority]
+    spec = yaml.get("spec")
+    spec["priority"] = priority_level[0]
+    # spec["SchedulingSpec"]["priorityClassName"] = priority_level
+    if "generictemplate" in item.keys():
+        head = item.get("generictemplate").get("spec").get("headGroupSpec")
+        worker = item.get("generictemplate").get("spec").get("workerGroupSpecs")[0]
+        head["template"]["spec"]["priorityClassName"] = priority_level[1]
+        worker["template"]["spec"]["priorityClassName"] = priority_level[1]
+
+
 def update_custompodresources(
     item, min_cpu, max_cpu, min_memory, max_memory, gpu, workers
 ):
@@ -175,6 +196,11 @@ def update_resources(spec, min_cpu, max_cpu, min_memory, max_memory, gpu):
             limits["nvidia.com/gpu"] = gpu
 
 
+def update_scheduling_spec(yaml, workers):
+    spec = yaml.get("spec")
+    spec["schedulingSpec"]["minAvailable"] = workers + 1
+
+
 def update_nodes(
     item,
     appwrapper_name,
@@ -346,6 +372,7 @@ def generate_appwrapper(
     env,
     local_interactive: bool,
     image_pull_secrets: list,
+    priority: str,
 ):
     user_yaml = read_template(template)
     appwrapper_name, cluster_name = gen_names(name)
@@ -354,6 +381,8 @@ def generate_appwrapper(
     route_item = resources["resources"].get("GenericItems")[1]
     update_names(user_yaml, item, appwrapper_name, cluster_name, namespace)
     update_labels(user_yaml, instascale, instance_types)
+    update_priority(user_yaml, item, priority)
+    update_scheduling_spec(user_yaml, workers)
     update_custompodresources(
         item, min_cpu, max_cpu, min_memory, max_memory, gpu, workers
     )

tests/test-case.yaml

@@ -6,7 +6,7 @@ metadata:
   name: unit-test-cluster
   namespace: ns
 spec:
-  priority: 9
+  priority: 1
   resources:
     GenericItems:
     - custompodresources:
@@ -176,6 +176,7 @@ spec:
                     do echo waiting for myservice; sleep 2; done
                   image: busybox:1.28
                   name: init-myservice
+                priorityClassName: low-priority
       replicas: 1
     - generictemplate:
         apiVersion: route.openshift.io/v1
@@ -193,3 +194,5 @@ spec:
             name: unit-test-cluster-head-svc
       replicas: 1
     Items: []
+  schedulingSpec:
+    minAvailable: 3

tests/unit_test.py

@@ -228,6 +228,7 @@ def test_config_creation():
         instascale=True,
         machine_types=["cpu.small", "gpu.large"],
         image_pull_secrets=["unit-test-pull-secret"],
+        priority="low",
     )
 
     assert config.name == "unit-test-cluster" and config.namespace == "ns"
@@ -240,11 +241,13 @@ def test_config_creation():
     assert config.instascale
     assert config.machine_types == ["cpu.small", "gpu.large"]
     assert config.image_pull_secrets == ["unit-test-pull-secret"]
+    assert config.priority == "low"
     return config
 
 
 def test_cluster_creation():
     cluster = Cluster(test_config_creation())
+    print(cluster.app_wrapper_yaml)
     assert cluster.app_wrapper_yaml == "unit-test-cluster.yaml"
     assert cluster.app_wrapper_name == "unit-test-cluster"
     assert filecmp.cmp(