feat: add distributed rank to logs

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",
         )
     )