diff --git a/aidge_core/mem_info.py b/aidge_core/mem_info.py index 6938ca35d11dad18cadca91963957f974fd45f55..c7ca85bbd73bd205850b19616e53fda210749a80 100644 --- a/aidge_core/mem_info.py +++ b/aidge_core/mem_info.py @@ -1,4 +1,5 @@ import os +import subprocess import shutil from pathlib import Path import aidge_core @@ -41,6 +42,18 @@ def compute_default_mem_info(scheduler: aidge_core.Scheduler) -> Tuple[int, List return mem_size, mem_info +def _gnuplot_installed(): + try: + # Run gnuplot with the --version flag and capture the output + subprocess.run(["gnuplot", "--version"]) + return True + except FileNotFoundError: + aidge_core.Log.warn("Gnuplot is not installed.") + return False + except subprocess.CalledProcessError: + aidge_core.Log.warn("Gnuplot command found but failed to run.") + return False + 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. @@ -73,21 +86,20 @@ 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()] - # 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.") + + if stats_folder is not None: + if _gnuplot_installed(): + # Use gnuplot to generate the log + 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") + else: + 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()