Repository: DataBiosphere/toil

Issue:
_switchOutPromisedTargetReturnValues is too intrusive
See the FIXMEs in https://github.com/BD2KGenomics/jobTree/commit/5953d890cbc46fca486f65028b8cd04a1425d8bd



Discussion:
Benedict mentioned to Hannes that custom unpickling would be the right way to approach this, and Hannes agrees.


Resolving patch:
diff --git a/src/toil/common.py b/src/toil/common.py
--- a/src/toil/common.py
+++ b/src/toil/common.py
@@ -145,7 +145,7 @@ def f(x):
         setOption("maxDisk", h2b, iC(1))
         
         #Retrying/rescuing jobs
-        setOption("retryCount", int, iC(1))
+        setOption("retryCount", int, iC(0))
         setOption("maxJobDuration", int, iC(1))
         setOption("rescueJobsFrequency", int, iC(1))
         
diff --git a/src/toil/job.py b/src/toil/job.py
--- a/src/toil/job.py
+++ b/src/toil/job.py
@@ -29,7 +29,9 @@
 import tempfile
 import uuid
 import time
+import copy_reg
 from toil.resource import ModuleDescriptor
+from toil.common import loadJobStore
 from bd2k.util.humanize import human2bytes
 
 try:
@@ -96,11 +98,7 @@ def run(self, fileStore):
         be used to create temporary files which can be shared between jobs.
         
         The return values of the function can be passed to other jobs
-        by means of the rv() function. 
-        
-        If the return value is a tuple, rv(i) would refer to the ith (indexed from 0)
-        member of the tuple. If the return value is not a tuple then rV(0) or rV() would
-        refer to the return value of the function. 
+        by means of the rv() function.  
         
         Note: We disallow return values to be PromisedJobReturnValue instances
         (generated by the Job.rv() function - see below).
@@ -215,15 +213,20 @@ def encapsulate(self):
     #job run functions
     ####################################################
     
-    def rv(self, argIndex=0):
+    def rv(self, argIndex=None):
         """
         Gets a PromisedJobReturnValue, representing the argIndex return
         value of the run function (see run method for description).
         This PromisedJobReturnValue, if a class attribute of a Job instance,
-        call it T, will be replaced by the actual return value just before the 
-        run function of T is called. The function rv therefore allows the output 
+        call it T, will be replaced by the actual return value when the 
+        T is loaded. The function rv therefore allows the output 
         from one Job to be wired as input to another Job before either
         is actually run.  
+        
+        :param argIndex: If None the complete return value will be returned, if argIndex
+        is an integer it is used to refer to the return value as indexable 
+        (tuple/list/dictionary, or in general object that implements __getitem__), 
+        hence rv(i) would refer to the ith (indexed from 0) member of return value.
         """
         #Check if the return value has already been promised and if it has
         #return it
@@ -231,7 +234,7 @@ def rv(self, argIndex=0):
             return self._rvs[argIndex]
         #Create, store, return new PromisedJobReturnValue
         self._rvs[argIndex] = PromisedJobReturnValue()
-        return self._rvs[argIndex]
+        return self._rvs[argIndex] 
     
     ####################################################
     #Cycle/connectivity checking
@@ -661,67 +664,25 @@ def _loadRootJob(jobStore):
     #Functions to pass Job.run return values to the
     #input arguments of other Job instances
     ####################################################
