dmlcloud.root_only

dmlcloud.root_only(fn, group=None, synchronize=True, timeout=86400)

Decorator for methods that should only be called on the root rank.

Can also be applied to individual callback methods of Pipeline and Stage, or to the whole class. In that case, Pipeline.gloo_group is used as process group.

If synchronize=True, a monitored_barrier before or after the function call depending on the rank. This can be important to prevent timeouts from future all_reduce operations if non-root ranks move on before the root rank has finished.

Parameters:

fn (Union[Callable, type]) – The function to decorate or a subclass of Pipeline or Stage.
group (ProcessGroup) – The process group to work on. If None (default), the default process group will be used.
synchronize (bool) – If True, a barrier is inserted before or after the function call depending on the rank. Default is True.
timeout (int) – Timeout in seconds for the monitored_barrier. Default is 24 hours.

Return type:

Union[Callable, type]

Returns:

The decorated function or class.

Examples

Annotating an individual function:

>>> @root_only
>>> def my_function():
>>>     print('Only the root rank prints this.')

Annotating a whole Stage subclass:

>>> @root_only
>>> class MyStage(Stage):
>>>     def pre_stage(self):
>>>         print('Only the root rank prints this.')
>>>
>>>     def run_epoch(self):
>>>         print('Only the root rank prints this.')
>>>
>>>     def post_stage(self):
>>>         print('Only the root rank prints this.')

Annotating individual methods of Stage:

>>> class MyStage(Stage):
>>>     def pre_stage(self):
>>>         print('All ranks print this.')
>>>
>>>     @root_only
>>>     def post_stage(self):
>>>         print('Only the root rank prints this.')