test: add timeout handling

app.py

@@ -609,191 +609,148 @@ def load_all_data(image_root, pkl_root):
 default_image_name = "christmas-imagenet"
 
 
-with gr.Blocks(
-    theme=gr.themes.Citrus(),
-    css="""
-    .image-row .gr-image { margin: 0 !important; padding: 0 !important; }
-    .image-row img { width: auto; height: 50px; } /* Set a uniform height for all images */
-""",
-) as demo:
-    with gr.Row():
-        with gr.Column():
-            # Left View: Image selection and click handling
-            gr.Markdown("## Select input image and patch on the image")
-            image_selector = gr.Dropdown(
-                choices=list(data_dict.keys()),
-                value=default_image_name,
-                label="Select Image",
-            )
-            image_display = gr.Image(
-                value=data_dict[default_image_name]["image"],
-                type="pil",
-                interactive=True,
-            )
-
-            # Update image display when a new image is selected
-            image_selector.change(
-                fn=lambda img_name: data_dict[img_name]["image"],
-                inputs=image_selector,
-                outputs=image_display,
-            )
-            image_display.select(
-                fn=highlight_grid, inputs=[image_selector], outputs=[image_display]
-            )
-
-        with gr.Column():
-            gr.Markdown("## SAE latent activations of CLIP and MaPLE")
-            model_options = [f"MaPLE-{dataset_name}" for dataset_name in DATASET_LIST]
-            model_selector = gr.Dropdown(
-                choices=model_options,
-                value=model_options[0],
-                label="Select adapted model (MaPLe)",
-            )
-            init_plot = plot_activation_distribution(
-                None, default_image_name, model_options[0]
-            )
-            neuron_plot = gr.Plot(
-                label="Neuron Activation", value=init_plot, show_label=False
-            )
-
-            image_selector.change(
-                fn=plot_activation_distribution,
-                inputs=[image_selector, model_selector],
-                outputs=neuron_plot,
-            )
-            image_display.select(
-                fn=plot_activation_distribution,
-                inputs=[image_selector, model_selector],
-                outputs=neuron_plot,
-            )
-            model_selector.change(
-                fn=load_image, inputs=[image_selector], outputs=image_display
-            )
-            model_selector.change(
-                fn=plot_activation_distribution,
-                inputs=[image_selector, model_selector],
-                outputs=neuron_plot,
-            )
-
-    with gr.Row():
-        with gr.Column():
-            radio_names = get_init_radio_options(default_image_name, model_options[0])
-
-            feautre_idx = radio_names[0].split("-")[-1]
-            markdown_display = gr.Markdown(
-                f"## Segmentation mask for the selected SAE latent - {feautre_idx}"
-            )
-            init_seg, init_tops, init_values = show_activation_heatmap(
-                default_image_name, radio_names[0], "CLIP"
-            )
-
-            gr.Markdown("### Localize SAE latent activation using CLIP")
-            seg_mask_display = gr.Image(value=init_seg, type="pil", show_label=False)
-            init_seg_maple, _, _ = show_activation_heatmap(
-                default_image_name, radio_names[0], model_options[0]
-            )
-            gr.Markdown("### Localize SAE latent activation using MaPLE")
-            seg_mask_display_maple = gr.Image(
-                value=init_seg_maple, type="pil", show_label=False
-            )
-
-        with gr.Column():
-            gr.Markdown("## Top activating SAE latent index")
-
-            radio_choices = gr.Radio(
-                choices=radio_names,
-                label="Top activating SAE latent",
-                interactive=True,
-                value=radio_names[0],
-            )
-            toggle_btn = gr.Checkbox(label="Show segmentation mask", value=False)
-
-            markdown_display_2 = gr.Markdown(
-                f"## Top reference images for the selected SAE latent - {feautre_idx}"
-            )
-
-            gr.Markdown("### ImageNet")
-            top_image_1 = gr.Image(
-                value=init_tops[0], type="pil", label="ImageNet", show_label=False
-            )
-            act_value_1 = gr.Markdown(init_values[0])
-
-            gr.Markdown("### ImageNet-Sketch")
-            top_image_2 = gr.Image(
-                value=init_tops[1],
-                type="pil",
-                label="ImageNet-Sketch",
-                show_label=False,
-            )
-            act_value_2 = gr.Markdown(init_values[1])
-
-            gr.Markdown("### Caltech101")
-            top_image_3 = gr.Image(
-                value=init_tops[2], type="pil", label="Caltech101", show_label=False
-            )
-            act_value_3 = gr.Markdown(init_values[2])
-
-            image_display.select(
-                fn=update_radio_options,
-                inputs=[image_selector, model_selector],
-                outputs=[radio_choices],
-            )
-
-            model_selector.change(
-                fn=update_radio_options,
-                inputs=[image_selector, model_selector],
-                outputs=[radio_choices],
-            )
-
-            image_selector.select(
-                fn=update_radio_options,
-                inputs=[image_selector, model_selector],
-                outputs=[radio_choices],
-            )
-
-        radio_choices.change(
-            fn=update_all,
-            inputs=[image_selector, radio_choices, toggle_btn, model_selector],
-            outputs=[
-                seg_mask_display,
-                seg_mask_display_maple,
-                top_image_1,
-                top_image_2,
-                top_image_3,
-                act_value_1,
-                act_value_2,
-                act_value_3,
-                markdown_display,
-                markdown_display_2,
-            ],
-        )
-
-        toggle_btn.change(
-            fn=show_activation_heatmap_clip,
-            inputs=[image_selector, radio_choices, toggle_btn],
-            outputs=[
-                seg_mask_display,
-                top_image_1,
-                top_image_2,
-                top_image_3,
-                act_value_1,
-                act_value_2,
-                act_value_3,
-            ],
-        )
+def safe_operation(func, *args, timeout=10, default=None):
+    """Execute function with timeout and return default value on failure."""
+    try:
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            future = executor.submit(func, *args)
+            return future.result(timeout=timeout)
+    except concurrent.futures.TimeoutError:
+        print(f"Operation timed out: {func.__name__}")
+        return default
+    except Exception as e:
+        print(f"Operation failed: {func.__name__}, Error: {e}")
+        return default
+
+def create_interface():
+    with gr.Blocks(
+        theme=gr.themes.Citrus(),
+        css="""
+        .image-row .gr-image { margin: 0 !important; padding: 0 !important; }
+        .image-row img { width: auto; height: 50px; }
+        """,
+    ) as demo:
+        # State management for preventing duplicate operations
+        state = gr.State({
+            'last_update': 0,
+            'processing': False
+        })
+
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("## Select input image and patch on the image")
+                image_selector = gr.Dropdown(
+                    choices=list(_CACHE['data_dict'].keys()),
+                    value=default_image_name,
+                    label="Select Image",
+                )
+                image_display = gr.Image(
+                    value=_CACHE['data_dict'][default_image_name]["image"],
+                    type="pil",
+                    interactive=True,
+                )
+
+                def safe_update_image(img_name):
+                    return safe_operation(
+                        lambda: _CACHE['data_dict'][img_name]["image"],
+                        timeout=5,
+                        default=Image.new('RGB', (IMAGE_SIZE, IMAGE_SIZE), 'gray')
+                    )
+
+                image_selector.change(
+                    fn=safe_update_image,
+                    inputs=image_selector,
+                    outputs=image_display,
+                )
+
+            with gr.Column():
+                gr.Markdown("## SAE latent activations of CLIP and MaPLE")
+                model_options = [f"MaPLE-{dataset_name}" for dataset_name in DATASET_LIST]
+                model_selector = gr.Dropdown(
+                    choices=model_options,
+                    value=model_options[0],
+                    label="Select adapted model (MaPLe)",
+                )
+
+                def safe_plot_activation(evt, selected_image, model_name):
+                    try:
+                        return safe_operation(
+                            plot_activation_distribution,
+                            evt, selected_image, model_name,
+                            timeout=10,
+                            default=go.Figure()  # Return empty figure on timeout
+                        )
+                    except Exception as e:
+                        print(f"Error in plot_activation: {e}")
+                        return go.Figure()
+
+                neuron_plot = gr.Plot(
+                    label="Neuron Activation",
+                    value=safe_plot_activation(None, default_image_name, model_options[0]),
+                    show_label=False,
+                )
+
+                def debounced_update(*args, minimum_interval=1.0):
+                    """Prevent updates that are too close together"""
+                    current_time = time.time()
+                    if current_time - state.value['last_update'] < minimum_interval:
+                        raise gr.Error("Please wait a moment before updating again")
+                    state.value['last_update'] = current_time
+                    return safe_plot_activation(*args)
+
+                # Add error handling and debouncing to all event handlers
+                image_selector.change(
+                    fn=debounced_update,
+                    inputs=[image_selector, model_selector],
+                    outputs=neuron_plot,
+                )
+                image_display.select(
+                    fn=debounced_update,
+                    inputs=[image_selector, model_selector],
+                    outputs=neuron_plot,
+                )
+                model_selector.change(
+                    fn=debounced_update,
+                    inputs=[image_selector, model_selector],
+                    outputs=neuron_plot,
+                )
+
+        # Rest of your interface code with similar error handling...
+        
+        return demo
 
