Skip to content

Commit eda66e1

Browse files
committed
Adds RayCluster.apply()
- Adds RayCluster.apply() implementation - Adds e2e tests for apply - Adds unit tests for apply - Exclude unit tests code from coverage - Add coverage to cluster.py - Adding coverage for the case of an openshift cluster
1 parent be9763a commit eda66e1

File tree

11 files changed

+567
-33
lines changed

11 files changed

+567
-33
lines changed

.coveragerc

+7
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
[run]
2+
omit =
3+
src/**/test_*.py,,src/codeflare_sdk/common/utils/unit_test_support.py
4+
5+
[report]
6+
exclude_lines =
7+
pragma: no cover

.github/workflows/coverage-badge.yaml

+1-1
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ jobs:
2626
poetry install --with test
2727
- name: Generate coverage report
2828
run: |
29-
coverage run -m pytest
29+
coverage run -m pytest --omit="src/**/test_*.py,src/codeflare_sdk/common/utils/unit_test_support.py"
3030
3131
- name: Coverage Badge
3232
uses: tj-actions/coverage-badge-py@v2

CONTRIBUTING.md

+1-1
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,7 @@ pytest -v src/codeflare_sdk
7676

7777
### Local e2e Testing
7878

79-
- Please follow the [e2e documentation](https://github.com/project-codeflare/codeflare-sdk/blob/main/docs/e2e.md)
79+
- Please follow the [e2e documentation](https://github.com/project-codeflare/codeflare-sdk/blob/main/docs/sphinx/user-docs/e2e.rst)
8080

8181
#### Code Coverage
8282

src/codeflare_sdk/common/kueue/test_kueue.py

+6-6
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
14-
from ..utils.unit_test_support import get_local_queue, createClusterConfig
14+
from ..utils.unit_test_support import get_local_queue, create_cluster_config
1515
from unittest.mock import patch
1616
from codeflare_sdk.ray.cluster.cluster import Cluster, ClusterConfiguration
1717
import yaml
@@ -46,7 +46,7 @@ def test_cluster_creation_no_aw_local_queue(mocker):
4646
"kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
4747
return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"),
4848
)
49-
config = createClusterConfig()
49+
config = create_cluster_config()
5050
config.name = "unit-test-cluster-kueue"
5151
config.write_to_file = True
5252
config.local_queue = "local-queue-default"
@@ -59,7 +59,7 @@ def test_cluster_creation_no_aw_local_queue(mocker):
5959
)
6060

6161
# With resources loaded in memory, no Local Queue specified.
62-
config = createClusterConfig()
62+
config = create_cluster_config()
6363
config.name = "unit-test-cluster-kueue"
6464
config.write_to_file = False
6565
cluster = Cluster(config)
@@ -79,7 +79,7 @@ def test_aw_creation_local_queue(mocker):
7979
"kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
8080
return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"),
8181
)
82-
config = createClusterConfig()
82+
config = create_cluster_config()
8383
config.name = "unit-test-aw-kueue"
8484
config.appwrapper = True
8585
config.write_to_file = True
@@ -93,7 +93,7 @@ def test_aw_creation_local_queue(mocker):
9393
)
9494

9595
# With resources loaded in memory, no Local Queue specified.
96-
config = createClusterConfig()
96+
config = create_cluster_config()
9797
config.name = "unit-test-aw-kueue"
9898
config.appwrapper = True
9999
config.write_to_file = False
@@ -114,7 +114,7 @@ def test_get_local_queue_exists_fail(mocker):
114114
"kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
115115
return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"),
116116
)
117-
config = createClusterConfig()
117+
config = create_cluster_config()
118118
config.name = "unit-test-aw-kueue"
119119
config.appwrapper = True
120120
config.write_to_file = True

src/codeflare_sdk/common/utils/unit_test_support.py

+55-11
Original file line numberDiff line numberDiff line change
@@ -26,32 +26,34 @@
2626
aw_dir = os.path.expanduser("~/.codeflare/resources/")
2727

2828

29-
def createClusterConfig():
29+
def create_cluster_config(num_workers=2, write_to_file=False):
3030
config = ClusterConfiguration(
3131
name="unit-test-cluster",
3232
namespace="ns",
33-
num_workers=2,
33+
num_workers=num_workers,
3434
worker_cpu_requests=3,
3535
worker_cpu_limits=4,
3636
worker_memory_requests=5,
3737
worker_memory_limits=6,
3838
appwrapper=True,
39-
write_to_file=False,
39+
write_to_file=write_to_file,
4040
)
4141
return config
4242

4343

