Allow for partial module loads in compiler cache.

Previously we would always attempt to link and validate all global
variables when loading a module from the cache. This caused linking
errors, or validation errors, or deserialization errors, and meant we
needed the mark_invalid mechanism for handling modules with outdated
global variables. Here we add double-serialised global variables, and
only deserialize,link&validate the subset required for the function
required (and its dependencies). This requires the cache to store
a function and global_var dependency graph. Also add utility methods
for GlobalVariableDefinition.

Author: William Grant

typed_python/compiler/binary_shared_object.py

@@ -26,8 +26,8 @@
 
 
 class LoadedBinarySharedObject(LoadedModule):
-    def __init__(self, binarySharedObject, diskPath, functionPointers, globalVariableDefinitions):
-        super().__init__(functionPointers, globalVariableDefinitions)
+    def __init__(self, binarySharedObject, diskPath, functionPointers, serializedGlobalVariableDefinitions):
+        super().__init__(functionPointers, serializedGlobalVariableDefinitions)
 
         self.binarySharedObject = binarySharedObject
         self.diskPath = diskPath
@@ -36,30 +36,32 @@ def __init__(self, binarySharedObject, diskPath, functionPointers, globalVariabl
 class BinarySharedObject:
     """Models a shared object library (.so) loadable on linux systems."""
 
-    def __init__(self, binaryForm, functionTypes, globalVariableDefinitions):
+    def __init__(self, binaryForm, functionTypes, serializedGlobalVariableDefinitions, globalDependencies):
         """
         Args:
-            binaryForm - a bytes object containing the actual compiled code for the module
-            globalVariableDefinitions - a map from name to GlobalVariableDefinition
+            binaryForm: a bytes object containing the actual compiled code for the module
+            serializedGlobalVariableDefinitions: a map from name to GlobalVariableDefinition
+            globalDependencies: a dict from function linkname to the list of global variables it depends on
         """
         self.binaryForm = binaryForm
         self.functionTypes = functionTypes
-        self.globalVariableDefinitions = globalVariableDefinitions
+        self.serializedGlobalVariableDefinitions = serializedGlobalVariableDefinitions
+        self.globalDependencies = globalDependencies
         self.hash = sha_hash(binaryForm)
 
     @property
     def definedSymbols(self):
         return self.functionTypes.keys()
 
     @staticmethod
-    def fromDisk(path, globalVariableDefinitions, functionNameToType):
+    def fromDisk(path, serializedGlobalVariableDefinitions, functionNameToType, globalDependencies):
         with open(path, "rb") as f:
             binaryForm = f.read()
 
-        return BinarySharedObject(binaryForm, functionNameToType, globalVariableDefinitions)
+        return BinarySharedObject(binaryForm, functionNameToType, serializedGlobalVariableDefinitions, globalDependencies)
 
     @staticmethod
-    def fromModule(module, globalVariableDefinitions, functionNameToType):
+    def fromModule(module, serializedGlobalVariableDefinitions, functionNameToType, globalDependencies):
         target_triple = llvm.get_process_triple()
         target = llvm.Target.from_triple(target_triple)
         target_machine_shared_object = target.create_target_machine(reloc='pic', codemodel='default')
@@ -80,7 +82,7 @@ def fromModule(module, globalVariableDefinitions, functionNameToType):
             )
 
             with open(os.path.join(tf, "module.so"), "rb") as so_file:
-                return BinarySharedObject(so_file.read(), functionNameToType, globalVariableDefinitions)
+                return BinarySharedObject(so_file.read(), functionNameToType, serializedGlobalVariableDefinitions, globalDependencies)
 
     def load(self, storageDir):
         """Instantiate this .so in temporary storage and return a dict from symbol -> integer function pointer"""
@@ -127,8 +129,7 @@ def loadFromPath(self, modulePath):
             self,
             modulePath,
             functionPointers,
-            self.globalVariableDefinitions
+            self.serializedGlobalVariableDefinitions
         )
-        loadedModule.linkGlobalVariables()
 
         return loadedModule

typed_python/compiler/compiler_cache.py

@@ -15,9 +15,12 @@
 import os
 import uuid
 import shutil
-from typed_python.compiler.loaded_module import LoadedModule
-from typed_python.compiler.binary_shared_object import BinarySharedObject
 
