Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[core] change GcsClient.get_all_node_info return type to proto type. #46057

Merged
merged 4 commits into from
Jun 17, 2024

Conversation

rynewang
Copy link
Contributor

@rynewang rynewang commented Jun 14, 2024

This PR removes a custom proto-to-dict conversion by directly returning protos. This avoids some arbitrary data type semantics and prepares for the removal of PythonGcsClient. Also returns dict keys with JobID and NodeID.

Before:

  • def get_all_job_info(self, timeout=None) -> Dict[bytes, JobTableData]
  • def get_all_node_info(self, timeout=None) -> Dict[bytes, dict] # dict has "node_name", "state", "labels", "resources"

After:

  • def get_all_job_info(self, timeout=None) -> Dict[JobID, JobTableData]
  • def get_all_node_info(self, timeout=None) -> Dict[NodeID, GcsNodeInfo]

Signed-off-by: Ruiyang Wang <rywang014@gmail.com>
Signed-off-by: Ruiyang Wang <rywang014@gmail.com>
@rynewang rynewang added the go add ONLY when ready to merge, run all tests label Jun 15, 2024
Signed-off-by: Ruiyang Wang <rywang014@gmail.com>
python/ray/_raylet.pyx Outdated Show resolved Hide resolved
python/ray/_raylet.pyx Outdated Show resolved Hide resolved
Signed-off-by: Ruiyang Wang <rywang014@gmail.com>
Comment on lines +2883 to 2899
def get_all_node_info(self, timeout=None) -> Dict[NodeID, GcsNodeInfo]:
cdef:
int64_t timeout_ms = round(1000 * timeout) if timeout else -1
CGcsNodeInfo node_info
c_vector[CGcsNodeInfo] node_infos
CGcsNodeInfo c_node_info
c_vector[CGcsNodeInfo] c_node_infos
c_vector[c_string] serialized_node_infos
with nogil:
check_status(self.inner.get().GetAllNodeInfo(timeout_ms, node_infos))
check_status(self.inner.get().GetAllNodeInfo(timeout_ms, c_node_infos))
for c_node_info in c_node_infos:
serialized_node_infos.push_back(c_node_info.SerializeAsString())

result = {}
for node_info in node_infos:
c_resources = PythonGetResourcesTotal(node_info)
result[node_info.node_id()] = {
"node_name": node_info.node_name(),
"state": node_info.state(),
"labels": PythonGetNodeLabels(node_info),
"resources": {key.decode(): value for key, value in c_resources}
}
for serialized in serialized_node_infos:
node_info = GcsNodeInfo()
node_info.ParseFromString(serialized)
result[NodeID.from_binary(node_info.node_id)] = node_info
return result
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We need to be careful about the performance since it now has ser + deser overhead.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I feel it should be ok because the current code just "manually" re-serializes by taking some C++ fields into python dicts and strings. Should be on par.

@jjyao jjyao merged commit b0d2cb2 into ray-project:master Jun 17, 2024
6 checks passed
chungen04 pushed a commit to chungen04/ray that referenced this pull request Jun 18, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
go add ONLY when ready to merge, run all tests
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

3 participants