RHOAIENG-34085: Fix ConfigMap mount

Author: kryanbeane

src/codeflare_sdk/ray/rayjobs/config.py

@@ -247,6 +247,7 @@ def build_ray_cluster_spec(self, cluster_name: str) -> Dict[str, Any]:
         """
         ray_cluster_spec = {
             "rayVersion": RAY_VERSION,
+            "enableInTreeAutoscaling": False,  # Required for Kueue-managed jobs
             "headGroupSpec": self._build_head_group_spec(),
             "workerGroupSpecs": [self._build_worker_group_spec(cluster_name)],
         }

src/codeflare_sdk/ray/rayjobs/rayjob.py

@@ -149,6 +149,14 @@ def submit(self) -> str:
 
         self._validate_ray_version_compatibility()
 
+        # Extract scripts to check if we need ConfigMaps
+        scripts = self._extract_script_files_from_entrypoint()
+
+        # Pre-declare ConfigMap in cluster config for new clusters
+        if scripts and self._cluster_config:
+            configmap_name = f"{self.name}-scripts"
+            self._cluster_config.add_script_volumes(configmap_name, MOUNT_PATH)
+
         rayjob_cr = self._build_rayjob_cr()
 
         logger.info(f"Submitting RayJob {self.name} to Kuberay operator")
@@ -157,20 +165,41 @@ def submit(self) -> str:
         if result:
             logger.info(f"Successfully submitted RayJob {self.name}")
 
-            # Handle script files after RayJob creation so we can set owner reference
-            if self._cluster_config is not None:
-                scripts = self._extract_script_files_from_entrypoint()
-                if scripts:
-                    self._handle_script_volumes_for_new_cluster(scripts, result)
-            elif self._cluster_name:
-                scripts = self._extract_script_files_from_entrypoint()
-                if scripts:
-                    self._handle_script_volumes_for_existing_cluster(scripts, result)
+            # Create ConfigMap with owner reference after RayJob exists
+            if scripts:
+                self._create_script_configmap(scripts, result)
 
             return self.name
         else:
             raise RuntimeError(f"Failed to submit RayJob {self.name}")
 
+    def _create_script_configmap(
+        self, scripts: Dict[str, str], rayjob_result: Dict[str, Any]
+    ):
+        """
+        Create ConfigMap with owner reference for script files.
+
+        For new clusters: ConfigMap volume was pre-declared, just create it.
+        For existing clusters: Create ConfigMap and patch the cluster.
+        """
+        # Get a config builder for utility methods
+        config_builder = (
+            self._cluster_config if self._cluster_config else ManagedClusterConfig()
+        )
+
+        # Validate and build ConfigMap spec
+        config_builder.validate_configmap_size(scripts)
+        configmap_spec = config_builder.build_script_configmap_spec(
+            job_name=self.name, namespace=self.namespace, scripts=scripts
+        )
+
+        # Create ConfigMap with owner reference
+        configmap_name = self._create_configmap_from_spec(configmap_spec, rayjob_result)
+
+        # For existing clusters, update the cluster with volumes
+        if self._cluster_name and not self._cluster_config:
+            self._update_existing_cluster_for_scripts(configmap_name, config_builder)
+
     def stop(self):
         """
         Suspend the Ray job.
