""" Trial Manager - Unified trial numbering and folder management ============================================================== Provides consistent trial_NNNN naming across all optimization methods (Optuna, Turbo, GNN, manual) with proper database integration. Usage: from optimization_engine.utils.trial_manager import TrialManager tm = TrialManager(study_dir) # Get next trial (creates folder, reserves DB row) trial = tm.new_trial(params={'rib_thickness': 10.5, ...}) # After FEA completes tm.complete_trial( trial_id=trial['trial_id'], objectives={'wfe_40_20': 5.63, 'mass_kg': 118.67}, metadata={'solve_time': 211.7} ) Key principles: - Trial numbers NEVER reset (monotonically increasing) - Folders NEVER get overwritten - Database is always in sync with filesystem - Surrogate predictions are NOT trials (only FEA results) """ import json import sqlite3 import shutil from pathlib import Path from datetime import datetime from typing import Dict, Any, Optional, List, Union from filelock import FileLock from .dashboard_db import DashboardDB class TrialManager: """Manages trial numbering, folders, and database for optimization studies.""" def __init__(self, study_dir: Union[str, Path], study_name: Optional[str] = None): """ Initialize trial manager for a study. Args: study_dir: Path to study directory (contains 1_setup/, 2_iterations/, 3_results/) study_name: Name of study (defaults to directory name) """ self.study_dir = Path(study_dir) self.study_name = study_name or self.study_dir.name self.iterations_dir = self.study_dir / "2_iterations" self.results_dir = self.study_dir / "3_results" self.db_path = self.results_dir / "study.db" self.lock_path = self.results_dir / ".trial_lock" # Ensure directories exist self.iterations_dir.mkdir(parents=True, exist_ok=True) self.results_dir.mkdir(parents=True, exist_ok=True) # Initialize database self.db = DashboardDB(self.db_path, self.study_name) def _get_next_trial_number(self) -> int: """Get next available trial number (never resets).""" # Check filesystem existing_folders = list(self.iterations_dir.glob("trial_*")) max_folder = 0 for folder in existing_folders: try: num = int(folder.name.split('_')[1]) max_folder = max(max_folder, num) except (IndexError, ValueError): continue # Check database conn = sqlite3.connect(self.db_path) cursor = conn.cursor() cursor.execute("SELECT COALESCE(MAX(number), -1) + 1 FROM trials") max_db = cursor.fetchone()[0] conn.close() # Return max of both + 1 (use 1-based for folders, 0-based for DB) return max(max_folder, max_db) + 1 def new_trial( self, params: Dict[str, float], source: str = "turbo", metadata: Optional[Dict[str, Any]] = None ) -> Dict[str, Any]: """ Start a new trial - creates folder and reserves DB row. Args: params: Design parameters for this trial source: How this trial was generated ("turbo", "optuna", "manual") metadata: Additional info (turbo_batch, predicted_ws, etc.) Returns: Dict with trial_id, trial_number, folder_path """ # Use file lock to prevent race conditions with FileLock(self.lock_path): trial_number = self._get_next_trial_number() # Create folder with zero-padded name folder_name = f"trial_{trial_number:04d}" folder_path = self.iterations_dir / folder_name folder_path.mkdir(exist_ok=True) # Save params to folder params_file = folder_path / "params.json" with open(params_file, 'w') as f: json.dump(params, f, indent=2) # Also save as .exp format for NX compatibility exp_file = folder_path / "params.exp" with open(exp_file, 'w') as f: for name, value in params.items(): f.write(f"[mm]{name}={value}\n") # Save metadata meta = { "trial_number": trial_number, "source": source, "status": "RUNNING", "datetime_start": datetime.now().isoformat(), "params": params, } if metadata: meta.update(metadata) meta_file = folder_path / "_meta.json" with open(meta_file, 'w') as f: json.dump(meta, f, indent=2) return { "trial_id": trial_number, # Will be updated after DB insert "trial_number": trial_number, "folder_path": folder_path, "folder_name": folder_name, } def complete_trial( self, trial_number: int, objectives: Dict[str, float], weighted_sum: Optional[float] = None, is_feasible: bool = True, metadata: Optional[Dict[str, Any]] = None ) -> int: """ Complete a trial - logs to database and updates folder metadata. Args: trial_number: Trial number from new_trial() objectives: Objective values from FEA weighted_sum: Combined objective for ranking is_feasible: Whether constraints are satisfied metadata: Additional info (solve_time, prediction_error, etc.) Returns: Database trial_id """ folder_path = self.iterations_dir / f"trial_{trial_number:04d}" # Load existing metadata meta_file = folder_path / "_meta.json" with open(meta_file, 'r') as f: meta = json.load(f) params = meta.get("params", {}) # Update metadata meta["status"] = "COMPLETE" meta["datetime_complete"] = datetime.now().isoformat() meta["objectives"] = objectives meta["weighted_sum"] = weighted_sum meta["is_feasible"] = is_feasible if metadata: meta.update(metadata) # Save results.json results_file = folder_path / "results.json" with open(results_file, 'w') as f: json.dump({ "objectives": objectives, "weighted_sum": weighted_sum, "is_feasible": is_feasible, "metadata": metadata or {} }, f, indent=2) # Update _meta.json with open(meta_file, 'w') as f: json.dump(meta, f, indent=2) # Log to database db_metadata = metadata or {} db_metadata["source"] = meta.get("source", "unknown") if "turbo_batch" in meta: db_metadata["turbo_batch"] = meta["turbo_batch"] if "predicted_ws" in meta: db_metadata["predicted_ws"] = meta["predicted_ws"] trial_id = self.db.log_trial( params=params, objectives=objectives, weighted_sum=weighted_sum, is_feasible=is_feasible, state="COMPLETE", datetime_start=meta.get("datetime_start"), datetime_complete=meta.get("datetime_complete"), metadata=db_metadata, ) # Check if this is the new best best = self.db.get_best_trial() if best and best['trial_id'] == trial_id: self.db.mark_best(trial_id) meta["is_best"] = True with open(meta_file, 'w') as f: json.dump(meta, f, indent=2) return trial_id def fail_trial(self, trial_number: int, error: str): """Mark a trial as failed.""" folder_path = self.iterations_dir / f"trial_{trial_number:04d}" meta_file = folder_path / "_meta.json" if meta_file.exists(): with open(meta_file, 'r') as f: meta = json.load(f) meta["status"] = "FAIL" meta["error"] = error meta["datetime_complete"] = datetime.now().isoformat() with open(meta_file, 'w') as f: json.dump(meta, f, indent=2) def get_trial_folder(self, trial_number: int) -> Path: """Get folder path for a trial number.""" return self.iterations_dir / f"trial_{trial_number:04d}" def get_all_trials(self) -> List[Dict[str, Any]]: """Get all completed trials from database.""" conn = sqlite3.connect(self.db_path) cursor = conn.cursor() cursor.execute(""" SELECT t.trial_id, t.number, tv.value FROM trials t JOIN trial_values tv ON t.trial_id = tv.trial_id WHERE t.state = 'COMPLETE' ORDER BY t.number """) trials = [] for row in cursor.fetchall(): trials.append({ "trial_id": row[0], "number": row[1], "value": row[2] }) conn.close() return trials def get_summary(self) -> Dict[str, Any]: """Get trial manager summary.""" summary = self.db.get_summary() # Add folder count folders = list(self.iterations_dir.glob("trial_*")) summary["folder_count"] = len(folders) return summary def copy_model_files(self, source_dir: Path, trial_number: int) -> Path: """Copy NX model files to trial folder.""" dest = self.get_trial_folder(trial_number) # Copy relevant files extensions = ['.prt', '.fem', '.sim', '.afm', '.op2', '.f06', '.dat'] for ext in extensions: for src_file in source_dir.glob(f"*{ext}"): shutil.copy2(src_file, dest / src_file.name) return dest