refactor: centralize logging in astrai.logging, replace ASTRAI_TIMED with log level

- move setup_logging to astrai/logging.py
- timed() now uses logger.isEnabledFor(DEBUG) instead of separate env var
- enable ASTR_LOG_LEVEL=DEBUG to see per-step timing logs
- call setup_logging() in stream_chat.py

astrai/__init__.py

@@ -1,9 +1,6 @@
 __version__ = "1.3.13"
 __author__ = "ViperEkura"
 
-import logging
-import os
-
 from astrai.config import (
     AutoRegressiveLMConfig,
     BaseModelConfig,
@@ -28,6 +25,7 @@ from astrai.inference import (
     run_server,
     sample,
 )
+from astrai.logging import setup_logging
 from astrai.model import (
     AutoModel,
     AutoRegressiveLM,
@@ -55,30 +53,6 @@ from astrai.trainer import (
     Trainer,
 )
 
-
-def setup_logging(level: str = "INFO"):
-    """Attach a handler to the ``astrai`` logger (only, not root).
-
-    Call once per process, e.g. at the top of CLI scripts.
-    Set ``ASTR_LOG_LEVEL`` to override the default ``INFO``.
-    """
-    _logger = logging.getLogger("astrai")
-    if _logger.handlers:
-        return
-    _level = getattr(
-        logging, os.environ.get("ASTR_LOG_LEVEL", level).upper(), logging.INFO
-    )
-    _logger.setLevel(_level)
-    _handler = logging.StreamHandler()
-    _handler.setFormatter(
-        logging.Formatter(
-            "%(asctime)s | %(levelname)-7s | %(name)s | %(message)s",
-            datefmt="%Y-%m-%d %H:%M:%S",
-        )
-    )
-    _logger.addHandler(_handler)
-
-
 __all__ = [
     "AutoRegressiveLM",
     "AutoRegressiveLMConfig",
@@ -122,3 +96,5 @@ __all__ = [
     "setup_logging",
     "spawn_parallel_fn",
 ]
+
+setup_logging()

astrai/inference/core/executor.py

@@ -1,5 +1,4 @@
 import logging
-import os
 import time
 from contextlib import contextmanager
 from dataclasses import dataclass
@@ -23,19 +22,19 @@ from astrai.model.automodel import AutoModel
 from astrai.tokenize.tokenizer import AutoTokenizer
 
 logger = logging.getLogger(__name__)
-_TIMED = os.environ.get("ASTRAI_TIMED", "") == "1"
 
 
 @contextmanager
 def timed(label: str, log: Optional[logging.Logger] = None):
-    """Wall-clock debug timer, enabled via ``ASTRAI_TIMED=1``."""
+    """Wall-clock debug timer, enabled when the logger level is DEBUG or lower."""
-    if not _TIMED:
+    log = log or logger
+    if not log.isEnabledFor(logging.DEBUG):
         yield
         return
     tic = time.perf_counter()
     yield
     elapsed_ms = (time.perf_counter() - tic) * 1000
-    (log or logger).info("%s %.1fms", label, elapsed_ms)
+    log.debug("%s %.1fms", label, elapsed_ms)
 
 
 @dataclass

astrai/logging.py

@@ -0,0 +1,27 @@
+import logging
+import os
+
+
+def setup_logging(level: str = "INFO"):
+    """Attach a StreamHandler to the ``astrai`` logger (idempotent).
+
+    Call once per process at the top of CLI scripts.
+    Set ``ASTR_LOG_LEVEL`` env var to override the default level.
+
+    Level names: ``DEBUG``, ``INFO``, ``WARNING``, ``ERROR``, ``CRITICAL``.
+    ``DEBUG`` enables per-step prefill/decode timing logs
+    (:func:`astrai.inference.core.executor.timed`).
+    """
+    logger = logging.getLogger("astrai")
+    if logger.handlers:
+        return
+    level_name = os.environ.get("ASTR_LOG_LEVEL", level).upper()
+    logger.setLevel(getattr(logging, level_name, logging.INFO))
+    handler = logging.StreamHandler()
+    handler.setFormatter(
+        logging.Formatter(
+            "%(asctime)s | %(levelname)-7s | %(name)s | %(message)s",
+            datefmt="%Y-%m-%d %H:%M:%S",
+        )
+    )
+    logger.addHandler(handler)

scripts/demo/stream_chat.py

@@ -3,7 +3,7 @@ from pathlib import Path
 
 import torch
 
-from astrai.inference import InferenceEngine
+from astrai import InferenceEngine
 from astrai.model import AutoModel
 from astrai.tokenize import AutoTokenizer
 

scripts/tools/benchmark.py

@@ -5,7 +5,6 @@ from typing import Optional, Union
 import click
 import torch
 
-from astrai import setup_logging
 from astrai.config import BaseModelConfig, ConfigFactory
 from astrai.extension import ATTN_BACKEND, AttentionBackendFactory, attn_backend
 from astrai.inference.core.cache import PagePool
@@ -478,5 +477,4 @@ def benchmark_command(
 
 
 if __name__ == "__main__":
-    setup_logging()
     benchmark_command()

scripts/tools/generate.py

@@ -6,7 +6,6 @@ import click
 import torch
 from tqdm import tqdm
 
-from astrai import setup_logging
 from astrai.inference import InferenceEngine
 from astrai.model import AutoModel
 from astrai.tokenize import AutoTokenizer
@@ -157,5 +156,4 @@ def generate_command(**kwargs):
 
 
 if __name__ == "__main__":
-    setup_logging()
     generate_command()

scripts/tools/preprocess.py

@@ -2,7 +2,6 @@
 
 import click
 
-from astrai import setup_logging
 from astrai.config.preprocess_config import PipelineConfig
 from astrai.preprocessing.pipeline import Pipeline
 
@@ -48,5 +47,4 @@ def preprocess_command(inputs, output_dir, pipeline_config, tokenizer_path, batc
 
 
 if __name__ == "__main__":
-    setup_logging()
     preprocess_command()

scripts/tools/server.py

@@ -3,7 +3,6 @@ from pathlib import Path
 import click
 import torch
 
-from astrai import setup_logging
 from astrai.inference import run_server
 
 _DTYPES = ["bfloat16", "float16", "float32"]
@@ -65,5 +64,4 @@ def server_command(
 
 
 if __name__ == "__main__":
-    setup_logging()
     server_command()

scripts/tools/train.py

@@ -8,7 +8,6 @@ import torch
 from click.core import ParameterSource
 from torch import optim
 
-from astrai import setup_logging
 from astrai.config import AutoRegressiveLMConfig, TrainConfig
 from astrai.dataset import DatasetFactory, dpo_collate_fn, grpo_collate_fn
 from astrai.model import AutoRegressiveLM
@@ -828,5 +827,4 @@ def train(
 
 
 if __name__ == "__main__":
-    setup_logging()
     train_command()