Skip to content

Commit 957fc35

Browse files
evaluate priority class name
1 parent e97c33b commit 957fc35

File tree

4 files changed

+45
-22
lines changed

4 files changed

+45
-22
lines changed

src/codeflare_sdk/cluster/cluster.py

+39
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,39 @@ def __init__(self, config: ClusterConfiguration):
6161
self.app_wrapper_yaml = self.create_app_wrapper()
6262
self.app_wrapper_name = self.app_wrapper_yaml.split(".")[0]
6363

64+
def evaluate_config(self):
65+
if not self.evaluate_dispatch_priority():
66+
return False
67+
else:
68+
return True
69+
70+
def evaluate_dispatch_priority(self):
71+
priority_class = self.config.dispatch_priority
72+
if priority_class is None:
73+
return True
74+
else:
75+
try:
76+
config_check()
77+
api_instance = client.CustomObjectsApi(api_config_handler())
78+
priority_classes = api_instance.list_cluster_custom_object(
79+
group="scheduling.k8s.io",
80+
version="v1",
81+
plural="priorityclasses",
82+
)
83+
available_priority_classes = [
84+
i["metadata"]["name"] for i in priority_classes["items"]
85+
]
86+
except Exception as e: # pragma: no cover
87+
return _kube_api_error_handling(e)
88+
89+
if priority_class in available_priority_classes:
90+
return True
91+
else:
92+
print(
93+
f"Priority class {priority_class} is not available in the cluster"
94+
)
95+
return False
96+
6497
def create_app_wrapper(self):
6598
"""
6699
Called upon cluster object creation, creates an AppWrapper yaml based on
@@ -117,6 +150,12 @@ def up(self):
117150
Applies the AppWrapper yaml, pushing the resource request onto
118151
the MCAD queue.
119152
"""
153+
154+
# Before attempting to bring up the cluster let's evaluate the ClusterConfig
155+
if not self.evaluate_config():
156+
print("Invalid Cluster Configuration")
157+
return False
158+
120159
namespace = self.config.namespace
121160
try:
122161
config_check()

src/codeflare_sdk/cluster/config.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -47,4 +47,4 @@ class ClusterConfiguration:
4747
image: str = "quay.io/project-codeflare/ray:2.5.0-py38-cu116"
4848
local_interactive: bool = False
4949
image_pull_secrets: list = field(default_factory=list)
50-
dispatch_priority: str = "default"
50+
dispatch_priority: str = None

src/codeflare_sdk/templates/base-template.yaml

-2
Original file line numberDiff line numberDiff line change
@@ -114,7 +114,6 @@ spec:
114114
operator: In
115115
values:
116116
- "aw-kuberay"
117-
priorityClassName: "default-priority"
118117
containers:
119118
# The Ray head pod
120119
- env:
@@ -224,7 +223,6 @@ spec:
224223
operator: In
225224
values:
226225
- "aw-kuberay"
227-
priorityClassName: "default-priority"
228226
initContainers:
229227
# the env var $RAY_IP is set by the operator if missing, with the value of the head service name
230228
- name: init-myservice

src/codeflare_sdk/utils/generate_yaml.py

+5-19
Original file line numberDiff line numberDiff line change
@@ -89,26 +89,12 @@ def update_labels(yaml, instascale, instance_types):
8989
metadata.pop("labels")
9090

9191

92-
PRIORITY_LEVELS = {
93-
"low": (1, "low-priority"),
94-
"default": (5, "default-priority"),
95-
"high": (10, "high-priority"),
96-
}
97-
98-
99-
def update_priority(yaml, item, priority):
100-
if priority not in PRIORITY_LEVELS:
101-
sys.exit("Priority must be 'low', 'default', or 'high'")
102-
103-
priority_level = PRIORITY_LEVELS[priority]
104-
spec = yaml.get("spec")
105-
spec["priority"] = priority_level[0]
106-
# spec["SchedulingSpec"]["priorityClassName"] = priority_level
107-
if "generictemplate" in item.keys():
92+
def update_priority(item, dispatch_priority):
93+
if dispatch_priority is not None:
10894
head = item.get("generictemplate").get("spec").get("headGroupSpec")
10995
worker = item.get("generictemplate").get("spec").get("workerGroupSpecs")[0]
110-
head["template"]["spec"]["priorityClassName"] = priority_level[1]
111-
worker["template"]["spec"]["priorityClassName"] = priority_level[1]
96+
head["template"]["spec"]["priorityClassName"] = dispatch_priority
97+
worker["template"]["spec"]["priorityClassName"] = dispatch_priority
11298

11399

114100
def update_custompodresources(
@@ -382,7 +368,7 @@ def generate_appwrapper(
382368
route_item = resources["resources"].get("GenericItems")[1]
383369
update_names(user_yaml, item, appwrapper_name, cluster_name, namespace)
384370
update_labels(user_yaml, instascale, instance_types)
385-
update_priority(user_yaml, item, dispatch_priority)
371+
update_priority(item, dispatch_priority)
386372
update_scheduling_spec(user_yaml, workers)
387373
update_custompodresources(
388374
item, min_cpu, max_cpu, min_memory, max_memory, gpu, workers

0 commit comments

Comments
 (0)