+from typing import Optional, List
+
+from typed_python.compiler.binary_shared_object import LoadedBinarySharedObject, BinarySharedObject
+from typed_python.compiler.directed_graph import DirectedGraph
+from typed_python.compiler.typed_call_target import TypedCallTarget
 from typed_python.SerializationContext import SerializationContext
 from typed_python import Dict, ListOf
 
@@ -52,146 +55,173 @@ def __init__(self, cacheDir):
 
         ensureDirExists(cacheDir)
 
-        self.loadedModules = Dict(str, LoadedModule)()
+        self.loadedBinarySharedObjects = Dict(str, LoadedBinarySharedObject)()
         self.nameToModuleHash = Dict(str, str)()
 
-        self.modulesMarkedValid = set()
-        self.modulesMarkedInvalid = set()
+        self.moduleManifestsLoaded = set()
 
         for moduleHash in os.listdir(self.cacheDir):
             if len(moduleHash) == 40:
                 self.loadNameManifestFromStoredModuleByHash(moduleHash)
 
-        self.targetsLoaded = {}
+        # the set of functions with an associated module in loadedBinarySharedObjects
+        self.targetsLoaded: Dict[str, TypedCallTarget] = {}
 
-    def hasSymbol(self, linkName):
-        return linkName in self.nameToModuleHash
+        # the set of functions with linked and validated globals (i.e. ready to be run).
+        self.targetsValidated = set()
 
-    def getTarget(self, linkName):
-        assert self.hasSymbol(linkName)
+        self.function_dependency_graph = DirectedGraph()
+        # dict from function linkname to list of global names (should be llvm keys in serialisedGlobalDefinitions)
+        self.global_dependencies = Dict(str, ListOf(str))()
 
-        self.loadForSymbol(linkName)
+    def hasSymbol(self, linkName: str) -> bool:
+        """NB this will return True even if the linkName is ultimately unretrievable."""
+        return linkName in self.nameToModuleHash
 
+    def getTarget(self, linkName: str) -> TypedCallTarget:
+        if not self.hasSymbol(linkName):
+            raise ValueError(f'symbol not found for linkName {linkName}')
+        self.loadForSymbol(linkName)
         return self.targetsLoaded[linkName]
 
-    def markModuleHashInvalid(self, hashstr):
-        with open(os.path.join(self.cacheDir, hashstr, "marked_invalid"), "w"):
-            pass
+    def dependencies(self, linkName: str) -> Optional[List[str]]:
+        """Returns all the function names that `linkName` depends on"""
+        return list(self.function_dependency_graph.outgoing(linkName))
 
-    def loadForSymbol(self, linkName):
+    def loadForSymbol(self, linkName: str) -> None:
+        """Loads the whole module, and any submodules, into LoadedBinarySharedObjects"""
         moduleHash = self.nameToModuleHash[linkName]
 
         self.loadModuleByHash(moduleHash)
 
-    def loadModuleByHash(self, moduleHash):
+        if linkName not in self.targetsValidated:
+            dependantFuncs = self.dependencies(linkName) + [linkName]
+            globalsToLink = {}  # dict from modulehash to list of globals.
+            for funcName in dependantFuncs:
+                if funcName not in self.targetsValidated:
+                    funcModuleHash = self.nameToModuleHash[funcName]
+                    # append to the list of globals to link for a given module.  TODO: optimise this, don't double-link.
+                    globalsToLink[funcModuleHash] = globalsToLink.get(funcModuleHash, []) + self.global_dependencies.get(funcName, [])
+
+            for moduleHash, globs in globalsToLink.items():  # this works because loadModuleByHash loads submodules too.
+                if globs:
+                    definitionsToLink = {x: self.loadedBinarySharedObjects[moduleHash].serializedGlobalVariableDefinitions[x]
+                                         for x in globs
+                                         }
+                    self.loadedBinarySharedObjects[moduleHash].linkGlobalVariables(definitionsToLink)
+                    if not self.loadedBinarySharedObjects[moduleHash].validateGlobalVariables(definitionsToLink):
+                        raise RuntimeError('failed to validate globals when loading:', linkName)
+
+            self.targetsValidated.update(dependantFuncs)
+
+    def loadModuleByHash(self, moduleHash: str) -> None:
         """Load a module by name.
 
         As we load, place all the newly imported typed call targets into
         'nameToTypedCallTarget' so that the rest of the system knows what functions
         have been uncovered.
         """
