add oom predictor

app.py

@@ -71,9 +71,10 @@ def format_config_display(config):
 
 def process_yaml_and_plot(config):
     if not config:
-        return None, None, "No configuration loaded"
-    fig1, fig2 = plot_memory_breakdown(**config)
-    return fig1, fig2, format_config_display(config)
+        return None, None, "No configuration loaded", None
+    fig1, fig2, memory_usage_peak_tbi = plot_memory_breakdown(**config)
+    oom_prediction = "OOM" if memory_usage_peak_tbi > 75000 else "No OOM"
+    return fig1, fig2, format_config_display(config), oom_prediction
 
 with gr.Blocks() as demo:
     with gr.Row():
@@ -109,9 +110,9 @@ with gr.Blocks() as demo:
                     zero_stage = gr.Radio([0, 1], value=0, label="ZeRO Stage")
 
                 manual_submit = gr.Button("Calculate Memory (Manual Input)")
-
         with gr.Column(scale=2):
             config_display = gr.Markdown(label="Configuration Values")
+            oom_display = gr.Text(label="OOM Prediction")
             plot1 = gr.Plot(label="Memory Component Breakdown")
             plot2 = gr.Plot(label="Aggregate Memory Metrics")
 
@@ -119,14 +120,14 @@ with gr.Blocks() as demo:
     yaml_file.change(
         lambda x: process_yaml_and_plot(load_config_from_yaml_file(x) if x else None),
         inputs=[yaml_file],
-        outputs=[plot1, plot2, config_display]
+        outputs=[plot1, plot2, config_display, oom_display]
     )
 
     # Handle YAML text input
     yaml_submit.click(
         lambda x: process_yaml_and_plot(load_config_from_yaml_content(x) if x else None),
         inputs=[yaml_text],
-        outputs=[plot1, plot2, config_display]
+        outputs=[plot1, plot2, config_display, oom_display]
     )
 
     # Handle manual input
@@ -156,7 +157,7 @@ with gr.Blocks() as demo:
             seq_len, mbs, batch_accum, tp, pp, dp, zero_stage,
             tie_word_embeddings
         ],
-        outputs=[plot1, plot2, config_display]
+        outputs=[plot1, plot2, config_display, oom_display]
     )
 
 if __name__ == "__main__":