Added RayJobSubmission Client wrapper for RCS not created by CodeFlare

Author: Bobbins228

src/codeflare_sdk/job/ray_jobs.py

@@ -0,0 +1,117 @@
+# Copyright 2022 IBM, Red Hat
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""
+The ray_jobs sub-module contains methods needed to submit jobs and connect to Ray Clusters that were not created by CodeFlare.
+The SDK acts as a wrapper for the Ray Job Submission Client.
+"""
+from ray.job_submission import JobSubmissionClient
+from typing import Iterator, Optional, Dict, Any, Union
+
+
+class RayJobClient:
+    """
+    An object for that acts as the Ray Job Submission Client.
+    """
+
+    def __init__(
+        self,
+        address: Optional[str] = None,
+        create_cluster_if_needed: bool = False,
+        cookies: Optional[Dict[str, Any]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        headers: Optional[Dict[str, Any]] = None,
+        verify: Optional[Union[str, bool]] = True,
+    ):
+        self.rayJobClient = JobSubmissionClient(
+            address=address,
+            create_cluster_if_needed=create_cluster_if_needed,
+            cookies=cookies,
+            metadata=metadata,
+            headers=headers,
+            verify=verify,
+        )
+
+    def submit_job(
+        self,
+        entrypoint: str,
+        job_id: Optional[str] = None,
+        runtime_env: Optional[Dict[str, Any]] = None,
+        metadata: Optional[Dict[str, str]] = None,
+        submission_id: Optional[str] = None,
+        entrypoint_num_cpus: Optional[Union[int, float]] = None,
+        entrypoint_num_gpus: Optional[Union[int, float]] = None,
+        entrypoint_resources: Optional[Dict[str, float]] = None,
+    ) -> str:
+        """
+        Method for submitting jobs to a Ray Cluster and returning the job id with entrypoint being a mandatory field.
+        """
+        return self.rayJobClient.submit_job(
+            entrypoint=entrypoint,
+            job_id=job_id,
+            runtime_env=runtime_env,
+            metadata=metadata,
+            submission_id=submission_id,
+            entrypoint_num_cpus=entrypoint_num_cpus,
+            entrypoint_num_gpus=entrypoint_num_gpus,
+            entrypoint_resources=entrypoint_resources,
+        )
+
+    def delete_job(self, job_id: str) -> bool:
+        """
+        Method for deleting jobs with the job id being a mandatory field.
+        """
+        return self.rayJobClient.delete_job(job_id=job_id)
+
+    def get_address(self) -> str:
+        """
+        Method for getting the address from the RayJobClient
+        """
+        return self.rayJobClient.get_address()
+
+    def get_job_info(self, job_id: str):
+        """
+        Method for getting the job info with the job id being a mandatory field.
+        """
+        return self.rayJobClient.get_job_info(job_id=job_id)
+
+    def get_job_logs(self, job_id: str) -> str:
+        """
+        Method for getting the job info with the job id being a mandatory field.
+        """
+        return self.rayJobClient.get_job_logs(job_id=job_id)
+
+    def get_job_status(self, job_id: str) -> str:
+        """
+        Method for getting the job's status with the job id being a mandatory field.
+        """
+        return self.rayJobClient.get_job_status(job_id=job_id)
+
+    def list_jobs(self):
+        """
+        Method for getting a list of current jobs in the Ray Cluster.
+        """
+        return self.rayJobClient.list_jobs()
+
+    def stop_job(self, job_id: str) -> bool:
+        """
+        Method for stopping a job with the job id being a mandatory field.
+        """
+        return self.rayJobClient.stop_job(job_id=job_id)
+
+    def tail_job_logs(self, job_id: str) -> Iterator[str]:
+        """
+        Method for getting an iterator that follows the logs of a job with the job id being a mandatory field.
+        """
+        return self.rayJobClient.tail_job_logs(job_id=job_id)

tests/unit_test.py

@@ -103,6 +103,7 @@
 from unittest.mock import MagicMock
 from pytest_mock import MockerFixture
 from ray.job_submission import JobSubmissionClient
+from codeflare_sdk.job.ray_jobs import RayJobClient
 
 # For mocking openshift client results
 fake_res = openshift.Result("fake")
@@ -2846,6 +2847,85 @@ def test_gen_app_wrapper_with_oauth(mocker: MockerFixture):
     )
 
 
+"""
+Ray Jobs tests
+"""
+# rjc == RayJobClient
+@pytest.fixture
+def ray_job_client(mocker):
+    # Creating a fixture to instantiate RayJobClient with a mocked JobSubmissionClient
+    mocker.patch.object(JobSubmissionClient, "__init__", return_value=None)
+    return RayJobClient(
+        "https://ray-dashboard-unit-test-cluster-ns.apps.cluster.awsroute.org"
+    )
+
+
+def test_rjc_submit_job(ray_job_client, mocker):
+    mocked_submit_job = mocker.patch.object(
+        JobSubmissionClient, "submit_job", return_value="mocked_submission_id"
+    )
+    submission_id = ray_job_client.submit_job(entrypoint={"pip": ["numpy"]})
+
+    mocked_submit_job.assert_called_once_with(
+        entrypoint={"pip": ["numpy"]},
+        job_id=None,
+        runtime_env=None,
+        metadata=None,
+        submission_id=None,
+        entrypoint_num_cpus=None,
+        entrypoint_num_gpus=None,
+        entrypoint_resources=None,
+    )
+
+    assert submission_id == "mocked_submission_id"
+
+
+def test_rjc_delete_job(ray_job_client, mocker):
+    mocked_delete_job = mocker.patch.object(
+        JobSubmissionClient, "delete_job", return_value=True
+    )
+    result = ray_job_client.delete_job(job_id="mocked_job_id")
+
+    mocked_delete_job.assert_called_once_with(job_id="mocked_job_id")
+    assert result is True
+
+
+def test_rjc_address(ray_job_client, mocker):
+    mocked_rjc_address = mocker.patch.object(
+        JobSubmissionClient,
+        "get_address",
+        return_value="https://ray-dashboard-unit-test-cluster-ns.apps.cluster.awsroute.org",
+    )
+    address = ray_job_client.get_address()
+
+    mocked_rjc_address.assert_called_once()
+    assert (
+        address
+        == "https://ray-dashboard-unit-test-cluster-ns.apps.cluster.awsroute.org"
+    )
+
+
+def test_rjc_get_job_logs(ray_job_client, mocker):
+    mocked_rjc_get_job_logs = mocker.patch.object(
+        JobSubmissionClient, "get_job_logs", return_value="Logs"
+    )
+    logs = ray_job_client.get_job_logs(job_id="mocked_job_id")
+
+    mocked_rjc_get_job_logs.assert_called_once_with(job_id="mocked_job_id")
+    assert logs == "Logs"
+
+
+def test_rjc_get_job_info(ray_job_client, mocker):
+    job_details_example = "JobDetails(type=<JobType.SUBMISSION: 'SUBMISSION'>, job_id=None, submission_id='mocked_submission_id', driver_info=None, status=<JobStatus.PENDING: 'PENDING'>, entrypoint='python test.py', message='Job has not started yet. It may be waiting for the runtime environment to be set up.', error_type=None, start_time=1701271760641, end_time=None, metadata={}, runtime_env={'working_dir': 'gcs://_ray_pkg_67de6f0e60d43b19.zip', 'pip': {'packages': ['numpy'], 'pip_check': False}, '_ray_commit': 'b4bba4717f5ba04ee25580fe8f88eed63ef0c5dc'}, driver_agent_http_address=None, driver_node_id=None)"
+    mocked_rjc_get_job_info = mocker.patch.object(
+        JobSubmissionClient, "get_job_info", return_value=job_details_example
+    )
+    job_details = ray_job_client.get_job_info(job_id="mocked_job_id")
+
+    mocked_rjc_get_job_info.assert_called_once_with(job_id="mocked_job_id")
+    assert job_details == job_details_example
+
+
 # Make sure to always keep this function last
 def test_cleanup():
     os.remove(f"{aw_dir}unit-test-cluster.yaml")