feat(replication): support single-node execution by disabling replication (#92)

Currently, ML Flashpoint requires an even number of nodes due to the
pairwise replication strategy. This prevents the library from running in
single-node environments, which is often required in some testing
scenarios.

This commit adds support for single-node execution by safely bypassing
replication logic while keeping the core save/load functionality intact.

Author: kkkapu

src/ml_flashpoint/replication/replication_manager.py

@@ -108,8 +108,12 @@ def __init__(self, replication_service_addresses: List[str], processes_per_node:
             )
 
         self._num_nodes = utils.get_num_of_nodes()
-        if self._num_nodes % 2 != 0:
-            raise ValueError(f"The total number of nodes ({self._num_nodes}) must be even.")
+        # Allow 1-node execution by bypassing the even-number node check.
+        # Pairwise replication requires an even number of nodes, but we allow exactly 1 node
+        # to run without any replication.
+        if self._num_nodes > 1 and self._num_nodes % 2 != 0:
+            raise ValueError(f"The total number of nodes ({self._num_nodes}) must be even for pairwise replication.")
+        self._disable_replication = self._num_nodes == 1
 
         if self._num_nodes != self._world_size // self._processes_per_node:
             raise ValueError(
@@ -119,6 +123,9 @@ def __init__(self, replication_service_addresses: List[str], processes_per_node:
 
     @override
     def get_destination_addresses(self, global_rank: int) -> List[str]:
+        if self._disable_replication:
+            return []
+
         if not 0 <= global_rank < self._world_size:
             raise ValueError(f"global_rank {global_rank} is out of valid range [0, {self._world_size - 1}].")
 
@@ -441,6 +448,7 @@ def sync_bulk_retrieve(
         Returns:
             `True` if the bulk retrieval was successful, `False` otherwise.
         """
+
         if self._retry_config is None:
             _LOGGER.error("ReplicationManager is not initialized. Cannot retrieve.")
             return False

tests/core/test_checkpoint_loader.py

@@ -998,6 +998,71 @@ def test_no_retrieval_needed(self, loader, mocker):
         assert result == ckpt_id
         mock_retrieve.assert_not_called()
 
+    def test_single_node_shared_storage_no_retrieval_needed(self, loader, mocker):
+        """
+        Tests that in a 1-node setup (e.g., fallback to NFS or local shared storage),
+        the loader correctly computes an empty retrieval plan. It does not attempt to
+        fetch from non-existent peer nodes, allowing the process to execute correctly
+        without throwing an error.
+        """
+        # Given: Simulate a single-node environment (1 node, 2 ranks)
+        self.mock_dist_get_rank.return_value = 0
+        self.mock_get_num_nodes.return_value = 1
+        self.mock_dist_get_world_size.return_value = 2
+
+        base_container = CheckpointContainerId("/tmp/checkpoints")
+        ckpt_id = CheckpointContainerId.create_child(base_container, "step-100_ckpt")
+        mocker.patch.object(loader, "get_candidate_checkpoints", return_value=[ckpt_id])
+
+        # Mock metadata indicating what needs to be loaded
+        mock_metadata = mocker.MagicMock()
+        mock_metadata.storage_data = {
+            0: _StorageInfo(relative_path="/src0/obj", offset=0, length=100),
+            1: _StorageInfo(relative_path="/src1/obj", offset=0, length=100),
+        }
+        mocker.patch.object(loader, "read_metadata", return_value=mock_metadata)
+
+        # Mock available objects: Rank 0 sees src0, Rank 1 sees src1.
+        # However, since they are on the same node (num_nodes=1), the underlying system
+        # can see all files via NFS/local disk.
+        mocker.patch.object(
+            loader,
+            "get_checkpoint_objects_by_rank",
+            return_value={
+                0: [
+                    CheckpointObjectId("/tmp/checkpoints/step-100_ckpt/src0/obj"),
+                    CheckpointObjectId("/tmp/checkpoints/step-100_ckpt/.metadata"),
+                    CheckpointObjectId("/tmp/checkpoints/step-100_ckpt/common.pt"),
+                ],
+                1: [
+                    CheckpointObjectId("/tmp/checkpoints/step-100_ckpt/src1/obj"),
+                    CheckpointObjectId("/tmp/checkpoints/step-100_ckpt/.metadata"),
+                    CheckpointObjectId("/tmp/checkpoints/step-100_ckpt/common.pt"),
+                ],
+            },
+        )
+
+        mock_retrieve = mocker.patch.object(loader, "retrieve_checkpoint", return_value=True)
+
+        # When: Execute the core logic to get the latest checkpoint
+        result = loader.get_latest_complete_checkpoint(base_container)
+
+        # Then:
+        # Verify that the method executes correctly instead of crashing with an error.
+        assert result == ckpt_id
+
+        # Verify that the generated retrieval plan is empty (no network fetch to non-existent peers).
+        args, _ = self.mock_dist_broadcast_object_list.call_args
+        plan_container = args[0]
+        plan = plan_container[0]
+
+        # Assert that the network retries wval plan for all ranks is empty (perfect fallback to NFS/Local read).
+        assert not plan.get(0)
+        assert not plan.get(1)
+
+        # Verify that the subsequent flow skips retrieval entirely since it's locally available.
+        mock_retrieve.assert_not_called()
+
     def test_non_rank0_success(self, loader, mocker):
         """Tests successful retrieval flow on non-Rank 0."""
         # Given

tests/replication/test_replication_manager.py

@@ -16,7 +16,11 @@
 
 import pytest
 
-from ml_flashpoint.replication.replication_manager import ReplicationManager, ReplicationRetryConfig
+from ml_flashpoint.replication.replication_manager import (
+    PairwiseReplicationStrategy,
+    ReplicationManager,
+    ReplicationRetryConfig,
+)
 
 
 @pytest.fixture
@@ -338,3 +342,59 @@ def test_sync_bulk_retrieve_invalid_rank(replication_manager):
 
     # Then
     assert result is False
+
+
+def test_pairwise_strategy_single_node_initialization(mocker):
+    """Tests that PairwiseReplicationStrategy successfully initializes for a single node without raising an error."""
+    # Given
+    mocker.patch("ml_flashpoint.core.utils.get_num_of_nodes", return_value=1)
+    # Simulate a single node with 2 processes (GPUs)
+    addresses = ["127.0.0.1:8000", "127.0.0.1:8001"]
+
+    # When
+    strategy = PairwiseReplicationStrategy(replication_service_addresses=addresses, processes_per_node=2)
+
+    # Then
+    assert getattr(strategy, "_disable_replication", False) is True
+
+
+def test_pairwise_strategy_single_node_get_destination(mocker):
+    """Tests that get_destination_addresses returns an empty list when running on a single node."""
+    # Given
+    mocker.patch("ml_flashpoint.core.utils.get_num_of_nodes", return_value=1)
+    addresses = ["127.0.0.1:8000"]
+    strategy = PairwiseReplicationStrategy(replication_service_addresses=addresses, processes_per_node=1)
+
+    # When
+    destinations = strategy.get_destination_addresses(global_rank=0)
+
+    # Then
+    assert destinations == []
+
+
+def test_async_replicate_single_node_skips(replication_manager, mocker):
+    """Tests that async_replicate does nothing and returns empty futures in a single-node environment."""
+    # Given
+    mocker.patch("ml_flashpoint.core.utils.get_num_of_nodes", return_value=1)
+    addresses = ["127.0.0.1:8000"]
+    # Initialize the strategy with 1 node
+    strategy = PairwiseReplicationStrategy(replication_service_addresses=addresses, processes_per_node=1)
+    replication_manager._repl_strategy = strategy
+
+    mocker.patch("torch.distributed.get_rank", return_value=0)
+
+    buffer_object = mocker.MagicMock()
+    buffer_object.get_id.return_value = "test_single_node_obj"
+    buffer_io = mocker.MagicMock(buffer_obj=buffer_object)
+
+    # When
+    result_futures = replication_manager.async_replicate(buffer_io)
+
+    # Then
+    assert result_futures == []
+    # Ensure transfer service is NOT called
+    replication_manager._transfer_service.async_put.assert_not_called()
+    # Ensure the buffer is closed properly
+    replication_manager._checkpoint_object_manager.close_buffer.assert_called_once_with(
+        buffer_io, skip_close_if_symlink=True
+    )