test_save_strategies.py

# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import pytest
import torch
from megatron.core.dist_checkpointing.mapping import ShardedTensor
from torch.distributed.checkpoint import metadata as torchdistmeta
from torch.distributed.checkpoint.planner import SavePlan, WriteItem, WriteItemType

from ml_flashpoint.adapter.megatron.save_strategies import MLFlashpointMegatronAsyncSaveStrategy
from ml_flashpoint.adapter.pytorch.memory_storage_writer import MemoryStorageWriter
from ml_flashpoint.checkpoint_object_manager.checkpoint_object_manager import CheckpointObjectManager
from ml_flashpoint.core.checkpoint_id_types import CheckpointContainerId, CheckpointObjectId
from ml_flashpoint.core.checkpoint_saver import (
    DefaultMLFlashpointCheckpointSaver,
    MLFlashpointCheckpointSaver,
    ObjectWriteBucket,
)
from ml_flashpoint.replication.replication_manager import ReplicationManager

_default_test_global_rank = 0
_default_test_local_rank = 0


@pytest.fixture
def checkpoint_saver() -> MLFlashpointCheckpointSaver:
    return DefaultMLFlashpointCheckpointSaver(
        global_rank_getter=lambda: _default_test_global_rank,
        local_rank_getter=lambda: _default_test_local_rank,
        global_barrier_func=lambda: None,
        ckpt_obj_manager=CheckpointObjectManager(),
        replication_manager=ReplicationManager(),
    )


@pytest.fixture
def storage_writer(mocker, checkpoint_saver) -> MemoryStorageWriter:
    # Using a real MemoryStorageWriter instance instead of a mock.
    # We can still spy on its methods if needed.
    # The mp_manager is mocked as it's not relevant to these tests.
    return MemoryStorageWriter(
        checkpoint_saver=checkpoint_saver,
        mp_manager_future=mocker.MagicMock(),
    )


