refactor(core): Move context recovery logic strictly to the NeMo layer (#19)

- Extracts context recovery logic from
DefaultMLFlashpointCheckpointLoader

- Introduces _get_extra_local_objects and _get_extra_needed_objects to
DefaultMLFlashpointCheckpointLoader

- Updates NeMo wrapper to instantiate the new
NeMoMLFlashpointCheckpointLoader

---------

Co-authored-by: g-husam <husameldawi@google.com>

Author: Leahlijuan

src/ml_flashpoint/adapter/nemo/nemo_checkpoint_loader.py

@@ -0,0 +1,80 @@
+# 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 os
+from pathlib import Path
+from typing import List, Set
+
+from typing_extensions import override
+
+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_loader import DefaultMLFlashpointCheckpointLoader
+from ml_flashpoint.replication.replication_manager import ReplicationManager
+
+
+class NeMoMLFlashpointCheckpointLoader(DefaultMLFlashpointCheckpointLoader):
+    """
+    NeMo-specific implementation of the MLFlashpointCheckpointLoader interface.
+    """
+
+    def __init__(
+        self,
+        checkpoint_object_manager: CheckpointObjectManager,
+        replication_manager: ReplicationManager,
+        recover_context: bool = False,
+    ):
+        """Initializes the NeMoMLFlashpointCheckpointLoader.
+
+        Args:
+            checkpoint_object_manager: The checkpoint object manager to use for
+                reading data.
+            replication_manager: The replication manager to use for retrieving
+                missing checkpoint objects from peer nodes.
+            recover_context: Whether to recover the context directory if missing.
+        """
+        super().__init__(checkpoint_object_manager, replication_manager)
+        self._recover_context = recover_context
+
+    @override
+    def _get_extra_local_objects(self, container_path: Path) -> List[CheckpointObjectId]:
+        local_objects = []
+        if self._recover_context:
+            context_path = container_path / "context"
+            if context_path.is_dir():
+                for root, _, files in os.walk(context_path):
+                    for file in files:
+                        local_objects.append(CheckpointObjectId(str(Path(root) / file)))
+        return local_objects
+
+    @override
+    def _get_extra_needed_objects(
+        self,
+        checkpoint: CheckpointContainerId,
+        available_objects_by_rank: dict[int, List[CheckpointObjectId]],
+    ) -> Set[str]:
+        extra_needed = set()
+        if self._recover_context:
+            # We assume that if a rank has the context dir, the content in the dir is complete.
+            # We assume that these are the files needed by all the nodes.
+            context_path = Path(checkpoint.data) / "context"
+            for objs in available_objects_by_rank.values():
+                for obj in objs:
+                    try:
+                        if Path(str(obj.data)).is_relative_to(context_path):
+                            extra_needed.add(str(obj.data))
+                    except ValueError:
+                        # Path.is_relative_to raises ValueError if it's not relative to the path
+                        pass
+        return extra_needed

src/ml_flashpoint/adapter/nemo/wrapper_util.py

@@ -27,6 +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.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
@@ -72,7 +73,7 @@ def wrap_trainer_and_auto_resume_with_mlflashpoint(
     replication_manager = ReplicationManager()
     replication_manager.initialize(checkpoint_object_manager=ckpt_obj_manager)
 
-    ckpt_loader = DefaultMLFlashpointCheckpointLoader(
+    ckpt_loader = NeMoMLFlashpointCheckpointLoader(
         checkpoint_object_manager=ckpt_obj_manager,
         replication_manager=replication_manager,
         recover_context=always_save_context,

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, Tuple, TypeVar, cast
+from typing import IO, List, Optional, Set, Tuple, TypeVar, cast
 
 import torch
 import torch.distributed as dist
@@ -128,7 +128,6 @@ def __init__(
         self,
         checkpoint_object_manager: CheckpointObjectManager,
         replication_manager: ReplicationManager,
-        recover_context: bool = False,
     ):
         """Initializes the DefaultMLFlashpointCheckpointLoader.
 
@@ -137,11 +136,9 @@ def __init__(
                 reading data.
             replication_manager: The replication manager to use for retrieving
                 missing checkpoint objects from peer nodes.
-            recover_context: Whether to recover the context directory if missing.
         """
         self._checkpoint_object_manager = checkpoint_object_manager
         self._replication_manager = replication_manager
-        self._recover_context = recover_context
         # Cache for available objects: CheckpointContainerId -> dict[object_path, list[rank]]
         self._available_objects_cache: dict[CheckpointContainerId, dict[str, List[int]]] = {}
 
@@ -417,6 +414,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.
@@ -450,14 +449,7 @@ def _compute_retrieval_plan(
             str(CheckpointObjectId.from_container(checkpoint, default_metadata_object_name()))
         )
 
-        if self._recover_context:
-            # We assume that if a rank has the context dir, the content in the dir is complete.
-            # We assume that are the files needed by all the nodes.
-            context_path = Path(checkpoint.data) / "context"
-            for objs in available_objects_by_rank.values():
-                for obj in objs:
-                    if Path(obj.data).parent == context_path:
-                        objects_needed_by_local_rank_0.add(str(obj.data))
+        objects_needed_by_local_rank_0.update(self._get_extra_needed_objects(checkpoint, available_objects_by_rank))
 
         world_size = dist.get_world_size()
         num_nodes = get_num_of_nodes()
@@ -581,25 +573,21 @@ def get_checkpoint_objects_by_rank(
             checkpoint_container_id: The ID of the checkpoint container to inspect.
 
         Returns:
-            A dictionary mapping each node's local rank 0 to a list of
-            `CheckpointObjectId`s available on that node.
+            A dictionary mapping each rank to a list of
+            `CheckpointObjectId`s available on that rank.
         """
         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))
 
-            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.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)
@@ -609,11 +597,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] = []
@@ -675,3 +660,37 @@ def retrieve_checkpoint(
         dist.all_gather_object(all_success_list, all_success)
         _LOGGER.debug("All success list: '%s'", all_success_list)
         return all(all_success_list)
+
+    def _get_extra_local_objects(self, container_path: Path) -> List[CheckpointObjectId]:
+        """Hook for subclasses to provide extra local objects that are available on the current rank,
+        which may be needed by other ranks. This would be called on every rank.
+
+        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.
+        """
+        return []
+
+    def _get_extra_needed_objects(
+        self,
+        checkpoint: CheckpointContainerId,
+        available_objects_by_rank: dict[int, List[CheckpointObjectId]],
+    ) -> Set[str]:
+        """Hook for subclasses to provide extra needed objects for each node.
+
+        The objects returned by this method are considered necessary for the first rank
+        (local rank 0) of every node.
+
+        This can leverage `available_objects_by_rank` to determine the set of additional objects
+        needed.
+
+        This should always be implemented alongside `_get_extra_local_objects`.
+
+        Returns:
+            Set of extra needed objects on each node (specifically local rank 0).
+        """
+        return set()

tests/adapter/nemo/test_nemo_checkpoint_loader.py

@@ -0,0 +1,162 @@
+# 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.
+
+from pathlib import Path
+from unittest.mock import MagicMock
+
+import pytest
+
+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
+
+
+class TestNeMoCheckpointLoaderContext:
+    @pytest.fixture
+    def loader(self):
+        ckpt_manager = CheckpointObjectManager()
+        repl_manager = MagicMock(spec=ReplicationManager)
+        return NeMoMLFlashpointCheckpointLoader(
+            checkpoint_object_manager=ckpt_manager, replication_manager=repl_manager, 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)
+
+        container_path = "/tmp/ckpt/step-1"
+        container_id = CheckpointContainerId(container_path)
+
+        # Mock fs
+        # base dir has "context" directory
+        # context dir has "file1.txt"
+        def mock_listdir(path):
+            if str(path) == container_path:
+                return ["context", "other_file"]
+            if str(path) == str(Path(container_path) / "context"):
+                return ["file1.txt", "file2.txt"]
+            return []
+
+        def mock_walk(path):
+            if str(path) == str(Path(container_path) / "context"):
+                # Simulation of:
+                # context/
+                #   file1.txt
+                #   file2.txt
+                #   subdir/
+                #     file3.txt
+                # Yields: (root, dirs, files)
+                yield (str(path), ["subdir"], ["file1.txt", "file2.txt"])
+                yield (str(Path(path) / "subdir"), [], ["file3.txt"])
+            return []
+
+        def mock_isdir(path):
+            if str(path) == container_path:
+                return True
+            if str(path) == str(Path(container_path) / "context"):
+                return True
+            return False
+
+        mocker.patch("os.walk", side_effect=mock_walk)
+        mocker.patch("os.listdir", side_effect=mock_listdir)
+        mocker.patch("pathlib.Path.is_dir", new=mock_isdir)
+
+        result = loader.get_checkpoint_objects_by_rank(container_id)
+
+        assert 0 in result
+        objs = result[0]
+
+        paths = [str(o.data) for o in objs]
+        expected_context_file1 = str(Path(container_path) / "context" / "file1.txt")
+        expected_context_file2 = str(Path(container_path) / "context" / "file2.txt")
+        expected_nested_file3 = str(Path(container_path) / "context" / "subdir" / "file3.txt")
+        assert expected_context_file1 in paths
+        assert expected_context_file2 in paths
+        assert expected_nested_file3 in paths
+
+    def test_compute_retrieval_plan_includes_context_optimized(self, loader, mocker):
+        """
+        Test that _compute_retrieval_plan includes context files ONLY for local rank 0 on each node.
+        Scenario:
+          - World Size: 4
+          - Nodes: 2 (Ranks 0,1 on Node 0; Ranks 2,3 on Node 1)
+          - Rank 0 has context files.
+          - Rank 2 needs context files (different node).
+          - Rank 1, 3 do NOT need context retrieval (same node as 0, 2).
+        """
+        checkpoint = CheckpointContainerId("/tmp/ckpt/step-1")
+
+        # Mock metadata read (empty storage data)
+        mock_metadata = MagicMock()
+        mock_metadata.storage_data = {}
+        mocker.patch.object(loader, "read_metadata", return_value=mock_metadata)
+
+        mocker.patch("torch.distributed.get_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)
+
+        ctx_file = str(Path(checkpoint.data) / "context" / "file1.txt")
+        nested_ctx_file = str(Path(checkpoint.data) / "context" / "subdir" / "file3.txt")
+        common_pt = str(Path(checkpoint.data) / "common.pt")
+        metadata_file = str(Path(checkpoint.data) / ".metadata")
+
+        # Available objects:
+        # Node 0 (Rank 0,1) has everything (Context + Nested + Common + Metadata)
+        # Node 1 (Rank 2,3) has nothing
+        available_objects = {
+            0: [
+                CheckpointObjectId(ctx_file),
+                CheckpointObjectId(nested_ctx_file),
+                CheckpointObjectId(common_pt),
+                CheckpointObjectId(metadata_file),
+            ],
+            1: [
+                CheckpointObjectId(ctx_file),
+                CheckpointObjectId(nested_ctx_file),
+                CheckpointObjectId(common_pt),
+                CheckpointObjectId(metadata_file),
+            ],
+            2: [],
+            3: [],
+        }
+
+        # Execute
+        plan = loader._compute_retrieval_plan(checkpoint, available_objects)
+
+        assert plan is not None
+
+        # Node 0: Already has files, no retrieval needed (or plan[0] is empty)
+        assert 0 not in plan or not plan[0]
+        assert 1 not in plan or not plan[1]
+
+        # Rank 2: Local rank 0 on Node 1. Needs Context + Common + Metadata.
+        assert 2 in plan
+        retrieved_objs_2 = [path for src, path in plan[2]]
+        assert ctx_file in retrieved_objs_2
+        assert common_pt in retrieved_objs_2
+        assert metadata_file in retrieved_objs_2
+
+        # Verify nested file
+        nested_ctx_file = str(Path(checkpoint.data) / "context" / "subdir" / "file3.txt")
+        assert nested_ctx_file in retrieved_objs_2
+
+        # Rank 3: Local rank 1 on Node 1. Shared FS with Rank 2. Should NOT retrieve.
+        assert 3 not in plan or not plan[3]