Update UI.

app.py

@@ -1,6 +1,6 @@
 import dash
 import dash_bootstrap_components as dbc
-from dash import dcc, html, Input, Output, State, callback_context
+from dash import dcc, html, Input, Output, State, callback_context, ALL, ClientsideFunction
 import plotly.graph_objects as go
 
 from src.execution_model import ScheduleConfig, Schedule
@@ -23,7 +23,7 @@ STRATEGIES = {
     "dualpipe": generate_dualpipe_schedule,
 }
 
-app = dash.Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP], suppress_callback_exceptions=True)
+app = dash.Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP, dbc.icons.BOOTSTRAP], suppress_callback_exceptions=True)
 app.title = "Pipeline Parallelism Schedule Visualizer"
 
 # Initial default values
@@ -36,107 +36,321 @@ default_values = {
     "op_time_backward_d": 1.0,
     "op_time_backward_w": 1.0,
     "op_time_backward": 2.0,
-    "strategy": "1f1b_interleave",
+    "strategy": ["1f1b_interleave"],
     "op_time_overlapped_fwd_bwd": None,
 }
 
 # Define input groups using dbc components
+card_style = {"marginBottom": "24px"}
+
 basic_params_card = dbc.Card(
     dbc.CardBody([
-        html.H5("Basic Parameters", className="card-title"),
-        html.Div([
-            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.H5("Basic Parameters", className="card-title mb-4"),
         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),
+            dbc.Label("Number of Devices (GPUs)", html_for='num_devices', className="form-label"),
+            dbc.Input(id='num_devices', type='number', value=default_values["num_devices"], min=1, step=1, required=True),
+            dbc.FormFeedback("Please provide a positive integer for the number of devices.", type="invalid", id="feedback-num_devices"),
         ], 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),
+            dbc.Label("Number of Stages (Model Chunks)", html_for='num_stages', className="form-label"),
+            dbc.Input(id='num_stages', type='number', value=default_values["num_stages"], min=1, step=1, required=True),
+            dbc.FormFeedback("Please provide a positive integer for the number of stages.", type="invalid", id="feedback-num_stages"),
         ], 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),
+            dbc.Label("Number of Microbatches", html_for='num_batches', className="form-label"),
+            dbc.Input(id='num_batches', type='number', value=default_values["num_batches"], min=1, step=1, required=True),
+            dbc.FormFeedback("Please provide a positive integer for the number of microbatches.", type="invalid", id="feedback-num_batches"),
         ], className="mb-3"),
