import os import threading 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, Select, TabbedContent, TabPane, ) from util.progress import ProgressTracker from util.screens import OverwriteConfirmScreen from util.tasks import PIPELINE_TASKS, SCENARIOS from util.ui_logger import PipelineLogger # ========================================== # TUI Application # ========================================== class CNNHarnessApp(App): CSS_PATH = "style.tcss" def __init__(self): super().__init__() self.stop_event = threading.Event() self.pause_event = threading.Event() # Initialize our consolidated, thread-aware logger self.ui_logger = PipelineLogger(self, widget_id="#live_log") def compose(self) -> ComposeResult: yield Header() with Vertical(id="main_container"): # --- TOP SECTION: Scrollable Configuration Tabs --- with TabbedContent(initial="tab-general"): with TabPane("General Config", id="tab-general"): with Horizontal(classes="row"): with Vertical(classes="input_group"): yield Label("Load Configuration:") yield Input(placeholder="./config.toml", id="config_path") yield Button("Load Config", variant="primary", id="load_btn") with Horizontal(classes="row"): with Vertical(classes="input_group"): yield Label("Save Directory:", classes="input_label") yield Input( placeholder="./models/experiment_1", id="save_dir" ) with Vertical(classes="input_group"): yield Label("Batch Size:", classes="input_label") yield Input(placeholder="32", id="batch_size") # --- BOTTOM SECTION: Controls, Progress, and Logs --- with Horizontal(id="bottom_area"): with Vertical(id="control_panel"): with Horizontal(id="status_row"): scenario_options = [ (data["label"], sc_id) for sc_id, data in SCENARIOS.items() ] yield Select( scenario_options, prompt="Select Pipeline Scenario...", id="scenario_select", ) yield Label("STATUS: READY", id="pipeline_status_label") yield Rule() yield Button( "Start Pipeline & Save Config", 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") yield Rule() 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" ) with Vertical(id="log_panel"): yield Label("Live Logs:") 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.") def _set_pipeline_status(self, text: str) -> None: try: self.query_one("#pipeline_status_label", Label).update(text) except Exception: pass # ========================================== # Configuration Management Methods # ========================================== def _load_config_file(self, filepath: str) -> Dict[str, Any]: with open(filepath, "r") as f: return toml.load(f) def _save_config_file(self, filepath: str, config: Dict[str, Any]) -> None: os.makedirs(os.path.dirname(filepath), exist_ok=True) with open(filepath, "w") as f: toml.dump(config, f) def _populate_ui_from_config(self, config: Dict[str, Any]) -> None: self.query_one("#save_dir", Input).value = str(config.get("save_dir", "")) self.query_one("#batch_size", Input).value = str(config.get("batch_size", "")) loaded_scenario = config.get("scenario", "") if loaded_scenario in SCENARIOS: self.query_one("#scenario_select", Select).value = loaded_scenario def _gather_config_from_ui(self) -> Dict[str, Any]: scenario_select = self.query_one("#scenario_select", Select) scenario_id = ( scenario_select.value if type(scenario_select.value) is str else None ) config = { "save_dir": self.query_one("#save_dir", Input).value or "./outputs/default", "batch_size": self.query_one("#batch_size", Input).value or "32", "scenario": scenario_id, } if scenario_id and scenario_id in SCENARIOS: config["steps"] = SCENARIOS[scenario_id]["tasks"] else: config["steps"] = [] return config # ========================================== # Event Handlers # ========================================== def on_button_pressed(self, event: Button.Pressed) -> None: if event.button.id == "load_btn": config_path = self.query_one("#config_path", Input).value if not config_path: self.ui_logger.error("Please specify a config path to load.") return try: loaded_config = self._load_config_file(config_path) self._populate_ui_from_config(loaded_config) self.ui_logger.info(f"Successfully loaded config from {config_path}") except Exception as e: self.ui_logger.error(f"Failed to load config: {str(e)}") elif event.button.id == "btn_start": config_dict = self._gather_config_from_ui() if not config_dict.get("scenario"): self.ui_logger.error( "No scenario selected! Please select a scenario from the dropdown." ) return save_dir = config_dict["save_dir"] if os.path.exists(save_dir) and os.listdir(save_dir): def check_overwrite_callback(proceed: bool) -> None: if proceed: self._execute_pipeline(config_dict) else: self.ui_logger.error( "Pipeline start cancelled by user (folder non-empty)." ) self.app.push_screen( OverwriteConfirmScreen(save_dir), check_overwrite_callback, # pyright: ignore ) else: self._execute_pipeline(config_dict) 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, config_dict: Dict[str, Any]) -> None: self.query_one(TabbedContent).disabled = True self.query_one("#scenario_select").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() self.query_one("#btn_pause", Button).label = "Pause" self.query_one("#btn_pause", Button).variant = "warning" self.run_background_pipeline(config_dict) @work(exclusive=True, thread=True) def run_background_pipeline(self, config: Dict[str, Any]) -> None: save_dir = config["save_dir"] steps = config["steps"] scenario_label = SCENARIOS[config["scenario"]]["label"] root_tracker = ProgressTracker(self, level=0) try: config_out_path = os.path.join(save_dir, "config.toml") self._save_config_file(config_out_path, config) self.ui_logger.info(f"Initiating Scenario: {scenario_label}") self.ui_logger.info(f"Configuration saved to {config_out_path}") total_tasks = len(steps) pipeline_state: Dict[str, Any] = { "events": {"stop": self.stop_event, "pause": self.pause_event} } 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"] # Generate a task-specific logger from our PipelineLogger 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: 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.query_one(TabbedContent).disabled = False self.query_one("#scenario_select").disabled = False if __name__ == "__main__": app = CNNHarnessApp() app.run()