Skip to content

Commit b6bac97

Browse files
make priority levels global
1 parent 87dea0a commit b6bac97

File tree

2 files changed

+9
-9
lines changed

2 files changed

+9
-9
lines changed

src/codeflare_sdk/utils/generate_yaml.py

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -89,17 +89,18 @@ 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+
9299
def update_priority(yaml, item, priority):
93-
if priority not in ["low", "default", "high"]:
100+
if priority not in PRIORITY_LEVELS:
94101
sys.exit("Priority must be 'low', 'default', or 'high'")
95102

96-
priority_levels = {
97-
"low": (1, "low-priority"),
98-
"default": (5, "default-priority"),
99-
"high": (10, "high-priority"),
100-
}
101-
102-
priority_level = priority_levels[priority]
103+
priority_level = PRIORITY_LEVELS[priority]
103104
spec = yaml.get("spec")
104105
spec["priority"] = priority_level[0]
105106
# spec["SchedulingSpec"]["priorityClassName"] = priority_level

tests/unit_test.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -247,7 +247,6 @@ def test_config_creation():
247247

248248
def test_cluster_creation():
249249
cluster = Cluster(test_config_creation())
250-
print(cluster.app_wrapper_yaml)
251250
assert cluster.app_wrapper_yaml == "unit-test-cluster.yaml"
252251
assert cluster.app_wrapper_name == "unit-test-cluster"
253252
assert filecmp.cmp(

0 commit comments

Comments
 (0)