refactor(core): abstract torch.distributed APIs in CheckpointLoader

Add injectable callable parameters to DefaultMLFlashpointCheckpointLoader
for get_rank, get_node_local_rank, broadcast_object_list, all_gather_object,
and get_world_size, mirroring the pattern already used in
DefaultMLFlashpointCheckpointSaver.

All parameters default to the corresponding torch.distributed functions,
preserving backwards compatibility. This makes the loader easier to test
via dependency injection and allows swapping implementations without
subclassing or monkey-patching torch.distributed.

Closes #30

Author: ronaldw07

src/ml_flashpoint/adapter/nemo/nemo_checkpoint_loader.py

@@ -14,7 +14,7 @@
 
 import os
 from pathlib import Path
-from typing import List, Set
+from typing import Callable, List, Set
 
 from typing_extensions import override
 
@@ -33,6 +33,12 @@ def __init__(
         self,
         checkpoint_object_manager: CheckpointObjectManager,
         replication_manager: ReplicationManager,
+        *,
+        global_rank_getter: Callable[[], int],
+        local_rank_getter: Callable[[], int],
+        broadcast_object_list_func: Callable[..., None],
+        all_gather_object_func: Callable[..., None],
+        world_size_getter: Callable[[], int],
         recover_context: bool = False,
     ):
         """Initializes the NeMoMLFlashpointCheckpointLoader.
@@ -42,9 +48,24 @@ def __init__(
                 reading data.
             replication_manager: The replication manager to use for retrieving
                 missing checkpoint objects from peer nodes.
+            global_rank_getter: A callable that returns the global rank.
+            local_rank_getter: A callable that returns the node-local rank.
+            broadcast_object_list_func: A callable with the same signature as
+                ``torch.distributed.broadcast_object_list``.
+            all_gather_object_func: A callable with the same signature as
+                ``torch.distributed.all_gather_object``.
+            world_size_getter: A callable that returns the world size.
             recover_context: Whether to recover the context directory if missing.
         """
-        super().__init__(checkpoint_object_manager, replication_manager)
+        super().__init__(
+            checkpoint_object_manager,
+            replication_manager,
+            global_rank_getter=global_rank_getter,
+            local_rank_getter=local_rank_getter,
+            broadcast_object_list_func=broadcast_object_list_func,
+            all_gather_object_func=all_gather_object_func,
+            world_size_getter=world_size_getter,
+        )
         self._recover_context = recover_context
 
     @override

src/ml_flashpoint/core/checkpoint_loader.py

@@ -22,7 +22,7 @@
 import struct
 from collections import defaultdict
 from pathlib import Path
-from typing import IO, List, Optional, Set, Tuple, TypeVar, cast
+from typing import IO, Callable, List, Optional, Set, Tuple, TypeVar, cast
 
 import torch
 import torch.distributed as dist
@@ -128,6 +128,12 @@ def __init__(
         self,
         checkpoint_object_manager: CheckpointObjectManager,
         replication_manager: ReplicationManager,
+        *,
+        global_rank_getter: Callable[[], int],
+        local_rank_getter: Callable[[], int],
+        broadcast_object_list_func: Callable[..., None],
+        all_gather_object_func: Callable[..., None],
+        world_size_getter: Callable[[], int],
     ):
         """Initializes the DefaultMLFlashpointCheckpointLoader.
 
@@ -136,9 +142,21 @@ def __init__(
                 reading data.
             replication_manager: The replication manager to use for retrieving
                 missing checkpoint objects from peer nodes.
+            global_rank_getter: A callable that returns the global rank.
+            local_rank_getter: A callable that returns the node-local rank.
+            broadcast_object_list_func: A callable with the same signature as
+                ``torch.distributed.broadcast_object_list``.
+            all_gather_object_func: A callable with the same signature as
+                ``torch.distributed.all_gather_object``.
+            world_size_getter: A callable that returns the world size.
         """
         self._checkpoint_object_manager = checkpoint_object_manager
         self._replication_manager = replication_manager
+        self._global_rank_getter = global_rank_getter
+        self._local_rank_getter = local_rank_getter
+        self._broadcast_object_list_func = broadcast_object_list_func
+        self._all_gather_object_func = all_gather_object_func
+        self._world_size_getter = world_size_getter
         # Cache for available objects: CheckpointContainerId -> dict[object_path, list[rank]]
         self._available_objects_cache: dict[CheckpointContainerId, dict[str, List[int]]] = {}
 
