[FIX] Issue 128: unregister TwinModel.__del__ from atexit and enhance…

… PR116 (#134)

Co-authored-by: chrpetre <cpetre@AAPtLkOSCC9TJDC.win.ansys.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Lucas Boucinha <108674439+lboucin@users.noreply.github.com>

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):
         """

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)

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

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