import os import threading import gc from typing import Any, Dict import toml from textual import work from textual.app import App, ComposeResult from textual.containers import Horizontal, Vertical from textual.widgets import ( Button, Footer, Header, Input, Label, ProgressBar, RichLog, Rule, ) from tasks import PIPELINE_TASKS, SCENARIOS from util.config_manager import handle_directory_config from util.progress import ProgressTracker from util.screens import OverwriteConfirmScreen from util.ui_logger import PipelineLogger class CNNHarnessApp(App): CSS_PATH = "style.tcss" def __init__(self): super().__init__() self.stop_event = threading.Event() self.pause_event = threading.Event() self.current_config: Dict[str, Any] = {} self.ui_logger = PipelineLogger(self, widget_id="#live_log") # ========================================== # UI Layout # ========================================== def compose(self) -> ComposeResult: yield Header() with Vertical(id="main_container"): # --- TOP SECTION: Two-Column Configuration --- with Horizontal(id="config_section"): # Left Column: Path entry and buttons with Vertical(id="config_left_col"): yield Label("Working Directory:") yield Input(placeholder="./outputs/experiment_1", id="work_dir") yield Button("Load / Init Dir", variant="primary", id="load_btn") yield Button("Reload Config", variant="default", id="reload_btn") # Right Column: Read-only parameters display with Vertical(id="config_right_col"): yield Label("Current Parameters:") yield RichLog(id="config_display", highlight=True, markup=True) yield Rule() # --- BOTTOM SECTION: Controls, Progress, and Logs --- with Horizontal(id="bottom_area"): # --- Pipeline Controls --- with Vertical(id="control_panel"): yield Button("Start Pipeline", variant="success", id="btn_start") with Horizontal(id="active_controls"): yield Button("Pause", variant="warning", id="btn_pause") yield Button("Stop", variant="error", id="btn_stop") # --- Progress Trackers --- with Vertical(id="progress_area"): for i in range(5): with Horizontal( id=f"progress_row_{i}", classes="progress_row" ): yield Label( id=f"progress_title_{i}", classes="progress_title" ) yield ProgressBar( id=f"progress_bar_{i}", classes="progress_bar", show_eta=True, ) yield Label( id=f"progress_stats_{i}", classes="progress_stats" ) # --- Live Log Output --- with Vertical(id="log_panel"): with Horizontal(id="log_status_row"): yield Label("STATUS: READY", id="pipeline_status_label") with Horizontal(id="log_header_row"): yield Label("Live Logs:", classes="section_label") yield RichLog(id="live_log", highlight=True, markup=True) yield Footer() def on_mount(self) -> None: self.query_one("#active_controls").display = False self.ui_logger.info("Application initialized and ready.") # ========================================== # Helper Methods # ========================================== def _set_pipeline_status(self, text: str) -> None: try: self.query_one("#pipeline_status_label", Label).update(text) except Exception: pass def _update_config_display(self) -> None: config_log = self.query_one("#config_display", RichLog) config_log.clear() if not self.current_config: config_log.write("[italic]No configuration loaded.[/italic]") return toml_string = toml.dumps(self.current_config) config_log.write(toml_string) def _toggle_config_inputs(self, disabled: bool) -> None: self.query_one("#work_dir", Input).disabled = disabled self.query_one("#load_btn", Button).disabled = disabled self.query_one("#reload_btn", Button).disabled = disabled # ========================================== # Event Handlers # ========================================== def on_button_pressed(self, event: Button.Pressed) -> None: if event.button.id in ("load_btn", "reload_btn"): work_dir = self.query_one("#work_dir", Input).value success, message, config_data = handle_directory_config(work_dir) if success: self.current_config = config_data self.current_config["work_dir"] = work_dir self._update_config_display() self.ui_logger.info(message) else: self.ui_logger.error(message) elif event.button.id == "btn_start": if not self.current_config: self.ui_logger.error( "No configuration loaded! Please load a directory first." ) return scenario_id = self.current_config.get("scenario") if not scenario_id or scenario_id not in SCENARIOS: self.ui_logger.error( f"Invalid or missing scenario '{scenario_id}' in config.toml!" ) return work_dir = self.current_config["work_dir"] # Check if the working directory exists and is non-empty (except for the config.toml file) if os.path.exists(work_dir) and any( f for f in os.listdir(work_dir) if f != "config.toml" ): def check_overwrite_callback(proceed: bool) -> None: if proceed: self._execute_pipeline() else: self.ui_logger.error( "Pipeline start cancelled by user (folder non-empty)." ) self.app.push_screen( OverwriteConfirmScreen(work_dir), check_overwrite_callback, # pyright: ignore ) else: self._execute_pipeline() elif event.button.id == "btn_pause": if self.pause_event.is_set(): self.pause_event.clear() event.button.label = "Pause" event.button.variant = "warning" self._set_pipeline_status("[bold green]PIPELINE RUNNING[/bold green]") self.ui_logger.info("Pipeline Resumed.") else: self.pause_event.set() event.button.label = "Resume" event.button.variant = "success" self._set_pipeline_status("[bold yellow]PIPELINE PAUSED[/bold yellow]") self.ui_logger.info( "Pipeline Paused. Waiting for current operation to yield..." ) elif event.button.id == "btn_stop": self._set_pipeline_status("[bold red]STOPPING PIPELINE...[/bold red]") self.ui_logger.error("Stop requested! Terminating gracefully...") event.button.disabled = True self.stop_event.set() # ========================================== # Pipeline Execution # ========================================== def _execute_pipeline(self) -> None: self._toggle_config_inputs(disabled=True) self.query_one("#btn_start").display = False self.query_one("#active_controls").display = True self._set_pipeline_status("[bold green]PIPELINE RUNNING[/bold green]") self.stop_event.clear() self.pause_event.clear() btn_pause = self.query_one("#btn_pause", Button) btn_pause.label = "Pause" btn_pause.variant = "warning" self.run_background_pipeline(self.current_config.copy()) @work(exclusive=True, thread=True) def run_background_pipeline(self, config: Dict[str, Any]) -> None: scenario_id = config["scenario"] steps = SCENARIOS[scenario_id]["tasks"] scenario_label = SCENARIOS[scenario_id]["label"] root_tracker = ProgressTracker(self, level=0) pipeline_state: Dict[str, Any] = { "events": {"stop": self.stop_event, "pause": self.pause_event} } try: self.ui_logger.info(f"Initiating Scenario: {scenario_label}") total_tasks = len(steps) root_tracker.reset() root_tracker.set_title("Pipeline Status") root_tracker.update(total=total_tasks, advance=0) for i, task_id in enumerate(steps): if self.stop_event.is_set(): break task_name = PIPELINE_TASKS[task_id]["task_name"] task_func = PIPELINE_TASKS[task_id]["task_func"] task_logger = self.ui_logger.get_task_logger(task_name) self.ui_logger.info(f"Starting Task {i + 1}/{total_tasks}: {task_name}") task_tracker = root_tracker.get_sub_tracker(task_name) task_func(task_tracker, task_logger, config, pipeline_state) if not self.stop_event.is_set(): task_logger.info("Task completed.") root_tracker.update(advance=1) if self.stop_event.is_set(): self.call_from_thread( self._set_pipeline_status, "[bold red]PIPELINE STOPPED[/bold red]" ) self.ui_logger.error("Pipeline stopped by user.") root_tracker.reset() else: self.call_from_thread( self._set_pipeline_status, "[bold blue]PIPELINE COMPLETE[/bold blue]", ) self.ui_logger.info("All tasks finished successfully!") except InterruptedError as e: self.ui_logger.error(f"STOPPED: {str(e)}") self.call_from_thread( self._set_pipeline_status, "[bold red]PIPELINE STOPPED[/bold red]" ) root_tracker.reset() except Exception as e: self.ui_logger.file_logger.error( f"Pipeline failed: {str(e)}", exc_info=True ) self.ui_logger.error(f"ERROR: {str(e)}") self.call_from_thread( self._set_pipeline_status, "[bold red]PIPELINE FAILED[/bold red]" ) finally: pipeline_state.clear() gc.collect() self.call_from_thread(self._reset_ui) def _reset_ui(self) -> None: self.query_one("#btn_stop", Button).disabled = False self.query_one("#btn_start").display = True self.query_one("#active_controls").display = False self._toggle_config_inputs(disabled=False) if __name__ == "__main__": app = CNNHarnessApp() app.run()