-        if moduleHash in self.loadedModules:
-            return True
+        if moduleHash in self.loadedBinarySharedObjects:
+            return
 
         targetDir = os.path.join(self.cacheDir, moduleHash)
 
-        try:
-            with open(os.path.join(targetDir, "type_manifest.dat"), "rb") as f:
-                callTargets = SerializationContext().deserialize(f.read())
+        # TODO (Will) - store these names as module consts, use one .dat only
+        with open(os.path.join(targetDir, "type_manifest.dat"), "rb") as f:
+            callTargets = SerializationContext().deserialize(f.read())
 
-            with open(os.path.join(targetDir, "globals_manifest.dat"), "rb") as f:
-                globalVarDefs = SerializationContext().deserialize(f.read())
+        with open(os.path.join(targetDir, "globals_manifest.dat"), "rb") as f:
+            serializedGlobalVarDefs = SerializationContext().deserialize(f.read())
 
-            with open(os.path.join(targetDir, "native_type_manifest.dat"), "rb") as f:
-                functionNameToNativeType = SerializationContext().deserialize(f.read())
+        with open(os.path.join(targetDir, "native_type_manifest.dat"), "rb") as f:
+            functionNameToNativeType = SerializationContext().deserialize(f.read())
+
+        with open(os.path.join(targetDir, "submodules.dat"), "rb") as f:
+            submodules = SerializationContext().deserialize(f.read(), ListOf(str))
 
-            with open(os.path.join(targetDir, "submodules.dat"), "rb") as f:
-                submodules = SerializationContext().deserialize(f.read(), ListOf(str))
-        except Exception:
-            self.markModuleHashInvalid(moduleHash)
-            return False
+        with open(os.path.join(targetDir, "function_dependencies.dat"), "rb") as f:
+            dependency_edgelist = SerializationContext().deserialize(f.read())
 
-        if not LoadedModule.validateGlobalVariables(globalVarDefs):
-            self.markModuleHashInvalid(moduleHash)
-            return False
+        with open(os.path.join(targetDir, "global_dependencies.dat"), "rb") as f:
+            globalDependencies = SerializationContext().deserialize(f.read())
 
         # load the submodules first
         for submodule in submodules:
-            if not self.loadModuleByHash(submodule):
-                return False
+            self.loadModuleByHash(submodule)
 
         modulePath = os.path.join(targetDir, "module.so")
 
         loaded = BinarySharedObject.fromDisk(
             modulePath,
-            globalVarDefs,
-            functionNameToNativeType
+            serializedGlobalVarDefs,
+            functionNameToNativeType,
+            globalDependencies
+
         ).loadFromPath(modulePath)
 
-        self.loadedModules[moduleHash] = loaded
+        self.loadedBinarySharedObjects[moduleHash] = loaded
 
         self.targetsLoaded.update(callTargets)
 
-        return True
+        assert not any(key in self.global_dependencies for key in globalDependencies)  # should only happen if there's a hash collision.
+        self.global_dependencies.update(globalDependencies)
+
+        # update the cache's dependency graph with our new edges.
+        for function_name, dependant_function_name in dependency_edgelist:
+            self.function_dependency_graph.addEdge(source=function_name, dest=dependant_function_name)
 
-    def addModule(self, binarySharedObject, nameToTypedCallTarget, linkDependencies):
+    def addModule(self, binarySharedObject, nameToTypedCallTarget, linkDependencies, dependencyEdgelist):
         """Add new code to the compiler cache.
 
         Args:
-            binarySharedObject - a BinarySharedObject containing the actual assembler
-                we've compiled
-            nameToTypedCallTarget - a dict from linkname to TypedCallTarget telling us
-                the formal python types for all the objects
-            linkDependencies - a set of linknames we depend on directly.
+            binarySharedObject: a BinarySharedObject containing the actual assembler
+                we've compiled.
+            nameToTypedCallTarget: a dict from linkname to TypedCallTarget telling us
+                the formal python types for all the objects.
+            linkDependencies: a set of linknames we depend on directly.
+            dependencyEdgelist (list): a list of source, dest pairs giving the set of dependency graph for the
+                module.
         """
         dependentHashes = set()
 
         for name in linkDependencies:
             dependentHashes.add(self.nameToModuleHash[name])
 
