add priorities and schedulingSpec to SDK

Author: MichaelClifford

src/codeflare_sdk/cluster/cluster.py

@@ -84,6 +84,7 @@ def create_app_wrapper(self):
         instascale = self.config.instascale
         instance_types = self.config.machine_types
         env = self.config.envs
+        priority = self.config.priority
         return generate_appwrapper(
             name=name,
             namespace=namespace,
@@ -98,6 +99,7 @@ def create_app_wrapper(self):
             instascale=instascale,
             instance_types=instance_types,
             env=env,
+            priority=priority,
         )
 
     # creates a new cluster with the provided or default spec

src/codeflare_sdk/cluster/config.py

@@ -48,3 +48,4 @@ class ClusterConfiguration:
     instascale: bool = False
     envs: dict = field(default_factory=dict)
     image: str = "ghcr.io/foundation-model-stack/base:ray2.1.0-py38-gpu-pytorch1.12.0cu116-20221213-193103"
+    priority: str = "low"

src/codeflare_sdk/utils/generate_yaml.py

@@ -76,6 +76,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
 ):
@@ -155,6 +176,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,
@@ -211,6 +237,7 @@ def generate_appwrapper(
     instascale: bool,
     instance_types: list,
     env,
+    priority: str,
 ):
     user_yaml = read_template(template)
     appwrapper_name, cluster_name = gen_names(name)
@@ -219,6 +246,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
     )
@@ -315,6 +344,12 @@ def main():  # pragma: no cover
         default="default",
         help="Set the kubernetes namespace you want to deploy your cluster to. Default. If left blank, uses the 'default' namespace",
     )
+    parser.add_argument(
+        "--priority",
+        required=False,
+        default="low",
+        help="Set the priority of the cluster. Default is 'low'. Options are 'low', 'default', 'high'",
+    )
 
     args = parser.parse_args()
     name = args.name
@@ -329,6 +364,7 @@ def main():  # pragma: no cover
     instascale = args.instascale
     instance_types = args.instance_types
     namespace = args.namespace
+    priority = args.priority
     env = {}
 
     outfile = generate_appwrapper(
@@ -345,6 +381,7 @@ def main():  # pragma: no cover
         instascale,
         instance_types,
         env,
+        priority,
     )
     return outfile
 

tests/test-case-cmd.yaml

@@ -4,7 +4,7 @@ metadata:
   name: unit-cmd-cluster
   namespace: default
 spec:
-  priority: 9
+  priority: 1
   resources:
     GenericItems:
     - custompodresources:
@@ -136,6 +136,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
@@ -151,5 +152,7 @@ spec:
           to:
             kind: Service
             name: unit-cmd-cluster-head-svc
-      replica: 1
+      replicas: 1
     Items: []
+  schedulingSpec:
+    minAvailable: 3

tests/test-case.yaml

@@ -6,7 +6,7 @@ metadata:
   name: unit-test-cluster
   namespace: ns
 spec:
-  priority: 9
+  priority: 1
   resources:
     GenericItems:
     - custompodresources:
@@ -156,6 +156,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
@@ -171,5 +172,7 @@ spec:
           to:
             kind: Service
             name: unit-test-cluster-head-svc
-      replica: 1
+      replicas: 1
     Items: []
+  schedulingSpec:
+    minAvailable: 3

tests/unit_test.py

@@ -213,6 +213,7 @@ def test_config_creation():
         gpu=7,
         instascale=True,
         machine_types=["cpu.small", "gpu.large"],
+        priority="low",
     )
 
     assert config.name == "unit-test-cluster" and config.namespace == "ns"
@@ -227,11 +228,13 @@ def test_config_creation():
     assert config.template == f"{parent}/src/codeflare_sdk/templates/base-template.yaml"
     assert config.instascale
     assert config.machine_types == ["cpu.small", "gpu.large"]
+    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(