Decorator to selectively run the decorated function on the main process only based on the `main_process_only` attribute in a class. Checks at function execution rather than initialization time, not triggering the initialization of the `PartialState`.
(function)
| 13 | |
| 14 | |
| 15 | def on_main_process(function): |
| 16 | """Decorator to selectively run the decorated function on the main process |
| 17 | only based on the `main_process_only` attribute in a class. |
| 18 | |
| 19 | Checks at function execution rather than initialization time, not |
| 20 | triggering the initialization of the `PartialState`. |
| 21 | """ |
| 22 | |
| 23 | def execute_on_main_process(self, *args, **kwargs): |
| 24 | if getattr(self, 'main_process_only', False): |
| 25 | if is_main_process(): |
| 26 | return function(self, *args, **kwargs) |
| 27 | else: |
| 28 | return do_nothing(self, *args, **kwargs) |
| 29 | else: |
| 30 | return function(self, *args, **kwargs) |
| 31 | |
| 32 | return execute_on_main_process |
| 33 | |
| 34 | |
| 35 | def get_local_rank(): |
nothing calls this directly
no outgoing calls
no test coverage detected