class TestMLFlashpointMegatronAsyncSaveStrategy:
    def test_init(self, storage_writer, checkpoint_saver):
        # Given
        strategy = MLFlashpointMegatronAsyncSaveStrategy(storage_writer=storage_writer)

        # When/Then
        assert strategy._storage_writer is storage_writer
        assert strategy._checkpoint_saver is checkpoint_saver
        assert strategy._use_cached_ckpt_structure is False

    def test_can_handle_sharded_objects(self, storage_writer):
        # Given
        strategy = MLFlashpointMegatronAsyncSaveStrategy(storage_writer=storage_writer)

        # When/Then
        assert strategy.can_handle_sharded_objects() is True

    class TestAsyncSave:
        @pytest.fixture(autouse=True)
        def mock_dist(self, mocker):
            mocker.patch("torch.distributed.is_initialized", return_value=True)
            mocker.patch("torch.distributed.get_rank", return_value=_default_test_global_rank)
            mocker.patch("torch.distributed.get_node_local_rank", return_value=_default_test_local_rank)
            mocker.patch("torch.distributed.get_world_size", return_value=1)

        @pytest.fixture
        def checkpoint_id(self, tmp_path):
            return CheckpointContainerId(str(tmp_path / "test_checkpoint"))

        @pytest.fixture
        def dummy_save_plan(self):
            return SavePlan(
                [
                    WriteItem(index=torchdistmeta.MetadataIndex("tensor1"), type=WriteItemType.TENSOR),
                    WriteItem(index=torchdistmeta.MetadataIndex("tensor2"), type=WriteItemType.TENSOR),
                    WriteItem(index=torchdistmeta.MetadataIndex("dir1/shard_0_0"), type=WriteItemType.BYTE_IO),
                ]
            )

        @pytest.fixture
        def dummy_metadata(self):
            return torchdistmeta.Metadata(
                state_dict_metadata={
                    "tensor1": torchdistmeta.TensorStorageMetadata(
                        size=torch.Size([10, 20]),
                        properties=torchdistmeta.TensorProperties(dtype=torch.float32),
                        chunks=[
                            torchdistmeta.ChunkStorageMetadata(offsets=torch.Size([0, 0]), sizes=torch.Size([5, 10]))
                        ],
                    ),
                    "tensor2": torchdistmeta.TensorStorageMetadata(
                        size=torch.Size([30, 40]),
                        properties=torchdistmeta.TensorProperties(dtype=torch.float32),
                        chunks=[
                            torchdistmeta.ChunkStorageMetadata(offsets=torch.Size([0, 0]), sizes=torch.Size([5, 10]))
                        ],
                    ),
                    "dir1/shard_0_0": torchdistmeta.BytesStorageMetadata(),
                }
            )

        @pytest.fixture
        def dummy_write_buckets(self, mocker, checkpoint_id):
            return [
                ObjectWriteBucket(
                    object_id=CheckpointObjectId(f"{checkpoint_id.data}/obj_{i}"),
                    object_name=f"obj_{i}",
                    bytesio_data=[],
                    tensor_data=[(mocker.MagicMock(), torch.tensor([i]))],
                )
                for i in range(3)
            ]

        @pytest.fixture
        def async_save_setup(self, mocker, monkeypatch, storage_writer, checkpoint_id):
            monkeypatch.setenv("MLFLASHPOINT_DISABLE_DIST", "true")

            strategy = MLFlashpointMegatronAsyncSaveStrategy(storage_writer=storage_writer)

            sharded_state_dict = {
                "layer1.weight": ShardedTensor(
                    key="layer1.weight",
                    data=torch.rand(10, 5),
                    dtype=torch.float32,
                    local_shape=(10, 5),
                    global_shape=(10, 5),
                    global_offset=(0, 0),
                    axis_fragmentations=(1, 1),
                )
            }
            pyt_state_dict = {"tensor1": torch.tensor([1, 2, 3]), "non_tensor": "test_string"}

            mocker.patch(
                "ml_flashpoint.adapter.megatron.save_strategies._replace_state_dict_keys_with_sharded_keys",
                return_value=(sharded_state_dict, None, None),
            )
            mocker.patch(
                "ml_flashpoint.adapter.megatron.save_strategies.mcore_to_pyt_state_dict", return_value=pyt_state_dict
            )

            return strategy, checkpoint_id, sharded_state_dict, pyt_state_dict

        def test_async_save_initialization_calls_success(
            self, mocker, async_save_setup, storage_writer, checkpoint_saver, dummy_write_buckets
        ):
            """Tests the initialization calls within async_save, including StorageWriter re-initialization."""
            # Given
            mock_statedictsaver = mocker.patch("ml_flashpoint.adapter.megatron.save_strategies.statedictsaver")
            (
                strategy,
                checkpoint_id,
                sharded_state_dict,
                _,
            ) = async_save_setup
            mock_statedictsaver.generate_plan.return_value = (
                dummy_write_buckets,
                mocker.MagicMock(),
                mocker.MagicMock(),
                mocker.MagicMock(),
                False,
            )

            mock_memory_storage_writer_cls = mocker.patch(
                "ml_flashpoint.adapter.megatron.save_strategies.MemoryStorageWriter"
            )
            mock_new_storage_writer_instance = mock_memory_storage_writer_cls.return_value

            initialize_checkpoint_spy = mocker.spy(checkpoint_saver, "initialize_checkpoint")

            # When
            strategy.async_save(sharded_state_dict, checkpoint_id.data)

            # Then
            initialize_checkpoint_spy.assert_called_once_with(checkpoint_id)

            mock_memory_storage_writer_cls.assert_called_once_with(
                checkpoint_saver=checkpoint_saver,
                mp_manager_future=storage_writer._main_process_torchmp_manager_future,
                files_per_rank=storage_writer._files_per_rank,
            )
            mock_new_storage_writer_instance.reset.assert_called_once_with(checkpoint_id.data)
            mock_new_storage_writer_instance.stage_write_data_buckets.assert_called_once_with(
                checkpoint_id, dummy_write_buckets, non_blocking=True
            )

        @pytest.mark.parametrize("expected_files_per_rank", [1, 2, 3, 5])
        def test_async_save_reinitializes_storage_writer_with_files_per_rank(
            self, mocker, async_save_setup, storage_writer, checkpoint_saver, dummy_write_buckets, expected_files_per_rank
        ):
            """Tests that the StorageWriter is re-initialized with the correct files_per_rank."""
            # Given
            mock_statedictsaver = mocker.patch("ml_flashpoint.adapter.megatron.save_strategies.statedictsaver")
            (
                strategy,
                checkpoint_id,
                sharded_state_dict,
                _,
            ) = async_save_setup
            mock_statedictsaver.generate_plan.return_value = (
                dummy_write_buckets,
                mocker.MagicMock(),
                mocker.MagicMock(),
                mocker.MagicMock(),
                False,
            )

            # Set a specific files_per_rank on the original storage_writer
            storage_writer._files_per_rank = expected_files_per_rank

            mock_memory_storage_writer_cls = mocker.patch(
                "ml_flashpoint.adapter.megatron.save_strategies.MemoryStorageWriter"
            )

            # When
            strategy.async_save(sharded_state_dict, checkpoint_id.data)

            # Then
            mock_memory_storage_writer_cls.assert_called_once_with(
                checkpoint_saver=checkpoint_saver,
                mp_manager_future=storage_writer._main_process_torchmp_manager_future,
                files_per_rank=expected_files_per_rank,
            )

        def test_initialize_checkpoint_failure(self, mocker, async_save_setup, checkpoint_saver):
            """Tests that the process terminates gracefully if initialize_checkpoint fails."""
            # Given
            strategy, checkpoint_id, sharded_state_dict, _ = async_save_setup
            mocker.patch.object(checkpoint_saver, "initialize_checkpoint", side_effect=Exception("Init Failed"))

            # When / Then
            with pytest.raises(Exception, match="Init Failed"):
                strategy.async_save(sharded_state_dict, checkpoint_id.data)

        def test_async_save_generate_plan_call_success(self, mocker, async_save_setup, storage_writer):
            """Tests that generate_plan is called correctly within async_save."""
            # Given
            mock_statedictsaver = mocker.patch("ml_flashpoint.adapter.megatron.save_strategies.statedictsaver")
            MockMCoreSavePlanner = mocker.patch("ml_flashpoint.adapter.megatron.save_strategies.MCoreSavePlanner")
            (
                strategy,
                checkpoint_id,
                sharded_state_dict,
                pyt_state_dict,
            ) = async_save_setup
            mock_planner = MockMCoreSavePlanner.return_value
            mock_statedictsaver.generate_plan.return_value = (
                mocker.MagicMock(),
                mocker.MagicMock(),
                mocker.MagicMock(),
                mocker.MagicMock(),
                False,
            )

            expected_kwarg_keys = {
                "checkpoint_id",
                "state_dict",
                "storage_writer",
                "planner",
                "world_dist_wrapper",
                "cached_ckpt_structure",
            }

            # When
            strategy.async_save(sharded_state_dict, checkpoint_id.data)

            # Then
            mock_statedictsaver.generate_plan.assert_called_once()
            _, kwargs = mock_statedictsaver.generate_plan.call_args
            actual_storage_writer_used = kwargs.get("storage_writer", None)
            assert set(kwargs.keys()) == expected_kwarg_keys
            assert kwargs["checkpoint_id"] == checkpoint_id
            assert kwargs["state_dict"] == pyt_state_dict
            assert actual_storage_writer_used is not None
            assert isinstance(actual_storage_writer_used, MemoryStorageWriter)
            assert (
                actual_storage_writer_used._main_process_torchmp_manager_future
                is storage_writer._main_process_torchmp_manager_future
            )
            assert kwargs["planner"] is mock_planner
            assert "world_dist_wrapper" in kwargs
            assert kwargs["world_dist_wrapper"].use_dist is False
            assert "cached_ckpt_structure" in kwargs
            assert "cached_global_metadata" not in kwargs

        def test_generate_plan_failure(self, mocker, async_save_setup):
            """Tests that an exception in generate_plan is propagated."""
            # Given
            mock_statedictsaver = mocker.patch("ml_flashpoint.adapter.megatron.save_strategies.statedictsaver")
            strategy, checkpoint_id, sharded_state_dict, _ = async_save_setup
            mock_statedictsaver.generate_plan.side_effect = ValueError("Plan Failed")

            # When / Then
            with pytest.raises(ValueError, match="Plan Failed"):
                strategy.async_save(sharded_state_dict, checkpoint_id.data)

        def test_async_save_async_fn_call_success(
            self, mocker, async_save_setup, storage_writer, dummy_save_plan, dummy_metadata, dummy_write_buckets
        ):
            """Tests that the async_fn returned by async_save calls write_data."""
            # Given
            from ml_flashpoint.core.checkpoint_id_types import CheckpointObjectId
            from ml_flashpoint.core.checkpoint_saver import ObjectWriteBucket

            mock_statedictsaver = mocker.patch("ml_flashpoint.adapter.megatron.save_strategies.statedictsaver")
            strategy, checkpoint_id, sharded_state_dict, _ = async_save_setup
            mock_statedictsaver.generate_plan.return_value = (
                dummy_write_buckets,
                dummy_metadata,
                mocker.MagicMock(),
                mocker.MagicMock(),
                False,
            )
            staged_write_buckets = [
                ObjectWriteBucket(
                    object_id=CheckpointObjectId(f"/test_checkpoint/staged_obj_{i}"),
                    object_name=f"staged_obj_{i}",
                    bytesio_data=[],
                    tensor_data=[(mocker.MagicMock(), torch.tensor([i + 10]))],
                )
                for i in range(2)
            ]

            mock_memory_storage_writer_cls = mocker.patch(
                "ml_flashpoint.adapter.megatron.save_strategies.MemoryStorageWriter"
            )
            mock_new_storage_writer_instance = mock_memory_storage_writer_cls.return_value
            mock_new_storage_writer_instance.stage_write_data_buckets.return_value = staged_write_buckets

            # When
            actual_async_request = strategy.async_save(sharded_state_dict, checkpoint_id.data)
            actual_async_request.async_fn(**actual_async_request.async_fn_kwargs)

            # Then
            mock_statedictsaver.write_data.assert_called_once_with(
                checkpoint_id=checkpoint_id,
                storage_writer=mock_new_storage_writer_instance,
                staged_write_buckets=staged_write_buckets,
                replicate_after_write=False,
            )

        def test_async_save_async_fn_failure(self, mocker, async_save_setup, checkpoint_saver):
            """Tests that finalize_checkpoint is not called when async_fn fails."""
            # Given
            finalize_checkpoint_spy = mocker.spy(checkpoint_saver, "finalize_checkpoint")
            mock_statedictsaver = mocker.patch("ml_flashpoint.adapter.megatron.save_strategies.statedictsaver")
            strategy, checkpoint_id, sharded_state_dict, _ = async_save_setup
            mock_statedictsaver.generate_plan.return_value = (
                mocker.MagicMock(),
                mocker.MagicMock(),
                mocker.MagicMock(),
                mocker.MagicMock(),
                False,
            )
            mock_statedictsaver.write_data.side_effect = Exception("Test Exception")

            # When
            actual_async_request = strategy.async_save(sharded_state_dict, checkpoint_id.data)
            with pytest.raises(Exception, match="Test Exception"):
                actual_async_request.async_fn(**actual_async_request.async_fn_kwargs)

            # Then
            finalize_checkpoint_spy.assert_not_called()

        def test_async_save_finalize_fns_calls(
            self,
            mocker,
            async_save_setup,
            storage_writer,
            checkpoint_saver,
            dummy_save_plan,
            dummy_metadata,
            dummy_write_buckets,
        ):
            """Tests that the finalize_fns returned by async_save call finish_write and finalize_checkpoint."""
            # Given
            finalize_checkpoint_spy = mocker.spy(checkpoint_saver, "finalize_checkpoint")
            mock_statedictsaver = mocker.patch("ml_flashpoint.adapter.megatron.save_strategies.statedictsaver")
            strategy, checkpoint_id, sharded_state_dict, _ = async_save_setup
            mock_statedictsaver.generate_plan.return_value = (
                dummy_write_buckets,
                dummy_metadata,
                mocker.MagicMock(),
                mocker.MagicMock(),
                False,
            )

            mock_memory_storage_writer_cls = mocker.patch(
                "ml_flashpoint.adapter.megatron.save_strategies.MemoryStorageWriter"
            )
            mock_storage_writer_instance = mock_memory_storage_writer_cls.return_value
            # We need to set _main_process_torchmp_manager_future on the mock because the test asserts on it later
            mock_storage_writer_instance._main_process_torchmp_manager_future = (
                storage_writer._main_process_torchmp_manager_future
            )
            mock_storage_writer_instance.stage_write_data_buckets.return_value = dummy_write_buckets

            expected_kwarg_keys = {"checkpoint_id", "storage_writer", "global_metadata", "world_dist_wrapper"}

            # When
            actual_async_request = strategy.async_save(sharded_state_dict, checkpoint_id.data)

            # Then
            assert len(actual_async_request.finalize_fns) == 3

            # Call 1st finalize function
            actual_async_request.finalize_fns[0]()

            # Then
            # Assert the actual storage_writer invoked replicate_written_objects with the entire set of object IDs
            expected_object_ids = {b.object_id for b in dummy_write_buckets}
            mock_storage_writer_instance.replicate_written_objects.assert_called_once_with(
                object_ids=expected_object_ids
            )

            # Call 2nd finalize function
            actual_async_request.finalize_fns[1]()

            # Then
            mock_statedictsaver.finish_write.assert_called_once()
            _, kwargs = mock_statedictsaver.finish_write.call_args
            actual_storage_writer_used = kwargs.get("storage_writer", None)
            assert set(kwargs.keys()) == expected_kwarg_keys
            assert kwargs["checkpoint_id"] == checkpoint_id
            assert actual_storage_writer_used is not None
            assert actual_storage_writer_used is mock_storage_writer_instance
            assert (
                actual_storage_writer_used._main_process_torchmp_manager_future
                is storage_writer._main_process_torchmp_manager_future
            )
            assert kwargs["global_metadata"] == dummy_metadata
            assert kwargs["world_dist_wrapper"].use_dist is False

            # Call 3rd finalize function
            actual_async_request.finalize_fns[2]()

            # Then
            finalize_checkpoint_spy.assert_called_once_with(checkpoint_id=checkpoint_id)

        def test_finalize_fns_failure(
            self, mocker, async_save_setup, checkpoint_saver, dummy_save_plan, dummy_metadata
        ):
            """Tests that a failure in the finish_write finalize_fn prevents finalize_checkpoint from running."""
            # Given
            finalize_checkpoint_spy = mocker.spy(checkpoint_saver, "finalize_checkpoint")
            mock_statedictsaver = mocker.patch("ml_flashpoint.adapter.megatron.save_strategies.statedictsaver")
            strategy, checkpoint_id, sharded_state_dict, _ = async_save_setup
            mock_statedictsaver.generate_plan.return_value = (
                mocker.MagicMock(),
                dummy_metadata,
                mocker.MagicMock(),
                mocker.MagicMock(),
                False,
            )
            mock_statedictsaver.finish_write.side_effect = ValueError("Finish Write Failed")

            # When
            actual_async_request = strategy.async_save(sharded_state_dict, checkpoint_id.data)
            with pytest.raises(ValueError, match="Finish Write Failed"):
                actual_async_request.finalize_fns[1]()

            # Then
            finalize_checkpoint_spy.assert_not_called()

        @pytest.mark.parametrize(
            "is_dist_initialized, dist_rank, expected_rank",
            [
                (True, 5, 5),
                (False, 0, -1),
            ],
        )
        def test_async_save_rank_determination(
            self,
            mocker,
            async_save_setup,
            is_dist_initialized,
            dist_rank,
            expected_rank,
        ):
            """Tests that the rank passed to async_fn is correct based on dist initialization."""
            # Given
            strategy, checkpoint_id, sharded_state_dict, _ = async_save_setup

            # Mock torch.distributed
            mocker.patch("torch.distributed.is_initialized", return_value=is_dist_initialized)
            if is_dist_initialized:
                mocker.patch("torch.distributed.get_rank", return_value=dist_rank)

            # Mock dependencies to ensure success path
            mock_statedictsaver = mocker.patch("ml_flashpoint.adapter.megatron.save_strategies.statedictsaver")
            mock_statedictsaver.generate_plan.return_value = (
                mocker.MagicMock(),
                mocker.MagicMock(),
                mocker.MagicMock(),
                mocker.MagicMock(),
                False,
            )

            # When
            actual_async_request = strategy.async_save(sharded_state_dict, checkpoint_id.data)

            # Then
            assert actual_async_request.async_fn_kwargs["rank"] == expected_rank

        def test_async_save_caching_flow(self, mocker, async_save_setup, storage_writer):
            """Tests the caching flow across multiple async_save calls."""
            # Given
            mock_statedictsaver = mocker.patch("ml_flashpoint.adapter.megatron.save_strategies.statedictsaver")
            strategy, checkpoint_id, sharded_state_dict, _ = async_save_setup
            cached_plan = mocker.MagicMock()
            cached_metadata = mocker.MagicMock()

            # --- Call 1: No cache ---
            # Given
            mock_statedictsaver.generate_plan.return_value = (
                [],
                mocker.MagicMock(),
                cached_plan,  # cached_central_plan returned
                mocker.MagicMock(),
                False,
            )

            # When
            strategy.async_save(sharded_state_dict, checkpoint_id.data)

            # Then
            assert strategy._cached_central_plan == cached_plan
            assert strategy._validated_cache_reuse is False

            # --- Call 2: Cache validation success ---
            # Given
            mock_statedictsaver.generate_plan.return_value = (
                [],
                cached_metadata,
                cached_plan,
                mocker.MagicMock(),
                True,
            )

            # When
            strategy.async_save(sharded_state_dict, checkpoint_id.data)

            # Then
            assert strategy._validated_cache_reuse is True
            assert strategy._cached_global_metadata == cached_metadata

            # --- Call 3: Reuse cache ---
            # Given
            mock_statedictsaver.generate_plan.return_value = (
                [],
                None,  # Returns None for metadata
                cached_plan,
                mocker.MagicMock(),
                True,
            )

            # When
            strategy.async_save(sharded_state_dict, checkpoint_id.data)

            # Then
            # Ensure generate_plan was called without cached_global_metadata
            _, kwargs = mock_statedictsaver.generate_plan.call_args
            assert "cached_global_metadata" not in kwargs
            # And cached_global_metadata in strategy should still be the same
            assert strategy._cached_global_metadata == cached_metadata

        def test_async_save_caching_disabled_by_default(self, mocker, async_save_setup, storage_writer):
            """Tests that caching is disabled by default."""
            # Given
            mock_statedictsaver = mocker.patch("ml_flashpoint.adapter.megatron.save_strategies.statedictsaver")
            strategy, checkpoint_id, sharded_state_dict, _ = async_save_setup

            cached_plan = mocker.MagicMock()

            # Call: Returns a plan that could be cached
            mock_statedictsaver.generate_plan.return_value = (
                [],
                None,
                cached_plan,
                mocker.MagicMock(),
                False,
            )
            # When
            strategy.async_save(sharded_state_dict, checkpoint_id.data)

            # Then
            # Should NOT have updated the specific cached plan attribute if we assume
            # generate_plan returns it regardless?
            # actually statedictsaver.generate_plan returns the plan to be cached.
            # But the strategy should NOT pass it back in the next call if use_cached_ckpt_structure is False.

            # Let's verify the next call doesn't pass it.
            strategy.async_save(sharded_state_dict, checkpoint_id.data)

            _, kwargs = mock_statedictsaver.generate_plan.call_args
            assert kwargs["cached_ckpt_structure"] is None
            assert strategy._use_cached_ckpt_structure is False

        def test_async_save_ensure_metadata_deepcopy(self, mocker, async_save_setup):
            """Tests that global_metadata is deepcopied to prevent cache pollution."""
            import copy

            # Given
            mock_statedictsaver = mocker.patch("ml_flashpoint.adapter.megatron.save_strategies.statedictsaver")
            strategy, checkpoint_id, sharded_state_dict, _ = async_save_setup

            # Use a real object that supports deepcopy and modification tracking
            class MockMetadata:
                def __init__(self, data):
                    self.data = data
                    self.storage_data = None

                def __eq__(self, other):
                    return self.data == other.data and self.storage_data == other.storage_data

                def __repr__(self):
                    return f"MockMetadata(data={self.data}, storage_data={self.storage_data})"

            original_metadata = MockMetadata({"key": "value"})

            # Spy on deepcopy
            deepcopy_spy = mocker.spy(copy, "deepcopy")

            # --- Call 1: Fresh metadata ---
            mock_statedictsaver.generate_plan.return_value = (
                [],
                original_metadata,
                mocker.MagicMock(),
                mocker.MagicMock(),
                False,
            )

            # When
            strategy.async_save(sharded_state_dict, checkpoint_id.data)

            # Then
            # 1. Verify deepcopy was called
            assert deepcopy_spy.call_count >= 1
            # 2. Verify cache holds a DIFFERENT object but with SAME content
            assert strategy._cached_global_metadata is not original_metadata
            assert strategy._cached_global_metadata == original_metadata

            # Simulate "dirtying" the metadata that was passed to finalize_fns
            original_metadata.storage_data = "DIRTY_DATA"
            assert strategy._cached_global_metadata.storage_data is None

            # --- Call 2: Reuse cached metadata ---
            # Reset mocks
            deepcopy_spy.reset_mock()
            mock_statedictsaver.generate_plan.return_value = (
                [],
                None,  # Returns None, triggering cache retrieval
                mocker.MagicMock(),
                mocker.MagicMock(),
                True,
            )

            # When
            strategy._use_cached_ckpt_structure = True

            request = strategy.async_save(sharded_state_dict, checkpoint_id.data)

            # Then
            # 1. Verify deepcopy called again (retrieving from cache)
            assert deepcopy_spy.call_count >= 1

            # 2. Verify the `global_metadata` passed to writing is a COPY
            finish_write_partial = request.finalize_fns[1]
            bound_metadata = finish_write_partial.keywords["global_metadata"]

            # Verify bound_metadata is a COPY of cache
            assert bound_metadata is not strategy._cached_global_metadata
            assert bound_metadata == strategy._cached_global_metadata

            # Verify modifications to bound_metadata don't affect cache
            bound_metadata.storage_data = "NEW_DIRTY"
            assert strategy._cached_global_metadata.storage_data is None