diff --git a/low_level.py b/low_level.py index f807c31..988c145 100644 --- a/low_level.py +++ b/low_level.py @@ -182,4 +182,19 @@ def print_ending_module(module:str, verbose:int=1): Print the ending message of a module. """ if verbose >= 1: - print(print_level(0, module)+f"Ending {module} module.") \ No newline at end of file + print(print_level(0, module)+f"Ending {module} module.") + + +def wait_until_file_exists(filename:str, verbose:int=1, limit:int=24*3600): + """ + Wait until a file exists. + """ + from time import sleep + from os.path import isfile + + k=0 + while not isfile(filename): + if k%60 == -1: + print_message(f"Waiting for {filename} to exist. {k//60} minutes elapsed.", 3, "low level", verbose=verbose) + sleep(1) + print_message(f"File {filename} exists. {k//60} minutes elapsed.", 3, "low level", verbose=verbose) \ No newline at end of file diff --git a/scola.py b/scola.py new file mode 100644 index 0000000..af94f25 --- /dev/null +++ b/scola.py @@ -0,0 +1,123 @@ +def main_scola(parsed_args): + from args_main import parse_arguments_main + from low_level import print_starting_module, print_message, print_ending_module + from os.path import isfile + import subprocess + + print_starting_module("scola", verbose=parsed_args.verbose) + main_dict = parse_arguments_main(parsed_args) + param_file=parsed_args.paramfile + if param_file is None: + param_file = main_dict["paramdir"]+"parameters_"+main_dict["simname"]+".sbmy" + log_file = main_dict["logdir"]+main_dict["mode"]+".log" + + nboxes_tot = int(parsed_args.N_tiles**3) + + if parsed_args.execution == "local": + print_message("Running sCOLA in local mode.", 1, "scola", verbose=parsed_args.verbose) + + for b in range(1,nboxes_tot+1): + print_message(f"Running box {b}/{nboxes_tot}.", 2, "scola", verbose=parsed_args.verbose) + command_args = ["scola", param_file, log_file, "-b", str(b)] + + if parsed_args.verbose < 2: + from io import BytesIO + from low_level import stdout_redirector, stderr_redirector + + f = BytesIO() + g = BytesIO() + with stdout_redirector(f): + with stderr_redirector(g): + subprocess.run(command_args) + g.close() + f.close() + else: + subprocess.run(command_args) + + print_message(f"Box {b}/{nboxes_tot} finished.", 3, "scola", verbose=parsed_args.verbose) + + print_message("sCOLA finished.", 1, "scola", verbose=parsed_args.verbose) + + elif parsed_args.execution == "slurm": + from slurm_submission import create_slurm_script, parse_arguments_slurm + from args_main import parse_arguments_main + print_message("Running monofonic in slurm mode.", 1, "scola", verbose=parsed_args.verbose) + slurm_dict=parse_arguments_slurm(parsed_args) + main_dict=parse_arguments_main(parsed_args) + slurm_script = slurm_dict["scripts"]+"scola.sh" + + if not isfile(slurm_script): + print_message(f"SLURM script {slurm_script} does not exist. Creating it.", 2, "scola", verbose=parsed_args.verbose) + create_slurm_script( + slurm_template=slurm_dict["scola_template"], + slurm_script=slurm_script, + job="scola", + job_config_file=param_file, + job_log=log_file, + array=(1, nboxes_tot+1), + ) + print_message(f"SLURM script written to {slurm_script}.", 3, "scola", verbose=parsed_args.verbose) + else: + print_message(f"SLURM script {slurm_script} exists.", 2, "scola", verbose=parsed_args.verbose) + + print_message(f"Submitting monofonic job to SLURM.", 1, "scola", verbose=parsed_args.verbose) + + command_args = ["sbatch", slurm_script] + + if parsed_args.verbose < 2: + from io import BytesIO + from low_level import stdout_redirector, stderr_redirector + + f = BytesIO() + g = BytesIO() + with stdout_redirector(f): + with stderr_redirector(g): + subprocess.run(command_args) + g.close() + f.close() + else: + subprocess.run(command_args) + + print_message("Monofonic job submitted.", 2, "scola", verbose=parsed_args.verbose) + + else: + raise ValueError(f"Execution mode {parsed_args.execution} not recognized.") + + print_ending_module("scola", verbose=parsed_args.verbose) + + +def main_pre_post_scola(parsed_args): + from simbelmyne import main_simbelmyne + main_simbelmyne(parsed_args) + +if __name__ == "__main__": + from argparse import ArgumentParser + from args_main import register_arguments_main + from timestepping import register_arguments_timestepping, main_timestepping + from parameters_card import register_arguments_card, main_parameter_card + from cosmo_params import register_arguments_cosmo + from parameters_monofonic import register_arguments_monofonic + from slurm_submission import register_arguments_slurm + from low_level import print_starting_module, print_message, print_ending_module, wait_until_file_exists + + parser = ArgumentParser(description="Run sCOLA.") + # TODO: reduce the volume of arguments + register_arguments_main(parser) + register_arguments_timestepping(parser) + register_arguments_monofonic(parser) + register_arguments_slurm(parser) + register_arguments_card(parser) + register_arguments_cosmo(parser) + parsed_args = parser.parse_args() + + card_dict = main_parameter_card(parsed_args) + nboxes_tot = int(parsed_args.N_tiles**3) + main_timestepping(parsed_args) + main_pre_post_scola(parsed_args) + if parsed_args.execution == "slurm": + wait_until_file_exists(card_dict["OutputLPTPotential2"]) + main_scola(parsed_args) + if parsed_args.execution == "slurm": + for b in range(1,nboxes_tot+1): + wait_until_file_exists(f"{card_dict['OutputTilesBase']}{b}.h5") + main_pre_post_scola(parsed_args) \ No newline at end of file diff --git a/simbelmyne.py b/simbelmyne.py index e5363f3..236a2f8 100644 --- a/simbelmyne.py +++ b/simbelmyne.py @@ -36,7 +36,7 @@ def main_simbelmyne(parsed_args): elif parsed_args.execution == "slurm": from slurm_submission import create_slurm_script, parse_arguments_slurm from args_main import parse_arguments_main - print_message("Running monofonic in slurm mode.", 1, "simbelmyne", verbose=parsed_args.verbose) + print_message("Running simbelmyne in slurm mode.", 1, "simbelmyne", verbose=parsed_args.verbose) slurm_dict=parse_arguments_slurm(parsed_args) main_dict=parse_arguments_main(parsed_args) slurm_script = slurm_dict["scripts"]+"simbelmyne.sh" @@ -54,7 +54,7 @@ def main_simbelmyne(parsed_args): else: print_message(f"SLURM script {slurm_script} exists.", 2, "simbelmyne", verbose=parsed_args.verbose) - print_message(f"Submitting monofonic job to SLURM.", 1, "simbelmyne", verbose=parsed_args.verbose) + print_message(f"Submitting simbelmyne job to SLURM.", 1, "simbelmyne", verbose=parsed_args.verbose) command_args = ["sbatch", slurm_script] @@ -86,14 +86,14 @@ if __name__ == "__main__": from timestepping import register_arguments_timestepping, main_timestepping from parameters_card import register_arguments_card, main_parameter_card from cosmo_params import register_arguments_cosmo - from parameters_monofonic import register_arguments_monofonic + from parameters_simbelmyne import register_arguments_simbelmyne from slurm_submission import register_arguments_slurm parser = ArgumentParser(description="Run Simbelmyne.") # TODO: reduce the volume of arguments register_arguments_main(parser) register_arguments_timestepping(parser) - register_arguments_monofonic(parser) + register_arguments_simbelmyne(parser) register_arguments_slurm(parser) register_arguments_card(parser) register_arguments_cosmo(parser) diff --git a/slurm_submission.py b/slurm_submission.py index f61e4ab..442a2e2 100644 --- a/slurm_submission.py +++ b/slurm_submission.py @@ -133,8 +133,10 @@ def create_slurm_script(slurm_template:str, match job: case "monofonic": f.write(f"{path_to_monofonic_binary} {job_config_file} > {job_log}") - case "simbelmyne" | "scola": + case "simbelmyne": f.write(f"{job} {job_config_file} {job_log}") + case "scola": + f.write(f"{job} {job_config_file} {job_log} "+"${SLURM_ARRAY_TASK_ID}") case _: raise ValueError(f"Job type {job} not recognized.")