modify saving

aimclub · maypink · Sep 12, 2023 · Mar 17, 2023 · Mar 17, 2023 · Mar 21, 2023
commit 8d664d2e8f76f12a6451d2cfbdba1ac5a578f588
diff --git a/golem/core/optimisers/adaptive/mab_agents/mab_agent.py b/golem/core/optimisers/adaptive/mab_agents/mab_agent.py
@@ -1,3 +1,4 @@
+import ntpath
 import os.path
 import pickle
 import random
@@ -61,11 +62,14 @@ def partial_fit(self, experience: ExperienceBuffer):
     def save(self, path_to_save: Optional[str] = None):
         """ Saves bandit to specified file. """
 
-        # if path was not specified
-        if not self._path_to_save and not path_to_save:
+        if not path_to_save:
+            path_to_save = self._path_to_save
+
+        # if path was not specified at all
+        if not path_to_save:
             path_to_save = os.path.join(default_data_dir(), 'MAB')
 
-        if not self._path_to_save.endswith('.pkl'):
+        if not path_to_save.endswith('.pkl'):
             os.makedirs(path_to_save, exist_ok=True)
             mabs_num = [int(name.split('_')[0]) for name in os.listdir(path_to_save)
                         if re.fullmatch(r'\d_mab.pkl', name)]
@@ -75,7 +79,9 @@ def save(self, path_to_save: Optional[str] = None):
                 max_saved_mab = max(mabs_num) + 1
             path_to_file = os.path.join(path_to_save, f'{max_saved_mab}_mab.pkl')
         else:
-            path_to_file = self._path_to_save
+            path_to_dir, _ = ntpath.split(path_to_save)
+            os.makedirs(path_to_dir, exist_ok=True)
+            path_to_file = path_to_save
         with open(path_to_file, 'wb') as f:
             pickle.dump(self, f)
         self._log.info(f"MAB was saved to {self._path_to_save}")

diff --git a/test/unit/adaptive/test_mab_agent.py b/test/unit/adaptive/test_mab_agent.py
@@ -1,24 +1,34 @@
 import os.path
 from pathlib import Path
 
+import pytest
+
 from golem.core.optimisers.adaptive.mab_agents.mab_agent import MultiArmedBanditAgent
 
 
-def test_save_mab():
-    """ Tests if MAB is saved. """
-    file_name = 'test_mab.pkl'
-    path_to_save = os.path.join(Path(__file__).parent, file_name)
+@pytest.mark.parametrize('path_to_save, file_name',
+                         [(os.path.join(Path(__file__).parent, 'test_mab.pkl'), 'test_mab.pkl'),
+                          (os.path.join(Path(__file__).parent), '0_mab.pkl')])
+def test_save_mab(path_to_save, file_name):
+    """ Tests if MAB is saved with specifying file_nam and without. """
     mab = MultiArmedBanditAgent(actions=[0, 1, 2],
                                 n_jobs=1,
                                 path_to_save=path_to_save)
     mab.save()
     assert file_name in os.listdir(Path(__file__).parent)
+    os.remove(path_to_save if path_to_save.endswith('pkl') else os.path.join(path_to_save, file_name))
 
 
 def test_load_mab():
     """ Tests if MAB is loaded. """
     file_name = 'test_mab.pkl'
     path_to_load = os.path.join(Path(__file__).parent, file_name)
+    # save mab to load it later
+    mab = MultiArmedBanditAgent(actions=[0, 1, 2],
+                                n_jobs=1,
+                                path_to_save=path_to_load)
+    mab.save()
+
     mab = MultiArmedBanditAgent.load(path=path_to_load)
     assert isinstance(mab, MultiArmedBanditAgent)
     os.remove(path_to_load)