Skip to content

Conversation

FollyCoolly
Copy link
Collaborator

No description provided.

Copy link
Contributor

Summary of Changes

Hello @FollyCoolly, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the Python-C++ interoperability for the yacl link library by enabling Python developers to define and utilize custom communication channels. By implementing the IChannel interface in Python, these channels can now be seamlessly integrated with the C++ yacl::link::Context, offering greater flexibility in network communication. The changes include robust pybind11 bindings, custom type casters for efficient zero-copy data transfer, and a new API for context creation, all validated by a new integration test.

Highlights

  • Python IChannel Binding: Introduced pybind11 bindings for yacl::link::transport::IChannel, allowing Python classes to implement this C++ interface and be used by the C++ backend.
  • Custom Channel Integration: Added a create_with_channels function to the Python link module, enabling the creation of yacl::link::Context instances using Python-implemented channels for flexible communication.
  • Zero-Copy Buffer Conversion: Implemented custom pybind11 type casters for yacl::Buffer and yacl::ByteContainerView to facilitate efficient, zero-copy data transfer between Python bytes/numpy.array and C++ buffers.
  • Dependency Updates: Updated various Bazel dependencies, including psi, yacl (with a git override), heu, platforms, and added liboqs and perfetto related entries, ensuring compatibility and leveraging newer versions.
  • New Integration Test: Added a Python test script (test_channel_integration.py) to demonstrate and verify the functionality of implementing IChannel in Python and using it with the C++ link.Context.
Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

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

Code Review

This pull request introduces pybind11 bindings for yacl::link::transport::IChannel, allowing Python classes to implement this interface. This is a significant feature that enables more flexible and custom communication channel implementations from Python. The changes include the necessary C++ binding code, a trampoline class, type casters for yacl::Buffer, and a new Python test to verify the integration.

The C++ implementation is well-done, especially the zero-copy buffer-to-numpy-array conversion and the safe handling of brpc flags. The Python test is a good starting point.

My review includes a few suggestions for cleanup and improving the robustness of the test mock. Overall, this is a solid contribution.


import spu.libspu.link as link
import spu.libspu as spu
from typing import Dict, Optional
Copy link
Contributor

Choose a reason for hiding this comment

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

medium

The Optional type is imported but not used in this file. It's good practice to remove unused imports to keep the code clean.

Suggested change
from typing import Dict, Optional
from typing import Dict

Comment on lines +10 to +11
print(dir(link))
print(spu.__file__)
Copy link
Contributor

Choose a reason for hiding this comment

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

medium

These print statements appear to be for debugging purposes. They should be removed before merging to keep the test output clean.

Comment on lines +48 to +52
def Recv(self, key: str) -> bytes:
"""Receive data"""
final_key = f"{key}_{self.remote_rank}_{self.local_rank}"
print(f"[{self.name}] Recv: key={final_key}")
return self.storage.get(final_key, b"")
Copy link
Contributor

Choose a reason for hiding this comment

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

medium

The Recv method in SimpleChannel is a mock that doesn't block and doesn't respect the recv_timeout. It immediately returns from the storage dictionary. While this works for the current synchronous tests, it's not a realistic simulation of a channel's behavior and could lead to race conditions or flaky tests if used in asynchronous scenarios. Consider implementing a blocking wait with a timeout to make the mock more robust.

};

} // namespace detail
} // namespace pybind11 No newline at end of file
Copy link
Contributor

Choose a reason for hiding this comment

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

medium

The file is missing a newline character at the end. It's a common convention and good practice to end files with a newline.

}  // namespace pybind11

PYBIND11_OVERRIDE_PURE(void, yacl::link::transport::IChannel,
SetChunkParallelSendSize, size);
}
}; No newline at end of file
Copy link
Contributor

Choose a reason for hiding this comment

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

medium

The file is missing a newline character at the end. It's a common convention and good practice to end files with a newline.

};

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants