diff --git a/nexus/examples/quantum_espresso/03_machine/run.py b/nexus/examples/quantum_espresso/03_machine/run.py new file mode 100755 index 0000000000..361fe6e286 --- /dev/null +++ b/nexus/examples/quantum_espresso/03_machine/run.py @@ -0,0 +1,220 @@ +#! /usr/bin/env python3 +""" +Nexus Enhanced Workflow Example: Diamond SCF with Conditional Execution + +This example demonstrates the new enhanced workflow features in Nexus: +1. Simulation dependencies (scf2 depends on scf1) +2. Error handlers with automatic resubmission +3. Conditional execution based on simulation results (energy threshold) +4. Branching workflow: scf3 or scf4 runs based on scf2's total energy + +The workflow: +- scf1: Initial SCF calculation (no dependencies) +- scf2: SCF calculation that depends on scf1's output +- scf3: Runs if scf2's energy is below threshold (depends on scf2) +- scf4: Runs if scf2's energy is above threshold (depends on scf2) +- Only one of scf3 or scf4 will execute based on the energy condition +- All simulations have error handlers that retry on failure +""" + +from nexus import settings, Job, run_project +from nexus import generate_physical_system +from nexus import generate_pwscf + +# Import enhanced workflow features +from enhanced_simulation import make_enhanced +from error_handler import RetryErrorHandler + +# Computer configuration +computer = 'ws16' + +if computer == 'baseline': + qe_modules = 'module purge; module load Core/25.05 gcc/12.4.0 openmpi/5.0.5 DefApps hdf5' + qe_bin = '/ccsopen/home/ksu/SOFTWARE/qe/q-e-qe-7.4.1/build/bin' + account = 'phy191' +elif computer == 'ws16': + qe_modules = '' + qe_bin = '/Users/ksu/Software/qe/q-e-qe-7.5/build/bin' + account = 'ks' +else: + print('Undefined computer') + exit() + +settings( + pseudo_dir = '../pseudopotentials', + results = '', + sleep = 3, + machine = computer, + account = account, + ) + +# Define physical system (diamond) +system = generate_physical_system( + units = 'A', + axes = '''1.785 1.785 0.000 + 0.000 1.785 1.785 + 1.785 0.000 1.785''', + elem_pos = ''' + C 0.0000 0.0000 0.0000 + C 0.8925 0.8925 0.8925 + ''', + C = 4, + ) + +# Simple example with two simulations and a dependency +# scf1 runs first, scf2 depends on scf1's output + +# Simulation 1: Initial SCF calculation (no dependencies) +# Testing without make_enhanced wrapper - this will work but without error handlers +qe_job = Job(nodes=1, + queue='batch', + hours=1, + presub=qe_modules, + app=qe_bin+'/pw.x') +scf1 = generate_pwscf( + identifier = 'scf', + path = 'diamond/scf', + job = qe_job, + input_type = 'generic', + calculation = 'scf', + input_dft = 'lda', + ecutwfc = 200, + conv_thr = 1e-8, + system = system, + pseudos = ['C.BFD.upf'], + kgrid = (4,4,4), + kshift = (0,0,0), + ) + +# Uncomment to add error handlers: +# scf1 = make_enhanced( +# scf1, +# error_handlers=[RetryErrorHandler(max_retries=3, delay=1.0)], +# ) + +# Simulation 2: SCF calculation that depends on scf1 +# This depends on scf1's charge_density output +scf2_base = generate_pwscf( + identifier = 'scf_dependent', + path = 'diamond/scf_dependent', + job = qe_job, + input_type = 'generic', + calculation = 'scf', + input_dft = 'lda', + ecutwfc = 200, + conv_thr = 1e-8, + system = system, + pseudos = ['C.BFD.upf'], + kgrid = (4,4,4), + kshift = (0,0,0), + dependencies = (scf1, 'charge_density'), + ) + +# Convert to enhanced simulation with error handler +scf2 = make_enhanced( + scf2_base, + error_handlers=[RetryErrorHandler(max_retries=2)], + required_machine='andes', + ) + +# Define energy threshold (in Ry, typical for Quantum ESPRESSO) +# Adjust this value based on your system's expected energy range +energy_threshold = -30.0 + +# Conditional function: scf3 runs if scf2 energy is below threshold +def energy_below_threshold(sim): + """Check if scf2's energy is below threshold.""" + try: + # Get scf2 from dependencies (scf3 depends on scf2) + scf2_sim = None + for dep in sim.dependencies.values(): + if dep.sim.identifier == 'scf_dependent': + scf2_sim = dep.sim + break + + if scf2_sim is None or not scf2_sim.finished: + return False # scf2 not finished yet + + # Load analyzer and get energy + analyzer = scf2_sim.load_analyzer_image() + if hasattr(analyzer, 'E') and analyzer.E != 0.0: + energy = analyzer.E + return energy < energy_threshold + else: + return False # Energy not available + except Exception as e: + # If anything fails, don't run + return False + +# Conditional function: scf4 runs if scf2 energy is above threshold +def energy_above_threshold(sim): + """Check if scf2's energy is above threshold.""" + try: + # Get scf2 from dependencies (scf4 depends on scf2) + scf2_sim = None + for dep in sim.dependencies.values(): + if dep.sim.identifier == 'scf_dependent': + scf2_sim = dep.sim + break + + if scf2_sim is None or not scf2_sim.finished: + return False # scf2 not finished yet + + # Load analyzer and get energy + analyzer = scf2_sim.load_analyzer_image() + if hasattr(analyzer, 'E') and analyzer.E != 0.0: + energy = analyzer.E + return energy >= energy_threshold + else: + return False # Energy not available + except Exception as e: + # If anything fails, don't run + return False + +# Simulation 3: Runs if scf2 energy is below threshold +scf3_base = generate_pwscf( + identifier = 'scf_low_energy', + path = 'diamond/scf_low_energy', + job = qe_job, + input_type = 'generic', + calculation = 'scf', + input_dft = 'lda', + ecutwfc = 200, + conv_thr = 1e-8, + system = system, + pseudos = ['C.BFD.upf'], + kgrid = (4,4,4), + kshift = (0,0,0), + dependencies = (scf2, 'charge_density'), + ) + +scf3 = make_enhanced( + scf3_base, + condition=energy_below_threshold, # Execution-time conditional + error_handlers=[RetryErrorHandler(max_retries=2)], + ) + +# Simulation 4: Runs if scf2 energy is above threshold +scf4_base = generate_pwscf( + identifier = 'scf_high_energy', + path = 'diamond/scf_high_energy', + job = qe_job, + input_type = 'generic', + calculation = 'scf', + input_dft = 'lda', + ecutwfc = 200, + conv_thr = 1e-8, + system = system, + pseudos = ['C.BFD.upf'], + kgrid = (4,4,4), + kshift = (0,0,0), + dependencies = (scf2, 'charge_density'), + ) + +scf4 = make_enhanced( + scf4_base, + condition=energy_above_threshold, # Execution-time conditional + error_handlers=[RetryErrorHandler(max_retries=2)], + ) + +run_project() diff --git a/nexus/examples/quantum_espresso/04_loop/run.py b/nexus/examples/quantum_espresso/04_loop/run.py new file mode 100755 index 0000000000..80b75faf20 --- /dev/null +++ b/nexus/examples/quantum_espresso/04_loop/run.py @@ -0,0 +1,151 @@ +#! /usr/bin/env python3 +""" +Enhanced loop workflow for Quantum ESPRESSO. + +This script demonstrates the loop support added to enhanced simulations. A +loop-enabled SCF step repeatedly resubmits itself until either a convergence +condition is met or the maximum number of iterations is reached. + +""" + +from nexus import settings, Job, run_project +from nexus import generate_physical_system +from nexus import generate_pwscf + +from enhanced_simulation import make_enhanced + +# Computer configuration +computer = 'ws16' + +if computer == 'baseline': + qe_modules = 'module purge; module load Core/25.05 gcc/12.4.0 openmpi/5.0.5 DefApps hdf5' + qe_bin = '/ccsopen/home/ksu/SOFTWARE/qe/q-e-qe-7.4.1/build/bin' + account = 'phy191' +elif computer == 'ws16': + qe_modules = '' + qe_bin = '/Users/ksu/Software/qe/q-e-qe-7.5/build/bin' + account = 'ks' +else: + print('Undefined computer') + exit() + +# Loop behavior notes: +# - Loops reuse the same directory and identifier (unlike error handlers which create +# new paths like "attempt_1", "attempt_2", etc.) +# - Simulation ID (simid) remains constant across all loop iterations +# - Each iteration gets a NEW scheduler job ID (process_id) when resubmitted +# (process_id is updated via update_process_id() when the new job is submitted, +# replacing the previous scheduler ID) +# - iteration_count tracks current loop iteration (0 = first run, 1 = second, etc.) +# - Status output shows [iter=N] for simulations with iteration_count > 0 +# - Access iteration_count programmatically: sim.iteration_count +# - Each iteration overwrites files in the same directory (no attempt subdirectories) +# - Loop continues until loop_condition returns False OR loop_max_iterations is reached + +LOOP_TARGET = 3 + +settings( + pseudo_dir = '../pseudopotentials', + results = '', + sleep = 1, + machine = computer, + account = account, + ) + + +system = generate_physical_system( + units = 'A', + axes = '''1.785 1.785 0.000 + 0.000 1.785 1.785 + 1.785 0.000 1.785''', + elem_pos = ''' + C 0.0000 0.0000 0.0000 + C 0.8925 0.8925 0.8925 + ''', + C = 4, + ) + + +# Seed calculation (plain DAG simulation) +qe_job = Job(nodes=1, + queue='batch', + hours=1, + presub=qe_modules, + app=qe_bin+'/pw.x') + +scf_seed = generate_pwscf( + identifier = 'scf_seed', + path = 'diamond/scf_seed', + job = qe_job, + input_type = 'generic', + calculation = 'scf', + input_dft = 'lda', + ecutwfc = 200, + conv_thr = 1e-3, + system = system, + pseudos = ['C.BFD.upf'], + kgrid = (4,4,4), + kshift = (0,0,0), + ) + + +def continue_loop(sim): + """ + Loop condition: keep iterating until the target count is reached. + The sim.iteration_count is incremented automatically after each pass. + """ + return sim.iteration_count < LOOP_TARGET + + +def modify_loop_input(sim): + """ + Loop modifier: modify simulation input at each iteration. + + This function is called before each loop iteration (after iteration_count + is incremented but before the simulation is resubmitted). You can modify + any simulation attributes here, such as: + - Input parameters (ecutwfc, conv_thr, etc.) + - System geometry + - K-point grid + - Any other simulation attributes + + Example: Gradually increase convergence threshold at each iteration + """ + # Example: Modify convergence threshold based on iteration + # Start with loose threshold, tighten as iterations progress + base_threshold = 1e-4 + iteration_factor = 10 ** (-sim.iteration_count) + sim.input.electrons.conv_thr = base_threshold * iteration_factor + # Example: Store iteration-specific values in loop_variables + sim.loop_variables['modified_conv_thr'] = sim.input.electrons.conv_thr + + + +loop_base = generate_pwscf( + identifier = 'scf_loop', + path = 'diamond/scf_loop', + job = qe_job, + input_type = 'generic', + calculation = 'scf', + input_dft = 'lda', + ecutwfc = 200, + # conv_thr = 1e-3, + system = system, + pseudos = ['C.BFD.upf'], + kgrid = (4,4,4), + kshift = (0,0,0), + dependencies = (scf_seed, 'charge_density'), + ) + +loop_base.sync_from_scf = True + +scf_loop = make_enhanced( + loop_base, + loop_enabled=True, + loop_condition=continue_loop, + loop_max_iterations=LOOP_TARGET, + loop_modifier=modify_loop_input, # Modify simulation at each iteration + ) + +run_project() + diff --git a/nexus/examples/quantum_espresso/05_conditionals/run.py b/nexus/examples/quantum_espresso/05_conditionals/run.py new file mode 100755 index 0000000000..fd799be059 --- /dev/null +++ b/nexus/examples/quantum_espresso/05_conditionals/run.py @@ -0,0 +1,149 @@ +#! /usr/bin/env python3 +""" +Using conditional utilities from conditionals.py + +This example demonstrates the conditional utility functions provided in +nexus/lib/conditionals.py: + +machine_conditional - Run simulation only on specific machines + +The workflow: +- scf1: Initial SCF calculation (no dependencies) +- scf2: SCF that only runs on 'baseline' machine (machine_conditional) +- scf3: SCF that checks scf2's energy result (custom conditional function) +""" + +from nexus import settings, Job, run_project +from nexus import generate_physical_system +from nexus import generate_pwscf + +from enhanced_simulation import make_enhanced + +# Computer configuration +computer = 'ws16' + +if computer == 'baseline': + qe_modules = 'module purge; module load Core/25.05 gcc/12.4.0 openmpi/5.0.5 DefApps hdf5' + qe_bin = '/ccsopen/home/ksu/SOFTWARE/qe/q-e-qe-7.4.1/build/bin' + account = 'phy191' +elif computer == 'ws16': + qe_modules = '' + qe_bin = '/Users/ksu/Software/qe/q-e-qe-7.5/build/bin' + account = 'ks' +else: + print('Undefined computer') + exit() + +settings( + pseudo_dir = '../pseudopotentials', + results = '', + sleep = 1, + machine = computer, + account = account, + ) + +# Define physical system (diamond) +system = generate_physical_system( + units = 'A', + axes = '''1.785 1.785 0.000 + 0.000 1.785 1.785 + 1.785 0.000 1.785''', + elem_pos = ''' + C 0.0000 0.0000 0.0000 + C 0.8925 0.8925 0.8925 + ''', + C = 4, + ) + +# Create reusable job definition for baseline +qe_job = Job(nodes=1, + queue='batch', + hours=1, + presub=qe_modules, + app=qe_bin+'/pw.x') + +# Simulation 1: Initial SCF (no dependencies) +scf1 = generate_pwscf( + identifier = 'scf1', + path = 'diamond/scf1', + job = qe_job, + input_type = 'generic', + calculation = 'scf', + input_dft = 'lda', + ecutwfc = 200, + conv_thr = 1e-8, + system = system, + pseudos = ['C.BFD.upf'], + kgrid = (4,4,4), + kshift = (0,0,0), + ) + +# Simulation 2: Only runs on 'baseline' machine (machine_conditional) +scf2_base = generate_pwscf( + identifier = 'scf2', + path = 'diamond/scf2', + job = qe_job, + input_type = 'generic', + calculation = 'scf', + input_dft = 'lda', + ecutwfc = 200, + conv_thr = 1e-8, + system = system, + pseudos = ['C.BFD.upf'], + kgrid = (4,4,4), + kshift = (0,0,0), + dependencies = (scf1, 'charge_density'), + ) + +# Use machine_conditional to only run on baseline +scf2 = make_enhanced( + scf2_base, + required_machine='ws16', + ) + +# Simulation 3: Checks scf2's energy result +scf3_base = generate_pwscf( + identifier = 'scf3', + path = 'diamond/scf3', + job = qe_job, + input_type = 'generic', + calculation = 'scf', + input_dft = 'lda', + ecutwfc = 200, + conv_thr = 1e-8, + system = system, + pseudos = ['C.BFD.upf'], + kgrid = (4,4,4), + kshift = (0,0,0), + dependencies = (scf2, 'charge_density'), + ) + +# Use a custom conditional to check if scf2's energy is below threshold +def scf2_energy_below_threshold(sim): + """Custom conditional that checks scf2's energy via analyzer.""" + try: + scf2_sim = None + for dep in sim.dependencies.values(): + if dep.sim.identifier == 'scf2': + scf2_sim = dep.sim + break + if scf2_sim is None or not scf2_sim.finished: + return False + analyzer = scf2_sim.load_analyzer_image() + if hasattr(analyzer, 'E') and analyzer.E != 0.0: + condition = analyzer.E < -30.0 + if condition: + print(analyzer.E, " is less than -30.0, condition is True") + else: + print(analyzer.E, " is greater than -30.0, condition is False") + return condition + return False + except Exception: + return False + +scf3 = make_enhanced( + scf3_base, + condition=scf2_energy_below_threshold, + ) + +run_project() \ No newline at end of file diff --git a/nexus/examples/quantum_espresso/06_branching/run.py b/nexus/examples/quantum_espresso/06_branching/run.py new file mode 100755 index 0000000000..53c21b0530 --- /dev/null +++ b/nexus/examples/quantum_espresso/06_branching/run.py @@ -0,0 +1,204 @@ +#! /usr/bin/env python3 +""" +Nexus Enhanced Workflow Example: Diamond SCF with Conditional Execution + +This example demonstrates the new enhanced workflow features in Nexus: +1. Simulation dependencies (scf2 depends on scf1) +2. Error handlers with automatic resubmission +3. Conditional execution based on simulation results (energy threshold) +4. Branching workflow: scf3 or scf4 runs based on scf2's total energy + +The workflow: +- scf1: Initial SCF calculation (no dependencies) +- scf2: SCF calculation that depends on scf1's output +- scf3: Runs if scf2's energy is below threshold (depends on scf2) +- scf4: Runs if scf2's energy is above threshold (depends on scf2) +- Only one of scf3 or scf4 will execute based on the energy condition +- All simulations have error handlers that retry on failure +""" + +from nexus import settings, Job, run_project +from nexus import generate_physical_system +from nexus import generate_pwscf + +# Add lib directory to path for enhanced workflow imports +import sys +import os +nexus_lib = os.path.join(os.path.dirname(__file__), '../../lib') +if nexus_lib not in sys.path: + sys.path.insert(0, nexus_lib) + +# Import enhanced workflow features +from enhanced_simulation import make_enhanced, create_branch +from error_handler import RetryErrorHandler + +# Computer configuration +computer = 'ws16' + +if computer == 'baseline': + qe_modules = 'module purge; module load Core/25.05 gcc/12.4.0 openmpi/5.0.5 DefApps hdf5' + qe_bin = '/ccsopen/home/ksu/SOFTWARE/qe/q-e-qe-7.4.1/build/bin' + account = 'phy191' +elif computer == 'ws16': + qe_modules = '' + qe_bin = '/Users/ksu/Software/qe/q-e-qe-7.5/build/bin' + account = 'ks' +else: + print('Undefined computer') + exit() + + +settings( + pseudo_dir = '../pseudopotentials', + results = '', + sleep = 1, + machine = computer, + account = account, + ) + +# Define physical system (diamond) +system = generate_physical_system( + units = 'A', + axes = '''1.785 1.785 0.000 + 0.000 1.785 1.785 + 1.785 0.000 1.785''', + elem_pos = ''' + C 0.0000 0.0000 0.0000 + C 0.8925 0.8925 0.8925 + ''', + C = 4, + ) + +# Simple example with two simulations and a dependency +# scf1 runs first, scf2 depends on scf1's output + +# Create reusable job definition for baseline +qe_job = Job(nodes=1, + queue='batch', + hours=1, + presub=qe_modules, + app=qe_bin+'/pw.x') + +# Simulation 1: Initial SCF calculation (no dependencies) +# Testing without make_enhanced wrapper - this will work but without error handlers +scf1 = generate_pwscf( + identifier = 'scf', + path = 'diamond/scf', + job = qe_job, + input_type = 'generic', + calculation = 'scf', + input_dft = 'lda', + ecutwfc = 200, + conv_thr = 1e-8, + system = system, + pseudos = ['C.BFD.upf'], + kgrid = (4,4,4), + kshift = (0,0,0), + ) + +# Simulation 2: SCF calculation that depends on scf1 +# This depends on scf1's charge_density output +scf2_base = generate_pwscf( + identifier = 'scf_dependent', + path = 'diamond/scf_dependent', + job = qe_job, + input_type = 'generic', + calculation = 'scf', + input_dft = 'lda', + ecutwfc = 200, + conv_thr = 1e-8, + system = system, + pseudos = ['C.BFD.upf'], + kgrid = (4,4,4), + kshift = (0,0,0), + dependencies = (scf1, 'charge_density'), + ) + +# Convert to enhanced simulation with error handler +scf2 = make_enhanced( + scf2_base, + error_handlers=[RetryErrorHandler(max_retries=2)], + # required_machine='andes', + ) + +# Define energy threshold (in Ry, typical for Quantum ESPRESSO) +# Adjust this value based on your system's expected energy range +energy_threshold = -30.0 + +# Conditional function: scf3 runs if scf2 energy is below threshold +def energy_below_threshold(sim): + """Check if scf2's energy is below threshold.""" + try: + # Get scf2 from dependencies (scf3 depends on scf2) + scf2_sim = None + for dep in sim.dependencies.values(): + if dep.sim.identifier == 'scf_dependent': + scf2_sim = dep.sim + break + + if scf2_sim is None or not scf2_sim.finished: + return False # scf2 not finished yet + + # Load analyzer and get energy + analyzer = scf2_sim.load_analyzer_image() + if hasattr(analyzer, 'E') and analyzer.E != 0.0: + energy = analyzer.E + condition = energy < energy_threshold + if condition: + print(energy, " is less than -30.0, condition is True") + else: + print(energy, " is greater than -30.0, condition is False") + return condition + else: + return False # Energy not available + except Exception: + # If anything fails, don't run + return False + + +# Simulation 3: Runs if scf2 energy is below threshold +scf3_base = generate_pwscf( + identifier = 'scf_low_energy', + path = 'diamond/scf_low_energy', + job = qe_job, + input_type = 'generic', + calculation = 'scf', + input_dft = 'lda', + ecutwfc = 200, + conv_thr = 1e-8, + system = system, + pseudos = ['C.BFD.upf'], + kgrid = (4,4,4), + kshift = (0,0,0), + # Note: dependencies will be set by create_branch + ) + +# Simulation 4: Runs if scf2 energy is above threshold +scf4_base = generate_pwscf( + identifier = 'scf_high_energy', + path = 'diamond/scf_high_energy', + job = qe_job, + input_type = 'generic', + calculation = 'scf', + input_dft = 'lda', + ecutwfc = 200, + conv_thr = 1e-8, + system = system, + pseudos = ['C.BFD.upf'], + kgrid = (4,4,4), + kshift = (0,0,0), + # Note: dependencies will be set by create_branch + ) + +# Create if/else branch using create_branch +scf3, scf4 = create_branch( + parent=scf2, + branches=[ + (scf3_base, lambda sim: not energy_below_threshold(sim), 'charge_density'), + (scf4_base, energy_below_threshold, 'charge_density'), + ], + error_handlers=[RetryErrorHandler(max_retries=2)], +) + +run_project() + diff --git a/nexus/examples/quantum_espresso/07_merging/run.py b/nexus/examples/quantum_espresso/07_merging/run.py new file mode 100755 index 0000000000..1c2176037d --- /dev/null +++ b/nexus/examples/quantum_espresso/07_merging/run.py @@ -0,0 +1,214 @@ +#! /usr/bin/env python3 +""" +Branching and Merging Workflows + +This example demonstrates: +1. Natural branching using create_branch() for if/else conditionals +2. Merging branches with different strategies (first, any, all, custom selector) + +Branching: +- scf1 → scf2 (if condition A) OR scf3 (if condition B) +- Both scf2 and scf3 depend on scf1 + +Merging: +- scf4 depends on scf2 and scf3 with different merge strategies +- Demonstrates 'first', 'any', 'all', and custom selector options +""" + +from nexus import settings, Job, run_project +from nexus import generate_physical_system +from nexus import generate_pwscf +import random + +from enhanced_simulation import make_enhanced, create_branch + +# Computer configuration +computer = 'ws16' + +if computer == 'baseline': + qe_modules = 'module purge; module load Core/25.05 gcc/12.4.0 openmpi/5.0.5 DefApps hdf5' + qe_bin = '/ccsopen/home/ksu/SOFTWARE/qe/q-e-qe-7.4.1/build/bin' + account = 'phy191' +elif computer == 'ws16': + qe_modules = '' + qe_bin = '/Users/ksu/Software/qe/q-e-qe-7.5/build/bin' + account = 'ks' +else: + print('Undefined computer') + exit() + +settings( + pseudo_dir = '../pseudopotentials', + results = '', + sleep = 1, + machine = computer, + account = account, + ) + +# Define physical system (diamond) +system = generate_physical_system( + units = 'A', + axes = '''1.785 1.785 0.000 + 0.000 1.785 1.785 + 1.785 0.000 1.785''', + elem_pos = ''' + C 0.0000 0.0000 0.0000 + C 0.8925 0.8925 0.8925 + ''', + C = 4, + ) + +# Create reusable job definition for baseline +qe_job = Job(nodes=1, + queue='batch', + hours=1, + presub=qe_modules, + app=qe_bin+'/pw.x') + +# Parent simulation +scf1 = generate_pwscf( + identifier = 'scf1', + path = 'diamond/branch_merge/scf1', + job = qe_job, + input_type = 'generic', + calculation = 'scf', + input_dft = 'lda', + ecutwfc = 200, + conv_thr = 1e-8, + system = system, + pseudos = ['C.BFD.upf'], + kgrid = (4,4,4), + kshift = (0,0,0), + ) + +# Create branches using create_branch (natural API) +scf2 = generate_pwscf( + identifier = 'scf2', + path = 'diamond/branch_merge/scf2', + job = qe_job, + input_type = 'generic', + calculation = 'scf', + input_dft = 'lda', + ecutwfc = 200, + conv_thr = 1e-8, + system = system, + pseudos = ['C.BFD.upf'], + kgrid = (4,4,4), + kshift = (0,0,0), + dependencies = (scf1, 'charge_density'), + ) + +scf3 = generate_pwscf( + identifier = 'scf3', + path = 'diamond/branch_merge/scf3', + job = qe_job, + input_type = 'generic', + calculation = 'scf', + input_dft = 'lda', + ecutwfc = 200, + conv_thr = 1e-8, + system = system, + pseudos = ['C.BFD.upf'], + kgrid = (4,4,4), + kshift = (0,0,0), + dependencies = (scf1, 'charge_density'), + ) + +# Now demonstrate merging strategies +# scf4 will merge scf2 and scf3 with different strategies + +# Example 1: Merge with 'first' strategy (race condition) +scf4_first = generate_pwscf( + identifier = 'scf4_first', + path = 'diamond/branch_merge/scf4_first', + job = qe_job, + input_type = 'generic', + calculation = 'scf', + input_dft = 'lda', + ecutwfc = 200, + conv_thr = 1e-8, + system = system, + pseudos = ['C.BFD.upf'], + kgrid = (4,4,4), + kshift = (0,0,0), + ) + +scf4_first_enhanced = make_enhanced(scf4_first) +# Run when FIRST of scf2 or scf3 completes +scf4_first_enhanced.depends((scf2, 'charge_density'), (scf3, 'charge_density'), strategy='first') + + +# Example 3: Merge with 'all' strategy (default, AND logic) +scf4_all = generate_pwscf( + identifier = 'scf4_all', + path = 'diamond/branch_merge/scf4_all', + job = qe_job, + input_type = 'generic', + calculation = 'scf', + input_dft = 'lda', + ecutwfc = 200, + conv_thr = 1e-8, + system = system, + pseudos = ['C.BFD.upf'], + kgrid = (4,4,4), + kshift = (0,0,0), + ) + +scf4_all_enhanced = make_enhanced(scf4_all) +# Run when ALL of scf2 and scf3 complete (default behavior) +scf4_all_enhanced.depends((scf2, 'other'), (scf3, 'other'), strategy='all') + +# Example 4: Merge with custom selector function +def select_lowest_energy(sims): + """ + Custom selector: choose simulation with lowest energy. + + Args: + sims: List of completed Simulation objects + + Returns: + Simulation with lowest energy + """ + if not sims: + return None + + # Load analyzer for each sim and compare energies + best_sim = None + best_energy = float('inf') + + for sim in sims: + try: + if sim.finished: + analyzer = sim.load_analyzer_image() + if hasattr(analyzer, 'E') and analyzer.E != 0.0: + energy = analyzer.E + if energy < best_energy: + best_energy = energy + best_sim = sim + except Exception: + continue + + # Return best sim, or first sim if no energies found + return best_sim if best_sim else sims[0] + +scf4_custom = generate_pwscf( + identifier = 'scf4_custom', + path = 'diamond/branch_merge/scf4_custom', + job = qe_job, + input_type = 'generic', + calculation = 'scf', + input_dft = 'lda', + ecutwfc = 200, + conv_thr = 1e-8, + system = system, + pseudos = ['C.BFD.upf'], + kgrid = (4,4,4), + kshift = (0,0,0), + ) + +scf4_custom_enhanced = make_enhanced(scf4_custom) +# Run when ALL complete, then select one with lowest energy +scf4_custom_enhanced.depends((scf2, 'charge_density'), (scf3, 'charge_density'), strategy='all', selector=select_lowest_energy) + +# Run the workflow +run_project() diff --git a/nexus/examples/quantum_espresso/08_error_handlers/run.py b/nexus/examples/quantum_espresso/08_error_handlers/run.py new file mode 100755 index 0000000000..a8fee64cdb --- /dev/null +++ b/nexus/examples/quantum_espresso/08_error_handlers/run.py @@ -0,0 +1,217 @@ +#! /usr/bin/env python3 +""" +Error Handler Demonstration with PwscfErrorHandler + +This example demonstrates the error handler system, specifically error handling for Quantum ESPRESSO pw.x simulations. + +Features demonstrated: +1. PwscfErrorHandler - QE-specific handler that modifies input on errors +2. RetryErrorHandler - Simple retry handler (for comparison) +3. Error strategies - Compact objects that combine diagnosis and handling +4. Modification strategies - Reusable strategies for input modification +5. Compatibility checking - Handlers check if they're compatible with simulations + +The workflow: +- scf1: Default PwscfErrorHandler (auto-modifies input on convergence failures) +- scf2: RetryErrorHandler (simple retry, no diagnosis) +- scf3: Multiple handlers (compatibility checked automatically) +- scf4: Custom modification strategies (user-defined input modifications) +""" + +from nexus import settings, Job, run_project +from nexus import generate_physical_system +from nexus import generate_pwscf + +import os +import sys + +nexus_lib = os.path.join(os.path.dirname(__file__), '../../lib') +if nexus_lib not in sys.path: + sys.path.insert(0, nexus_lib) + +from enhanced_simulation import make_enhanced +from error_handler import ( + RetryErrorHandler, +) +from pwscf_error_handler import ( + PwscfErrorHandler, + RelaxConvergenceThreshold, + IncreaseMaxSteps, + AdjustMixingBeta, + UseConjugateGradient, + CombinedSimulationModifier, +) + +# Computer configuration +computer = 'ws16' + +if computer == 'baseline': + qe_modules = 'module purge; module load Core/25.05 gcc/12.4.0 openmpi/5.0.5 DefApps hdf5' + qe_bin = '/ccsopen/home/ksu/SOFTWARE/qe/q-e-qe-7.4.1/build/bin' + account = 'phy191' +elif computer == 'ws16': + qe_modules = '' + qe_bin = '/Users/ksu/Software/qe/q-e-qe-7.5/build/bin' + account = 'ks' +else: + print('Undefined computer') + exit() + +settings( + pseudo_dir = '../pseudopotentials', + results = '', + sleep = 1, + machine = computer, + account = account, + ) + +# Define physical system (diamond) +system = generate_physical_system( + units = 'A', + axes = '''1.785 1.785 0.000 + 0.000 1.785 1.785 + 1.785 0.000 1.785''', + elem_pos = ''' + C 0.0000 0.0000 0.0000 + C 0.8925 0.8925 0.8925 + ''', + C = 4, + ) + +qe_job = Job(nodes=1, + queue='batch', + hours=0.25, + presub=qe_modules, + app=qe_bin+'/pw.x') + +# Simulation 1: PwscfErrorHandler (default configuration) +# This handler uses error strategies to diagnose QE errors and automatically +# modifies input parameters to help the simulation succeed on retry. +# Default behavior: On convergence failures, it will: +# - Relax conv_thr by factor of 10 +# - Increase electron_maxstep by 200 +# - Reduce mixing_beta by 0.1 +scf1_base = generate_pwscf( + identifier = 'scf1', + path = 'diamond/error_handlers/scf1', + job = qe_job, + input_type = 'generic', + calculation = 'scf', + input_dft = 'lda', + ecutwfc = 200, + conv_thr = 1e-8, + electron_maxstep = 100, + system = system, + pseudos = ['C.BFD.upf'], + kgrid = (4,4,4), + kshift = (0,0,0), + ) + +# Use PwscfErrorHandler with default modification strategies +# If convergence fails, input will be automatically modified before retry +scf1 = make_enhanced( + scf1_base, + error_handlers=[ + PwscfErrorHandler( + max_retries=3, + delay=1.0, # Wait 1 second before retry + ) + ], +) + +# Simulation 2: RetryErrorHandler (simple) +# Simple retry handler - no error diagnosis, just retries up to max_retries +scf2_base = generate_pwscf( + identifier = 'scf2', + path = 'diamond/error_handlers/scf2', + job = qe_job, + input_type = 'generic', + calculation = 'scf', + input_dft = 'lda', + ecutwfc = 200, + conv_thr = 1e-8, + system = system, + pseudos = ['C.BFD.upf'], + kgrid = (4,4,4), + kshift = (0,0,0), + dependencies = (scf1, 'charge_density'), + ) + +scf2 = make_enhanced( + scf2_base, + error_handlers=[ + RetryErrorHandler(max_retries=2, delay=1.0) + ], +) + +# Simulation 3: Multiple handlers (compatibility checked) +# Demonstrates that handlers check compatibility automatically and are called in order. +# The FIRST handler that returns True will trigger resubmission and stop the chain. +scf3_base = generate_pwscf( + identifier = 'scf3', + path = 'diamond/error_handlers/scf3', + job = qe_job, + input_type = 'generic', + calculation = 'scf', + input_dft = 'lda', + ecutwfc = 200, + conv_thr = 1e-8, + system = system, + pseudos = ['C.BFD.upf'], + kgrid = (4,4,4), + kshift = (0,0,0), + dependencies = (scf2, 'charge_density'), + ) + +# Multiple handlers - compatibility is checked automatically +scf3 = make_enhanced( + scf3_base, + error_handlers=[ + PwscfErrorHandler(max_retries=2), # Tried first: QE-specific with input modification + RetryErrorHandler(max_retries=1), # Last resort: simple retry + ], +) + +# Simulation 4: Custom modification strategies +# Demonstrates how to create custom modification strategies for specific error types +scf4_base = generate_pwscf( + identifier = 'scf4', + path = 'diamond/error_handlers/scf4', + job = qe_job, + input_type = 'generic', + calculation = 'scf', + input_dft = 'lda', + ecutwfc = 200, + conv_thr = 1e-8, # Tight convergence - might fail + electron_maxstep = 100, # Limited steps - might need more + system = system, + pseudos = ['C.BFD.upf'], + kgrid = (4,4,4), + kshift = (0,0,0), + dependencies = (scf3, 'charge_density'), + ) + +# Use PwscfErrorHandler with custom modification strategies +# This demonstrates the compact strategy-based API: +# - Error strategies define what errors to detect +# - Modification strategies define how to fix them +# - Strategies can be combined and customized +scf4 = make_enhanced( + scf4_base, + error_handlers=[ + PwscfErrorHandler( + max_retries=3, + delay=1.0, + modification_strategies={ + 'convergence_failure': CombinedSimulationModifier([ + RelaxConvergenceThreshold(factor=10.0), # conv_thr *= 10 + IncreaseMaxSteps(increase=200), # electron_maxstep += 200 + AdjustMixingBeta(reduction=0.1), # mixing_beta -= 0.1 (min 0.1) + UseConjugateGradient(), # Use CG diagonalization + ]), + } + ) + ], +) + +run_project() diff --git a/nexus/lib/conditionals.py b/nexus/lib/conditionals.py new file mode 100644 index 0000000000..7b34667145 --- /dev/null +++ b/nexus/lib/conditionals.py @@ -0,0 +1,150 @@ +""" +Conditional utilities for enhanced workflow simulations. + +Provides utility functions for conditional execution and machine filtering. +""" + +from typing import Callable, Optional, Any, Dict + + +def _add_conditional_metadata(func: Callable, description: str, metadata: Optional[Dict[str, Any]] = None) -> Callable: + """ + Add metadata to a conditional function for better status reporting. + + Args: + func: Conditional function to annotate + description: Human-readable description of the conditional + metadata: Additional metadata dictionary + + Returns: + Function with metadata attached + """ + func._conditional_description = description + func._conditional_metadata = metadata or {} + return func + + +def get_conditional_description(condition: Optional[Callable]) -> str: + """ + Get a human-readable description of a conditional function. + + Args: + condition: Conditional function (may have metadata attached) + + Returns: + Description string, or default if no metadata available + """ + if condition is None: + return "None" + + # Handle non-callable objects gracefully + if not callable(condition): + return f"" + + try: + # Check for attached metadata + if hasattr(condition, '_conditional_description'): + return condition._conditional_description + + # Try to infer from function name + if hasattr(condition, '__name__'): + name = condition.__name__ + if name != '': + return name + + # For lambda functions, try to get source if available + if hasattr(condition, '__code__'): + # Could extract more info from __code__ if needed + pass + + # Default fallback + return "custom" + except Exception: + # If anything goes wrong, return a safe fallback + return "custom" + + +def machine_conditional(machine_name: str, *, required: bool = True) -> Callable: + """ + Create a conditional function that checks machine compatibility. + + Args: + machine_name: Name of machine to check + required: If True, must match; if False, must not match + + Returns: + Conditional function + """ + def check_machine(sim: Any) -> bool: + current_machine = getattr(sim, 'job', None) + if current_machine is None: + return not required + + machine_name_attr = getattr(current_machine, 'machine', None) + if machine_name_attr is None: + return not required + + matches = machine_name_attr == machine_name + return matches if required else not matches + + description = f"machine={'==' if required else '!='}{machine_name}" + return _add_conditional_metadata(check_machine, description, {'machine_name': machine_name, 'required': required}) + + +def result_conditional(dependency: str, check_func: Callable[[Any], bool]) -> Callable: + """ + Create a conditional function that checks dependency results. + + Args: + dependency: Name of dependency to check + check_func: Function that takes dependency result and returns bool + + Returns: + Conditional function + """ + def check_result(sim: Any) -> bool: + if not hasattr(sim, 'dependencies'): + return False + + for dep in sim.dependencies.values(): + dep_sim = dep.sim + if hasattr(dep_sim, 'outputs') and dep_sim.outputs is not None: + result = getattr(dep_sim.outputs, dependency, None) + if result is not None: + return check_func(result) + + return False + + check_func_name = getattr(check_func, '__name__', 'custom') + description = f"result:{dependency}({check_func_name})" + return _add_conditional_metadata(check_result, description, {'dependency': dependency, 'check_func': check_func_name}) + + +def combine_conditionals(*conditionals: Callable, + operator: str = 'and') -> Callable: + """ + Combine multiple conditionals with AND or OR logic. + + Args: + *conditionals: Conditional functions to combine + operator: 'and' or 'or' + + Returns: + Combined conditional function + """ + allowed_operators = ['and', 'or'] + if operator not in allowed_operators: + raise ValueError(f"Unknown operator: {operator}. Use one of: {allowed_operators}") + + def combined(sim: Any) -> bool: + if operator == 'and': + return all(cond(sim) for cond in conditionals) + elif operator == 'or': + return any(cond(sim) for cond in conditionals) + + cond_descriptions = [get_conditional_description(cond) for cond in conditionals] + op_str = f' {operator.upper()} ' + description = f"({op_str.join(cond_descriptions)})" + return _add_conditional_metadata(combined, description, {'operator': operator, 'conditionals': conditionals}) + + diff --git a/nexus/lib/enhanced_project_manager.py b/nexus/lib/enhanced_project_manager.py new file mode 100644 index 0000000000..309aec3a8d --- /dev/null +++ b/nexus/lib/enhanced_project_manager.py @@ -0,0 +1,623 @@ +""" +EnhancedProjectManager class extending ProjectManager with enhanced DAG workflow support. + +Provides conditional filtering, enhanced progress tracking, and blocked simulation +handling for DAG workflows with loops, conditionals, and error handlers. +""" + +from typing import Dict, Set, List, Optional, Any +from developer import obj +from project_manager import ProjectManager, trivial +from simulation import Simulation +from enhanced_simulation import EnhancedSimulation + + +class EnhancedProjectManager(ProjectManager): + """ + Extended ProjectManager with enhanced DAG workflow support. + + Inheritance Pattern: + - Inherits from ProjectManager using standard Python inheritance + - Calls super().__init__() to initialize base attributes (simulations, cascades, progressing_cascades) + - Adds enhanced tracking: filtered_sims, blocked simulation reporting + + Method Overrides: + + 1. add_simulations(): + Base: Adds all non-fake simulations to cascades + Override: Adds pre-filtering for machine compatibility, definition-time conditionals, + and transitive dependency blocking. Snapshot blocked sims before detaching. + Calls super().add_simulations(final_verified) at end. + + 2. run_project(): + Base: Initializes cascades and runs project + Override: Detaches/prunes blocked sims BEFORE calling super().run_project() + + 3. init_cascades(): + Base: Screens fake sims, resolves collisions, propagates blockages + Override: Detaches/prunes blocked sims BEFORE calling super().init_cascades() + + 4. screen_fake_sims(): + Base: Errors if any fake sims found + Override: Allows intentionally blocked EnhancedSimulations (skips them in error check) + + 5. traverse_cascades(): + Base: Standard DAG traversal + Override: Uses standard DAG traversal (all workflows are DAGs) + + 6. write_simulation_status(): + Base: Writes status for all simulations + Override: Filters blocked sims from main status, adds separate "blocked simulations" section + + 7. status_line(): + Base: Writes single status line + Override: Adds [iter=N] annotation for simulations with loops (iteration_count > 0) + + New Methods (Not in Base): + - _detach_blocked_simulations(): Removes blocked sims from dependency graph + - _prune_blocked_from_tracking(): Cleans up tracking dictionaries + - _snapshot_blocked(): Captures blocked sim info before graph mutation + - _get_blocked_simulations(): Identifies all blocked sims (direct + transitive) + + Extension Strategy: + - Pre-processing: Filters/validates before calling parent methods + - Post-processing: Adds tracking/reporting after parent methods + - Side-by-side: Adds new functionality (blocked sim reporting, enhanced status) + - Selective override: Only overrides methods needed for enhanced DAG workflows + + Backward Compatibility: + - DAG workflows unchanged: If no EnhancedSimulation instances exist, behavior matches base class (see run_project in nexus.py) + - Base class methods called via super() to preserve original functionality + - New features opt-in: Only EnhancedSimulation instances trigger enhanced functionalities + """ + + def __init__(self): + super().__init__() + self.filtered_sims: List[Simulation] = [] # Track simulations filtered out + self._blocked_report: Dict[int, Dict[str, Any]] = {} + self._blocked_ids: Set[int] = set() + + def add_simulations(self, *simulations): + """Override to filter by definition-time conditionals and blocked dependencies.""" + if len(simulations) == 0: + simulations = Simulation.all_sims + if len(simulations) > 0 and not isinstance(simulations[0], Simulation): + simulations = simulations[0] + + # FIRST: Check and mark all simulations as fake if they're machine incompatible + # This must happen before any other processing + # Use direct comparison to avoid any issues with is_machine_compatible() + for sim in simulations: + if isinstance(sim, EnhancedSimulation): + required = getattr(sim, 'required_machine', None) + allowed = getattr(sim, 'allowed_machines', None) + + if required or (allowed and len(allowed) > 0): + # Get current machine + machine_name = None + if hasattr(sim, 'job') and sim.job is not None: + machine_name = getattr(sim.job, 'machine', None) + if machine_name is None: + from machines import Job + machine_name = Job.machine + if machine_name is None: + from project_manager import ProjectManager + if ProjectManager.machine is not None: + machine_name = getattr(ProjectManager.machine, 'name', None) + + # Direct comparison + if required: + if machine_name is None or str(machine_name).strip().lower() != str(required).strip().lower(): + sim.fake_sim = True + continue + if allowed and len(allowed) > 0: + if machine_name is None: + sim.fake_sim = True + continue + machine_normalized = str(machine_name).strip().lower() + allowed_normalized = {str(a).strip().lower() for a in allowed} + if machine_normalized not in allowed_normalized: + sim.fake_sim = True + continue + + # Filter by definition-time conditionals and machine compatibility + filtered_sims = [] + self.filtered_sims = [] # Reset filtered list + blocked_ids = set() # Track sim_ids of blocked simulations + + # First pass: identify directly blocked simulations + for sim in simulations: + # Check fake status first - use multiple methods to be sure + fake_status = sim.fake() if hasattr(sim, 'fake') else False + fake_attr = getattr(sim, 'fake_sim', False) + + # Also check machine compatibility directly (in case fake_sim wasn't set) + if isinstance(sim, EnhancedSimulation): + required = getattr(sim, 'required_machine', None) + if required: + machine_name = None + if hasattr(sim, 'job') and sim.job is not None: + machine_name = getattr(sim.job, 'machine', None) + if machine_name is None: + from machines import Job + machine_name = Job.machine + if machine_name is None: + from project_manager import ProjectManager + if ProjectManager.machine is not None: + machine_name = getattr(ProjectManager.machine, 'name', None) + if machine_name and str(machine_name).strip().lower() != str(required).strip().lower(): + fake_attr = True # Force fake if machine doesn't match + + if fake_status or fake_attr: + # Ensure it's marked as fake + sim.fake_sim = True + self.filtered_sims.append(sim) + blocked_ids.add(sim.simid) + continue + + # Check if it's a EnhancedSimulation with definition-time conditionals + if isinstance(sim, EnhancedSimulation): + # Machine filtering - check compatibility + # This is critical for make_enhanced instances that bypass __init__ + if not sim.is_machine_compatible(): + sim.fake_sim = True # Ensure it's marked as fake + self.filtered_sims.append(sim) + blocked_ids.add(sim.simid) + continue + + + filtered_sims.append(sim) + # Second pass: filter out simulations that depend on blocked simulations + # Iterate until no new blocked simulations are found (transitive closure) + # Create lookup maps for blocked simulations + blocked_by_id = {} # simid -> sim + blocked_by_identifier = {} # identifier -> sim + for sim in self.filtered_sims: + blocked_by_id[sim.simid] = sim + blocked_by_identifier[sim.identifier] = sim + + # Also check all simulations to see if any dependency would be blocked + # (dependencies might reference simulations not in the input list) + all_sims_check = Simulation.all_sims if hasattr(Simulation, 'all_sims') else [] + for sim in all_sims_check: + if sim.simid in blocked_ids: + blocked_by_id[sim.simid] = sim + blocked_by_identifier[sim.identifier] = sim + elif isinstance(sim, EnhancedSimulation): + # Check if this simulation would be blocked + if not sim.is_machine_compatible(): + blocked_by_id[sim.simid] = sim + blocked_by_identifier[sim.identifier] = sim + if sim not in self.filtered_sims: + self.filtered_sims.append(sim) + blocked_ids.add(sim.simid) + + changed = True + while changed: + changed = False + remaining_sims = [] + for sim in filtered_sims: + # Check if any dependency is blocked + has_blocked_dep = False + blocked_dep_identifier = None + + # Check dependency sim objects directly (most reliable) + for dep in sim.dependencies.values(): + if hasattr(dep, 'sim'): + dep_sim = dep.sim + # Check by simid + if dep_sim.simid in blocked_ids or dep_sim.simid in blocked_by_id: + has_blocked_dep = True + blocked_dep_identifier = dep_sim.identifier + break + # Check by identifier + if dep_sim.identifier in blocked_by_identifier: + has_blocked_dep = True + blocked_dep_identifier = dep_sim.identifier + break + # Check if dependency sim is in filtered list + if dep_sim in self.filtered_sims: + has_blocked_dep = True + blocked_dep_identifier = dep_sim.identifier + break + + # Also check dependency_ids set + if not has_blocked_dep: + for dep_id in sim.dependency_ids: + if dep_id in blocked_ids or dep_id in blocked_by_id: + has_blocked_dep = True + # Find the identifier for reporting + if dep_id in blocked_by_id: + blocked_dep_identifier = blocked_by_id[dep_id].identifier + break + + # Also check dependencies dict keys + if not has_blocked_dep: + for dep_id in sim.dependencies.keys(): + if dep_id in blocked_ids or dep_id in blocked_by_id: + has_blocked_dep = True + if dep_id in blocked_by_id: + blocked_dep_identifier = blocked_by_id[dep_id].identifier + break + + if has_blocked_dep: + self.filtered_sims.append(sim) + blocked_ids.add(sim.simid) + blocked_by_id[sim.simid] = sim + blocked_by_identifier[sim.identifier] = sim + changed = True + else: + remaining_sims.append(sim) + + filtered_sims = remaining_sims + + # Final safety check: remove any blocked simulations from filtered_sims + # (in case they slipped through) + final_filtered = [] + for sim in filtered_sims: + if sim.simid not in blocked_ids and not sim.fake() and not getattr(sim, 'fake_sim', False): + # Double-check machine compatibility for EnhancedSimulations + if isinstance(sim, EnhancedSimulation): + if sim.is_machine_compatible(): + final_filtered.append(sim) + else: + # Shouldn't happen, but be safe + self.filtered_sims.append(sim) + blocked_ids.add(sim.simid) + else: + final_filtered.append(sim) + else: + # Shouldn't happen, but be safe + if sim not in self.filtered_sims: + self.filtered_sims.append(sim) + blocked_ids.add(sim.simid) + + # Final verification: ensure no fake/blocked simulations in final_filtered + # This is a critical safety check + verified_filtered = [] + for sim in final_filtered: + # Triple-check: fake status, blocked_ids, and machine compatibility + if (sim.simid not in blocked_ids and + not sim.fake() and + not getattr(sim, 'fake_sim', False)): + if isinstance(sim, EnhancedSimulation): + # One more machine compatibility check + if sim.is_machine_compatible(): + verified_filtered.append(sim) + else: + # Shouldn't happen, but be safe + if sim not in self.filtered_sims: + self.filtered_sims.append(sim) + blocked_ids.add(sim.simid) + else: + verified_filtered.append(sim) + else: + # Shouldn't happen, but be safe + if sim not in self.filtered_sims: + self.filtered_sims.append(sim) + blocked_ids.add(sim.simid) + + # Call parent method with verified filtered simulations + # But first, double-check that all simulations are not fake + final_verified = [] + for sim in verified_filtered: + # CRITICAL: Check fake status one more time before adding + if sim.fake() or getattr(sim, 'fake_sim', False): + # Should never happen, but skip if fake + continue + if isinstance(sim, EnhancedSimulation): + # Double-check machine compatibility + required = getattr(sim, 'required_machine', None) + if required: + machine_name = None + if hasattr(sim, 'job') and sim.job is not None: + machine_name = getattr(sim.job, 'machine', None) + if machine_name is None: + from machines import Job + machine_name = Job.machine + if machine_name is None: + from project_manager import ProjectManager + if ProjectManager.machine is not None: + machine_name = getattr(ProjectManager.machine, 'name', None) + if machine_name and str(machine_name).strip().lower() != str(required).strip().lower(): + # Machine doesn't match - skip + continue + # Explicitly ensure fake_sim is False + sim.fake_sim = False + final_verified.append(sim) + + # Snapshot blocked sims (includes downstream dependents) before detaching + _, snapshot_blocked_ids = self._snapshot_blocked() + + # Eliminate blocked simulations from dependency graph before adding new ones + self._detach_blocked_simulations(snapshot_blocked_ids) + + # Final check: ensure parent method won't add any fake simulations + # The parent checks sim.fake(), so we need to make sure fake_sim is False + super().add_simulations(final_verified) + + # Remove any blocked simulations that might have been added + # (e.g., if parent method added them despite our filtering) + self._prune_blocked_from_tracking(snapshot_blocked_ids) + + def screen_fake_sims(self): + """Override to allow blocked simulations (they're intentionally fake).""" + # For EnhancedProjectManager, blocked simulations are intentionally marked as fake + # and should not cause an error. We've already filtered them out in add_simulations + # and run_project, so if any remain, they're legitimate fake sims that should error. + # But we need to check only simulations that are actually in cascades. + def collect_fake(sim, fake): + # Only collect if it's fake AND not intentionally blocked + if sim.fake(): + # Check if it's a blocked EnhancedSimulation + if isinstance(sim, EnhancedSimulation): + # If it's in our filtered list, it's intentionally blocked - skip + if sim in getattr(self, 'filtered_sims', []): + return + fake.append(sim) + #end def collect_fake + fake = [] + self.traverse_cascades(collect_fake, fake) + if len(fake) > 0: + msg = 'fake/temporary simulation objects detected in cascade\nthis is a developer error\nlist of fake sims and directories:\n' + for sim in fake: + msg += ' {0:>8} {1}\n'.format(sim.simid, sim.locdir) + #end for + self.error(msg) + #end if + #end def screen_fake_sims + + def run_project(self, status=False, status_only=False): + """Override to ensure blocked simulations are removed before running.""" + # Before running, ensure all blocked simulations are removed + # This is critical for generate_only and other operations + _, blocked_ids = self._snapshot_blocked() + self._detach_blocked_simulations(blocked_ids) + self._prune_blocked_from_tracking(blocked_ids) + # Now call parent method + super().run_project(status=status, status_only=status_only) + + def init_cascades(self): + """Override to ensure blocked simulations are removed before initializing cascades.""" + # Remove blocked simulations BEFORE initializing cascades + # This is critical because cascades are created in add_simulations + _, blocked_ids = self._snapshot_blocked() + self._detach_blocked_simulations(blocked_ids) + self._prune_blocked_from_tracking(blocked_ids) + # Now call parent method + super().init_cascades() + + def traverse_cascades(self, operation=None, *args, **kwargs): + """Override to use standard DAG traversal.""" + # Reset wait_ids for all cascades + for cascade in self.cascades.values(): + cascade.reset_wait_ids() + + # Use standard DAG traversal (all workflows are DAGs) + for cascade in self.cascades.values(): + cascade.traverse_cascade(operation or trivial, *args, **kwargs) + + def _detach_blocked_simulations(self, blocked_ids: Set[int]): + """Remove blocked simulations from dependency graph (dependents and dependencies).""" + if not blocked_ids: + return + + sim_lookup = {} + all_sims = getattr(Simulation, 'all_sims', []) + for sim in all_sims: + if hasattr(sim, 'simid'): + sim_lookup[sim.simid] = sim + + for sim_id in blocked_ids: + blocked_sim = sim_lookup.get(sim_id) + if blocked_sim is None: + continue + + if hasattr(blocked_sim, 'eliminate'): + blocked_sim.eliminate() + else: + # Manual removal as fallback + dependencies = list(getattr(blocked_sim, 'dependencies', {}).values()) + for dep in dependencies: + upstream = getattr(dep, 'sim', None) + if upstream and hasattr(upstream, 'dependents'): + upstream.dependents.pop(sim_id, None) + dependents = list(getattr(blocked_sim, 'dependents', {}).values()) + for dependent in dependents: + if hasattr(dependent, 'dependencies') and sim_id in dependent.dependencies: + dependent.undo_depends(blocked_sim) + + # Ensure blocked simulations remain blocked/fake + blocked_sim.block = True + blocked_sim.block_subcascade = True + blocked_sim.fake_sim = True + + def _prune_blocked_from_tracking(self, blocked_ids: Set[int]): + """Remove blocked simulations from manager tracking dictionaries.""" + if not blocked_ids: + return + + for sim_id in list(self.simulations.keys()): + sim = self.simulations[sim_id] + if (sim_id in blocked_ids or + sim.fake() or + getattr(sim, 'fake_sim', False) or + (isinstance(sim, EnhancedSimulation) and not sim.is_machine_compatible())): + del self.simulations[sim_id] + + for sim_id in list(self.cascades.keys()): + if sim_id in blocked_ids: + del self.cascades[sim_id] + + for sim_id in list(self.progressing_cascades.keys()): + if sim_id in blocked_ids: + del self.progressing_cascades[sim_id] + + def _snapshot_blocked(self): + """Capture blocked simulations before mutating the workflow graph.""" + blocked_sims, blocked_ids = self._get_blocked_simulations() + for sim, reasons in blocked_sims: + entry = dict( + simid=sim.simid, + identifier=str(sim.identifier), + path=str(sim.path), + reasons=list(reasons) + ) + self._blocked_report[sim.simid] = entry + self._blocked_ids.update(blocked_ids) + return blocked_sims, blocked_ids + + def _get_blocked_simulations(self): + """Identify all blocked simulations (directly and by dependencies).""" + blocked_sims = [] + directly_blocked = set() # Track sim_ids of directly blocked simulations + + # Check simulations that were filtered out during add_simulations + sims_to_check = [] + sims_to_check.extend(getattr(self, 'filtered_sims', [])) + + # Also check all_sims for any EnhancedSimulations that might be blocked + all_sims = Simulation.all_sims if hasattr(Simulation, 'all_sims') else [] + # Create a set of simids to avoid recursion issues with 'in' operator + sims_to_check_simids = {s.simid for s in sims_to_check if hasattr(s, 'simid')} + for sim in all_sims: + if isinstance(sim, EnhancedSimulation): + # Use simid comparison to avoid recursion issues with __eq__ + if hasattr(sim, 'simid') and sim.simid not in sims_to_check_simids: + # Check if it would be blocked + # Check if it would be blocked by definition-time conditional + is_definition_time = False + if not sim.is_machine_compatible() or (is_definition_time and not sim.evaluate_condition()): + sims_to_check.append(sim) + sims_to_check_simids.add(sim.simid) + + # First pass: identify directly blocked simulations + for sim in sims_to_check: + reasons = [] + + if isinstance(sim, EnhancedSimulation): + # Check machine compatibility + if not sim.is_machine_compatible(): + machine_name = getattr(sim.job, 'machine', None) if hasattr(sim, 'job') and sim.job else 'None' + required = getattr(sim, 'required_machine', None) + allowed = getattr(sim, 'allowed_machines', None) + + if required: + reasons.append(f"required_machine='{required}' (current: '{machine_name}')") + elif allowed and len(allowed) > 0: + reasons.append(f"allowed_machines={allowed} (current: '{machine_name}')") + else: + reasons.append(f"machine incompatible (current: '{machine_name}')") + if getattr(sim, 'skipped', False): + reasons.append("skipped") + + if reasons: + blocked_sims.append((sim, reasons)) + directly_blocked.add(sim.simid) + + # Second pass: check for simulations blocked by dependencies + # Iterate until no new blocked simulations are found (transitive closure) + changed = True + while changed: + changed = False + for sim in all_sims: + if sim.simid in directly_blocked: + continue + + # Check if any dependency is blocked + blocked_deps = [] + for dep_id, dep in sim.dependencies.items(): + if dep_id in directly_blocked: + dep_sim = dep.sim + blocked_deps.append(dep_sim.identifier) + + if blocked_deps: + reasons = [f"dependency blocked: {', '.join(blocked_deps)}"] + blocked_sims.append((sim, reasons)) + directly_blocked.add(sim.simid) + changed = True + + return blocked_sims, directly_blocked + + def write_simulation_status(self): + """Override to filter blocked simulations from normal status and add blocked section.""" + # Snapshot current state to ensure any new blocked sims are captured + self._snapshot_blocked() + blocked_ids = set(self._blocked_ids) + + # Temporarily remove blocked simulations from self.simulations for status output + original_simulations = self.simulations.copy() + filtered_simulations = obj() + for sim_id, sim in self.simulations.items(): + if sim_id not in blocked_ids: + filtered_simulations[sim_id] = sim + + # Temporarily replace simulations dict + self.simulations = filtered_simulations + + # Call parent method (will only show non-blocked simulations) + super().write_simulation_status() + + # Restore original simulations dict + self.simulations = original_simulations + + # Report blocked simulations + if self._blocked_report: + blocked_entries = sorted(self._blocked_report.values(), key=lambda x: x['identifier']) + self.log('\n==== blocked simulations (will not run) ====', n=1) + for entry in blocked_entries: + reason_str = '; '.join(entry['reasons']) + self.log(f" {entry['identifier']:30s} ({entry['path']:40s}) - {reason_str}", n=2) + self.log('', n=1) + + def status_line(self, sim, extra=''): + """Override to show iteration_count for simulations with loops.""" + from enhanced_simulation import EnhancedSimulation + indicators = ('setup', 'sent_files', 'submitted', 'finished', 'got_output', 'analyzed') + stats = sim.tuple(*indicators) + status = '' + for stat in stats: + status += str(int(stat)) + if sim.process_id is None: + pid = '------' + else: + pid = sim.process_id + # Add enhanced information for EnhancedSimulation instances + if isinstance(sim, EnhancedSimulation): + info_parts = [] + + # Loop information + if sim.loop_enabled: + iter_info = f"iter={sim.iteration_count}" + if sim.loop_max_iterations is not None: + iter_info += f"/{sim.loop_max_iterations}" + info_parts.append(iter_info) + + # Conditional information + cond_info = sim.get_conditional_info() + if cond_info: + info_parts.append(cond_info) + + # Merge strategy information + if sim.merge_selector is not None: + merge_info = f"merge={sim.merge_strategy}" + selector_name = getattr(sim.merge_selector, '__name__', 'custom') + merge_info += f":{selector_name}" + info_parts.append(merge_info) + + # Build status line with enhanced info + if info_parts: + enhanced_info = ' '.join(info_parts) + sline = '{0} {1} {2:<8} {3:<6} {4} [{5}]'.format( + status, str(int(sim.failed)), pid, sim.identifier, sim.locdir, enhanced_info + ) + else: + sline = '{0} {1} {2:<8} {3:<6} {4}'.format( + status, str(int(sim.failed)), pid, sim.identifier, sim.locdir + ) + else: + sline = '{0} {1} {2:<8} {3:<6} {4}'.format( + status, str(int(sim.failed)), pid, sim.identifier, sim.locdir + ) + self.log(sline, extra, n=2) + diff --git a/nexus/lib/enhanced_simulation.py b/nexus/lib/enhanced_simulation.py new file mode 100644 index 0000000000..12960b0eb3 --- /dev/null +++ b/nexus/lib/enhanced_simulation.py @@ -0,0 +1,1004 @@ +""" +EnhancedSimulation class extending Simulation with enhanced DAG workflow support. + +Provides error handlers, conditionals, machine filtering, and loops +for DAG workflows while maintaining backward compatibility. +""" + +from typing import Dict, Set, List, Optional, Callable, Any, Tuple +import os +import shutil +from datetime import datetime + +from developer import obj +from nexus_base import nexus_core +from simulation import Simulation +from error_handler import ErrorHandler + + +class EnhancedSimulation(Simulation): + """ + Extended Simulation class with enhanced DAG workflow support. + + Inheritance Pattern: + - Inherits from Simulation using standard Python inheritance + - Calls super().__init__(**kwargs) after extracting enhanced-specific kwargs + - Adds enhanced attributes: attempt_number, iteration_count, error_handlers, + condition, required_machine, allowed_machines, loop_enabled, etc. + + Method Overrides: + + 1. __init__(): + Base: Initializes standard simulation attributes (path, job, dependencies, etc.) + Override: Extracts enhanced kwargs first, initializes enhanced state, calls super(), + then performs machine filtering and definition-time conditional checks. + Early returns if marked fake (machine incompatible or conditional failed). + + 2. set_directories(): + Base: Sets locdir, remdir, resdir based on path + Override: Generates attempt-specific identifier/path for resubmissions before + calling super().set_directories() + + 3. save_attempt(): + Base: Saves attempt information (if implemented in base) + Override: Calls super() then stores enhanced metadata including timestamp, + attempt_number, iteration_count, error messages in attempt_history + + 4. resubmit(): + Base: Resubmits simulation (if implemented in base) + Override: Saves current attempt, increments attempt_number, generates new + path/identifier, resets indicators, updates directories/files. + Preserves dependencies automatically (stored by simid). + + 5. check_status(): + Base: Checks simulation status and sets failed flag + Override: Calls super() then invokes error handlers if failed. Can + trigger automatic resubmission with error handling by resubmit(). + + 6. progress(): + Base: Progresses simulation through stages (setup, submit, analyze, etc.) + and triggers dependent simulations when finished + Override: Checks execution-time conditionals first (skips if false), manages + loop conditions/limits, then calls super().progress(). After parent + call, triggers next_iteration() if loop enabled and finished. + + New Methods (Not in Base): + - is_machine_compatible(): Checks if simulation can run on current machine + - evaluate_condition(): Evaluates conditional function (returns bool) + - skip(): Marks simulation as skipped (sets skipped=True, block=True) + - generate_resubmission_path(): Generates new path/identifier for resubmission + - next_iteration(): Prepares simulation for next loop iteration (increments count, resets indicators) + - depends(): Override with strategy and selector parameters for merge strategies ('all', 'first') + + New Attributes (Not in Base): + - attempt_number: Tracks resubmission attempts (0 = first run) + - base_identifier/base_path: Original identifier/path before resubmissions + - iteration_count: Current loop iteration (0 = first) + - error_handlers: List of error handler objects/functions + - attempt_history: List of attempt metadata dictionaries + - condition: Conditional function (Callable[[Simulation], bool]) + - conditional_enabled: Whether conditional execution is active + - skipped: Whether simulation was skipped + - required_machine/allowed_machines: Machine filtering constraints + - loop_enabled: Whether loop support is active + - loop_condition: Function that determines if loop should continue + - loop_max_iterations: Maximum loop iterations allowed + - loop_modifier: Function called before each iteration to modify input/state + - loop_variables: Dict for storing loop state + - merge_strategy: Merge strategy for multiple dependencies ('all', 'first') + - merge_selector: Selector function for merge strategy + - _completed_dependencies: Set tracking completed dependencies for merge strategies + + Extension Strategy: + - Pre-processing: Extracts enhanced kwargs, validates machine compatibility, evaluates + definition-time conditionals before calling super().__init__() + - Post-processing: Adds error handler invocation after check_status(), loop iteration + handling after progress() + - Side-by-side: Adds new functionality (machine filtering, loops, conditionals) + without modifying base behavior + - Selective override: Only overrides methods needed for enhanced features + + Backward Compatibility: + - DAG simulations unchanged: If no enhanced features used (no error_handlers, condition, + required_machine, loop_enabled), behavior matches base Simulation class + - Base class methods called via super() to preserve original functionality + - New features opt-in: Enhanced features only active when explicitly configured + - make_enhanced() function: Converts any Simulation to EnhancedSimulation dynamically, + preserving all original attributes and behavior + """ + + # Allowed inputs extended from Simulation + enhanced_allowed_inputs = set([ + 'error_handlers', 'condition', 'conditional_execution_time', + 'required_machine', 'allowed_machines', 'loop_enabled', + 'loop_condition', 'loop_max_iterations', 'loop_modifier', + 'merge_strategy', 'merge_selector' + ]) + + def __init__(self, **kwargs): + # Extract enhanced-specific kwargs + enhanced_kwargs = {} + for key in list(kwargs.keys()): + if key in self.enhanced_allowed_inputs: + enhanced_kwargs[key] = kwargs.pop(key) + + # Initialize enhanced state + self.attempt_number: int = 0 + self.base_identifier: Optional[str] = None + self.base_path: Optional[str] = None + self.iteration_count: int = 0 + self.error_handlers: List[ErrorHandler] = list(enhanced_kwargs.get('error_handlers', [])) + self.attempt_history: List[Dict[str, Any]] = [] + self.condition: Optional[Callable[[Any], bool]] = enhanced_kwargs.get('condition', None) + self.conditional_enabled: bool = self.condition is not None + # Three-state system: True=evaluated and passed, False=evaluated and failed, None=NotSet (keep checking) + # All conditionals are execution-time (dynamic) - this tracks evaluation state, not timing + self.conditional_execution_time: Optional[bool] = enhanced_kwargs.get('conditional_execution_time', None) + self.skipped: bool = False + self.required_machine: Optional[str] = enhanced_kwargs.get('required_machine', None) + self.allowed_machines: Set[str] = set(enhanced_kwargs.get('allowed_machines', [])) + self.loop_enabled: bool = enhanced_kwargs.get('loop_enabled', False) + self.loop_condition: Optional[Callable[[Any], bool]] = enhanced_kwargs.get('loop_condition', None) + self.loop_max_iterations: Optional[int] = enhanced_kwargs.get('loop_max_iterations', None) + self.loop_modifier: Optional[Callable[[Any], None]] = enhanced_kwargs.get('loop_modifier', None) + self.loop_variables: Dict[str, Any] = {} + self.merge_strategy: str = enhanced_kwargs.get('merge_strategy', 'all') # 'all', 'first' + self.merge_selector: Optional[Callable[[List[Simulation]], Simulation]] = enhanced_kwargs.get('merge_selector', None) + self._completed_dependencies: Set[int] = set() # Track completed dependencies for merge strategies + self._selected_dependency_id: Optional[int] = None # Store selected dependency ID when merge_selector is used + + # Check for decorator-based configuration + cls = self.__class__ + if hasattr(cls, '_required_machine'): + self.required_machine = cls._required_machine + if hasattr(cls, '_allowed_machines'): + self.allowed_machines = cls._allowed_machines + if hasattr(cls, '_loop_enabled'): + self.loop_enabled = cls._loop_enabled + if hasattr(cls, '_loop_condition'): + self.loop_condition = cls._loop_condition + if hasattr(cls, '_loop_max_iterations'): + self.loop_max_iterations = cls._loop_max_iterations + if hasattr(cls, '_loop_modifier'): + self.loop_modifier = cls._loop_modifier + if hasattr(cls, '_conditional_func'): + self.condition = cls._conditional_func + self.conditional_enabled = True + + # Call parent __init__ first to initialize base attributes + super().__init__(**kwargs) + + # Machine filtering at definition time (after parent init so job is available) + if self.required_machine or self.allowed_machines: + if not self.is_machine_compatible(): + self.fake_sim = True # Mark as fake to exclude from workflow + return + + # Definition-time conditional check + # Store base identifier and path for resubmissions + if self.base_identifier is None: + self.base_identifier = self.identifier + if self.base_path is None: + self.base_path = self.path + + def is_machine_compatible(self) -> bool: + """Check if simulation is compatible with current machine.""" + # If required_machine or allowed_machines is set, we must check compatibility + required = getattr(self, 'required_machine', None) + allowed = getattr(self, 'allowed_machines', None) + + # If no machine restrictions, always compatible + if not required and (not allowed or len(allowed) == 0): + return True + + # Get current machine name from multiple sources + machine_name = None + + # First try: job instance machine + if hasattr(self, 'job') and self.job is not None: + machine_name = getattr(self.job, 'machine', None) + + # Second try: Job class variable + if machine_name is None: + from machines import Job + machine_name = Job.machine + + # Third try: ProjectManager machine + if machine_name is None: + from project_manager import ProjectManager + if ProjectManager.machine is not None: + machine_name = getattr(ProjectManager.machine, 'name', None) + + # If still no machine, can't verify - return False to be safe + if machine_name is None: + return False + + # Normalize machine names for comparison (strip whitespace, lowercase) + machine_name = str(machine_name).strip().lower() if machine_name else None + required = str(required).strip().lower() if required else None + allowed = {str(a).strip().lower() for a in allowed} if allowed else set() + + # Check required_machine if set + if required: + compatible = machine_name == required + if not compatible: + # Explicitly mark as fake if incompatible + self.fake_sim = True + return compatible + + # Check allowed_machines if set + if allowed and len(allowed) > 0: + compatible = machine_name in allowed + if not compatible: + # Explicitly mark as fake if incompatible + self.fake_sim = True + return compatible + + return True + + def evaluate_condition(self) -> bool: + """Evaluate conditional function.""" + if self.condition is None: + return True + try: + return self.condition(self) + except Exception as e: + self.warn(f'Error evaluating condition: {e}') + return False + + def get_conditional_info(self) -> str: + """ + Get a human-readable description of the conditional configuration. + + Returns: + String describing the conditional status and type + """ + if not self.conditional_enabled: + return "" + + # Handle edge case: conditional_enabled but condition is None + if self.condition is None: + return "cond:None[invalid]" + + # Three-state system: True=evaluated and passed, False=evaluated and failed, None=NotSet (not yet evaluated) + # All conditionals are execution-time (dynamic) - this shows evaluation state + if self.conditional_execution_time is True: + state = "passed" + elif self.conditional_execution_time is False: + state = "failed" + else: + state = "pending" # NotSet - not yet evaluated + + # Get conditional description (with error handling) + try: + from conditionals import get_conditional_description + cond_desc = get_conditional_description(self.condition) + + # Truncate very long descriptions to keep status line readable + max_desc_len = 40 + if len(cond_desc) > max_desc_len: + cond_desc = cond_desc[:max_desc_len-3] + "..." + except Exception: + cond_desc = "custom" + + # Check if skipped + if self.skipped: + return f"cond:{cond_desc}[dynamic,{state},skipped]" + + # Return conditional info with evaluation state + return f"cond:{cond_desc}[dynamic,{state}]" + + def skip(self): + """Mark simulation as skipped.""" + self.skipped = True + self.block = True + + def set_directories(self): + """Override to support resubmissions with new identifiers.""" + # Generate identifier with attempt number if resubmission + if self.attempt_number > 0: + self.identifier = f"{self.base_identifier}_attempt{self.attempt_number}" + if self.base_path: + self.path = f"{self.base_path}/attempt_{self.attempt_number}" + + # Call parent method + super().set_directories() + + def generate_resubmission_path(self) -> Tuple[str, str]: + """ + Generate new path and identifier for resubmission. + + Returns: + Tuple of (new_path, new_identifier) + """ + new_attempt = self.attempt_number + 1 + new_identifier = f"{self.base_identifier}_attempt{new_attempt}" + new_path = f"{self.base_path}/attempt_{new_attempt}" if self.base_path else f"attempt_{new_attempt}" + return new_path, new_identifier + + def save_attempt(self): + """Enhanced save_attempt that stores more metadata.""" + # Call parent method + super().save_attempt() + + # Store additional metadata + attempt_data = { + 'timestamp': datetime.now().isoformat(), + 'attempt_number': self.attempt_number, + 'identifier': self.identifier, + 'path': self.path, + 'failed': self.failed, + 'finished': self.finished, + 'iteration_count': self.iteration_count, + } + + # Try to capture error information + if self.failed: + attempt_data['error_type'] = 'simulation_failed' + if hasattr(self, 'errfile') and self.errfile: + errfile_path = os.path.join(self.locdir, self.errfile) + if os.path.exists(errfile_path): + try: + with open(errfile_path, 'r', encoding='utf-8', errors='replace') as f: + attempt_data['error_message'] = f.read()[:1000] # Limit size + except: + pass + + self.attempt_history.append(attempt_data) + + def resubmit(self): + """ + Resubmit simulation after error. + + Increments attempt_number, generates new identifier/path, + resets indicators, and preserves dependencies. + """ + # Save current attempt + self.save_attempt() + + # Increment attempt number + self.attempt_number += 1 + + # Generate new path and identifier + new_path, new_identifier = self.generate_resubmission_path() + self.path = new_path + self.identifier = new_identifier + + # Reset indicators + self.reset_indicators() + + # Update directories and files + self.set_directories() + self.set_files() + + # Dependencies are preserved automatically (they're stored by simid) + + def check_status(self): + """Override to call error handlers on failure.""" + # Call parent method + super().check_status() + + # If failed, call error handlers + if self.failed and self.error_handlers: + error_info = { + 'error_type': 'simulation_failed', + 'attempt_number': self.attempt_number, + 'iteration_count': self.iteration_count, + 'identifier': self.identifier, + } + + should_resubmit = False + handler_results = [] # Track which handlers were called and their results + + for handler in self.error_handlers: + try: + # Check compatibility if handler supports it + if hasattr(handler, 'is_compatible'): + if not handler.is_compatible(self): + handler_name = getattr(handler, '__class__', type(handler)).__name__ + self.warn(f'Error handler {handler_name} is not compatible with {self.__class__.__name__}, skipping') + handler_results.append((handler_name, 'skipped', 'incompatible')) + continue + + # Call the handler + if hasattr(handler, 'handle_error'): + result = handler.handle_error(self, error_info) + elif callable(handler): + result = handler(self, error_info) + else: + handler_results.append((str(handler), 'skipped', 'not_callable')) + continue + + handler_name = getattr(handler, '__class__', type(handler)).__name__ + handler_results.append((handler_name, 'retry' if result else 'no_retry', result)) + + # First handler that recommends retry triggers resubmission + # This implements a "first handler wins" strategy + # Handlers are tried in order, and the first one that says "retry" wins + if result: + should_resubmit = True + self.log(f'Error handler {handler_name} recommended retry, resubmitting', n=2) + break + except Exception as e: + handler_name = getattr(handler, '__class__', type(handler)).__name__ + self.warn(f'Error in error handler {handler_name}: {e}') + handler_results.append((handler_name, 'error', str(e))) + + # Log handler results for debugging + if handler_results: + results_str = ', '.join([f'{name}:{status}' for name, status, _ in handler_results]) + self.log(f'Error handler results: {results_str}', n=3) + + if should_resubmit: + self.resubmit() + elif handler_results: + # All handlers were tried but none recommended retry + self.log('All error handlers processed, none recommended retry', n=2) + + def progress(self, dependency_id=None): + """Override to handle conditionals, loops, and merge strategies.""" + # Check execution-time conditional + # Three-state system: True/False = already evaluated (use cached), None = NotSet (evaluate now) + if self.conditional_enabled and not self.skipped: + # If NotSet (None), evaluate the conditional and cache the result + if self.conditional_execution_time is None: + result = self.evaluate_condition() + self.conditional_execution_time = result + else: + # Use cached result (True or False) + result = self.conditional_execution_time + + # If conditional failed, skip the simulation + if not result: + self.skip() + return + + # Handle merge strategies for multiple dependencies + # This must happen BEFORE calling super().progress() because + # super() checks if wait_ids is empty + handled_dependency = False # Track if we've already handled dependency_id removal + if dependency_id is not None and len(self.dependency_ids) > 1: + if not hasattr(self, '_completed_dependencies'): + self._completed_dependencies = set() + self._completed_dependencies.add(dependency_id) + + if self.merge_strategy == 'first': + # First dependency completed - clear all waits and proceed + # Remove the completed one from wait_ids, then clear all + if dependency_id in self.wait_ids: + self.wait_ids.remove(dependency_id) + self.wait_ids.clear() # Clear all to proceed immediately + handled_dependency = True # We've already handled dependency_id removal + elif self.merge_strategy == 'all': + # For 'all' strategy, wait for all dependencies + # Don't remove dependency_id here - let super() handle it + # Check if all completed + if len(self._completed_dependencies) == len(self.dependency_ids): + # All completed - if selector is set, use it to choose which dependency to use + if self.merge_selector is not None: + # Get all completed simulations + completed_sims = [] + for dep_id in self._completed_dependencies: + if dep_id in self.dependencies: + dep_sim = self.dependencies[dep_id].sim + # Selector will filter for completed status + completed_sims.append(dep_sim) + + if completed_sims: + try: + # Selector should filter for completed and return one + selected = self.merge_selector(completed_sims) + # Store the selected dependency ID for use in get_dependencies() + if selected is not None: + # Find the dependency ID for the selected simulation + for dep_id, dep in self.dependencies.items(): + if dep.sim is selected: + self._selected_dependency_id = dep_id + break + # Use selected simulation's results when get_dependencies() is called + # Remove dependency_id before clearing (super() will try to remove it) + if dependency_id in self.wait_ids: + self.wait_ids.remove(dependency_id) + self.wait_ids.clear() + handled_dependency = True # We've cleared wait_ids, don't pass dependency_id to super() + except Exception as e: + self.warn(f'Error in merge_selector: {e}') + # Remove dependency_id before clearing (super() will try to remove it) + if dependency_id in self.wait_ids: + self.wait_ids.remove(dependency_id) + self.wait_ids.clear() + handled_dependency = True # We've cleared wait_ids, don't pass dependency_id to super() + else: + # No selector - standard AND behavior, all dependencies are used + # Let super() handle dependency_id removal normally + pass + + # Handle loops + if self.loop_enabled: + if self.loop_condition is not None: + if not self.loop_condition(self): + # Loop condition false, exit loop + self.loop_enabled = False + elif self.loop_max_iterations is not None: + if self.iteration_count >= self.loop_max_iterations: + self.loop_enabled = False + + # Call parent progress (handles standard AND logic and dependency removal) + # If we've already handled dependency_id removal (e.g., cleared wait_ids), pass None + super().progress(None if handled_dependency else dependency_id) + + # Handle loop iteration + if self.loop_enabled and self.finished and not self.failed: + if self.loop_condition is None or self.loop_condition(self): + self.next_iteration() + + def next_iteration(self): + """ + Prepare simulation for next loop iteration. + + This method: + 1. Increments iteration_count + 2. Backs up input and output files from previous iteration (with iteration_count - 1) + 3. Resets indicators for next iteration + 4. Calls loop_modifier (if provided) to modify input/state + """ + self.iteration_count += 1 + prev_iteration = self.iteration_count - 1 + + # Backup input file from previous iteration + if self.infile is not None: + infile_path = os.path.join(self.locdir, self.infile) + if os.path.exists(infile_path): + # Create backup filename with iteration number + infile_base, infile_ext = os.path.splitext(self.infile) + if infile_ext == '': + backup_name = f'{self.infile}.iter{prev_iteration}' + else: + backup_name = f'{infile_base}.iter{prev_iteration}{infile_ext}' + backup_path = os.path.join(self.locdir, backup_name) + try: + shutil.copy2(infile_path, backup_path) + if hasattr(self, 'log'): + self.log(f'Backed up input file to {backup_name} (iteration {prev_iteration})', n=2) + except Exception as e: + if hasattr(self, 'warn'): + self.warn(f'Failed to backup input file: {e}') + + # Backup output files from previous iteration + for file_attr in ['outfile', 'errfile']: + if hasattr(self, file_attr) and getattr(self, file_attr) is not None: + file_name = getattr(self, file_attr) + file_path = os.path.join(self.locdir, file_name) + if os.path.exists(file_path): + # Create backup filename with iteration number + file_base, file_ext = os.path.splitext(file_name) + if file_ext == '': + backup_name = f'{file_name}.iter{prev_iteration}' + else: + backup_name = f'{file_base}.iter{prev_iteration}{file_ext}' + backup_path = os.path.join(self.locdir, backup_name) + try: + shutil.copy2(file_path, backup_path) + if hasattr(self, 'log'): + self.log(f'Backed up {file_attr} to {backup_name} (iteration {prev_iteration})', n=2) + except Exception as e: + if hasattr(self, 'warn'): + self.warn(f'Failed to backup {file_attr}: {e}') + + # Reset indicators for next iteration + self.reset_indicators() + + # Call loop_modifier if provided (allows modifying input at each iteration) + if self.loop_modifier is not None: + try: + self.loop_modifier(self) + except Exception as e: + self.warn(f'Error in loop_modifier: {e}') + + # Update loop variables if needed + # (subclasses can override to modify state) + + def reset_indicators(self): + """ + Override to preserve got_dependencies for loop iterations after the first one. + + For loop iterations after the first one, dependencies are already incorporated + and don't need to be re-incorporated. + """ + # Preserve got_dependencies for loop iterations after the first one + preserve_got_deps = self.loop_enabled and self.iteration_count > 0 + if preserve_got_deps: + got_deps = self.got_dependencies + + # Call parent method + super().reset_indicators() + + # Restore got_dependencies if we're in a loop iteration after the first + if preserve_got_deps: + self.got_dependencies = got_deps + + def incorporate_result(self, result_name, result, sim): + """ + Override to skip dependency incorporation for loop iterations after the first one. + + For loop iterations after the first one, dependencies are already in the same + directory and don't need to be re-incorporated. + """ + # Skip incorporation for loop iterations after the first one + if self.loop_enabled and self.iteration_count > 0: + # Dependencies are already in the same directory from the first iteration + # No need to incorporate again + if hasattr(self, 'log'): + self.log(f'Skipping dependency incorporation for loop iteration {self.iteration_count} (files already in directory)', n=2) + return + + # Call parent method for first iteration or non-loop simulations + super().incorporate_result(result_name, result, sim) + + def depends(self, *dependencies, strategy = 'all', selector=None): + """ + Override to add dependencies with merge strategy support. + + Args: + *dependencies: Dependency tuples (sim, result_name, ...) + strategy: Merge strategy for multiple dependencies: + - 'all' (default): Wait for all dependencies (AND logic) + - 'first': Run when first dependency completes (race condition) + selector: + This function should take a list of completed Simulation objects and return the selected one. + Example: + selector=lambda sims: min(sims, key=lambda s: s.energy) + """ + # Store selector for merge strategy + self.merge_strategy = strategy + if selector is not None: + self.merge_selector = selector + else: + # Default selector: return first completed simulation + # This ensures we only work with simulations that have finished + def completed_sims_selector(sims): + """Default selector: returns first completed simulation.""" + if not sims: + return None + # Filter to only completed simulations + completed_sims = [s for s in sims if hasattr(s, 'finished')] + # If none completed yet, return first one (will be checked later) + return completed_sims + self.merge_selector = completed_sims_selector + + # Call parent method to add dependencies + super().depends(*dependencies) + + def get_dependencies(self): + """ + Override to support merge_selector by filtering dependencies. + + When a merge_selector has selected a specific dependency, + only that dependency's results are used. + + For loop iterations after the first one (iteration_count > 0), + skips dependency incorporation since files are already in the same directory. + """ + # For loop iterations after the first one, skip dependency incorporation + # Dependencies are already in the same directory from the first iteration + if self.loop_enabled and self.iteration_count > 0: + # Still need to mark dependencies as satisfied, but skip incorporation + # Get results but don't incorporate them + if nexus_core.generate_only or self.finished: + for dep in self.dependencies: + for result_name in dep.result_names: + dep.results[result_name] = result_name + #end for + #end for + else: + for dep in self.dependencies: + sim = dep.sim + for result_name in dep.result_names: + if result_name!='other': + if sim.has_generic_input(): + self.error('a simulation result cannot be inferred from generic formatted or template input\nplease use {0} instead of {1}\nsim id: {2}\ndirectory: {3}\nresult: {4}'.format(self.input_type.__class__.__name__,sim.input.__class__.__name__,sim.id,sim.locdir,result_name)) + #end if + dep.results[result_name] = sim.get_result(result_name,sim) + else: + dep.results['other'] = obj() + #end if + #end for + #end for + # Skip incorporation for loop iterations after the first + # Files are already in the same directory + if not self.got_dependencies: + # Mark as got_dependencies without incorporating + pass + #end if + #end if + if self.renew_app_command: + self.job.renew_app_command(self) + #end if + self.got_dependencies = True + return + #end if + + # If a dependency was selected via merge_selector, filter to only that one + if self._selected_dependency_id is not None and self._selected_dependency_id in self.dependencies: + # Temporarily store original dependencies + original_dependencies = self.dependencies + # Create filtered dependencies dict with only the selected one + filtered_deps = obj() + filtered_deps[self._selected_dependency_id] = self.dependencies[self._selected_dependency_id] + # Temporarily replace dependencies + self.dependencies = filtered_deps + try: + # Call parent method with filtered dependencies + super().get_dependencies() + finally: + # Restore original dependencies + self.dependencies = original_dependencies + else: + # No selector or no selection made - use all dependencies (standard behavior) + super().get_dependencies() + + +def make_enhanced(sim: Simulation, **enhanced_kwargs) -> EnhancedSimulation: + """ + Convert any Simulation instance to an EnhancedSimulation. + + This is a fully generic function that works with any Simulation subclass + (Pwscf, Qmcpack, Vasp, etc.) without needing type-specific implementations. + + Args: + sim: Any Simulation instance to convert + **enhanced_kwargs: Enhanced workflow arguments: + - error_handlers: List of error handlers + - condition: Conditional function + - required_machine: Required machine name + - allowed_machines: Set of allowed machine names + - loop_enabled: Enable loop support + - loop_condition: Loop condition function + - loop_max_iterations: Maximum loop iterations + - loop_modifier: Function to modify input/state before each iteration + + Returns: + A new instance that inherits from both EnhancedSimulation and sim's class + + Example: + scf = generate_pwscf(...) + scf_enhanced = make_enhanced(scf, error_handlers=[...], loop_max_iterations=5) + """ + # Check if already an EnhancedSimulation + if isinstance(sim, EnhancedSimulation): + # Just update enhanced kwargs + for key, value in enhanced_kwargs.items(): + if hasattr(sim, key): + setattr(sim, key, value) + return sim + + # Create dynamic class that inherits from both EnhancedSimulation and original class + class_name = f"Enhanced{sim.__class__.__name__}" + + # Check if class already exists in module (avoid creating duplicates) + import sys + module = sys.modules[__name__] + if hasattr(module, class_name): + EnhancedClass = getattr(module, class_name) + else: + EnhancedClass = type(class_name, (EnhancedSimulation, sim.__class__), {}) + # Store in module for reuse + setattr(module, class_name, EnhancedClass) + + # Create new instance without calling __init__ (to avoid re-initialization) + enhanced_sim = EnhancedClass.__new__(EnhancedClass) + + # Copy all attributes from original simulation + # Save fake_sim status first (will be checked/updated after) + enhanced_sim.__dict__.update(sim.__dict__) + + # Initialize ALL enhanced attributes with defaults + # This ensures all attributes exist even if not in enhanced_kwargs + # Note: fake_sim is preserved from original (will be updated if machine incompatible) + enhanced_sim.attempt_number = 0 + enhanced_sim.base_identifier = enhanced_sim.identifier + enhanced_sim.base_path = enhanced_sim.path + enhanced_sim.iteration_count = 0 + enhanced_sim.error_handlers = [] + enhanced_sim.attempt_history = [] + enhanced_sim.skipped = False + enhanced_sim.loop_variables = {} + enhanced_sim.condition = None + enhanced_sim.conditional_enabled = False + enhanced_sim.conditional_execution_time = None # NotSet - not yet evaluated + enhanced_sim.required_machine = None + enhanced_sim.allowed_machines = set() + enhanced_sim.loop_enabled = False + enhanced_sim.loop_condition = None + enhanced_sim.loop_max_iterations = None + enhanced_sim.loop_modifier = None + enhanced_sim.merge_strategy = 'all' + enhanced_sim.merge_selector = None + enhanced_sim._completed_dependencies = set() + enhanced_sim._selected_dependency_id = None + + # Apply enhanced kwargs (overrides defaults) + enhanced_keys = [ + 'error_handlers', 'condition', 'conditional_execution_time', + 'required_machine', 'allowed_machines', 'loop_enabled', + 'loop_condition', 'loop_max_iterations', 'loop_modifier', + 'merge_strategy', 'merge_selector' + ] + + for key in enhanced_keys: + if key in enhanced_kwargs: + value = enhanced_kwargs[key] + setattr(enhanced_sim, key, value) + # Special handling for allowed_machines (convert to set if list) + if key == 'allowed_machines' and isinstance(value, (list, tuple)): + enhanced_sim.allowed_machines = set(value) + + # Set conditional_enabled based on condition (after applying kwargs) + if enhanced_sim.condition is not None: + enhanced_sim.conditional_enabled = True + + # Helper function to replace original in all_sims and update dependency references + def _replace_in_all_sims(): + if hasattr(Simulation, 'all_sims'): + # Find and replace the original simulation in all_sims + # Use identity comparison (is) to avoid recursion issues with __eq__ + for i, existing_sim in enumerate(Simulation.all_sims): + if existing_sim is sim: + Simulation.all_sims[i] = enhanced_sim + break + + # Also update any dependency references in other simulations + # This is critical: if scf3 depends on scf2_base, we need to update it to scf2 + for other_sim in Simulation.all_sims: + if other_sim is not enhanced_sim and hasattr(other_sim, 'dependencies'): + # Check dependencies dict + for dep_id, dep in list(other_sim.dependencies.items()): + if hasattr(dep, 'sim') and dep.sim is sim: + dep.sim = enhanced_sim + # Update the key if simid changed (though it shouldn't) + if dep_id != enhanced_sim.simid: + other_sim.dependencies[enhanced_sim.simid] = dep + del other_sim.dependencies[dep_id] + other_sim.dependency_ids.discard(dep_id) + other_sim.dependency_ids.add(enhanced_sim.simid) + other_sim.wait_ids.discard(dep_id) + other_sim.wait_ids.add(enhanced_sim.simid) + + # Check dependents dict + if hasattr(other_sim, 'dependents'): + for dep_id, dependent in list(other_sim.dependents.items()): + if dependent is sim: + other_sim.dependents[enhanced_sim.simid] = enhanced_sim + if dep_id != enhanced_sim.simid: + del other_sim.dependents[dep_id] + + # Update dependency_ids and wait_ids sets + if hasattr(other_sim, 'dependency_ids') and sim.simid in other_sim.dependency_ids: + other_sim.dependency_ids.discard(sim.simid) + other_sim.dependency_ids.add(enhanced_sim.simid) + if hasattr(other_sim, 'wait_ids') and sim.simid in other_sim.wait_ids: + other_sim.wait_ids.discard(sim.simid) + other_sim.wait_ids.add(enhanced_sim.simid) + + # Check machine compatibility IMMEDIATELY after setting required_machine/allowed_machines + # This is critical since we bypassed __init__ where this normally happens + if enhanced_sim.required_machine or enhanced_sim.allowed_machines: + # Get current machine for explicit comparison + machine_name = None + if hasattr(enhanced_sim, 'job') and enhanced_sim.job is not None: + machine_name = getattr(enhanced_sim.job, 'machine', None) + if machine_name is None: + from machines import Job + machine_name = Job.machine + if machine_name is None: + from project_manager import ProjectManager + if ProjectManager.machine is not None: + machine_name = getattr(ProjectManager.machine, 'name', None) + + # Explicit check: if required_machine is set, verify it matches + if enhanced_sim.required_machine: + if machine_name is None: + # Can't determine machine - mark as fake to be safe + enhanced_sim.fake_sim = True + _replace_in_all_sims() + return enhanced_sim + elif str(machine_name).strip().lower() != str(enhanced_sim.required_machine).strip().lower(): + # Machine doesn't match - mark as fake + enhanced_sim.fake_sim = True + _replace_in_all_sims() + return enhanced_sim + + # Explicit check: if allowed_machines is set, verify machine is in list + if enhanced_sim.allowed_machines and len(enhanced_sim.allowed_machines) > 0: + if machine_name is None: + # Can't determine machine - mark as fake to be safe + enhanced_sim.fake_sim = True + _replace_in_all_sims() + return enhanced_sim + else: + machine_normalized = str(machine_name).strip().lower() + allowed_normalized = {str(a).strip().lower() for a in enhanced_sim.allowed_machines} + if machine_normalized not in allowed_normalized: + # Machine not in allowed list - mark as fake + enhanced_sim.fake_sim = True + _replace_in_all_sims() + return enhanced_sim + + # Also call is_machine_compatible for consistency (it will also set fake_sim) + compatible = enhanced_sim.is_machine_compatible() + if not compatible: + enhanced_sim.fake_sim = True + _replace_in_all_sims() + return enhanced_sim + + # CRITICAL: Replace original simulation in Simulation.all_sims with enhanced version + # This ensures that when add_simulations() is called without arguments, + # it uses the enhanced version (which may be blocked) instead of the original + _replace_in_all_sims() + + return enhanced_sim + + +def create_branch(parent: Simulation, branches: List[Tuple[Simulation, Optional[Callable], str]], **enhanced_kwargs) -> List[EnhancedSimulation]: + """ + Create conditional branches from a parent simulation. + + This provides a natural way to create if/else or multi-way conditional branches + without needing separate make_enhanced() calls for each branch. + + Args: + parent: Parent simulation that branches depend on + branches: List of (sim_base, condition, result_name) tuples where: + - sim_base: Base simulation instance (before make_enhanced) + - condition: Optional conditional function (Callable[[Simulation], bool]) + If None, branch always runs (unconditional) + - result_name: Result name from parent (e.g., 'charge_density') + **enhanced_kwargs: Additional enhanced kwargs applied to all branches + (error_handlers, required_machine, etc.) + + Returns: + List of EnhancedSimulation instances (one for each branch) + + Example: + # Create if/else branch + scf3, scf4 = create_branch( + parent=scf2, + branches=[ + (scf3_base, energy_below_threshold, 'charge_density'), + (scf4_base, energy_above_threshold, 'charge_density'), + ], + error_handlers=[RetryErrorHandler(max_retries=2)] + ) + + # Create multi-way branch + scf_a, scf_b, scf_c = create_branch( + parent=scf1, + branches=[ + (scf_a_base, condition_a, 'charge_density'), + (scf_b_base, condition_b, 'charge_density'), + (scf_c_base, None, 'charge_density'), # Always runs + ] + ) + """ + enhanced_branches = [] + + for i, branch_spec in enumerate(branches): + if len(branch_spec) < 3: + raise ValueError(f'Branch {i} must be (sim_base, condition, result_name) tuple') + + sim_base, condition, result_name = branch_spec[0], branch_spec[1], branch_spec[2] + + # Ensure sim_base depends on parent + if not hasattr(sim_base, 'dependencies') or parent.simid not in sim_base.dependency_ids: + sim_base.depends((parent, result_name)) + + # Create enhanced simulation with condition + branch_kwargs = enhanced_kwargs.copy() + if condition is not None: + branch_kwargs['condition'] = condition + + enhanced_branch = make_enhanced(sim_base, **branch_kwargs) + enhanced_branches.append(enhanced_branch) + + return enhanced_branches + diff --git a/nexus/lib/error_handler.py b/nexus/lib/error_handler.py new file mode 100644 index 0000000000..cb970600c6 --- /dev/null +++ b/nexus/lib/error_handler.py @@ -0,0 +1,241 @@ +""" +Error handler infrastructure for DAG simulations. +""" + +from typing import Protocol, Callable, Dict, Any, Optional, List +from dataclasses import dataclass +import time +import os + +class ErrorHandler(Protocol): + """Protocol for error handlers.""" + + def diagnose_error(self, sim: Any) -> Optional[Dict[str, Any]]: + """ + Diagnose an error for a simulation. + + Args: + sim: Simulation instance that failed + """ + ... + + def handle_error(self, sim: Any, error_info: Dict[str, Any]) -> bool: + """ + Handle an error for a simulation. + + Args: + sim: Simulation instance that failed + error_info: Dictionary with error information (error_type, message, etc.) + + Returns: + True if simulation should be resubmitted, False otherwise + """ + ... + + def is_compatible(self, sim: Any) -> bool: + """ + Check if this error handler is compatible with the given simulation. + + Args: + sim: Simulation instance to check compatibility with + + Returns: + True if compatible, False otherwise + """ + # Default implementation: compatible with all simulations + # Override in subclasses for simulation-specific handlers + return True + + +@dataclass +class RetryErrorHandler: + """Error handler that retries a simulation up to max_retries times.""" + + max_retries: int = 3 + delay: float = 0.0 # Delay in seconds before retry + + def handle_error(self, sim: Any, error_info: Dict[str, Any]) -> bool: + """Handle error by retrying if under max_retries.""" + attempt_count = getattr(sim, 'attempt_number', 0) + if attempt_count < self.max_retries: + if self.delay > 0: + time.sleep(self.delay) + return True # Resubmit + return False # Don't resubmit + + def is_compatible(self, sim: Any) -> bool: + """Generic retry handler is compatible with all simulations.""" + return True + + +@dataclass +class ModifyInputErrorHandler: + """Error handler that modifies simulation input before resubmission.""" + + modification_func: Callable[[Any, Dict[str, Any]], None] + + def handle_error(self, sim: Any, error_info: Dict[str, Any]) -> bool: + """Handle error by modifying input and resubmitting.""" + self.modification_func(sim, error_info) + return True # Resubmit after modification + + def is_compatible(self, sim: Any) -> bool: + """Generic input modifier is compatible with all simulations.""" + return True + + +@dataclass +class SkipDependentErrorHandler: + """Error handler that skips dependent simulations when this one fails.""" + + def handle_error(self, sim: Any, error_info: Dict[str, Any]) -> bool: + """Handle error by blocking dependents.""" + if hasattr(sim, 'block_dependents'): + sim.block_dependents(block_self=False) + return False # Don't resubmit + + def is_compatible(self, sim: Any) -> bool: + """Generic handler is compatible with all simulations.""" + return True + + +@dataclass +class ConditionalErrorHandler: + """Error handler that conditionally handles errors based on a function.""" + + condition_func: Callable[[Any, Dict[str, Any]], bool] + true_handler: ErrorHandler + false_handler: Optional[ErrorHandler] = None + + def handle_error(self, sim: Any, error_info: Dict[str, Any]) -> bool: + """Handle error conditionally.""" + if self.condition_func(sim, error_info): + return self.true_handler.handle_error(sim, error_info) + elif self.false_handler is not None: + return self.false_handler.handle_error(sim, error_info) + return False + + def is_compatible(self, sim: Any) -> bool: + """Compatible if both handlers are compatible.""" + compatible = self.true_handler.is_compatible(sim) if hasattr(self.true_handler, 'is_compatible') else True + if self.false_handler is not None and hasattr(self.false_handler, 'is_compatible'): + compatible = compatible and self.false_handler.is_compatible(sim) + return compatible + + +def read_file_safe(filepath: str, max_lines: Optional[int] = None) -> Optional[str]: + """ + Safely read a file, returning None if it doesn't exist or can't be read. + + Args: + filepath: Path to the file to read + max_lines: If provided, only return the last max_lines lines + + Returns: + File contents as string, or None if file doesn't exist or can't be read + """ + if not os.path.exists(filepath): + return None + try: + with open(filepath, 'r', encoding='utf-8', errors='replace') as f: + if max_lines is not None: + lines = f.readlines() + return ''.join(lines[-max_lines:]) + else: + return f.read() + except Exception: + return None + + +def read_simulation_files(sim: Any, max_lines: Optional[int] = None) -> tuple[Optional[str], Optional[str]]: + """ + Read output and error files from a simulation. + + Args: + sim: Simulation instance + max_lines: If provided, only return the last max_lines lines from each file + + Returns: + Tuple of (output_content, error_content), either may be None + """ + output_content = None + error_content = None + + if hasattr(sim, 'outfile') and hasattr(sim, 'locdir'): + outfile_path = os.path.join(sim.locdir, sim.outfile) + output_content = read_file_safe(outfile_path, max_lines) + + if hasattr(sim, 'errfile') and hasattr(sim, 'locdir'): + errfile_path = os.path.join(sim.locdir, sim.errfile) + error_content = read_file_safe(errfile_path, max_lines) + + return output_content, error_content + + +class SimulationModifier(Protocol): + """Protocol for input modification strategies.""" + + def modify(self, sim: Any) -> None: + """ + Modify simulation input to address an error. + + Args: + sim: Simulation instance to modify + """ + ... + +@dataclass +class CombinedSimulationModifier: + """Combine multiple modification strategies.""" + + strategies: List[SimulationModifier] + + def modify(self, sim: Any) -> None: + """Apply all modification strategies in order.""" + for strategy in self.strategies: + if hasattr(strategy, 'modify'): + strategy.modify(sim) + + +@dataclass +class SimulationError: + """ + Base class for error strategies that combine diagnosis and handling. + + Each error strategy encapsulates: + - Pattern to detect in output/error files + - Whether it's recoverable + - How to handle it (modify input, retry, etc.) + """ + + pattern: str # Pattern to search for in output/error files + error_type: str # Name of error type + recoverable: bool = True + should_retry: bool = True + modification_strategy: Optional[SimulationModifier] = None + + def matches(self, output: str, error: str) -> bool: + """Check if this error strategy matches the output.""" + in_output = self.pattern in output if output else False + in_error = self.pattern in error if error else False + # Also check additional files if specified + return in_output or in_error + + def handle(self, sim: Any, error_info: Dict[str, Any]) -> bool: + """ + Handle the error by modifying input if needed and deciding whether to retry. + + Returns: + True if should retry, False otherwise + """ + # Apply modification strategy if provided + if self.modification_strategy is not None: + try: + self.modification_strategy.modify(sim) + except Exception as e: + if hasattr(sim, 'warn'): + sim.warn(f'Error in modification strategy: {e}') + + return self.should_retry + + diff --git a/nexus/lib/nexus.py b/nexus/lib/nexus.py index 20ce8e4c18..257b18c835 100644 --- a/nexus/lib/nexus.py +++ b/nexus/lib/nexus.py @@ -556,7 +556,20 @@ def run_project(*args,**kwargs): if nexus_core.graph_sims: graph_sims() #end if - pm = ProjectManager() + # Support for enhanced project manager + from simulation import Simulation + use_enhanced_pm = False + try: + from enhanced_simulation import EnhancedSimulation + except ImportError: + EnhancedSimulation = None + else: + use_enhanced_pm = any(isinstance(sim, EnhancedSimulation) for sim in Simulation.all_sims) + if use_enhanced_pm: + from enhanced_project_manager import EnhancedProjectManager + pm = EnhancedProjectManager() + else: + pm = ProjectManager() pm.add_simulations(*args,**kwargs) pm.run_project() return pm diff --git a/nexus/lib/pwscf_error_handler.py b/nexus/lib/pwscf_error_handler.py new file mode 100644 index 0000000000..ecf508a84f --- /dev/null +++ b/nexus/lib/pwscf_error_handler.py @@ -0,0 +1,251 @@ +""" +Pwscf-specific error handler for Quantum ESPRESSO pw.x simulations. + +Provides error strategies and modification strategies specifically designed +for Pwscf simulations, including automatic input parameter modification. +""" + +from typing import Dict, Any, Optional, List +from dataclasses import dataclass +import time + +from error_handler import ( + SimulationError, + SimulationModifier, + CombinedSimulationModifier, + read_simulation_files, +) + + +@dataclass +class PwscfError(SimulationError): + """ + Pwscf-specific error class that inherits from SimulationError. + + This class could be extended in the future with Pwscf-specific + error handling capabilities. + - QE-specific error metadata + - Pwscf input parameter extraction + - Advanced error diagnosis based on QE output structure + - Integration with Pwscf-specific recovery strategies + + """ + + +# Pre-defined QE error strategies factory +def _create_default_qe_strategies() -> Dict[str, PwscfError]: + """Create default QE error strategies using PwscfError.""" + return { + 'convergence_failure': PwscfError( + pattern='convergence NOT achieved', + error_type='convergence_failure', + recoverable=True, + should_retry=True, + ), + 'time_limit': PwscfError( + pattern='Maximum CPU time exceeded', + error_type='time_limit', + recoverable=True, + should_retry=True, + ), + 'user_stop': PwscfError( + pattern='Program stopped by user request', + error_type='user_stop', + recoverable=False, + should_retry=False, + ), + 'error_in_routine': PwscfError( + pattern='Error in routine', + error_type='error_in_routine', + recoverable=False, + should_retry=False, + ) + } + + +QE_ERROR_STRATEGIES = _create_default_qe_strategies() + +# Pwscf-specific input modification strategies +@dataclass +class RelaxConvergenceThreshold: + """Modification strategy: relax convergence threshold.""" + + factor: float = 10.0 + + def modify(self, sim: Any) -> None: + """Relax conv_thr by multiplying by factor.""" + if not hasattr(sim, 'input') or not hasattr(sim.input, 'electrons'): + return + electrons = sim.input.electrons + if hasattr(electrons, 'conv_thr'): + current = getattr(electrons, 'conv_thr', 1e-6) + new = current * self.factor + electrons.conv_thr = new + if hasattr(sim, 'log'): + sim.log(f'Relaxed conv_thr: {current:.2e} -> {new:.2e}', n=2) + + +@dataclass +class IncreaseMaxSteps: + """Modification strategy: increase maximum SCF steps.""" + + increase: int = 200 + + def modify(self, sim: Any) -> None: + """Increase electron_maxstep by increase amount.""" + if not hasattr(sim, 'input') or not hasattr(sim.input, 'electrons'): + return + electrons = sim.input.electrons + if hasattr(electrons, 'electron_maxstep'): + current = getattr(electrons, 'electron_maxstep', 100) + new = current + self.increase + electrons.electron_maxstep = new + if hasattr(sim, 'log'): + sim.log(f'Increased electron_maxstep: {current} -> {new}', n=2) + + +@dataclass +class AdjustMixingBeta: + """Modification strategy: adjust mixing beta for stability.""" + + reduction: float = 0.1 + min_beta: float = 0.1 + + def modify(self, sim: Any) -> None: + """Reduce mixing_beta by reduction amount (with minimum).""" + if not hasattr(sim, 'input') or not hasattr(sim.input, 'electrons'): + return + electrons = sim.input.electrons + if hasattr(electrons, 'mixing_beta'): + current = getattr(electrons, 'mixing_beta', 0.7) + new = max(self.min_beta, current - self.reduction) + electrons.mixing_beta = new + if hasattr(sim, 'log'): + sim.log(f'Reduced mixing_beta: {current:.3f} -> {new:.3f}', n=2) + + +@dataclass +class UseConjugateGradient: + """Modification strategy: use conjugate gradient for convergence.""" + + def modify(self, sim: Any) -> None: + """Use conjugate gradient for convergence.""" + if not hasattr(sim, 'input') or not hasattr(sim.input, 'electrons'): + return + electrons = sim.input.electrons + if hasattr(electrons, 'diagonalization'): + electrons.diagonalization = 'cg' + if hasattr(sim, 'log'): + sim.log('Using conjugate gradient for convergence', n=2) + + + +@dataclass +class PwscfErrorHandler: + """ + Error handler for Pwscf that modifies input parameters based on error type. + + Uses error strategies to diagnose and handle errors. Modification strategies + can be provided to customize how inputs are modified. + + Args: + max_retries: Maximum number of retry attempts + delay: Delay in seconds before retry + error_strategies: List of PwscfError objects (defaults to QE_ERROR_STRATEGIES) + modification_strategies: Dict mapping error_type to SimulationModifier + """ + + max_retries: int = 3 + delay: float = 0.0 + error_strategies: Optional[Dict[str, PwscfError]] = None + modification_strategies: Optional[Dict[str, SimulationModifier]] = None + + def __post_init__(self): + """Initialize with default QE error strategies if not provided.""" + if self.error_strategies is None: + self.error_strategies = _create_default_qe_strategies() + + # Apply default modification strategies for convergence failures + if self.modification_strategies is None: + self.modification_strategies = {} + + # Default convergence failure modifications + if 'convergence_failure' not in self.modification_strategies: + self.modification_strategies['convergence_failure'] = CombinedSimulationModifier([ + RelaxConvergenceThreshold(factor=10.0), + IncreaseMaxSteps(increase=200), + AdjustMixingBeta(reduction=0.1), + ]) + + def diagnose_error(self, sim: Any) -> Optional[PwscfError]: + """ + Diagnose error by checking all error strategies. + + Returns: + Matching PwscfError or None if no match + """ + # Read both output and error files using shared utility + output, error = read_simulation_files(sim) + + # If no files found, can't diagnose + if not output and not error: + return None + + # Check strategies in order (job_done first, then others) + for strategy in self.error_strategies.values(): + if strategy.matches(output or '', error or ''): + return strategy + + return None + + def handle_error(self, sim: Any, error_info: Dict[str, Any]) -> bool: + """ + Handle QE errors by modifying input and deciding whether to retry. + + Returns: + True if should retry, False otherwise + """ + attempt_count = getattr(sim, 'attempt_number', 0) + + # Check retry limit + if attempt_count >= self.max_retries: + return False + + # Diagnose error + error_strategy = self.diagnose_error(sim) + if error_strategy is None: + return False # Unknown error, don't retry + + # If job is done, no need to retry + if error_strategy.error_type == 'completed': + return False + + # Apply modification strategy if available + if error_strategy.error_type in self.modification_strategies: + mod_strategy = self.modification_strategies[error_strategy.error_type] + if hasattr(mod_strategy, 'modify'): + mod_strategy.modify(sim) + + # Also check if strategy has its own modification + if error_strategy.modification_strategy is not None: + error_strategy.modification_strategy.modify(sim) + + # Handle error using strategy + should_retry = error_strategy.handle(sim, error_info) + + if should_retry and self.delay > 0: + time.sleep(self.delay) + + return should_retry + + def is_compatible(self, sim: Any) -> bool: + """Check if this handler is compatible with the simulation.""" + if hasattr(sim, '__class__'): + class_name = sim.__class__.__name__ + if class_name == 'Pwscf' or (class_name.startswith('Enhanced') and 'Pwscf' in class_name): + return True + module_name = getattr(sim.__class__, '__module__', '') + if 'pwscf' in module_name.lower(): + return True + return False +