src/server.py

import dash
from dash import dcc, html, Input, Output, State, callback_context
import plotly.graph_objects as go
import webbrowser
from threading import Timer

from src.execution_model import ScheduleConfig, Schedule
from src.strategies import (
    generate_1f1b_schedule,
    generate_zero_bubble_1p_schedule,
    generate_1f1b_overlap_schedule,
    generate_1f1b_interleave_schedule,
    generate_1f1b_interleave_overlap_schedule,
    generate_dualpipe_schedule
)
from src.visualizer import convert_schedule_to_visualization_format, create_pipeline_figure

def open_browser(port):
    webbrowser.open_new(f"http://127.0.0.1:{port}")

STRATEGIES = {
    "1f1b": generate_1f1b_schedule,
    "zb1p": generate_zero_bubble_1p_schedule,
    "1f1b_overlap": generate_1f1b_overlap_schedule,
    "1f1b_interleave": generate_1f1b_interleave_schedule,
    "1f1b_interleave_overlap": generate_1f1b_interleave_overlap_schedule,
    "dualpipe": generate_dualpipe_schedule,
}

app = dash.Dash(__name__, suppress_callback_exceptions=True)
app.title = "Pipeline Parallelism Visualizer"

# Initial default values
default_values = {
    "num_devices": 4,
    "num_stages": 8,
    "num_batches": 16,
    "p2p_latency": 0.1,
    "op_time_forward": 1.0,
    "op_time_backward_d": 1.0,
    "op_time_backward_w": 1.0,
    "op_time_backward": 2.0,
    "strategy": "1f1b_interleave",
    "split_backward": False,
    "placement_strategy": "interleave"
}

app.layout = html.Div([
    html.H1("Pipeline Parallelism Schedule Visualizer", style={'textAlign': 'center'}),

    html.Div([
        html.Div([
            html.Label("Number of Devices (GPUs):"),
            dcc.Input(id='num_devices', type='number', value=default_values["num_devices"], min=1, step=1, style={'width': '100%'}),

            html.Label("Number of Stages (Model Chunks):"),
            dcc.Input(id='num_stages', type='number', value=default_values["num_stages"], min=1, step=1, style={'width': '100%'}),

            html.Label("Number of Microbatches:"),
            dcc.Input(id='num_batches', type='number', value=default_values["num_batches"], min=1, step=1, style={'width': '100%'}),

            html.Label("P2P Latency (ms):"),
            dcc.Input(id='p2p_latency', type='number', value=default_values["p2p_latency"], min=0, step=0.01, style={'width': '100%'}),

        ], style={'padding': 10, 'flex': 1}),

        html.Div([
            html.Label("Scheduling Strategy:"),
            dcc.Dropdown(
                id='strategy',
                options=[{'label': k, 'value': k} for k in STRATEGIES.keys()],
                value=default_values["strategy"],
                clearable=False,
                style={'width': '100%'}
            ),

            html.Label("Placement Strategy:"),
            dcc.Dropdown(
                id='placement_strategy',
                options=[
                    {'label': 'Standard', 'value': 'standard'},
                    {'label': 'Interleave', 'value': 'interleave'},
                    {'label': 'DualPipe', 'value': 'dualpipe'}
                ],
                value=default_values["placement_strategy"],
                clearable=False,
                style={'width': '100%'}
            ),

            html.Div([ # Wrap checkbox and label
                dcc.Checklist(
                    id='split_backward',
                    options=[{'label': ' Split Backward Pass (for ZB-1P, DualPipe)', 'value': 'True'}],
                    value=['True'] if default_values["split_backward"] else [],
                    style={'display': 'inline-block'}
                ),
            ], style={'marginTop': '20px'}),

        ], style={'padding': 10, 'flex': 1}),

        html.Div([
            html.Label("Operation Time - Forward (ms):"),
            dcc.Input(id='op_time_forward', type='number', value=default_values["op_time_forward"], min=0.01, step=0.01, style={'width': '100%'}),

            html.Label("Operation Time - Backward (ms):"),
            dcc.Input(id='op_time_backward', type='number', value=default_values["op_time_backward"], min=0.01, step=0.01, style={'width': '100%'}),

            html.Label("Operation Time - Backward D (Data Grad) (ms):"),
            dcc.Input(id='op_time_backward_d', type='number', value=default_values["op_time_backward_d"], min=0.01, step=0.01, style={'width': '100%'}),

            html.Label("Operation Time - Backward W (Weight Grad) (ms):"),
            dcc.Input(id='op_time_backward_w', type='number', value=default_values["op_time_backward_w"], min=0.01, step=0.01, style={'width': '100%'}),
        ], style={'padding': 10, 'flex': 1}),

    ], style={'display': 'flex', 'flexDirection': 'row'}),

    html.Div([
        html.Button('Generate Schedule', id='generate-button', n_clicks=0, style={'margin': '20px auto', 'display': 'block'}),
    ]),

    html.Div(id='error-message', style={'color': 'red', 'textAlign': 'center', 'marginTop': '10px'}),

    dcc.Loading(
        id="loading-graph",
        type="circle",
        children=dcc.Graph(id='pipeline-graph', figure=go.Figure())
    )
])

