Repository: DataBiosphere/toil

Issue:
Service-based jobs need user script resolution
Continued from #268



Discussion:
Tentatively scheduling this for the release.


Resolving patch:
diff --git a/src/toil/common.py b/src/toil/common.py
--- a/src/toil/common.py
+++ b/src/toil/common.py
@@ -16,7 +16,6 @@
 from contextlib import contextmanager
 import logging
 import os
-from subprocess import CalledProcessError
 import sys
 import cPickle
 from argparse import ArgumentParser
@@ -82,7 +81,7 @@ def __init__(self):
     def setOptions(self, options):
         """
         Creates a config object from the options object.
-        """
+        """ 
         from bd2k.util.humanize import human2bytes #This import is used to convert
         #from human readable quantites to integers 
         def setOption(varName, parsingFn=None, checkFn=None):
diff --git a/src/toil/job.py b/src/toil/job.py
--- a/src/toil/job.py
+++ b/src/toil/job.py
@@ -24,16 +24,14 @@
 import uuid
 import time
 import copy_reg
+import cPickle 
+import logging 
+
+from bd2k.util.humanize import human2bytes
+
 from toil.resource import ModuleDescriptor
 from toil.common import loadJobStore
-from bd2k.util.humanize import human2bytes
 
-try:
-    import cPickle 
-except ImportError:
-    import pickle as cPickle
-    
-import logging
 logger = logging.getLogger( __name__ )
 
 from toil.lib.bioio import (setLoggingFromOptions,
@@ -656,6 +654,43 @@ def _loadRootJob(jobStore):
         if not jobStore.exists(rootJobID):
             raise JobException("No root job (%s) left in toil workflow (workflow has finished successfully?)" % rootJobID)
         return jobStore.load(rootJobID)
+    
+    @staticmethod
+    def _loadClass(className, userModule):
+        """
+        Loads class so that instance can be unpickled.
+
+        :type userModule: ModuleDescriptor
+        """
+        if not userModule.belongsToToil:
+            userModule = userModule.localize()
+        if userModule.dirPath not in sys.path:
+            # FIXME: prepending to sys.path will probably fix #103
+            sys.path.append(userModule.dirPath)
+        userModule = importlib.import_module(userModule.name)
+        thisModule = sys.modules[__name__]
+        #TODO: Document what this magic is doing
+        thisModule.__dict__[className] = userModule.__dict__[className]
+    
+    @staticmethod
+    def _loadJob(command, jobStore):
+        """
+        Unpickles a job.Job instance by decoding the command. See job.Job._serialiseFirstJob and
+        job.Job._makeJobWrappers to see how the Job is encoded in the command. Essentially the
+        command is a reference to a jobStoreFileID containing the pickle file for the job and a
+        list of modules which must be imported so that the Job can be successfully unpickled.
+        """
+        commandTokens = command.split()
+        assert "scriptTree" == commandTokens[0]
+        userModule = ModuleDescriptor(*(commandTokens[3:]))
+        Job._loadClass(commandTokens[2], userModule)
+        pickleFile = commandTokens[1]
+        if pickleFile == "firstJob":
+            openFileStream = jobStore.readSharedFileStream( pickleFile )
+        else:
+            openFileStream = jobStore.readFileStream( pickleFile )
+        with openFileStream as fileHandle:
+            return cPickle.load( fileHandle )
 
     ####################################################
     #Functions to pass Job.run return values to the
@@ -903,7 +938,9 @@ def __init__(self, userFunction, *args, **kwargs):
     def _getUserFunction(self):
         #If dill is installed unpickle the user function directly
         
-        userFunctionModule = self.userFunctionModule.localize()
+        userFunctionModule = self.userFunctionModule
+        if not userFunctionModule.belongsToToil:
+            userFunctionModule = userFunctionModule.localize()
         if userFunctionModule.dirPath not in sys.path:
             # FIXME: prepending to sys.path will probably fix #103
             sys.path.append(userFunctionModule.dirPath)
@@ -942,12 +979,18 @@ def run(self, fileStore):
     
 class ServiceJob(Job):
     """
