resolve comments

Author: Leahlijuan

src/ml_flashpoint/adapter/pytorch/custom_state_dict_saver.py

@@ -131,9 +131,9 @@ def global_step(all_local_plans: list[SavePlan]) -> list[SavePlan]:
 
     return (
         (final_local_plan, write_buckets, global_metadata),
-        central_plan,  # cached_central_plan
-        local_plan,  # cached_local_plan
-        cached_central_plan == central_plan,  # validated_cache_reuse
+        central_plan,
+        local_plan,
+        cached_central_plan == central_plan,
     )
 
 

tests/adapter/megatron/test_save_strategies.py

@@ -512,56 +512,54 @@ def test_async_save_caching_flow(self, mocker, async_save_setup, storage_writer)
             # Given
             mock_statedictsaver = mocker.patch("ml_flashpoint.adapter.megatron.save_strategies.statedictsaver")
             strategy, checkpoint_id, sharded_state_dict, _ = async_save_setup
-            # Enable caching significantly for this test
-            strategy.use_cached_ckpt_structure = True
-
             cached_plan = mocker.MagicMock()
             cached_metadata = mocker.MagicMock()
 
-            # First call: No cache
+            # --- Call 1: No cache ---
+            # Given
             mock_statedictsaver.generate_plan.return_value = (
                 (mocker.MagicMock(), [], mocker.MagicMock()),
                 cached_plan,  # cached_central_plan returned
                 mocker.MagicMock(),
                 False,
             )
 
-            # When 1
+            # When
             strategy.async_save(sharded_state_dict, checkpoint_id.data)
 
-            # Then 1
+            # Then
             assert strategy.cached_central_plan == cached_plan
             assert strategy.validated_cache_reuse is False
 
-            # Second call: Cache validation success
+            # --- Call 2: Cache validation success ---
+            # Given
             mock_statedictsaver.generate_plan.return_value = (
                 (mocker.MagicMock(), [], cached_metadata),
                 cached_plan,
                 mocker.MagicMock(),
                 True,  # validated_cache_reuse
             )
 
-            # When 2
+            # When
             strategy.async_save(sharded_state_dict, checkpoint_id.data)
 
-            # Then 2
+            # Then
             assert strategy.validated_cache_reuse is True
             assert strategy.cached_global_metadata == cached_metadata
 
-            # Third call: Reuse cache
+            # --- Call 3: Reuse cache ---
+            # Given
             mock_statedictsaver.generate_plan.return_value = (
                 (mocker.MagicMock(), [], None),  # Returns None for metadata
                 cached_plan,
                 mocker.MagicMock(),
                 True,
             )
 
-            # During third call, async_save should use self.cached_global_metadata
-
-            # When 3
+            # When
             strategy.async_save(sharded_state_dict, checkpoint_id.data)
 
-            # Then 3
+            # Then
             # Ensure generate_plan was called without cached_global_metadata
             _, kwargs = mock_statedictsaver.generate_plan.call_args
             assert "cached_global_metadata" not in kwargs