microsoft/mscclpp
1
0
Fork 0
mirror of https://github.com/microsoft/mscclpp.git synced 2026-05-13 09:46:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

working on connect

Browse Source
This commit is contained in:
Crutcher Dunnavant
2023-04-04 23:34:27 +00:00
committed by Crutcher Dunnavant
parent b6ea0ca266
commit 7753c38eb1
4 changed files with 115 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
python/src/_py_mscclpp.cpp
View File

@@ -136,6 +136,11 @@ NB_MODULE(_py_mscclpp, m) {
mscclppSetLogHandler(mscclppDefaultLogHandler);
});
nb::enum_<mscclppTransport_t>(m, "TransportType")
.value("P2P", mscclppTransport_t::mscclppTransportP2P)
.value("SHM", mscclppTransport_t::mscclppTransportSHM)
.value("IB", mscclppTransport_t::mscclppTransportIB);
nb::class_<mscclppUniqueId>(m, "MscclppUniqueId")
.def_ro_static("__doc__", &DOC_MscclppUniqueId)
.def_static(
@@ -215,6 +220,33 @@ NB_MODULE(_py_mscclpp, m) {
"Is this comm object closed?")
.def_ro("rank", &_Comm::_rank)
.def_ro("world_size", &_Comm::_world_size)
.def(
"connect",
[](_Comm& self,
int remote_rank,
int tag,
uint64_t local_buff,
uint64_t buff_size,
mscclppTransport_t transport_type) -> void {
checkResult(
mscclppConnect(
self._handle,
remote_rank,
tag,
reinterpret_cast<void*>(local_buff),
buff_size,
transport_type,
0 // ibDev
),
"Connect failed");
},
"remote_rank"_a,
"tag"_a,
"local_buf"_a,
"buff_size"_a,
"transport_type"_a,
nb::call_guard<nb::gil_scoped_release>(),
"Attach a local buffer to a remote connection.")
.def(
"connection_setup",
[](_Comm& comm) {
@@ -222,7 +254,7 @@ NB_MODULE(_py_mscclpp, m) {
return maybe(
mscclppConnectionSetup(comm._handle),
true,
"Failed to settup MSCCLPP connection");
"Failed to setup MSCCLPP connection");
},
nb::call_guard<nb::gil_scoped_release>(),
"Run connection setup for MSCCLPP.")