Skip to content

Commit 089d185

Browse files
authored
feat: docs for force unregister (#63)
* feat: force unregister the shared memory pool * feat: tests added for force unregister shared memory pool * feat: add docs * misc
1 parent f69e116 commit 089d185

File tree

1 file changed

+6
-0
lines changed

1 file changed

+6
-0
lines changed

checkpoint_engine/ps.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -986,6 +986,8 @@ def register_checkpoint(
986986
use_shared_memory_pool: If True, uses a reusable shared pin memory pool instead of allocating new memory.
987987
Only one checkpoint can use the shared pool at a time. The pool's shape is fixed on first use and
988988
cannot accommodate checkpoints with different memory requirements.
989+
To free the actual memory of the shared pool or to modify its shape,
990+
please unregister the current user of the shared memory pool using `unregister_checkpoint` with `force=True`.
989991
"""
990992
try:
991993
if use_shared_memory_pool:
@@ -1031,6 +1033,10 @@ def unregister_checkpoint(self, checkpoint_name: str, force: bool = False) -> No
10311033
"""
10321034
Unregister a checkpoint from the parameter server. This function will also unregister the checkpoint
10331035
from p2p store if p2p store is initialized.
1036+
Args:
1037+
checkpoint_name: The name of the checkpoint.
1038+
force: This flag is designed for shared memory pool user. If True, the memory for shared memory pool itself will be freed.
1039+
If False, only the checkpoint name will be unregistered, and the shared memory pool will be kept for future use.
10341040
"""
10351041
if (
10361042
checkpoint_name not in self._memory_pool

0 commit comments

Comments
 (0)