fix(core/saver): create checkpoint_dir in each rank when initialize_checkpoint (#63)

Fix the error:
```
[MLF 2026-02-26 19:09:24,334 ERROR Step=110 Rank=5 ml_flashpoint.adapter.megatron.save_utils:78] Failed to save ML Flashpoint checkpoint. Skipping saving and continuing.
Traceback (most recent call last):
  File "/tmp/ml-flashpoint/src/ml_flashpoint/adapter/megatron/save_utils.py", line 67, in save_local_aware_megatron_checkpoint
    return save_strategy.async_save(
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.12/contextlib.py", line 81, in inner
    return func(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^
  File "/tmp/ml-flashpoint/src/ml_flashpoint/adapter/megatron/save_strategies.py", line 220, in async_save
    with open(os.path.join(checkpoint_dir, "metadata.json"), "w") as f:
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
FileNotFoundError: [Errno 2] No such file or directory: '/tmp/logs/ml-flashpoint/job-plan/step-110_ckpt/metadata.json'
```

Author: Leahlijuan

src/ml_flashpoint/core/checkpoint_saver.py

@@ -339,9 +339,8 @@ def __setstate__(self, state):
     @log_execution_time(logger=_LOGGER, name="initialize_checkpoint")
     def initialize_checkpoint(self, checkpoint_id: CheckpointContainerId) -> None:
         self._create_dirty_checkpoint_marker(checkpoint_id)
-        if self._local_rank_getter() == 0:
-            os.makedirs(checkpoint_id.data, exist_ok=True)
-            _LOGGER.info("Created checkpoint directory: '%s'", checkpoint_id.data)
+        os.makedirs(checkpoint_id.data, exist_ok=True)
+        _LOGGER.info("Created checkpoint directory: '%s'", checkpoint_id.data)
 
     @override
     @log_execution_time(logger=_LOGGER, name="stage_data", level=logging.INFO)

tests/core/test_checkpoint_saver.py

@@ -406,7 +406,7 @@ def mock_open(file, mode):
             assert not os.path.exists(checkpoint_id.data)
 
         @pytest.mark.parametrize("local_rank", [0, 1, 8])
-        def test_initialize_checkpoint_creates_container_dir_when_not_exists_local_rank_0(
+        def test_initialize_checkpoint_creates_container_dir_when_not_exists(
             self,
             local_rank,
             temp_dir_path,
@@ -427,12 +427,9 @@ def test_initialize_checkpoint_creates_container_dir_when_not_exists_local_rank_
             saver.initialize_checkpoint(checkpoint_id)
 
             # Then
-            # Check for directory creation only on local rank 0
-            if local_rank == 0:
-                assert os.path.exists(checkpoint_id.data)
-                assert os.path.isdir(checkpoint_id.data)
-            else:
-                assert not os.path.exists(checkpoint_id.data)
+            # Check for directory creation on all ranks
+            assert os.path.exists(checkpoint_id.data)
+            assert os.path.isdir(checkpoint_id.data)
 
         @pytest.mark.parametrize("local_rank", [0, 1])
         def test_initialize_checkpoint_leaves_container_dir_when_exists(
@@ -494,11 +491,8 @@ def test_initialize_checkpoint_idempotent(
 
             # Then
             assert os.path.exists(expected_dirty_marker_path)
-            if local_rank == 0:
-                assert os.path.exists(checkpoint_id.data)
-                assert os.path.isdir(checkpoint_id.data)
-            else:
-                assert not os.path.exists(checkpoint_id.data)
+            assert os.path.exists(checkpoint_id.data)
+            assert os.path.isdir(checkpoint_id.data)
 
         def test_initialize_checkpoint_ensures_parent_dir_exists(
             self,