Expose RDMA support through Python APIs #462
Open
+403
−20
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
facebook-github-bot
added
the
CLA Signed
|
This pull request was exported from Phabricator. Differential Revision: D76937776 |
allenwang28
force-pushed
the
export-D76937776
branch
from
be4affa to
63ce0c9
Compare
allenwang28
added a commit
to allenwang28/monarch-1
that referenced
this pull request
Jul 9, 2025
Summary: This change brings in RDMA support from Rust to the Python APIs, overwriting the existing stand-in RDMABuffer APIs. There are quite a few changes, so here's a high level overview: ## `monarch_hyperactor` (rust) Within `PyProcMesh`, adds an optional RDMAMangerActor attribute. This mimics the prior functionality - before, a `proc_mesh` (in Python) would unconditionally spin up an `RdmaManagerActor`. Now, a `proc_mesh` will spin up an `RdmaManagerActor` if and only if `tensor_engine` is enabled and `ibverbs` APIs are supported. A few other pieces must be accessible within the RDMA extension library - specifically `Mailbox` / the send caps capabiliites. ## RDMA extension Adds an `extension` folder within `monarch_rdma` which contains the actual bindings. Consistent with other bindings, provides a blocking and non-blocking version. ## RDMA components Corrects behavior of a few edge cases, i.e. `ibverbs_supported` based on the number of devices, supporting a "loopback" case, wherein two actors share an RDMA buffer but are spawned on the same proc. ## Python Removes the stand-in RDMAManagerActor In rdma.py, adds in basic support for RDMABuffer CPU. While CUDA<>CPU|CUDA is implemented in `monarch_rdma`, this was running into issues with MR registration. Logging shows that the code path correctly interprets the address as CUDA, but `cuMemGetHandleForAddressRange` returns the handle/fd as `-1` and a null pointer MR. RDMABuffer will only support CPU for now with GPU support as a follow up, cc dstaay-fb # Test Updates tests for anything touching RDMABuffer. Splits out `test_python_actor` into `test_python_actor` and `test_rdma` now that RDMA is utilizing the backend network. Differential Revision: D76937776
|
This pull request was exported from Phabricator. Differential Revision: D76937776 |
allenwang28
added a commit
to allenwang28/monarch-1
that referenced
this pull request
Jul 9, 2025
Summary: This change introduces an `extension` folder for `monarch_rdma`, prepping for Python<>RDMA APIs: - Consistent with other bindings, provides a blocking and non-blocking version. - Corrects behavior of a few edge cases for RDMA components, i.e. `ibverbs_supported` based on the number of devices, supporting a "loopback" case, wherein two actors share an RDMA buffer but are spawned on the same proc. Differential Revision: D76937776
allenwang28
force-pushed
the
export-D76937776
branch
from
63ce0c9 to
47406f1
Compare
|
This pull request was exported from Phabricator. Differential Revision: D76937776 |
allenwang28
force-pushed
the
export-D76937776
branch
from
47406f1 to
aca7f36
Compare
allenwang28
added a commit
to allenwang28/monarch-1
that referenced
this pull request
Jul 10, 2025
Summary: This change introduces an `extension` folder for `monarch_rdma`, prepping for Python<>RDMA APIs: - Consistent with other bindings, provides a blocking and non-blocking version. - Corrects behavior of a few edge cases for RDMA components, i.e. `ibverbs_supported` based on the number of devices, supporting a "loopback" case, wherein two actors share an RDMA buffer but are spawned on the same proc. Reviewed By: dstaay-fb Differential Revision: D76937776
|
This pull request was exported from Phabricator. Differential Revision: D76937776 |
allenwang28
added a commit
to allenwang28/monarch-1
that referenced
this pull request
Jul 10, 2025
Summary: This change introduces an `extension` folder for `monarch_rdma`, prepping for Python<>RDMA APIs: - Consistent with other bindings, provides a blocking and non-blocking version. - Corrects behavior of a few edge cases for RDMA components, i.e. `ibverbs_supported` based on the number of devices, supporting a "loopback" case, wherein two actors share an RDMA buffer but are spawned on the same proc. Reviewed By: dstaay-fb Differential Revision: D76937776
allenwang28
force-pushed
the
export-D76937776
branch
from
aca7f36 to
ef4acef
Compare
|
This pull request was exported from Phabricator. Differential Revision: D76937776 |
allenwang28
added a commit
to allenwang28/monarch-1
that referenced
this pull request
Jul 10, 2025
Summary: This change introduces an `extension` folder for `monarch_rdma`, prepping for Python<>RDMA APIs: - Consistent with other bindings, provides a blocking and non-blocking version. - Corrects behavior of a few edge cases for RDMA components, i.e. `ibverbs_supported` based on the number of devices, supporting a "loopback" case, wherein two actors share an RDMA buffer but are spawned on the same proc. Reviewed By: dstaay-fb Differential Revision: D76937776
allenwang28
force-pushed
the
export-D76937776
branch
from
ef4acef to
1e20ff9
Compare
|
This pull request was exported from Phabricator. Differential Revision: D76937776 |
Summary: This change introduces an `extension` folder for `monarch_rdma`, prepping for Python<>RDMA APIs: - Consistent with other bindings, provides a blocking and non-blocking version. - Corrects behavior of a few edge cases for RDMA components, i.e. `ibverbs_supported` based on the number of devices, supporting a "loopback" case, wherein two actors share an RDMA buffer but are spawned on the same proc. Reviewed By: dstaay-fb Differential Revision: D76937776
allenwang28
force-pushed
the
export-D76937776
branch
from
1e20ff9 to
f2d0f91
Compare
|
This pull request was exported from Phabricator. Differential Revision: D76937776 |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary:
This change brings in RDMA support from Rust to the Python APIs, overwriting the existing stand-in RDMABuffer APIs.
There are quite a few changes, so here's a high level overview:
monarch_hyperactor(rust)Within
PyProcMesh, adds an optional RDMAMangerActor attribute.This mimics the prior functionality - before, a
proc_mesh(in Python) would unconditionally spin up anRdmaManagerActor.Now, a
proc_meshwill spin up anRdmaManagerActorif and only iftensor_engineis enabled andibverbsAPIs are supported.A few other pieces must be accessible within the RDMA extension library - specifically
Mailbox/ the send caps capabiliites.RDMA extension
Adds an
extensionfolder withinmonarch_rdmawhich contains the actual bindings.Consistent with other bindings, provides a blocking and non-blocking version.
RDMA components
Corrects behavior of a few edge cases, i.e.
ibverbs_supportedbased on the number of devices, supporting a "loopback" case, wherein two actors share an RDMA buffer but are spawned on the same proc.Python
Removes the stand-in RDMAManagerActor
In rdma.py, adds in basic support for RDMABuffer CPU.
While CUDA<>CPU|CUDA is implemented in
monarch_rdma, this was running into issues with MR registration.Logging shows that the code path correctly interprets the address as CUDA, but
cuMemGetHandleForAddressRangereturns the handle/fd as-1and a null pointer MR.RDMABuffer will only support CPU for now with GPU support as a follow up, cc dstaay-fb
Test
Updates tests for anything touching RDMABuffer. Splits out
test_python_actorintotest_python_actorandtest_rdmanow that RDMA is utilizing the backend network.Differential Revision: D76937776