-    ])
+    ]),
+    style=card_style
 )
 
 scheduling_params_card = dbc.Card(
     dbc.CardBody([
-        html.H5("Scheduling Parameters", className="card-title"),
-        html.Div([
-            dbc.Label("Scheduling Strategies:"),
-            dbc.Checklist(
-                id='strategy-checklist',
-                options=[{'label': k, 'value': k} for k in STRATEGIES.keys()],
-                value=list(STRATEGIES.keys()),
-                inline=False,
-            ),
-        ], className="mb-3"),
-    ])
+        html.H5("Scheduling Strategy", className="card-title mb-4"),
+        dbc.ButtonGroup(
+            [
+                dbc.Button(
+                    strategy,
+                    id={"type": "strategy-button", "index": strategy},
+                    color="secondary",
+                    outline=True,
+                    active=strategy in default_values["strategy"],
+                    className="me-1"
+                 )
+                for strategy in STRATEGIES.keys()
+            ],
+            className="d-flex flex-wrap"
+        ),
+        dcc.Store(id='selected-strategies-store', data=default_values["strategy"]),
+        html.Div(id='strategy-selection-feedback', className='invalid-feedback d-block mt-2')
+    ]),
+    style=card_style
 )
 
 timing_params_card = dbc.Card(
     dbc.CardBody([
-        html.H5("Operation Timing (ms)", className="card-title"),
-        html.Div([
-            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.H5("Operation Timing (ms)", className="card-title mb-4"),
         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)."),
+            html.Div([
+                dbc.Label("P2P Latency", html_for='p2p_latency', className="form-label d-inline-block me-1"),
+                html.I(className="bi bi-info-circle", id="tooltip-target-p2p", style={"cursor": "pointer"})
+            ]),
+            dbc.Input(id='p2p_latency', type='number', value=default_values["p2p_latency"], min=0, step=0.01, required=True),
+            dbc.FormFeedback("P2P latency must be a number >= 0.", type="invalid", id="feedback-p2p_latency"),
+            dbc.Tooltip(
+                "Time (ms) for point-to-point communication between adjacent devices.",
+                target="tooltip-target-p2p",
+                placement="right"
+            )
         ], 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)."),
+            html.Div([
+                dbc.Label("Forward Operation Time", html_for='op_time_forward', className="form-label d-inline-block me-1"),
+                html.I(className="bi bi-info-circle", id="tooltip-target-fwd", style={"cursor": "pointer"})
+            ]),
+            dbc.Input(id='op_time_forward', type='number', value=default_values["op_time_forward"], min=0.01, step=0.01, required=True),
+            dbc.FormFeedback("Forward time must be a number > 0.", type="invalid", id="feedback-op_time_forward"),
+            dbc.Tooltip(
+                "Time (ms) for a single forward pass of one microbatch through one stage.",
+                target="tooltip-target-fwd",
+                placement="right"
+            )
         ], className="mb-3"),
         html.Div([
-            dbc.Label("Overlapped Forward+Backward:"),
-            dbc.Input(id='op_time_overlapped_fwd_bwd', type='number', placeholder="Optional: Defaults to Fwd + Bwd times", min=0.01, step=0.01, value=default_values["op_time_overlapped_fwd_bwd"]),
-            dbc.FormText("Specify a custom duration if Forward and Backward ops overlap completely."),
+            html.Div([
+                dbc.Label("Backward (Combined)", html_for='op_time_backward', className="form-label d-inline-block me-1"),
+                html.I(className="bi bi-info-circle", id="tooltip-target-bwd", style={"cursor": "pointer"})
+            ]),
+            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."),
+            dbc.FormFeedback("Backward time must be > 0 if specified.", type="invalid", id="feedback-op_time_backward"),
+            dbc.Tooltip(
+                "Time (ms) for a combined backward pass (data gradient + weight gradient) of one microbatch through one stage.",
+                target="tooltip-target-bwd",
+                placement="right"
+            )
         ], className="mb-3"),
