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
This commit is contained in:
2026-08-08 12:39:27 +08:00
parent cb60713a72
commit cbc584470d
9 changed files with 35 additions and 43 deletions
+4 -5
View File
@@ -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``."""
if not _TIMED:
"""Wall-clock debug timer, enabled when the logger level is DEBUG or lower."""
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