Skip to content

Commit 680cd06

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 6b0a3cc commit 680cd06

File tree

9 files changed

+678
-42
lines changed

9 files changed

+678
-42
lines changed

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

+124-12
Original file line numberDiff line numberDiff line change
@@ -11,19 +11,14 @@
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 (
15-
apply_template,
16-
get_local_queue,
17-
createClusterConfig,
18-
get_template_variables,
19-
)
14+
from ..utils.unit_test_support import get_local_queue, create_cluster_config, get_template_variables, apply_template
2015
from unittest.mock import patch
2116
from codeflare_sdk.ray.cluster.cluster import Cluster, ClusterConfiguration
2217
import yaml
2318
import os
2419
import filecmp
2520
from pathlib import Path
26-
from .kueue import list_local_queues
21+
from .kueue import list_local_queues, local_queue_exists, add_queue_label
2722

2823
parent = Path(__file__).resolve().parents[4] # project directory
2924
aw_dir = os.path.expanduser("~/.codeflare/resources/")
@@ -51,7 +46,7 @@ def test_cluster_creation_no_aw_local_queue(mocker):
5146
"kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
5247
return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"),
5348
)
54-
config = createClusterConfig()
49+
config = create_cluster_config()
5550
config.name = "unit-test-cluster-kueue"
5651
config.write_to_file = True
5752
config.local_queue = "local-queue-default"
@@ -67,7 +62,7 @@ def test_cluster_creation_no_aw_local_queue(mocker):
6762
assert cluster_kueue == expected_rc
6863

6964
# With resources loaded in memory, no Local Queue specified.
70-
config = createClusterConfig()
65+
config = create_cluster_config()
7166
config.name = "unit-test-cluster-kueue"
7267
config.write_to_file = False
7368
cluster = Cluster(config)
@@ -84,7 +79,7 @@ def test_aw_creation_local_queue(mocker):
8479
"kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
8580
return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"),
8681
)
87-
config = createClusterConfig()
82+
config = create_cluster_config()
8883
config.name = "unit-test-aw-kueue"
8984
config.appwrapper = True
9085
config.write_to_file = True
@@ -101,7 +96,7 @@ def test_aw_creation_local_queue(mocker):
10196
assert aw_kueue == expected_rc
10297

10398
# With resources loaded in memory, no Local Queue specified.
104-
config = createClusterConfig()
99+
config = create_cluster_config()
105100
config.name = "unit-test-aw-kueue"
106101
config.appwrapper = True
107102
config.write_to_file = False
@@ -120,7 +115,7 @@ def test_get_local_queue_exists_fail(mocker):
120115
"kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
121116
return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"),
122117
)
123-
config = createClusterConfig()
118+
config = create_cluster_config()
124119
config.name = "unit-test-aw-kueue"
125120
config.appwrapper = True
126121
config.write_to_file = True
@@ -175,6 +170,123 @@ def test_list_local_queues(mocker):
175170
assert lqs == []
176171

177172

173+
def test_local_queue_exists_found(mocker):
174+
# Mock Kubernetes client and list_namespaced_custom_object method
175+
mocker.patch("kubernetes.config.load_kube_config", return_value="ignore")
176+
mock_api_instance = mocker.Mock()
177+
mocker.patch("kubernetes.client.CustomObjectsApi", return_value=mock_api_instance)
178+
mocker.patch("codeflare_sdk.ray.cluster.cluster.config_check")
179+
180+
# Mock return value for list_namespaced_custom_object
181+
mock_api_instance.list_namespaced_custom_object.return_value = {
182+
"items": [
183+
{"metadata": {"name": "existing-queue"}},
184+
{"metadata": {"name": "another-queue"}},
185+
]
186+
}
187+
188+
# Call the function
189+
namespace = "test-namespace"
190+
local_queue_name = "existing-queue"
191+
result = local_queue_exists(namespace, local_queue_name)
192+
193+
# Assertions
194+
assert result is True
195+
mock_api_instance.list_namespaced_custom_object.assert_called_once_with(
196+
group="kueue.x-k8s.io",
197+
version="v1beta1",
198+
namespace=namespace,
199+
plural="localqueues",
200+
)
201+
202+
203+
def test_local_queue_exists_not_found(mocker):
204+
# Mock Kubernetes client and list_namespaced_custom_object method
205+
mocker.patch("kubernetes.config.load_kube_config", return_value="ignore")
206+
mock_api_instance = mocker.Mock()
207+
mocker.patch("kubernetes.client.CustomObjectsApi", return_value=mock_api_instance)
208+
mocker.patch("codeflare_sdk.ray.cluster.cluster.config_check")
209+
210+
# Mock return value for list_namespaced_custom_object
211+
mock_api_instance.list_namespaced_custom_object.return_value = {
212+
"items": [
213+
{"metadata": {"name": "another-queue"}},
214+
{"metadata": {"name": "different-queue"}},
215+
]
216+
}
217+
218+
# Call the function
219+
namespace = "test-namespace"
220+
local_queue_name = "non-existent-queue"
221+
result = local_queue_exists(namespace, local_queue_name)
222+
223+
# Assertions
224+
assert result is False
225+
mock_api_instance.list_namespaced_custom_object.assert_called_once_with(
226+
group="kueue.x-k8s.io",
227+
version="v1beta1",
228+
namespace=namespace,
229+
plural="localqueues",
230+
)
231+
232+
233+
import pytest
234+
from unittest import mock # If you're also using mocker from pytest-mock
235+
236+
237+
def test_add_queue_label_with_valid_local_queue(mocker):
238+
# Mock the kubernetes.client.CustomObjectsApi and its response
239+
mock_api_instance = mocker.patch("kubernetes.client.CustomObjectsApi")
240+
mock_api_instance.return_value.list_namespaced_custom_object.return_value = {
241+
"items": [
242+
{"metadata": {"name": "valid-queue"}},
243+
]
244+
}
245+
246+
# Mock other dependencies
247+
mocker.patch("codeflare_sdk.common.kueue.local_queue_exists", return_value=True)
248+
mocker.patch(
249+
"codeflare_sdk.common.kueue.get_default_kueue_name",
250+
return_value="default-queue",
251+
)
252+
253+
# Define input item and parameters
254+
item = {"metadata": {}}
255+
namespace = "test-namespace"
256+
local_queue = "valid-queue"
257+
258+
# Call the function
259+
add_queue_label(item, namespace, local_queue)
260+
261+
# Assert that the label is added to the item
262+
assert item["metadata"]["labels"] == {"kueue.x-k8s.io/queue-name": "valid-queue"}
263+
264+
265+
def test_add_queue_label_with_invalid_local_queue(mocker):
266+
# Mock the kubernetes.client.CustomObjectsApi and its response
267+
mock_api_instance = mocker.patch("kubernetes.client.CustomObjectsApi")
268+
mock_api_instance.return_value.list_namespaced_custom_object.return_value = {
269+
"items": [
270+
{"metadata": {"name": "valid-queue"}},
271+
]
272+
}
273+
274+
# Mock the local_queue_exists function to return False
275+
mocker.patch("codeflare_sdk.common.kueue.local_queue_exists", return_value=False)
276+
277+
# Define input item and parameters
278+
item = {"metadata": {}}
279+
namespace = "test-namespace"
280+
local_queue = "invalid-queue"
281+
282+
# Call the function and expect a ValueError
283+
with pytest.raises(
284+
ValueError,
285+
match="local_queue provided does not exist or is not in this namespace",
286+
):
287+
add_queue_label(item, namespace, local_queue)
288+
289+
178290
# Make sure to always keep this function last
179291
def test_cleanup():
180292
os.remove(f"{aw_dir}unit-test-cluster-kueue.yaml")

