Repository: DataBiosphere/toil

Issue:
Toil can't unpickle classes defined in user scripts
I have a class, `AzureIOStore`, defined in my user script. It overrides `__getstate__` and `__setstate__` for custom pickling behavior.

Toil can't seem to unpickle it on my Mesos workers:

```
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:   ---TOIL WORKER OUTPUT LOG---
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:   Next available file descriptor: 5
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:   DEBUG:__main__:Next available file descriptor: 5
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:   Parsed job
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:   DEBUG:__main__:Parsed job
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:   Traceback (most recent call last):
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:     File "/usr/local/lib/python2.7/dist-packages/toil/worker.py", line 238, in main
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:       jobStore)._execute( jobWrapper=job,
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:     File "/usr/local/lib/python2.7/dist-packages/toil/job.py", line 555, in _loadJob
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:       return cls._unpickle(userModule, fileHandle)
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:     File "/usr/local/lib/python2.7/dist-packages/toil/job.py", line 576, in _unpickle
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:       return unpickler.load()
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:     File "/usr/local/lib/python2.7/dist-packages/toil/job.py", line 571, in filter_main
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:       return getattr(userModule, class_name)
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:   AttributeError: 'module' object has no attribute 'AzureIOStore'
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:   Exiting the worker because of a failed job on host hgvmclusteragent1
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:   ERROR:__main__:Exiting the worker because of a failed job on host hgvmclusteragent1
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:   WARNING:toil.jobWrapper:Due to failure we are reducing the remaining retry count of job 732fd38e_d585_4405_8d6b_c37a982f1ba5 to 0
```

Toil can't unpickle classes defined in user scripts
I have a class, `AzureIOStore`, defined in my user script. It overrides `__getstate__` and `__setstate__` for custom pickling behavior.

Toil can't seem to unpickle it on my Mesos workers:

```
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:   ---TOIL WORKER OUTPUT LOG---
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:   Next available file descriptor: 5
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:   DEBUG:__main__:Next available file descriptor: 5
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:   Parsed job
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:   DEBUG:__main__:Parsed job
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:   Traceback (most recent call last):
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:     File "/usr/local/lib/python2.7/dist-packages/toil/worker.py", line 238, in main
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:       jobStore)._execute( jobWrapper=job,
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:     File "/usr/local/lib/python2.7/dist-packages/toil/job.py", line 555, in _loadJob
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:       return cls._unpickle(userModule, fileHandle)
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:     File "/usr/local/lib/python2.7/dist-packages/toil/job.py", line 576, in _unpickle
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:       return unpickler.load()
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:     File "/usr/local/lib/python2.7/dist-packages/toil/job.py", line 571, in filter_main
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:       return getattr(userModule, class_name)
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:   AttributeError: 'module' object has no attribute 'AzureIOStore'
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:   Exiting the worker because of a failed job on host hgvmclusteragent1
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:   ERROR:__main__:Exiting the worker because of a failed job on host hgvmclusteragent1
WARNING:toil.leader:732fd38e_d585_4405_8d6b_c37a982f1ba5:   WARNING:toil.jobWrapper:Due to failure we are reducing the remaining retry count of job 732fd38e_d585_4405_8d6b_c37a982f1ba5 to 0
```



Discussion:
This is the code that handles the references to the "**main**" module in the pickled data. When you run the user script on the leader **main** is the user script. On a worker, **main** is worker.py. I don't know why this is failing though. I would like to know that userModule is at the time this exception is raised. It might be worth trying to remove custom pickling methods and see if that changes anything.

I've tested it without the custom pickling code. Even if I make my class pickle and unpickle normally, I can't unpickle it on the worker.

In my old jobTree scripts I would sometimes need to re-import things:

```
# Make sure we've given everything an absolute module name.
# Don't try to import * because that's illegal.
if __name__ == "__main__":
    from compareSchemes import SchemeAssessmentTarget, BedSplitTarget, \
        MafGeneCheckerTarget
```

I tried that here, but Toil seems to be trying to prevent it:

