From 399839744f8ea6ce55f70687b310c74c35dcbad7 Mon Sep 17 00:00:00 2001 From: Ignas Baranauskas Date: Tue, 10 Sep 2024 09:30:06 +0100 Subject: [PATCH] Add entrypoint_memory parameter in Ray Job Submission Client --- src/codeflare_sdk/job/ray_jobs.py | 4 ++++ tests/unit_test.py | 1 + 2 files changed, 5 insertions(+) diff --git a/src/codeflare_sdk/job/ray_jobs.py b/src/codeflare_sdk/job/ray_jobs.py index ff1ebdfe7..9f0022951 100644 --- a/src/codeflare_sdk/job/ray_jobs.py +++ b/src/codeflare_sdk/job/ray_jobs.py @@ -16,6 +16,7 @@ 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 ray.dashboard.modules.job.pydantic_models import JobDetails from typing import Iterator, Optional, Dict, Any, Union, List @@ -62,6 +63,7 @@ def submit_job( submission_id: Optional[str] = None, entrypoint_num_cpus: Optional[Union[int, float]] = None, entrypoint_num_gpus: Optional[Union[int, float]] = None, + entrypoint_memory: Optional[int] = None, entrypoint_resources: Optional[Dict[str, float]] = None, ) -> str: """ @@ -75,6 +77,7 @@ def submit_job( job_id -- DEPRECATED. This has been renamed to submission_id entrypoint_num_cpus -- The quantity of CPU cores to reserve for the execution of the entrypoint command, separately from any tasks or actors launched by it. Defaults to 0. entrypoint_num_gpus -- The quantity of GPUs to reserve for the execution of the entrypoint command, separately from any tasks or actors launched by it. Defaults to 0. + entrypoint_memory –- The quantity of memory to reserve for the execution of the entrypoint command, separately from any tasks or actors launched by it. Defaults to 0. entrypoint_resources -- The quantity of custom resources to reserve for the execution of the entrypoint command, separately from any tasks or actors launched by it. """ return self.rayJobClient.submit_job( @@ -85,6 +88,7 @@ def submit_job( submission_id=submission_id, entrypoint_num_cpus=entrypoint_num_cpus, entrypoint_num_gpus=entrypoint_num_gpus, + entrypoint_memory=entrypoint_memory, entrypoint_resources=entrypoint_resources, ) diff --git a/tests/unit_test.py b/tests/unit_test.py index fd0be3dfd..a5f9cae96 100644 --- a/tests/unit_test.py +++ b/tests/unit_test.py @@ -2724,6 +2724,7 @@ def test_rjc_submit_job(ray_job_client, mocker): submission_id=None, entrypoint_num_cpus=None, entrypoint_num_gpus=None, + entrypoint_memory=None, entrypoint_resources=None, )