imt-ring 1.6.35__tar.gz → 1.6.37__tar.gz

{imt_ring-1.6.35 → imt_ring-1.6.37}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: imt-ring
-Version: 1.6.35
+Version: 1.6.37
 Summary: RING: Recurrent Inertial Graph-based Estimator
 Author-email: Simon Bachhuber <simon.bachhuber@fau.de>
 Project-URL: Homepage, https://github.com/SimiPixel/ring

{imt_ring-1.6.35 → imt_ring-1.6.37}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "imt-ring"
-version = "1.6.35"
+version = "1.6.37"
 authors = [
   { name="Simon Bachhuber", email="simon.bachhuber@fau.de" },
 ]

{imt_ring-1.6.35 → imt_ring-1.6.37}/src/imt_ring.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: imt-ring
-Version: 1.6.35
+Version: 1.6.37
 Summary: RING: Recurrent Inertial Graph-based Estimator
 Author-email: Simon Bachhuber <simon.bachhuber@fau.de>
 Project-URL: Homepage, https://github.com/SimiPixel/ring

{imt_ring-1.6.35 → imt_ring-1.6.37}/src/ring/ml/training_loop.py

@@ -2,11 +2,12 @@ import random
 from typing import Optional
 
 import jax
-from ring.algorithms import Generator
-from ring.ml import ml_utils
 import tqdm
 import tree_utils
 
+from ring.algorithms import Generator
+from ring.ml import ml_utils
+
 _KILL_RUN = False
 
 
@@ -83,14 +84,20 @@ class TrainingLoop:
         # reset the kill_run flag from previous runs
         send_kill_run_signal(value=False)
 
-        for _ in tqdm.tqdm(range(n_episodes)):
-            self.step()
-
-            if recv_kill_run_signal():
-                break
+        try:
+            for _ in tqdm.tqdm(range(n_episodes)):
+                self.step()
 
-        if close_afterwards:
-            self.close()
+                if recv_kill_run_signal():
+                    break
+        except Exception as e:
+            print(
+                "Exception occured, attemping to .close() all callbacks before raising."
+            )
+            raise e
+        finally:
+            if close_afterwards:
+                self.close()
 
         return recv_kill_run_signal()
 

{imt_ring-1.6.35 → imt_ring-1.6.37}/src/ring/utils/utils.py

@@ -1,5 +1,6 @@
 from importlib import import_module as _import_module
 import io
+from pathlib import Path
 import pickle
 import random
 from typing import Optional
@@ -152,13 +153,28 @@ def import_lib(
 
 def pickle_save(obj, path, overwrite: bool = False):
     path = parse_path(path, extension="pickle", file_exists_ok=overwrite)
-    with open(path, "wb") as file:
-        pickle.dump(obj, file, protocol=5)
+    try:
+        with open(path, "wb") as file:
+            pickle.dump(obj, file, protocol=5)
+    except OSError as e:
+        print(
+            f"saving with `pickle` throws exception {e}. "
+            + "Attempting to save using `joblib`"
+        )
+        path = parse_path(path, extension="joblib", file_exists_ok=overwrite)
+        import joblib
+
+        joblib.dump(obj, path)
 
 
 def pickle_load(
     path,
 ):
+    if Path(path).suffix == ".joblib":
+        import joblib
+
+        return joblib.load(path)
+
     path = parse_path(path, extension="pickle", require_is_file=True)
     with open(path, "rb") as file:
         obj = pickle.load(file)