```
Traceback (most recent call last):
  File "./parallelMappingEvaluation.py", line 975, in <module>
    sys.exit(main(sys.argv))
  File "./parallelMappingEvaluation.py", line 962, in main
    cores=16, memory="100G", disk="50G")
  File "/usr/local/lib/python2.7/dist-packages/toil/job.py", line 183, in wrapJobFn
    return JobFunctionWrappingJob(fn, *args, **kwargs)
  File "/usr/local/lib/python2.7/dist-packages/toil/job.py", line 995, in __init__
    super(JobFunctionWrappingJob, self).__init__(userFunction, *args, **kwargs)
  File "/usr/local/lib/python2.7/dist-packages/toil/job.py", line 964, in __init__
    self.userFunctionModule = ModuleDescriptor.forModule(userFunction.__module__).globalize()
  File "/usr/local/lib/python2.7/dist-packages/toil/resource.py", line 343, in forModule
    cls._check_conflict(dirPath, name)
  File "/usr/local/lib/python2.7/dist-packages/toil/resource.py", line 371, in _check_conflict
    name, colliding_module.__name__, colliding_module.__file__))
RuntimeError: The user module 'parallelMappingEvaluation' collides with module 'parallelMappingEvaluation from '/home/anovak/hgvm/parallelMappingEvaluation.py'.
```

That's the code that makes sure a user module can only be loaded from one entry on the sys.path. We did this to ensure that someone naming their user script in a way that collides with a module already installed, `test` was a prominent example.

@adamnovak could you provide the user script?

I think this is it:

https://github.com/adamnovak/hgvm/blob/a0e9be6f544984cff3b6d8622783de8700ab42f2/parallelMappingEvaluation.py

@adamnovak, let me know if that is wrong. The commit above is the one before the one where you switched from sticking  the IOStore instance in the options instance in the root job.

I can reproduce this.

This is the code that handles the references to the "**main**" module in the pickled data. When you run the user script on the leader **main** is the user script. On a worker, **main** is worker.py. I don't know why this is failing though. I would like to know that userModule is at the time this exception is raised. It might be worth trying to remove custom pickling methods and see if that changes anything.

I've tested it without the custom pickling code. Even if I make my class pickle and unpickle normally, I can't unpickle it on the worker.

In my old jobTree scripts I would sometimes need to re-import things:

```
# Make sure we've given everything an absolute module name.
# Don't try to import * because that's illegal.
if __name__ == "__main__":
    from compareSchemes import SchemeAssessmentTarget, BedSplitTarget, \
        MafGeneCheckerTarget
```

I tried that here, but Toil seems to be trying to prevent it:

```
Traceback (most recent call last):
  File "./parallelMappingEvaluation.py", line 975, in <module>
    sys.exit(main(sys.argv))
  File "./parallelMappingEvaluation.py", line 962, in main
    cores=16, memory="100G", disk="50G")
  File "/usr/local/lib/python2.7/dist-packages/toil/job.py", line 183, in wrapJobFn
    return JobFunctionWrappingJob(fn, *args, **kwargs)
  File "/usr/local/lib/python2.7/dist-packages/toil/job.py", line 995, in __init__
    super(JobFunctionWrappingJob, self).__init__(userFunction, *args, **kwargs)
  File "/usr/local/lib/python2.7/dist-packages/toil/job.py", line 964, in __init__
    self.userFunctionModule = ModuleDescriptor.forModule(userFunction.__module__).globalize()
  File "/usr/local/lib/python2.7/dist-packages/toil/resource.py", line 343, in forModule
    cls._check_conflict(dirPath, name)
  File "/usr/local/lib/python2.7/dist-packages/toil/resource.py", line 371, in _check_conflict
    name, colliding_module.__name__, colliding_module.__file__))
RuntimeError: The user module 'parallelMappingEvaluation' collides with module 'parallelMappingEvaluation from '/home/anovak/hgvm/parallelMappingEvaluation.py'.
```

That's the code that makes sure a user module can only be loaded from one entry on the sys.path. We did this to ensure that someone naming their user script in a way that collides with a module already installed, `test` was a prominent example.

