Improve UI components.

pyproject.toml

@@ -24,6 +24,7 @@ dependencies = [
     "pandas>=2.1.0",
     "numpy>=1.26.0",
     "tqdm>=4.67.0",
+    "dash-bootstrap-components>=1.7.1",
 ]
 
 [project.optional-dependencies]
@@ -64,4 +65,4 @@ disallow_incomplete_defs = true
 
 [tool.pytest]
 testpaths = ["tests"]
-pythonpath = ["."] 
+pythonpath = ["."] 

src/server.py

@@ -1,4 +1,5 @@
 import dash
+import dash_bootstrap_components as dbc
 from dash import dcc, html, Input, Output, State, callback_context
 import plotly.graph_objects as go
 import webbrowser
@@ -27,8 +28,8 @@ STRATEGIES = {
     "dualpipe": generate_dualpipe_schedule,
 }
 
-app = dash.Dash(__name__, suppress_callback_exceptions=True)
-app.title = "Pipeline Parallelism Visualizer"
+app = dash.Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP], suppress_callback_exceptions=True)
+app.title = "Pipeline Parallelism Schedule Visualizer"
 
 # Initial default values
 default_values = {
@@ -45,91 +46,98 @@ default_values = {
     "placement_strategy": "interleave"
 }
 
-app.layout = html.Div([
-    html.H1("Pipeline Parallelism Schedule Visualizer", style={'textAlign': 'center'}),
-
-    html.Div([
+# Define input groups using dbc components
+basic_params_card = dbc.Card(
+    dbc.CardBody([
+        html.H5("Basic Parameters", className="card-title"),
         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}),
+            dbc.Label("Number of Devices (GPUs):"),
+            dbc.Input(id='num_devices', type='number', value=default_values["num_devices"], min=1, step=1),
+        ], className="mb-3"),
+        html.Div([
+            dbc.Label("Number of Stages (Model Chunks):"),
+            dbc.Input(id='num_stages', type='number', value=default_values["num_stages"], min=1, step=1),
+        ], className="mb-3"),
+        html.Div([
+            dbc.Label("Number of Microbatches:"),
+            dbc.Input(id='num_batches', type='number', value=default_values["num_batches"], min=1, step=1),
+        ], className="mb-3"),
+        html.Div([
+            dbc.Label("P2P Latency (ms):"),
+            dbc.Input(id='p2p_latency', type='number', value=default_values["p2p_latency"], min=0, step=0.01),
+        ], className="mb-3"),
+    ])
+)
 
+scheduling_params_card = dbc.Card(
+    dbc.CardBody([
+        html.H5("Scheduling Parameters", className="card-title"),
         html.Div([
-            html.Label("Scheduling Strategy:"),
-            dcc.Dropdown(
-                id='strategy',
+            dbc.Label("Scheduling Strategies:"),
+            dbc.Checklist(
+                id='strategy-checklist',
                 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%'}
+                value=[default_values["strategy"]],
+                inline=False,
             ),
+        ], className="mb-3"),
+    ])
+)
 
-            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}),
-
+timing_params_card = dbc.Card(
+    dbc.CardBody([
+        html.H5("Operation Timing (ms)", className="card-title"),
         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}),
+            dbc.Label("Forward:"),
+            dbc.Input(id='op_time_forward', type='number', value=default_values["op_time_forward"], min=0.01, step=0.01),
+        ], className="mb-3"),
+        html.Div([
+            dbc.Label("Backward (Combined):"),
+            dbc.Input(id='op_time_backward', type='number', value=default_values["op_time_backward"], min=0.01, step=0.01),
+            dbc.FormText("Used when strategy does NOT require split backward."),
+        ], className="mb-3"),
+        html.Div([
+            dbc.Label("Backward D (Data Grad):"),
+            dbc.Input(id='op_time_backward_d', type='number', value=default_values["op_time_backward_d"], min=0.01, step=0.01),
+            dbc.FormText("Used when strategy requires split backward (e.g., ZB-1P, DualPipe)."),
+        ], className="mb-3"),
+        html.Div([
+            dbc.Label("Backward W (Weight Grad):"),
+            dbc.Input(id='op_time_backward_w', type='number', value=default_values["op_time_backward_w"], min=0.01, step=0.01),
+            dbc.FormText("Used when strategy requires split backward (e.g., ZB-1P, DualPipe)."),
+        ], className="mb-3"),
+    ])
+)
 
