|
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