44-
def createClusterWithConfig(mocker):
45-
mocker.patch("kubernetes.config.load_kube_config", return_value="ignore")
46-
mocker.patch(
47-
"kubernetes.client.CustomObjectsApi.get_cluster_custom_object",
48-
return_value={"spec": {"domain": "apps.cluster.awsroute.org"}},
49-
)
50-
cluster = Cluster(createClusterConfig())
44+
def create_cluster(mocker, num_workers=2, write_to_file=False):
45+
cluster = Cluster(create_cluster_config(num_workers, write_to_file))
5146
return cluster
5247

5348

54-
def createClusterWrongType():
49+
def patch_cluster_with_dynamic_client(mocker, cluster, dynamic_client=None):
50+
mocker.patch.object(cluster, "get_dynamic_client", return_value=dynamic_client)
51+
mocker.patch.object(cluster, "down", return_value=None)
52+
mocker.patch.object(cluster, "config_check", return_value=None)
53+
# mocker.patch.object(cluster, "_throw_for_no_raycluster", return_value=None)
54+
55+
56+
def create_cluster_wrong_type():
5557
config = ClusterConfiguration(
5658
name="unit-test-cluster",
5759
namespace="ns",
@@ -383,6 +385,48 @@ def mocked_ingress(port, cluster_name="unit-test-cluster", annotations: dict = N
383385
return mock_ingress
384386

385387

388+
# Global dictionary to maintain state in the mock
389+
cluster_state = {}
390+
391+
392+
# The mock side_effect function for server_side_apply
393+
def mock_server_side_apply(resource, body=None, name=None, namespace=None, **kwargs):
394+
# Simulate the behavior of server_side_apply:
395+
# Update a mock state that represents the cluster's current configuration.
396+
# Stores the state in a global dictionary for simplicity.
397+
398+
global cluster_state
399+
400+
if not resource or not body or not name or not namespace:
401+
raise ValueError("Missing required parameters for server_side_apply")
402+
403+
# Extract worker count from the body if it exists
404+
try:
405+
worker_count = (
406+
body["spec"]["workerGroupSpecs"][0]["replicas"]
407+
if "spec" in body and "workerGroupSpecs" in body["spec"]
408+
else None
409+
)
410+
except KeyError:
411+
worker_count = None
412+
413+
# Apply changes to the cluster_state mock
414+
cluster_state[name] = {
415+
"namespace": namespace,
416+
"worker_count": worker_count,
417+
"body": body,
418+
}
419+
420+
# Return a response that mimics the behavior of a successful apply
421+
return {
422+
"status": "success",
423+
"applied": True,
424+
"name": name,
425+
"namespace": namespace,
426+
"worker_count": worker_count,
427+
}
428+
429+
386430
@patch.dict("os.environ", {"NB_PREFIX": "test-prefix"})
387431
def create_cluster_all_config_params(mocker, cluster_name, is_appwrapper) -> Cluster:
388432
mocker.patch(

src/codeflare_sdk/common/widgets/test_widgets.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
import codeflare_sdk.common.widgets.widgets as cf_widgets
1616
import pandas as pd
1717
from unittest.mock import MagicMock, patch
18-
from ..utils.unit_test_support import get_local_queue, createClusterConfig
18+
from ..utils.unit_test_support import get_local_queue, create_cluster_config
1919
from codeflare_sdk.ray.cluster.cluster import Cluster
2020
from codeflare_sdk.ray.cluster.status import (
2121
RayCluster,
@@ -38,7 +38,7 @@ def test_cluster_up_down_buttons(mocker):
3838
"kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
3939
return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"),
4040
)
41-
cluster = Cluster(createClusterConfig())
41+
cluster = Cluster(create_cluster_config())
4242

4343
with patch("ipywidgets.Button") as MockButton, patch(
4444
"ipywidgets.Checkbox"

src/codeflare_sdk/ray/cluster/cluster.py

+86-4
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,10 @@
5252
import requests
5353

5454
from kubernetes import config
55+
from kubernetes.dynamic import DynamicClient
56+
from kubernetes import client as k8s_client
57+
from kubernetes.client.rest import ApiException
58+
5559
from kubernetes.client.rest import ApiException
5660
import warnings
5761

@@ -84,6 +88,14 @@ def __init__(self, config: ClusterConfiguration):
8488
if is_notebook():
8589
cluster_up_down_buttons(self)
8690

91+
def get_dynamic_client(self):
92+
"""Return a dynamic client, optionally mocked in tests."""
93+
return DynamicClient(get_api_client())
94+
95+
def config_check(self):
96+
"""Return a dynamic client, optionally mocked in tests."""
97+
return config_check()
98+
8799
@property
88100
def _client_headers(self):
89101
k8_client = get_api_client()
@@ -130,7 +142,6 @@ def create_resource(self):
130142
raise TypeError(
131143
f"Namespace {self.config.namespace} is of type {type(self.config.namespace)}. Check your Kubernetes Authentication."
132144
)
133-
134145
return build_ray_cluster(self)
135146

136147
# creates a new cluster with the provided or default spec
@@ -139,10 +150,11 @@ def up(self):
139150
Applies the Cluster yaml, pushing the resource request onto
140151
the Kueue localqueue.
141152
"""
142-
153+
print(
154+
"WARNING: The up() function is planned for deprecation in favor of apply()."
155+
)
143156
# check if RayCluster CustomResourceDefinition exists if not throw RuntimeError
144157
self._throw_for_no_raycluster()
145-
146158
namespace = self.config.namespace
147159

148160
try:
@@ -176,6 +188,52 @@ def up(self):
176188
except Exception as e: # pragma: no cover
177189
return _kube_api_error_handling(e)
178190

191+
# Applies a new cluster with the provided or default spec
192+
def apply(self, force=False):
193+
"""
194+
Applies the Cluster yaml using server-side apply.
195+
If 'force' is set to True, conflicts will be forced.
196+
"""
197+
# check if RayCluster CustomResourceDefinition exists if not throw RuntimeError
198+
self._throw_for_no_raycluster()
199+
namespace = self.config.namespace
200+
201+
try:
202+
self.config_check()
203+
api_instance = client.CustomObjectsApi(get_api_client())
204+
crds = self.get_dynamic_client().resources
205+
api_instance = crds.get(
206+
api_version="workload.codeflare.dev/v1beta2", kind="AppWrapper"
207+
)
208+
if self.config.appwrapper:
209+
if self.config.write_to_file:
210+
with open(self.resource_yaml) as f:
211+
aw = yaml.load(f, Loader=yaml.FullLoader)
212+
api_instance.server_side_apply(
213+
group="workload.codeflare.dev",
214+
version="v1beta2",
215+
namespace=namespace,
216+
plural="appwrappers",
217+
body=aw,
218+
)
219+
else:
220+
api_instance.server_side_apply(
221+
group="workload.codeflare.dev",
222+
version="v1beta2",
223+
namespace=namespace,
224+
plural="appwrappers",
225+
body=self.resource_yaml,
226+
)
227+
print(f"AppWrapper: '{self.config.name}' has successfully been created")
228+
else:
229+
api_instance = crds.get(api_version="ray.io/v1", kind="RayCluster")
230+
self._component_resources_apply(namespace, api_instance)
231+
print(
232+
f"Ray Cluster: '{self.config.name}' has successfully been applied"
233+
)
234+
except Exception as e: # pragma: no cover
235+
return _kube_api_error_handling(e)
236+
179237
def _throw_for_no_raycluster(self):
180238
api_instance = client.CustomObjectsApi(get_api_client())
181239
try:
@@ -204,7 +262,7 @@ def down(self):
204262
resource_name = self.config.name
205263
self._throw_for_no_raycluster()
206264
try:
207-
config_check()
265+
self.config_check()
208266
api_instance = client.CustomObjectsApi(get_api_client())
209267
if self.config.appwrapper:
210268
api_instance.delete_namespaced_custom_object(
@@ -507,6 +565,16 @@ def _component_resources_up(
507565
else:
508566
_create_resources(self.resource_yaml, namespace, api_instance)
509567

568+
def _component_resources_apply(
569+
self, namespace: str, api_instance: client.CustomObjectsApi
570+
):
571+
if self.config.write_to_file:
572+
with open(self.resource_yaml) as f:
573+
ray_cluster = yaml.safe_load(f)
574+
_apply_resources(ray_cluster, namespace, api_instance)
575+
else:
576+
_apply_resources(self.resource_yaml, namespace, api_instance)
577+
510578
def _component_resources_down(
511579
self, namespace: str, api_instance: client.CustomObjectsApi
512580
):
@@ -744,6 +812,20 @@ def _create_resources(yamls, namespace: str, api_instance: client.CustomObjectsA
744812
)
745813

746814

815+
def _apply_resources(
816+
yamls, namespace: str, api_instance: client.CustomObjectsApi, force=False
817+
):
818+
api_instance.server_side_apply(
819+
field_manager="cluster-manager",
820+
group="ray.io",
821+
version="v1",
822+
namespace=namespace,
823+
plural="rayclusters",
824+
body=yamls,
825+
force_conflicts=force, # Allow forcing conflicts if needed
826+
)
827+
828+
747829
def _check_aw_exists(name: str, namespace: str) -> bool:
748830
try:
749831
config_check()

0 commit comments

Comments
 (0)