Enhancement: Allow overriding Channel hehavior in Python

MODULE.bazel

@@ -38,6 +38,17 @@ local_path_override(
 
 bazel_dep(name = "psi", version = "0.6.0.dev250507")
 bazel_dep(name = "yacl", version = "0.4.5b10-nightly-20250110")
+git_override(
+    module_name = "yacl",
+    remote = "https://github.com/secretflow/yacl.git",
+    commit = "a8af9f85816139f62712e29f0e5421da6f7f1e32",
+)
+git_override(
+    module_name = "psi",
+    remote = "https://github.com/secretflow/psi.git",
+    commit = "b992382f2889003d39b52e525904b932dca09d37",
+)
+
 bazel_dep(name = "grpc", version = "1.66.0.bcr.4")
 
 # pin [email protected]

spu/BUILD.bazel

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-load("@pybind11_bazel//:build_defs.bzl", "pybind_extension")
+load("@pybind11_bazel//:build_defs.bzl", "pybind_extension", "pybind_library")
 
 package(default_visibility = ["//visibility:public"])
 
@@ -37,6 +37,7 @@ pybind_extension(
     }),
     deps = [
         ":exported_symbols.lds",
+        ":pychannel",
         ":version_script.lds",
         "@spulib//libspu:version",
         "@spulib//libspu/compiler:compile",
@@ -68,3 +69,12 @@ pybind_extension(
         "@yacl//yacl/link",
     ],
 )
+
+pybind_library(
+    name = "pychannel",
+    hdrs = ["pychannel.h"],
+    visibility = ["//visibility:private"],
+    deps = [
+        "@yacl//yacl/link/transport:channel",
+    ],
+)

spu/libspu.cc

@@ -20,6 +20,7 @@
 #include "pybind11/numpy.h"
 #include "pybind11/pybind11.h"
 #include "pybind11/stl.h"
+#include "spu/pychannel.h"
 #include "yacl/link/algorithm/allgather.h"
 #include "yacl/link/algorithm/barrier.h"
 #include "yacl/link/algorithm/broadcast.h"
@@ -76,6 +77,7 @@ void BindLink(py::module& m) {
   using yacl::link::RetryOptions;
   using yacl::link::SSLOptions;
   using yacl::link::VerifyOptions;
+  using yacl::link::transport::IChannel;
 
   // TODO(jint) expose this tag to python?
   constexpr char PY_CALL_TAG[] = "PY_CALL";
@@ -84,6 +86,23 @@ void BindLink(py::module& m) {
               SPU Link Library
                   )pbdoc";
 
+  py::class_<IChannel, std::shared_ptr<IChannel>, PyChannel>(m, "IChannel")
+      .def(py::init<>())
+      .def("send_async",
+           static_cast<void (IChannel::*)(const std::string&, yacl::Buffer)>(
+               &IChannel::SendAsync))
+      .def("send_async_throttled",
+           static_cast<void (IChannel::*)(const std::string&, yacl::Buffer)>(
+               &IChannel::SendAsyncThrottled))
+      .def("send", &IChannel::Send)
+      .def("recv", &IChannel::Recv)
+      .def("test_send", &IChannel::TestSend)
+      .def("test_recv", &IChannel::TestRecv)
+      .def("set_throttle_window_size", &IChannel::SetThrottleWindowSize)
+      .def("set_chunk_parallel_send_size", &IChannel::SetChunkParallelSendSize)
+      .def("wait_link_task_finish", &IChannel::WaitLinkTaskFinish)
+      .def("abort", &IChannel::Abort);
+
   py::class_<CertInfo>(m, "CertInfo", "The config info used for certificate")
       .def_readwrite("certificate_path", &CertInfo::certificate_path,
                      "certificate file path")
@@ -302,6 +321,17 @@ void BindLink(py::module& m) {
           ctx->ConnectToMesh();
           return ctx;
         });
+  m.def(
+      "create_with_channels",
+      [](const ContextDesc& desc, size_t self_rank,
+         std::vector<std::shared_ptr<IChannel>> channels) {
+        py::gil_scoped_release release;
+        auto ctx = std::make_shared<yacl::link::Context>(
+            desc, self_rank, std::move(channels), nullptr, false);
+        ctx->ConnectToMesh();
+        return ctx;
+      },
+      py::arg("desc"), py::arg("self_rank"), py::arg("channels"));
 }
 
 struct PyBindShare {