diff --git a/src/scgenerator/cli/cli.py b/src/scgenerator/cli/cli.py index f803f8a..0ab6a47 100644 --- a/src/scgenerator/cli/cli.py +++ b/src/scgenerator/cli/cli.py @@ -52,6 +52,12 @@ def create_parser(): "data_dir", help="path to the directory where the initial_config.toml and the data is stored", ) + resume_parser.add_argument( + "configs", + nargs="*", + default=[], + help="list of subsequent config files (excluding the resumed one)", + ) resume_parser.set_defaults(func=resume_sim) merge_parser = subparsers.add_parser("merge", help="merge simulation results") @@ -99,9 +105,11 @@ def prep_ray(args): def resume_sim(args): + method = prep_ray(args) - sim = resume_simulations(args.data_dir, args.id, method=method) + sim = resume_simulations(args.data_dir, method=method) sim.run() + run_simulation_sequence(*args.configs, method=method, prev_data_folder=sim.data_folder) if __name__ == "__main__": diff --git a/src/scgenerator/data/submit_job_template.txt b/src/scgenerator/data/submit_job_template.txt index b872a5e..609b86f 100644 --- a/src/scgenerator/data/submit_job_template.txt +++ b/src/scgenerator/data/submit_job_template.txt @@ -50,5 +50,5 @@ done ############################################################################################## #### call your code below -scgenerator run {configs_list} +scgenerator {command} {configs_list} exit \ No newline at end of file diff --git a/src/scgenerator/initialize.py b/src/scgenerator/initialize.py index 59687d3..e43d786 100644 --- a/src/scgenerator/initialize.py +++ b/src/scgenerator/initialize.py @@ -175,7 +175,7 @@ def validate_config_sequence(*configs: os.PathLike) -> Dict[str, Any]: Parameters ---------- configs : os.PathLike - sequence of paths to toml config files + sequence of paths to toml config files. The first element may be a folder containing data intead Returns ------- @@ -184,6 +184,8 @@ def validate_config_sequence(*configs: os.PathLike) -> Dict[str, Any]: """ previous = None for config in configs: + if (p := Path(config)).is_dir(): + config = p / "initial_config.toml" dico = io.load_toml(config) previous = override_config(dico, previous) validate(previous) diff --git a/src/scgenerator/physics/simulate.py b/src/scgenerator/physics/simulate.py index c6f5284..a1402f7 100644 --- a/src/scgenerator/physics/simulate.py +++ b/src/scgenerator/physics/simulate.py @@ -710,8 +710,13 @@ class RaySimulations(Simulations, priority=2): self.p_bars.print() -def run_simulation_sequence(*config_files: os.PathLike, method=None, final_name: str = None): - prev = None +def run_simulation_sequence( + *config_files: os.PathLike, + method=None, + final_name: str = None, + prev_data_folder: os.PathLike = None, +): + prev = prev_data_folder for config_file in config_files: sim = new_simulation(config_file, prev, method) sim.run() @@ -738,10 +743,9 @@ def new_simulation( return _new_simulations(param_seq, task_id, method) -def resume_simulations( - data_folder: str, task_id: int = 0, method: Type[Simulations] = None -) -> Simulations: +def resume_simulations(data_folder: str, method: Type[Simulations] = None) -> Simulations: + task_id = np.random.randint(1e9, 1e12) config = io.load_toml(os.path.join(data_folder, "initial_config.toml")) io.set_data_folder(task_id, data_folder) param_seq = initialize.RecoveryParamSequence(config, task_id) diff --git a/src/scgenerator/scripts/slurm_submit.py b/src/scgenerator/scripts/slurm_submit.py index 372b18a..e1e774b 100644 --- a/src/scgenerator/scripts/slurm_submit.py +++ b/src/scgenerator/scripts/slurm_submit.py @@ -34,6 +34,9 @@ def create_parser(): default=f"source {os.path.expanduser('~/anaconda3/etc/profile.d/conda.sh')} && conda activate sc", help="commands to run to setup the environement (default : activate the sc environment with conda)", ) + parser.add_argument( + "--command", default="run", choices=["run", "resume"], help="command to run" + ) return parser @@ -47,6 +50,9 @@ def copy_starting_files(): def main(): + + command_map = dict(run="Propagate", resume="Resuming") + parser = create_parser() template = Paths.gets("submit_job_template") args = parser.parse_args() @@ -69,13 +75,13 @@ def main(): ) job_name = f"supercontinuum {final_config['name']}" submit_sh = template.format( - job_name=job_name, configs_list=" ".join(args.configs), **vars(args) + job_name=job_name, configs_list=" ".join(f'"{c}"' for c in args.configs), **vars(args) ) with open(file_name, "w") as file: file.write(submit_sh) subprocess.run(["sbatch", "--test-only", file_name]) submit = input( - f"Propagate {sim_num} pulses from configs {args.configs} with {args.cpus_per_node} cpus" + f"{command_map[args.command]} {sim_num} pulses from configs {args.configs} with {args.cpus_per_node} cpus" + f" per node on {args.nodes} nodes for {format_time(args.time)} ? (y/[n])\n" ) if submit.lower() in ["y", "yes"]: diff --git a/src/scgenerator/spectra.py b/src/scgenerator/spectra.py index ed4613e..a45596b 100644 --- a/src/scgenerator/spectra.py +++ b/src/scgenerator/spectra.py @@ -32,7 +32,7 @@ class Spectrum(np.ndarray): class Pulse(Sequence): def __init__(self, path: str, ensure_2d=True): self.logger = get_logger(__name__) - self.path = path + self.path = str(path) self.__ensure_2d = ensure_2d if not os.path.isdir(self.path):