feat: Pre-migration checkpoint - updated docs and utilities
Updates before optimization_engine migration: - Updated migration plan to v2.1 with complete file inventory - Added OP_07 disk optimization protocol - Added SYS_16 self-aware turbo protocol - Added study archiver and cleanup utilities - Added ensemble surrogate module - Updated NX solver and session manager - Updated zernike HTML generator - Added context engineering plan - LAC session insights updates 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
411
optimization_engine/utils/study_cleanup.py
Normal file
411
optimization_engine/utils/study_cleanup.py
Normal file
@@ -0,0 +1,411 @@
|
||||
"""
|
||||
Study Cleanup Utility
|
||||
====================
|
||||
|
||||
Cleans up completed optimization studies to save disk space by removing
|
||||
large intermediate files (NX models, FEM meshes, solver results) while
|
||||
preserving essential data (parameters, extracted results, database).
|
||||
|
||||
Usage:
|
||||
python -m optimization_engine.utils.study_cleanup <study_path> [options]
|
||||
|
||||
Options:
|
||||
--dry-run Show what would be deleted without actually deleting
|
||||
--keep-best N Keep iteration folders for the top N best trials
|
||||
--keep-pareto Keep all Pareto-optimal iterations (for multi-objective)
|
||||
--aggressive Delete ALL iteration data (only keep DB and config)
|
||||
|
||||
The database (study.db) contains all optimization results and can regenerate
|
||||
any analysis. The original NX model in 1_setup is always preserved.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import shutil
|
||||
import sqlite3
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
|
||||
# Files to ALWAYS keep in iteration folders (tiny, essential)
|
||||
ESSENTIAL_FILES = {
|
||||
'params.exp', # Design parameters for this iteration
|
||||
'_temp_mass.txt', # Extracted mass
|
||||
'_temp_part_properties.json', # Part properties
|
||||
'_temp_zernike.json', # Zernike coefficients (if exists)
|
||||
'results.json', # Any extracted results
|
||||
}
|
||||
|
||||
# Extensions to DELETE (large, regenerable/already extracted)
|
||||
DELETABLE_EXTENSIONS = {
|
||||
'.op2', # Nastran binary results (~65 MB each)
|
||||
'.prt', # NX Part files (~30-35 MB each)
|
||||
'.fem', # FEM mesh files (~15 MB each)
|
||||
'.dat', # Nastran input deck (~15 MB each)
|
||||
'.sim', # Simulation file (~7 MB each)
|
||||
'.afm', # FEA auxiliary (~4 MB each)
|
||||
'.f04', # Nastran log
|
||||
'.f06', # Nastran output
|
||||
'.log', # Solver log
|
||||
'.diag', # Diagnostics
|
||||
}
|
||||
|
||||
|
||||
def get_study_info(study_path: Path) -> dict:
|
||||
"""Get study metadata from config and database."""
|
||||
config_path = study_path / 'optimization_config.json'
|
||||
# Try both possible DB locations
|
||||
db_path = study_path / '3_results' / 'study.db'
|
||||
if not db_path.exists():
|
||||
db_path = study_path / '2_results' / 'study.db'
|
||||
|
||||
info = {
|
||||
'name': study_path.name,
|
||||
'has_config': config_path.exists(),
|
||||
'has_db': db_path.exists(),
|
||||
'trial_count': 0,
|
||||
'best_trials': [],
|
||||
'pareto_trials': [],
|
||||
}
|
||||
|
||||
if config_path.exists():
|
||||
with open(config_path) as f:
|
||||
info['config'] = json.load(f)
|
||||
|
||||
if db_path.exists():
|
||||
conn = sqlite3.connect(db_path)
|
||||
cursor = conn.cursor()
|
||||
|
||||
# Get trial count
|
||||
cursor.execute("SELECT COUNT(*) FROM trials WHERE state = 'COMPLETE'")
|
||||
info['trial_count'] = cursor.fetchone()[0]
|
||||
|
||||
# Try to get best trials (for single objective)
|
||||
try:
|
||||
cursor.execute("""
|
||||
SELECT trial_id, value FROM trial_values
|
||||
WHERE objective = 0
|
||||
ORDER BY value ASC LIMIT 10
|
||||
""")
|
||||
info['best_trials'] = [row[0] for row in cursor.fetchall()]
|
||||
except Exception as e:
|
||||
pass
|
||||
|
||||
# Check for Pareto attribute
|
||||
try:
|
||||
cursor.execute("""
|
||||
SELECT DISTINCT trial_id FROM trial_system_attrs
|
||||
WHERE key = 'pareto_optimal' AND value = '1'
|
||||
""")
|
||||
info['pareto_trials'] = [row[0] for row in cursor.fetchall()]
|
||||
except:
|
||||
pass
|
||||
|
||||
conn.close()
|
||||
|
||||
return info
|
||||
|
||||
|
||||
def calculate_cleanup_savings(study_path: Path, keep_iters: set = None) -> dict:
|
||||
"""Calculate how much space would be saved by cleanup."""
|
||||
iterations_path = study_path / '2_iterations'
|
||||
if not iterations_path.exists():
|
||||
iterations_path = study_path / '1_working' # Legacy structure
|
||||
|
||||
if not iterations_path.exists():
|
||||
return {'total_size': 0, 'deletable_size': 0, 'keep_size': 0}
|
||||
|
||||
total_size = 0
|
||||
deletable_size = 0
|
||||
keep_size = 0
|
||||
keep_iters = keep_iters or set()
|
||||
|
||||
for iter_folder in iterations_path.iterdir():
|
||||
if not iter_folder.is_dir():
|
||||
continue
|
||||
|
||||
# Extract iteration number
|
||||
try:
|
||||
iter_num = int(iter_folder.name.replace('iter', ''))
|
||||
except:
|
||||
continue
|
||||
|
||||
for f in iter_folder.iterdir():
|
||||
if not f.is_file():
|
||||
continue
|
||||
size = f.stat().st_size
|
||||
total_size += size
|
||||
|
||||
# Keep entire folder if in keep_iters
|
||||
if iter_num in keep_iters:
|
||||
keep_size += size
|
||||
continue
|
||||
|
||||
# Keep essential files
|
||||
if f.name.lower() in {e.lower() for e in ESSENTIAL_FILES}:
|
||||
keep_size += size
|
||||
elif f.suffix.lower() in DELETABLE_EXTENSIONS:
|
||||
deletable_size += size
|
||||
else:
|
||||
keep_size += size # Keep unknown files by default
|
||||
|
||||
return {
|
||||
'total_size': total_size,
|
||||
'deletable_size': deletable_size,
|
||||
'keep_size': keep_size,
|
||||
}
|
||||
|
||||
|
||||
def cleanup_study(
|
||||
study_path: Path,
|
||||
dry_run: bool = True,
|
||||
keep_best: int = 0,
|
||||
keep_pareto: bool = False,
|
||||
aggressive: bool = False,
|
||||
) -> dict:
|
||||
"""
|
||||
Clean up a study to save disk space.
|
||||
|
||||
Args:
|
||||
study_path: Path to study folder
|
||||
dry_run: If True, only report what would be deleted
|
||||
keep_best: Number of best iterations to keep completely
|
||||
keep_pareto: Keep all Pareto-optimal iterations
|
||||
aggressive: Delete ALL iteration folders (only keep DB)
|
||||
|
||||
Returns:
|
||||
dict with cleanup statistics
|
||||
"""
|
||||
study_path = Path(study_path)
|
||||
if not study_path.exists():
|
||||
raise ValueError(f"Study path does not exist: {study_path}")
|
||||
|
||||
# Get study info
|
||||
info = get_study_info(study_path)
|
||||
|
||||
# Determine which iterations to keep
|
||||
keep_iters = set()
|
||||
if keep_best > 0 and info['best_trials']:
|
||||
keep_iters.update(info['best_trials'][:keep_best])
|
||||
if keep_pareto and info['pareto_trials']:
|
||||
keep_iters.update(info['pareto_trials'])
|
||||
|
||||
# Find iterations folder
|
||||
iterations_path = study_path / '2_iterations'
|
||||
if not iterations_path.exists():
|
||||
iterations_path = study_path / '1_working'
|
||||
|
||||
if not iterations_path.exists():
|
||||
return {'status': 'no_iterations', 'deleted_bytes': 0, 'deleted_files': 0}
|
||||
|
||||
# Calculate savings
|
||||
savings = calculate_cleanup_savings(study_path, keep_iters)
|
||||
|
||||
deleted_bytes = 0
|
||||
deleted_files = 0
|
||||
deleted_folders = 0
|
||||
|
||||
if aggressive:
|
||||
# Delete entire iterations folder
|
||||
if not dry_run:
|
||||
shutil.rmtree(iterations_path)
|
||||
deleted_bytes = savings['total_size']
|
||||
deleted_folders = 1
|
||||
else:
|
||||
deleted_bytes = savings['total_size']
|
||||
else:
|
||||
# Selective cleanup
|
||||
for iter_folder in iterations_path.iterdir():
|
||||
if not iter_folder.is_dir():
|
||||
continue
|
||||
|
||||
# Extract iteration number
|
||||
try:
|
||||
iter_num = int(iter_folder.name.replace('iter', ''))
|
||||
except:
|
||||
continue
|
||||
|
||||
# Skip kept iterations
|
||||
if iter_num in keep_iters:
|
||||
continue
|
||||
|
||||
for f in iter_folder.iterdir():
|
||||
if not f.is_file():
|
||||
continue
|
||||
|
||||
# Keep essential files
|
||||
if f.name.lower() in {e.lower() for e in ESSENTIAL_FILES}:
|
||||
continue
|
||||
|
||||
# Delete deletable extensions
|
||||
if f.suffix.lower() in DELETABLE_EXTENSIONS:
|
||||
size = f.stat().st_size
|
||||
if not dry_run:
|
||||
f.unlink()
|
||||
deleted_bytes += size
|
||||
deleted_files += 1
|
||||
|
||||
return {
|
||||
'status': 'dry_run' if dry_run else 'completed',
|
||||
'study_name': info['name'],
|
||||
'trial_count': info['trial_count'],
|
||||
'kept_iterations': list(keep_iters),
|
||||
'total_size_before': savings['total_size'],
|
||||
'deleted_bytes': deleted_bytes,
|
||||
'deleted_files': deleted_files,
|
||||
'deleted_folders': deleted_folders,
|
||||
'space_saved_gb': deleted_bytes / (1024**3),
|
||||
}
|
||||
|
||||
|
||||
def cleanup_batch(
|
||||
parent_path: Path,
|
||||
pattern: str = "*",
|
||||
dry_run: bool = True,
|
||||
keep_best: int = 3,
|
||||
keep_pareto: bool = False,
|
||||
aggressive: bool = False,
|
||||
) -> list:
|
||||
"""
|
||||
Clean up multiple studies matching a pattern.
|
||||
|
||||
Args:
|
||||
parent_path: Parent directory containing studies
|
||||
pattern: Glob pattern to match study folders (e.g., "m1_mirror_*")
|
||||
dry_run: If True, only report
|
||||
keep_best: Keep N best iterations per study
|
||||
keep_pareto: Keep Pareto-optimal iterations
|
||||
aggressive: Delete all iteration folders
|
||||
|
||||
Returns:
|
||||
List of cleanup results
|
||||
"""
|
||||
parent_path = Path(parent_path)
|
||||
results = []
|
||||
|
||||
for study_path in sorted(parent_path.glob(pattern)):
|
||||
if not study_path.is_dir():
|
||||
continue
|
||||
# Check if it looks like a study (has iterations folder)
|
||||
if not (study_path / '2_iterations').exists() and not (study_path / '1_working').exists():
|
||||
continue
|
||||
|
||||
try:
|
||||
result = cleanup_study(
|
||||
study_path,
|
||||
dry_run=dry_run,
|
||||
keep_best=keep_best,
|
||||
keep_pareto=keep_pareto,
|
||||
aggressive=aggressive,
|
||||
)
|
||||
results.append(result)
|
||||
except Exception as e:
|
||||
results.append({
|
||||
'study_name': study_path.name,
|
||||
'status': 'error',
|
||||
'error': str(e),
|
||||
})
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description='Clean up completed optimization studies to save disk space.',
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog=__doc__
|
||||
)
|
||||
parser.add_argument('study_path', type=Path, help='Path to study folder or parent directory')
|
||||
parser.add_argument('--dry-run', action='store_true', default=True,
|
||||
help='Show what would be deleted without deleting (default)')
|
||||
parser.add_argument('--execute', action='store_true',
|
||||
help='Actually delete files (opposite of --dry-run)')
|
||||
parser.add_argument('--keep-best', type=int, default=3,
|
||||
help='Keep N best iterations completely (default: 3)')
|
||||
parser.add_argument('--keep-pareto', action='store_true',
|
||||
help='Keep all Pareto-optimal iterations')
|
||||
parser.add_argument('--aggressive', action='store_true',
|
||||
help='Delete ALL iteration data (only keep DB)')
|
||||
parser.add_argument('--batch', type=str, metavar='PATTERN',
|
||||
help='Clean multiple studies matching pattern (e.g., "m1_mirror_*")')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
dry_run = not args.execute
|
||||
|
||||
if args.batch:
|
||||
# Batch cleanup mode
|
||||
print(f"\n{'='*60}")
|
||||
print(f"BATCH CLEANUP: {args.study_path}")
|
||||
print(f"Pattern: {args.batch}")
|
||||
print(f"{'='*60}")
|
||||
print(f"Mode: {'DRY RUN' if dry_run else 'EXECUTE'}")
|
||||
|
||||
results = cleanup_batch(
|
||||
args.study_path,
|
||||
pattern=args.batch,
|
||||
dry_run=dry_run,
|
||||
keep_best=args.keep_best,
|
||||
keep_pareto=args.keep_pareto,
|
||||
aggressive=args.aggressive,
|
||||
)
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print("BATCH RESULTS")
|
||||
print(f"{'='*60}")
|
||||
print(f"{'Study':<45} {'Trials':>7} {'Size':>8} {'Savings':>8}")
|
||||
print("-" * 75)
|
||||
|
||||
total_saved = 0
|
||||
for r in results:
|
||||
if r.get('status') == 'error':
|
||||
print(f"{r['study_name']:<45} ERROR: {r.get('error', 'Unknown')}")
|
||||
else:
|
||||
saved = r.get('space_saved_gb', 0)
|
||||
total_saved += saved
|
||||
print(f"{r['study_name']:<45} {r.get('trial_count', 0):>7} "
|
||||
f"{r.get('total_size_before', 0)/(1024**3):>7.1f}G {saved:>7.1f}G")
|
||||
|
||||
print("-" * 75)
|
||||
print(f"{'TOTAL SAVINGS:':<45} {' '*15} {total_saved:>7.1f}G")
|
||||
|
||||
if dry_run:
|
||||
print(f"\n[!] This was a dry run. Run with --execute to actually delete files.")
|
||||
|
||||
return results
|
||||
|
||||
else:
|
||||
# Single study cleanup
|
||||
print(f"\n{'='*60}")
|
||||
print(f"STUDY CLEANUP: {args.study_path.name}")
|
||||
print(f"{'='*60}")
|
||||
print(f"Mode: {'DRY RUN (no files deleted)' if dry_run else 'EXECUTE (files WILL be deleted)'}")
|
||||
print(f"Keep best: {args.keep_best} iterations")
|
||||
print(f"Keep Pareto: {args.keep_pareto}")
|
||||
print(f"Aggressive: {args.aggressive}")
|
||||
|
||||
result = cleanup_study(
|
||||
args.study_path,
|
||||
dry_run=dry_run,
|
||||
keep_best=args.keep_best,
|
||||
keep_pareto=args.keep_pareto,
|
||||
aggressive=args.aggressive,
|
||||
)
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print("RESULTS")
|
||||
print(f"{'='*60}")
|
||||
print(f"Trials in study: {result['trial_count']}")
|
||||
print(f"Iterations kept: {len(result['kept_iterations'])} {result['kept_iterations'][:5]}{'...' if len(result['kept_iterations']) > 5 else ''}")
|
||||
print(f"Total size before: {result['total_size_before'] / (1024**3):.2f} GB")
|
||||
print(f"{'Would delete' if dry_run else 'Deleted'}: {result['deleted_files']} files")
|
||||
print(f"Space {'to save' if dry_run else 'saved'}: {result['space_saved_gb']:.2f} GB")
|
||||
|
||||
if dry_run:
|
||||
print(f"\n[!] This was a dry run. Run with --execute to actually delete files.")
|
||||
|
||||
return result
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in New Issue
Block a user