[Feature] Add context manager to allow main process first.

src/accelerate/accelerator.py

@@ -14,6 +14,7 @@
 
 import gc
 import os
+from contextlib import contextmanager
 from typing import List, Optional, Union
 
 import torch
@@ -183,6 +184,33 @@ def use_fp16(self):
             use_fp16 = self.state.use_fp16
         return use_fp16
 
+    @contextmanager
+    def local_main_process_first(self):
+        """
+        Lets the local main process go inside a with block.
+
+        The other processes will enter the with block after the main process exits.
+        """
+        yield from self._goes_first(self.is_local_main_process)
+
+    @contextmanager
+    def main_process_first(self):
+        """
+        Lets the main process go first inside a with block.
+
+        The other processes will enter the with block after the main process exits.
+        """
+        yield from self._goes_first(self.is_main_process)
+
+    def _goes_first(self, is_master):
+        if not is_master:
+            self.wait_for_everyone()
+
+        yield
+
+        if is_master:
+            self.wait_for_everyone()
+
     def print(self, *args, **kwargs):
         """
         Use in replacement of :obj:`print()` to only print once per server.