[tab] convert_wandb_logs_to_tb.py (python)
#!/usr/bin/env python3
"""
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 json
import logging
from pathlib import Path
from typing import List, Optional, Dict, Any
from dataclasses import dataclass
import wandb
import tensorflow as tf
import pandas as pd
from datetime import datetime
import tyro

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


@dataclass
class Args:
    # Run selection
    project: str = "instant-uv/jafar"
    all_runs: bool = True
    
    # Output settings
    output_dir: str = "./logs"
    
    # Listing options
    list_runs: bool = False
    limit: int = 1000


class WandbToTensorboardConverter:
    """Convert wandb runs to TensorBoard format for local storage."""
    
    def __init__(self, output_dir: str = "./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]
                            # Clean column name for TensorBoard (remove special characters)
                            clean_column = column.replace('/', '_').replace('\\', '_').replace(' ', '_')
                            
                            try:
                                if isinstance(value, (int, float)):
                                    tf.summary.scalar(clean_column, float(value), step=int(step))
                                elif isinstance(value, str):
                                    tf.summary.text(clean_column, value, step=int(step))
                                # Add more data type handlers as needed
                            except Exception as e:
                                logger.warning(f"Could not log {clean_column}: {e}")
                                continue
        
        # Add config as text
        if run_data['config']:
            try:
                config_text = json.dumps(run_data['config'], indent=2, default=str)
                with writer.as_default():
                    tf.summary.text("config", config_text, step=0)
            except Exception as e:
                logger.warning(f"Could not serialize config: {e}")
                # Try to convert config to string representation
                config_str = str(run_data['config'])
                with writer.as_default():
                    tf.summary.text("config", config_str, step=0)
        
        # Add summary as text
        if run_data['summary']:
            # Convert wandb summary to a serializable format
            summary_dict = {}
            for key, value in run_data['summary'].items():
                try:
                    # Try to convert to a basic type
                    if hasattr(value, '__dict__'):
                        summary_dict[key] = str(value)
                    else:
                        summary_dict[key] = value
                except:
                    summary_dict[key] = str(value)
            
            summary_text = json.dumps(summary_dict, indent=2, default=str)
            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 get_all_run_ids(self, project_name: str) -> List[str]:
        """
        Get all run IDs from a project.
        
        Args:
            project_name: Project name
            
        Returns:
            List of run IDs
        """
        try:
            api = wandb.Api()
            runs = api.runs(project_name)
            run_ids = [run.id for run in runs]
            logger.info(f"Found {len(run_ids)} runs in project {project_name}")
            return run_ids
            
        except Exception as e:
            logger.error(f"Error getting run IDs: {e}")
            return []


def main(args: Args):
    """Download wandb runs and convert to TensorBoard format."""
    
    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.all_runs:
        logger.info(f"Getting all run IDs from project {args.project}...")
        all_run_ids = converter.get_all_run_ids(args.project)
        
        if all_run_ids:
            logger.info(f"Downloading and converting all {len(all_run_ids)} runs...")
            converted_paths = converter.download_and_convert(all_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:
            print("No runs found in the project.")
    
    
    else:
        print("Please specify one of the following options:")
        print("  --all-runs: Download all runs from the project")
        print("  --run-ids <id1> <id2> ...: Download specific run IDs")
        print("  --list-runs: List available runs")
        print("Use --help for more information.")


if __name__ == "__main__":
    args = tyro.cli(Args)
    main(args)