Adds RayCluster.apply()

- Adds  RayCluster.apply() implementation
- Adds e2e tests for apply
- Adds unit tests for apply

Author: akram

CONTRIBUTING.md

@@ -76,7 +76,7 @@ pytest -v src/codeflare_sdk
 
 ### Local e2e Testing
 
-- Please follow the [e2e documentation](https://github.com/project-codeflare/codeflare-sdk/blob/main/docs/e2e.md)
+- Please follow the [e2e documentation](https://github.com/project-codeflare/codeflare-sdk/blob/main/docs/sphinx/user-docs/e2e.rst)
 
 #### Code Coverage
 

src/codeflare_sdk/common/kueue/test_kueue.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from ..utils.unit_test_support import get_local_queue, createClusterConfig
+from ..utils.unit_test_support import get_local_queue, create_cluster_config
 from unittest.mock import patch
 from codeflare_sdk.ray.cluster.cluster import Cluster, ClusterConfiguration
 import yaml
@@ -46,7 +46,7 @@ def test_cluster_creation_no_aw_local_queue(mocker):
         "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
         return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"),
     )
-    config = createClusterConfig()
+    config = create_cluster_config()
     config.name = "unit-test-cluster-kueue"
     config.write_to_file = True
     config.local_queue = "local-queue-default"
@@ -59,7 +59,7 @@ def test_cluster_creation_no_aw_local_queue(mocker):
     )
 
     # With resources loaded in memory, no Local Queue specified.
-    config = createClusterConfig()
+    config = create_cluster_config()
     config.name = "unit-test-cluster-kueue"
     config.write_to_file = False
     cluster = Cluster(config)
@@ -79,7 +79,7 @@ def test_aw_creation_local_queue(mocker):
         "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
         return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"),
     )
-    config = createClusterConfig()
+    config = create_cluster_config()
     config.name = "unit-test-aw-kueue"
     config.appwrapper = True
     config.write_to_file = True
@@ -93,7 +93,7 @@ def test_aw_creation_local_queue(mocker):
     )
 
     # With resources loaded in memory, no Local Queue specified.
-    config = createClusterConfig()
+    config = create_cluster_config()
     config.name = "unit-test-aw-kueue"
     config.appwrapper = True
     config.write_to_file = False
@@ -114,7 +114,7 @@ def test_get_local_queue_exists_fail(mocker):
         "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
         return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"),
     )
-    config = createClusterConfig()
+    config = create_cluster_config()
     config.name = "unit-test-aw-kueue"
     config.appwrapper = True
     config.write_to_file = True

src/codeflare_sdk/common/utils/unit_test_support.py

@@ -26,11 +26,11 @@
 aw_dir = os.path.expanduser("~/.codeflare/resources/")
 
 
