import argparse import os import threading import gc from typing import Any, Dict, Optional # Must be set before first CUDA context initialization. os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") 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, describe_protected_conflicts from util.config_manager import handle_directory_config from util.progress import MAX_LEVELS, ProgressTracker from util.screens import OverwriteConfirmScreen from util.ui_logger import PipelineLogger class CNNHarnessApp(App[None]): CSS_PATH = "style.tcss" def __init__( self, work_dir: Optional[str] = None, autorun: bool = False, autoquit: bool = False, force: bool = False, ): 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") # CLI overrides. These force the corresponding behavior ON; when not # passed, the effective value falls back to the loaded config's [run] # section (see _effective_* helpers). self._cli_work_dir = work_dir self._cli_autorun = autorun self._cli_autoquit = autoquit self._cli_force = force self._last_pipeline_ok = False # ========================================== # 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(MAX_LEVELS): 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.") # Script mode: a CLI work_dir auto-loads its config; autorun then starts # the pipeline automatically once the config is available. if self._cli_work_dir: self.query_one("#work_dir", Input).value = self._cli_work_dir if self._load_directory(self._cli_work_dir): if self._effective_autorun(): self.call_after_refresh(self._request_start, False) elif self._effective_autoquit(): # Load failed and we're running headless: exit with error code. self.call_after_refresh(self._quit_after_run, False) # ========================================== # 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 # ------------------------------------------ # Script-mode helpers (autorun / autoquit / force) # ------------------------------------------ def _run_cfg(self) -> Dict[str, Any]: run = self.current_config.get("run", {}) if self.current_config else {} return run if isinstance(run, dict) else {} def _effective_autorun(self) -> bool: return self._cli_autorun or bool(self._run_cfg().get("autorun", False)) def _effective_autoquit(self) -> bool: return self._cli_autoquit or bool(self._run_cfg().get("autoquit", False)) def _effective_force(self) -> bool: return self._cli_force or bool(self._run_cfg().get("force_overwrite", False)) def _load_directory(self, work_dir: str) -> bool: """Load-or-init a work_dir's config into current_config. Returns success.""" 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) return True self.ui_logger.error(message) return False def _quit_after_run(self, ok: bool) -> None: """Exit the app with a script-friendly return code (0 ok, 1 failure).""" self.exit(return_code=0 if ok else 1) def _request_start(self, interactive: bool) -> None: """Validate and start the pipeline. In non-interactive (autorun) mode a protected-overwrite conflict aborts unless force is set, since a modal prompt cannot be answered by a script.""" if not self.current_config: self.ui_logger.error( "No configuration loaded! Please load a directory first." ) if not interactive and self._effective_autoquit(): self._quit_after_run(False) 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!" ) if not interactive and self._effective_autoquit(): self._quit_after_run(False) return work_dir = self.current_config["work_dir"] # Only PROTECTED (hard-to-regenerate) outputs trigger a warning. An # analysis-only rerun writes nothing protected, so it proceeds silently. conflicts = describe_protected_conflicts(work_dir, scenario_id) if conflicts and not self._effective_force(): if interactive: def check_overwrite_callback(proceed: bool) -> None: if proceed: self._execute_pipeline() else: self.ui_logger.error("Pipeline start cancelled by user.") self.app.push_screen( OverwriteConfirmScreen(work_dir, conflicts), check_overwrite_callback, # pyright: ignore ) return # Non-interactive with a protected conflict and no force: refuse. self.ui_logger.error( "Autorun aborted: scenario would overwrite protected outputs " f"({'; '.join(conflicts)}). Set run.force_overwrite = true or pass " "--force to allow." ) if self._effective_autoquit(): self._quit_after_run(False) return self._execute_pipeline() # ========================================== # 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 self._load_directory(work_dir) elif event.button.id == "btn_start": self._request_start(interactive=True) 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} } self._last_pipeline_ok = False 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!") self._last_pipeline_ok = True 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) # Script mode: quit once the run finishes, with a return code that # reflects success/failure so it can be used in shell pipelines. if self._effective_autoquit(): self.call_from_thread(self._quit_after_run, self._last_pipeline_ok) 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) def _parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description=( "ALNN evaluation harness (TUI). Pass a working directory to auto-load " "it; combine with autorun/autoquit (CLI flags or the config [run] " "section) to use it as a simple script." ) ) parser.add_argument( "work_dir", nargs="?", default=None, help="Working/output directory to load on startup (its config.toml is " "loaded, or the default config is copied in if absent).", ) parser.add_argument( "--autorun", action="store_true", help="Start the pipeline automatically once the config loads " "(overrides run.autorun).", ) parser.add_argument( "--autoquit", action="store_true", help="Exit the app when the pipeline finishes, with a 0/1 return code " "(overrides run.autoquit).", ) parser.add_argument( "--force", action="store_true", help="In autorun, overwrite protected outputs (models/evaluations) " "without prompting (overrides run.force_overwrite).", ) return parser.parse_args() if __name__ == "__main__": args = _parse_args() app = CNNHarnessApp( work_dir=args.work_dir, autorun=args.autorun, autoquit=args.autoquit, force=args.force, ) app.run() raise SystemExit(app.return_code or 0)