Skip to content

Commit

Permalink
Adds RayCluster.apply()
Browse files Browse the repository at this point in the history
- 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
  • Loading branch information
akram committed Dec 16, 2024
1 parent be9763a commit d02ec50
Show file tree
Hide file tree
Showing 11 changed files with 568 additions and 37 deletions.
7 changes: 7 additions & 0 deletions .coveragerc
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
[run]
omit =
src/**/test_*.py,,src/codeflare_sdk/common/utils/unit_test_support.py

[report]
exclude_lines =
pragma: no cover
2 changes: 1 addition & 1 deletion .github/workflows/coverage-badge.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ jobs:
poetry install --with test
- name: Generate coverage report
run: |
coverage run -m pytest
coverage run -m pytest --omit="src/**/test_*.py,src/codeflare_sdk/common/utils/unit_test_support.py"
- name: Coverage Badge
uses: tj-actions/coverage-badge-py@v2
Expand Down
2 changes: 1 addition & 1 deletion CONTRIBUTING.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
12 changes: 6 additions & 6 deletions src/codeflare_sdk/common/kueue/test_kueue.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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"
Expand All @@ -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)
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand Down
66 changes: 55 additions & 11 deletions src/codeflare_sdk/common/utils/unit_test_support.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,32 +26,34 @@
aw_dir = os.path.expanduser("~/.codeflare/resources/")


def createClusterConfig():
def create_cluster_config(num_workers=2, write_to_file=False):
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,
worker_memory_limits=6,
appwrapper=True,
write_to_file=False,
write_to_file=write_to_file,
)
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, write_to_file=False):
cluster = Cluster(create_cluster_config(num_workers, write_to_file))
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",
Expand Down Expand Up @@ -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(
Expand Down
4 changes: 2 additions & 2 deletions src/codeflare_sdk/common/widgets/test_widgets.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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"
Expand Down
95 changes: 87 additions & 8 deletions src/codeflare_sdk/ray/cluster/cluster.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -84,6 +88,14 @@ def __init__(self, config: ClusterConfiguration):
if is_notebook():
cluster_up_down_buttons(self)

def get_dynamic_client(self): # pragma: no cover
"""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()
Expand All @@ -95,9 +107,7 @@ def _client_headers(self):

@property
def _client_verify_tls(self):
if not _is_openshift_cluster or not self.config.verify_tls:
return False
return True
return _is_openshift_cluster and self.config.verify_tls

@property
def job_client(self):
Expand All @@ -121,7 +131,6 @@ def create_resource(self):
Called upon cluster object creation, creates an AppWrapper yaml based on
the specifications of the ClusterConfiguration.
"""

if self.config.namespace is None:
self.config.namespace = get_current_namespace()
if self.config.namespace is None:
Expand All @@ -130,7 +139,6 @@ def create_resource(self):
raise TypeError(
f"Namespace {self.config.namespace} is of type {type(self.config.namespace)}. Check your Kubernetes Authentication."
)

return build_ray_cluster(self)

# creates a new cluster with the provided or default spec
Expand All @@ -139,10 +147,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:
Expand Down Expand Up @@ -176,6 +185,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)
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:
Expand Down Expand Up @@ -204,7 +259,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(
Expand Down Expand Up @@ -507,6 +562,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
):
Expand Down Expand Up @@ -744,6 +809,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()
Expand Down
Loading

0 comments on commit d02ec50

Please sign in to comment.