@adamnovak could you provide the user script?

I think this is it:

https://github.com/adamnovak/hgvm/blob/a0e9be6f544984cff3b6d8622783de8700ab42f2/parallelMappingEvaluation.py

@adamnovak, let me know if that is wrong. The commit above is the one before the one where you switched from sticking  the IOStore instance in the options instance in the root job.

I can reproduce this.


Resolving patch:
diff --git a/src/toil/job.py b/src/toil/job.py
--- a/src/toil/job.py
+++ b/src/toil/job.py
@@ -204,7 +204,7 @@ def encapsulate(self):
     #The following function is used for passing return values between
     #job run functions
     ####################################################
-    
+
     def rv(self, argIndex=None):
         """
         Gets a PromisedJobReturnValue, representing the argIndex return
@@ -366,7 +366,7 @@ class FileStore( object ):
         Class used to manage temporary files and log messages, 
         passed as argument to the Job.run method.
         """
-        def __init__(self, jobStore, jobWrapper, localTempDir, 
+        def __init__(self, jobStore, jobWrapper, localTempDir,
                      inputBlockFn, jobStoreFileIDToCacheLocation, terminateEvent):
             """
             This constructor should not be called by the user, 
@@ -382,7 +382,7 @@ def __init__(self, jobStore, jobWrapper, localTempDir,
             #Asynchronous writes stuff
             self.workerNumber = 2
             self.queue = Queue()
-            self.updateSemaphore = Semaphore() 
+            self.updateSemaphore = Semaphore()
             self.terminateEvent = terminateEvent
             #Function to write files to job store
             def asyncWrite():
@@ -392,7 +392,7 @@ def asyncWrite():
                             #Block for up to two seconds waiting for a file
                             args = self.queue.get(timeout=2)
                         except Empty:
-                            #Check if termination event is signaled 
+                            #Check if termination event is signaled
                             #(set in the event of an exception in the worker)
                             if terminateEvent.isSet():
                                 raise RuntimeError("The termination flag is set, exiting")
@@ -401,8 +401,8 @@ def asyncWrite():
                         if args == None:
                             break
                         inputFileHandle, jobStoreFileID = args
-                        #We pass in a fileHandle, rather than the file-name, in case 
-                        #the file itself is deleted. The fileHandle itself should persist 
+                        #We pass in a fileHandle, rather than the file-name, in case
+                        #the file itself is deleted. The fileHandle itself should persist
                         #while we maintain the open file handle
                         with jobStore.updateFileStream(jobStoreFileID) as outputFileHandle:
                             bufferSize=1000000 #TODO: This buffer number probably needs to be modified/tuned
@@ -415,15 +415,15 @@ def asyncWrite():
                 except:
                     terminateEvent.set()
                     raise
-                    
-            self.workers = map(lambda i : Thread(target=asyncWrite), 
+
+            self.workers = map(lambda i : Thread(target=asyncWrite),
                                range(self.workerNumber))
             for worker in self.workers:
                 worker.start()
             self.inputBlockFn = inputBlockFn
-            #Caching 
+            #Caching
             #
-            #For files in jobStore that are on the local disk, 
+            #For files in jobStore that are on the local disk,
             #map of jobStoreFileIDs to locations in localTempDir.
             self.jobStoreFileIDToCacheLocation = jobStoreFileIDToCacheLocation
 
@@ -435,7 +435,7 @@ def getLocalTempDir(self):
             the job terminates, removing all files it contains recursively. 
             """
             return os.path.abspath(tempfile.mkdtemp(prefix="t", dir=self.localTempDir))
-        
+
         def getLocalTempFile(self):
             """
             Get an absolute path to a local temporary file. 
@@ -445,7 +445,7 @@ def getLocalTempFile(self):
             handle, tmpFile = tempfile.mkstemp(prefix="tmp",
                                                suffix=".tmp", dir=self.localTempDir)
             os.close(handle)
-            return os.path.abspath(tmpFile) 
+            return os.path.abspath(tmpFile)
 
         def writeGlobalFile(self, localFileName, cleanup=False):
             """
@@ -476,7 +476,7 @@ def writeGlobalFile(self, localFileName, cleanup=False):
                 #Write the file directly to the file store
                 jobStoreFileID = self.jobStore.writeFile(localFileName, cleanupID)
             return jobStoreFileID
-        
+
         def writeGlobalFileStream(self, cleanup=False):
             """
             Similar to writeGlobalFile, but returns a context manager yielding a 
@@ -488,7 +488,7 @@ def writeGlobalFileStream(self, cleanup=False):
             """
             #TODO: Make this work with the caching??
             return self.jobStore.writeFileStream(None if not cleanup else self.jobWrapper.jobStoreID)
-        
+
         def readGlobalFile(self, fileStoreID, userPath=None, cache=True):
             """
             Returns an absolute path to a local, temporary copy of the file 
@@ -523,11 +523,11 @@ def readGlobalFile(self, fileStoreID, userPath=None, cache=True):
                 userPath = os.path.abspath(userPath) #Make an absolute path
                 #Turn off caching if user file is not in localTempDir
                 if cache and not userPath.startswith(self.localTempDir):
-                    cache = False  
+                    cache = False
             #When requesting a new file from the jobStore first check if fileStoreID
             #is a key in jobStoreFileIDToCacheLocation.
             if fileStoreID in self.jobStoreFileIDToCacheLocation:
-                cachedAbsFilePath = self.jobStoreFileIDToCacheLocation[fileStoreID]   
+                cachedAbsFilePath = self.jobStoreFileIDToCacheLocation[fileStoreID]
                 if cache:
                     #If the user specifies a location and it is not the current location
                     #return a hardlink to the location, else return the original location
@@ -544,7 +544,7 @@ def readGlobalFile(self, fileStoreID, userPath=None, cache=True):
                     shutil.copyfile(cachedAbsFilePath, localFilePath)
                     return localFilePath
             else:
-                #If it is not in the cache read it from the jobStore to the 
+                #If it is not in the cache read it from the jobStore to the
                 #desired location
                 localFilePath = userPath if userPath != None else self.getLocalTempFile()
                 self.jobStore.readFile(fileStoreID, localFilePath)
@@ -564,11 +564,11 @@ def readGlobalFileStream(self, fileStoreID):
             """
             if fileStoreID in self.filesToDelete:
                 raise RuntimeError("Trying to access a file in the jobStore you've deleted: %s" % fileStoreID)
-            
+
             #If fileStoreID is in the cache provide a handle from the local cache
             if fileStoreID in self.jobStoreFileIDToCacheLocation:
                 #This leaks file handles (but the commented out code does not work properly)
-                return open(self.jobStoreFileIDToCacheLocation[fileStoreID], 'r') 
+                return open(self.jobStoreFileIDToCacheLocation[fileStoreID], 'r')
                 #with open(self.jobStoreFileIDToCacheLocation[fileStoreID], 'r') as fH:
                 #        yield fH
             else:
@@ -596,9 +596,9 @@ def logToMaster(self, string, level=logging.INFO):
             """
             logger.log(level=level, msg=("LOG-TO-MASTER: " + string))
             self.loggingMessages.append((str(string), level))
-            
-        #Private methods 
-        
+
+        #Private methods
+
         def _updateJobWhenDone(self):
             """
             Asynchronously update the status of the job on the disk, first waiting
@@ -610,33 +610,33 @@ def asyncUpdate():
                     #Wait till all file writes have completed
                     for i in xrange(len(self.workers)):
                         self.queue.put(None)
-            
+
                     for thread in self.workers:
                         thread.join()
-                    
+
                     #Wait till input block-fn returns - in the event of an exception
-                    #this will eventually terminate 
+                    #this will eventually terminate
                     self.inputBlockFn()
-                    
+
                     #Check the terminate event, if set we can not guarantee
                     #that the workers ended correctly, therefore we exit without
                     #completing the update
                     if self.terminateEvent.isSet():
                         raise RuntimeError("The termination flag is set, exiting before update")
-                    
-                    #Indicate any files that should be deleted once the update of 
+
+                    #Indicate any files that should be deleted once the update of
                     #the job wrapper is completed.
                     self.jobWrapper.filesToDelete = list(self.filesToDelete)
-                    
+
                     #Complete the job
                     self.jobStore.update(self.jobWrapper)
-                    
+
                     #Delete any remnant jobs
                     map(self.jobStore.delete, self.jobsToDelete)
-                    
+
                     #Delete any remnant files
                     map(self.jobStore.deleteFile, self.filesToDelete)
-                    
+
                     #Remove the files to delete list, having successfully removed the files
                     if len(self.filesToDelete) > 0:
                         self.jobWrapper.filesToDelete = []
@@ -657,7 +657,7 @@ def asyncUpdate():
             except: #This is to ensure that the semaphore is released in a crash to stop a deadlock scenario
                 self.updateSemaphore.release()
                 raise
-            
+
         def _cleanLocalTempDir(self, cacheSize):
             """
             At the end of the job, remove all localTempDir files except those whose value is in
@@ -666,7 +666,7 @@ def _cleanLocalTempDir(self, cacheSize):
             :param cacheSize: the total number of bytes of files allowed in the cache.
             """
             #Remove files so that the total cached files are smaller than a cacheSize
-            
+
             #List of pairs of (fileCreateTime, fileStoreID) for cached files
             cachedFileCreateTimes = map(lambda x : (os.stat(self.jobStoreFileIDToCacheLocation[x]).st_ctime, x),
                                         self.jobStoreFileIDToCacheLocation.keys())
@@ -686,13 +686,13 @@ def _cleanLocalTempDir(self, cacheSize):
                 os.remove(filePath)
                 totalCachedFileSizes -= fileSize
                 assert totalCachedFileSizes >= 0
-            
-            #Iterate from the base of localTempDir and remove all 
+
+            #Iterate from the base of localTempDir and remove all
             #files/empty directories, recursively
             cachedFiles = set(self.jobStoreFileIDToCacheLocation.values())
-            
+
             def clean(dirOrFile):
-                canRemove = True 
+                canRemove = True
                 if os.path.isdir(dirOrFile):
                     for f in os.listdir(dirOrFile):
                         canRemove = canRemove and clean(os.path.join(dirOrFile, f))
@@ -702,20 +702,20 @@ def clean(dirOrFile):
                 if dirOrFile in cachedFiles:
                     return False
                 os.remove(dirOrFile)
-                return True    
+                return True
             clean(self.localTempDir)
-        
+
         def _blockFn(self):
             """
             Blocks while _updateJobWhenDone is running.
-            """ 
+            """
             self.updateSemaphore.acquire()
             self.updateSemaphore.release() #Release so that the block function can be recalled
             #This works, because once acquired the semaphore will not be acquired
             #by _updateJobWhenDone again.
             return
-        
-        def __del__(self): 
+
+        def __del__(self):
             """Cleanup function that is run when destroying the class instance 
             that ensures that all the file writing threads exit.
             """
@@ -836,7 +836,7 @@ def filter_main(module_name, class_name):
 
         unpickler.find_global = filter_main
         return unpickler.load()
-    
+
     def getUserScript(self):
         return self.userModule
 
@@ -912,7 +912,7 @@ def _getImpliedEdges(roots):
                 for descendant in reacheable:
                     extraEdges[descendant] += job._followOns[:]
         return extraEdges
-    
+
     ####################################################
     #The following functions are used to serialise
     #a job graph to the jobStore
@@ -974,7 +974,7 @@ def _createEmptyJobForJob(self, jobStore, command=None,
                                disk=(self.disk if self.disk is not None
                                     else float(jobStore.config.defaultDisk)),
                                predecessorNumber=predecessorNumber)
-        
+
     def _makeJobWrappers(self, jobWrapper, jobStore):
         """
         Creates a job for each job in the job graph, recursively.
@@ -998,14 +998,14 @@ def _makeJobWrappers2(self, jobStore, jobsToJobWrappers):
                 jobWrapper.stack.append(jobs)
         else:
             jobWrapper = jobsToJobWrappers[self]
-        #The return is a tuple stored within a job.stack 
+        #The return is a tuple stored within a job.stack
         #The tuple is jobStoreID, memory, cores, disk, predecessorID
         #The predecessorID is used to establish which predecessors have been
         #completed before running the given Job - it is just a unique ID
         #per predecessor
         return (jobWrapper.jobStoreID, jobWrapper.memory, jobWrapper.cores, jobWrapper.disk,
                 None if jobWrapper.predecessorNumber <= 1 else str(uuid.uuid4()))
-        
+
     def getTopologicalOrderingOfJobs(self):
         """
         Get a list of jobs such that for all pairs of indices i, j for which i < j, 
@@ -1025,7 +1025,7 @@ def getRunOrder(job):
                 map(getRunOrder, job._children + job._followOns)
         getRunOrder(self)
         return ordering
-    
+
     def _serialiseJob(self, jobStore, jobsToJobWrappers, rootJobWrapper):
         """
         Pickle a job and its jobWrapper to disk.
@@ -1043,11 +1043,18 @@ def _serialiseJob(self, jobStore, jobsToJobWrappers, rootJobWrapper):
         #method.
         with jobStore.writeFileStream(rootJobWrapper.jobStoreID) as (fileHandle, fileStoreID):
             cPickle.dump(self, fileHandle, cPickle.HIGHEST_PROTOCOL)
-        jobsToJobWrappers[self].command = ' '.join( ('_toil', fileStoreID) + self.userModule.globalize())      
+        # Note that getUserScript() may have beeen overridden. This is intended. If we used
+        # self.userModule directly, we'd be getting a reference to job.py if the job was
+        # specified as a function (as opposed to a class) since that is where FunctionWrappingJob
+        #  is defined. What we really want is the module that was loaded as __main__,
+        # and FunctionWrappingJob overrides getUserScript() to give us just that. Only then can
+        # filter_main() in _unpickle( ) do its job of resolveing any user-defined type or function.
+        userScript = self.getUserScript().globalize()
+        jobsToJobWrappers[self].command = ' '.join( ('_toil', fileStoreID) + userScript)
         #Update the status of the jobWrapper on disk
         jobStore.update(jobsToJobWrappers[self])
-    
-    def _serialiseJobGraph(self, jobWrapper, jobStore, returnValues, firstJob):  
+
+    def _serialiseJobGraph(self, jobWrapper, jobStore, returnValues, firstJob):
         """
         Pickle the graph of jobs in the jobStore. The graph is not fully serialised
         until the jobWrapper itself is written to disk, this is not performed by this
@@ -1060,13 +1067,13 @@ def _serialiseJobGraph(self, jobWrapper, jobStore, returnValues, firstJob):
         self.checkJobGraphForDeadlocks()
         #Create the jobWrappers for followOns/children
         jobsToJobWrappers = self._makeJobWrappers(jobWrapper, jobStore)
-        #Get an ordering on the jobs which we use for pickling the jobs in the 
+        #Get an ordering on the jobs which we use for pickling the jobs in the
         #correct order to ensure the promises are properly established
         ordering = self.getTopologicalOrderingOfJobs()
         assert len(ordering) == len(jobsToJobWrappers)
         #Temporarily set the jobStore strings for the promise call back functions
         for job in ordering:
-            job._promiseJobStore = jobStore 
+            job._promiseJobStore = jobStore
         ordering.reverse()
         assert self == ordering[-1]
         if firstJob:
@@ -1096,7 +1103,7 @@ def _serialiseJobGraph(self, jobWrapper, jobStore, returnValues, firstJob):
                 jobWrapper.stack.append(combinedFollowOns)
             if len(combinedChildren) > 0:
                 jobWrapper.stack.append(combinedChildren)
-            
+
     def _serialiseFirstJob(self, jobStore):
         """
         Serialises the root job. Returns the wrapping job.
@@ -1106,7 +1113,7 @@ def _serialiseFirstJob(self, jobStore):
                                                 predecessorNumber=0)
         #Write the graph of jobs to disk
         self._serialiseJobGraph(jobWrapper, jobStore, None, True)
