From 7d627a79dcdfabb668bb22c3161826ca0dd6106e Mon Sep 17 00:00:00 2001 From: Laura Fitzgerald Date: Fri, 9 May 2025 15:28:51 +0100 Subject: [PATCH 1/2] make local_queue optional --- src/codeflare_sdk/common/kueue/kueue.py | 19 ++- src/codeflare_sdk/common/kueue/test_kueue.py | 108 ++++++++++++++++-- .../ray/cluster/build_ray_cluster.py | 23 +++- 3 files changed, 132 insertions(+), 18 deletions(-) diff --git a/src/codeflare_sdk/common/kueue/kueue.py b/src/codeflare_sdk/common/kueue/kueue.py index 00f3364a4..3db7ce905 100644 --- a/src/codeflare_sdk/common/kueue/kueue.py +++ b/src/codeflare_sdk/common/kueue/kueue.py @@ -17,6 +17,7 @@ from codeflare_sdk.common.kubernetes_cluster.auth import config_check, get_api_client from kubernetes import client from kubernetes.client.exceptions import ApiException +import warnings def get_default_kueue_name(namespace: str) -> Optional[str]: @@ -157,18 +158,24 @@ def add_queue_label(item: dict, namespace: str, local_queue: Optional[str]): The namespace of the local queue. local_queue (str, optional): The name of the local queue to use. Defaults to None. - - Raises: - ValueError: - If the provided or default local queue does not exist in the namespace. """ lq_name = local_queue or get_default_kueue_name(namespace) if lq_name == None: return elif not local_queue_exists(namespace, lq_name): - raise ValueError( - "local_queue provided does not exist or is not in this namespace. Please provide the correct local_queue name in Cluster Configuration" + available_queues = list_local_queues(namespace) + if available_queues is None: + warnings.warn( + f"Local queue '{local_queue}' does not exist in namespace '{namespace}'. " + "Unable to retrieve list of available queues." + ) + return + available_queue_names = [q["name"] for q in available_queues] + warnings.warn( + f"Local queue '{local_queue}' does not exist in namespace '{namespace}'. " + f"Available queues are: {', '.join(available_queue_names)}" ) + return if not "labels" in item["metadata"]: item["metadata"]["labels"] = {} item["metadata"]["labels"].update({"kueue.x-k8s.io/queue-name": lq_name}) diff --git a/src/codeflare_sdk/common/kueue/test_kueue.py b/src/codeflare_sdk/common/kueue/test_kueue.py index bbc54e9e0..3ebefa93c 100644 --- a/src/codeflare_sdk/common/kueue/test_kueue.py +++ b/src/codeflare_sdk/common/kueue/test_kueue.py @@ -118,7 +118,14 @@ def test_get_local_queue_exists_fail(mocker): ) mocker.patch( "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object", - return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"), + return_value={ + "items": [ + { + "metadata": {"name": "local-queue-default"}, + "status": {"flavors": [{"name": "default"}]}, + } + ] + }, ) config = create_cluster_config() config.name = "unit-test-aw-kueue" @@ -267,15 +274,99 @@ def test_add_queue_label_with_valid_local_queue(mocker): assert item["metadata"]["labels"] == {"kueue.x-k8s.io/queue-name": "valid-queue"} -def test_add_queue_label_with_invalid_local_queue(mocker): +def test_add_queue_label_with_invalid_local_queue_shows_available_queues(mocker): # Mock the kubernetes.client.CustomObjectsApi and its response mock_api_instance = mocker.patch("kubernetes.client.CustomObjectsApi") mock_api_instance.return_value.list_namespaced_custom_object.return_value = { "items": [ - {"metadata": {"name": "valid-queue"}}, + {"metadata": {"name": "queue1"}, "status": {"flavors": [{"name": "default"}]}}, + {"metadata": {"name": "queue2"}, "status": {"flavors": [{"name": "default"}]}}, + ] + } + + # Mock the local_queue_exists function to return False + mocker.patch("codeflare_sdk.common.kueue.local_queue_exists", return_value=False) + + # Define input item and parameters + item = {"metadata": {}} + namespace = "test-namespace" + local_queue = "invalid-queue" + + # Call the function and expect a warning with available queues in the message + with pytest.warns( + UserWarning, + match=f"Local queue '{local_queue}' does not exist in namespace '{namespace}'. Available queues are: queue1, queue2", + ): + add_queue_label(item, namespace, local_queue) + + # Assert that no label is added + assert "labels" not in item["metadata"] + + +def test_add_queue_label_with_no_local_queue(mocker): + # Mock the kubernetes.client.CustomObjectsApi and its response + mock_api_instance = mocker.patch("kubernetes.client.CustomObjectsApi") + mock_api_instance.return_value.list_namespaced_custom_object.return_value = { + "items": [] + } + + # Mock get_default_kueue_name to return None + mocker.patch( + "codeflare_sdk.common.kueue.get_default_kueue_name", + return_value=None, + ) + + # Define input item and parameters + item = {"metadata": {}} + namespace = "test-namespace" + local_queue = None + + # Call the function + add_queue_label(item, namespace, local_queue) + + # Assert that no label is added + assert "labels" not in item["metadata"] + + +def test_add_queue_label_with_default_queue(mocker): + # Mock the kubernetes.client.CustomObjectsApi and its response + mock_api_instance = mocker.patch("kubernetes.client.CustomObjectsApi") + mock_api_instance.return_value.list_namespaced_custom_object.return_value = { + "items": [ + { + "metadata": { + "name": "default-queue", + "annotations": {"kueue.x-k8s.io/default-queue": "true"} + } + } ] } + # Mock get_default_kueue_name to return a default queue + mocker.patch( + "codeflare_sdk.common.kueue.get_default_kueue_name", + return_value="default-queue", + ) + + # Define input item and parameters + item = {"metadata": {}} + namespace = "test-namespace" + local_queue = None + + # Call the function + add_queue_label(item, namespace, local_queue) + + # Assert that the default queue label is added + assert item["metadata"]["labels"] == {"kueue.x-k8s.io/queue-name": "default-queue"} + + +def test_add_queue_label_with_invalid_local_queue_and_no_available_queues(mocker): + # Mock the kubernetes.client.CustomObjectsApi and its response + mock_api_instance = mocker.patch("kubernetes.client.CustomObjectsApi") + mock_api_instance.return_value.list_namespaced_custom_object.return_value = { + "items": [] # Empty list instead of None + } + # Mock the local_queue_exists function to return False mocker.patch("codeflare_sdk.common.kueue.local_queue_exists", return_value=False) @@ -284,13 +375,16 @@ def test_add_queue_label_with_invalid_local_queue(mocker): namespace = "test-namespace" local_queue = "invalid-queue" - # Call the function and expect a ValueError - with pytest.raises( - ValueError, - match="local_queue provided does not exist or is not in this namespace", + # Call the function and expect a warning about unavailable queues + with pytest.warns( + UserWarning, + match=f"Local queue '{local_queue}' does not exist in namespace '{namespace}'. Available queues are:", ): add_queue_label(item, namespace, local_queue) + # Assert that no label is added + assert "labels" not in item["metadata"] + # Make sure to always keep this function last def test_cleanup(): diff --git a/src/codeflare_sdk/ray/cluster/build_ray_cluster.py b/src/codeflare_sdk/ray/cluster/build_ray_cluster.py index 2a3436b26..5ef9e6e91 100644 --- a/src/codeflare_sdk/ray/cluster/build_ray_cluster.py +++ b/src/codeflare_sdk/ray/cluster/build_ray_cluster.py @@ -17,6 +17,8 @@ (in the cluster sub-module) for RayCluster/AppWrapper generation. """ from typing import List, Union, Tuple, Dict + +from ...common.kueue.kueue import list_local_queues from ...common import _kube_api_error_handling from ...common.kubernetes_cluster import get_api_client, config_check from kubernetes.client.exceptions import ApiException @@ -482,15 +484,26 @@ def head_worker_extended_resources_from_cluster( # Local Queue related functions def add_queue_label(cluster: "codeflare_sdk.ray.cluster.Cluster", labels: dict): """ - The add_queue_label() function updates the given base labels with the local queue label if Kueue exists on the Cluster + The add_queue_label() function updates the given base labels with the local queue label if Kueue exists on the Cluster. + If no local_queue is provided, no queue label will be added. """ lq_name = cluster.config.local_queue or get_default_local_queue(cluster, labels) - if lq_name == None: + if lq_name is None: return - elif not local_queue_exists(cluster): - raise ValueError( - "local_queue provided does not exist or is not in this namespace. Please provide the correct local_queue name in Cluster Configuration" + elif cluster.config.local_queue and not local_queue_exists(cluster): + available_queues = list_local_queues(cluster.config.namespace) + if available_queues is None: + print( + f"WARNING: Local queue '{cluster.config.local_queue}' does not exist in namespace '{cluster.config.namespace}'. " + "Unable to retrieve list of available queues." + ) + return + available_queue_names = [q["name"] for q in available_queues] + print( + f"WARNING: Local queue '{cluster.config.local_queue}' does not exist in namespace '{cluster.config.namespace}'. " + f"Available queues are: {', '.join(available_queue_names)}" ) + return labels.update({"kueue.x-k8s.io/queue-name": lq_name}) From 7a9755153f9963236750225f69b7f46f4ad043af Mon Sep 17 00:00:00 2001 From: Pat O'Connor Date: Tue, 27 May 2025 12:23:15 +0100 Subject: [PATCH 2/2] updated formatting Signed-off-by: Pat O'Connor --- src/codeflare_sdk/common/kueue/test_kueue.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/src/codeflare_sdk/common/kueue/test_kueue.py b/src/codeflare_sdk/common/kueue/test_kueue.py index 3ebefa93c..5e039b4d7 100644 --- a/src/codeflare_sdk/common/kueue/test_kueue.py +++ b/src/codeflare_sdk/common/kueue/test_kueue.py @@ -279,8 +279,14 @@ def test_add_queue_label_with_invalid_local_queue_shows_available_queues(mocker) mock_api_instance = mocker.patch("kubernetes.client.CustomObjectsApi") mock_api_instance.return_value.list_namespaced_custom_object.return_value = { "items": [ - {"metadata": {"name": "queue1"}, "status": {"flavors": [{"name": "default"}]}}, - {"metadata": {"name": "queue2"}, "status": {"flavors": [{"name": "default"}]}}, + { + "metadata": {"name": "queue1"}, + "status": {"flavors": [{"name": "default"}]}, + }, + { + "metadata": {"name": "queue2"}, + "status": {"flavors": [{"name": "default"}]}, + }, ] } @@ -336,7 +342,7 @@ def test_add_queue_label_with_default_queue(mocker): { "metadata": { "name": "default-queue", - "annotations": {"kueue.x-k8s.io/default-queue": "true"} + "annotations": {"kueue.x-k8s.io/default-queue": "true"}, } } ]