[FEAT] Pytwin cleanup temporary directory at python process exit (#111)

Co-authored-by: Lucas Boucinha <lboucin@ansys.com>

src/ansys/pytwin/evaluate/tbrom.py

@@ -173,44 +173,42 @@ def input_field_size(self, fieldname: str):
         return len(self._infbasis[fieldname][0])
 
     @staticmethod
-    def _read_basis(fn):
-        fr = open(fn, "rb")
-        var = struct.unpack("cccccccccccccccc", fr.read(16))[0]
-        nb_val = struct.unpack("Q", fr.read(8))[0]
-        nb_mc = struct.unpack("Q", fr.read(8))[0]
-        basis = []
-        for i in range(nb_mc):
-            vec = []
-            for j in range(nb_val):
-                vec.append(struct.unpack("d", fr.read(8))[0])
-            basis.append(vec)
-        fr.close()
+    def _read_basis(filepath):
+        with open(filepath, "rb") as f:
+            var = struct.unpack("cccccccccccccccc", f.read(16))[0]
+            nb_val = struct.unpack("Q", f.read(8))[0]
+            nb_mc = struct.unpack("Q", f.read(8))[0]
+            basis = []
+            for i in range(nb_mc):
+                vec = []
+                for j in range(nb_val):
+                    vec.append(struct.unpack("d", f.read(8))[0])
+                basis.append(vec)
         return basis
 
     @staticmethod
-    def _read_binary(file):
-        fr = open(file, "rb")
-        nbdof = struct.unpack("Q", fr.read(8))[0]
-        vec = []
-        for i in range(nbdof):
-            vec.append(struct.unpack("d", fr.read(8))[0])
-        fr.close()
+    def _read_binary(filepath):
+        with open(filepath, "rb") as f:
+            nbdof = struct.unpack("Q", f.read(8))[0]
+            vec = []
+            for i in range(nbdof):
+                vec.append(struct.unpack("d", f.read(8))[0])
         return vec
 
     @staticmethod
-    def _write_binary(fn, vec):
-        if os.path.exists(fn):
-            os.remove(fn)
-        with open(fn, "xb") as fw:
-            fw.write(struct.pack("Q", len(vec)))
+    def _write_binary(filepath, vec):
+        if os.path.exists(filepath):
+            os.remove(filepath)
+        with open(filepath, "xb") as f:
+            f.write(struct.pack("Q", len(vec)))
             for i in vec:
-                fw.write(struct.pack("d", i))
+                f.write(struct.pack("d", i))
         return True
 
     @staticmethod
-    def _read_settings(settingspath):
-        f = open(settingspath)
-        data = json.load(f)
+    def _read_settings(filepath):
+        with open(filepath) as f:
+            data = json.load(f)
 
         namedselection = {}
         dimensionality = None
@@ -244,9 +242,9 @@ def _read_settings(settingspath):
         return [tbromns, dimensionality, outputname, unit]
 
     @staticmethod
-    def read_snapshot_size(file):
-        fr = open(file, "rb")
-        nbdof = struct.unpack("Q", fr.read(8))[0]
+    def read_snapshot_size(filepath):
+        with open(filepath, "rb") as f:
+            nbdof = struct.unpack("Q", f.read(8))[0]
         return nbdof
 
     @property

src/ansys/pytwin/evaluate/twin_model.py

@@ -950,7 +950,7 @@ def evaluate_batch(self, inputs_df: pd.DataFrame, field_inputs: dict = None):
         field_inputs : dict (optional)
             Dictionary of snapshot file paths that must be used as field input at all time instants
             given by the 'inputs_df' argument. One file path must be given per time instant, for a field input
-             of a TBROM included in the twin model, using following dictionary format:
+            of a TBROM included in the twin model, using following dictionary format:
             {"tbrom_name": {"field_input_name": [snapshotpath_t0, snapshotpath_t1, ... ]}}
 
         Returns

src/ansys/pytwin/settings.py

@@ -1,3 +1,4 @@
+import atexit
 from enum import Enum
 import logging
 import os
@@ -142,10 +143,15 @@ def modify_pytwin_working_dir(new_path: str, erase: bool = True):
     """
     Modify the global PyTwin working directory.
 
+    By default, a temporary directory is used by PyTwin as working directory. This temporary directory is automatically
+    cleaned up at exit of the python process that imported pytwin. When this method is used, the new PyTwin working
+    directory won't be deleted at python process exit. Note that this may lead to an overflow of the working directory.
+
     Parameters
     ----------
     new_path: str
         Absolute path to the working directory to use for PyTwin. The directory is created if it does not exist.
+        This directory is kept alive at python process exit.
     erase: bool, optional
         Whether to erase a non-empty existing working directory. The default is ``True``,
         in which case the existing working directory is erased and a new one is created.
@@ -231,14 +237,24 @@ def get_pytwin_working_dir():
     return PYTWIN_SETTINGS.working_dir
 
 
-def reinit_settings_for_unit_tests():
+def reinit_settings_for_unit_tests(create_new_temp_dir: bool = False):
     # Mutable attributes init
     _PyTwinSettings.LOGGING_OPTION = None
     _PyTwinSettings.LOGGING_LEVEL = None
     _PyTwinSettings.WORKING_DIRECTORY_PATH = None
-    _PyTwinSettings.SESSION_ID = None
     logging.getLogger(_PyTwinSettings.LOGGER_NAME).handlers.clear()
-    _PyTwinSettings().__init__()
+    if create_new_temp_dir:
+        PYTWIN_SETTINGS._initialize(keep_session_id=False)
+    else:
+        PYTWIN_SETTINGS._initialize(keep_session_id=True)
+    return PYTWIN_SETTINGS.SESSION_ID
+
+
+def reinit_settings_session_id_for_unit_tests(session_id: int):
+    PYTWIN_SETTINGS.SESSION_ID = session_id
+    PYTWIN_SETTINGS.TEMP_WORKING_DIRECTORY_PATH = os.path.join(
+        tempfile.gettempdir(), _PyTwinSettings.WORKING_DIRECTORY_NAME, _PyTwinSettings.SESSION_ID
+    )
 
 
 class _PyTwinSettings(object):
@@ -253,12 +269,15 @@ class _PyTwinSettings(object):
     LOGGING_LEVEL = None
     SESSION_ID = None
     WORKING_DIRECTORY_PATH = None
+    TEMP_WORKING_DIRECTORY_PATH = None
 
     # Immutable constants
     LOGGER_NAME = "pytwin_logger"
     LOGGING_FILE_NAME = "pytwin.log"
     WORKING_DIRECTORY_NAME = "pytwin"
     TEMP_WD_NAME = ".temp"
+    PYTWIN_START_MSG = "pytwin starts!"
+    PYTWIN_END_MSG = "pytwin ends!"
 
     @property
     def logfile(self):
@@ -289,8 +308,9 @@ def working_dir(self):
             raise PyTwinSettingsError(msg)
         return _PyTwinSettings.WORKING_DIRECTORY_PATH
 
-    def __init__(self):
-        self._initialize()
+    def __init__(self, keep_session_id: bool = False):
+        self._initialize(keep_session_id)
+        self.logger.info(_PyTwinSettings.PYTWIN_START_MSG)
 
     @staticmethod
     def _add_default_file_handler_to_pytwin_logger(filepath: str, level: PyTwinLogLevel, mode: str = "w"):
@@ -321,9 +341,13 @@ def _add_default_stream_handler_to_pytwin_logger(level: PyTwinLogLevel):
         logger.addHandler(log_handler)
 
     @staticmethod
-    def _initialize():
+    def _initialize(keep_session_id: bool):
         pytwin_logger = logging.getLogger(_PyTwinSettings.LOGGER_NAME)
         pytwin_logger.handlers.clear()
+
+        if not keep_session_id:
+            _PyTwinSettings.SESSION_ID = f"{uuid.uuid4()}"[0:24].replace("-", "")
+
         _PyTwinSettings._initialize_wd()
         _PyTwinSettings._initialize_logging()
 
@@ -346,12 +370,11 @@ def _initialize_wd():
         Provides default settings for the PyTwin working directory.
         """
         # Create a unique working directory for each python process that imports pytwin
-        _PyTwinSettings.SESSION_ID = f"{uuid.uuid4()}"[0:24].replace("-", "")
-        pytwin_wd_dir = os.path.join(
+        _PyTwinSettings.TEMP_WORKING_DIRECTORY_PATH = os.path.join(
             tempfile.gettempdir(), _PyTwinSettings.WORKING_DIRECTORY_NAME, _PyTwinSettings.SESSION_ID
         )
-        os.makedirs(pytwin_wd_dir)
-        _PyTwinSettings.WORKING_DIRECTORY_PATH = pytwin_wd_dir
+        os.makedirs(_PyTwinSettings.TEMP_WORKING_DIRECTORY_PATH, exist_ok=True)
+        _PyTwinSettings.WORKING_DIRECTORY_PATH = _PyTwinSettings.TEMP_WORKING_DIRECTORY_PATH
 
     @staticmethod
     def _migration_due_to_new_wd(old_path: str, new_path: str):
@@ -447,3 +470,18 @@ def modify_logging(new_option: PyTwinLogOption, new_level: PyTwinLogLevel):
 
 
 PYTWIN_SETTINGS = _PyTwinSettings()  # This instance is here to launch default settings initialization.
+
+
+@atexit.register
+def cleanup_temp_pytwin_working_directory():
+    pytwin_logger = PYTWIN_SETTINGS.logger
+    pytwin_logger.info(PYTWIN_SETTINGS.PYTWIN_END_MSG)
+    for handler in pytwin_logger.handlers:
+        handler.close()
+        pytwin_logger.removeHandler(handler)
+    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)}"
+        print(msg)

