Skip to content

Commit 6c8f3db

Browse files
committedDec 16, 2024
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 6c8f3db

File tree

9 files changed

+678
-37
lines changed

9 files changed

+678
-37
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-7
Original file line numberDiff line numberDiff line change
@@ -11,14 +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 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
1818
import os
1919
import filecmp
2020
from pathlib import Path
21-
from .kueue import list_local_queues
21+
from .kueue import list_local_queues, local_queue_exists, add_queue_label
2222

2323
parent = Path(__file__).resolve().parents[4] # project directory
2424
aw_dir = os.path.expanduser("~/.codeflare/resources/")
@@ -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
@@ -169,6 +169,123 @@ def test_list_local_queues(mocker):
169169
assert lqs == []
170170

171171

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

0 commit comments

Comments
 (0)