[content] convert_wandb_logs_to_tb.py (python)
"""
Script to download specific wandb runs and convert them to TensorBoard format for local storage.
This provides much faster access to your experiment logs compared to the wandb web interface.
"""

import os
import argparse
import json
import logging
from pathlib import Path
from typing import List, Optional, Dict, Any
import wandb
import tensorflow as tf
from tensorboard.backend.event_processing.event_accumulator import EventAccumulator
import pandas as pd
from datetime import datetime

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class WandbToTensorboardConverter:
    """Convert wandb runs to TensorBoard format for local storage."""
    
    def __init__(self, output_dir: str = "./tensorboard_logs"):
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        
    def download_run(self, run_id: str, project_name: Optional[str] = None) -> Optional[Dict[str, Any]]:
        """
        Download a specific wandb run by ID.
        
        Args:
            run_id: The wandb run ID
            project_name: Optional project name (if not provided, will search all accessible projects)
            
        Returns:
            Run data dictionary or None if not found
        """
        try:
            api = wandb.Api()
            
            if project_name:
                run = api.run(f"{project_name}/{run_id}")
            else:
                # Search across all accessible projects
                runs = api.runs()
                run = None
                for r in runs:
                    if r.id == run_id:
                        run = r
                        break
                
                if not run:
                    logger.error(f"Run {run_id} not found in any accessible project")
                    return None
            
            logger.info(f"Downloading run {run_id} from project {run.project}")
            
            # Get run history
            history = run.history()
            
            # Get run config
            config = run.config
            
            # Get run summary
            summary = run.summary
            
            run_data = {
                'id': run_id,
                'name': run.name,
                'project': run.project,
                'config': config,
                'summary': summary,
                'history': history,
                'tags': run.tags,
                'created_at': run.created_at,
                'state': run.state
            }
            
            logger.info(f"Successfully downloaded run {run_id}")
            return run_data
            
        except Exception as e:
            logger.error(f"Error downloading run {run_id}: {e}")
            return None
    
    def convert_to_tensorboard(self, run_data: Dict[str, Any], output_subdir: Optional[str] = None) -> str:
        """
        Convert wandb run data to TensorBoard format.
        
        Args:
            run_data: The run data from download_run
            output_subdir: Optional subdirectory name (defaults to run name)
            
        Returns:
            Path to the TensorBoard log directory
        """
        if not run_data:
            raise ValueError("No run data provided")
        
        # Create output directory
        if output_subdir is None:
            output_subdir = f"{run_data['name']}_{run_data['id']}"
        
        tb_dir = self.output_dir / output_subdir
        tb_dir.mkdir(parents=True, exist_ok=True)
        
        # Create TensorBoard writer
        writer = tf.summary.create_file_writer(str(tb_dir))
        
        # Convert history data to TensorBoard format
        history = run_data['history']
        if not history.empty:
            with writer.as_default():
                for step, row in history.iterrows():
                    for column in row.index:
                        if pd.notna(row[column]) and column not in ['_step', '_runtime', '_timestamp']:
                            # Handle different data types
                            value = row[column]
                            if isinstance(value, (int, float)):
                                tf.summary.scalar(column, value, step=step)
                            elif isinstance(value, str):
                                tf.summary.text(column, value, step=step)
                            # Add more data type handlers as needed
        
        # Add config as text
        if run_data['config']:
            config_text = json.dumps(run_data['config'], indent=2)
            with writer.as_default():
                tf.summary.text("config", config_text, step=0)
        
        # Add summary as text
        if run_data['summary']:
            summary_text = json.dumps(run_data['summary'], indent=2)
            with writer.as_default():
                tf.summary.text("summary", summary_text, step=0)
        
        writer.close()
        
        logger.info(f"Converted run {run_data['id']} to TensorBoard format at {tb_dir}")
        return str(tb_dir)
    
    def download_and_convert(self, run_ids: List[str], project_name: Optional[str] = None) -> List[str]:
        """
        Download and convert multiple runs to TensorBoard format.
        
        Args:
            run_ids: List of wandb run IDs
            project_name: Optional project name
            
        Returns:
            List of paths to TensorBoard log directories
        """
        converted_paths = []
        
        for run_id in run_ids:
            logger.info(f"Processing run {run_id}")
            
            # Download run
            run_data = self.download_run(run_id, project_name)
            if run_data is None:
                logger.warning(f"Skipping run {run_id} - could not download")
                continue
            
            # Convert to TensorBoard
            try:
                tb_path = self.convert_to_tensorboard(run_data)
                converted_paths.append(tb_path)
            except Exception as e:
                logger.error(f"Error converting run {run_id}: {e}")
        
        return converted_paths
    
    def list_available_runs(self, project_name: Optional[str] = None, limit: int = 50) -> List[Dict[str, Any]]:
        """
        List available runs for easier selection.
        
        Args:
            project_name: Optional project name
            limit: Maximum number of runs to list
            
        Returns:
            List of run information dictionaries
        """
        try:
            api = wandb.Api()
            
            if project_name:
                runs = api.runs(project_name)
            else:
                runs = api.runs()
            
            run_info = []
            for run in runs[:limit]:
                run_info.append({
                    'id': run.id,
                    'name': run.name,
                    'project': run.project,
                    'state': run.state,
                    'created_at': run.created_at,
                    'tags': run.tags
                })
            
            return run_info
            
        except Exception as e:
            logger.error(f"Error listing runs: {e}")
            return []


def main():
    parser = argparse.ArgumentParser(description="Download wandb runs and convert to TensorBoard format")
    parser.add_argument("--run-ids", nargs="+", help="List of wandb run IDs to download")
    parser.add_argument("--project", help="Wandb project name (optional)")
    parser.add_argument("--output-dir", default="./tensorboard_logs", help="Output directory for TensorBoard logs")
    parser.add_argument("--list-runs", action="store_true", help="List available runs instead of downloading")
    parser.add_argument("--limit", type=int, default=50, help="Limit number of runs to list")
    
    args = parser.parse_args()
    
    converter = WandbToTensorboardConverter(args.output_dir)
    
    if args.list_runs:
        logger.info("Listing available runs...")
        runs = converter.list_available_runs(args.project, args.limit)
        
        if runs:
            print("
Available runs:")
            print("-" * 80)
            for run in runs:
                created = datetime.fromtimestamp(run['created_at']).strftime('%Y-%m-%d %H:%M:%S')
                print(f"ID: {run['id']}")
                print(f"Name: {run['name']}")
                print(f"Project: {run['project']}")
                print(f"State: {run['state']}")
                print(f"Created: {created}")
                if run['tags']:
                    print(f"Tags: {', '.join(run['tags'])}")
                print("-" * 80)
        else:
            print("No runs found or error occurred.")
    
    elif args.run_ids:
        logger.info(f"Downloading and converting {len(args.run_ids)} runs...")
        converted_paths = converter.download_and_convert(args.run_ids, args.project)
        
        if converted_paths:
            print(f"
Successfully converted {len(converted_paths)} runs:")
            for path in converted_paths:
                print(f"  {path}")
            
            print(f"
To view with TensorBoard, run:")
            print(f"tensorboard --logdir {args.output_dir}")
        else:
            print("No runs were successfully converted.")
    
    else:
        parser.print_help()


if __name__ == "__main__":
    main()