Resolve comments.

Author: Leahlijuan

…hpoint/adapter/nemo/checkpoint_loader.py …t/adapter/nemo/nemo_checkpoint_loader.pysrc/ml_flashpoint/adapter/nemo/checkpoint_loader.py renamed to src/ml_flashpoint/adapter/nemo/nemo_checkpoint_loader.py src/ml_flashpoint/adapter/nemo/checkpoint_loader.py renamed to src/ml_flashpoint/adapter/nemo/nemo_checkpoint_loader.py

@@ -48,14 +48,14 @@ def __init__(
         self._recover_context = recover_context
 
     @override
-    def _get_extra_local_objects(self, container_path: Path) -> List[str]:
+    def _get_extra_local_objects(self, container_path: Path) -> List[CheckpointObjectId]:
         """Returns extra local objects to include, specifically context files."""
         local_objects = []
         if self._recover_context:
             context_path = container_path / "context"
             if context_path.is_dir():
                 for entry in os.listdir(context_path):
-                    local_objects.append(os.path.join("context", entry))
+                    local_objects.append(CheckpointObjectId(str(context_path / entry)))
         return local_objects
 
     @override

src/ml_flashpoint/adapter/nemo/wrapper_util.py

@@ -27,7 +27,7 @@
 from ml_flashpoint.adapter.nemo.auto_resume import MLFlashpointAutoResume
 from ml_flashpoint.adapter.nemo.checkpoint_callback import MLFlashpointCheckpointCallback
 from ml_flashpoint.adapter.nemo.checkpoint_io import MLFlashpointAsyncFinalizableCheckpointIO, MLFlashpointCheckpointIO
-from ml_flashpoint.adapter.nemo.checkpoint_loader import NeMoMLFlashpointCheckpointLoader
+from ml_flashpoint.adapter.nemo.nemo_checkpoint_loader import NeMoMLFlashpointCheckpointLoader
 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

src/ml_flashpoint/core/checkpoint_loader.py

@@ -379,6 +379,8 @@ def _compute_retrieval_plan(
     ) -> Optional[dict[int, List[Tuple[int, str]]]]:
         """Computes the retrieval plan.
 
+        The plan assumes an even number of ranks (accelerator processes) on each node in the training cluster.
+
         Args:
             checkpoint: The checkpoint container ID.
             available_objects_by_rank: Map of rank to available objects on that rank.
@@ -540,15 +542,15 @@ def get_checkpoint_objects_by_rank(
             `CheckpointObjectId`s available on that node.
         """
         container_path = Path(checkpoint_container_id.data)
-        local_objects = []
+        local_objects: List[CheckpointObjectId] = []
         if not container_path.is_dir():
             _LOGGER.debug(
                 "Checkpoint container path '%s' is not a directory. Returning empty list.",
                 container_path,
             )
         else:
             for entry in os.listdir(container_path):
-                local_objects.append(entry)
+                local_objects.append(CheckpointObjectId.from_container(checkpoint_container_id, entry))
 
             local_objects.extend(self._get_extra_local_objects(container_path))
 
@@ -560,11 +562,8 @@ def get_checkpoint_objects_by_rank(
         if all_objects_by_rank_paths:
             for rank, objects in enumerate(all_objects_by_rank_paths):
                 if objects:
-                    # Convert filenames to full paths and then to CheckpointObjectId
-                    full_paths = [str(container_path / obj) for obj in objects]
-                    checkpoint_objects = [CheckpointObjectId(p) for p in full_paths]
-                    result[rank] = checkpoint_objects
-                    for obj in checkpoint_objects:
+                    result[rank] = objects
+                    for obj in objects:
                         object_locations[obj.data].append(rank)
                 else:
                     result[rank] = []
@@ -628,13 +627,13 @@ def retrieve_checkpoint(
         return all(all_success_list)
 
     def _get_extra_local_objects(self, container_path: Path) -> List[str]:
-        """Hook for subclasses to provide extra local objects that are available and relevant, 
+        """Hook for subclasses to provide extra local objects that are available and relevant,
         which may be needed by other hosts.
-        This can be used when additional objects beyond the standard checkpoint data are needed, 
+        This can be used when additional objects beyond the standard checkpoint data are needed,
         such as framework-specific context data.
-        
+
         This should always be implemented alongside `_get_extra_needed_objects`.
-        
+
         Returns:
             List of additional locally available objects.
         """
@@ -646,11 +645,11 @@ def _get_extra_needed_objects(
         available_objects_by_rank: dict[int, List[CheckpointObjectId]],
     ) -> Set[str]:
         """Hook for subclasses to provide extra needed objects on any given host (each local rank 0).
-        This can leverage `available_objects_by_rank` to determine the set of additional objects 
+        This can leverage `available_objects_by_rank` to determine the set of additional objects
         each host needs.
-        
+
         This should always be implemented alongside `_get_extra_local_objects`.
-        
+
         Returns:
             Set of extra needed objects on any given node.
         """

tests/adapter/nemo/test_nemo_checkpoint_loader.py

@@ -17,7 +17,7 @@
 
 import pytest
 
-from ml_flashpoint.adapter.nemo.checkpoint_loader import NeMoMLFlashpointCheckpointLoader
+from ml_flashpoint.adapter.nemo.nemo_checkpoint_loader import NeMoMLFlashpointCheckpointLoader
 from ml_flashpoint.checkpoint_object_manager.checkpoint_object_manager import CheckpointObjectManager
 from ml_flashpoint.core.checkpoint_id_types import CheckpointContainerId, CheckpointObjectId
 from ml_flashpoint.replication.replication_manager import ReplicationManager

tests/adapter/nemo/test_wrapper_util.py

@@ -256,6 +256,32 @@ def test_write_thread_count_forwarding(
         _, kwargs = spy_memory_storage_writer_init.call_args  # Capture kwargs
         assert kwargs["thread_count"] == expected_thread_count
 
+    @pytest.mark.parametrize("always_save_context", [True, False])
+    def test_loader_initialization_arguments(self, mocker, always_save_context):
+        """Tests that NeMoMLFlashpointCheckpointLoader is initialized with correct arguments."""
+        # Given
+        mocker.patch("ml_flashpoint.adapter.nemo.wrapper_util.ReplicationManager")
+        mocker.patch("ml_flashpoint.adapter.nemo.wrapper_util.wrap_trainer_checkpoint_io_with_mlflashpoint")
+        mock_loader_cls = mocker.patch("ml_flashpoint.adapter.nemo.wrapper_util.NeMoMLFlashpointCheckpointLoader")
+
+        trainer = mocker.MagicMock(spec=nl_trainer.Trainer)
+        flashpoint_base_container = "/tmp/test_container"
+        default_auto_resume = nl.AutoResume()
+
+        # When
+        wrap_trainer_and_auto_resume_with_mlflashpoint(
+            trainer,
+            flashpoint_base_container,
+            async_save=True,
+            default_auto_resume=default_auto_resume,
+            always_save_context=always_save_context,
+        )
+
+        # Then
+        mock_loader_cls.assert_called_once()
+        _, kwargs = mock_loader_cls.call_args
+        assert kwargs["recover_context"] == always_save_context
+
 
 class TestWrapTrainerCheckpointIOWithMLFlashpoint:
     """Tests for the wrap_trainer_checkpoint_io_with_mlflashpoint function."""

tests/core/test_checkpoint_loader.py

@@ -1725,8 +1725,9 @@ def test_get_checkpoint_objects_by_rank_optimization(self, loader, mocker, check
 
         # Mock all_gather to return some data for rank 1
         def side_effect_all_gather(out_list, in_obj):
-            out_list[0] = ["obj1", "obj2"]  # Rank 0 data
-            out_list[1] = ["obj3"]  # Rank 1 data (filename only)
+            out_list[0] = in_obj  # Rank 0 data (passed in)
+            # Rank 1 data
+            out_list[1] = [CheckpointObjectId(str(ckpt_dir / "obj3"))]
 
         mock_all_gather.side_effect = side_effect_all_gather
 
@@ -1735,11 +1736,13 @@ def side_effect_all_gather(out_list, in_obj):
         # Call method
         result = loader.get_checkpoint_objects_by_rank(container_id)
 
-        # Verify all_gather called with filenames only
+        # Verify all_gather called
         mock_all_gather.assert_called_once()
         args, _ = mock_all_gather.call_args
-        # args[1] is the input object (list of filenames for rank 0)
-        assert set(args[1]) == {"obj1", "obj2"}
+        # args[1] is the input object (list of CheckpointObjectIds for rank 0)
+        # Should be full paths now
+        expected_local = {str(ckpt_dir / "obj1"), str(ckpt_dir / "obj2")}
+        assert set(str(o.data) for o in args[1]) == expected_local
 
         # Verify result has full paths
         assert 0 in result
@@ -1913,9 +1916,10 @@ def test_get_checkpoint_objects_by_rank_sync(self, mocker, loader, tmp_path):
         (container_path / "file1").touch()
 
         def side_effect_all_gather(obj_list, local_obj):
-            obj_list[0] = ["file1"]
-            obj_list[1] = ["file1"]
-            obj_list[2] = ["file2"]
+            # Manually constructing full paths for the mock, mimicking what other ranks would send
+            obj_list[0] = [CheckpointObjectId(str(container_path / "file1"))]
+            obj_list[1] = [CheckpointObjectId(str(container_path / "file1"))]
+            obj_list[2] = [CheckpointObjectId(str(container_path / "file2"))]
             obj_list[3] = []
 
         mock_all_gather.side_effect = side_effect_all_gather