Repository: DataBiosphere/toil

Issue:
Encapsulating root jobs always returns None
… independently of whether the encapsulated job returns a value.


Discussion:

Resolving patch:
diff --git a/src/toil/common.py b/src/toil/common.py
--- a/src/toil/common.py
+++ b/src/toil/common.py
@@ -595,22 +595,21 @@ def start(self, rootJob):
             self._serialiseEnv()
             self._cacheAllJobs()
 
-            # Make a file to store the root job's return value in
-            rootJobReturnValueID = self._jobStore.getEmptyFileStoreID()
-
-            # Add the root job return value as a promise
-            rootJob._rvs[()].append(rootJobReturnValueID)
-
-            # Write the name of the promise file in a shared file
-            with self._jobStore.writeSharedFileStream("rootJobReturnValue") as fH:
-                fH.write(rootJobReturnValueID)
+            # Pickle the promised return value of the root job, then write the pickled promise to
+            # a shared file, where we can find and unpickle it at the end of the workflow.
+            # Unpickling the promise will automatically substitute the promise for the actual
+            # return value.
+            with self._jobStore.writeSharedFileStream('rootJobReturnValue') as fH:
+                rootJob.prepareForPromiseRegistration(self._jobStore)
+                promise = rootJob.rv()
+                cPickle.dump(promise, fH)
 
             # Setup the first wrapper and cache it
-            job = rootJob._serialiseFirstJob(self._jobStore)
-            self._cacheJob(job)
+            rootJobWrapper = rootJob._serialiseFirstJob(self._jobStore)
+            self._cacheJob(rootJobWrapper)
 
             self._setProvisioner()
-            return self._runMainLoop(job)
+            return self._runMainLoop(rootJobWrapper)
         finally:
             self._shutdownBatchSystem()
 
@@ -633,8 +632,8 @@ def restart(self):
             self._serialiseEnv()
             self._cacheAllJobs()
             self._setProvisioner()
-            rootJob = self._jobStore.clean(jobCache=self._jobCache)
-            return self._runMainLoop(rootJob)
+            rootJobWrapper = self._jobStore.clean(jobCache=self._jobCache)
+            return self._runMainLoop(rootJobWrapper)
         finally:
             self._shutdownBatchSystem()
 