@@ -337,8 +355,7 @@ def get_latest_complete_checkpoint(
             else continue to the next candidate checkpoint
             - return the checkpoint container id of the latest complete checkpoint
         """
-        # TODO: use global_rank_getter and local_rank_getter.
-        rank = dist.get_rank()
+        rank = self._global_rank_getter()
         _LOGGER.debug(
             "Rank %s: Getting latest complete checkpoint for '%s'",
             rank,
@@ -382,7 +399,7 @@ def get_latest_complete_checkpoint(
                 retrieval_plan = self._compute_retrieval_plan(checkpoint, available_objects_by_rank)
             # Broadcast the retrieval plan to all ranks.
             plan_container = [retrieval_plan]
-            dist.broadcast_object_list(plan_container, src=planner_rank)
+            self._broadcast_object_list_func(plan_container, src=planner_rank)
             retrieval_plan = plan_container[0]
 
             if retrieval_plan is None:
@@ -451,7 +468,7 @@ def _compute_retrieval_plan(
 
         objects_needed_by_local_rank_0.update(self._get_extra_needed_objects(checkpoint, available_objects_by_rank))
 
-        world_size = dist.get_world_size()
+        world_size = self._world_size_getter()
         num_nodes = get_num_of_nodes()
         ranks_per_node = world_size // num_nodes
 
@@ -507,8 +524,8 @@ def get_candidate_checkpoints(
 
         # Scan locally only on the first rank of each node
         base_path = Path(checkpoint_base_container.data)
-        rank = dist.get_rank()
-        local_rank = dist.get_node_local_rank()
+        rank = self._global_rank_getter()
+        local_rank = self._local_rank_getter()
 
         local_candidate_ckpt_ids = []
 
@@ -532,8 +549,8 @@ def get_candidate_checkpoints(
         else:
             _LOGGER.debug("Rank %s: Base path '%s' is not a directory or does not exist.", rank, base_path)
 
-        all_checkpoint_container_path_lists = [None for _ in range(dist.get_world_size())]
-        dist.all_gather_object(all_checkpoint_container_path_lists, local_candidate_ckpt_ids)
+        all_checkpoint_container_path_lists = [None for _ in range(self._world_size_getter())]
+        self._all_gather_object_func(all_checkpoint_container_path_lists, local_candidate_ckpt_ids)
         _LOGGER.debug(
             "Rank %s: Gathered checkpoint container paths from all ranks: '%s'",
             rank,
@@ -589,8 +606,8 @@ def get_checkpoint_objects_by_rank(
 
             local_objects.extend(self._get_extra_local_objects(container_path))
 
-        all_objects_by_rank_paths = [None for _ in range(dist.get_world_size())]
-        dist.all_gather_object(all_objects_by_rank_paths, local_objects)
+        all_objects_by_rank_paths = [None for _ in range(self._world_size_getter())]
+        self._all_gather_object_func(all_objects_by_rank_paths, local_objects)
 
         result = {}
         object_locations = defaultdict(list)
@@ -620,7 +637,7 @@ def retrieve_checkpoint(
                             If empty for this rank, no retrieval is needed.
         """
 
-        rank = dist.get_rank()
+        rank = self._global_rank_getter()
         all_success = True
 
         # Only proceed with retrieval if we have items to retrieve
@@ -656,8 +673,8 @@ def retrieve_checkpoint(
 
         # Gather success status from all ranks
         _LOGGER.debug("Gathering success status from all ranks")
-        all_success_list = [None for _ in range(dist.get_world_size())]
-        dist.all_gather_object(all_success_list, all_success)
+        all_success_list = [None for _ in range(self._world_size_getter())]
+        self._all_gather_object_func(all_success_list, all_success)
         _LOGGER.debug("All success list: '%s'", all_success_list)
         return all(all_success_list)
 

tests/adapter/nemo/test_nemo_checkpoint_loader.py

@@ -25,22 +25,32 @@
 
 class TestNeMoCheckpointLoaderContext:
     @pytest.fixture
-    def loader(self):
+    def _setup_mocks(self, mocker):
+        self.mock_global_rank = MagicMock(return_value=0)
+        self.mock_local_rank = MagicMock(return_value=0)
+        self.mock_world_size = MagicMock(return_value=1)
+        self.mock_all_gather = MagicMock()
+        self.mock_broadcast = MagicMock()
+
+    @pytest.fixture
+    def loader(self, mocker, _setup_mocks):
         ckpt_manager = CheckpointObjectManager()
         repl_manager = MagicMock(spec=ReplicationManager)
         return NeMoMLFlashpointCheckpointLoader(
-            checkpoint_object_manager=ckpt_manager, replication_manager=repl_manager, recover_context=True
+            checkpoint_object_manager=ckpt_manager,
+            replication_manager=repl_manager,
+            global_rank_getter=self.mock_global_rank,
+            local_rank_getter=self.mock_local_rank,
+            broadcast_object_list_func=self.mock_broadcast,
+            all_gather_object_func=self.mock_all_gather,
+            world_size_getter=self.mock_world_size,
+            recover_context=True,
         )
 
     def test_get_checkpoint_objects_by_rank_finds_context(self, loader, mocker):
         """Test that get_checkpoint_objects_by_rank finds files in context/ dir when recover_context=True."""
-        mocker.patch("torch.distributed.get_world_size", return_value=1)
-        mocker.patch(
-            "torch.distributed.all_gather_object",
-            side_effect=lambda obj_list, local_obj: obj_list.__setitem__(0, local_obj),
-        )
-        # Mock get_node_local_rank to avoid external dependency issues if called
-        mocker.patch("torch.distributed.get_node_local_rank", return_value=0)
+        self.mock_world_size.return_value = 1
+        self.mock_all_gather.side_effect = lambda obj_list, local_obj: obj_list.__setitem__(0, local_obj)
 
         container_path = "/tmp/ckpt/step-1"
         container_id = CheckpointContainerId(container_path)
@@ -109,9 +119,9 @@ def test_compute_retrieval_plan_includes_context_optimized(self, loader, mocker)
         mock_metadata.storage_data = {}
         mocker.patch.object(loader, "read_metadata", return_value=mock_metadata)
 
-        mocker.patch("torch.distributed.get_world_size", return_value=4)
+        self.mock_world_size.return_value = 4
         mocker.patch("ml_flashpoint.core.checkpoint_loader.get_num_of_nodes", return_value=2)
-        mocker.patch("torch.distributed.get_rank", return_value=0)
+        self.mock_global_rank.return_value = 0
 
         ctx_file = str(Path(checkpoint.data) / "context" / "file1.txt")
         nested_ctx_file = str(Path(checkpoint.data) / "context" / "subdir" / "file3.txt")