Merge pull request #3 from clumsy/main

Add support to wandb and tensorboard loggers, add legacy pytorch_lightning.loggers support

Author: azzhipa

.gitignore

@@ -147,7 +147,9 @@ venv.bak/
 runs/
 
 # mlflow
-mlruns/
+examples/mlflow/
+examples/tensorboard/
+examples/wandb/
 
 # Pyre type checker
 .pyre/

examples/conf/hf.yaml

@@ -19,10 +19,17 @@ trainer:
   val_check_interval: 5
   limit_val_batches: 1
   logger:
-    _target_: lightning.pytorch.loggers.MLFlowLogger
-    experiment_name: ${model_name}-train
-    tracking_uri: ./mlruns
-    synchronous: false
+    - _target_: lightning.pytorch.loggers.MLFlowLogger
+      experiment_name: ${model_name}-train
+      tracking_uri: ./mlflow
+      synchronous: false
+    - _target_: lightning.pytorch.loggers.TensorBoardLogger
+      save_dir: ./tensorboard
+      name: ${model_name}-train
+    - _target_: lightning.pytorch.loggers.WandbLogger
+      project: ${model_name}-train
+      save_dir: ./wandb
+      offline: true
   callbacks:
     - _target_: fkat.pytorch.callbacks.cuda.memory.MemoryObserver
     - _target_: fkat.pytorch.callbacks.monitoring.HardwareStats

examples/sync_wandb_offline.py

@@ -0,0 +1,105 @@
+#!/usr/bin/env python3
+"""Sync wandb offline runs by manually uploading metrics from summary files."""
+
+import json
+import os
+import sys
+from pathlib import Path
+
+import wandb
+
+
+def _extract_project_name(config_file: Path) -> str:
+    """Extract project name from wandb config file."""
+    if not config_file.exists():
+        return "unknown"
+
+    import yaml
+
+    with open(config_file) as f:
+        config = yaml.safe_load(f)
+
+    # Try to extract model_name from various locations
+    if "model_name" in config and isinstance(config["model_name"], dict):
+        return f"{config['model_name']['value']}-train"
+
+    if "_wandb" in config and "value" in config["_wandb"]:
+        wandb_val = config["_wandb"]["value"]
+        e_dict = wandb_val.get("e", {})
+        if e_dict:
+            first_exec = list(e_dict.values())[0]
+            args = first_exec.get("args", [])
+            for arg in args:
+                if "model_name=" in str(arg):
+                    model = str(arg).split("model_name=")[1]
+                    return f"{model}-train"
+
+    return "unknown"
+
+
+def _upload_files(run_path: Path, run_dir_path: str) -> None:
+    """Upload files from offline run to wandb."""
+    import shutil
+
+    files_dir = run_path / "files"
+    if not files_dir.exists():
+        return
+
+    for file_path in files_dir.rglob("*"):
+        if file_path.is_file() and not file_path.name.startswith("wandb-"):
+            try:
+                rel_path = file_path.relative_to(files_dir)
+                dest = Path(run_dir_path) / rel_path
+                dest.parent.mkdir(parents=True, exist_ok=True)
+                shutil.copy2(file_path, dest)
+                wandb.save(str(rel_path), base_path=run_dir_path, policy="now")
+            except Exception:
+                pass
+
+
+def sync_offline_run(run_dir: str, base_url: str, api_key: str) -> None:
+    """Sync an offline run by reading summary and uploading metrics."""
+    run_path = Path(run_dir)
+    summary_file = run_path / "files" / "wandb-summary.json"
+
+    if not summary_file.exists():
+        print(f"No summary file found in {run_dir}")
+        return
+
+    run_id = run_path.name.split("-")[-1]
+
+    with open(summary_file) as f:
+        summary = json.load(f)
+
+    config_file = run_path / "files" / "config.yaml"
+    project = _extract_project_name(config_file)
+
+    os.environ["WANDB_BASE_URL"] = base_url
+    os.environ["WANDB_API_KEY"] = api_key
+
+    run = wandb.init(project=project, id=run_id, resume="allow", mode="online")
+
+    for key, value in summary.items():
+        if not key.startswith("_"):
+            run.log({key: value})
+
+    _upload_files(run_path, run.dir)
+
+    run.finish()
+    print(f"✓ Synced {run_dir}")
+
+
+if __name__ == "__main__":
+    if len(sys.argv) < 4:
+        print("Usage: sync_wandb_offline.py <base_url> <api_key> <run_dir1> [run_dir2 ...]")
+        sys.exit(1)
+
+    base_url = sys.argv[1]
+    api_key = sys.argv[2]
+    run_dirs = sys.argv[3:]
+
+    for run_dir in run_dirs:
+        try:
+            sync_offline_run(run_dir, base_url, api_key)
+        except Exception as e:
+            print(f"✗ Failed to sync {run_dir}: {e}")

