main.py

#!/usr/bin/env python3
"""
main.py
Main entry point for GPU unschedulable analysis and visualization.

Usage:
    python main.py                    # Analyze jobs2.csv and create monthly reports
    python main.py --csv jobs.csv     # Use different input file
    python main.py --no-show          # Don't open browser, just save files
    python main.py --heatmap-only     # Generate monthly heatmaps (skip full reports)
    python main.py --single-report    # Generate a single full report HTML
    python main.py --dashboard        # Also generate summary dashboard
"""

import argparse
import pickle
import sys
from pathlib import Path

from gpu_analysis import analyze_gpu_unschedulable
from plot_interactive import (
    plot_unschedulable_heatmap,
    plot_summary_dashboard,
    plot_full_report,
    plot_monthly_heatmaps,
    plot_monthly_reports
)


def main():
    parser = argparse.ArgumentParser(
        description='Analyze GPU scheduling bottlenecks in SLURM job data'
    )
    parser.add_argument(
        '--csv', '-c',
        default='jobs2.csv',
        help='Path to the SLURM jobs CSV file (pipe-delimited)'
    )
    parser.add_argument(
        '--output', '-o',
        default='gpu_unschedulable_report.html',
        help='Output HTML file for the report'
    )
    parser.add_argument(
        '--no-show',
        action='store_true',
        help="Don't open visualizations in browser"
    )
    parser.add_argument(
        '--dashboard',
        action='store_true',
        help='Also generate summary dashboard'
    )
    parser.add_argument(
        '--heatmap-only',
        action='store_true',
        help='Generate heatmaps only (skip full reports)'
    )
    parser.add_argument(
        '--single-report',
        action='store_true',
        help='Generate a single full report HTML instead of monthly heatmaps'
    )
    parser.add_argument(
        '--monthly-output-dir',
        default='monthly_plots',
        help='Output directory for monthly heatmap HTML files'
    )
    parser.add_argument(
        '--save-pickle',
        action='store_true',
        help='Save analysis results to pickle file'
    )
    parser.add_argument(
        '--load-pickle',
        type=str,
        default=None,
        help='Load analysis results from pickle file instead of CSV'
    )
    parser.add_argument(
        '--quiet', '-q',
        action='store_true',
        help='Suppress progress messages'
    )
    
    args = parser.parse_args()
    
    # Run analysis or load from pickle
    if args.load_pickle:
        print(f"Loading analysis results from {args.load_pickle}...")
        with open(args.load_pickle, 'rb') as f:
            results = pickle.load(f)
        print("Results loaded successfully.")
    else:
        # Check if input file exists
        if not Path(args.csv).exists():
            print(f"Error: Input file '{args.csv}' not found.")
            sys.exit(1)
        
        # Run analysis
        results = analyze_gpu_unschedulable(args.csv, verbose=not args.quiet)
        
        # Save to pickle if requested
        if args.save_pickle:
            pickle_file = Path(args.csv).stem + '_analysis_results.pkl'
            with open(pickle_file, 'wb') as f:
                pickle.dump(results, f)
            print(f"\nResults saved to {pickle_file}")
    
    # Create report or heatmap visualization
    if args.single_report:
        if args.heatmap_only:
            print("\nGenerating interactive heatmap...")
            plot_unschedulable_heatmap(
                results,
                output_file=args.output,
                show=not args.no_show
            )
        else:
            print("\nGenerating full report...")
            plot_full_report(
                results,
                output_file=args.output,
                show=not args.no_show
            )
    else:
        if args.heatmap_only:
            print("\nGenerating monthly heatmaps...")
            plot_monthly_heatmaps(
                results,
                output_dir=args.monthly_output_dir,
                show=not args.no_show
            )
        else:
            print("\nGenerating monthly reports...")
            plot_monthly_reports(
                results,
                output_dir=args.monthly_output_dir,
                show=not args.no_show
            )
    
    # Create dashboard if requested
    if args.dashboard:
        print("\nGenerating summary dashboard...")
        dashboard_file = Path(args.output).stem + '_dashboard.html'
        plot_summary_dashboard(
            results,
            output_file=dashboard_file,
            show=not args.no_show
        )
    
    print("\nDone!")


if __name__ == '__main__':
    main()