"""Unified logger that writes to multiple backends simultaneously. This logger ensures all experimental data is preserved by writing to: 1. Weights & Biases (when available) 2. Local disk (JSON files, model checkpoints, run.log) 3. stdout (for real-time monitoring) """ from enum import Enum import logging import yaml import sys import time from pathlib import Path from typing import Any, Dict, List, Optional import flax import jax.numpy as jnp import numpy as np from experiment_logger.wandb_utils import finish_wandb, init_wandb from experiment_logger.config_logger import LoggingConfig # Global storage for the active logger and the proxy singleton _active_logger: Optional[Any] = None _proxy_instance: Optional["LoggerProxy"] = None def _sanitize_for_yaml(obj: Any) -> Any: """Convert non-primitive values into YAML-safe structures. In particular, avoids PyYAML serializing Enums as ``!!python/object/apply:...`` which OmegaConf will not load. """ if isinstance(obj, Enum): return obj.name if isinstance(obj, Path): return str(obj) if isinstance(obj, (np.generic, jnp.ndarray)): try: return obj.item() except Exception: pass if isinstance(obj, np.ndarray): return obj.tolist() if isinstance(obj, dict): return {str(k): _sanitize_for_yaml(v) for k, v in obj.items()} if isinstance(obj, list): return [_sanitize_for_yaml(v) for v in obj] if isinstance(obj, tuple): return [_sanitize_for_yaml(v) for v in obj] return obj def get_logger() -> "LoggerProxy": """Retrieve the global LoggerProxy. This should be used for all logging calls. It returns a proxy that delegates to the active logger (defaulting to a SimpleLogger until init_logger is called). """ global _proxy_instance, _active_logger if _proxy_instance is None: if _active_logger is None: # Fallback to SimpleLogger to avoid premature directory creation from experiment_logger.simple_logger import SimpleLogger _active_logger = SimpleLogger(run_name="pre_init") _proxy_instance = LoggerProxy() return _proxy_instance def init_logger(**kwargs) -> "UnifiedLogger": """Initialize the full UnifiedLogger and set it as the active logger. This should be called once the configuration is ready. It will create the output directories and set up all logging backends. """ global _active_logger logger = UnifiedLogger(**kwargs) _active_logger = logger return logger class LoggerProxy: """Proxy that delegates all method calls to the active logger instance. This allows the logger to be swapped out (e.g., from a SimpleLogger to a UnifiedLogger) without any clients needing to update their references. """ def _get_logger(self) -> Any: global _active_logger if _active_logger is None: # This shouldn't normally happen since get_logger handles it from experiment_logger.simple_logger import SimpleLogger _active_logger = SimpleLogger(run_name="pre_init_fallback") return _active_logger def __getattr__(self, name: str) -> Any: return getattr(self._get_logger(), name) def __enter__(self): return self._get_logger().__enter__() def __exit__(self, exc_type, exc_val, exc_tb): return self._get_logger().__exit__(exc_type, exc_val, exc_tb) class UnifiedLogger: """Unified logger for scientific experiments with redundant backup.""" def __init__( self, run_name: str, full_config: Dict[str, Any], logging_cfg: LoggingConfig, base_dir: str = "runs", save_code: bool = True, log_level: int = logging.INFO, ): """Initialize the unified logger. Args: run_name: Unique name for this run full_config: Full configuration dictionary with hyperparameters to be saved logging_cfg: Structured logging configuration dataclass base_dir: Base directory for local storage save_code: Whether to save code to WandB """ self.run_name = run_name self.full_config = full_config self.use_wandb = logging_cfg.track self.upload_final_model = logging_cfg.upload_final_model self.upload_checkpoints = logging_cfg.upload_checkpoints self.wandb_available = False self.wandb_run = None self.is_interactive = sys.stdout.isatty() # Setup local storage self.run_dir = Path(base_dir) / run_name self.run_dir.mkdir(parents=True, exist_ok=True) self.checkpoints_dir = self.run_dir / "checkpoints" self.checkpoints_dir.mkdir(exist_ok=True) self.metrics_dir = self.run_dir / "metrics" self.metrics_dir.mkdir(exist_ok=True) self.config_file = self.run_dir / "config.yaml" # Setup standard Python logging mirror self.text_log_file = self.run_dir / "run.log" self._text_logger = logging.getLogger(f"UnifiedLogger_{self.run_name}") self._text_logger.setLevel(log_level) self._text_logger.propagate = False # Avoid duplicate handlers if re-instantiated if not self._text_logger.handlers: fh = logging.FileHandler(self.text_log_file) ch = logging.StreamHandler() formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s") fh.setFormatter(formatter) ch.setFormatter(formatter) self._text_logger.addHandler(fh) self._text_logger.addHandler(ch) # Save config to disk self._save_config() # Setup TensorBoard self.writer = None try: from torch.utils.tensorboard import SummaryWriter self.writer = SummaryWriter(self.run_dir) self.info("TensorBoard SummaryWriter initialized.") except ImportError: self.warning("tensorboard not installed. Skipping SummaryWriter.") # Initialize WandB if requested if self.use_wandb: self._init_wandb(logging_cfg.wandb_project_name, logging_cfg.wandb_entity, save_code) # Initialize metrics storage self.metrics_buffer: List[Dict[str, Any]] = [] self.step_counter = 0 self.info(f"Initialized UnifiedLogger for run: {run_name}") self.info(f"Local storage: {self.run_dir.absolute()}") self.info(f"WandB logging: {self.wandb_available}") def set_level(self, level: int): """Dynamically update the verbosity of the stdout/text logger.""" self._text_logger.setLevel(level) def log_non_interactive(self, msg: str, *args, **kwargs): """Log an info message only if running in a non-interactive environment.""" if not self.is_interactive: self.info(msg, *args, **kwargs) def progress_bar(self, iterable=None, *args, **kwargs): """Wrapper around tqdm that automatically disables in non-interactive environments.""" import tqdm kwargs.setdefault("disable", not self.is_interactive) return tqdm.tqdm(iterable, *args, **kwargs) def info(self, msg: str, *args, **kwargs): """Log an info message to stdout and disk.""" self._text_logger.info(msg, *args, **kwargs) def warning(self, msg: str, *args, **kwargs): """Log a warning message to stdout and disk.""" self._text_logger.warning(msg, *args, **kwargs) def error(self, msg: str, *args, **kwargs): """Log an error message to stdout and disk.""" self._text_logger.error(msg, *args, **kwargs) def debug(self, msg: str, *args, **kwargs): """Log a debug message to stdout and disk.""" self._text_logger.debug(msg, *args, **kwargs) def _init_wandb(self, project_name: str, entity: Optional[str], save_code: bool): """Initialize Weights & Biases logging.""" self.wandb_run = init_wandb( project=project_name, entity=entity, name=self.run_name, config=self.full_config, save_code=save_code, resume="allow", ) self.wandb_available = self.wandb_run is not None def _save_config(self): """Save configuration to disk.""" try: with open(self.config_file, "w") as f: yaml.safe_dump( _sanitize_for_yaml(self.full_config), f, default_flow_style=False, indent=2, sort_keys=False, ) self.info(f"Config saved to {self.config_file}") except Exception as e: self.error(f"Error saving config: {e}") def log(self, metrics: Dict[str, Any], step: Optional[int] = None, commit: bool = True): """Log metrics to all backends. Args: metrics: Dictionary of metric name -> value step: Global step counter (auto-incremented if None) commit: Whether to commit to WandB immediately """ if step is None: step = self.step_counter self.step_counter += 1 # Add timestamp metrics_with_metadata = { "step": step, "timestamp": time.time(), **metrics, } # Log to stdout self._log_to_stdout(metrics_with_metadata) # Log to WandB if self.wandb_run is not None: try: self.wandb_run.log(metrics, step=step, commit=commit) except Exception as e: self.warning(f"WandB logging failed: {e}") # Log to TensorBoard if self.writer is not None: for k, v in metrics.items(): if isinstance(v, (int, float, np.floating, np.integer)): self.writer.add_scalar(k, v, step) elif hasattr(v, "item"): self.writer.add_scalar(k, v.item(), step) elif isinstance(v, (np.ndarray, jnp.ndarray)) and v.size == 1: self.writer.add_scalar(k, v.item(), step) # Buffer for disk storage self.metrics_buffer.append(metrics_with_metadata) # Periodically flush to disk if len(self.metrics_buffer) >= 100: self._flush_metrics() def _log_to_stdout(self, metrics: Dict[str, Any]): """Log metrics to stdout for real-time monitoring.""" step = metrics.get("step", "?") metric_str = ", ".join( f"{k}={v:.6f}" if isinstance(v, (float, np.floating)) else f"{k}={v}" for k, v in metrics.items() if k not in ["step", "timestamp"] ) self.info(f"[Step {step}] {metric_str}") def _flush_metrics(self): """Flush buffered metrics to disk.""" if not self.metrics_buffer: return try: metrics_file = self.metrics_dir / "metrics.yaml" with open(metrics_file, "a") as f: for metric in self.metrics_buffer: # Convert numpy/jax types to native Python types for YAML serialization serializable_metric = {} for k, v in metric.items(): if hasattr(v, "item"): # numpy/jax scalar serializable_metric[k] = v.item() elif isinstance(v, (np.ndarray, jnp.ndarray)): serializable_metric[k] = v.tolist() else: serializable_metric[k] = v f.write("---\n") yaml.safe_dump( _sanitize_for_yaml(serializable_metric), f, default_flow_style=False, sort_keys=False, ) self.metrics_buffer.clear() except Exception as e: self.error(f"Error flushing metrics: {e}") def save_checkpoint( self, params: Any, step: int, prefix: str = "checkpoint", metadata: Optional[Dict[str, Any]] = None, ): """Save model checkpoint to disk and optionally to WandB.""" checkpoint_name = f"{prefix}_step_{step}.flax" checkpoint_path = self.checkpoints_dir / checkpoint_name try: # Save to disk using Flax serialization with open(checkpoint_path, "wb") as f: f.write(flax.serialization.to_bytes(params)) # Save metadata if provided if metadata: metadata_path = self.checkpoints_dir / f"{prefix}_step_{step}_metadata.yaml" with open(metadata_path, "w") as f: yaml.safe_dump( _sanitize_for_yaml(metadata), f, default_flow_style=False, indent=2, sort_keys=False, ) self.info(f"Checkpoint saved: {checkpoint_path}") # Log to WandB as artifact if self.wandb_run is not None and self.upload_checkpoints: try: import wandb artifact = wandb.Artifact( name=f"{self.run_name}_{prefix}", type="model", metadata=metadata or {}, ) artifact.add_file(str(checkpoint_path)) if metadata: artifact.add_file(str(metadata_path)) self.wandb_run.log_artifact(artifact) self.info("Checkpoint uploaded to WandB") except Exception as e: self.warning(f"Could not upload checkpoint to WandB: {e}") except Exception as e: self.error(f"Error saving checkpoint: {e}") def save_final_model(self, params: Any, metadata: Optional[Dict[str, Any]] = None): """Save the final trained model.""" final_model_path = self.run_dir / "final_model.flax" try: with open(final_model_path, "wb") as f: f.write(flax.serialization.to_bytes(params)) if metadata: metadata_path = self.run_dir / "final_model_metadata.yaml" with open(metadata_path, "w") as f: yaml.safe_dump( _sanitize_for_yaml(metadata), f, default_flow_style=False, indent=2, sort_keys=False, ) self.info(f"Final model saved: {final_model_path}") # Log to WandB if self.wandb_run is not None and self.upload_final_model: try: import wandb artifact = wandb.Artifact( name=f"{self.run_name}_final_model", type="model", metadata=metadata or {}, ) artifact.add_file(str(final_model_path)) if metadata: artifact.add_file(str(metadata_path)) self.wandb_run.log_artifact(artifact) except Exception as e: self.warning(f"Could not upload final model to WandB: {e}") except Exception as e: self.error(f"Error saving final model: {e}") def sync_file(self, path: Path) -> None: """Upload a file to W&B if tracking is enabled. Best-effort: logs a warning on failure, never raises. """ if self.wandb_run is None: return try: import wandb # "Simple sync" behavior: wandb will copy this file into the run. wandb.save(str(path), base_path=str(path.parent)) except Exception as e: self.warning(f"Failed to sync file to W&B: {e}") def finish(self): """Finalize logging and cleanup.""" # Flush remaining metrics self._flush_metrics() if self.writer is not None: self.writer.close() self.info(f"Run complete. Results saved to: {self.run_dir.absolute()}") # Finish WandB run if self.wandb_available: finish_wandb() def __enter__(self): """Context manager entry.""" return self def __exit__(self, exc_type, exc_val, exc_tb): """Context manager exit.""" self.finish()