microsoft/mscclpp
1
0
Fork 0
mirror of https://github.com/microsoft/mscclpp.git synced 2026-04-19 22:39:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

Adding Support to Setting Message Size Range in Native Algorithm API (#758)

Browse Source
This commit is contained in:
Caio Rocha
2026-02-27 17:50:43 -08:00
committed by GitHub
parent ab49386839
commit 4bc1999001
4 changed files with 36 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/csrc/algorithm.cpp
View File

@@ -60,6 +60,12 @@ void register_algorithm(nb::module_& m) {
.def_prop_ro("name", &Algorithm::name)
.def_prop_ro("collective", &Algorithm::collective)
.def_prop_ro("message_range", &Algorithm::messageRange)
.def(
"set_message_size_range",
[](Algorithm& self, size_t minMessageSize, size_t maxMessageSize) {
self.setMessageSizeRange(minMessageSize, maxMessageSize);
},
nb::arg("min_message_size"), nb::arg("max_message_size"))
.def_prop_ro("tags", &Algorithm::tags)
.def_prop_ro("buffer_mode", &Algorithm::bufferMode)
.def_prop_ro("constraint", &Algorithm::constraint)

15
python/mscclpp/_core/algorithm.py
View File

@@ -114,11 +114,24 @@ class Algorithm:
"""The collective operation this algorithm implements (e.g., "allreduce", "allgather")."""
return self._algorithm.collective
@cached_property
@property
def message_size_range(self) -> Tuple[int, int]:
"""The valid message size range (min_size, max_size) in bytes."""
return (self._algorithm.message_range[0], self._algorithm.message_range[1])
def set_message_size_range(self, min_message_size: int, max_message_size: int):
"""Set the valid message size range in bytes.
Args:
min_message_size: Minimum supported message size in bytes.
max_message_size: Maximum supported message size in bytes.
Only supported for native algorithms. Raises TypeError for DSL algorithms.
"""
if self.is_dsl_algorithm():
raise TypeError("set_message_size_range is only supported for native algorithms")
self._algorithm.set_message_size_range(min_message_size, max_message_size)
@cached_property
def tags(self) -> Dict[str, int]:
"""Dictionary of tag names to tag values for algorithm selection hints."""