diff --git a/aidge_core/mem_info.py b/aidge_core/mem_info.py index 1a800f9488905a4aed3a907deca8327979effbe4..6938ca35d11dad18cadca91963957f974fd45f55 100644 --- a/aidge_core/mem_info.py +++ b/aidge_core/mem_info.py @@ -41,7 +41,7 @@ def compute_default_mem_info(scheduler: aidge_core.Scheduler) -> Tuple[int, List return mem_size, mem_info -def generate_optimized_memory_info(scheduler: aidge_core.Scheduler, stats_folder: Path, wrapping: bool = False) -> Tuple[int, List[dict]]: +def generate_optimized_memory_info(scheduler: aidge_core.Scheduler, stats_folder: Path = None, wrapping: bool = False) -> Tuple[int, List[dict]]: """Generates optimized memory information for a computation graph managed by a scheduler. This function analyzes the memory usage of a computation graph, determining the memory peak @@ -52,8 +52,8 @@ def generate_optimized_memory_info(scheduler: aidge_core.Scheduler, stats_folder nodes and facilitates memory planning by invoking its `generate_memory` method. :type scheduler: aidge_core.Scheduler :param stats_folder: Directory path to store memory statistics and plots generated by `mem_manager`. - If provided as a string, it is converted to a `Path` object. - :type stats_folder: Path + If provided as a string, it is converted to a `Path` object, default=None. + :type stats_folder: Path, optional :param wrapping: Boolean flag to enable or disable wrap-around buffer optimization. Defaults to `False`. :type wrapping: bool, optional @@ -73,17 +73,21 @@ def generate_optimized_memory_info(scheduler: aidge_core.Scheduler, stats_folder # List of nodes which are connected at the input of the graph (None if input is not connected) nodes_at_input = [n[0] for n in scheduler.graph_view().inputs()] - # Use gnuplot to generate the log - if isinstance(stats_folder, str): - stats_folder = Path(stats_folder) - os.makedirs(str(Path(stats_folder) / "graph"), exist_ok=True) - mem_manager.log("memory_info") - os.chmod("memory_info_plot.gnu", 0o777) - os.system("./memory_info_plot.gnu") - shutil.move("memory_info", str(Path(stats_folder) / "graph" / "memory_info")) - shutil.move("memory_info_plot.png", str( - Path(stats_folder) / "graph" / "memory_info_plot.png")) - os.remove("memory_info_plot.gnu") + # Return 0 ==> gnuplot installed + if os.system("gnuplot --version > /dev/null 2>&1") == 0: + # Use gnuplot to generate the log + if isinstance(stats_folder, str): + stats_folder = Path(stats_folder) + os.makedirs(str(Path(stats_folder) / "graph"), exist_ok=True) + mem_manager.log("memory_info") + os.chmod("memory_info_plot.gnu", 0o777) + os.system("./memory_info_plot.gnu") + shutil.move("memory_info", str(Path(stats_folder) / "graph" / "memory_info")) + shutil.move("memory_info_plot.png", str( + Path(stats_folder) / "graph" / "memory_info_plot.png")) + os.remove("memory_info_plot.gnu") + elif stats_folder is not None: + aidge_core.Log.warn("Warning: gnuplot is not installed, could not generate stat folder.") # In the export, we currently use an unified memory buffer whose size # is determined by the memory peak usage mem_size = mem_manager.get_peak_usage()