Ensure that serializer can use the 'name' pickle protocol.

Pickle supports a protocol where __reduce__returns a string giving
the global name. Implementing this behaviour lets us serialize numpy
ufuncs. Also adjust installInflightFunctions to handle new load
behaviour, fix an instability caused by not leaving LoadedModule
objects in memory, and adjust alternative test.

Author: William Grant

typed_python/SerializationContext.py

@@ -31,11 +31,24 @@
 import types
 import traceback
 import logging
+import numpy
+import pickle
 
 
 _badModuleCache = set()
 
 
+def pickledByStr(module_name: str, name: str) -> None:
+    """Generate the object given the module_name and name.
+
+    This mimics pickle's behavior when given a string from __reduce__. The
+    string is interpreted as the name of a global variable, and pickle.whichmodules
+    is used to search the module namespace, generating module_name.
+    """
+    module = importlib.import_module(module_name)
+    return getattr(module, name)
+
+
 def createFunctionWithLocalsAndGlobals(code, globals):
     if globals is None:
         globals = {}
@@ -708,26 +721,30 @@ def walkCodeObject(code):
             return (createFunctionWithLocalsAndGlobals, args, representation)
 
         if not isinstance(inst, type) and hasattr(type(inst), '__reduce_ex__'):
-            res = inst.__reduce_ex__(4)
+            if isinstance(inst, numpy.ufunc):
+                res = inst.__name__
+            else:
+                res = inst.__reduce_ex__(4)
 
-            # pickle supports a protocol where __reduce__ can return a string
-            # giving a global name. We'll already find that separately, so we
-            # don't want to handle it here. We ought to look at this in more detail
-            # however
+            # mimic pickle's behaviour when a string is received.
             if isinstance(res, str):
-                return None
+                name_tuple = (inst, res)
+                module_name = pickle.whichmodule(*name_tuple)
+                res = (pickledByStr, (module_name, res,), pickledByStr)
 
             return res
 
         if not isinstance(inst, type) and hasattr(type(inst), '__reduce__'):
-            res = inst.__reduce__()
+            if isinstance(inst, numpy.ufunc):
+                res = inst.__name__
+            else:
+                res = inst.__reduce()
 
-            # pickle supports a protocol where __reduce__ can return a string
-            # giving a global name. We'll already find that separately, so we
-            # don't want to handle it here. We ought to look at this in more detail
-            # however
+            # mimic pickle's behaviour when a string is received.
             if isinstance(res, str):
-                return None
+                name_tuple = (inst, res)
+                module_name = pickle.whichmodule(*name_tuple)
+                res = (pickledByStr, (module_name, res,), pickledByStr)
 
             return res
 
@@ -736,6 +753,9 @@ def walkCodeObject(code):
     def setInstanceStateFromRepresentation(
         self, instance, representation=None, itemIt=None, kvPairIt=None, setStateFun=None
     ):
+        if representation is pickledByStr:
+            return
+
         if representation is reconstructTypeFunctionType:
             return
 

typed_python/compiler/global_variable_definition.py

@@ -87,5 +87,4 @@ def __eq__(self, other):
         return self.name == other.name and self.type == other.type and self.metadata == other.metadata
 
     def __str__(self):
-        metadata_str = str(self.metadata) if len(str(self.metadata)) < 100 else str(self.metadata)[:100] + "..."
-        return f"GlobalVariableDefinition(name={self.name}, type={self.type}, metadata={metadata_str})"
+        return f"GlobalVariableDefinition(name={self.name}, type={self.type}, metadata={pad(str(self.metadata))})"

typed_python/compiler/llvm_compiler_test.py

@@ -20,6 +20,8 @@
 from typed_python.compiler.module_definition import ModuleDefinition
 from typed_python.compiler.global_variable_definition import GlobalVariableMetadata
 
+from typed_python.test_util import evaluateExprInFreshProcess
+
 import pytest
 import ctypes
 
@@ -131,3 +133,28 @@ def test_create_binary_shared_object():
         pointers[0].set(5)
 
         assert loaded.functionPointers['__test_f_2']() == 5
+
+
+@pytest.mark.skipif('sys.platform=="darwin"')
+def test_loaded_modules_persist():
+    """
+    Make sure that loaded modules are persisted in the converter state.
+
+    We have to maintain these references to avoid surprise segfaults - if this test fails,
+    it should be because the GlobalVariableDefinition memory management has been refactored.
+    """
+
+    # compile a module
+    xmodule = "\n".join([
+        "@Entrypoint",
+        "def f(x):",
+        "    return x + 1",
+        "@Entrypoint",
+        "def g(x):",
+        "    return f(x) * 100",
+        "g(1000)",
+        "def get_loaded_modules():",
+        "    return len(Runtime.singleton().converter.loadedUncachedModules)"
+    ])
+    VERSION1 = {'x.py': xmodule}
+    assert evaluateExprInFreshProcess(VERSION1, 'x.get_loaded_modules()') == 1

