gogo2/web/tensorboard_component.py

#!/usr/bin/env python3
"""
TensorBoard Component for Dashboard

This module provides a Dash component that embeds TensorBoard in the dashboard.
"""

import dash
from dash import html, dcc
import dash_bootstrap_components as dbc
import logging
from typing import Optional, Dict, Any

logger = logging.getLogger(__name__)

def create_tensorboard_tab(tensorboard_url: str = "http://localhost:6006") -> html.Div:
    """
    Create a dashboard tab that embeds TensorBoard
    
    Args:
        tensorboard_url: URL of the TensorBoard server
        
    Returns:
        html.Div: Dash component containing TensorBoard iframe
    """
    return html.Div([
        dbc.Alert([
            html.I(className="fas fa-chart-line me-2"),
            "TensorBoard Training Visualization",
            html.A(
                "Open in New Window",
                href=tensorboard_url,
                target="_blank",
                className="ms-2 btn btn-sm btn-primary"
            )
        ], color="info", className="mb-3"),
        
        # TensorBoard iframe
        html.Iframe(
            src=tensorboard_url,
            style={
                'width': '100%',
                'height': '800px',
                'border': 'none'
            }
        ),
        
        # Training metrics summary
        html.Div([
            html.H5("Training Metrics Summary", className="mt-3"),
            html.Div(id="training-metrics-summary", className="mt-2")
        ], className="mt-3")
    ])

def create_training_metrics_card() -> dbc.Card:
    """
    Create a card displaying key training metrics
    
    Returns:
        dbc.Card: Dash Bootstrap card component
    """
    return dbc.Card([
        dbc.CardHeader([
            html.I(className="fas fa-brain me-2"),
            "Training Metrics"
        ]),
        dbc.CardBody([
            dbc.Row([
                dbc.Col([
                    html.H6("Model Status"),
                    html.Div(id="model-training-status", children="Initializing...")
                ], width=6),
                dbc.Col([
                    html.H6("Training Progress"),
                    dbc.Progress(id="training-progress-bar", value=0, className="mb-2"),
                    html.Div(id="training-progress-text", children="0%")
                ], width=6)
            ], className="mb-3"),
            
            dbc.Row([
                dbc.Col([
                    html.H6("Loss"),
                    html.Div(id="training-loss-value", children="N/A")
                ], width=4),
                dbc.Col([
                    html.H6("Reward"),
                    html.Div(id="training-reward-value", children="N/A")
                ], width=4),
                dbc.Col([
                    html.H6("State Quality"),
                    html.Div(id="training-state-quality", children="N/A")
                ], width=4)
            ], className="mb-3"),
            
            dbc.Row([
                dbc.Col([
                    html.A(
                        dbc.Button([
                            html.I(className="fas fa-chart-line me-2"),
                            "Open TensorBoard"
                        ], color="primary", size="sm", className="w-100"),
                        href="http://localhost:6006",
                        target="_blank"
                    )
                ], width=12)
            ])
        ])
    ], className="mb-3")

def create_tensorboard_status_indicator(tensorboard_url: str = "http://localhost:6006") -> html.Div:
    """
    Create a status indicator for TensorBoard
    
    Args:
        tensorboard_url: URL of the TensorBoard server
        
    Returns:
        html.Div: Dash component showing TensorBoard status
    """
    return html.Div([
        dbc.Button([
            html.I(className="fas fa-chart-line me-2"),
            "TensorBoard"
        ], 
        id="tensorboard-status-button",
        color="success",
        size="sm",
        href=tensorboard_url,
        target="_blank",
        external_link=True,
        className="ms-2")
    ], id="tensorboard-status-container")

def update_training_metrics_card(metrics: Dict[str, Any]) -> Dict[str, Any]:
    """
    Update training metrics card with latest data
    
    Args:
        metrics: Dictionary of training metrics
        
    Returns:
        Dict: Dictionary of Dash component updates
    """
    # Extract metrics
    training_active = metrics.get("training_active", False)
    loss = metrics.get("loss", None)
    reward = metrics.get("reward", None)
    state_quality = metrics.get("state_quality", None)
    progress = metrics.get("progress", 0)
    
    # Format values
    loss_str = f"{loss:.4f}" if loss is not None else "N/A"
    reward_str = f"{reward:.4f}" if reward is not None else "N/A"
    state_quality_str = f"{state_quality:.1%}" if state_quality is not None else "N/A"
    progress_str = f"{progress:.1%}"
    
    # Determine status
    if training_active:
        status = "Training Active"
        status_class = "text-success"
    else:
        status = "Training Inactive"
        status_class = "text-warning"
    
    # Return updates
    return {
        "model-training-status": html.Span(status, className=status_class),
        "training-progress-bar": progress * 100,
        "training-progress-text": progress_str,
        "training-loss-value": loss_str,
        "training-reward-value": reward_str,
        "training-state-quality": state_quality_str
    }