[not for merge] Parsl serializer plugins development

Makefile

@@ -84,7 +84,7 @@ radical_local_test:
 
 .PHONY: config_local_test
 config_local_test: $(CCTOOLS_INSTALL)
-	pip3 install ".[monitoring,visualization,proxystore]"
+	pip3 install ".[monitoring,visualization,proxystore,serializer_demos]"
 	PYTHONPATH=/tmp/cctools/lib/python3.8/site-packages pytest parsl/tests/ -k "not cleannet" --config local --random-order --durations 10
 
 .PHONY: site_test

docs/userguide/plugins.rst

@@ -87,6 +87,20 @@ from invoking a Parsl app. This includes as the return value of a
 An specific example of this is integrating Globus Compute tasks into a Parsl
 task graph. See :ref:`label-join-globus-compute`
 
+
+Serialization
+-------------
+
+By default Parsl will serialize objects with either ``pickle`` or ``dill``, in
+some cases applying an LRU cache.
+
+Parsl has an unstable API to register new serialization methods
+
+Additional serialization methods can be registered by XXXXX TODO XXXXX
+
+Limitations:
+TODO
+
 Dependency resolution
 ---------------------
 

mypy.ini

@@ -69,6 +69,12 @@ check_untyped_defs = True
 disallow_subclassing_any = True
 disallow_untyped_defs = True
 
+[mypy-parsl.serialize.plugin_proxystore_deep_pickle.*]
+disallow_subclassing_any = False
+
+[mypy-parsl.serialize.plugin_codeprotector.*]
+disallow_subclassing_any = False
+
 [mypy-parsl.serialize.proxystore.*]
 # parsl/serialize/proxystore.py:9: error: Class cannot subclass "Pickler" (has type "Any") 
 disallow_subclassing_any = False
@@ -183,6 +189,9 @@ ignore_missing_imports = True
 #[mypy-multiprocessing.synchronization.*]
 #ignore_missing_imports = True
 
+[mypy-proxystore.connectors.file]
+ignore_missing_imports = True
+
 [mypy-pandas.*]
 ignore_missing_imports = True
 
@@ -204,5 +213,11 @@ ignore_missing_imports = True
 [mypy-setproctitle.*]
 ignore_missing_imports = True
 
+
+# for serializer plugin demos:
+
+[mypy-serpent.*]
+ignore_missing_imports = True
+
 [mypy-proxystore.*]
 ignore_missing_imports = True

parsl/executors/high_throughput/executor.py

@@ -493,6 +493,9 @@ def _queue_management_worker(self):
 
                             if 'result' in msg:
                                 result = deserialize(msg['result'])
+                                # ^ if this raises an exception, queue management worker fails and
+                                # parsl hangs. this is more relevant now allowing user-pluggable
+                                # serialization and so user code exceptions can be raised here.
                                 task_fut.set_result(result)
 
                             elif 'exception' in msg:

parsl/serialize/base.py

@@ -1,12 +1,9 @@
-import logging
-from abc import abstractmethod
+from abc import ABCMeta, abstractmethod
 from functools import cached_property
 from typing import Any
 
-logger = logging.getLogger(__name__)
 
-
-class SerializerBase:
+class SerializerBase(metaclass=ABCMeta):
     """ Adds shared functionality for all serializer implementations
     """
 

parsl/serialize/facade.py

@@ -8,34 +8,80 @@
 
 logger = logging.getLogger(__name__)
 
+methods_for_code: List[SerializerBase]
+methods_for_code = []
+
+methods_for_data: List[SerializerBase]
+methods_for_data = []
+
+deserializers: Dict[bytes, SerializerBase]
+deserializers = {}
+
+
+# maybe it's weird to want to clear away all serializers rather than just the
+# data or just code ones? eg in proxystore test, clearing serializers means
+# there is no code serializer... so just data sreializer should be cleared
+# (or if only having one kind of serializer, no code/data distinction, then
+# this is more consistent, but clearing serializer is still a bit weird in
+#  that case (maybe for security?) - with inserting one near the start more
+# usual?
+# def clear_serializers() -> None:
+#    # does not clear deserializers because remote sending back will have a
+#    # different serializer list (and wants to send back results with one of
+#    # the default 4) so clearing all the deserializers means we cannot receive
+#    # results...
+#    global methods_for_data, methods_for_code
+#    methods_for_code = []
+#    methods_for_data = []
+
+
+def unregister_serializer(serializer: SerializerBase) -> None:
+    logger.info(f"BENC: deserializers {deserializers}, serializer {serializer}")
+    logger.info(f"BENC: unregistering serializer {serializer}")
+    if serializer.identifier in deserializers:
+        del deserializers[serializer.identifier]
+    else:
+        logger.warning("BENC: not found in deserializers list")
+    if serializer in methods_for_code:
+        logger.info("BENC: removing serializer from methods_for_code")
+        methods_for_code.remove(serializer)
+    else:
+        logger.warning("BENC: not found in methods for code")
+    if serializer in methods_for_data:
+        logger.info("BENC: removing serializer from methods_for_data")
+        methods_for_data.remove(serializer)
+    else:
+        logger.warning("BENC: not found in methods for data")
 
-methods_for_code: Dict[bytes, SerializerBase] = {}
 
+# structuring it this way is probably wrong - should perhaps be a single
+# Pickle-variant (or dill-variant) that is used, with all pluggable hooked
+# in - eg proxystore should hook into the same Pickle/Dill subclass as
+# CodeProtector?
 
 def register_method_for_code(s: SerializerBase) -> None:
-    methods_for_code[s.identifier] = s
+    deserializers[s.identifier] = s
+    methods_for_code.insert(0, s)
 
 
 register_method_for_code(concretes.DillCallableSerializer())
-
-
-methods_for_data: Dict[bytes, SerializerBase] = {}
+# register_method_for_code(CodeProtectorSerializer())
 
 
 def register_method_for_data(s: SerializerBase) -> None:
-    methods_for_data[s.identifier] = s
+    deserializers[s.identifier] = s
+    methods_for_data.insert(0, s)
 
 
+# These must be registered in reverse order of
+# importance: later registered serializers
+# will take priority over earlier ones. This is
+# to facilitate user registered serializers
+
 register_method_for_data(concretes.PickleSerializer())
 register_method_for_data(concretes.DillSerializer())
 
 
-# When deserialize dynamically loads a deserializer, it will be stored here,
-# rather than in the methods_for_* dictionaries, so that loading does not
-# cause it to be used for future serializations.
-additional_methods_for_deserialization: Dict[bytes, SerializerBase] = {}
-
-
 def pack_apply_message(func: Any, args: Any, kwargs: Any, buffer_threshold: int = int(128 * 1e6)) -> bytes:
     """Serialize and pack function and parameters
 
@@ -106,24 +152,35 @@ def serialize(obj: Any, buffer_threshold: int = int(1e6)) -> bytes:
     Individual serialization methods might raise a TypeError (eg. if objects are non serializable)
     This method will raise the exception from the last method that was tried, if all methods fail.
     """
-    result: Union[bytes, Exception]
+    logger.info(f"BENC: Trying to serialize {obj}")
+    result: Union[bytes, Exception, None]
+    result = None
     if callable(obj):
         methods = methods_for_code
     else:
         methods = methods_for_data
 
-    for method in methods.values():
+    if methods == []:
+        raise RuntimeError("There are no configured serializers")
+
+    for method in methods:
         try:
+            logger.info(f"BENC: trying serializer {method}")
             result = method.identifier + b'\n' + method.serialize(obj)
         except Exception as e:
+            logger.warning(f"BENC: serializer {method} skipping, with exception: {e}")
             result = e
             continue
         else:
             break
 
-    if isinstance(result, BaseException):
+    if result is None:
+        raise RuntimeError("No serializer returned a result")
+    elif isinstance(result, BaseException):
+        logger.error("Serializer returned an excepton, reraise")
         raise result
     else:
+        logger.debug("Serialization complete")
         if len(result) > buffer_threshold:
             logger.warning(f"Serialized object exceeds buffer threshold of {buffer_threshold} bytes, this could cause overflows")
         return result