-   
-    def _switchOutPromisedJobReturnValues(self, jobStore):
-        """
-        Replaces each PromisedJobReturnValue instance that is a class
-        attribute of the job with PromisedJobReturnValue's stored value.
-        Will do this also for PromisedJobReturnValue instances within lists,
-        tuples, sets or dictionaries that are class attributes of the Job.
-        
-        This function is called just before the run method.
-        """
-
-        # FIXME: why is the substitution only done one level deep? IOW, what about, say, a dict of lists
-
-        # FIXME: This replaces subclasses of list, tuple and dict with the instances of the respective base ...
-        # FIXME: ... class. For a potential fix, see my quick fix for tuples.
-
-        # FIXME: This unnecessarily touches attributes that don't have a promise. It touches attributes of the ...
-        # FIXME: ... job base classes which provably have no PromisedJobReturnValues in them
-
-        #Iterate on the class attributes of the Job instance.
-        for attr, value in self.__dict__.iteritems():
-            #If the variable is a PromisedJobReturnValue replace with the
-            #actual stored return value of the PromisedJobReturnValue
-            #else if the variable is a list, tuple or set or dict replace any 
-            #PromisedJobReturnValue instances within
-            #the container with the stored return value.
-
-            # FIXME: This lambda might become more readable if "value" wasn't closed over but passed in.
-
-            # FIXME: I find lambdas awkward. A simple nested def should suffice
-
-            f = lambda : map(lambda x : x.loadValue(jobStore) if
-                        isinstance(x, PromisedJobReturnValue) else x, value)
-            if isinstance(value, PromisedJobReturnValue):
-                self.__dict__[attr] = value.loadValue(jobStore)
-            elif isinstance(value, list):
-                self.__dict__[attr] = f()
-            elif value.__class__ == tuple:
-                self.__dict__[attr] = tuple(f())
-            elif isinstance(value, tuple):
-                self.__dict__[attr] = value.__class__(*f())
-            elif isinstance(value, set):
-                self.__dict__[attr] = set(f())
-            elif isinstance(value, dict):
-                self.__dict__[attr] = dict(map(lambda x : (x, value[x].loadValue(jobStore) if
-                        isinstance(x, PromisedJobReturnValue) else value[x]), value))
       
     def _setFileIDsForPromisedValues(self, jobStore, jobStoreID, visited):
         """
         Sets the jobStoreFileID for each PromisedJobReturnValue in the
         graph of jobs created.
         """
-        #Replace any None references with valid jobStoreFileIDs. We
+        #Replace None references with valid jobStoreFileIDs. We
         #do this here, rather than within the original constructor of the
         #promised value because we don't necessarily have access to the jobStore when
         #the PromisedJobReturnValue instances are created.
         if self not in visited:
             visited.add(self)
-            for PromisedJobReturnValue in self._rvs.values():
-                if PromisedJobReturnValue.jobStoreFileID == None:
-                    PromisedJobReturnValue.jobStoreFileID = jobStore.getEmptyFileStoreID(jobStoreID)
+            for i in self._rvs.keys():
+                promisedJobReturnValue = self._rvs[i]
+                #Instances of PromisedJobReturnValue are replaced with jobStoreFileIDs
+                if isinstance(promisedJobReturnValue, PromisedJobReturnValue):
+                    promisedJobReturnValue.jobStoreFileID = jobStore.getEmptyFileStoreID(jobStoreID)
+                    promisedJobReturnValue.jobStoreString = jobStore.config.jobStore
+                    self._rvs[i] = promisedJobReturnValue.jobStoreFileID
             #Now recursively do the same for the children and follow ons.
             for successorJob in self._children + self._followOns + self._services:
                 successorJob._setFileIDsForPromisedValues(jobStore, jobStoreID, visited)
@@ -733,14 +694,14 @@ def _setReturnValuesForPromises(job, returnValues, jobStore):
         run function.
         """
         for i in job._rvs.keys():
-            if isinstance(returnValues, tuple):
-                argToStore = returnValues[i]
-            else:
-                if i != 0:
-                    raise RuntimeError("Referencing return value index (%s)"
-                                " that is out of range: %s" % (i, returnValues))
+            if i == None:
                 argToStore = returnValues
-            job._rvs[i]._storeValue(argToStore, jobStore)
+            else:
+                argToStore = returnValues[i]
+            with jobStore.updateFileStream(job._rvs[i]) as fileHandle:
+                if isinstance(argToStore, PromisedJobReturnValue):
+                    raise RuntimeError("A nested PromisedJobReturnValue has been found.") #We do not allow the return of PromisedJobReturnValue instance from the run function
+                cPickle.dump(argToStore, fileHandle, cPickle.HIGHEST_PROTOCOL)
     
     ####################################################
     #Functions associated with Job.checkJobGraphAcyclic to establish
@@ -853,8 +814,6 @@ def _execute(self, jobWrapper, stats, localTempDir, jobStore):
             startClock = getTotalCpuTime()
         
         baseDir = os.getcwd()
-        #Switch out any promised return value instances with the actual values
-        self._switchOutPromisedJobReturnValues(jobStore)
         #Run the job, first cleanup then run.
         fileStore = Job.FileStore(jobStore, jobWrapper, localTempDir)
         returnValues = self.run(fileStore)
@@ -920,7 +879,8 @@ class FunctionWrappingJob(Job):
     """
     Job used to wrap a function.
     
-    Function can not be nested function or class function, currently.
+    If dill is installed 
+    Function can be nested function class function, currently.
     *args and **kwargs are used as the arguments to the function.
     """
     def __init__(self, userFunction, *args, **kwargs):
@@ -929,12 +889,17 @@ def __init__(self, userFunction, *args, **kwargs):
         disk = kwargs.pop("disk") if "disk" in kwargs else sys.maxint
         memory = kwargs.pop("memory") if "memory" in kwargs else sys.maxint
         Job.__init__(self, memory=memory, cpu=cpu, disk=disk)
+        #If dill is installed pickle the user function directly
+        
+        #else use indirect method
         self.userFunctionModule = ModuleDescriptor.forModule(userFunction.__module__).globalize()
         self.userFunctionName = str(userFunction.__name__)
         self._args=args
         self._kwargs=kwargs
         
     def _getUserFunction(self):
+        #If dill is installed unpickle the user function directly
+        
         userFunctionModule = self.userFunctionModule.localize()
         if userFunctionModule.dirPath not in sys.path:
             # FIXME: prepending to sys.path will probably fix #103
@@ -1059,40 +1024,47 @@ def addFollowOn(self, followOnJob):
     def rv(self, argIndex=0):
         return self.followOn.rv(argIndex)
 
-class PromisedJobReturnValue():
+class PromisedJobReturnValue(object):
     """
     References a return value from a Job's run function. Let T be a job.
     Instances of PromisedJobReturnValue are created by
-    T.rv(i), where i is an integer reference to a return value of T's run function
-    (casting the return value as a tuple). 
-    When passed to the constructor of a different Job the PromisedJobReturnValue
-    will be replaced by the actual referenced return value after the Job's run function
-    has finished (see Job._switchOutPromisedJobReturnValues).
+    T.rv(), which is used to reference the return value of T's run function.
+    When passed to the constructor of a different, successor Job the PromisedJobReturnValue
+    will be replaced by the actual referenced return value.
     This mechanism allows a return values from one Job's run method to be input
     argument to Job before the former Job's run function has been executed.
     """ 
     def __init__(self):
-        self.jobStoreFileID = None #The None value is
-        #replaced with a real jobStoreFileID by the Job object.
-        
-    def loadValue(self, jobStore):
-        """
-        Unpickles the promised value and returns it. 
-        """
-        assert self.jobStoreFileID != None
-        with jobStore.readFileStream(self.jobStoreFileID) as fileHandle:
-            value = cPickle.load(fileHandle) #If this doesn't work then the file containing the promise may not exist or be corrupted.
-            if isinstance(value, PromisedJobReturnValue):
-                raise RuntimeError("A nested PromisedJobReturnValue has been found.") #We do not allow the return of PromisedJobReturnValue instance from the run function
-            return value
+        #The None values are replaced with a actual values in Job._setFileIDsForPromisedValues
+        self.jobStoreFileID = None #A file containing the actual pickled return value
+        self.jobStoreString = None #This is the string used to load the jobStore. 
+            
+def promisedJobReturnValuePickleFunction(promise):
+    """
+    This function and promisedJobReturnValueUnpickleFunction are used as custom pickle/unpickle 
+    functions to ensure that when the PromisedJobReturnValue instance p is unpickled it is replaced with 
+    the object pickled in p.jobStoreFileID
+    """
+    return promisedJobReturnValueUnpickleFunction, (promise.jobStoreString, promise.jobStoreFileID)
 
-    def _storeValue(self, valueToStore, jobStore):
-        """
-        Pickle the promised value. This is done by the job.
-        """
-        assert self.jobStoreFileID != None
-        with jobStore.updateFileStream(self.jobStoreFileID) as fileHandle:
-            cPickle.dump(valueToStore, fileHandle, cPickle.HIGHEST_PROTOCOL)
+def promisedJobReturnValueUnpickleFunction(jobStoreString, jobStoreFileID):
+    """The PromisedJobReturnValue custom unpickle function.
+    """
+    #If the attributes jobStoreFileID and jobStoreString are None then we return
+    #a new empty PromisedJobReturnValue (this can happen if Promise is serialised between its creation
+    #and the running of Job._setFileIDsForPromisedValues.
+    if jobStoreString == None:
+        assert jobStoreFileID == None
+        return PromisedJobReturnValue()
+    jobStore = loadJobStore(jobStoreString)
+    with jobStore.readFileStream(jobStoreFileID) as fileHandle:
+        value = cPickle.load(fileHandle) #If this doesn't work then the file containing the promise may not exist or be corrupted.
+        return value
+
+#This sets up the custom magic for pickling/unpickling a PromisedJobReturnValue
+copy_reg.pickle(PromisedJobReturnValue, 
+                promisedJobReturnValuePickleFunction, 
+                promisedJobReturnValueUnpickleFunction)
 
 def deleteFileStoreIDs(job, jobStoreFileIDsToDelete):
     """