diff --git a/src/toil/job.py b/src/toil/job.py
--- a/src/toil/job.py
+++ b/src/toil/job.py
@@ -429,7 +429,7 @@ def rv(self, *path):
         """
         return Promise(self, path)
 
-    def allocatePromiseFile(self, path):
+    def registerPromise(self, path):
         if self._promiseJobStore is None:
             raise RuntimeError('Trying to pass a promise from a promising job that is not a '
                                'predecessor of the job receiving the promise')
@@ -437,6 +437,18 @@ def allocatePromiseFile(self, path):
         self._rvs[path].append(jobStoreFileID)
         return self._promiseJobStore.config.jobStore, jobStoreFileID
 
+    def prepareForPromiseRegistration(self, jobStore):
+        """
+        Ensure that a promise by this job (the promissor) can register with the promissor when
+        another job referring to the promise (the promissee) is being serialized. The promissee
+        holds the reference to the promise (usually as part of the the job arguments) and when it
+        is being pickled, so will the promises it refers to. Pickling a promise triggers it to be
+        registered with the promissor.
+
+        :return:
+        """
+        self._promiseJobStore = jobStore
+
     ####################################################
     #Cycle/connectivity checking
     ####################################################
@@ -1046,9 +1058,9 @@ def _serialiseJobGraph(self, jobWrapper, jobStore, returnValues, firstJob):
 
         # Temporarily set the jobStore locators for the promise call back functions
         for job in ordering:
-            job._promiseJobStore = jobStore
+            job.prepareForPromiseRegistration(jobStore)
             def setForServices(serviceJob):
-                serviceJob._promiseJobStore = jobStore
+                serviceJob.prepareForPromiseRegistration(jobStore)
                 for childServiceJob in serviceJob.service._childServices:
                     setForServices(childServiceJob)
             for serviceJob in self._services:
@@ -1344,15 +1356,16 @@ class EncapsulatedJob(Job):
     Let A be the root job of a job subgraph and B be another job we'd like to run after A
     and all its successors have completed, for this use encapsulate::
 
-        A, B = A(), B() #Job A and subgraph, Job B
+        #  Job A and subgraph, Job B
+        A, B = A(), B()
         A' = A.encapsulate()
-        A'.addChild(B) #B will run after A and all its successors have
-        # completed, A and its subgraph of successors in effect appear
-        # to be just one job.
+        A'.addChild(B)
+        #  B will run after A and all its successors have completed, A and its subgraph of
+        # successors in effect appear to be just one job.
 
-    The return value of an encapsulatd job (as accessed by the :func:`toil.job.Job.rv` function) \
-    is the return value of the root job, e.g. A().encapsulate().rv() and A().rv() \
-    will resolve to the same value after A or A.encapsulate() has been run.
+    The return value of an encapsulatd job (as accessed by the :func:`toil.job.Job.rv` function)
+    is the return value of the root job, e.g. A().encapsulate().rv() and A().rv() will resolve to
+    the same value after A or A.encapsulate() has been run.
     """
     def __init__(self, job):
         """
@@ -1369,8 +1382,8 @@ def __init__(self, job):
     def addChild(self, childJob):
         return Job.addChild(self.encapsulatedFollowOn, childJob)
 
-    def addService(self, service):
-        return Job.addService(self.encapsulatedFollowOn, service)
+    def addService(self, service, parentService=None):
+        return Job.addService(self.encapsulatedFollowOn, service, parentService=parentService)
 
     def addFollowOn(self, followOnJob):
         return Job.addFollowOn(self.encapsulatedFollowOn, followOnJob)
@@ -1378,6 +1391,13 @@ def addFollowOn(self, followOnJob):
     def rv(self, *path):
         return self.encapsulatedJob.rv(*path)
 
+    def prepareForPromiseRegistration(self, jobStore):
+        super(EncapsulatedJob, self).prepareForPromiseRegistration(jobStore)
+        self.encapsulatedJob.prepareForPromiseRegistration(jobStore)
+
+    def getUserScript(self):
+        return self.encapsulatedJob.getUserScript()
+
 
 class ServiceJob(Job):
     """
@@ -1502,21 +1522,23 @@ class Promise(object):
     def __init__(self, job, path):
         """
         :param Job job: the job whose return value this promise references
+        :param path: see :meth:`Job.rv`
         """
         self.job = job
         self.path = path
 
     def __reduce__(self):
         """
-        Called during pickling when a promise of this class is about to be be pickled. Returns
-        the Promise class and construction arguments that will be evaluated during unpickling,
-        namely the job store coordinates of a file that will hold the promised return value. By
-        the time the promise is about to be unpickled, that file should be populated.
+        Called during pickling when a promise (an instance of this class) is about to be be
+        pickled. Returns the Promise class and construction arguments that will be evaluated
+        during unpickling, namely the job store coordinates of a file that will hold the promised
+        return value. By the time the promise is about to be unpickled, that file should be
+        populated.
         """
         # The allocation of the file in the job store is intentionally lazy, we only allocate an
         # empty file in the job store if the promise is actually being pickled. This is done so
         # that we do not allocate files for promises that are never used.
-        jobStoreLocator, jobStoreFileID = self.job.allocatePromiseFile(self.path)
+        jobStoreLocator, jobStoreFileID = self.job.registerPromise(self.path)
         # Returning a class object here causes the pickling machinery to attempt to instantiate
         # the class. We will catch that with __new__ and return an the actual return value instead.
         return self.__class__, (jobStoreLocator, jobStoreFileID)