-    ], style={'display': 'flex', 'flexDirection': 'row'}),
+# Updated app layout using dbc components and structure
+app.layout = dbc.Container([
+    html.H1("Pipeline Parallelism Schedule Visualizer", className="my-4 text-center"),
 
-    html.Div([
-        html.Button('Generate Schedule', id='generate-button', n_clicks=0, style={'margin': '20px auto', 'display': 'block'}),
+    dbc.Row([
+        dbc.Col(basic_params_card, md=4),
+        dbc.Col(scheduling_params_card, md=4),
+        dbc.Col(timing_params_card, md=4),
     ]),
 
-    html.Div(id='error-message', style={'color': 'red', 'textAlign': 'center', 'marginTop': '10px'}),
+    dbc.Row([
+        dbc.Col([
+            dbc.Button('Generate Schedule', id='generate-button', n_clicks=0, color="primary", className="mt-4"),
+        ], className="text-center")
+    ]),
 
-    dcc.Loading(
-        id="loading-graph",
-        type="circle",
-        children=dcc.Graph(id='pipeline-graph', figure=go.Figure())
-    )
-])
+    dbc.Row([
+        dbc.Col([
+            dcc.Loading(
+                id="loading-graph-area",
+                type="circle",
+                children=html.Div(id='graph-output-container', className="mt-4")
+            )
+        ])
+    ])
+], fluid=True)
 
 @app.callback(
-    Output('pipeline-graph', 'figure'),
-    Output('error-message', 'children'),
+    Output('graph-output-container', 'children'),
     Input('generate-button', 'n_clicks'),
     State('num_devices', 'value'),
     State('num_stages', 'value'),
@@ -139,83 +147,99 @@ app.layout = html.Div([
     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'),
+    State('strategy-checklist', '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):
+                 selected_strategies):
 
-    error_message = ""
-    fig = go.Figure()
+    output_components = []
 
-    split_backward = 'True' in split_backward_list
+    if not selected_strategies:
+        return [dbc.Alert("Please select at least one scheduling strategy.", color="warning")]
 
-    # 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 [dbc.Alert("Missing required basic input values (Devices, Stages, Batches, Forward Time).", color="danger")]
 
-    return fig, error_message
+    for strategy in selected_strategies:
+        error_message = ""
+        fig = go.Figure()
+        placement_strategy = ""
+
+        split_backward = strategy in ["zb1p", "dualpipe"]
+
+        if split_backward and not all([op_time_backward_d, op_time_backward_w]):
+            error_message = f"Strategy '{strategy}': Backward D and Backward W times are required."
+        elif not split_backward and not op_time_backward:
+            error_message = f"Strategy '{strategy}': Combined Backward time is required."
+
+        if not error_message:
+            if strategy in ["1f1b", "1f1b_overlap", "zb1p"]:
+                placement_strategy = "standard"
+                if num_devices != num_stages:
+                    error_message = f"Strategy '{strategy}': Requires Number of Stages == Number of Devices."
+            elif strategy in ["1f1b_interleave", "1f1b_interleave_overlap"]:
+                placement_strategy = "interleave"
+                if num_stages % num_devices != 0:
+                    error_message = f"Strategy '{strategy}': Requires Number of Stages to be divisible by Number of Devices."
+            elif strategy == "dualpipe":
+                placement_strategy = "dualpipe"
+                if num_stages % 2 != 0:
+                    error_message = f"Strategy '{strategy}' (DualPipe): Requires an even number of stages."
+                elif num_stages != num_devices:
+                    error_message = f"Strategy '{strategy}' (DualPipe): Requires Number of Stages == Number of Devices."
+
+        if not error_message:
+            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)
+                    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 function for: {strategy}")
+
+                schedule = schedule_func(config)
+                schedule.execute()
+
+                vis_data = convert_schedule_to_visualization_format(schedule)
+                fig = create_pipeline_figure(vis_data, show_progress=False)
+
+                output_components.append(html.Div([
+                    html.H4(f"Schedule: {strategy}", className="text-center mt-3 mb-2"),
+                    dcc.Graph(figure=fig)
+                ]))
+
+            except (AssertionError, ValueError, TypeError) as e:
+                 error_message = f"Error generating schedule for '{strategy}': {e}"
+                 import traceback
+                 traceback.print_exc()
+            except Exception as e:
+                 error_message = f"An unexpected error occurred for '{strategy}': {e}"
+                 import traceback
+                 traceback.print_exc()
+
+        if error_message:
+             output_components.append(
+                 dbc.Alert(error_message, color="danger", className="mt-3")
+             )
+
+    return output_components
 
 if __name__ == '__main__':
     port = 8050