-    ])
+
+        # --- Collapsible Advanced Options (Item 3) ---
+        html.Hr(className="my-3"),
+        dbc.Switch(
+            id="advanced-timing-switch",
+            label="Show Advanced Timing Options",
+            value=False,
+            className="mb-3"
+        ),
+        dbc.Collapse(
+            id="advanced-timing-collapse",
+            is_open=False,
+            children=[
+                html.Div([
+                    html.Div([
+                        dbc.Label("Backward D (Data Grad)", html_for='op_time_backward_d', className="form-label d-inline-block me-1"),
+                        html.I(className="bi bi-info-circle", id="tooltip-target-bwd-d", style={"cursor": "pointer"})
+                    ]),
+                    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)."),
+                    dbc.FormFeedback("Backward D time must be > 0 if specified.", type="invalid", id="feedback-op_time_backward_d"),
+                    dbc.Tooltip(
+                        "Time (ms) for the data gradient part of the backward pass.",
+                        target="tooltip-target-bwd-d",
+                        placement="right"
+                    )
+                ], className="mb-3"),
+                html.Div([
+                    html.Div([
+                        dbc.Label("Backward W (Weight Grad)", html_for='op_time_backward_w', className="form-label d-inline-block me-1"),
+                        html.I(className="bi bi-info-circle", id="tooltip-target-bwd-w", style={"cursor": "pointer"})
+                    ]),
+                    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)."),
+                    dbc.FormFeedback("Backward W time must be > 0 if specified.", type="invalid", id="feedback-op_time_backward_w"),
+                    dbc.Tooltip(
+                        "Time (ms) for the weight gradient part of the backward pass.",
+                        target="tooltip-target-bwd-w",
+                        placement="right"
+                    )
+                ], className="mb-3"),
+                html.Div([
+                    html.Div([
+                        dbc.Label("Overlapped Forward+Backward", html_for='op_time_overlapped_fwd_bwd', className="form-label d-inline-block me-1"),
+                        html.I(className="bi bi-info-circle", id="tooltip-target-overlap", style={"cursor": "pointer"})
+                    ]),
+                    dbc.Input(id='op_time_overlapped_fwd_bwd', type='number', placeholder="Defaults to Fwd + Bwd", min=0.01, step=0.01, value=default_values["op_time_overlapped_fwd_bwd"]),
+                    dbc.FormText("Specify if Forward and Backward ops overlap completely."),
+                    dbc.FormFeedback("Overlapped time must be > 0 if specified.", type="invalid", id="feedback-op_time_overlapped_fwd_bwd"),
+                    dbc.Tooltip(
+                        "Optional: Specify a single time (ms) if the forward and backward passes for a microbatch can be fully overlapped within the same stage execution slot.",
+                        target="tooltip-target-overlap",
+                        placement="right"
+                    )
+                ], className="mb-3"),
+            ]
+        )
+    ]),
+    style=card_style
 )
 
 # Updated app layout using dbc components and structure
 app.layout = dbc.Container([
     html.H1("Pipeline Parallelism Schedule Visualizer", className="my-4 text-center"),
 
+    # Main Row with Left (Graphs) and Right (Controls) Columns
     dbc.Row([
-        dbc.Col(basic_params_card, md=4),
-        dbc.Col(scheduling_params_card, md=4),
-        dbc.Col(timing_params_card, md=4),
-    ]),
-
-    dbc.Row([
-        dbc.Col([
-            dbc.Button('Generate Schedule', id='generate-button', n_clicks=0, color="primary", className="mt-4"),
-        ], className="text-center")
-    ]),
-
-    dbc.Row([
+        # --- Left Column (Graphs Area) ---
         dbc.Col([
+            # Output Area for Graphs
             dcc.Loading(
                 id="loading-graph-area",
                 type="circle",
-                children=html.Div(id='graph-output-container', className="mt-4")
+                children=html.Div(id='graph-output-container', style={"minHeight": "600px"})
             )
-        ])
-    ])
-], fluid=True)
+        ], lg=8, md=7, sm=12, className="mb-4 mb-lg-0"),
+
+        # --- Right Column (Controls Area) ---
+        dbc.Col([
+            # Parameter Cards Stacked Vertically
+            basic_params_card,
+            scheduling_params_card,
+            timing_params_card,
+
+            # Generate Button below the cards in the right column
+            dbc.Row([
+                dbc.Col(
+                    dbc.Button(
+                        'Generate Schedule',
+                        id='generate-button',
+                        n_clicks=0,
+                        color="primary",
+                        className="w-100",
+                        disabled=False
+                    ),
+                )
+            ], className="mt-3")
+        ], lg=4, md=5, sm=12)
+    ]),
+
+    # --- Toast Container (Positioned Fixed) ---
+    html.Div(id="toast-container", style={"position": "fixed", "top": 20, "right": 20, "zIndex": 1050})
 
