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

@@ -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
 
 

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