src/codeflare_sdk/common/utils/unit_test_support.py

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

3030

31-
def createClusterConfig():
31+
def create_cluster_config(num_workers=2, write_to_file=False):
3232
config = ClusterConfiguration(
3333
name="unit-test-cluster",
3434
namespace="ns",
35-
num_workers=2,
35+
num_workers=num_workers,
3636
worker_cpu_requests=3,
3737
worker_cpu_limits=4,
3838
worker_memory_requests=5,
3939
worker_memory_limits=6,
4040
appwrapper=True,
41-
write_to_file=False,
41+
write_to_file=write_to_file,
4242
)
4343
return config
4444

4545

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

5550

56-
def createClusterWrongType():
51+
def patch_cluster_with_dynamic_client(mocker, cluster, dynamic_client=None):
52+
mocker.patch.object(cluster, "get_dynamic_client", return_value=dynamic_client)
53+
mocker.patch.object(cluster, "down", return_value=None)
54+
mocker.patch.object(cluster, "config_check", return_value=None)
55+
# mocker.patch.object(cluster, "_throw_for_no_raycluster", return_value=None)
56+
57+
58+
def create_cluster_wrong_type():
5759
config = ClusterConfiguration(
5860
name="unit-test-cluster",
5961
namespace="ns",
@@ -411,6 +413,48 @@ def mocked_ingress(port, cluster_name="unit-test-cluster", annotations: dict = N
411413
return mock_ingress
412414

413415

416+
# Global dictionary to maintain state in the mock
417+
cluster_state = {}
418+
419+
420+
# The mock side_effect function for server_side_apply
421+
def mock_server_side_apply(resource, body=None, name=None, namespace=None, **kwargs):
422+
# Simulate the behavior of server_side_apply:
423+
# Update a mock state that represents the cluster's current configuration.
424+
# Stores the state in a global dictionary for simplicity.
425+
426+
global cluster_state
427+
428+
if not resource or not body or not name or not namespace:
429+
raise ValueError("Missing required parameters for server_side_apply")
430+
431+
# Extract worker count from the body if it exists
432+
try:
433+
worker_count = (
434+
body["spec"]["workerGroupSpecs"][0]["replicas"]
435+
if "spec" in body and "workerGroupSpecs" in body["spec"]
436+
else None
437+
)
438+
except KeyError:
439+
worker_count = None
440+
441+
# Apply changes to the cluster_state mock
442+
cluster_state[name] = {
443+
"namespace": namespace,
444+
"worker_count": worker_count,
445+
"body": body,
446+
}
447+
448+
# Return a response that mimics the behavior of a successful apply
449+
return {
450+
"status": "success",
451+
"applied": True,
452+
"name": name,
453+
"namespace": namespace,
454+
"worker_count": worker_count,
455+
}
456+
457+
414458
@patch.dict("os.environ", {"NB_PREFIX": "test-prefix"})
415459
def create_cluster_all_config_params(mocker, cluster_name, is_appwrapper) -> Cluster:
416460
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"

0 commit comments

Comments
 (0)