@app.callback(
    Output('pipeline-graph', 'figure'),
    Output('error-message', 'children'),
    Input('generate-button', 'n_clicks'),
    State('num_devices', 'value'),
    State('num_stages', 'value'),
    State('num_batches', 'value'),
    State('p2p_latency', 'value'),
    State('op_time_forward', 'value'),
    State('op_time_backward', 'value'),
    State('op_time_backward_d', 'value'),
    State('op_time_backward_w', 'value'),
    State('strategy', 'value'),
    State('split_backward', 'value'),
    State('placement_strategy', 'value'),
    prevent_initial_call=True
)
def update_graph(n_clicks, num_devices, num_stages, num_batches, p2p_latency,
                 op_time_forward, op_time_backward, op_time_backward_d, op_time_backward_w,
                 strategy, split_backward_list, placement_strategy):

    error_message = ""
    fig = go.Figure()

    split_backward = 'True' in split_backward_list

    # Basic Validations
    if not all([num_devices, num_stages, num_batches, op_time_forward]):
        return fig, "Missing required input values."
    if split_backward and not all([op_time_backward_d, op_time_backward_w]):
        return fig, "Backward D and Backward W times are required when 'Split Backward' is checked."
    if not split_backward and not op_time_backward:
        return fig, "Backward time is required when 'Split Backward' is unchecked."
    if num_stages % num_devices != 0 and placement_strategy != 'dualpipe':
         return fig, "Number of Stages must be divisible by Number of Devices for standard/interleave placement."
    if placement_strategy == 'dualpipe' and num_stages % 2 != 0:
        return fig, "DualPipe requires an even number of stages."
    if placement_strategy == 'dualpipe' and num_stages != num_devices:
        return fig, "DualPipe requires Number of Stages to be equal to Number of Devices."
    if strategy == 'dualpipe' and not split_backward:
        return fig, "DualPipe strategy currently requires 'Split Backward' to be checked."
    if strategy == 'dualpipe' and placement_strategy != 'dualpipe':
        return fig, "DualPipe strategy requires 'DualPipe' placement strategy."
    if strategy == 'zb1p' and not split_backward:
        return fig, "ZB-1P strategy requires 'Split Backward' to be checked."

    try:
        op_times = {
            "forward": float(op_time_forward),
        }
        if split_backward:
            op_times["backward_D"] = float(op_time_backward_d)
            op_times["backward_W"] = float(op_time_backward_w)
            # Add combined backward time for compatibility if needed by some visualization or calculation
            op_times["backward"] = float(op_time_backward_d) + float(op_time_backward_w)
        else:
            op_times["backward"] = float(op_time_backward)

        config = ScheduleConfig(
            num_devices=int(num_devices),
            num_stages=int(num_stages),
            num_batches=int(num_batches),
            p2p_latency=float(p2p_latency),
            placement_strategy=placement_strategy,
            split_backward=split_backward,
            op_times=op_times,
        )

        schedule_func = STRATEGIES.get(strategy)
        if not schedule_func:
            raise ValueError(f"Invalid strategy selected: {strategy}")

        schedule = schedule_func(config)
        schedule.execute() # Calculate start/end times

        vis_data = convert_schedule_to_visualization_format(schedule)
        fig = create_pipeline_figure(vis_data, show_progress=False) # Disable progress bar in server mode

    except AssertionError as e:
        error_message = f"Configuration Error: {e}"
        fig = go.Figure() # Return empty figure on error
    except ValueError as e:
        error_message = f"Input Error: {e}"
        fig = go.Figure()
    except Exception as e:
        error_message = f"An unexpected error occurred: {e}"
        fig = go.Figure()

    return fig, error_message

if __name__ == '__main__':
    port = 8050
    # Timer(1, open_browser, args=(port,)).start() # Optional: automatically open browser
    print(f"Dash server running on http://127.0.0.1:{port}")
    app.run_server(debug=True, port=port)