[BugFix] loguru coverage

Author: RandallBalestriero

stable_pretraining/utils/distributed.py

@@ -10,7 +10,6 @@
 """
 
 import functools
-import logging
 import os
 import random
 import warnings
@@ -19,8 +18,7 @@
 import torch
 import torch.distributed as dist
 import torch.distributed.nn
-
-log = logging.getLogger(__name__)
+from loguru import logger
 
 # Same precedence Lightning uses (see ``lightning.fabric.utilities.rank_zero``):
 # LOCAL_RANK before SLURM_PROCID because SLURM_PROCID can be set even when SLURM
@@ -116,10 +114,9 @@ def rank_zero_warn(message: Union[str, Warning], stacklevel: int = 4, **kwargs)
 
 
 @rank_zero_only
-def rank_zero_info(*args, stacklevel: int = 4, **kwargs) -> None:
+def rank_zero_info(message, *args, **kwargs) -> None:
     """Emit an info-level log message only on global rank 0."""
-    kwargs["stacklevel"] = stacklevel
-    log.info(*args, **kwargs)
+    logger.info(message, *args, **kwargs)
 
 
 class _DummyExperiment:
@@ -217,7 +214,7 @@ def seed_everything(
         )
 
     if verbose:
-        log.info(f"Seed set to {seed}")
+        logger.info(f"Seed set to {seed}")
 
     os.environ["PL_GLOBAL_SEED"] = str(seed)
     random.seed(seed)

stable_pretraining/utils/log_reader.py

@@ -1,6 +1,5 @@
 """Unified log reader for local and wandb logs."""
 
-import logging
 import re
 from abc import ABC, abstractmethod
 from multiprocessing import Pool
@@ -9,25 +8,26 @@
 
 import numpy as np
 import pandas as pd
+from loguru import logger
 from tqdm import tqdm
 
 try:
     import jsonlines
 except ModuleNotFoundError:
-    logging.warning(
+    logger.warning(
         "jsonlines module is not installed, local log reading will not work."
     )
 
 try:
     import omegaconf
 except ModuleNotFoundError:
-    logging.warning("omegaconf module is not installed, config loading will not work.")
+    logger.warning("omegaconf module is not installed, config loading will not work.")
 
 try:
     import wandb as wandbapi
     from tqdm.contrib.logging import logging_redirect_tqdm
 except ModuleNotFoundError:
-    logging.warning(
+    logger.warning(
         "Wandb module is not installed, make sure to not use wandb for logging "
         "or an error will be thrown."
     )
@@ -113,16 +113,16 @@ def read(self, path: Union[str, Path]) -> List[Dict[str, Any]]:
 
         values = []
         logs_files = list(_path.glob("logs_rank_*.jsonl"))
-        logging.info(f"Reading .jsonl files from {_path}")
-        logging.info(f"\t=> {len(logs_files)} ranks detected")
+        logger.info(f"Reading .jsonl files from {_path}")
+        logger.info(f"\t=> {len(logs_files)} ranks detected")
 
         for log_file in logs_files:
             rank = int(log_file.stem.split("rank_")[1])
             for obj in jsonlines.open(log_file).iter(type=dict, skip_invalid=True):
                 obj["rank"] = rank
                 values.append(obj)
 
-        logging.info(f"\t=> total length of logs: {len(values)}")
+        logger.info(f"\t=> total length of logs: {len(values)}")
         return values
 
     def read_project(
@@ -143,7 +143,6 @@ def read_project(
         configs = []
         values = []
 
-        logging.basicConfig(level=logging.INFO)
         if logging_redirect_tqdm:
             with logging_redirect_tqdm():
                 args = [run.parent for run in runs]
@@ -294,7 +293,7 @@ def read_project(
             per_page=per_page,
             include_sweeps=include_sweeps,
         )
-        logging.info(f"Found {len(runs)} runs for project {project}")
+        logger.info(f"Found {len(runs)} runs for project {project}")
 
         if return_summary:
             data = []
@@ -368,7 +367,7 @@ def create_table(
         Returns:
             Formatted table as DataFrame
         """
-        logging.info(f"Creating table from {len(configs)} runs.")
+        logger.info(f"Creating table from {len(configs)} runs.")
         filters = filters or {}
 
         df = pd.DataFrame(configs).T
@@ -381,12 +380,12 @@ def create_table(
                 v = [v]
             s = df[k].isin(v)
             df = df.loc[s]
-            logging.info(f"After filtering {k}, {len(df)} runs are left.")
+            logger.info(f"After filtering {k}, {len(df)} runs are left.")
 
         rows = natural_sort(df[row].astype(str).unique())
-        logging.info(f"Found rows: {rows}")
+        logger.info(f"Found rows: {rows}")
         columns = natural_sort(df[column].astype(str).unique())
-        logging.info(f"Found columns: {columns}")
+        logger.info(f"Found columns: {columns}")
 
         output = pd.DataFrame(columns=columns, index=rows)
 
@@ -395,11 +394,11 @@ def create_table(
                 cell_runs = (df[row].astype(str) == r) & (df[column].astype(str) == c)
                 n = np.count_nonzero(cell_runs)
                 samples = []
-                logging.info(f"Number of runs for cell ({r}, {c}): {n}")
+                logger.info(f"Number of runs for cell ({r}, {c}): {n}")
 
                 for id in df[cell_runs].index.values:
                     if value not in dfs[id].columns:
-                        logging.info(f"Run {id} missing {value}, skipping....")
+                        logger.info(f"Run {id} missing {value}, skipping....")
                         continue
                     samples.append(dfs[id][value].values.reshape(-1))