-    # Launch the app
-    # demo.queue()
-    # demo.launch()
+demo = create_interface()
 
-    
 if __name__ == "__main__":
-    demo.queue()  # Enable queuing for better handling of concurrent users
-    demo.launch(
-        server_name="0.0.0.0",  # Allow external access
-        server_port=7860,
-        share=False,  # Set to True if you want to create a public URL
-        show_error=True,
-        # Optimize concurrency
-        max_threads=8,  # Adjust based on your CPU cores
+    # Configure logging
+    import logging
+    logging.basicConfig(
+        level=logging.INFO,
+        format='%(asctime)s - %(levelname)s - %(message)s'
     )
+
+    # Initialize the interface with error handling
+    try:
+        demo = create_interface()
+        demo.queue(concurrency_count=2)  # Limit concurrent operations
+        demo.launch(
+            server_name="0.0.0.0",
+            server_port=7860,
+            share=False,
+            show_error=True,
+            max_threads=4,  # Reduced for better stability
+            prevent_thread_lock=True,
+            # Add health check endpoint
+            root_path="/healthz",
+            # Add automatic cleanup
+            _teardown=lambda: (
+                clear_caches(),
+                logging.info("Caches cleared during teardown")
+            ),
+            # Configure timeouts
+            api_open_timeout=60,
+            api_call_timeout=300,
+        )
+    except Exception as e:
+        logging.error(f"Failed to launch interface: {e}")
+        raise