Recursively go through context/ dir.

Author: Leahlijuan

src/ml_flashpoint/adapter/nemo/nemo_checkpoint_loader.py

@@ -54,8 +54,9 @@ def _get_extra_local_objects(self, container_path: Path) -> List[CheckpointObjec
         if self._recover_context:
             context_path = container_path / "context"
             if context_path.is_dir():
-                for entry in os.listdir(context_path):
-                    local_objects.append(CheckpointObjectId(str(context_path / entry)))
+                for root, _, files in os.walk(context_path):
+                    for file in files:
+                        local_objects.append(CheckpointObjectId(str(Path(root) / file)))
         return local_objects
 
     @override
@@ -72,6 +73,10 @@ def _get_extra_needed_objects(
             context_path = Path(checkpoint.data) / "context"
             for objs in available_objects_by_rank.values():
                 for obj in objs:
-                    if Path(str(obj.data)).parent == context_path:
-                        extra_needed.add(str(obj.data))
+                    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

tests/adapter/nemo/test_nemo_checkpoint_loader.py

@@ -55,28 +55,42 @@ def mock_listdir(path):
                 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]
-        # Should contain: .../context/file1.txt, .../context/file2.txt, .../context, .../other_file
-        # Note: 'context' is from top-level listdir. 'context/file*.txt' is from recursive check.
 
         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):
         """
@@ -100,15 +114,26 @@ def test_compute_retrieval_plan_includes_context_optimized(self, loader, mocker)
         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 + Common + Metadata)
+        # 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(common_pt), CheckpointObjectId(metadata_file)],
-            1: [CheckpointObjectId(ctx_file), CheckpointObjectId(common_pt), CheckpointObjectId(metadata_file)],
+            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: [],
         }
@@ -129,5 +154,9 @@ def test_compute_retrieval_plan_includes_context_optimized(self, loader, mocker)
         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]