@@ -488,47 +517,6 @@ def _find_local_imports(
         except (SyntaxError, ValueError) as e:
             logger.debug(f"Could not parse imports from {script_path}: {e}")
 
-    def _handle_script_volumes_for_new_cluster(
-        self, scripts: Dict[str, str], rayjob_result: Dict[str, Any] = None
-    ):
-        """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
-        )
-
-        # Create ConfigMap via Kubernetes API with owner reference
-        configmap_name = self._create_configmap_from_spec(configmap_spec, rayjob_result)
-
-        # Add volumes to cluster config (config.py handles spec building)
-        self._cluster_config.add_script_volumes(
-            configmap_name=configmap_name, mount_path=MOUNT_PATH
-        )
-
-    def _handle_script_volumes_for_existing_cluster(
-        self, scripts: Dict[str, str], rayjob_result: Dict[str, Any] = None
-    ):
-        """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(
-            job_name=self.name, namespace=self.namespace, scripts=scripts
-        )
-
-        # Create ConfigMap via Kubernetes API with owner reference
-        configmap_name = self._create_configmap_from_spec(configmap_spec, rayjob_result)
-
-        # Update existing RayCluster
-        self._update_existing_cluster_for_scripts(configmap_name, config_builder)
-
     def _create_configmap_from_spec(
         self, configmap_spec: Dict[str, Any], rayjob_result: Dict[str, Any] = None
     ) -> str:

src/codeflare_sdk/ray/rayjobs/test_rayjob.py

@@ -586,6 +586,8 @@ def test_build_ray_cluster_spec_function():
 
     spec = cluster_config.build_ray_cluster_spec("test-cluster")
     assert "rayVersion" in spec
+    assert "enableInTreeAutoscaling" in spec
+    assert spec["enableInTreeAutoscaling"] is False  # Required for Kueue
     assert "headGroupSpec" in spec
     assert "workerGroupSpecs" in spec
 
@@ -1304,11 +1306,13 @@ def func2(): pass
         os.chdir(original_cwd)
 
 
-def test_script_handling_timing_after_rayjob_submission(
-    mocker, auto_mock_setup, tmp_path
-):
-    """Test that script handling happens after RayJob is submitted (not before)."""
-    mock_api_instance = auto_mock_setup["rayjob_api"]
+def test_script_handling_kubernetes_best_practice_flow(mocker, tmp_path):
+    """Test the Kubernetes best practice flow: pre-declare volume, submit, create ConfigMap."""
+    mocker.patch("kubernetes.config.load_kube_config")
+
+    mock_api_class = mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayjobApi")
+    mock_api_instance = MagicMock()
+    mock_api_class.return_value = mock_api_instance
 
     submit_result = {
         "metadata": {
@@ -1319,9 +1323,8 @@ def test_script_handling_timing_after_rayjob_submission(
     }
     mock_api_instance.submit_job.return_value = submit_result
 
-    mock_handle_new = mocker.patch.object(
-        RayJob, "_handle_script_volumes_for_new_cluster"
-    )
+    mock_create_cm = mocker.patch.object(RayJob, "_create_script_configmap")
+    mock_add_volumes = mocker.patch.object(ManagedClusterConfig, "add_script_volumes")
 
     # RayClusterApi is already mocked by auto_mock_setup
 
@@ -1330,17 +1333,22 @@ def test_script_handling_timing_after_rayjob_submission(
 
     call_order = []
 
+    def track_add_volumes(*args, **kwargs):
+        call_order.append("add_volumes")
+        # Should be called with ConfigMap name
+        assert args[0] == "test-job-scripts"
+
     def track_submit(*args, **kwargs):
         call_order.append("submit_job")
         return submit_result
 
-    def track_handle_scripts(*args, **kwargs):
-        call_order.append("handle_scripts")
-        assert len(args) >= 2
+    def track_create_cm(*args, **kwargs):
+        call_order.append("create_configmap")
         assert args[1] == submit_result  # rayjob_result should be second arg
 
+    mock_add_volumes.side_effect = track_add_volumes
     mock_api_instance.submit_job.side_effect = track_submit
-    mock_handle_new.side_effect = track_handle_scripts
+    mock_create_cm.side_effect = track_create_cm
 
     original_cwd = os.getcwd()
     try:
@@ -1359,12 +1367,14 @@ def track_handle_scripts(*args, **kwargs):
     finally:
         os.chdir(original_cwd)
 
-    assert call_order == ["submit_job", "handle_scripts"]
+    # Verify the order: add volumes → submit → create ConfigMap
+    assert call_order == ["add_volumes", "submit_job", "create_configmap"]
 
+    mock_add_volumes.assert_called_once()
     mock_api_instance.submit_job.assert_called_once()
-    mock_handle_new.assert_called_once()
+    mock_create_cm.assert_called_once()
 
-    mock_handle_new.assert_called_with({"test.py": "print('test')"}, submit_result)
+    mock_create_cm.assert_called_with({"test.py": "print('test')"}, submit_result)
 
 
 def test_rayjob_submit_with_scripts_new_cluster(auto_mock_setup, tmp_path):

tests/e2e/rayjob/rayjob_lifecycled_cluster_test.py

@@ -2,11 +2,14 @@
 import sys
 import os
 from time import sleep
+import tempfile
 
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
 from support import *
 
 from codeflare_sdk import RayJob, ManagedClusterConfig
+
+from kubernetes import client
 from python_client.kuberay_job_api import RayjobApi
 from python_client.kuberay_cluster_api import RayClusterApi
 
@@ -22,7 +25,7 @@ def teardown_method(self):
         delete_kueue_resources(self)
 
     def test_lifecycled_kueue_managed(self):
-        """Test RayJob with Kueue-managed lifecycled cluster."""
+        """Test RayJob with Kueue-managed lifecycled cluster with ConfigMap validation."""
         self.setup_method()
         create_namespace(self)
         create_kueue_resources(self)
@@ -46,17 +49,36 @@ def test_lifecycled_kueue_managed(self):
             worker_memory_limits=resources["worker_memory_limits"],
         )
 
-        rayjob = RayJob(
-            job_name=job_name,
-            namespace=self.namespace,
-            cluster_config=cluster_config,
-            entrypoint="python -c \"import ray; ray.init(); print('Kueue job done')\"",
-            runtime_env={"env_vars": get_setup_env_variables(ACCELERATOR="cpu")},
-            local_queue=self.local_queues[0],
-        )
+        # Create a temporary script file to test ConfigMap functionality
+        with tempfile.NamedTemporaryFile(
+            mode="w", suffix=".py", delete=False, dir=os.getcwd()
+        ) as script_file:
+            script_file.write(
+                """
+                import ray
+                ray.init()
+                print('Kueue job with ConfigMap done')
+                ray.shutdown()
+                """
+            )
+            script_file.flush()
+            script_filename = os.path.basename(script_file.name)
 
         try:
+            rayjob = RayJob(
+                job_name=job_name,
+                namespace=self.namespace,
+                cluster_config=cluster_config,
+                entrypoint=f"python {script_filename}",
+                runtime_env={"env_vars": get_setup_env_variables(ACCELERATOR="cpu")},
+                local_queue=self.local_queues[0],
+            )
+
             assert rayjob.submit() == job_name
+
+            # Verify ConfigMap was created with owner reference
+            self.verify_configmap_with_owner_reference(rayjob)
+
             assert self.job_api.wait_until_job_running(
                 name=rayjob.name, k8s_namespace=rayjob.namespace, timeout=600
             )
@@ -70,6 +92,12 @@ def test_lifecycled_kueue_managed(self):
             except Exception:
                 pass  # Job might already be deleted
             verify_rayjob_cluster_cleanup(cluster_api, rayjob.name, rayjob.namespace)
+            # Clean up the temporary script file
+            if "script_filename" in locals():
+                try:
+                    os.remove(script_filename)
+                except:
+                    pass
 
     def test_lifecycled_kueue_resource_queueing(self):
         """Test Kueue resource queueing with lifecycled clusters."""
@@ -161,3 +189,61 @@ def test_lifecycled_kueue_resource_queueing(self):
                         )
                     except:
                         pass
+
+    def verify_configmap_with_owner_reference(self, rayjob: RayJob):
+        """Verify that the ConfigMap was created with proper owner reference to the RayJob."""
+        v1 = client.CoreV1Api()
+        configmap_name = f"{rayjob.name}-scripts"
+
+        try:
+            # Get the ConfigMap
+            configmap = v1.read_namespaced_config_map(
+                name=configmap_name, namespace=rayjob.namespace
+            )
+
+            # Verify ConfigMap exists
+            assert configmap is not None, f"ConfigMap {configmap_name} not found"
+
+            # Verify it contains the script
+            assert configmap.data is not None, "ConfigMap has no data"
+            assert len(configmap.data) > 0, "ConfigMap data is empty"
+
+            # Verify owner reference
+            assert (
+                configmap.metadata.owner_references is not None
+            ), "ConfigMap has no owner references"
+            assert (
+                len(configmap.metadata.owner_references) > 0
+            ), "ConfigMap owner references list is empty"
+
+            owner_ref = configmap.metadata.owner_references[0]
+            assert (
+                owner_ref.api_version == "ray.io/v1"
+            ), f"Wrong API version: {owner_ref.api_version}"
+            assert owner_ref.kind == "RayJob", f"Wrong kind: {owner_ref.kind}"
+            assert owner_ref.name == rayjob.name, f"Wrong owner name: {owner_ref.name}"
+            assert (
+                owner_ref.controller is True
+            ), "Owner reference controller not set to true"
+            assert (
+                owner_ref.block_owner_deletion is True
+            ), "Owner reference blockOwnerDeletion not set to true"
+
+            # Verify labels
+            assert configmap.metadata.labels.get("ray.io/job-name") == rayjob.name
+            assert (
+                configmap.metadata.labels.get("app.kubernetes.io/managed-by")
+                == "codeflare-sdk"
+            )
+            assert (
+                configmap.metadata.labels.get("app.kubernetes.io/component")
+                == "rayjob-scripts"
+            )
+
+            print(f"✓ ConfigMap {configmap_name} verified with proper owner reference")
+
+        except client.rest.ApiException as e:
+            if e.status == 404:
+                raise AssertionError(f"ConfigMap {configmap_name} not found")
+            else:
+                raise e