Tiny fix single-gpu dumper and add tests for dumper (#16285)

This commit is contained in:
fzyzcjy
2026-01-02 13:30:09 +08:00
committed by GitHub
parent 7254986342
commit 749736ba95
4 changed files with 415 additions and 4 deletions

View File

@@ -250,8 +250,8 @@ class _DumperRpcHandler:
def _create_zmq_rpc_handles(handler, base_port: int) -> Optional[List["_ZmqRpcHandle"]]:
import zmq
rank = dist.get_rank()
world_size = dist.get_world_size()
rank = _get_rank()
world_size = dist.get_world_size() if dist.is_initialized() else 1
port = base_port + rank
local_addr = f"tcp://{_get_local_ip_by_remote()}:{port}"
@@ -274,8 +274,11 @@ def _create_zmq_rpc_handles(handler, base_port: int) -> Optional[List["_ZmqRpcHa
thread.start()
print(f"[Dumper.ZmqRpc] rank={rank} server started at {local_addr}")
all_addresses = [None] * world_size
dist.all_gather_object(all_addresses, local_addr)
if dist.is_initialized():
all_addresses = [None] * world_size
dist.all_gather_object(all_addresses, local_addr)
else:
all_addresses = [local_addr]
print(f"[Dumper.ZmqRpc] rank={rank} all_addresses={all_addresses}")
if rank == 0: