[patch] Fall back on cloudpickle (#411)

* Refactor storage test

* Remove empty line

* Rename workflow in tests

With the introduction of a pickle backend, this test wound up failing. I haven't been able to work out exactly why, but since pickle is much faster my guess is that something in the test parallelization gets in a race condition with deleting a saved workflow from another test. Simply renaming it here solved the issue, which is an innocuous change.

* Introduce `pickle` as a backend for saving

* Fix root cause of storage conflict

Ok, the reason workflow storage tests were failing after introducing pickle was that the original version of the test had the wrong order of the try/finally scope and the for-loop scope. I got away with it earlier because of interplay between the outer-most member of the loop and the default storage backend. Actually fixing the problem is as simple as ensuring the "finally delete" clause happens after _each_ loop step. This does that and reverts the renaming of the workflow.

* Again, correctly order try/finally and for-loops

* Remove keyword argument from pure-decorator

You're only supposed to use it as a decorator to start with, so the kwarg was senseless

* Add factory import source

This is necessary (but not sufficient) to get `as_dataclass_node` decorated classes to pickle. The field name is a bit off compared to `Function` and `Macro`, as now we decorate a class definition instead of a function definition, but it's close.

* Bring the dataclass node in line with function and macro

* Leverage new pyiron_snippets.factory stuff to find the class

* Mangle the stored dataclass qualname so it can be found later

* Add tests

* Update docs examples to reflect new naming

* Update snippets dependency

* [dependabot skip] Update env file

* Use new pyiron_snippets syntax consistently

* Expose `as_dataclass_node` in the API

Now that it's pickling as well as anything else

* [patch] Fall back on cloudpickle

When the pickle backend fails

* Format black

---------

Co-authored-by: pyiron-runner <[email protected]>

pyiron_workflow/mixin/storage.py

@@ -12,6 +12,7 @@
 import sys
 from typing import Optional
 
+import cloudpickle
 import h5io
 from pyiron_snippets.files import FileObject, DirectoryObject
 
@@ -88,6 +89,7 @@ def _delete(self):
 class PickleStorage(StorageInterface):
 
     _PICKLE_STORAGE_FILE_NAME = "pickle.pckl"
+    _CLOUDPICKLE_STORAGE_FILE_NAME = "cloudpickle.cpckl"
 
     def __init__(self, owner: HasPickleStorage):
         super().__init__(owner=owner)
@@ -97,12 +99,22 @@ def owner(self) -> HasPickleStorage:
         return self._owner
 
     def _save(self):
-        with open(self._pickle_storage_file_path, "wb") as file:
-            pickle.dump(self.owner, file)
+        try:
+            with open(self._pickle_storage_file_path, "wb") as file:
+                pickle.dump(self.owner, file)
+        except Exception:
+            self._delete()
+            with open(self._cloudpickle_storage_file_path, "wb") as file:
+                cloudpickle.dump(self.owner, file)
 
     def _load(self):
-        with open(self._pickle_storage_file_path, "rb") as file:
-            inst = pickle.load(file)
+        if self._has_pickle_contents:
+            with open(self._pickle_storage_file_path, "rb") as file:
+                inst = pickle.load(file)
+        elif self._has_cloudpickle_contents:
+            with open(self._cloudpickle_storage_file_path, "rb") as file:
+                inst = cloudpickle.load(file)
+
         if inst.__class__ != self.owner.__class__:
             raise TypeError(
                 f"{self.owner.label} cannot load, as it has type "
@@ -111,24 +123,38 @@ def _load(self):
             )
         self.owner.__setstate__(inst.__getstate__())
 
+    def _delete_file(self, file: str):
+        FileObject(file, self.owner.storage_directory).delete()
+
     def _delete(self):
-        if self.has_contents:
-            FileObject(
-                self._PICKLE_STORAGE_FILE_NAME, self.owner.storage_directory
-            ).delete()
+        if self._has_pickle_contents:
+            self._delete_file(self._PICKLE_STORAGE_FILE_NAME)
+        elif self._has_cloudpickle_contents:
+            self._delete_file(self._CLOUDPICKLE_STORAGE_FILE_NAME)
+
+    def _storage_path(self, file: str):
+        return str((self.owner.storage_directory.path / file).resolve())
 
     @property
     def _pickle_storage_file_path(self) -> str:
-        return str(
-            (
-                self.owner.storage_directory.path / self._PICKLE_STORAGE_FILE_NAME
-            ).resolve()
-        )
+        return self._storage_path(self._PICKLE_STORAGE_FILE_NAME)
+
+    @property
+    def _cloudpickle_storage_file_path(self) -> str:
+        return self._storage_path(self._CLOUDPICKLE_STORAGE_FILE_NAME)
 
     @property
     def _has_contents(self) -> bool:
+        return self._has_pickle_contents or self._has_cloudpickle_contents
+
+    @property
+    def _has_pickle_contents(self) -> bool:
         return os.path.isfile(self._pickle_storage_file_path)
 
+    @property
+    def _has_cloudpickle_contents(self) -> bool:
+        return os.path.isfile(self._cloudpickle_storage_file_path)
+
 
 class H5ioStorage(StorageInterface):
 

pyiron_workflow/node.py

@@ -143,7 +143,8 @@ class Node(
         - As long as you haven't put anything unpickleable on them, or defined them in
             an unpicklable place (e.g. in the `<locals>` of another function), you can
             simple (un)pickle nodes. There is no save/load interface for this right
-            now, just import pickle and do it.
+            now, just import pickle and do it. The "pickle" backend to the `Node.save`
+            method will fall back on `cloudpickle` as needed to overcome this.
         - Saving is triggered manually, or by setting a flag to save after the nodes
             runs.
         - At the end of instantiation, nodes will load automatically if they find saved
@@ -168,8 +169,8 @@ class Node(
         - [ALPHA ISSUE] There are three possible back-ends for saving: one leaning on
             `tinybase.storage.GenericStorage` (in practice,
             `H5ioStorage(GenericStorage)`), and the other that uses the `h5io` module
-            directly. The third (default) option is to use `pickle`. The backend used
-            is always the one on the graph root.
+            directly. The third (default) option is to use `(cloud)pickle`. The backend
+            used is always the one on the graph root.
         - [ALPHA ISSUE] The `h5io` backend is deprecated -- it can't handle custom
             reconstructors (i.e. when `__reduce__` returns a tuple with some
             non-standard callable as its first entry), and basically all our nodes do

tests/unit/test_node.py

@@ -20,7 +20,7 @@ class ANode(Node):
     """To de-abstract the class"""
 
     def _setup_node(self) -> None:
-        self._inputs = Inputs(InputData("x", self, type_hint=int))
+        self._inputs = Inputs(InputData("x", self, type_hint=int),)
         self._outputs = OutputsWithInjection(
             OutputDataWithInjection("y", self, type_hint=int),
         )
@@ -472,7 +472,29 @@ def test_storage(self):
                         force_run.outputs.y.value,
                         msg="Destroying the save should allow immediate re-running"
                     )
+
+                    hard_input = ANode(label="hard", storage_backend=backend)
+                    hard_input.inputs.x.type_hint = callable
+                    hard_input.inputs.x = lambda x: x * 2
+                    if backend == "pickle":
+                        hard_input.save()
+                        reloaded = ANode(
+                            label=hard_input.label,
+                            storage_backend=backend
+                        )
+                        self.assertEqual(
+                            reloaded.inputs.x.value(4),
+                            hard_input.inputs.x.value(4),
+                            msg="Cloud pickle should be strong enough to recover this"
+                        )
+                    else:
+                        with self.assertRaises(
+                            (TypeError, AttributeError),
+                            msg="Other backends are not powerful enough for some values"
+                        ):
+                            hard_input.save()
                 finally:
+                    hard_input.delete_storage()
                     self.n1.delete_storage()
 
     @unittest.skipIf(sys.version_info < (3, 11), "Storage will only work in 3.11+")