RHOAIENG-27792: Auto tear down training config map when job is deleted

Author: kryanbeane

poetry.lock

@@ -33,7 +33,7 @@ cryptography = "43.0.3"
 executing = "1.2.0"
 pydantic = "< 2"
 ipywidgets = "8.1.2"
-python-client = { git = "https://github.com/ray-project/kuberay.git", subdirectory = "clients/python-client", rev = "49419654418865a5838adc7f323f13d82454aa18" }
+python-client = { git = "https://github.com/ray-project/kuberay.git", subdirectory = "clients/python-client", rev = "a16c0365e3b19a202d835097e1139eca9406b383" }
 
 [[tool.poetry.source]]
 name = "pypi"

pyproject.toml

@@ -585,7 +585,15 @@ def build_script_configmap_spec(
         return {
             "apiVersion": "v1",
             "kind": "ConfigMap",
-            "metadata": {"name": configmap_name, "namespace": namespace},
+            "metadata": {
+                "name": configmap_name,
+                "namespace": namespace,
+                "labels": {
+                    "ray.io/job-name": job_name,
+                    "app.kubernetes.io/managed-by": "codeflare-sdk",
+                    "app.kubernetes.io/component": "rayjob-scripts",
+                },
+            },
             "data": scripts,
         }
 

src/codeflare_sdk/ray/rayjobs/config.py

@@ -160,23 +160,24 @@ def submit(self) -> str:
 
         self._validate_ray_version_compatibility()
 
-        # Automatically handle script files for new clusters
-        if self._cluster_config is not None:
-            scripts = self._extract_script_files_from_entrypoint()
-            if scripts:
-                self._handle_script_volumes_for_new_cluster(scripts)
-        elif self._cluster_name:
-            scripts = self._extract_script_files_from_entrypoint()
-            if scripts:
-                self._handle_script_volumes_for_existing_cluster(scripts)
-
         rayjob_cr = self._build_rayjob_cr()
 
         logger.info(f"Submitting RayJob {self.name} to Kuberay operator")
         result = self._api.submit_job(k8s_namespace=self.namespace, job=rayjob_cr)
 
         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)
+
             if self.shutdown_after_job_finishes:
                 logger.info(
                     f"Cluster will be automatically cleaned up {self.ttl_seconds_after_finished}s after job completion"
@@ -206,6 +207,17 @@ def resubmit(self):
         else:
             raise RuntimeError(f"Failed to resubmit the RayJob {self.name}")
 
+    def delete(self):
+        """
+        Delete the Ray job.
+        """
+        deleted = self._api.delete_job(name=self.name, k8s_namespace=self.namespace)
+        if deleted:
+            logger.info(f"Successfully deleted the RayJob {self.name}")
+            return True
+        else:
+            raise RuntimeError(f"Failed to delete the RayJob {self.name}")
+
     def _build_rayjob_cr(self) -> Dict[str, Any]:
         """
         Build the RayJob custom resource specification using native RayJob capabilities.
@@ -466,7 +478,9 @@ 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]):
+    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)
@@ -476,15 +490,17 @@ def _handle_script_volumes_for_new_cluster(self, scripts: Dict[str, str]):
             job_name=self.name, namespace=self.namespace, scripts=scripts
         )
 
-        # Create ConfigMap via Kubernetes API
-        configmap_name = self._create_configmap_from_spec(configmap_spec)
+        # 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="/home/ray/scripts"
         )
 
-    def _handle_script_volumes_for_existing_cluster(self, scripts: Dict[str, str]):
+    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()
@@ -497,28 +513,57 @@ def _handle_script_volumes_for_existing_cluster(self, scripts: Dict[str, str]):
             job_name=self.name, namespace=self.namespace, scripts=scripts
         )
 
-        # Create ConfigMap via Kubernetes API
-        configmap_name = self._create_configmap_from_spec(configmap_spec)
+        # 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]) -> str:
+    def _create_configmap_from_spec(
+        self, configmap_spec: Dict[str, Any], rayjob_result: Dict[str, Any] = None
+    ) -> str:
         """
         Create ConfigMap from specification via Kubernetes API.
 
         Args:
             configmap_spec: ConfigMap specification dictionary
+            rayjob_result: The result from RayJob creation containing UID
 
         Returns:
             str: Name of the created ConfigMap
         """
 
         configmap_name = configmap_spec["metadata"]["name"]
 
+        metadata = client.V1ObjectMeta(**configmap_spec["metadata"])
+
+        # Add owner reference if we have the RayJob result
+        if (
+            rayjob_result
+            and isinstance(rayjob_result, dict)
+            and rayjob_result.get("metadata", {}).get("uid")
+        ):
+            logger.info(
+                f"Adding owner reference to ConfigMap '{configmap_name}' with RayJob UID: {rayjob_result['metadata']['uid']}"
+            )
+            metadata.owner_references = [
+                client.V1OwnerReference(
+                    api_version="ray.io/v1",
+                    kind="RayJob",
+                    name=self.name,
+                    uid=rayjob_result["metadata"]["uid"],
+                    controller=True,
+                    block_owner_deletion=True,
+                )
+            ]
+        else:
+            logger.warning(
+                f"No valid RayJob result with UID found, ConfigMap '{configmap_name}' will not have owner reference. Result: {rayjob_result}"
+            )
+
         # Convert dict spec to V1ConfigMap
         configmap = client.V1ConfigMap(
-            metadata=client.V1ObjectMeta(**configmap_spec["metadata"]),
+            metadata=metadata,
             data=configmap_spec["data"],
         )
 

src/codeflare_sdk/ray/rayjobs/rayjob.py

@@ -170,3 +170,26 @@ def test_add_script_volumes_existing_mount_early_return():
     # Should still have only one mount, no volume added
     assert len(config.volumes) == 0
     assert len(config.volume_mounts) == 1
+
+
+def test_build_script_configmap_spec_labels():
+    """Test that build_script_configmap_spec creates ConfigMap with correct labels."""
+    config = ManagedClusterConfig()
+
+    job_name = "test-job"
+    namespace = "test-namespace"
+    scripts = {"script.py": "print('hello')", "helper.py": "# helper code"}
+
+    configmap_spec = config.build_script_configmap_spec(job_name, namespace, scripts)
+
+    assert configmap_spec["apiVersion"] == "v1"
+    assert configmap_spec["kind"] == "ConfigMap"
+    assert configmap_spec["metadata"]["name"] == f"{job_name}-scripts"
+    assert configmap_spec["metadata"]["namespace"] == namespace
+
+    labels = configmap_spec["metadata"]["labels"]
+    assert labels["ray.io/job-name"] == job_name
+    assert labels["app.kubernetes.io/managed-by"] == "codeflare-sdk"
+    assert labels["app.kubernetes.io/component"] == "rayjob-scripts"
+
+    assert configmap_spec["data"] == scripts