fkat/__version__.py

@@ -1,3 +1,3 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
-__version__ = "0.1.1"
+__version__ = "0.1.2"

fkat/pytorch/callbacks/cuda/memory.py

@@ -18,14 +18,15 @@
 )
 from fkat.pytorch.loggers import LightningLogger
 from fkat.pytorch.callbacks.loggers import CallbackLogger
+from fkat.utils import safe_timestamp
 
 logger: logging.Logger = logging.getLogger(__name__)
 
 
 def _artifact_path(root_dir: str, rank: int, file_type: str, ext: str) -> tuple[str, str]:
     base_dir = os.path.join(root_dir, "torch.cuda.memory")
-    now = datetime.now(timezone.utc).isoformat()
-    file_path = os.path.join(base_dir, f"rank{rank}/{file_type}/rank{rank}_{now}.{ext}")
+    timestamp = safe_timestamp()
+    file_path = os.path.join(base_dir, f"rank{rank}/{file_type}/rank{rank}_{timestamp}.{ext}")
     os.makedirs(os.path.dirname(file_path), exist_ok=True)
     return base_dir, file_path
 

fkat/pytorch/callbacks/cuda/test/memory_test.py

@@ -14,10 +14,13 @@
 
 class TestArtifactPath(unittest.TestCase):
     @patch("os.makedirs")
-    @patch(f"{memory.__name__}.datetime")
+    @patch("fkat.utils.datetime")
     def test_artifact_path_creates_correct_structure(self, mock_datetime, mock_makedirs):
         # Arrange
-        mock_datetime.now.return_value.isoformat.return_value = "2025-06-18T20:00:00"
+        mock_now = MagicMock()
+        mock_now.strftime.return_value = "2025-06-18_20-00-00-"
+        mock_now.microsecond = 0
+        mock_datetime.now.return_value = mock_now
         root_dir = "/tmp/test"
         rank = 1
         file_type = "snapshot"