@@ -139,12 +196,8 @@ def deserialize(payload: bytes) -> Any:
     """
     header, body = payload.split(b'\n', 1)
 
-    if header in methods_for_code:
-        deserializer = methods_for_code[header]
-    elif header in methods_for_data:
-        deserializer = methods_for_data[header]
-    elif header in additional_methods_for_deserialization:
-        deserializer = additional_methods_for_deserialization[header]
+    if header in deserializers:
+        deserializer = deserializers[header]
     else:
         logger.info("Trying to dynamically load deserializer: {!r}".format(header))
         # This is a user plugin point, so expect exceptions to happen.
@@ -154,7 +207,7 @@ def deserialize(payload: bytes) -> Any:
             module = importlib.import_module(decoded_module_name)
             deserializer_class = getattr(module, class_name.decode('utf-8'))
             deserializer = deserializer_class()
-            additional_methods_for_deserialization[header] = deserializer
+            deserializers[header] = deserializer
         except Exception as e:
             raise DeserializerPluginError(header) from e
 

parsl/serialize/plugin_codeprotector.py

@@ -0,0 +1,44 @@
+import io
+import logging
+import sys
+import types
+from typing import Any
+
+import dill
+
+from parsl.serialize.base import SerializerBase
+
+logger = logging.getLogger(__name__)
+
+
+def _deprotectCode(major: int, minor: int, b: bytes) -> Any:
+    if sys.version_info.major != major:
+        raise RuntimeError("Major version mismatch deserializing code")
+    if sys.version_info.minor != minor:
+        raise RuntimeError("Major version mismatch deserializing code")
+
+    return dill.loads(b)
+
+
+class CodeProtectorPickler(dill.Pickler):
+
+    def reducer_override(self, o: Any) -> Any:
+        logger.info(f"BENC: reducing object {o!r} of type {type(o)}")
+        if isinstance(o, types.CodeType):
+            logger.info(f"BENC: special casing code object {o!r} of type {type(o)}")
+            return (_deprotectCode, (sys.version_info.major, sys.version_info.minor, dill.dumps(o)))
+
+        return NotImplemented
+
+
+class CodeProtectorSerializer(SerializerBase):
+
+    def serialize(self, data: Any) -> bytes:
+
+        f = io.BytesIO()
+        pickler = CodeProtectorPickler(file=f)
+        pickler.dump(data)
+        return f.getvalue()
+
+    def deserialize(self, body: bytes) -> Any:
+        return dill.loads(body)

parsl/serialize/plugin_proxystore.py

@@ -0,0 +1,44 @@
+# parsl/serialize/concretes.py:10: error: Module "proxystore.store" does not explicitly export attribute "Store"  [attr-defined]
+import pickle
+from typing import Any, Optional
+
+from proxystore.connectors.file import FileConnector
+from proxystore.store import Store, register_store
+
+from parsl.serialize.base import SerializerBase
+from parsl.serialize.facade import register_method_for_data
+
+
+class ProxyStoreSerializer(SerializerBase):
+
+    def __init__(self, store: Optional[Store] = None) -> None:
+        """Because of jumbled use of this class for init-time configurable
+        serialization, and non-configurable remote deserializer loading, the
+        store field can be None... TODO: this would go away if serializer and
+        deserializer were split into different objects/classes/functions."""
+        self._store = store
+
+    def serialize(self, data: Any) -> bytes:
+        assert self._store is not None
+        assert data is not None
+        p = self._store.proxy(data)
+        return pickle.dumps(p)
+
+    def deserialize(self, body: bytes) -> Any:
+        return pickle.loads(body)
+
+
+def register_proxystore_serializer() -> None:
+    """Initializes proxystore and registers it as a serializer with parsl"""
+    serializer = create_proxystore_serializer()
+    register_method_for_data(serializer)
+
+
+def create_proxystore_serializer() -> ProxyStoreSerializer:
+    """Creates a serializer but does not register with global system - so this
+    can be used in testing."""
+
+    import uuid
+    store = Store(name='parsl_store_' + str(uuid.uuid4()), connector=FileConnector(store_dir="/tmp"))
+    register_store(store)
+    return ProxyStoreSerializer(store)