diff --git a/mphys/network/remote_component.py b/mphys/network/remote_component.py index bfb26bf1..ba27aa9d 100644 --- a/mphys/network/remote_component.py +++ b/mphys/network/remote_component.py @@ -1,7 +1,9 @@ import json import os +import re import time from functools import wraps +from glob import glob import numpy as np import openmdao.api as om @@ -70,6 +72,11 @@ def initialize(self): default=False, desc="dump a separate input/output json file for each evaluation", ) + self.options.declare( + "reuse_dumped_json", + default=False, + desc="try to reuse existing output json files instead of running the remote component", + ) self.options.declare( "var_naming_dot_replacement", default=":", @@ -136,6 +143,7 @@ def setup(self): ] self.dump_json = self.options["dump_json"] self.dump_separate_json = self.options["dump_separate_json"] + self.reuse_dumped_json = self.options["reuse_dumped_json"] self.additional_remote_inputs = self.options["additional_remote_inputs"] self.additional_remote_outputs = self.options["additional_remote_outputs"] self.additional_remote_constants = self.options[ @@ -144,20 +152,16 @@ def setup(self): self.last_analysis_completed_time = ( time.time() ) # for tracking down time between function/gradient calls + if self.reuse_dumped_json: + self.dump_separate_json = True if self.dump_separate_json: self.dump_json = True - - self._setup_server_manager() + self.server_manager = None # for tracking model times, and determining whether to relaunch servers self.times_function = np.array([]) self.times_gradient = np.array([]) - # get baseline model - print( - f"CLIENT (subsystem {self.name}): Running model from setup to get design problem info", - flush=True, - ) output_dict = self.evaluate_model( command="initialize", remote_input_dict={ @@ -214,7 +218,20 @@ def compute_partials(self, inputs, partials): self._assign_additional_partials_from_remote_output(remote_dict, partials) def evaluate_model(self, remote_input_dict=None, command="initialize"): - if self._need_to_restart_server(command): + + # first check if able to reuse dumped json file + remote_output_dict = self._reuse_dumped_json(remote_input_dict, command) + if remote_output_dict is not None: + return remote_output_dict + + if self.server_manager is None: + self._setup_server_manager() + if command == "initialize": + self._print_status_message( + "Running model from setup to get design problem info" + ) + + elif self._need_to_restart_server(command): self.server_manager.stop_server() self.server_manager.start_server() @@ -246,13 +263,16 @@ def evaluate_model(self, remote_input_dict=None, command="initialize"): and self._doing_derivative_evaluation(command) ): if self.comm.rank == 0: - print( - f"CLIENT (subsystem {self.name}): Stopping server's HPC job for down time" + self._print_status_message( + "Stopping server's HPC job for down time" ) self.server_manager.stop_server() return remote_output_dict + def _print_status_message(self, message): + print(f"CLIENT (subsystem {self.name}): {message}", flush=True) + def _assign_objective_partials_from_remote_output(self, remote_dict, partials): for obj in remote_dict["objective"].keys(): for dv in remote_dict["design_vars"].keys(): @@ -374,6 +394,110 @@ def _need_to_restart_server(self, command: str): ) return not self.server_manager.enough_time_is_remaining(estimated_model_time) + def _reuse_dumped_json(self, remote_input_dict, command): + def extract_number(filepath): + # for sorting filenames of json files + match = re.search(r"(\d+)\.json$", filepath) + return int(match.group(1)) + + save_dir = "remote_json_files" + if not self.reuse_dumped_json or not os.path.isdir(save_dir): + return None + + if command == "initialize": + + # assume *_function0.json contains info needed for design problem setup + filename = f"{save_dir}/{self.name}_outputs_function0.json" + if not os.path.isfile(filename): + return None + else: + with open(filename, "r") as file: + remote_output_dict = json.load(file) + model_time_elapsed = remote_output_dict["wall_time"] + if self._doing_derivative_evaluation(command): + self.times_gradient = np.hstack( + [self.times_gradient, model_time_elapsed] + ) + else: + self.times_function = np.hstack( + [self.times_function, model_time_elapsed] + ) + if self.comm.rank == 0: + self._print_status_message( + f"Obtained design problem info from dumped json file '{filename}'" + ) + return remote_output_dict + + else: + + # possible filenames to read through + filenames = [] + if not self._doing_derivative_evaluation(command): + filenames += sorted( + glob(f"{save_dir}/{self.name}_outputs_function*.json"), + key=extract_number, + ) + filenames += sorted( + glob(f"{save_dir}/{self.name}_outputs_derivative*.json"), + key=extract_number, + ) + + # check each json file for design of interest + for filename in filenames: + with open(filename, "r") as file: + remote_output_dict = json.load(file) + if self._designs_match(remote_input_dict, remote_output_dict): + model_time_elapsed = remote_output_dict["wall_time"] + if self._doing_derivative_evaluation(command): + self._print_status_message( + f"Found design derivatives in dumped json file '{filename}'" + ) + self.times_gradient = np.hstack( + [self.times_gradient, model_time_elapsed] + ) + else: + self._print_status_message( + f"Found design responses in dumped json file '{filename}'" + ) + self.times_function = np.hstack( + [self.times_function, model_time_elapsed] + ) + return remote_output_dict + + return None + + def _designs_match(self, input_dict, output_dict): + if not self._check_for_consistent_inputs(input_dict, output_dict): + self._print_status_message( + "Inconsistent inputs and outputs found in dumped json file... skipping" + ) + return False + for input_type in ["design_vars", "additional_constants", "additional_inputs"]: + for input_name in input_dict[input_type].keys(): + # TODO: worth having a tolerance on this? + if not np.allclose( + input_dict[input_type][input_name]["val"], + output_dict[input_type][input_name]["val"], + ): + return False + return True + + def _check_for_consistent_inputs(self, input_dict, output_dict): + input_keys = ( + list(input_dict["design_vars"].keys()) + + list(input_dict["additional_constants"].keys()) + + list(input_dict["additional_inputs"].keys()) + ) + output_keys = ( + list(output_dict["design_vars"].keys()) + + list(output_dict["additional_constants"].keys()) + + list(output_dict["additional_inputs"].keys()) + ) + if set(input_keys) != set(output_keys): + return False + else: + return True + def _dump_json(self, remote_dict: dict, command: str): if "objective" in remote_dict.keys(): dict_type = "outputs" @@ -386,7 +510,9 @@ def _dump_json(self, remote_dict: dict, command: str): os.mkdir(save_dir) except Exception: pass # may have been created by now, by a parallel server - if self._doing_derivative_evaluation(command): + if self._doing_derivative_evaluation( + command + ): # TODO: change len(times) to something that won't overwrite existing json files? filename = f"{save_dir}/{self.name}_{dict_type}_derivative{len(self.times_gradient)}.json" else: filename = f"{save_dir}/{self.name}_{dict_type}_function{len(self.times_function)}.json"