RHOAIENG-27792: Add stop and resubmit functions to RayJob

Author: kryanbeane

pyproject.toml

@@ -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 = "d1e750d9beac612ad455b951c1a789f971409ab3" }
+python-client = { git = "https://github.com/ray-project/kuberay.git", subdirectory = "clients/python-client", rev = "49419654418865a5838adc7f323f13d82454aa18" }
 
 [[tool.poetry.source]]
 name = "pypi"

src/codeflare_sdk/ray/rayjobs/rayjob.py

@@ -154,29 +154,24 @@ def __init__(
         logger.info(f"Initialized RayJob: {self.name} in namespace: {self.namespace}")
 
     def submit(self) -> str:
-        # Validate required parameters
         if not self.entrypoint:
-            raise ValueError("entrypoint must be provided to submit a RayJob")
+            raise ValueError("Entrypoint must be provided to submit a RayJob")
 
-        # Validate Ray version compatibility for both cluster_config and runtime_env
         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)
-
-        # Handle script files for existing clusters
         elif self._cluster_name:
             scripts = self._extract_script_files_from_entrypoint()
             if scripts:
                 self._handle_script_volumes_for_existing_cluster(scripts)
 
-        # Build the RayJob custom resource
         rayjob_cr = self._build_rayjob_cr()
 
-        # Submit the job - KubeRay operator handles everything else
-        logger.info(f"Submitting RayJob {self.name} to KubeRay operator")
+        logger.info(f"Submitting RayJob {self.name} to Kuberay operator")
         result = self._api.submit_job(k8s_namespace=self.namespace, job=rayjob_cr)
 
         if result:
@@ -189,11 +184,31 @@ def submit(self) -> str:
         else:
             raise RuntimeError(f"Failed to submit RayJob {self.name}")
 
+    def stop(self):
+        """
+        Suspend the Ray job.
+        """
+        stopped = self._api.suspend_job(name=self.name, k8s_namespace=self.namespace)
+        if stopped:
+            logger.info(f"Successfully stopped the RayJob {self.name}")
+            return True
+        else:
+            raise RuntimeError(f"Failed to stop the RayJob {self.name}")
+
+    def resubmit(self):
+        """
+        Resubmit the Ray job.
+        """
+        if self._api.resubmit_job(name=self.name, k8s_namespace=self.namespace):
+            logger.info(f"Successfully resubmitted the RayJob {self.name}")
+            return True
+        else:
+            raise RuntimeError(f"Failed to resubmit the RayJob {self.name}")
+
     def _build_rayjob_cr(self) -> Dict[str, Any]:
         """
         Build the RayJob custom resource specification using native RayJob capabilities.
         """
-        # Basic RayJob custom resource structure
         rayjob_cr = {
             "apiVersion": "ray.io/v1",
             "kind": "RayJob",

src/codeflare_sdk/ray/rayjobs/test_rayjob.py

@@ -390,7 +390,7 @@ def test_submit_validation_no_entrypoint(mocker):
     )
 
     with pytest.raises(
-        ValueError, match="entrypoint must be provided to submit a RayJob"
+        ValueError, match="Entrypoint must be provided to submit a RayJob"
     ):
         rayjob.submit()
 
@@ -1878,3 +1878,123 @@ def test_add_script_volumes_existing_mount_skip():
     # Should still have only one mount and no volume added
     assert len(config.volumes) == 0  # Volume not added due to mount skip
     assert len(config.volume_mounts) == 1
+
+
+def test_rayjob_stop_success(mocker, caplog):
+    """Test successful RayJob stop operation."""
+    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
+
+    mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayClusterApi")
+
+    mock_api_instance.suspend_job.return_value = {
+        "metadata": {"name": "test-rayjob"},
+        "spec": {"suspend": True},
+    }
+
+    rayjob = RayJob(
+        job_name="test-rayjob",
+        cluster_name="test-cluster",
+        namespace="test-namespace",
+        entrypoint="python script.py",
+    )
+
+    with caplog.at_level("INFO"):
+        result = rayjob.stop()
+
+    assert result is True
+
+    mock_api_instance.suspend_job.assert_called_once_with(
+        name="test-rayjob", k8s_namespace="test-namespace"
+    )
+
+    # Verify success message was logged
+    assert "Successfully stopped the RayJob test-rayjob" in caplog.text
+
+
+def test_rayjob_stop_failure(mocker):
+    """Test RayJob stop operation when API call fails."""
+    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
+
+    mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayClusterApi")
+
+    mock_api_instance.suspend_job.return_value = None
+
+    rayjob = RayJob(
+        job_name="test-rayjob",
+        cluster_name="test-cluster",
+        namespace="test-namespace",
+        entrypoint="python script.py",
+    )
+
+    with pytest.raises(RuntimeError, match="Failed to stop the RayJob test-rayjob"):
+        rayjob.stop()
+
+    mock_api_instance.suspend_job.assert_called_once_with(
+        name="test-rayjob", k8s_namespace="test-namespace"
+    )
+
+
+def test_rayjob_resubmit_success(mocker):
+    """Test successful RayJob resubmit operation."""
+    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
+
+    mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayClusterApi")
+
+    mock_api_instance.resubmit_job.return_value = {
+        "metadata": {"name": "test-rayjob"},
+        "spec": {"suspend": False},
+    }
+
+    rayjob = RayJob(
+        job_name="test-rayjob",
+        cluster_name="test-cluster",
+        namespace="test-namespace",
+        entrypoint="python script.py",
+    )
+
+    result = rayjob.resubmit()
+
+    assert result is True
+
+    mock_api_instance.resubmit_job.assert_called_once_with(
+        name="test-rayjob", k8s_namespace="test-namespace"
+    )
+
+
+def test_rayjob_resubmit_failure(mocker):
+    """Test RayJob resubmit operation when API call fails."""
+    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
+
+    mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayClusterApi")
+
+    mock_api_instance.resubmit_job.return_value = None
+
+    rayjob = RayJob(
+        job_name="test-rayjob",
+        cluster_name="test-cluster",
+        namespace="test-namespace",
+        entrypoint="python script.py",
+    )
+
+    with pytest.raises(RuntimeError, match="Failed to resubmit the RayJob test-rayjob"):
+        rayjob.resubmit()
+
+    mock_api_instance.resubmit_job.assert_called_once_with(
+        name="test-rayjob", k8s_namespace="test-namespace"
+    )