Removed atomic operations from mesh_consistency.save_process_metadata

PiperOrigin-RevId: 759460002

Author: Orbax Authors

checkpoint/orbax/checkpoint/experimental/emergency/mesh_consistency.py

@@ -20,13 +20,10 @@
 from absl import logging
 from etils import epath
 import jax
-from orbax.checkpoint import options as options_lib
 from orbax.checkpoint.experimental.emergency import multihost as emergency_multihost
-from orbax.checkpoint.path import atomicity_defaults
 from orbax.checkpoint.path import step as step_lib
 
 
-
 _PROCESS_METADATA_FOLDER = 'process_metadata'
 _PROCESS_METADATA_FILE_NAME = 'process_metadata.json'
 _GLOBAL_PROCESS_METADATA_FILE_NAME = 'global_process_metadata.json'
@@ -76,29 +73,17 @@ async def save_process_metadata(
   """Saves process metadata to local storage. Runs on every process."""
   metadata_folder = process_metadata_folder(directory)
   logging.info('Saving process index metadata at %s', metadata_folder)
-  if metadata_folder.exists():
-    logging.warning(
-        'Process metadata folder already exists at %s. Overwriting.',
-        metadata_folder,
+  if not metadata_folder.exists():
+    raise FileNotFoundError(
+        f'Process metadata folder does not exist at {metadata_folder}.'
     )
-    metadata_folder.rmtree()
 
-  multiprocessing_options = options_lib.MultiprocessingOptions(
-      primary_host=None
-  )
-  tmp_path = atomicity_defaults.get_default_temporary_path_class(
-      metadata_folder
-  ).from_final(metadata_folder, multiprocessing_options=multiprocessing_options)
-  await tmp_path.create()
-
-  (tmp_path.get() / _GLOBAL_PROCESS_METADATA_FILE_NAME).write_text(
+  (metadata_folder / _GLOBAL_PROCESS_METADATA_FILE_NAME).write_text(
       json.dumps(distributed_to_device_ids)
   )
-  (tmp_path.get() / _MESH_METADATA_FILE_NAME).write_text(
+  (metadata_folder / _MESH_METADATA_FILE_NAME).write_text(
       json.dumps([int(id) for id in global_mesh.device_ids.flatten()])
   )
-  tmp_path.finalize(
-  )
 
 
 def consistent_restore_mesh_from_metadata(