[GCS] Optimize `GetAllJobInfo` API for performance #47530

liuxsh9 · 2024-09-06T07:47:33Z

Why are these changes needed?

The GetAllJobInfo API is frequently called by job, log management, and dashboard data collection interfaces. Currently, the API retrieves three fields: 1) job table data, 2) is_running_tasks, and 3) JobInfo. The latter two fields are costly, but they appear to be unused (based on my knowledge and search). To improve performance, we added two optional parameters to control whether to retrieve those two fields (defaulting to not). This optimization reduces the API response time from 1s to approximately 0.05s when there are hundreds of jobs in the cluster. If the latter two fields are indeed unnecessary, consider removing the collection operations altogether.

Related issue number

Checks

I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
I've run scripts/format.sh to lint the changes in this PR.
I've included any doc changes needed for https://docs.ray.io/en/master/.
- I've added any new APIs to the API Reference. For example, if I added a
  method in Tune, I've added it in doc/source/tune/api/ under the
  corresponding .rst file.
I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
Testing Strategy
- Unit tests
- Release tests
- This PR is not tested :(

Signed-off-by: liuxsh9 <[email protected]>

rynewang

Generally LGTM! One issue though: the field is used by places like:

get_total_num_running_jobs_to_report

@routes.get("/api/jobs/")

/api/jobs/{job_or_submission_id}

get_job_info

And those APIs may be consumed by external code. To keep backwards compatibility can you add args to the web APIs with default fetch=True? Then in call sites like

ray/python/ray/dashboard/modules/log/log_manager.py

Line 150 in c083ea5

job_infos = await self.client.get_job_info(timeout=DEFAULT_RPC_TIMEOUT)

we can set fetch=False.

Especially, we have a @PublicAPI , @jjyao can you help evaluate if this PublicAPI contained such fields, and shall we change its semantics for this better perf?

rynewang · 2024-09-06T16:08:01Z

python/ray/_private/gcs_aio_client.py

@@ -86,7 +87,8 @@ def __init__(self, inner, loop, executor):

    def _function_to_async(self, func):
        async def wrapper(*args, **kwargs):
-            return await self.loop.run_in_executor(self.executor, func, *args, **kwargs)
+            partial_func = partial(func, *args, **kwargs)


why is this change?

+1

Let's avoid bundling up unrelated changes

We discovered that run_in_executor doesn't support **kwargs ref: asyncio.loop.run_in_executor. The implicit params passing was exposing this issue. For now, we've changed to explicit params passing. So have reverted this change.
Anyway, should we adjust it to accommodate run_in_executor's limitations?

rynewang · 2024-09-06T17:20:58Z

python/ray/_raylet.pyx

@@ -2907,18 +2907,23 @@ cdef class OldGcsClient:
        return result

    @_auto_reconnect
-    def get_all_job_info(self, timeout=None) -> Dict[JobID, JobTableData]:
+    def get_all_job_info(self, timeout=None, **kwargs) -> Dict[JobID, JobTableData]:


wonder if we can make the fields just defined here, instead of **kwargs:

def get_all_job_info(self, query_job_info_field, c_query_is_running_tasks_field, timeout=None) -> Dict[JobID, JobTableData]:

+1

Let's err on the side of properly codifying params

Good suggestion, have updated to use explicit field definitions instead.

rynewang · 2024-09-06T17:21:25Z

python/ray/includes/gcs_client.pxi

@@ -433,20 +433,28 @@ cdef class NewGcsClient:
    #############################################################

    def get_all_job_info(
-        self, timeout: Optional[float] = None
+        self, timeout: Optional[float] = None, **kwargs


ditto, explicit args

rynewang · 2024-09-06T17:22:21Z

src/ray/gcs/gcs_server/gcs_job_manager.cc

-        reply->mutable_job_info_list(i)->set_is_running_tasks(false);
-        core_worker_clients_.Disconnect(worker_id);
+      if (request.has_query_is_running_tasks_field() &&
+          request.query_is_running_tasks_field()) {


you can omit the has_ check. if it's not set, query_is_running_tasks_field() defaults to false.

Good suggestion!

Signed-off-by: liuxsh9 <[email protected]>

liuxsh9 · 2024-09-07T09:22:38Z

Generally LGTM! One issue though: the field is used by places like:

get_total_num_running_jobs_to_report

@routes.get("/api/jobs/")

/api/jobs/{job_or_submission_id}

get_job_info

And those APIs may be consumed by external code. To keep backwards compatibility can you add args to the web APIs with default fetch=True? Then in call sites like

ray/python/ray/dashboard/modules/log/log_manager.py

Line 150 in c083ea5

job_infos = await self.client.get_job_info(timeout=DEFAULT_RPC_TIMEOUT)

we can set fetch=False.

Sorry, I am a little confused. Do you mean we're still fetching all 3 fields by default, and then setting query_x_field=False in individual APIs that definitely don't need them?

Signed-off-by: liuxsh9 <[email protected]>

jjyao · 2024-09-07T15:47:15Z

Sorry, I am a little confused. Do you mean we're still fetching all 3 fields by default, and then setting query_x_field=False in individual APIs that definitely don't need them?

Yes.

…d by default. Signed-off-by: liuxsh9 <[email protected]>

liuxsh9 · 2024-09-09T11:26:54Z

Updated the default to retrieving all fields, like before. To accommodate proto3's boolean fields defaulting to false, we've changed the params to skip_x_field. So now we can set skip_x_field=True when method doesn't need the field. Additionally, we've optimized functions like get_driver_jobs to use GetAllJobInfo with selective field retrieval.

Signed-off-by: liuxsh9 <[email protected]>

jjyao

lg, thanks!

python/ray/_private/gcs_aio_client.py

jjyao · 2024-09-10T05:23:05Z

src/ray/gcs/gcs_client/accessor.cc

+                             /*skip_job_info_field=*/false,
+                             /*skip_is_running_tasks_field=*/false,


Can those be True as well If they are not used?

jjyao · 2024-09-10T05:26:32Z

src/ray/gcs/gcs_client/global_state_accessor.cc

+        /*skip_job_info_field=*/false,
+        /*skip_is_running_tasks_field=*/false,


Can those be True? Are those two fields used?

GlobalStateAccessor::GetAllJobInfo() is invoked by multiple functions, including

ray/python/ray/_private/state.py

Line 170 in f40313b

def job_table(self):

which explicitly sets two params to True as it doesn't utilize those fields.

And

ray/java/runtime/src/main/java/io/ray/runtime/gcs/GlobalStateAccessor.java

Line 50 in f40313b

public List<byte[]> getAllJobInfo() {

may be called by external methods, so we're leaving it unchanged and defaulting to False.

To accommodate these varying callers, have added explicit skip params to GetAllJobInfo().

jjyao · 2024-09-10T05:38:04Z

Failed tests are probably related.

…field. Signed-off-by: liuxsh9 <[email protected]>

Signed-off-by: liuxsh9 <[email protected]>

jjyao

Looks great!

The remaining comments are about changing them to keyword arguments.

python/ray/_private/gcs_aio_client.py

jjyao · 2024-09-11T05:23:39Z

python/ray/_private/state.py

@@ -181,7 +181,7 @@ def job_table(self):
        """
        self._check_connected()

-        job_table = self.global_state_accessor.get_job_table()
+        job_table = self.global_state_accessor.get_job_table(True, True)


Using keyword arguments otherwise it's hard to know what these two True are

jjyao · 2024-09-11T05:30:58Z

python/ray/_raylet.pyx

-    def get_all_job_info(self, job_or_submission_id: str = None,
-                         timeout=None) -> Dict[JobID, JobTableData]:
+    def get_all_job_info(
+        self, job_or_submission_id: str = None, skip_submission_job_info_field=False,


Suggested change

self, job_or_submission_id: str = None, skip_submission_job_info_field=False,

self, *, job_or_submission_id: str = None, skip_submission_job_info_field=False,

jjyao · 2024-09-11T05:32:59Z

python/ray/includes/global_state_accessor.pxi

@@ -48,10 +48,16 @@ cdef class GlobalStateAccessor:
        with nogil:
            self.inner.get().Disconnect()

-    def get_job_table(self):
+    def get_job_table(
+        self, skip_submission_job_info_field=False, skip_is_running_tasks_field=False


Suggested change

self, skip_submission_job_info_field=False, skip_is_running_tasks_field=False

self, *, skip_submission_job_info_field=False, skip_is_running_tasks_field=False

jjyao · 2024-09-11T05:36:53Z

src/ray/gcs/gcs_server/gcs_job_manager.cc

+          reply->mutable_job_info_list(i)->set_is_running_tasks(false);
+          core_worker_clients_.Disconnect(worker_id);
+          (*num_processed_jobs)++;
+          ;


This can be deleted

jjyao · 2024-09-11T05:40:20Z

python/ray/includes/gcs_client.pxi

@@ -434,11 +434,15 @@ cdef class NewGcsClient:

    def get_all_job_info(
        self, job_or_submission_id: Optional[str] = None,


Suggested change

self, job_or_submission_id: Optional[str] = None,

self, *, job_or_submission_id: Optional[str] = None,

jjyao · 2024-09-11T05:40:31Z

python/ray/includes/gcs_client.pxi

@@ -447,17 +451,22 @@ cdef class NewGcsClient:
                make_optional[c_string](c_job_or_submission_id)
        with nogil:
            status = self.inner.get().Jobs().GetAll(
-                c_optional_job_or_submission_id, reply, timeout_ms)
+                c_optional_job_or_submission_id, c_skip_submission_job_info_field,
+                c_skip_is_running_tasks_field, reply, timeout_ms)
        return raise_or_return((convert_get_all_job_info(status, move(reply))))

    def async_get_all_job_info(
        self, job_or_submission_id: Optional[str] = None,


Suggested change

self, job_or_submission_id: Optional[str] = None,

self, *, job_or_submission_id: Optional[str] = None,

Signed-off-by: liuxsh9 <[email protected]>

…arameters. Signed-off-by: liuxsh9 <[email protected]>

liuxsh9 · 2024-09-11T11:36:36Z

Updates applied, kindly review at your convenience. @jjyao

Signed-off-by: liuxsh9 <[email protected]> Signed-off-by: ujjawal-khare <[email protected]>

Add two boolean parameters to control whether to retrieve two fields

7c40d7b

Signed-off-by: liuxsh9 <[email protected]>

liuxsh9 requested review from a team, ericl, pcmoritz and raulchen as code owners September 6, 2024 07:47

liuxsh9 mentioned this pull request Sep 6, 2024

[core][dashboard] push down job_or_submission_id to GCS. #47492

Merged

Fix test case errors

f762acb

Signed-off-by: liuxsh9 <[email protected]>

jjyao assigned rynewang Sep 6, 2024

rynewang reviewed Sep 6, 2024

View reviewed changes

Modify parameter transmission method

3def496

Signed-off-by: liuxsh9 <[email protected]>

Modify the order of function parameters.

c327e16

Signed-off-by: liuxsh9 <[email protected]>

Change query_xx_field to skip_xx_field, so that two fields are querie…

319450c

…d by default. Signed-off-by: liuxsh9 <[email protected]>

liuxsh9 requested review from richardliaw, rkooo567 and jjyao as code owners September 9, 2024 06:31

jjyao added the go add ONLY when ready to merge, run all tests label Sep 9, 2024

fix test errors

04c3065

Signed-off-by: liuxsh9 <[email protected]>

jjyao reviewed Sep 10, 2024

View reviewed changes

liuxsh9 added 4 commits September 10, 2024 16:58

Rename the skip_job_info_field parameter to skip_submission_job_info_…

4929e2c

…field. Signed-off-by: liuxsh9 <[email protected]>

Merge branch 'master' into optimize-get_all_job_info

900e99b

fix errors in test cases

62cfbd3

Signed-off-by: liuxsh9 <[email protected]>

Merge branch 'master' into optimize-get_all_job_info

631440f

Signed-off-by: liuxsh9 <[email protected]>

jjyao reviewed Sep 11, 2024

View reviewed changes

liuxsh9 added 2 commits September 11, 2024 14:33

force to use the keyword parameters.

700cf0d

Signed-off-by: liuxsh9 <[email protected]>

Fixed the error of passing OldGcsAioClient.get_all_job_info keyword p…

7b8bd87

…arameters. Signed-off-by: liuxsh9 <[email protected]>

jjyao approved these changes Sep 11, 2024

View reviewed changes

jjyao merged commit 432dbce into ray-project:master Sep 11, 2024
5 checks passed

ujjawal-khare pushed a commit to ujjawal-khare-27/ray that referenced this pull request Oct 15, 2024

[GCS] Optimize GetAllJobInfo API for performance (ray-project#47530)

98ce570

Signed-off-by: liuxsh9 <[email protected]> Signed-off-by: ujjawal-khare <[email protected]>

ujjawal-khare pushed a commit to ujjawal-khare-27/ray that referenced this pull request Oct 15, 2024

[GCS] Optimize GetAllJobInfo API for performance (ray-project#47530)

0e567a1

Signed-off-by: liuxsh9 <[email protected]> Signed-off-by: ujjawal-khare <[email protected]>

ujjawal-khare pushed a commit to ujjawal-khare-27/ray that referenced this pull request Oct 15, 2024

[GCS] Optimize GetAllJobInfo API for performance (ray-project#47530)

419af7b

Signed-off-by: liuxsh9 <[email protected]> Signed-off-by: ujjawal-khare <[email protected]>

ujjawal-khare pushed a commit to ujjawal-khare-27/ray that referenced this pull request Oct 15, 2024

[GCS] Optimize GetAllJobInfo API for performance (ray-project#47530)

95e078e

Signed-off-by: liuxsh9 <[email protected]> Signed-off-by: ujjawal-khare <[email protected]>

ujjawal-khare pushed a commit to ujjawal-khare-27/ray that referenced this pull request Oct 15, 2024

[GCS] Optimize GetAllJobInfo API for performance (ray-project#47530)

8c73f30

Signed-off-by: liuxsh9 <[email protected]> Signed-off-by: ujjawal-khare <[email protected]>

ujjawal-khare pushed a commit to ujjawal-khare-27/ray that referenced this pull request Oct 15, 2024

[GCS] Optimize GetAllJobInfo API for performance (ray-project#47530)

48784b1

Signed-off-by: liuxsh9 <[email protected]> Signed-off-by: ujjawal-khare <[email protected]>

ujjawal-khare pushed a commit to ujjawal-khare-27/ray that referenced this pull request Oct 15, 2024

[GCS] Optimize GetAllJobInfo API for performance (ray-project#47530)

1528145

Signed-off-by: liuxsh9 <[email protected]> Signed-off-by: ujjawal-khare <[email protected]>

ujjawal-khare pushed a commit to ujjawal-khare-27/ray that referenced this pull request Oct 15, 2024

[GCS] Optimize GetAllJobInfo API for performance (ray-project#47530)

b1c7caa

Signed-off-by: liuxsh9 <[email protected]> Signed-off-by: ujjawal-khare <[email protected]>

ujjawal-khare pushed a commit to ujjawal-khare-27/ray that referenced this pull request Oct 15, 2024

[GCS] Optimize GetAllJobInfo API for performance (ray-project#47530)

649148c

Signed-off-by: liuxsh9 <[email protected]> Signed-off-by: ujjawal-khare <[email protected]>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[GCS] Optimize `GetAllJobInfo` API for performance #47530

[GCS] Optimize `GetAllJobInfo` API for performance #47530

liuxsh9 commented Sep 6, 2024 •

edited

Loading

rynewang left a comment

rynewang Sep 6, 2024

alexeykudinkin Sep 6, 2024

liuxsh9 Sep 7, 2024 •

edited

Loading

rynewang Sep 6, 2024

alexeykudinkin Sep 6, 2024

liuxsh9 Sep 7, 2024

rynewang Sep 6, 2024

liuxsh9 Sep 7, 2024

rynewang Sep 6, 2024

liuxsh9 Sep 7, 2024

liuxsh9 commented Sep 7, 2024

jjyao commented Sep 7, 2024

liuxsh9 commented Sep 9, 2024

jjyao left a comment

jjyao Sep 10, 2024

liuxsh9 Sep 10, 2024

jjyao Sep 10, 2024

liuxsh9 Sep 10, 2024

jjyao commented Sep 10, 2024

jjyao left a comment

jjyao Sep 11, 2024

jjyao Sep 11, 2024

jjyao Sep 11, 2024

jjyao Sep 11, 2024

jjyao Sep 11, 2024

jjyao Sep 11, 2024

liuxsh9 commented Sep 11, 2024

		/skip_job_info_field=/false,
		/skip_is_running_tasks_field=/false,

	self, job_or_submission_id: str = None, skip_submission_job_info_field=False,
	self, *, job_or_submission_id: str = None, skip_submission_job_info_field=False,

	self, skip_submission_job_info_field=False, skip_is_running_tasks_field=False
	self, *, skip_submission_job_info_field=False, skip_is_running_tasks_field=False

		@@ -434,11 +434,15 @@ cdef class NewGcsClient:

		def get_all_job_info(
		self, job_or_submission_id: Optional[str] = None,

	self, job_or_submission_id: Optional[str] = None,
	self, *, job_or_submission_id: Optional[str] = None,

[GCS] Optimize GetAllJobInfo API for performance #47530

[GCS] Optimize GetAllJobInfo API for performance #47530

Conversation

liuxsh9 commented Sep 6, 2024 • edited Loading

Why are these changes needed?

Related issue number

Checks

rynewang left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

liuxsh9 Sep 7, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

liuxsh9 commented Sep 7, 2024

jjyao commented Sep 7, 2024

liuxsh9 commented Sep 9, 2024

jjyao left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jjyao commented Sep 10, 2024

jjyao left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

liuxsh9 commented Sep 11, 2024

[GCS] Optimize `GetAllJobInfo` API for performance #47530

[GCS] Optimize `GetAllJobInfo` API for performance #47530

liuxsh9 commented Sep 6, 2024 •

edited

Loading

liuxsh9 Sep 7, 2024 •

edited

Loading