+], fluid=True, className="py-4")
+
+# --- Callback for Input Validation and Generate Button State ---
+@app.callback(
+    Output('generate-button', 'disabled'),
+    # Outputs to control the 'invalid' state of Inputs
+    Output('num_devices', 'invalid'),
+    Output('num_stages', 'invalid'),
+    Output('num_batches', 'invalid'),
+    Output('p2p_latency', 'invalid'),
+    Output('op_time_forward', 'invalid'),
+    Output('op_time_backward', 'invalid'),
+    Output('op_time_backward_d', 'invalid'),
+    Output('op_time_backward_w', 'invalid'),
+    Output('op_time_overlapped_fwd_bwd', 'invalid'),
+    # Outputs to control the visibility/content of FormFeedback (can also just control Input's invalid state)
+    # We are primarily using the Input's `invalid` prop which automatically shows/hides associated FormFeedback
+    # Output('feedback-num_devices', 'children'), ... (Add if more specific messages needed per validation type)
+    Output('strategy-selection-feedback', 'children', allow_duplicate=True), # Update feedback from validation callback too
+    # Inputs: Trigger validation whenever any relevant input changes
+    Input('num_devices', 'value'),
+    Input('num_stages', 'value'),
+    Input('num_batches', 'value'),
+    Input('p2p_latency', 'value'),
+    Input('op_time_forward', 'value'),
+    Input('op_time_backward', 'value'),
+    Input('op_time_backward_d', 'value'),
+    Input('op_time_backward_w', 'value'),
+    Input('op_time_overlapped_fwd_bwd', 'value'),
+    Input('selected-strategies-store', 'data'), # Validate strategy selection
+    prevent_initial_call=True # Prevent callback running on page load before user interaction
+)
+def validate_inputs(num_devices, num_stages, num_batches, p2p_latency,
+                    op_time_forward, op_time_backward, op_time_backward_d, op_time_backward_w,
+                    op_time_overlapped_fwd_bwd, selected_strategies):
+    is_invalid = {
+        "num_devices": num_devices is None or num_devices < 1,
+        "num_stages": num_stages is None or num_stages < 1,
+        "num_batches": num_batches is None or num_batches < 1,
+        "p2p_latency": p2p_latency is None or p2p_latency < 0,
+        "op_time_forward": op_time_forward is None or op_time_forward <= 0,
+        "op_time_backward": op_time_backward is not None and op_time_backward <= 0,
+        "op_time_backward_d": op_time_backward_d is not None and op_time_backward_d <= 0,
+        "op_time_backward_w": op_time_backward_w is not None and op_time_backward_w <= 0,
+        "op_time_overlapped_fwd_bwd": op_time_overlapped_fwd_bwd is not None and op_time_overlapped_fwd_bwd <= 0,
+    }
+
+    # Validate strategy selection
+    strategy_feedback = "" # Default empty feedback
+    if not selected_strategies or len(selected_strategies) == 0:
+        is_invalid["strategies"] = True
+        strategy_feedback = "Please select at least one strategy."
+    else:
+        is_invalid["strategies"] = False
+        # Additional validation: Check if required timings are provided for selected strategies
+        needs_split_backward = any(s in ["zb1p", "dualpipe"] for s in selected_strategies)
+        needs_combined_backward = any(s not in ["zb1p", "dualpipe"] for s in selected_strategies)
+
+        if needs_split_backward and (op_time_backward_d is None or op_time_backward_w is None):
+            is_invalid["op_time_backward_d"] = op_time_backward_d is None or op_time_backward_d <= 0
+            is_invalid["op_time_backward_w"] = op_time_backward_w is None or op_time_backward_w <= 0
+            # We might want specific feedback here, but setting invalid=True is often enough
+
+        if needs_combined_backward and op_time_backward is None:
+            is_invalid["op_time_backward"] = op_time_backward is None or op_time_backward <= 0
+
+    # Check if any input is invalid
+    overall_invalid = any(is_invalid.values())
+
+    # Disable button if any validation fails
+    disable_button = overall_invalid
+
+    # Return button state and invalid states for each input
+    return (
+        disable_button,
+        is_invalid["num_devices"],
+        is_invalid["num_stages"],
+        is_invalid["num_batches"],
+        is_invalid["p2p_latency"],
+        is_invalid["op_time_forward"],
+        is_invalid["op_time_backward"],
+        is_invalid["op_time_backward_d"],
+        is_invalid["op_time_backward_w"],
+        is_invalid["op_time_overlapped_fwd_bwd"],
+        strategy_feedback # Update strategy feedback based on validation
+    )
+
+# --- Callback to toggle Advanced Options Collapse ---
 @app.callback(
+    Output("advanced-timing-collapse", "is_open"),
+    Input("advanced-timing-switch", "value"),
+    prevent_initial_call=True,
+)
+def toggle_advanced_options(switch_value):
+    return switch_value
+
+# --- Client-side Callback for Strategy ButtonGroup ---
+app.clientside_callback(
+    ClientsideFunction(
+        namespace='clientside',
+        function_name='update_strategy_selection'
+    ),
+    Output('selected-strategies-store', 'data'),
+    Output({'type': 'strategy-button', 'index': ALL}, 'active'),
+    Output({'type': 'strategy-button', 'index': ALL}, 'color'),
+    Output({'type': 'strategy-button', 'index': ALL}, 'outline'),
+    Output('strategy-selection-feedback', 'children'),
+    Input({'type': 'strategy-button', 'index': ALL}, 'n_clicks'),
+    State('selected-strategies-store', 'data'),
+    prevent_initial_call=True
+)
+
+# --- Main Graph Update Callback ---
+@app.callback(
+    # Output graph container and toast container separately
     Output('graph-output-container', 'children'),
+    Output('toast-container', 'children'), # Output for toasts
     Input('generate-button', 'n_clicks'),
     State('num_devices', 'value'),
     State('num_stages', 'value'),
@@ -147,7 +361,7 @@ app.layout = dbc.Container([
     State('op_time_backward_d', 'value'),
     State('op_time_backward_w', 'value'),
     State('op_time_overlapped_fwd_bwd', 'value'),
-    State('strategy-checklist', 'value'),
+    State('selected-strategies-store', 'data'),
     prevent_initial_call=True
 )
 def update_graph(n_clicks, num_devices, num_stages, num_batches, p2p_latency,
@@ -155,19 +369,39 @@ def update_graph(n_clicks, num_devices, num_stages, num_batches, p2p_latency,
                  op_time_overlapped_fwd_bwd,
                  selected_strategies):
 
-    # Define the desired display order for strategies
     strategy_display_order = ["1f1b", "1f1b_interleave", "1f1b_overlap", "1f1b_interleave_overlap", "dualpipe", "zb1p"]
-    
-    output_components = []
-    valid_results = []  # Store (strategy_name, schedule, vis_data) for valid schedules
-    error_messages = []  # Store (strategy_name, error_message) for errors
-    automatic_adjustments = []  # Store messages about automatic parameter adjustments
+
+    graph_components = [] # Renamed from output_components
+    toast_components = [] # New list for toasts
+    valid_results = []
+    error_messages = []
+    automatic_adjustments = []
+
+    # Use a variable to track if initial validation fails
+    initial_validation_error = None
 
     if not selected_strategies:
-        return [dbc.Alert("Please select at least one scheduling strategy.", color="warning")]
+        initial_validation_error = dbc.Toast(
+            "Please select at least one scheduling strategy.",
+            header="Input Error",
+            icon="warning",
+            duration=4000,
+            is_open=True,
+            className="border-warning"
+        )
+    elif not all([num_devices, num_stages, num_batches, op_time_forward]):
+        initial_validation_error = dbc.Toast(
+            "Missing required basic input values (Devices, Stages, Batches, Forward Time).",
+            header="Input Error",
+            icon="danger",
+            duration=4000,
+            is_open=True,
+            className="border-danger"
+        )
 
-    if not all([num_devices, num_stages, num_batches, op_time_forward]):
-         return [dbc.Alert("Missing required basic input values (Devices, Stages, Batches, Forward Time).", color="danger")]
+    if initial_validation_error:
+        # Return empty graph list and the validation error toast
+        return [], [initial_validation_error]
 
     for strategy in selected_strategies:
         error_message = ""
@@ -179,17 +413,15 @@ def update_graph(n_clicks, num_devices, num_stages, num_batches, p2p_latency,
         
         # Apply automatic adjustments for dualpipe
         if strategy == "dualpipe" and num_stages != num_devices:
-            current_num_stages = num_devices  # Force num_stages = num_devices for dualpipe
-            automatic_adjustments.append(
-                f"Strategy '{strategy}': Number of Stages automatically adjusted to {num_devices} to match Number of Devices."
-            )
+            current_num_stages = num_devices
+            adjustment_msg = f"Strategy '{strategy}': Number of Stages auto-adjusted to {num_devices} to match Devices."
+            automatic_adjustments.append(adjustment_msg)
 
         # Apply automatic adjustments for strategies that require num_stages == num_devices
         if strategy in ["1f1b", "1f1b_overlap", "zb1p"] and num_stages != num_devices:
             current_num_stages = num_devices
-            automatic_adjustments.append(
-                f"Strategy '{strategy}': Number of Stages automatically adjusted to {num_devices} to match Number of Devices."
-            )
+            adjustment_msg = f"Strategy '{strategy}': Number of Stages auto-adjusted to {num_devices} to match Devices."
+            automatic_adjustments.append(adjustment_msg)
 
         split_backward = strategy in ["zb1p", "dualpipe"]
 
@@ -201,41 +433,32 @@ def update_graph(n_clicks, num_devices, num_stages, num_batches, p2p_latency,
         if not error_message:
             if strategy in ["1f1b", "1f1b_overlap", "zb1p"]:
                 placement_strategy = "standard"
-                # No need to check num_stages == num_devices as we've enforced it above
             elif strategy in ["1f1b_interleave", "1f1b_interleave_overlap"]:
                 placement_strategy = "interleave"
                 if current_num_stages % current_num_devices != 0:
-                    error_message = f"Strategy '{strategy}': Requires Number of Stages to be divisible by Number of Devices."
+                    error_message = f"Strategy '{strategy}': Requires Stages divisible by Devices."
             elif strategy == "dualpipe":
                 placement_strategy = "dualpipe"
                 if current_num_stages % 2 != 0:
-                    error_message = f"Strategy '{strategy}' (DualPipe): Requires an even number of stages."
+                    error_message = f"Strategy '{strategy}': Requires an even number of stages."
 
         # Create adjusted operation times based on placement strategy
         if not error_message:
             try:
-                # Calculate number of stages per device for time adjustment
                 stages_per_device = current_num_stages // current_num_devices
-                
-                # Calculate scaling factor - this normalizes operation time by stages per device
-                # For standard placement (1:1 stage:device mapping), this remains 1.0
-                # For interleaved, this scales down the time proportionally
                 time_scale_factor = 1.0 / stages_per_device if stages_per_device > 0 else 1.0
-                
+
                 if stages_per_device > 1:
-                    automatic_adjustments.append(
-                        f"Strategy '{strategy}': Operation times scaled by 1/{stages_per_device} to account for {stages_per_device} stages per device."
-                    )
-                
-                # Apply scaling to operation times
-                op_times = { 
-                    "forward": float(op_time_forward) * time_scale_factor 
-                }
-                
+                     adjustment_msg = f"Strategy '{strategy}': Op times scaled by 1/{stages_per_device} ({stages_per_device} stages/device)."
+                     # Avoid adding duplicate adjustment messages if already added above
+                     if adjustment_msg not in automatic_adjustments:
+                         automatic_adjustments.append(adjustment_msg)
+
+                op_times = { "forward": float(op_time_forward) * time_scale_factor }
+
                 if split_backward:
                     op_times["backward_D"] = float(op_time_backward_d) * time_scale_factor
                     op_times["backward_W"] = float(op_time_backward_w) * time_scale_factor
-                    # Keep combined for compatibility
                     op_times["backward"] = (float(op_time_backward_d) + float(op_time_backward_w)) * time_scale_factor
                 else:
                     op_times["backward"] = float(op_time_backward) * time_scale_factor
@@ -244,14 +467,13 @@ def update_graph(n_clicks, num_devices, num_stages, num_batches, p2p_latency,
                     try:
                         overlapped_val = float(op_time_overlapped_fwd_bwd)
                         if overlapped_val > 0:
-                             # Scale overlapped time too
                              op_times["overlapped_forward_backward"] = overlapped_val * time_scale_factor
                     except (ValueError, TypeError):
                          pass
 
                 config = ScheduleConfig(
                     num_devices=int(current_num_devices),
-                    num_stages=int(current_num_stages),  # Use adjusted value
+                    num_stages=int(current_num_stages),
                     num_batches=int(num_batches),
                     p2p_latency=float(p2p_latency),
                     placement_strategy=placement_strategy,
@@ -265,73 +487,62 @@ def update_graph(n_clicks, num_devices, num_stages, num_batches, p2p_latency,
 
                 schedule = schedule_func(config)
                 schedule.execute()
-
-                # Store valid results instead of creating figure immediately
                 vis_data = convert_schedule_to_visualization_format(schedule)
                 valid_results.append((strategy, schedule, vis_data))
 
             except (AssertionError, ValueError, TypeError) as e:
-                 error_message = f"Error generating schedule for '{strategy}': {e}"
-                 import traceback
-                 traceback.print_exc()
+                 error_message = f"Error for '{strategy}': {e}"
             except Exception as e:
-                 error_message = f"An unexpected error occurred for '{strategy}': {e}"
-                 import traceback
-                 traceback.print_exc()
+                 error_message = f"Unexpected error for '{strategy}': {e}"
 
         if error_message:
              error_messages.append((strategy, error_message))
 
-    # Add alerts for any automatic parameter adjustments
+    # --- Generate Toasts --- 
+    # Add toasts for automatic adjustments
     for adjustment in automatic_adjustments:
-        output_components.append(
-            dbc.Alert(adjustment, color="info", dismissable=True)
+        toast_components.append(
+            dbc.Toast(
+                adjustment,
+                header="Parameter Adjustment",
+                icon="info",
+                duration=5000, # Slightly longer duration for info
+                is_open=True,
+                className="border-info"
+            )
         )
 
-    # If we have valid results, calculate the maximum execution time across all schedules
+    # Add toasts for errors
+    for strategy, msg in error_messages:
+        toast_components.append(
+            dbc.Toast(
+                msg,
+                header=f"Error: {strategy}",
+                icon="danger",
+                duration=8000, # Longer duration for errors
+                is_open=True,
+                className="border-danger"
+            )
+        )
+    
+    # --- Generate Graphs ---
     if valid_results:
-        # Find global maximum execution time
         max_execution_time = max(schedule.get_total_execution_time() for _, schedule, _ in valid_results)
-        
-        # Sort valid results according to the display order
-        sorted_valid_results = []
-        
-        # First add strategies in the predefined order
-        for strategy_name in strategy_display_order:
-            for result in valid_results:
-                if result[0] == strategy_name:
-                    sorted_valid_results.append(result)
-        
-        # Then add any remaining strategies that might not be in the predefined order
-        for result in valid_results:
-            if result[0] not in strategy_display_order:
-                sorted_valid_results.append(result)
-        
-        # Create figures with aligned x-axis, using the sorted results
+        sorted_valid_results = sorted(valid_results, key=lambda x: strategy_display_order.index(x[0]) if x[0] in strategy_display_order else float('inf'))
+
         for strategy, _, vis_data in sorted_valid_results:
             fig = create_pipeline_figure(vis_data, max_time=max_execution_time, show_progress=False)
-            
-            # Force the x-axis range to be the same for all figures
-            # Add a small margin (5%) for better visualization
             margin = max_execution_time * 0.05
             fig.update_layout(
-                xaxis=dict(
-                    range=[0, max_execution_time + margin]
-                )
+                xaxis=dict(range=[0, max_execution_time + margin])
             )
-            
-            output_components.append(html.Div([
+            graph_components.append(html.Div([
                 html.H4(f"Schedule: {strategy}", className="text-center mt-3 mb-2"),
                 dcc.Graph(figure=fig)
             ]))
-    
-    # Add error messages to output
-    for strategy, msg in error_messages:
-        output_components.append(
-            dbc.Alert(msg, color="danger", className="mt-3")
-        )
 
-    return output_components
+    # Return graph components and toast components
+    return graph_components, toast_components
 
 # For Hugging Face Spaces deployment
 server = app.server

assets/clientside.js

@@ -0,0 +1,62 @@
+// assets/clientside.js
+
+// Make sure the assets folder is configured correctly in Dash for this to be loaded.
+// Dash automatically serves files from a folder named 'assets' in the root directory.
+
+if (!window.dash_clientside) { window.dash_clientside = {}; }
+
+window.dash_clientside.clientside = {
+    update_strategy_selection: function(n_clicks_all, current_selection) {
+        // Determine which button triggered the callback
+        const ctx = dash_clientside.callback_context;
+        if (!ctx.triggered || ctx.triggered.length === 0) {
+            // Should not happen with prevent_initial_call=True, but handle defensively
+            return dash_clientside.no_update;
+        }
+
+        const triggered_id_str = ctx.triggered[0].prop_id.split('.')[0];
+        if (!triggered_id_str) {
+            // If we can't parse the ID, don't update
+            return dash_clientside.no_update;
+        }
+
+        // Parse the JSON ID string to get the actual index (strategy name)
+        let triggered_index;
+        try {
+            const triggered_id_obj = JSON.parse(triggered_id_str);
+            triggered_index = triggered_id_obj.index;
+        } catch (e) {
+            console.error("Error parsing callback context ID:", e);
+            return dash_clientside.no_update; // Don't update if ID parsing fails
+        }
+
+        // --- Update Selection Logic ---
+        // Initialize new_selection as a copy of the current selection
+        let new_selection = current_selection ? [...current_selection] : [];
+
+        // Toggle the selected state
+        const index_in_selection = new_selection.indexOf(triggered_index);
+        if (index_in_selection > -1) {
+            // If already selected, remove it (allow deselecting all for now)
+            new_selection.splice(index_in_selection, 1);
+        } else {
+            // If not selected, add it
+            new_selection.push(triggered_index);
+        }
+
+        // --- Prepare Outputs ---
+        const all_indices = ctx.inputs_list[0].map(input => input.id.index); // Get all strategy names from the Input IDs
+
+        // Generate active states, colors, and outlines for ALL buttons
+        const active_states = all_indices.map(index => new_selection.includes(index));
+        const colors = active_states.map(active => active ? 'primary' : 'secondary'); // 'primary' for active, 'secondary' for inactive
+        const outlines = active_states.map(active => !active); // Outline=true for inactive, false for active
+
+        // Generate validation message
+        const feedback = new_selection.length === 0 ? "Please select at least one strategy." : "";
+
+        // Return updated store data, button states, and feedback
+        return [new_selection, active_states, colors, outlines, feedback];
+    }
+    // Add other clientside functions here if needed
+}; 

assets/custom.css

@@ -0,0 +1,129 @@
+/* assets/custom.css */
+
+/* --- General & Typography (Item 7, 11) --- */
+body {
+    background-color: #F7F9FC; /* Neutral background */
+    color: #212B36; /* Dark text */
+    font-family: -apple-system, BlinkMacSystemFont, \"Segoe UI\", Roboto, \"Helvetica Neue\", Arial, sans-serif;
+    font-size: 14px;
+}
+
+/* Use H1 from dbc.Container/app.layout directly */
+.h1, h1 {
+    font-size: 24px; /* H2 equivalent in request */
+    font-weight: 600;
+}
+
+/* Card titles */
+.card-title.h5, .h5.card-title {
+    font-size: 18px; /* H3 equivalent */
+    font-weight: 600;
+    margin-bottom: 1rem; /* Add space below title */
+}
+
+/* Form labels (Item 2) */
+.form-label {
+    font-size: 14px;
+    font-weight: 500;
+    margin-bottom: 0.3rem; /* Space between label and input */
+    display: block; /* Ensure it takes full width */
+}
+
+/* Form inputs (Item 2) */
+.form-control,
+.form-select {
+    font-size: 14px;
+    /* width: 100%; Ensure inputs take full width - Bootstrap usually handles this in columns */
+    padding: 0.5rem 0.75rem;
+    border-radius: 0.375rem; /* Softer corners */
+}
+
+/* Form help text */
+.form-text {
+    font-size: 12px;
+    color: #6c757d; /* Muted color */
+}
+
+/* --- Layout & Spacing (Item 1, 7) --- */
+.container-fluid {
+    padding-top: 2rem;
+    padding-bottom: 2rem;
+}
+
+/* Spacing between form rows inside cards */
+.card-body .mb-3 {
+    margin-bottom: 1rem !important; /* Default is 1rem, ensure consistency */
+}
+
+/* Spacing between cards */
+.card {
+    margin-bottom: 24px;
+    border: 1px solid #dee2e6; /* Subtle border */
+    border-radius: 0.5rem; /* Consistent radius */
+    box-shadow: 0 2px 4px rgba(0,0,0,0.05); /* Subtle shadow */
+    /* Padding is handled by card-body */
+}
+
+/* --- Button Styling (Item 4, 5, 11) --- */
+
+/* Primary Action Button (Generate Schedule) */
+#generate-button.btn-primary {
+    background-color: #0A74DA; /* Accent color */
+    border-color: #0A74DA;
+    font-weight: 500;
+    padding: 0.6rem 1.2rem; /* Slightly larger padding */
+}
+
+#generate-button.btn-primary:hover,
+#generate-button.btn-primary:focus {
+    background-color: #085ead; /* Darker accent on hover/focus */
+    border-color: #085ead;
+}
+
+#generate-button.btn-primary:disabled {
+    background-color: #a0cff7; /* Lighter, muted accent when disabled */
+    border-color: #a0cff7;
+}
+
+/* Strategy Toggle Buttons */
+.btn-group .btn {
+    margin-right: 0.5rem; /* Space between buttons */
+    margin-bottom: 0.5rem; /* Space for wrapping */
+    border-radius: 1rem; /* Pill shape */
+    padding: 0.4rem 0.8rem;
+    font-size: 13px;
+}
+
+/* Active strategy button */
+.btn-group .btn.btn-primary:not(.disabled):not(:disabled).active,
+.btn-group .btn.btn-primary:not(.disabled):not(:disabled):active {
+    background-color: #0A74DA; /* Accent color */
+    border-color: #0A74DA;
+    color: white;
+    box-shadow: none; /* Remove default active shadow if needed */
+}
+
+/* Inactive strategy button (using outline secondary) */
+.btn-group .btn.btn-outline-secondary {
+    border-color: #ced4da;
+    color: #495057;
+}
+
+.btn-group .btn.btn-outline-secondary:hover {
+    background-color: #e9ecef;
+}
+
+/* --- Validation Feedback --- */
+.invalid-feedback {
+    font-size: 12px;
+    margin-top: 0.25rem;
+}
+
+/* --- Responsive Adjustments (Item 10) --- */
+/* Bootstrap handles column stacking. We might need more specific rules later */
+/* e.g., adjust chart container width/scrolling on smaller screens */
+
+/* Chart Container - Add basic styles, will be refined (Item 9) */
+#graph-output-container .plotly.graph-div {
+    /* Add styles for the chart itself if needed */
+}