zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Expose per-backend availability in C++ and python

Browse Source
This commit is contained in:
Angelos Katharopoulos
2025-11-20 15:26:59 -08:00
parent 1216afdc91
commit bfdddd644b
4 changed files with 42 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
python/src/distributed.cpp
View File

@@ -52,9 +52,25 @@ void init_distributed(nb::module_& parent_module) {
m.def(
"is_available",
&mx::distributed::is_available,
[](const std::string& backend) {
return mx::distributed::is_available(backend);
},
"backend"_a = "any",
nb::sig("def is_available(backend: str = 'any') -> bool"),
R"pbdoc(
Check if a communication backend is available.
Note, this function returns whether MLX has the capability of
instantiating that distributed backend not whether it is possible to
create a communication group. For that purpose one should use
``init(strict=True)``.
Args:
backend (str, optional): The name of the backend to check for availability.
It takes the same values as ``init()``. Default: ``any``.
Returns:
bool: Whether the distributed backend is available.
)pbdoc");
m.def(
@@ -79,10 +95,10 @@ void init_distributed(nb::module_& parent_module) {
in case ``mx.distributed.is_available()`` returns False otherwise
it throws a runtime error. Default: ``False``
backend (str, optional): Which distributed backend to initialize.
Possible values ``mpi``, ``ring``, ``nccl``, ``any``. If set to ``any`` all
available backends are tried and the first one that succeeds
becomes the global group which will be returned in subsequent
calls. Default: ``any``
Possible values ``mpi``, ``ring``, ``nccl``, ``ibv``, ``any``. If
set to ``any`` all available backends are tried and the first one
that succeeds becomes the global group which will be returned in
subsequent calls. Default: ``any``
Returns:
Group: The group representing all the launched processes.