@@ -28,17 +31,17 @@ def test_artifact_path_creates_correct_structure(self, mock_datetime, mock_maked
 
         # Assert
         expected_base_dir = "/tmp/test/torch.cuda.memory"
-        expected_file_path = "/tmp/test/torch.cuda.memory/rank1/snapshot/rank1_2025-06-18T20:00:00.pickle"
+        expected_file_path = "/tmp/test/torch.cuda.memory/rank1/snapshot/rank1_2025-06-18_20-00-00-000.pickle"
 
         assert base_dir == expected_base_dir
         assert file_path == expected_file_path
         mock_makedirs.assert_called_once_with("/tmp/test/torch.cuda.memory/rank1/snapshot", exist_ok=True)
 
     @patch("os.makedirs")
-    @patch(f"{memory.__name__}.datetime")
-    def test_artifact_path_different_parameters(self, mock_datetime, mock_makedirs):
+    @patch(f"{memory.__name__}.safe_timestamp")
+    def test_artifact_path_different_parameters(self, mock_safe_timestamp, mock_makedirs):
         # Arrange
-        mock_datetime.now.return_value.isoformat.return_value = "2025-06-18T15:30:45"
+        mock_safe_timestamp.return_value = "2025-06-18_15-30-45-123"
         root_dir = "/var/logs"
         rank = 0
         file_type = "flamegraph"
@@ -49,7 +52,7 @@ def test_artifact_path_different_parameters(self, mock_datetime, mock_makedirs):
 
         # Assert
         expected_base_dir = "/var/logs/torch.cuda.memory"
-        expected_file_path = "/var/logs/torch.cuda.memory/rank0/flamegraph/rank0_2025-06-18T15:30:45.svg"
+        expected_file_path = "/var/logs/torch.cuda.memory/rank0/flamegraph/rank0_2025-06-18_15-30-45-123.svg"
 
         assert base_dir == expected_base_dir
         assert file_path == expected_file_path

fkat/pytorch/callbacks/loggers.py

@@ -4,22 +4,22 @@
 from typing_extensions import override
 
 import lightning as L
-from lightning.pytorch.loggers import MLFlowLogger
+from lightning.pytorch.loggers import TensorBoardLogger, WandbLogger
 from mlflow.entities import Metric, RunTag, Param
 from mlflow.tracking import MlflowClient  # type: ignore[possibly-unbound-import]
 
 if TYPE_CHECKING:
-    from lightning.pytorch.loggers import MLFlowLogger
+    pass
 
-from fkat.pytorch.loggers import LightningLogger
+from fkat.pytorch.loggers import LightningLogger, _is_logger_type
 from fkat.utils import assert_not_none
 from fkat.utils.logging import rank0_logger
 from fkat.utils.mlflow import broadcast_mlflow_run_id, mlflow_logger
 
 log = rank0_logger(__name__)
 
 
-class MLFlowCallbackLogger:
+class MLFlowCallbackLogger(LightningLogger):
     """
     Mlflow logger class that supports distributed logging of tags, metrics and artifacts.
 
@@ -86,6 +86,69 @@ def log_artifact(self, local_path: str, artifact_path: str | None = None) -> Non
         )
 
 
+class TensorBoardCallbackLogger(LightningLogger):
+    """TensorBoard logger for distributed logging."""
+
+    def __init__(self, logger: TensorBoardLogger) -> None:
+        self._logger = logger
+
+    def log_tag(self, key: str, value: str) -> None:
+        self._logger.experiment.add_text(key, value)
+
+    def tags(self) -> dict[str, Any]:
+        return {}
+
+    def log_batch(
+        self,
+        metrics: dict[str, float] | None = None,
+        params: dict[str, Any] | None = None,
+        tags: dict[str, str] | None = None,
+        timestamp: int | None = None,
+        step: int | None = None,
+    ) -> None:
+        if metrics:
+            for k, v in metrics.items():
+                self._logger.experiment.add_scalar(k, v, step)
+        if tags:
+            for k, v in tags.items():
+                self._logger.experiment.add_text(k, v, step)
+
+    def log_artifact(self, local_path: str, artifact_path: str | None = None) -> None:
+        pass
+
+
+class WandbCallbackLogger(LightningLogger):
+    """WandB logger for distributed logging."""
+
+    def __init__(self, logger: WandbLogger) -> None:
+        self._logger = logger
+
+    def log_tag(self, key: str, value: str) -> None:
+        self._logger.experiment.config.update({key: value})
+
+    def tags(self) -> dict[str, Any]:
+        return dict(self._logger.experiment.config)
+
+    def log_batch(
+        self,
+        metrics: dict[str, float] | None = None,
+        params: dict[str, Any] | None = None,
+        tags: dict[str, str] | None = None,
+        timestamp: int | None = None,
+        step: int | None = None,
+    ) -> None:
+        log_dict = {}
+        if metrics:
+            log_dict.update(metrics)
+        if tags:
+            log_dict.update(tags)
+        if log_dict:
+            self._logger.experiment.log(log_dict, step=step)
+
+    def log_artifact(self, local_path: str, artifact_path: str | None = None) -> None:
+        self._logger.experiment.save(local_path)
+
+
 class CallbackLogger(LightningLogger):
     """
     A wrapper on top of the collection of Logger instances,
@@ -104,9 +167,13 @@ class CallbackLogger(LightningLogger):
     def __init__(self, trainer: "L.Trainer | None", loggers: list[LightningLogger] | None = None) -> None:
         if trainer:
             self.loggers = []
-            for logger in trainer.logger if isinstance(trainer.logger, list) else [trainer.logger]:
-                if isinstance(logger, MLFlowLogger):
+            for logger in trainer.loggers:
+                if _is_logger_type(logger, "MLFlowLogger"):
                     self.loggers.append(MLFlowCallbackLogger(trainer=trainer))
+                elif _is_logger_type(logger, "TensorBoardLogger"):
+                    self.loggers.append(TensorBoardCallbackLogger(logger=logger))  # type: ignore[arg-type]
+                elif _is_logger_type(logger, "WandbLogger"):
+                    self.loggers.append(WandbCallbackLogger(logger=logger))  # type: ignore[arg-type]
         else:
             assert loggers
             self.loggers = loggers