-        jobStore.update(jobWrapper) 
+        jobStore.update(jobWrapper)
         #Store the name of the first job in a file in case of restart
         #Up to this point the root-job is not recoverable
         with jobStore.writeSharedFileStream("rootJobStoreID") as f:
@@ -1134,9 +1141,9 @@ def _execute(self, jobWrapper, stats, localTempDir, jobStore, fileStore):
         #Add the promise files to delete to the list of jobStoreFileIDs to delete
         for jobStoreFileID in promiseFilesToDelete:
             fileStore.deleteGlobalFile(jobStoreFileID)
-        promiseFilesToDelete.clear() 
+        promiseFilesToDelete.clear()
         #Now indicate the asynchronous update of the job can happen
-            
+
         fileStore._updateJobWhenDone()
         #Change dir back to cwd dir, if changed by job (this is a safety issue)
         if os.getcwd() != baseDir:
@@ -1150,35 +1157,6 @@ def _execute(self, jobWrapper, stats, localTempDir, jobStore, fileStore):
             stats.attrib["class"] = self._jobName()
             stats.attrib["memory"] = str(totalMemoryUsage)
 
-    ####################################################
-    #Method used to resolve the module in which an inherited job instances
-    #class is defined
-    ####################################################
-
-    @staticmethod
-    def _resolveMainModule( moduleName ):
-        """
-        Returns a tuple of two elements, the first element being the path 
-        to the directory containing the given
-        module and the second element being the name of the module. 
-        If the given module name is "__main__",
-        then that is translated to the actual file name of the top-level 
-        script without .py or .pyc extensions. The
-        caller can then add the first element of the returned tuple to 
-        sys.path and load the module from there. See also worker.loadJob().
-        """
-        # looks up corresponding module in sys.modules, gets base name, drops .py or .pyc
-        moduleDirPath, moduleName = os.path.split(os.path.abspath(sys.modules[moduleName].__file__))
-        if moduleName.endswith('.py'):
-            moduleName = moduleName[:-3]
-        elif moduleName.endswith('.pyc'):
-            moduleName = moduleName[:-4]
-        else:
-            raise RuntimeError(
-                "Can only handle main modules loaded from .py or .pyc files, but not '%s'" %
-                moduleName)
-        return moduleDirPath, moduleName
-
     def _jobName(self):
         """
         :rtype : string, used as identifier of the job class in the stats report. 
@@ -1361,8 +1339,8 @@ class PromisedJobReturnValue(object):
     argument to Job before the former Job's run function has been executed.
     """
     def __init__(self, promiseCallBackFunction):
