293 lines
9.7 KiB
Python
293 lines
9.7 KiB
Python
|
"""
|
||
|
|
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
|