diff --git a/src/toil/leader.py b/src/toil/leader.py
--- a/src/toil/leader.py
+++ b/src/toil/leader.py
@@ -753,13 +753,12 @@ def mainLoop(config, batchSystem, provisioner, jobStore, rootJobWrapper, jobCach
         raise FailedJobsException(config.jobStore, toilState.totalFailedJobs, jobStore)
 
     # Parse out the return value from the root job
-    with jobStore.readSharedFileStream("rootJobReturnValue") as jobStoreFileID:
-        with jobStore.readFileStream(jobStoreFileID.read()) as fH:
-            try:
-                return cPickle.load(fH)  # rootJobReturnValue
-            except EOFError:
-                logger.exception("Failed to unpickle root job return value")
-                raise FailedJobsException(jobStoreFileID, toilState.totalFailedJobs, jobStore)
+    with jobStore.readSharedFileStream('rootJobReturnValue') as fH:
+        try:
+            return cPickle.load(fH)
+        except EOFError:
+            logger.exception('Failed to unpickle root job return value')
+            raise FailedJobsException(config.jobStore, toilState.totalFailedJobs, jobStore)
 
 
 


Test patch:
diff --git a/src/toil/test/src/hotDeploymentTest.py b/src/toil/test/src/hotDeploymentTest.py
--- a/src/toil/test/src/hotDeploymentTest.py
+++ b/src/toil/test/src/hotDeploymentTest.py
@@ -28,7 +28,7 @@ def _venvApplianceCluster(self):
                                       '--system-site-packages',
                                       '--never-download',  # prevent silent upgrades to pip etc
                                       'venv')
-                leader.runOnAppliance('venv/bin/pip', 'list') # For diagnostic purposes
+                leader.runOnAppliance('venv/bin/pip', 'list')  # For diagnostic purposes
                 yield leader, worker
             finally:
                 # Without this step, we would leak files owned by root on the host's file system
@@ -141,3 +141,51 @@ def job(job, disk='10M', cores=1, memory='10M'):
                                   '/data/jobstore')
             # Assert that out output file is there
             worker.runOnAppliance('test', '-f', '/data/foo.txt')
+
+    def testUserTypesInJobFunctionArgs(self):
+        """
+        Test encapsulated, function-wrapping jobs where the function arguments reference
+        user-defined types.
+
+        Mainly written to cover https://github.com/BD2KGenomics/toil/issues/1259 but then also
+        revealed https://github.com/BD2KGenomics/toil/issues/1278.
+        """
+        with self._venvApplianceCluster() as (leader, worker):
+            def userScript():
+                from toil.job import Job
+                from toil.common import Toil
+
+                # A user-defined type, i.e. a type defined in the user script
+                class X(object):
+                    pass
+
+                # noinspection PyUnusedLocal
+                def job(job, x, disk='10M', cores=1, memory='10M'):
+                    return x
+
+                if __name__ == '__main__':
+                    options = Job.Runner.getDefaultArgumentParser().parse_args()
+                    x = X()
+                    with Toil(options) as toil:
+                        r = toil.start(Job.wrapJobFn(job, x).encapsulate())
+                    # Assert that the return value is of type X, but not X from the __main__
+                    # module but X from foo.bar, the canonical name for the user module. The
+                    # translation from __main__ to foo.bar is a side effect of hot-deployment.
+                    assert r.__class__ is not X
+                    import foo.bar
+                    assert r.__class__ is foo.bar.X
+                    # Assert that a copy was made. This is a side effect of pickling/unpickling.
+                    assert x is not r
+
+            userScript = self._getScriptSource(userScript)
+
+            leader.deployScript(path=self.sitePackages,
+                                packagePath='foo.bar',
+                                script=userScript)
+
+            leader.runOnAppliance('venv/bin/python', '-m', 'foo.bar',
+                                  '--logDebug',
+                                  '--batchSystem=mesos',
+                                  '--mesosMaster=localhost:5050',
+                                  '--defaultMemory=10M',
+                                  '/data/jobstore')