fix: preserve registered checkpoints when duplicate registration is rejected

checkpoint_engine/ps.py

@@ -335,16 +335,21 @@ def register_checkpoint(
                 f"[rank{self._rank}] Only cuda devices support in-place pin memory, set use_inplace_pin_memory to False"
             )
             use_inplace_pin_memory = False
+        if use_shared_memory_pool:
+            assert self._current_shared_memory_pool_user == "", (
+                f"cannot register checkpoint {checkpoint_name} to shared memory pool, "
+                f"since checkpoint {self._current_shared_memory_pool_user} is already using shared memory pool. "
+                f"This registration may cause unexpected conflicts."
+            )
+        else:
+            assert checkpoint_name not in self._memory_pool, (
+                f"checkpoint {checkpoint_name} already registered"
+            )
         try:
             if use_shared_memory_pool:
                 logger.info(
                     f"[rank{self._rank}] checkpoint {checkpoint_name} use shared memory pool"
                 )
-                assert self._current_shared_memory_pool_user == "", (
-                    f"cannot register checkpoint {checkpoint_name} to shared memory pool, "
-                    f"since checkpoint {self._current_shared_memory_pool_user} is already using shared memory pool. "
-                    f"This registration may cause unexpected conflicts."
-                )
                 # Since we set the uninitialized shared memory pool to empty list,
                 # we can check whether this is the first time to use shared memory pool
                 _is_first_time = not self._memory_pool[self.shared_memory_pool_name]
@@ -359,9 +364,6 @@ def register_checkpoint(
                 if self._p2p_store is not None and _is_first_time:
                     self._register_parameters_to_p2p_store(checkpoint_name)
             else:
-                assert checkpoint_name not in self._memory_pool, (
-                    f"checkpoint {checkpoint_name} already registered"
-                )
                 self._memory_pool[checkpoint_name] = _register_checkpoint(
                     files=files or [],
                     named_tensors=named_tensors or {},

tests/test_registration_rejection.py

@@ -0,0 +1,32 @@
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from checkpoint_engine.data_types import MemoryBuffer
+from checkpoint_engine.ps import ParameterServer
+
+
+@pytest.mark.parametrize("shared", [False, True])
+def test_duplicate_registration_keeps_existing_checkpoint(
+    monkeypatch: pytest.MonkeyPatch, shared: bool
+) -> None:
+    ps = ParameterServer.__new__(ParameterServer)
+    ps._rank = 0
+    ps._p2p_store = None
+    ps._current_shared_memory_pool_user = ""
+    ps._memory_pool = {ps.shared_memory_pool_name: []}
+    ps.device_manager = MagicMock()
+    pool = [MemoryBuffer(buffer=torch.arange(16, dtype=torch.uint8), size=16, metas=[])]
+    register = MagicMock(return_value=pool)
+    monkeypatch.setattr("checkpoint_engine.ps._register_checkpoint", register)
+
+    ps.register_checkpoint("weights", use_shared_memory_pool=shared)
+    with pytest.raises(AssertionError, match=r"already (registered|using shared memory pool)"):
+        ps.register_checkpoint("weights", use_shared_memory_pool=shared)
+
+    assert ps._get_memory_pool("weights") is pool
+    assert torch.equal(pool[0].buffer, torch.arange(16, dtype=torch.uint8))
+    assert ps._current_shared_memory_pool_user == ("weights" if shared else "")
+    register.assert_called_once()
+    ps.device_manager.host_empty_cache.assert_not_called()