torchtnt.utils.distributed.rank_zero_fn

torchtnt.utils.distributed.rank_zero_fn(fn: Callable[[...], TReturn]) → Callable[[...], Optional[TReturn]]

Function that can be used as a decorator to enable a function to be called on global rank 0 only.

Note

This decorator should be used judiciously. it should never be used on functions that need synchronization. It should be used very carefully with functions that mutate local state as well

Example

>>> from torchtnt.utilities.distributed import rank_zero_fn
>>> @rank_zero_fn
... def foo():
...     return 1
...
>>> x = foo() # x is 1 if global rank is 0 else x is None

Parameters:fn – the desired function to be executed on rank 0 only Returns:the wrapped function that executes only if the global rank is 0 Return type:wrapped_fn