From 8d25cc2024e916686db9f6b493fb976863a4d4cd Mon Sep 17 00:00:00 2001 From: Jackie2049 Date: Tue, 28 Jul 2026 10:25:47 +0800 Subject: [PATCH] fix: rank=0 falsy-value in ParameterServer.__init__ MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `self._rank = rank or int(os.environ["RANK"])` evaluates `rank=0` as falsy, so a rank-0 process always falls back to the environment variable. If `RANK` is unset in that process, it raises `KeyError` instead of using the zero passed by the caller. Fix: use `rank if rank is not None else ...` — only fall back to the env when the caller explicitly passes `None`. Co-Authored-By: Claude Opus 4.8 --- checkpoint_engine/ps.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/checkpoint_engine/ps.py b/checkpoint_engine/ps.py index 1d8c5cf..c531c26 100644 --- a/checkpoint_engine/ps.py +++ b/checkpoint_engine/ps.py @@ -188,7 +188,7 @@ def __init__( Notice that if auto_pg is True, will destroy the process group after update. It is recommended to set auto_pg to True! mem_fraction: The proportion (as a fraction) of the current free device memory for allocation. """ - self._rank = rank or int(os.environ["RANK"]) + self._rank = rank if rank is not None else int(os.environ["RANK"]) self._world_size = world_size or int(os.environ["WORLD_SIZE"]) self.device_manager = DeviceManager() self._gpu_count = gpu_count or self.device_manager.device_module.device_count()