Upload logger.py

logger.py

@@ -0,0 +1,310 @@
+"""
+Training logger: wandb + CSV + stdout + run summary JSON.
+
+All logging is gated on ``enabled`` (typically ``is_main_process()``).
+Wandb is optional -- if ``wandb`` is not installed or fails to init,
+logging falls back to CSV + stdout silently.
+
+CSV columns (one row per logged event):
+  run_id, step, samples_seen, wall_time_sec, event_type,
+  train_loss, loss_fine, loss_coarse, loss_ratio,
+  grad_norm, lr_connector, lr_dino, lr_llm,
+  throughput_samples_sec, gpu_mem_gb,
+  val_loss, val_loss_fine, val_loss_coarse, val_loss_ratio,
+  attention_entropy
+"""
+
+import csv
+import json
+import os
+import subprocess
+import time
+from datetime import datetime
+from pathlib import Path
+from typing import Optional
+
+
+def _get_git_hash() -> str:
+    """Get current git commit hash, or 'unknown' if not in a repo."""
+    try:
+        result = subprocess.run(
+            ["git", "rev-parse", "--short", "HEAD"],
+            capture_output=True, text=True, timeout=5,
+        )
+        return result.stdout.strip() if result.returncode == 0 else "unknown"
+    except Exception:
+        return "unknown"
+
+
+def _gpu_memory_gb() -> float:
+    """Get current GPU memory allocated in GB, or 0 if no GPU."""
+    try:
+        import torch
+        if torch.cuda.is_available():
+            return torch.cuda.memory_allocated() / (1024 ** 3)
+    except Exception:
+        pass
+    return 0.0
+
+
+CSV_COLUMNS = [
+    "run_id", "step", "samples_seen", "wall_time_sec", "event_type",
+    "train_loss", "loss_fine", "loss_coarse", "loss_ratio",
+    "grad_norm", "lr_connector", "lr_dino", "lr_llm",
+    "throughput_samples_sec", "gpu_mem_gb",
+    "val_loss", "val_loss_fine", "val_loss_coarse", "val_loss_ratio",
+    "attention_entropy",
+]
+
+
+class TrainingLogger:
+    """
+    Unified logger that writes to wandb, structured CSV, and stdout.
+
+    Parameters
+    ----------
+    project : str
+        wandb project name.
+    config : dict
+        Training config to log as wandb config / CSV header metadata.
+    enabled : bool
+        If False, all log calls are no-ops (use for non-rank-0 processes).
+    log_dir : str
+        Directory for the CSV log file.
+    """
+
+    def __init__(
+        self,
+        project: str = "foveated-vlm",
+        config: Optional[dict] = None,
+        enabled: bool = True,
+        log_dir: Optional[str] = None,
+    ):
+        self.enabled = enabled
+        self._wandb_run = None
+        self._csv_path = None
+        self._csv_writer = None
+        self._csv_file = None
+        self._start_time = time.time()
+        self._config = config or {}
+        self._run_id = ""
+        self._best_val_loss = float("inf")
+        self._best_step = 0
+        self._last_step = 0
+        self._last_samples = 0
+        self._git_hash = _get_git_hash()
+
+        if not enabled:
+            return
+
+        # ---- Run ID ----
+        run_name = (config or {}).get("wandb", {}).get("run_name", "run")
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        self._run_id = f"{run_name}_{timestamp}"
+
+        # ---- wandb ----
+        try:
+            import wandb
+            self._wandb_run = wandb.init(
+                project=project,
+                name=run_name,
+                config=config or {},
+                resume="allow",
+            )
+        except Exception:
+            pass
+
+        # ---- CSV ----
+        if log_dir is None:
+            log_dir = (config or {}).get("checkpoint", {}).get(
+                "save_dir", "/workspace/logs"
+            )
+        self._log_dir = Path(log_dir)
+        self._log_dir.mkdir(parents=True, exist_ok=True)
+
+        self._csv_path = self._log_dir / f"metrics_{self._run_id}.csv"
+        self._csv_file = open(self._csv_path, "w", newline="")
+        self._csv_writer = csv.DictWriter(
+            self._csv_file, fieldnames=CSV_COLUMNS, extrasaction="ignore",
+        )
+        self._csv_writer.writeheader()
+        self._csv_file.flush()
+
+    def _write_csv_row(self, row: dict):
+        if self._csv_writer is not None:
+            row.setdefault("run_id", self._run_id)
+            row.setdefault("wall_time_sec", f"{time.time() - self._start_time:.1f}")
+            self._csv_writer.writerow(row)
+            self._csv_file.flush()
+
+    def log_step(
+        self,
+        step: int,
+        loss: float,
+        fine_loss: float = 0.0,
+        coarse_loss: float = 0.0,
+        lr: float = 0.0,
+        grad_norm: float = 0.0,
+        samples_seen: int = 0,
+        samples_per_sec: float = 0.0,
+        lr_groups: Optional[dict] = None,
+    ):
+        """Log a training step with full metrics."""
+        if not self.enabled:
+            return
+
+        self._last_step = step
+        self._last_samples = samples_seen
+
+        loss_ratio = fine_loss / max(coarse_loss, 1e-8) if coarse_loss > 0 else 0.0
+        gpu_mem = _gpu_memory_gb()
+
+        # Parse per-group LRs
+        lr_connector = lr
+        lr_dino = lr
+        lr_llm = lr
+        if lr_groups:
+            lr_connector = lr_groups.get("connector", lr)
+            lr_dino = lr_groups.get("dino", lr)
+            lr_llm = lr_groups.get("llm", lr)
+
+        # stdout
+        print(
+            f"  step {step:6d} | loss {loss:.4f} | "
+            f"fine {fine_loss:.4f} | ratio {loss_ratio:.3f} | "
+            f"lr {lr:.2e} | gnorm {grad_norm:.2f} | "
+            f"{samples_per_sec:.0f} samp/s | {gpu_mem:.1f}GB",
+            flush=True,
+        )
+
+        # wandb
+        if self._wandb_run is not None:
+            try:
+                import wandb
+                log_dict = {
+                    "train/loss": loss,
+                    "train/fine_loss": fine_loss,
+                    "train/coarse_loss": coarse_loss,
+                    "train/loss_ratio": loss_ratio,
+                    "train/lr": lr,
+                    "train/lr_connector": lr_connector,
+                    "train/lr_dino": lr_dino,
+                    "train/lr_llm": lr_llm,
+                    "train/grad_norm": grad_norm,
+                    "train/samples_seen": samples_seen,
+                    "train/throughput": samples_per_sec,
+                    "train/gpu_mem_gb": gpu_mem,
+                }
+                wandb.log(log_dict, step=step)
+            except Exception:
+                pass
+
+        # CSV
+        self._write_csv_row({
+            "step": step,
+            "samples_seen": samples_seen,
+            "event_type": "train",
+            "train_loss": f"{loss:.6f}",
+            "loss_fine": f"{fine_loss:.6f}",
+            "loss_coarse": f"{coarse_loss:.6f}",
+            "loss_ratio": f"{loss_ratio:.4f}",
+            "grad_norm": f"{grad_norm:.4f}",
+            "lr_connector": f"{lr_connector:.2e}",
+            "lr_dino": f"{lr_dino:.2e}",
+            "lr_llm": f"{lr_llm:.2e}",
+            "throughput_samples_sec": f"{samples_per_sec:.1f}",
+            "gpu_mem_gb": f"{gpu_mem:.2f}",
+        })
+
+    def log_eval(
+        self,
+        step: int,
+        val_loss: float,
+        val_fine_loss: float = 0.0,
+        val_coarse_loss: float = 0.0,
+        attention_entropy: float = 0.0,
+    ):
+        """Log a validation result with extended metrics."""
+        if not self.enabled:
+            return
+
+        val_ratio = val_fine_loss / max(val_coarse_loss, 1e-8) if val_coarse_loss > 0 else 0.0
+
+        if val_loss < self._best_val_loss:
+            self._best_val_loss = val_loss
+            self._best_step = step
+
+        print(
+            f"  [eval] step {step:6d} | val_loss {val_loss:.4f} | "
+            f"fine {val_fine_loss:.4f} | ratio {val_ratio:.3f} | "
+            f"entropy {attention_entropy:.4f} | "
+            f"best {self._best_val_loss:.4f}@{self._best_step}",
+            flush=True,
+        )
+
+        if self._wandb_run is not None:
+            try:
+                import wandb
+                wandb.log({
+                    "eval/val_loss": val_loss,
+                    "eval/val_fine_loss": val_fine_loss,
+                    "eval/val_coarse_loss": val_coarse_loss,
+                    "eval/val_loss_ratio": val_ratio,
+                    "eval/attention_entropy": attention_entropy,
+                    "eval/best_val_loss": self._best_val_loss,
+                }, step=step)
+            except Exception:
+                pass
+
+        self._write_csv_row({
+            "step": step,
+            "samples_seen": self._last_samples,
+            "event_type": "eval",
+            "val_loss": f"{val_loss:.6f}",
+            "val_loss_fine": f"{val_fine_loss:.6f}",
+            "val_loss_coarse": f"{val_coarse_loss:.6f}",
+            "val_loss_ratio": f"{val_ratio:.4f}",
+            "attention_entropy": f"{attention_entropy:.6f}",
+        })
+
+    def save_run_summary(self, final_loss: float = 0.0, total_samples: int = 0):
+        """Save run summary JSON at end of training."""
+        if not self.enabled:
+            return
+
+        elapsed = time.time() - self._start_time
+        summary = {
+            "run_id": self._run_id,
+            "git_hash": self._git_hash,
+            "config_file": self._config.get("_config_path", ""),
+            "final_train_loss": final_loss,
+            "best_val_loss": self._best_val_loss,
+            "best_val_step": self._best_step,
+            "total_steps": self._last_step,
+            "total_samples": total_samples,
+            "wall_time_sec": elapsed,
+            "wall_time_hours": elapsed / 3600,
+            "csv_path": str(self._csv_path) if self._csv_path else "",
+            "timestamp": datetime.now().isoformat(),
+        }
+
+        summary_path = self._log_dir / f"run_summary_{self._run_id}.json"
+        with open(summary_path, "w") as f:
+            json.dump(summary, f, indent=2)
+        print(f"  Run summary saved to {summary_path}", flush=True)
+
+    def finish(self):
+        """Flush and close all logging backends."""
+        if not self.enabled:
+            return
+
+        if self._wandb_run is not None:
+            try:
+                import wandb
+                wandb.finish()
+            except Exception:
+                pass
+
+        if self._csv_file is not None:
+            self._csv_file.close()
+            self._csv_file = None