-    Job used to wrap a Job.Service instance. This constructor should
-    not be called by a user.
+    Job used to wrap a Job.Service instance. This constructor should not be called by a user.
     """
     def __init__(self, service):
+        """
+        :type service: Job.Service
+        """
         Job.__init__(self, memory=service.memory, cores=service.cores)
-        self.service = service
+        # service.__module__ is the module defining the class service is an instance of.
+        self.serviceModule = ModuleDescriptor.forModule(service.__module__).globalize()
+        self.serviceClassName = service.__class__.__name__
+        #The service to run, pickled
+        self.pickledService = cPickle.dumps(service)
         #An empty file in the jobStore which when deleted is used to signal
         #that the service should cease, is initialised in 
         #Job._modifyJobGraphForServices
@@ -957,8 +1000,11 @@ def __init__(self, service):
         self.startFileStoreID = None
         
     def run(self, fileStore):
+        #Unpickle the service
+        self._loadClass(self.serviceClassName, self.serviceModule) #This gets the class loaded 
+        service = cPickle.loads(self.pickledService)
         #Start the service
-        startCredentials = self.service.start()
+        startCredentials = service.start()
         #The start credentials  must be communicated to processes connecting to
         #the service, to do this while the run method is running we 
         #cheat and set the return value promise within the run method
@@ -977,8 +1023,12 @@ def run(self, fileStore):
         while fileStore.globalFileExists(self.stopFileStoreID):
             time.sleep(1) #Avoid excessive polling
         #Now kill the service
-        self.service.stop()
+        service.stop()
         
+    def getUserScript(self):
+        return self.serviceModule
+
+
 class EncapsulatedJob(Job):
     """
     An convenience Job class used to make a job subgraph appear to
diff --git a/src/toil/worker.py b/src/toil/worker.py
--- a/src/toil/worker.py
+++ b/src/toil/worker.py
@@ -24,7 +24,6 @@
     sys.path.remove(os.path.dirname(os.path.abspath(__file__)))
 
 import traceback
-import importlib
 import time
 import socket
 import logging
@@ -47,36 +46,6 @@ def truncateFile(fileNameString, tooBig=50000):
         fh.write(data)
         fh.truncate()
         fh.close()
-
-def loadJob(command, jobStore):
-    """
-    Unpickles a job.Job instance by decoding the command.
-    See job.Job._serialiseFirstJob and job.Job._serialiseFirstJob
-    job.Job._makeJobWrappers to see how the Job is encoded in the command.
-    Essentially the command is a reference to a jobStoreFileID containing 
-    the pickle file for the job and a list of modules which must be imported
-    so that the Job can be successfully unpickled.
-    """
-    commandTokens = command.split()
-    assert "scriptTree" == commandTokens[0]
-    pickleFile = commandTokens[1]
-    jobClassName = commandTokens[2]
-    # must import lazily because toil might not be on sys.path when the top-level of this module is run
-    from toil.resource import ModuleDescriptor
-    userModule = ModuleDescriptor(*commandTokens[3:])
-    if not userModule.belongsToToil:
-        userModule = userModule.localize()
-    if userModule.dirPath not in sys.path:
-        sys.path.append(userModule.dirPath)
-    userModule = importlib.import_module(userModule.name)
-    thisModule = sys.modules[__name__]
-    thisModule.__dict__[jobClassName] = userModule.__dict__[jobClassName]
-    if pickleFile == "firstJob":
-        openFileStream = jobStore.readSharedFileStream( pickleFile )
-    else:
-        openFileStream = jobStore.readFileStream( pickleFile )
-    with openFileStream as fileHandle:
-        return cPickle.load( fileHandle )
     
 def nextOpenDescriptor():
     """Gets the number of the next available file descriptor.
@@ -104,6 +73,7 @@ def main():
     from toil.lib.bioio import makePublicDir
     from toil.lib.bioio import system
     from toil.common import loadJobStore
+    from toil.job import Job
     
     ########################################## 
     #Input args
@@ -262,7 +232,7 @@ def main():
                     localTempDir = makePublicDir(os.path.join(localWorkerTempDir, "localTempDir"))
                     
                     #Is a job command
-                    messages = loadJob(job.command, jobStore)._execute( jobWrapper=job,
+                    messages = Job._loadJob(job.command, jobStore)._execute( jobWrapper=job,
                                     stats=stats, localTempDir=localTempDir, 
                                     jobStore=jobStore)
                     


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
@@ -24,7 +24,7 @@
 from toil.job import Job
 from toil.test.sort.lib import merge, sort, copySubRangeOfFile, getMidPoint
 
-success_ratio = 1.0 #0.5
+success_ratio = 0.5
 
 def setup(job, inputFile, N):
     """Sets up the sort.
diff --git a/src/toil/test/src/jobServiceTest.py b/src/toil/test/src/jobServiceTest.py
--- a/src/toil/test/src/jobServiceTest.py
+++ b/src/toil/test/src/jobServiceTest.py
@@ -19,7 +19,6 @@
 from toil.job import Job
 from toil.test import ToilTest
 
