diff --git a/src/ansys/pytwin/evaluate/twin_model.py b/src/ansys/pytwin/evaluate/twin_model.py index 119724cf..e57295c5 100644 --- a/src/ansys/pytwin/evaluate/twin_model.py +++ b/src/ansys/pytwin/evaluate/twin_model.py @@ -1,9 +1,10 @@ -import atexit import json import os from pathlib import Path +import shutil import time from typing import Union +import weakref import numpy as np import pandas as pd @@ -73,26 +74,34 @@ def __init__(self, model_filepath: str): self._twin_runtime = None self._tbrom_info = None self._tbroms = None - if self._check_model_filepath_is_valid(model_filepath): self._model_filepath = model_filepath self._instantiate_twin_model() + self._finalizer = weakref.finalize(self, self._cleanup, self._twin_runtime, self.model_dir) - # We are registering the __del__ method to atexit module at the end of the TwinModel instantiation - # in order to avoid it to be called after the settings.cleanup_temp_pytwin_working_directory method - # that is deleting PyTwin temporary working directories. - # Otherwise, an error could be raised at python process exit because the TwinModel log file won't be freed - # and the settings.cleanup_temp_pytwin_working_directory will try to delete it. - # This happens when a TwinModel is instantiated into a script file, like in the examples. - atexit.register(self.__del__) + @staticmethod + def _cleanup(twin_runtime, model_dir): + """ + Close twin runtime and remove model temporary folder. + """ + if twin_runtime is not None: + if twin_runtime.is_model_opened: + twin_runtime.twin_close() + # Delete model directory + if os.path.exists(model_dir): + shutil.rmtree(model_dir) - def __del__(self): + def close(self): """ - Close twin runtime when object is garbage collected. + Cleanup object when user asks to close it. """ - if self._twin_runtime is not None: - if self._twin_runtime.is_model_opened: - self._twin_runtime.twin_close() + self._cleanup(self._twin_runtime, self.model_dir) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.close() def _check_model_filepath_is_valid(self, model_filepath): """ diff --git a/src/ansys/pytwin/settings.py b/src/ansys/pytwin/settings.py index 652f3af4..73a2f645 100644 --- a/src/ansys/pytwin/settings.py +++ b/src/ansys/pytwin/settings.py @@ -472,6 +472,9 @@ def cleanup_temp_pytwin_working_directory(): try: shutil.rmtree(PYTWIN_SETTINGS.TEMP_WORKING_DIRECTORY_PATH) except BaseException as e: - msg = "Something went wrong while trying to cleanup pytwin temporary directory!" - msg += f"error message:\n{str(e)}" + msg = ( + "Something went wrong while trying to cleanup pytwin temporary directory! You might have to clean it up " + "manually. " + ) + msg += f" Error message:\n{str(e)}" print(msg) diff --git a/tests/evaluate/test_twin_model.py b/tests/evaluate/test_twin_model.py index cb113914..77492bd0 100644 --- a/tests/evaluate/test_twin_model.py +++ b/tests/evaluate/test_twin_model.py @@ -365,12 +365,15 @@ def test_each_twin_model_has_a_subfolder_in_wd(self): # Verify a subfolder is created each time a new twin model is instantiated m_count = 5 wd = get_pytwin_working_dir() - ref_count = len(os.listdir(wd)) + ref_dir = os.listdir(wd) + cur_dir = os.listdir(wd) for m in range(m_count): model = TwinModel(model_filepath=COUPLE_CLUTCHES_FILEPATH) time.sleep(1) - wd = get_pytwin_working_dir() - assert len(os.listdir(wd)) == ref_count + m_count + for x in os.listdir(wd): + if x not in cur_dir: + cur_dir.append(x) + assert len(cur_dir) == len(ref_dir) + m_count def test_model_dir_migration_after_modifying_wd_dir(self): # Init unit test diff --git a/tests/evaluate/test_twin_model_finalizer.py b/tests/evaluate/test_twin_model_finalizer.py new file mode 100644 index 00000000..179d1f02 --- /dev/null +++ b/tests/evaluate/test_twin_model_finalizer.py @@ -0,0 +1,62 @@ +import os +import shutil +import time +import tracemalloc + +from pytwin import TwinModel +from pytwin.settings import get_pytwin_working_dir + +TBROM_MODEL_FILEPATH = os.path.join(os.path.dirname(__file__), "data", "ThermalTBROM_FieldInput_23R1.twin") +UNIT_TEST_WD = os.path.join(os.path.dirname(__file__), "unit_test_wd") + + +def reinit_settings(): + from pytwin.settings import reinit_settings_for_unit_tests + + reinit_settings_for_unit_tests() + if os.path.exists(UNIT_TEST_WD): + try: + shutil.rmtree(UNIT_TEST_WD) + except Exception as e: + pass + return UNIT_TEST_WD + + +class TestTwinModelFinalize: + def test_twin_model_finalizer_free_memory(self): + # Init unit test + reinit_settings() + # TwinModel memory is freed at the end of a loop and its model directory is deleted + tracemalloc.start() + snapshot = tracemalloc.take_snapshot() + allocated_mem_size = "" + model_dir = "" + for i in range(3): + twin_model = TwinModel(model_filepath=TBROM_MODEL_FILEPATH) + model_dir_old = model_dir + model_dir = twin_model.model_dir + snapshot2 = tracemalloc.take_snapshot() + top_stats = snapshot2.compare_to(snapshot, "lineno") + allocated_mem_size_old = allocated_mem_size + allocated_mem_size = f"{top_stats[0]}".split("size=")[1].split(",")[0].split("+")[1].split(" ")[0] + time.sleep(0.25) + if i > 0: + # Current twin_model directory exists + assert os.path.exists(model_dir) + # Previous twin_model directory as been deleted + assert not os.path.exists(model_dir_old) + # Previous twin_model memory as been freed (allow for +/- 0.5% difference of memory + assert ( + 1.005 * int(allocated_mem_size_old) > int(allocated_mem_size) > 0.995 * int(allocated_mem_size_old) + ) + + def test_clean_unit_test(self): + reinit_settings() + temp_wd = get_pytwin_working_dir() + parent_dir = os.path.dirname(temp_wd) + try: + for dir_name in os.listdir(parent_dir): + if dir_name not in temp_wd: + shutil.rmtree(os.path.join(parent_dir, dir_name)) + except Exception as e: + pass