typed_python/compiler/loaded_module.py

@@ -28,6 +28,8 @@ def __init__(self, functionPointers, serializedGlobalVariableDefinitions):
 
         self.functionPointers[ModuleDefinition.GET_GLOBAL_VARIABLES_NAME](self.pointers.pointerUnsafe(0))
 
+        self.installedGlobalVariableDefinitions = {}
+
     @staticmethod
     def validateGlobalVariables(serializedGlobalVariableDefinitions: Dict[str, bytes]) -> bool:
         """Check that each global variable definition is sensible.
@@ -83,6 +85,8 @@ def linkGlobalVariables(self, variable_names: List[str] = None) -> None:
 
             meta = SerializationContext().deserialize(self.orderedDefs[i]).metadata
 
+            self.installedGlobalVariableDefinitions[i] = meta
+
             if meta.matches.StringConstant:
                 self.pointers[i].cast(str).initialize(meta.value)
 

typed_python/compiler/python_to_native_converter.py

@@ -125,6 +125,12 @@ def __init__(self, llvmCompiler, compilerCache):
         self.llvmCompiler = llvmCompiler
         self.compilerCache = compilerCache
 
+        # all LoadedModule objects that we have created. We need to keep them alive so
+        # that any python metadata objects the've created stay alive as well. Ultimately, this
+        # may not be the place we put these objects (for instance, you could imagine a
+        # 'dummy' compiler cache or something). But for now, we need to keep them alive.
+        self.loadedUncachedModules = []
+
         # if True, then insert additional code to check for undefined behavior.
         self.generateDebugChecks = False
 
@@ -191,6 +197,7 @@ def buildAndLinkNewModule(self):
         if self.compilerCache is None:
             loadedModule = self.llvmCompiler.buildModule(targets)
             loadedModule.linkGlobalVariables()
+            self.loadedUncachedModules.append(loadedModule)
             return
 
         # get a set of function names that we depend on
@@ -926,7 +933,11 @@ def _installInflightFunctions(self):
             outboundTargets = []
             for outboundFuncId in self._dependencies.getNamesDependedOn(identifier):
                 name = self._link_name_for_identity[outboundFuncId]
-                outboundTargets.append(self._targets[name])
+                target = self.getTarget(name)
+                if target is not None:
+                    outboundTargets.append(target)
+                else:
+                    raise RuntimeError(f'dependency not found for {name}.')
 
             nativeFunction, actual_output_type = self._inflight_definitions.get(identifier)
 

typed_python/compiler/tests/numpy_interaction_test.py

@@ -1,4 +1,4 @@
-from typed_python import ListOf, Entrypoint
+from typed_python import ListOf, Entrypoint, SerializationContext
 import numpy
 import numpy.linalg
 
@@ -44,3 +44,12 @@ def test_listof_from_sliced_numpy_array():
     y = x[::2]
 
     assert ListOf(int)(y) == [0, 2]
+
+
+def test_can_serialize_numpy_ufunc():
+    assert numpy.sin == SerializationContext().deserialize(SerializationContext().serialize(numpy.sin))
+
+
+def test_can_serialize_numpy_array():
+    x = numpy.ones(10)
+    assert (x == SerializationContext().deserialize(SerializationContext().serialize(x))).all()

typed_python/compiler/tests/type_of_instances_compilation_test.py

@@ -17,13 +17,13 @@ def typeOfArg(x: C):
 
 def test_type_of_alternative_is_specific():
     for members in [{}, {'a': int}]:
-        A = Alternative("A", A=members)
+        Alt = Alternative("Alt", A=members)
 
         @Entrypoint
-        def typeOfArg(x: A):
+        def typeOfArg(x: Alt):
             return type(x)
 
-        assert typeOfArg(A.A()) is A.A
+        assert typeOfArg(Alt.A()) is Alt.A
 
 
 def test_type_of_concrete_alternative_is_specific():

typed_python/types_serialization_test.py

@@ -3061,3 +3061,27 @@ def f(self):
         print(x)
         # TODO: make this True
         # assert x[0].f.__closure__[0].cell_contents is x
+
+    def test_serialize_pyobj_with_custom_reduce(self):
+        class CustomReduceObject:
+            def __reduce__(self):
+                return 'CustomReduceObject'
+
+        assert CustomReduceObject == SerializationContext().deserialize(SerializationContext().serialize(CustomReduceObject))
+
+    def test_serialize_pyobj_in_MRTG_with_custom_reduce(self):
+        def getX():
+            class InnerCustomReduceObject:
+                def __reduce__(self):
+                    return 'InnerCustomReduceObject'
+
+                def f(self):
+                    return x
+
+            x = (InnerCustomReduceObject, InnerCustomReduceObject)
+
+            return x
+
+        x = callFunctionInFreshProcess(getX, (), showStdout=True)
+
+        assert x == SerializationContext().deserialize(SerializationContext().serialize(x))