|
14 | 14 |
|
15 | 15 | import string
|
16 | 16 | import sys
|
| 17 | +from codeflare_sdk.common.utils import constants |
17 | 18 | from codeflare_sdk.ray.cluster.cluster import (
|
18 | 19 | Cluster,
|
19 | 20 | ClusterConfiguration,
|
@@ -68,7 +69,7 @@ def create_cluster_wrong_type():
|
68 | 69 | worker_extended_resource_requests={"nvidia.com/gpu": 7},
|
69 | 70 | appwrapper=True,
|
70 | 71 | image_pull_secrets=["unit-test-pull-secret"],
|
71 |
| - image="quay.io/modh/ray@sha256:a5b7c04a14f180d7ca6d06a5697f6bb684e40a26b95a0c872cac23b552741707", |
| 72 | + image=constants.CUDA_RUNTIME_IMAGE, |
72 | 73 | write_to_file=True,
|
73 | 74 | labels={1: 1},
|
74 | 75 | )
|
@@ -148,19 +149,29 @@ def get_cluster_object(file_a, file_b):
|
148 | 149 |
|
149 | 150 | def get_ray_obj(group, version, namespace, plural):
|
150 | 151 | # To be used for mocking list_namespaced_custom_object for Ray Clusters
|
151 |
| - rc_a_path = f"{parent}/tests/test_cluster_yamls/support_clusters/test-rc-a.yaml" |
152 |
| - rc_b_path = f"{parent}/tests/test_cluster_yamls/support_clusters/test-rc-b.yaml" |
153 |
| - rc_a, rc_b = get_cluster_object(rc_a_path, rc_b_path) |
| 152 | + rc_a = apply_template( |
| 153 | + f"{parent}/tests/test_cluster_yamls/support_clusters/test-rc-a.yaml", |
| 154 | + get_template_variables(), |
| 155 | + ) |
| 156 | + rc_b = apply_template( |
| 157 | + f"{parent}/tests/test_cluster_yamls/support_clusters/test-rc-b.yaml", |
| 158 | + get_template_variables(), |
| 159 | + ) |
154 | 160 |
|
155 | 161 | rc_list = {"items": [rc_a, rc_b]}
|
156 | 162 | return rc_list
|
157 | 163 |
|
158 | 164 |
|
159 | 165 | def get_ray_obj_with_status(group, version, namespace, plural):
|
160 | 166 | # To be used for mocking list_namespaced_custom_object for Ray Clusters with statuses
|
161 |
| - rc_a_path = f"{parent}/tests/test_cluster_yamls/support_clusters/test-rc-a.yaml" |
162 |
| - rc_b_path = f"{parent}/tests/test_cluster_yamls/support_clusters/test-rc-b.yaml" |
163 |
| - rc_a, rc_b = get_cluster_object(rc_a_path, rc_b_path) |
| 167 | + rc_a = apply_template( |
| 168 | + f"{parent}/tests/test_cluster_yamls/support_clusters/test-rc-a.yaml", |
| 169 | + get_template_variables(), |
| 170 | + ) |
| 171 | + rc_b = apply_template( |
| 172 | + f"{parent}/tests/test_cluster_yamls/support_clusters/test-rc-b.yaml", |
| 173 | + get_template_variables(), |
| 174 | + ) |
164 | 175 |
|
165 | 176 | rc_a.update(
|
166 | 177 | {
|
@@ -205,19 +216,29 @@ def get_ray_obj_with_status(group, version, namespace, plural):
|
205 | 216 |
|
206 | 217 | def get_aw_obj(group, version, namespace, plural):
|
207 | 218 | # To be used for mocking list_namespaced_custom_object for AppWrappers
|
208 |
| - aw_a_path = f"{parent}/tests/test_cluster_yamls/support_clusters/test-aw-a.yaml" |
209 |
| - aw_b_path = f"{parent}/tests/test_cluster_yamls/support_clusters/test-aw-b.yaml" |
210 |
| - aw_a, aw_b = get_cluster_object(aw_a_path, aw_b_path) |
| 219 | + aw_a = apply_template( |
| 220 | + f"{parent}/tests/test_cluster_yamls/support_clusters/test-aw-a.yaml", |
| 221 | + get_template_variables(), |
| 222 | + ) |
| 223 | + aw_b = apply_template( |
| 224 | + f"{parent}/tests/test_cluster_yamls/support_clusters/test-aw-b.yaml", |
| 225 | + get_template_variables(), |
| 226 | + ) |
211 | 227 |
|
212 | 228 | aw_list = {"items": [aw_a, aw_b]}
|
213 | 229 | return aw_list
|
214 | 230 |
|
215 | 231 |
|
216 | 232 | def get_aw_obj_with_status(group, version, namespace, plural):
|
217 | 233 | # To be used for mocking list_namespaced_custom_object for AppWrappers with statuses
|
218 |
| - aw_a_path = f"{parent}/tests/test_cluster_yamls/support_clusters/test-aw-a.yaml" |
219 |
| - aw_b_path = f"{parent}/tests/test_cluster_yamls/support_clusters/test-aw-b.yaml" |
220 |
| - aw_a, aw_b = get_cluster_object(aw_a_path, aw_b_path) |
| 234 | + aw_a = apply_template( |
| 235 | + f"{parent}/tests/test_cluster_yamls/support_clusters/test-aw-a.yaml", |
| 236 | + get_template_variables(), |
| 237 | + ) |
| 238 | + aw_b = apply_template( |
| 239 | + f"{parent}/tests/test_cluster_yamls/support_clusters/test-aw-b.yaml", |
| 240 | + get_template_variables(), |
| 241 | + ) |
221 | 242 |
|
222 | 243 | aw_a.update(
|
223 | 244 | {
|
@@ -273,9 +294,8 @@ def apply_template(yaml_file_path, variables):
|
273 | 294 |
|
274 | 295 |
|
275 | 296 | def get_expected_image():
|
276 |
| - # TODO: Add Python 3.12 support |
277 |
| - python_version = sys.version_info |
278 |
| - return "quay.io/modh/ray@sha256:a5b7c04a14f180d7ca6d06a5697f6bb684e40a26b95a0c872cac23b552741707" |
| 297 | + # TODO: Select image based on Python version |
| 298 | + return constants.CUDA_RUNTIME_IMAGE |
279 | 299 |
|
280 | 300 |
|
281 | 301 | def get_template_variables():
|
|
0 commit comments