Changes as per review

Signed-off-by: Pat O'Connor <[email protected]>

Author: chipspeak

src/codeflare_sdk/ray/rayjobs/config.py

@@ -559,6 +559,13 @@ def add_script_volumes(
         logger.info(
             f"Added script volume '{configmap_name}' to cluster config: mount_path={mount_path}"
         )
+        
+    def validate_configmap_size(self, scripts: Dict[str, str]) -> None:
+        total_size = sum(len(content.encode("utf-8")) for content in scripts.values())
+        if total_size > 1024 * 1024:  # 1MB
+            raise ValueError(
+                f"ConfigMap size exceeds 1MB limit. Total size: {total_size} bytes"
+            )
 
     def build_script_configmap_spec(
         self, job_name: str, namespace: str, scripts: Dict[str, str]

src/codeflare_sdk/ray/rayjobs/rayjob.py

@@ -24,7 +24,7 @@
 from kubernetes import client
 from ...common.kubernetes_cluster.auth import get_api_client
 from python_client.kuberay_job_api import RayjobApi
-
+from python_client.kuberay_cluster_api import RayClusterApi
 from codeflare_sdk.ray.rayjobs.config import ManagedClusterConfig
 
 from ...common.utils import get_current_namespace
@@ -39,6 +39,8 @@
 
 logger = logging.getLogger(__name__)
 
+mount_path = "/home/ray/scripts"
+
 
 class RayJob:
     """
@@ -146,6 +148,7 @@ def __init__(
             logger.info(f"Using existing cluster: {self.cluster_name}")
 
         self._api = RayjobApi()
+        self._cluster_api = RayClusterApi()
 
         logger.info(f"Initialized RayJob: {self.name} in namespace: {self.namespace}")
 
@@ -311,7 +314,7 @@ def _extract_script_files_from_entrypoint(self) -> Optional[Dict[str, str]]:
             return None
 
         scripts = {}
-        mount_path = "/home/ray/scripts"
+        # mount_path = "/home/ray/scripts"
         processed_files = set()  # Avoid infinite loops
 
         # Look for Python file patterns in entrypoint (e.g., "python script.py", "python /path/to/script.py")
@@ -410,6 +413,9 @@ def _find_local_imports(
 
     def _handle_script_volumes_for_new_cluster(self, scripts: Dict[str, str]):
         """Handle script volumes for new clusters (uses ManagedClusterConfig)."""
+        # Validate ConfigMap size before creation
+        self._cluster_config.validate_configmap_size(scripts)
+        
         # Build ConfigMap spec using config.py
         configmap_spec = self._cluster_config.build_script_configmap_spec(
             job_name=self.name, namespace=self.namespace, scripts=scripts
@@ -427,6 +433,9 @@ def _handle_script_volumes_for_existing_cluster(self, scripts: Dict[str, str]):
         """Handle script volumes for existing clusters (updates RayCluster CR)."""
         # Create config builder for utility methods
         config_builder = ManagedClusterConfig()
+        
+        # Validate ConfigMap size before creation
+        config_builder.validate_configmap_size(scripts)
 
         # Build ConfigMap spec using config.py
         configmap_spec = config_builder.build_script_configmap_spec(
@@ -495,12 +504,9 @@ def _update_existing_cluster_for_scripts(
         # Get existing RayCluster
         api_instance = client.CustomObjectsApi(get_api_client())
         try:
-            ray_cluster = api_instance.get_namespaced_custom_object(
-                group="ray.io",
-                version="v1",
-                namespace=self.namespace,
-                plural="rayclusters",
+            ray_cluster = self._cluster_api.get_ray_cluster(
                 name=self.cluster_name,
+                k8s_namespace=self.namespace,
             )
         except client.ApiException as e:
             raise RuntimeError(f"Failed to get RayCluster '{self.cluster_name}': {e}")
@@ -546,13 +552,10 @@ def mount_exists(mounts_list, mount_name):
 
         # Update the RayCluster
         try:
-            api_instance.patch_namespaced_custom_object(
-                group="ray.io",
-                version="v1",
-                namespace=self.namespace,
-                plural="rayclusters",
+            self._cluster_api.patch_ray_cluster(
                 name=self.cluster_name,
-                body=ray_cluster,
+                ray_cluster=ray_cluster,
+                k8s_namespace=self.namespace,
             )
             logger.info(
                 f"Updated RayCluster '{self.cluster_name}' with script volumes from ConfigMap '{configmap_name}'"

src/codeflare_sdk/ray/rayjobs/test_rayjob.py

@@ -31,6 +31,9 @@ def test_rayjob_submit_success(mocker):
     mock_api_instance = MagicMock()
     mock_api_class.return_value = mock_api_instance
 
+    # Mock the RayClusterApi class
+    mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayClusterApi")
+
     # Configure the mock to return success when submit is called
     mock_api_instance.submit.return_value = {"metadata": {"name": "test-rayjob"}}
 
@@ -75,6 +78,9 @@ def test_rayjob_submit_failure(mocker):
     mock_api_instance = MagicMock()
     mock_api_class.return_value = mock_api_instance
 
+    # Mock the RayClusterApi class
+    mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayClusterApi")
+
     # Configure the mock to return failure (False/None) when submit_job is called
     mock_api_instance.submit_job.return_value = None
 
@@ -1580,21 +1586,19 @@ def test_create_configmap_api_error_non_409(mocker):
 def test_update_existing_cluster_get_cluster_error(mocker):
     """Test _update_existing_cluster_for_scripts handles get cluster errors."""
     mocker.patch("kubernetes.config.load_kube_config")
-    mock_api_class = mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayjobApi")
+    mock_rayjob_api = mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayjobApi")
 
-    # Mock CustomObjectsApi with error
-    mock_custom_api = mocker.patch("kubernetes.client.CustomObjectsApi")
-    mock_api_instance = mocker.Mock()
-    mock_custom_api.return_value = mock_api_instance
+    # Mock RayClusterApi with error
+    mock_cluster_api_class = mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayClusterApi")
+    mock_cluster_api_instance = mocker.Mock()
+    mock_cluster_api_class.return_value = mock_cluster_api_instance
 
     from kubernetes.client import ApiException
 
-    mock_api_instance.get_namespaced_custom_object.side_effect = ApiException(
+    mock_cluster_api_instance.get_ray_cluster.side_effect = ApiException(
         status=404
     )
 
-    mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.get_api_client")
-
     from codeflare_sdk.ray.rayjobs.config import ManagedClusterConfig
 
     config_builder = ManagedClusterConfig()
@@ -1614,15 +1618,15 @@ def test_update_existing_cluster_get_cluster_error(mocker):
 def test_update_existing_cluster_patch_error(mocker):
     """Test _update_existing_cluster_for_scripts handles patch errors."""
     mocker.patch("kubernetes.config.load_kube_config")
-    mock_api_class = mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayjobApi")
+    mock_rayjob_api = mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayjobApi")
 
-    # Mock CustomObjectsApi
-    mock_custom_api = mocker.patch("kubernetes.client.CustomObjectsApi")
-    mock_api_instance = mocker.Mock()
-    mock_custom_api.return_value = mock_api_instance
+    # Mock RayClusterApi
+    mock_cluster_api_class = mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayClusterApi")
+    mock_cluster_api_instance = mocker.Mock()
+    mock_cluster_api_class.return_value = mock_cluster_api_instance
 
     # Mock successful get but failed patch
-    mock_api_instance.get_namespaced_custom_object.return_value = {
+    mock_cluster_api_instance.get_ray_cluster.return_value = {
         "spec": {
             "headGroupSpec": {
                 "template": {
@@ -1641,12 +1645,10 @@ def test_update_existing_cluster_patch_error(mocker):
 
     from kubernetes.client import ApiException
 
-    mock_api_instance.patch_namespaced_custom_object.side_effect = ApiException(
+    mock_cluster_api_instance.patch_ray_cluster.side_effect = ApiException(
         status=500
     )
 
-    mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.get_api_client")
-
     from codeflare_sdk.ray.rayjobs.config import ManagedClusterConfig
 
     config_builder = ManagedClusterConfig()