-        self.promiseCallBackFunction = promiseCallBackFunction 
-        
+        self.promiseCallBackFunction = promiseCallBackFunction
+
 def promisedJobReturnValuePickleFunction(promise):
     """
     This function and promisedJobReturnValueUnpickleFunction are used as custom pickle/unpickle 


Test patch:
diff --git a/src/toil/test/src/helloWorldTest.py b/src/toil/test/src/helloWorldTest.py
--- a/src/toil/test/src/helloWorldTest.py
+++ b/src/toil/test/src/helloWorldTest.py
@@ -1,40 +1,52 @@
-'''
-Created on Sep 16, 2015
+# Copyright (C) 2015 UCSC Computational Genomics Lab
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 
-@author: benedictpaten
-'''
+from __future__ import absolute_import
 from toil.job import Job
-from argparse import ArgumentParser
-
-import unittest
 from toil.test import ToilTest
 
+
 class HelloWorldTest(ToilTest):
     def testHelloWorld(self):
         options = Job.Runner.getDefaultOptions(self._getTestJobStorePath())
         options.logLevel = "INFO"
-        # Run the workflow, the return value being the number of failed jobs
         Job.Runner.startToil(HelloWorld(), options)
 
+
 class HelloWorld(Job):
     def __init__(self):
-        Job.__init__(self,  memory=100000, cores=2, disk=20000)
+        Job.__init__(self, memory=100000, cores=2, disk=20000)
+
     def run(self, fileStore):
         fileID = self.addChildJobFn(childFn, cores=1, memory="1M", disk="10M").rv()
         self.addFollowOn(FollowOn(fileID))
 
+
 def childFn(job):
     with job.fileStore.writeGlobalFileStream() as (fH, fileID):
         fH.write("Hello, World!")
         return fileID
-    
+
+
 class FollowOn(Job):
-    def __init__(self,fileId):
+    def __init__(self, fileId):
         Job.__init__(self)
-        self.fileId=fileId
+        self.fileId = fileId
+
     def run(self, fileStore):
         tempDir = fileStore.getLocalTempDir()
-        tempFilePath = "/".join([tempDir,"LocalCopy"])
+        tempFilePath = "/".join([tempDir, "LocalCopy"])
         with fileStore.readGlobalFileStream(self.fileId) as globalFile:
             with open(tempFilePath, "w") as localFile:
                 localFile.write(globalFile.read())
diff --git a/src/toil/test/src/userDefinedJobArgTypeTest.py b/src/toil/test/src/userDefinedJobArgTypeTest.py
new file mode 100644
--- /dev/null
+++ b/src/toil/test/src/userDefinedJobArgTypeTest.py
@@ -0,0 +1,96 @@
+# Copyright (C) 2015 UCSC Computational Genomics Lab
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import absolute_import
+import sys
+from subprocess import check_call
+
+from toil.job import Job
+from toil.test import ToilTest
+
+
+class UserDefinedJobArgTypeTest(ToilTest):
+    """
+    Test for issue #423 (Toil can't unpickle classes defined in user scripts) and variants
+    thereof.
+
+    https://github.com/BD2KGenomics/toil/issues/423
+    """
+
+    def setUp(self):
+        super(UserDefinedJobArgTypeTest, self).setUp()
+        options = Job.Runner.getDefaultOptions(self._getTestJobStorePath())
+        options.logLevel = "INFO"
+        options.foo = Foo()
+        self.options = options
+
+    def testJobFunction(self):
+        """Test with first job being a function"""
+        Job.Runner.startToil(Job.wrapJobFn(jobFunction, 0, Foo()), self.options)
+
+    def testJobClass(self):
+        """Test with first job being an instance of a class"""
+        Job.Runner.startToil(JobClass(0, Foo()), self.options)
+
+    def testJobFunctionFromMain(self):
+        """Test with first job being a function defined in __main__"""
+        self._testFromMain()
+
+    def testJobClassFromMain(self):
+        """Test with first job being an instance of a class defined in __main__"""
+        self._testFromMain()
+
+    def _testFromMain(self):
+        testMethodName = self.id().split('.')[-1]
+        self.assertTrue(testMethodName.endswith('FromMain'))
+        check_call([sys.executable, __file__, testMethodName[:-8]])
+
+
+class JobClass(Job):
+    def __init__(self, level, foo):
+        Job.__init__(self, memory=100000, cores=2, disk=20000)
+        self.level = level
+        self.foo = foo
+
+    def run(self, fileStore):
+        self.foo.assertIsCopy()
+        if self.level < 2:
+            self.addChildJobFn(jobFunction, self.level + 1, Foo(), cores=1, memory="1M", disk="10M")
+
+
+def jobFunction(job, level, foo):
+    foo.assertIsCopy()
+    if level < 2:
+        job.addChild(JobClass(level + 1, Foo()))
+
+
+class Foo(object):
+    def __init__(self):
+        super(Foo, self).__init__()
+        self.original_id = id(self)
+
+    def assertIsCopy(self):
+        assert self.original_id != id(self)
+
+
+def main():
+    testMethodName = sys.argv[1]
+    test = UserDefinedJobArgTypeTest(testMethodName)
+    test.setUpClass()
+    test.debug()
+    test.tearDownClass()
+
+
+if __name__ == '__main__':
+    main()