tests/evaluate/test_tbrom.py

@@ -776,7 +776,7 @@ def test_generate_points_with_tbrom_exceptions(self):
 
     def test_tbrom_getters_that_do_not_need_initialization(self):
         reinit_settings()
-        model_filepath = download_file("ThermalTBROM_23R1_other.twin", "twin_files")
+        model_filepath = download_file("ThermalTBROM_23R1_other.twin", "twin_files", force_download=True)
         twin = TwinModel(model_filepath=model_filepath)
 
         # Test rom name

tests/evaluate/test_twin_model.py

@@ -1,4 +1,5 @@
 import os
+import shutil
 import time
 
 import pandas as pd
@@ -14,15 +15,13 @@
 UNIT_TEST_WD = os.path.join(os.path.dirname(__file__), "unit_test_wd")
 
 
-def reinit_settings():
-    import shutil
-
+def reinit_settings(create_new_temp_dir: bool = False):
     from pytwin.settings import reinit_settings_for_unit_tests
 
-    reinit_settings_for_unit_tests()
+    session_id = reinit_settings_for_unit_tests(create_new_temp_dir)
     if os.path.exists(UNIT_TEST_WD):
         shutil.rmtree(UNIT_TEST_WD)
-    return UNIT_TEST_WD
+    return UNIT_TEST_WD, session_id
 
 
 class TestTwinModel:
@@ -357,24 +356,29 @@ def test_close_method(self):
         twin = TwinModel(model_filepath=model_filepath)
 
     def test_each_twin_model_has_a_subfolder_in_wd(self):
+        from pytwin.settings import reinit_settings_session_id_for_unit_tests
+
         # Init unit test
-        reinit_settings()
+        wd, session_id = reinit_settings(create_new_temp_dir=True)
         logger = get_pytwin_logger()
         # Verify a subfolder is created each time a new twin model is instantiated
         m_count = 5
         for m in range(m_count):
             model = TwinModel(model_filepath=COUPLE_CLUTCHES_FILEPATH)
             time.sleep(1)
         wd = get_pytwin_working_dir()
-        temp = os.listdir(wd)
         assert len(os.listdir(wd)) == m_count + 2
+        reinit_settings_session_id_for_unit_tests(session_id)
 
     def test_model_dir_migration_after_modifying_wd_dir(self):
+        from pytwin.settings import reinit_settings_session_id_for_unit_tests
+
         # Init unit test
-        wd = reinit_settings()
+        wd, session_id = reinit_settings(create_new_temp_dir=True)
         assert not os.path.exists(wd)
         model = TwinModel(model_filepath=COUPLE_CLUTCHES_FILEPATH)
         assert os.path.split(model.model_dir)[0] == get_pytwin_working_dir()
+
         # Run test
         modify_pytwin_working_dir(new_path=wd)
         assert os.path.split(model.model_dir)[0] == wd
@@ -383,6 +387,46 @@ def test_model_dir_migration_after_modifying_wd_dir(self):
         assert os.path.split(model2.model_dir)[0] == wd
         assert len(os.listdir(wd)) == 2 + 1 + 1  # 2 models + pytwin log + .temp
 
