From bbb2d95256f8d8de271d41a700b5cbc623e86686 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Thu, 27 Aug 2026 13:39:38 +0800 Subject: [PATCH] fix: report true rank in logs via dist-aware helpers - get_rank/get_world_size fall back to RANK/WORLD_SIZE env instead of hardcoded 0/1, matching torchrun's env-before-init contract - log filter reuses the helpers so initialized groups show real ranks; local-spawn children previously logged rank=0/8 --- astrai/logging.py | 6 ++++-- astrai/parallel/setup.py | 6 ++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/astrai/logging.py b/astrai/logging.py index 668ab78..996d038 100644 --- a/astrai/logging.py +++ b/astrai/logging.py @@ -1,11 +1,13 @@ import logging import os +from astrai.parallel.setup import get_rank, get_world_size + class _DistributedContextFilter(logging.Filter): def filter(self, record: logging.LogRecord) -> bool: - record.rank = os.environ.get("RANK", "0") - record.world_size = os.environ.get("WORLD_SIZE", "1") + record.rank = str(get_rank()) + record.world_size = str(get_world_size()) return True diff --git a/astrai/parallel/setup.py b/astrai/parallel/setup.py index e79eacb..4709d75 100644 --- a/astrai/parallel/setup.py +++ b/astrai/parallel/setup.py @@ -30,15 +30,13 @@ def get_current_device(): def get_world_size() -> int: if dist.is_available() and dist.is_initialized(): return dist.get_world_size() - else: - return 1 + return int(os.environ.get("WORLD_SIZE", "1")) def get_rank() -> int: if dist.is_available() and dist.is_initialized(): return dist.get_rank() - else: - return 0 + return int(os.environ.get("RANK", "0")) @contextmanager