-def createClusterConfig():
+def create_cluster_config(num_workers=2):
     config = ClusterConfiguration(
         name="unit-test-cluster",
         namespace="ns",
-        num_workers=2,
+        num_workers=num_workers,
         worker_cpu_requests=3,
         worker_cpu_limits=4,
         worker_memory_requests=5,
@@ -41,17 +41,19 @@ def createClusterConfig():
     return config
 
 
-def createClusterWithConfig(mocker):
-    mocker.patch("kubernetes.config.load_kube_config", return_value="ignore")
-    mocker.patch(
-        "kubernetes.client.CustomObjectsApi.get_cluster_custom_object",
-        return_value={"spec": {"domain": "apps.cluster.awsroute.org"}},
-    )
-    cluster = Cluster(createClusterConfig())
+def create_cluster(mocker, num_workers=2):
+    cluster = Cluster(create_cluster_config(num_workers))
     return cluster
 
 
-def createClusterWrongType():
+def patch_cluster_with_dynamic_client(mocker, cluster, dynamic_client=None):
+    mocker.patch.object(cluster, "get_dynamic_client", return_value=dynamic_client)
+    mocker.patch.object(cluster, "down", return_value=None)
+    mocker.patch.object(cluster, "config_check", return_value=None)
+    # mocker.patch.object(cluster, "_throw_for_no_raycluster", return_value=None)
+
+
+def create_cluster_wrong_type():
     config = ClusterConfiguration(
         name="unit-test-cluster",
         namespace="ns",
@@ -383,6 +385,48 @@ def mocked_ingress(port, cluster_name="unit-test-cluster", annotations: dict = N
     return mock_ingress
 
 
+# Global dictionary to maintain state in the mock
+cluster_state = {}
+
+
+# The mock side_effect function for server_side_apply
+def mock_server_side_apply(resource, body=None, name=None, namespace=None, **kwargs):
+    # Simulate the behavior of server_side_apply:
+    # Update a mock state that represents the cluster's current configuration.
+    # Stores the state in a global dictionary for simplicity.
+
+    global cluster_state
+
+    if not resource or not body or not name or not namespace:
+        raise ValueError("Missing required parameters for server_side_apply")
+
+    # Extract worker count from the body if it exists
+    try:
+        worker_count = (
+            body["spec"]["workerGroupSpecs"][0]["replicas"]
+            if "spec" in body and "workerGroupSpecs" in body["spec"]
+            else None
+        )
+    except KeyError:
+        worker_count = None
+
+    # Apply changes to the cluster_state mock
+    cluster_state[name] = {
+        "namespace": namespace,
+        "worker_count": worker_count,
+        "body": body,
+    }
+
+    # Return a response that mimics the behavior of a successful apply
+    return {
+        "status": "success",
+        "applied": True,
+        "name": name,
+        "namespace": namespace,
+        "worker_count": worker_count,
+    }
+
+
 @patch.dict("os.environ", {"NB_PREFIX": "test-prefix"})
 def create_cluster_all_config_params(mocker, cluster_name, is_appwrapper) -> Cluster:
     mocker.patch(

src/codeflare_sdk/common/widgets/test_widgets.py

@@ -15,7 +15,7 @@
 import codeflare_sdk.common.widgets.widgets as cf_widgets
 import pandas as pd
 from unittest.mock import MagicMock, patch
-from ..utils.unit_test_support import get_local_queue, createClusterConfig
+from ..utils.unit_test_support import get_local_queue, create_cluster_config
 from codeflare_sdk.ray.cluster.cluster import Cluster
 from codeflare_sdk.ray.cluster.status import (
     RayCluster,
@@ -38,7 +38,7 @@ def test_cluster_up_down_buttons(mocker):
         "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
         return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"),
     )
-    cluster = Cluster(createClusterConfig())
+    cluster = Cluster(create_cluster_config())
 
     with patch("ipywidgets.Button") as MockButton, patch(
         "ipywidgets.Checkbox"

src/codeflare_sdk/ray/cluster/cluster.py

@@ -52,6 +52,10 @@
 import requests
 
 from kubernetes import config
+from kubernetes.dynamic import DynamicClient
+from kubernetes import client as k8s_client
+from kubernetes.client.rest import ApiException
+
 from kubernetes.client.rest import ApiException
 import warnings
 
@@ -84,6 +88,14 @@ def __init__(self, config: ClusterConfiguration):
         if is_notebook():
             cluster_up_down_buttons(self)
 
+    def get_dynamic_client(self):
+        """Return a dynamic client, optionally mocked in tests."""
+        return DynamicClient(get_api_client())
+
+    def config_check(self):
+        """Return a dynamic client, optionally mocked in tests."""
+        return config_check()
+
     @property
     def _client_headers(self):
         k8_client = get_api_client()
@@ -139,10 +151,11 @@ def up(self):
         Applies the Cluster yaml, pushing the resource request onto
         the Kueue localqueue.
         """
-
+        print(
+            "WARNING: The up() function is planned for deprecation in favor of apply()."
+        )
         # check if RayCluster CustomResourceDefinition exists if not throw RuntimeError
         self._throw_for_no_raycluster()
-
         namespace = self.config.namespace
 
         try:
@@ -176,6 +189,52 @@ def up(self):
         except Exception as e:  # pragma: no cover
             return _kube_api_error_handling(e)
 
+    # Applies a new cluster with the provided or default spec
+    def apply(self, force=False):
+        """
+        Applies the Cluster yaml using server-side apply.
+        If 'force' is set to True, conflicts will be forced.
+        """
+        # check if RayCluster CustomResourceDefinition exists if not throw RuntimeError
+        self._throw_for_no_raycluster()
+        namespace = self.config.namespace
+
+        try:
+            self.config_check()
+            api_instance = client.CustomObjectsApi(get_api_client())
+            crds = self.get_dynamic_client().resources
+            api_instance = crds.get(
+                api_version="workload.codeflare.dev/v1beta2", kind="AppWrapper"
+            )
+            if self.config.appwrapper:
+                if self.config.write_to_file:
+                    with open(self.resource_yaml) as f:
+                        aw = yaml.load(f, Loader=yaml.FullLoader)
+                        api_instance.server_side_apply(
+                            group="workload.codeflare.dev",
+                            version="v1beta2",
+                            namespace=namespace,
+                            plural="appwrappers",
+                            body=aw,
+                        )
+                else:
+                    api_instance.server_side_apply(
+                        group="workload.codeflare.dev",
+                        version="v1beta2",
+                        namespace=namespace,
+                        plural="appwrappers",
+                        body=self.resource_yaml,
+                    )
+                print(f"AppWrapper: '{self.config.name}' has successfully been created")
+            else:
+                api_instance = crds.get(api_version="ray.io/v1", kind="RayCluster")
+                self._component_resources_apply(namespace, api_instance, force)
+                print(
+                    f"Ray Cluster: '{self.config.name}' has successfully been applied"
+                )
+        except Exception as e:  # pragma: no cover
+            return _kube_api_error_handling(e)
+
     def _throw_for_no_raycluster(self):
         api_instance = client.CustomObjectsApi(get_api_client())
         try:
@@ -204,7 +263,7 @@ def down(self):
         resource_name = self.config.name
         self._throw_for_no_raycluster()
         try:
-            config_check()
+            self.config_check()
             api_instance = client.CustomObjectsApi(get_api_client())
             if self.config.appwrapper:
                 api_instance.delete_namespaced_custom_object(
@@ -507,6 +566,16 @@ def _component_resources_up(
         else:
             _create_resources(self.resource_yaml, namespace, api_instance)
 
+    def _component_resources_apply(
+        self, namespace: str, api_instance: client.CustomObjectsApi
+    ):
+        if self.config.write_to_file:
+            with open(self.resource_yaml) as f:
+                ray_cluster = yaml.safe_load(f)
+                _apply_resources(ray_cluster, namespace, api_instance)
+        else:
+            _apply_resources(self.resource_yaml, namespace, api_instance)
+
     def _component_resources_down(
         self, namespace: str, api_instance: client.CustomObjectsApi
     ):
@@ -744,6 +813,20 @@ def _create_resources(yamls, namespace: str, api_instance: client.CustomObjectsA
     )
 
 
+def _apply_resources(
+    yamls, namespace: str, api_instance: client.CustomObjectsApi, force=False
+):
+    api_instance.server_side_apply(
+        field_manager="cluster-manager",
+        group="ray.io",
+        version="v1",
+        namespace=namespace,
+        plural="rayclusters",
+        body=yamls,
+        force_conflicts=force,  # Allow forcing conflicts if needed
+    )
+
+
 def _check_aw_exists(name: str, namespace: str) -> bool:
     try:
         config_check()