Update Wall-X to 1.1.0 (#104)

wall_x/utils/logger.py

@@ -0,0 +1,95 @@
+"""Rank-aware text logger for distributed training."""
+
+from __future__ import annotations
+
+import logging
+import os
+import sys
+from pathlib import Path
+from typing import Optional
+
+from torch.distributed import get_rank, is_initialized
+
+
+class DistributedLogger:
+    def __init__(
+        self,
+        name: str = "wallx",
+        save_path: Optional[str] = None,
+        level: int = logging.INFO,
+    ):
+        if is_initialized():
+            self._rank = get_rank()
+        else:
+            self._rank = int(os.environ.get("RANK", 0))
+            logging.warning(
+                "DistributedLogger created before init_process_group; "
+                "falling back to RANK env var (rank=%d).",
+                self._rank,
+            )
+        logger = logging.getLogger(f"{name}.rank{self._rank}")
+        logger.setLevel(level)
+        logger.propagate = False
+        for h in logger.handlers[:]:
+            h.close()
+            logger.removeHandler(h)
+
+        fmt = logging.Formatter(
+            f"%(asctime)s - [rank{self._rank}] - %(levelname)s - %(message)s"
+        )
+
+        # All ranks write a file log (if a save_path was given).
+        if save_path:
+            log_dir = Path(save_path) / "logs"
+            log_dir.mkdir(parents=True, exist_ok=True)
+            file_handler = logging.FileHandler(log_dir / f"rank_{self._rank}.log")
+            file_handler.setFormatter(fmt)
+            logger.addHandler(file_handler)
+
+        # Only rank 0 writes to stdout.
+        if self._rank == 0:
+            stream_handler = logging.StreamHandler(sys.stdout)
+            stream_handler.setFormatter(fmt)
+            logger.addHandler(stream_handler)
+
+        self._logger = logger
+
+    # Thin pass-throughs - callers use standard logging verbs.
+    def info(self, msg, *args, **kwargs):
+        self._logger.info(msg, *args, **kwargs)
+
+    def warning(self, msg, *args, **kwargs):
+        self._logger.warning(msg, *args, **kwargs)
+
+    def error(self, msg, *args, **kwargs):
+        self._logger.error(msg, *args, **kwargs)
+
+    def debug(self, msg, *args, **kwargs):
+        self._logger.debug(msg, *args, **kwargs)
+
+    @property
+    def rank(self) -> int:
+        return self._rank
+
+    # --- backward-compatible shim ----------------------------------------
+    # Some legacy call sites still invoke `.log(msg, level=...,
+    # main_process_only=...)`. Forward those to the standard logger so we
+    # don't break them during the migration; new code should use .info /
+    # .warning / .error / .debug directly.
+    def log(
+        self,
+        message,
+        level: int = logging.INFO,
+        main_process_only: bool = False,
+    ):
+        if main_process_only and self._rank != 0:
+            return
+        self._logger.log(level, message)
+
+    # Pytorch's ``accelerate``-style fallback for code that still passes an
+    # accelerator object to the old constructor; accept and ignore it.
+    # Older codepaths can be migrated incrementally.
+    @classmethod
+    def legacy(cls, name: str, level: int = logging.INFO, accelerator=None):
+        del accelerator  # ignored
+        return cls(name=name, save_path=None, level=level)