From 53a7149577a07c80032ee76ad9c4e50723273264 Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Thu, 20 Aug 2026 01:37:00 +0800
Subject: [PATCH] feat: add distributed rank to logs

---
 astrai/logging.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/astrai/logging.py b/astrai/logging.py
index 5ea2c46..668ab78 100644
--- a/astrai/logging.py
+++ b/astrai/logging.py
@@ -2,6 +2,13 @@ import logging
 import os
 
 
+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")
+        return True
+
+
 def setup_logging(level: str = "INFO"):
     """Attach a StreamHandler to the ``astrai`` logger (idempotent).
 
@@ -18,9 +25,10 @@ def setup_logging(level: str = "INFO"):
     level_name = os.environ.get("ASTR_LOG_LEVEL", level).upper()
     logger.setLevel(getattr(logging, level_name, logging.INFO))
     handler = logging.StreamHandler()
+    handler.addFilter(_DistributedContextFilter())
     handler.setFormatter(
         logging.Formatter(
-            "%(asctime)s | %(levelname)-7s | %(name)s | %(message)s",
+            "%(asctime)s | %(levelname)-8s | rank=%(rank)2s/%(world_size)-2s | %(name)-32s | %(message)s",
             datefmt="%Y-%m-%d %H:%M:%S",
         )
     )