Skip to content

Commit 962d918

Browse files
evaluate priority class name
1 parent 20e6d93 commit 962d918

File tree

4 files changed

+45
-22
lines changed

4 files changed

+45
-22
lines changed

src/codeflare_sdk/cluster/cluster.py

Lines changed: 39 additions & 0 deletions
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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,4 +48,4 @@ class ClusterConfiguration:
4848
image: str = "quay.io/project-codeflare/ray:2.5.0-py38-cu116"
4949
local_interactive: bool = False
5050
image_pull_secrets: list = field(default_factory=list)
51-
dispatch_priority: str = "default"
51+
dispatch_priority: str = None

src/codeflare_sdk/templates/base-template.yaml

Lines changed: 0 additions & 2 deletions
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

Lines changed: 5 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -88,26 +88,12 @@ def update_labels(yaml, instascale, instance_types):
8888
metadata.pop("labels")
8989

9090

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

11298

11399
def update_custompodresources(
@@ -346,7 +332,7 @@ def generate_appwrapper(
346332
route_item = resources["resources"].get("GenericItems")[1]
347333
update_names(user_yaml, item, appwrapper_name, cluster_name, namespace)
348334
update_labels(user_yaml, instascale, instance_types)
349-
update_priority(user_yaml, item, dispatch_priority)
335+
update_priority(item, dispatch_priority)
350336
update_scheduling_spec(user_yaml, workers)
351337
update_custompodresources(
352338
item, min_cpu, max_cpu, min_memory, max_memory, gpu, workers

0 commit comments

Comments
 (0)