Test patch:
diff --git a/src/toil/test/sort/sort.py b/src/toil/test/sort/sort.py
--- a/src/toil/test/sort/sort.py
+++ b/src/toil/test/sort/sort.py
@@ -40,13 +40,13 @@ def down(job, inputFile, fileStart, fileEnd, N, outputFileStoreID):
         tempFileStoreID1 = job.fileStore.getEmptyFileStoreID()
         tempFileStoreID2 = job.fileStore.getEmptyFileStoreID()
         #The use of rv here is for testing purposes
-        #The rv(0) of the first child job is tempFileStoreID1,
-        #similarly rv(0) of the second child is tempFileStoreID2
+        #The rv() of the first child job is tempFileStoreID1,
+        #similarly rv() of the second child is tempFileStoreID2
         job.addFollowOnJobFn(up,
                                    job.addChildJobFn(down, inputFile, fileStart,
-                                                           midPoint+1, N, tempFileStoreID1).rv(0),
+                                                           midPoint+1, N, tempFileStoreID1).rv(),
                                    job.addChildJobFn(down, inputFile, midPoint+1,
-                                                           fileEnd, N, tempFileStoreID2).rv(0), #Add one to avoid the newline
+                                                           fileEnd, N, tempFileStoreID2).rv(), #Add one to avoid the newline
                                    outputFileStoreID)                
     else:
         #We can sort this bit of the file
diff --git a/src/toil/test/src/jobTest.py b/src/toil/test/src/jobTest.py
--- a/src/toil/test/src/jobTest.py
+++ b/src/toil/test/src/jobTest.py
@@ -30,11 +30,11 @@ def testStatic(self):
 
             # Create the jobs
             A = Job.wrapFn(f, "A", outFile)
-            B = Job.wrapFn(f, A.rv(0), outFile)
-            C = Job.wrapFn(f, B.rv(0), outFile)
-            D = Job.wrapFn(f, C.rv(0), outFile)
-            E = Job.wrapFn(f, D.rv(0), outFile)
-            F = Job.wrapFn(f, E.rv(0), outFile)
+            B = Job.wrapFn(f, A.rv(), outFile)
+            C = Job.wrapFn(f, B.rv(), outFile)
+            D = Job.wrapFn(f, C.rv(), outFile)
+            E = Job.wrapFn(f, D.rv(), outFile)
+            F = Job.wrapFn(f, E.rv(), outFile)
 
             # Connect them into a workflow
             A.addChild(B)