-        path, hashToUse = self.writeModuleToDisk(binarySharedObject, nameToTypedCallTarget, dependentHashes)
+        path, hashToUse = self.writeModuleToDisk(binarySharedObject, nameToTypedCallTarget, dependentHashes, dependencyEdgelist)
 
-        self.loadedModules[hashToUse] = (
+        self.loadedBinarySharedObjects[hashToUse] = (
             binarySharedObject.loadFromPath(os.path.join(path, "module.so"))
         )
 
         for n in binarySharedObject.definedSymbols:
             self.nameToModuleHash[n] = hashToUse
 
-    def loadNameManifestFromStoredModuleByHash(self, moduleHash):
-        if moduleHash in self.modulesMarkedValid:
-            return True
-
-        targetDir = os.path.join(self.cacheDir, moduleHash)
+        # link & validate all globals for the new module
+        self.loadedBinarySharedObjects[hashToUse].linkGlobalVariables()
+        if not self.loadedBinarySharedObjects[hashToUse].validateGlobalVariables(
+                self.loadedBinarySharedObjects[hashToUse].serializedGlobalVariableDefinitions):
+            raise RuntimeError('failed to validate globals in new module:', hashToUse)
 
-        # ignore 'marked invalid'
-        if os.path.exists(os.path.join(targetDir, "marked_invalid")):
-            # just bail - don't try to read it now
+    def loadNameManifestFromStoredModuleByHash(self, moduleHash) -> None:
+        if moduleHash in self.moduleManifestsLoaded:
+            return
 
-            # for the moment, we don't try to clean up the cache, because
-            # we can't be sure that some process is not still reading the
-            # old files.
-            self.modulesMarkedInvalid.add(moduleHash)
-            return False
+        targetDir = os.path.join(self.cacheDir, moduleHash)
 
         with open(os.path.join(targetDir, "submodules.dat"), "rb") as f:
             submodules = SerializationContext().deserialize(f.read(), ListOf(str))
 
         for subHash in submodules:
-            if not self.loadNameManifestFromStoredModuleByHash(subHash):
-                self.markModuleHashInvalid(subHash)
-                return False
+            self.loadNameManifestFromStoredModuleByHash(subHash)
 
         with open(os.path.join(targetDir, "name_manifest.dat"), "rb") as f:
             self.nameToModuleHash.update(
                 SerializationContext().deserialize(f.read(), Dict(str, str))
             )
 
-        self.modulesMarkedValid.add(moduleHash)
-
-        return True
+        self.moduleManifestsLoaded.add(moduleHash)
 
-    def writeModuleToDisk(self, binarySharedObject, nameToTypedCallTarget, submodules):
+    def writeModuleToDisk(self, binarySharedObject, nameToTypedCallTarget, submodules, dependencyEdgelist):
         """Write out a disk representation of this module.
 
         This includes writing both the shared object, a manifest of the function names
@@ -244,11 +274,17 @@ def writeModuleToDisk(self, binarySharedObject, nameToTypedCallTarget, submodule
 
         # write the type manifest
         with open(os.path.join(tempTargetDir, "globals_manifest.dat"), "wb") as f:
-            f.write(SerializationContext().serialize(binarySharedObject.globalVariableDefinitions))
+            f.write(SerializationContext().serialize(binarySharedObject.serializedGlobalVariableDefinitions))
 
         with open(os.path.join(tempTargetDir, "submodules.dat"), "wb") as f:
             f.write(SerializationContext().serialize(ListOf(str)(submodules), ListOf(str)))
 
+        with open(os.path.join(tempTargetDir, "function_dependencies.dat"), "wb") as f:
+            f.write(SerializationContext().serialize(dependencyEdgelist))  # might need a listof
+
+        with open(os.path.join(tempTargetDir, "global_dependencies.dat"), "wb") as f:
+            f.write(SerializationContext().serialize(binarySharedObject.globalDependencies))
+
         try:
             os.rename(tempTargetDir, targetDir)
         except IOError:
@@ -264,7 +300,7 @@ def function_pointer_by_name(self, linkName):
         if moduleHash is None:
             raise Exception("Can't find a module for " + linkName)
 
-        if moduleHash not in self.loadedModules:
+        if moduleHash not in self.loadedBinarySharedObjects:
             self.loadForSymbol(linkName)
 
-        return self.loadedModules[moduleHash].functionPointers[linkName]
+        return self.loadedBinarySharedObjects[moduleHash].functionPointers[linkName]