-
 class JobServiceTest(ToilTest):
     """
     Tests testing the Job.Service class
@@ -33,33 +32,7 @@ def testService(self):
         try:
             # Wire up the services/jobs
             t = Job.wrapFn(f, "1", outFile)
-
-            # This is a temporary fix for https://github.com/BD2KGenomics/toil/issues/268. The
-            # current implementation of ServiceJob ignores the fact that a Job.Service subclass
-            # in a user script requires that user script to be imported first. This is very
-            # similar to how the FunctionWrappingJob first imports the module containing the user
-            #  function before referencing it. The pickled job references the module containing
-            # the service class, likely causing the unpickling to fail since the referenced
-            # module hasn't been loaded yet. The reason this test worked without this temporary
-            # fix was that the test runner imports this test module using the fully-qualified
-            # package name, i.e. "toil.test.src.jobServiceTest". That package name can be
-            # resolved if the top-level src directory is present on the sys.path, which is the
-            # case when the test is run from the command line, e.g. on Jenkins. OTOH, when this
-            # test is run from PyCharm, the test module is imported from the directory containing
-            #  the module, yielding the unqalified module name "jobServiceTest". The directory
-            # containing the module is not on the worker's sys.path and the module name can't
-            # therefore be resolved. But don't be fooled: the problem is not limited to PyCharm,
-            # it will happen in every case where the user script is not part of the Toil source
-            # tree or not loaded from the top-level source tree.
-            #
-            # By aliasing the service class it will be referenced in the pickled job using its
-            # fully qualified name which, as mentioned above, *can* be resolved in the worker
-            # process.
-
-            from toil.test.src.jobServiceTest import TestService as _TestService
-
-
-            t.addChildFn(f, t.addService(_TestService("2", "3", outFile)), outFile)
+            t.addChildFn(f, t.addService(TestService("2", "3", outFile)), outFile)
             # Create the runner for the workflow.
             options = Job.Runner.getDefaultOptions()
             options.logLevel = "INFO"
@@ -71,7 +44,6 @@ def testService(self):
         finally:
             os.remove(outFile)
 
-
 class TestService(Job.Service):
     def __init__(self, startString, stopString, outFile):
         Job.Service.__init__(self)
diff --git a/src/toil/test/utils/utilsTest.py b/src/toil/test/utils/utilsTest.py
--- a/src/toil/test/utils/utilsTest.py
+++ b/src/toil/test/utils/utilsTest.py
@@ -66,28 +66,34 @@ def testUtilsSort(self):
         self.assertRaises(CalledProcessError, system, toilCommandString + " --restart")
         # Check that trying to run it in restart mode does not create the jobStore
         self.assertFalse(os.path.exists(self.toilDir))
+        
+        # Status command
+        rootPath = os.path.join(toilPackageDirPath(), "utils")
+        toilStatusCommandString = ("{rootPath}/toilMain.py status "
+                                   "--jobStore {self.toilDir} --failIfNotComplete".format(**locals()))
 
         # Run the script for the first time
-        system(toilCommandString)
+        try:
+            system(toilCommandString)
+            finished = True
+        except CalledProcessError: #This happens when the script fails due to having unfinished jobs
+            self.assertRaises(CalledProcessError, system, toilStatusCommandString)
+            finished = False
         self.assertTrue(os.path.exists(self.toilDir))
 
         # Try running it without restart and check an exception is thrown
         self.assertRaises(CalledProcessError, system, toilCommandString)
 
         # Now restart it until done
-        while True:
-            # Run the status command
-            rootPath = os.path.join(toilPackageDirPath(), "utils")
-            toilStatusString = ("{rootPath}/toilMain.py status "
-                                "--jobStore {self.toilDir} --failIfNotComplete".format(**locals()))
+        while not finished:
             try:
-                system(toilStatusString)
-            except CalledProcessError:
-                # If there is an exception there are still failed jobs then restart
                 system(toilCommandString + " --restart")
-            else:
-                break
-
+                finished = True
+            except CalledProcessError: #This happens when the script fails due to having unfinished jobs
+                self.assertRaises(CalledProcessError, system, toilStatusCommandString)
+                
+        #Check the toil status command does not issue an exception
+        system(toilStatusCommandString)
 
         # Check if we try to launch after its finished that we get a JobException
         self.assertRaises(CalledProcessError, system, toilCommandString + " --restart")