+        # Finalize unit test
+        reinit_settings_session_id_for_unit_tests(session_id)
+
+    def test_multiprocess_execution_modify_wd_dir(self):
+        import subprocess
+        import sys
+
+        # Init unit test
+        wd, session_id = reinit_settings()
+        assert not os.path.exists(wd)
+        current_wd_dir_count = len(os.listdir(os.path.dirname(get_pytwin_working_dir())))
+
+        # In another process, modify working dir before having instantiating a twin model
+        subprocess_code = "import pytwin, os\n"
+        subprocess_code += f'pytwin.modify_pytwin_working_dir(new_path=r"{wd}")\n'
+        subprocess_code += f'model = pytwin.TwinModel(model_filepath=r"{COUPLE_CLUTCHES_FILEPATH}")\n'
+        subprocess_code += f'assert os.path.split(model.model_dir)[0] == r"{wd}"\n'
+        result = subprocess.run([sys.executable, "-c", subprocess_code], capture_output=True)
+        new_wd_dir_count = len(os.listdir(os.path.dirname(get_pytwin_working_dir())))
+
+        assert new_wd_dir_count == current_wd_dir_count
+        assert len(result.stderr) == 0
+        assert os.path.exists(wd)
+
+        # In another process, modify working dir after having instantiating a twin model
+        subprocess_code = "import pytwin, os\n"
+        subprocess_code += f'model = pytwin.TwinModel(model_filepath=r"{COUPLE_CLUTCHES_FILEPATH}")\n'
+        subprocess_code += "assert os.path.split(model.model_dir)[0] == pytwin.get_pytwin_working_dir()\n"
+        subprocess_code += f'pytwin.modify_pytwin_working_dir(new_path=r"{wd}")\n'
+        subprocess_code += f'assert os.path.split(model.model_dir)[0] == r"{wd}"\n'
+        result = subprocess.run([sys.executable, "-c", subprocess_code], capture_output=True)
+        new_wd_dir_count = len(os.listdir(os.path.dirname(get_pytwin_working_dir())))
+
+        if sys.platform != "linux":
+            assert new_wd_dir_count == current_wd_dir_count + 1
+        else:
+            assert new_wd_dir_count == current_wd_dir_count
+        assert len(result.stderr) == 0
+        assert os.path.exists(wd)
+
     def test_model_warns_at_initialization(self):
         # Init unit test
         wd = reinit_settings()
@@ -539,3 +583,11 @@ def test_save_and_load_state_with_rc_heat_circuit(self):
 
     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

tests/evaluate/test_twin_model_logging.py

@@ -1,27 +1,28 @@
 import os
+import shutil
 
 from pytwin import PYTWIN_LOGGING_OPT_NOLOGGING, TwinModel
-from pytwin.settings import get_pytwin_log_file, modify_pytwin_logging
+from pytwin.settings import get_pytwin_log_file, get_pytwin_working_dir, modify_pytwin_logging
 
 COUPLE_CLUTCHES_FILEPATH = os.path.join(os.path.dirname(__file__), "data", "CoupleClutches_22R2_other.twin")
 UNIT_TEST_WD = os.path.join(os.path.dirname(__file__), "unit_test_wd")
 
 
-def reinit_settings():
-    import shutil
-
+def reinit_settings(create_new_temp_dir: bool = False):
     from pytwin.settings import reinit_settings_for_unit_tests
 
-    reinit_settings_for_unit_tests()
+    session_id = reinit_settings_for_unit_tests(create_new_temp_dir)
     if os.path.exists(UNIT_TEST_WD):
         shutil.rmtree(UNIT_TEST_WD)
-    return UNIT_TEST_WD
+    return UNIT_TEST_WD, session_id
 
 
 class TestTwinModelLogging:
     def test_twin_model_no_logging(self):
+        from pytwin.settings import reinit_settings_session_id_for_unit_tests
+
         # Init unit test
-        reinit_settings()
+        wd, session_id = reinit_settings(create_new_temp_dir=True)
         # Twin Model does not log anything if PYTWIN_LOGGING_OPT_NOLOGGING
         modify_pytwin_logging(new_option=PYTWIN_LOGGING_OPT_NOLOGGING)
         log_file = get_pytwin_log_file()
@@ -34,3 +35,15 @@ def test_twin_model_no_logging(self):
         temp_dir = twin.model_temp
         assert os.path.exists(temp_dir)
         assert len(os.listdir(temp_dir)) == 0
+        reinit_settings_session_id_for_unit_tests(session_id)
+
+    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