Repository: DataBiosphere/toil

Issue:
Docker containers can become zombies
If a job terminates but that job was running docker containers, the docker containers can persist in the background and not get killed off.

Depends on #1189 (PR #1199)


Discussion:
As suggested by @hannes-ucsc, look at https://docs.docker.com/swarm/ before doing anything.

@benedictpaten to write up his proposal here.

Renaming to zombie as that is established Unix terminology for processes and Docker containers are processes.

@benedictpaten and I had a long chat yesterday and we built on the chat @hannes-ucsc and I had last week. We will use the inherent caching mechanism to track pids of jobs and then use the filestore context manager to allow jobs to clean up behind themselves on completion in a finally clause. On entering the filestore, jobs will check for other jobs' integrity and purge any dead jobs. Users will have the ability to register cleanup functions with toil that will be run in these situations. We will add a `docker_call` and a corresponding `docker_cleanup` function in toil-lib that can be used by any user.

Try to design this mechanism in a way that let's us do arbitrary per-job cleanup: Let's say we have a job function `foo`:

``` py
def foo(job):
    job.defer(myCleanUp) # register a function to be called when this function returns or raises
    job.recover(myRecovery) # register a function to be invoked before this job is to be retried
    …
```

The particulars would have to be discussed further. For example, we might want Toil to guarantee that either the deferred function or the clean-up function are going to be invoked, but never both. Not sure.

`docker_call()` would then have to be extended to exploit this mechanism. The deferred function would `docker rm` one specific container—the one that will be created by this invocation of `docker_call`, the recovery function would list all containers and kill the ones that could have been created by any number of docker_call invocations.

```
```

can we discuss this tomorrow after lunch?

@hannes-ucsc one thing to keep in mind is that the clean up function would need to run on the same node as the original job, right? This is hard to do on some batch systems, without "going around" the batch system (i.e., ssh'ing directly to the node).

yes which is why we need to chat. We were thinking of a generic cleanup that could be registered with toil. Then any job on a node can check for failed jobs and handle them

If you find a crashed job (a job whose worker process crashed before being able to run the deferred functions), you can load its recovery function and run it. You are guaranteed to run it on the same node as the one the job was run on because otherwise you wouldn't have found the crashed job by looking at its left-over caching state. Note that this looking for crashed jobs  happens whenever any job is run, not just on retries of the crashed job. IOW, as long as there are any jobs to run on a node, cleanup will happen.

The problem arises during the end game when only zombies are left for retry. The retries may not be run on the same node as a previous attempt so on some worker nodes there is nothing triggering the code that looks for crashed jobs. I feel like this is where BatchSystem.supportsWorkerCleanup() comes in. The recovery, or garbage collection, of crashed jobs happens one last time during worker cleanup, on every node.

And yes, discuss over lunch.

This generally sounds good. The one thing that I would make clear when this is actually implemented is that the cleanup function isn't guaranteed to be executed (that whole `BatchSystem.supportsWorkerCleanup()`) bit. Essentially, what we're providing isn't a _job_ cleanup function, it is a _job invocation_ cleanup function. In a job cleanup function, one might do something like try to clean up a failed upload to S3, or etc. A job invocation cleanup function would just clean up the state of the job left on the node where it ran. Catch my drift?

> And yes, discuss over lunch.

Obviously, I won't catch you all at lunch ;). Can you post a summary of whatever you discuss back here? TIA!

>  In a job cleanup function, one might do something like try to clean up a failed upload to S3, or etc. A job invocation cleanup function would just clean up the state of the job left on the node where it ran. Catch my drift?

Then we need to think about "how failed is a job". Either that or we write a robust cleanup function.

> Then we need to think about "how failed is a job". Either that or we write a robust cleanup function.

What exactly do you mean here? I didn't quite follow.

> clean up a failed upload to S3
> you will clean up an upload  only if there was actually an upload that failed, right? Was there an upload in the job fn? what stage of the job did the job fail on? Did it reach an upload stage? If so, how do we find and handle the upload failure; If not what do we do?

Things like that.

> Things like that.

Oh yeah, I agree. That being said, you could imagine having a _job_ cleanup function that is guaranteed to run on _any_ host after the job fails (perhaps, guaranteed to run before the next retry of said failed job), and then having a _job invocation_ cleanup function that is guaranteed to run the next time a job starts on the host where said job invocation failed. The former would be sufficient for cleaning up things like failed uploads, while the latter would be sufficient for cleaning up things like ghost containers.

> you could imagine having a job cleanup function that is guaranteed to run on any host after the job fails

That is what we are targeting, i believe. The plan is to have a check at every job start that looks at the state of all jobs on the worker, and deals with dead/killed jobs and their containers(through the cleanup fn). This sort of ties up to 

>  a job invocation cleanup function that is guaranteed to run the next time a job starts on the host where said job invocation failed

as well

SGTM!

Sounds like you guys don't think we need the distinction between deferred and recovery functions. I tend to agree.

Hannes and I had a talk today and came up with a great idea of pickling/dilling job specific cleanup functions into the cache state file in the form of a list such that the cleanup of a job is on the node and is not integrated into the job store. Jobs detected to have been failed can be handled easily since the cleanup function is already defined on the node.

+1, thanks for the update @arkal 


Resolving patch:
diff --git a/src/toil/batchSystems/abstractBatchSystem.py b/src/toil/batchSystems/abstractBatchSystem.py
--- a/src/toil/batchSystems/abstractBatchSystem.py
+++ b/src/toil/batchSystems/abstractBatchSystem.py
@@ -23,6 +23,7 @@
 from bd2k.util.objects import abstractclassmethod
 
 from toil.common import Toil, cacheDirName
+from toil.fileStore import shutdownCache
 
 # A class containing the information required for worker cleanup on shutdown of the batch system.
 WorkerCleanupInfo = namedtuple('WorkerCleanupInfo', (
@@ -290,6 +291,7 @@ def workerCleanup(info):
         assert isinstance(info, WorkerCleanupInfo)
         workflowDir = Toil.getWorkflowDir(info.workflowID, info.workDir)
         workflowDirContents = os.listdir(workflowDir)
+        shutdownCache(os.path.join(workflowDir, cacheDirName(info.workflowID)))
         if (info.cleanWorkDir == 'always'
             or info.cleanWorkDir in ('onSuccess', 'onError')
             and workflowDirContents in ([], [cacheDirName(info.workflowID)])):
diff --git a/src/toil/fileStore.py b/src/toil/fileStore.py
--- a/src/toil/fileStore.py
+++ b/src/toil/fileStore.py
@@ -15,7 +15,7 @@
 from __future__ import absolute_import, print_function
 
 import base64
-import cPickle
+import dill
 import collections
 import errno
 import logging
@@ -38,6 +38,14 @@
 
 logger = logging.getLogger( __name__ )
 
+DeferredFunction = collections.namedtuple('DeferredFunction', (
+    # The deferred action to take in the form of a function
+    'callable',
+    # Non-keyword arguments to the function
+    'args',
+    # Keyword arguments to the function
+    'kwargs'))
+
 
 class FileStore(object):
     '''
@@ -82,7 +90,9 @@ def __init__(self, jobStore, jobWrapper, localTempDir, inputBlockFn):
         # a time on a worker, we can bookkeep the job's file store operated files in a
         # dictionary.
         self.jobSpecificFiles = {}
-        self.hashedJobCommand = sha1(self.jobWrapper.command.split()[1]).hexdigest()
+        self.jobName = self.jobWrapper.command.split()[1]
+        self.jobID = sha1(self.jobName).hexdigest()
+        logger.info('Starting job (%s) with ID (%s).', self.jobName, self.jobID)
         # A variable to describe how many hard links an unused file in the cache will have.
         self.nlinkThreshold = None
         self.workflowAttemptNumber = self.jobStore.config.workflowAttemptNumber
@@ -108,18 +118,11 @@ def open(self, job):
         # Create a working directory for the job
         startingDir = os.getcwd()
         self.localTempDir = makePublicDir(os.path.join(self.localTempDir, str(uuid.uuid4())))
-        # Figure out if this operation job already went through the process of writing it's info
-        # into the cache lock file.  If it has, restore the cache file to a state where the job
-        # doesn't exist.
+        # Check the status of all jobs on this node. If there are jobs that started and died before
+        # cleaning up their presence from the cache state file, restore the cache file to a state
+        # where the jobs don't exist.
         with self._CacheState.open(self) as cacheInfo:
-            if cacheInfo.jobState[self.hashedJobCommand]:
-                # Delete the old work directory if it still exists, to remove unwanted nlinks.
-                jobState = self._JobState(cacheInfo.jobState[self.hashedJobCommand])
-                assert jobState.jobDir != self.localTempDir
-                if os.path.exists(jobState.jobDir):
-                    shutil.rmtree(jobState.jobDir)  # Ignore_errors?
-                cacheInfo.sigmaJob -= jobState.jobReqs
-                cacheInfo.jobState.pop(self.hashedJobCommand)
+            self.findAndHandleDeadJobs(cacheInfo)
         # Get the requirements for the job and clean the cache if necessary. cleanCache will
         # ensure that the requirements for this job are stored in the state file.
         jobReqs = job.effectiveRequirements(self.jobStore.config).disk
@@ -131,7 +134,16 @@ def open(self, job):
         finally:
             os.chdir(startingDir)
             self.cleanupInProgress = True
+            # Delete all the job specific files and return sizes to jobReqs
             self.returnJobReqs(jobReqs)
+            with self._CacheState.open(self) as cacheInfo:
+                # Carry out any user-defined cleanup actions
+                deferredFunctions = cacheInfo.jobState[self.jobID]['deferredFunctions']
+                failures = self._runDeferredFunctions(deferredFunctions)
+                for failure in failures:
+                    self.logToMaster('Deferred function "%s" failed.' % failure, logging.WARN)
+                # Finally delete the job from the cache state file
+                cacheInfo.jobState.pop(self.jobID)
 
     # Functions related to temp file and directories
     def getLocalTempDir(self):
@@ -201,7 +213,7 @@ def writeGlobalFile(self, localFileName, cleanup=False):
             # the two have distinct nlink counts.
             # Can read without a lock because we're only reading job-specific info.
             jobSpecificFiles = self._CacheState._load(self.cacheStateFile).jobState[
-                self.hashedJobCommand]['filesToFSIDs'].keys()
+                self.jobID]['filesToFSIDs'].keys()
             # Saying nlink is 2 implicitly means we are using the job file store, and it is on
             # the same device as the work dir.
             if self.nlinkThreshold == 2 and absLocalFileName not in jobSpecificFiles:
@@ -328,9 +340,9 @@ def readGlobalFile(self, fileStoreID, userPath=None, cache=True, mutable=None):
                 if mutable:
                     shutil.copyfile(cachedFileName, localFilePath)
                     cacheInfo = self._CacheState._load(self.cacheStateFile)
-                    jobState = self._JobState(cacheInfo.jobState[self.hashedJobCommand])
+                    jobState = self._JobState(cacheInfo.jobState[self.jobID])
                     jobState.addToJobSpecFiles(fileStoreID, localFilePath, -1, None)
-                    cacheInfo.jobState[self.hashedJobCommand] = jobState.__dict__
+                    cacheInfo.jobState[self.jobID] = jobState.__dict__
                     cacheInfo.write(self.cacheStateFile)
                 else:
                     os.link(cachedFileName, localFilePath)
@@ -446,7 +458,7 @@ def deleteLocalFile(self, fileStoreID):
         # dict item having key = fileStoreID. If it was cached, it holds the value True else
         # False.
         with self._CacheState.open(self) as cacheInfo:
-            jobState = self._JobState(cacheInfo.jobState[self.hashedJobCommand])
+            jobState = self._JobState(cacheInfo.jobState[self.jobID])
             if fileStoreID not in jobState.jobSpecificFiles.keys():
                 # EOENT indicates that the file did not exist
                 raise OSError(errno.ENOENT, "Attempting to delete a non-local file")
@@ -459,7 +471,7 @@ def deleteLocalFile(self, fileStoreID):
                 if fileToDelete is None:
                     filesToDelete.pop(fileToDelete)
                     allOwnedFiles[fileToDelete].remove(fileStoreID)
-                    cacheInfo.jobState[self.hashedJobCommand] = jobState.__dict__
+                    cacheInfo.jobState[self.jobID] = jobState.__dict__
                     cacheInfo.write(self.cacheStateFile)
                     continue
                 # If the file size is zero (copied into the local temp dir) or -1 (mutable), we
@@ -477,7 +489,7 @@ def deleteLocalFile(self, fileStoreID):
                                 raise IllegalDeletionCacheError(fileToDelete)
                     allOwnedFiles[fileToDelete].remove(fileStoreID)
                     filesToDelete.pop(fileToDelete)
-                    cacheInfo.jobState[self.hashedJobCommand] = jobState.__dict__
+                    cacheInfo.jobState[self.jobID] = jobState.__dict__
                     cacheInfo.write(self.cacheStateFile)
                     continue
                 # If not, we need to do bookkeeping
@@ -496,7 +508,7 @@ def deleteLocalFile(self, fileStoreID):
                 filesToDelete.pop(fileToDelete)
                 allOwnedFiles[fileToDelete].remove(fileStoreID)
                 jobState.updateJobReqs(fileSize, 'remove')
-                cacheInfo.jobState[self.hashedJobCommand] = jobState.__dict__
+                cacheInfo.jobState[self.jobID] = jobState.__dict__
             # If the job is not in the process of cleaning up, then we may need to remove the
             # cached copy of the file as well.
             if not self.cleanupInProgress:
@@ -523,9 +535,12 @@ def deleteGlobalFile(self, fileStoreID):
 
         :param fileStoreID: the job store ID of the file to be deleted.
         """
+        jobStateIsPopulated = False
         with self._CacheState.open(self) as cacheInfo:
-            jobState = self._JobState(cacheInfo.jobState[self.hashedJobCommand])
-        if jobState.isPopulated() and fileStoreID in jobState.jobSpecificFiles.keys():
+            if self.jobID in cacheInfo.jobState:
+                jobState = self._JobState(cacheInfo.jobState[self.jobID])
+                jobStateIsPopulated = True
+        if jobStateIsPopulated and fileStoreID in jobState.jobSpecificFiles.keys():
             # Use deleteLocalFile in the backend to delete the local copy of the file.
             self.deleteLocalFile(fileStoreID)
             # At this point, the local file has been deleted, and possibly the cached copy. If
@@ -633,7 +648,7 @@ def _createCacheLockFile(self, tempCacheDir):
             'cached': 0,
             'sigmaJob': 0,
             'cacheDir': self.localCacheDir,
-            'jobState': collections.defaultdict(dict)})
+            'jobState': {}})
         cacheInfo.write(personalCacheStateFile)
 
     def encodedFileID(self, jobStoreFileID):
@@ -714,9 +729,9 @@ def addToCache(self, localFilePath, jobStoreFileID, callingFunc, mutable=None):
                 else:
                     logger.info('CACHE: Added file with ID \'%s\' to the cache.' %
                                 jobStoreFileID)
-                jobState = self._JobState(cacheInfo.jobState[self.hashedJobCommand])
+                jobState = self._JobState(cacheInfo.jobState[self.jobID])
                 jobState.addToJobSpecFiles(jobStoreFileID, localFilePath, -1, False)
-                cacheInfo.jobState[self.hashedJobCommand] = jobState.__dict__
+                cacheInfo.jobState[self.jobID] = jobState.__dict__
                 cacheInfo.write(self.cacheStateFile)
             else:
                 # There are two possibilities, read and immutable, and write. both cases do
@@ -782,9 +797,9 @@ def returnFileSize(self, fileStoreID, cachedFileSource, lockFileHandle,
             self.logToMaster('CACHE: The cache was not balanced on returning file size',
                              logging.WARN)
         # Add the info to the job specific cache info
-        jobState = self._JobState(cacheInfo.jobState[self.hashedJobCommand])
+        jobState = self._JobState(cacheInfo.jobState[self.jobID])
         jobState.addToJobSpecFiles(fileStoreID, cachedFileSource, fileSize, True)
-        cacheInfo.jobState[self.hashedJobCommand] = jobState.__dict__
+        cacheInfo.jobState[self.jobID] = jobState.__dict__
         cacheInfo.write(self.cacheStateFile)
 
     @staticmethod
@@ -817,13 +832,16 @@ def cleanCache(self, newJobReqs):
             # Initialize the job state here. we use a partial in the jobSpecificFiles call so
             # that this entire thing is pickleable. Based on answer by user Nathaniel Gentile at
             # http://stackoverflow.com/questions/2600790
-            assert not cacheInfo.jobState[self.hashedJobCommand]
-            cacheInfo.jobState[self.hashedJobCommand] = {
+            assert self.jobID not in cacheInfo.jobState
+            cacheInfo.jobState[self.jobID] = {
+                'jobName': self.jobName,
                 'jobReqs': newJobReqs,
                 'jobDir': self.localTempDir,
                 'jobSpecificFiles': collections.defaultdict(partial(collections.defaultdict,
                                                                     int)),
-                'filesToFSIDs': collections.defaultdict(set)}
+                'filesToFSIDs': collections.defaultdict(set),
+                'pid': os.getpid(),
+                'deferredFunctions': []}
             # If the caching equation is balanced, do nothing.
             if cacheInfo.isBalanced():
                 return None
@@ -906,12 +924,11 @@ def returnJobReqs(self, jobReqs):
         # Since we are only reading this job's specific values from the state file, we don't
         # need a lock
         jobState = self._JobState(self._CacheState._load(self.cacheStateFile
-                                                         ).jobState[self.hashedJobCommand])
+                                                         ).jobState[self.jobID])
         for x in jobState.jobSpecificFiles.keys():
             self.deleteLocalFile(x)
         with self._CacheState.open(self) as cacheInfo:
             cacheInfo.sigmaJob -= jobReqs
-            cacheInfo.jobState.pop(self.hashedJobCommand)
             # assert cacheInfo.isBalanced() # commenting this out for now. God speed
 
     def _abspath(self, key):
@@ -962,7 +979,7 @@ def _load(cls, fileName):
             # Read the value from the cache state file then initialize and instance of
             # _CacheState with it.
             with open(fileName, 'r') as fH:
-                cacheInfoDict = cPickle.load(fH)
+                cacheInfoDict = dill.load(fH)
             return cls(cacheInfoDict)
 
         def write(self, fileName):
@@ -977,7 +994,7 @@ def write(self, fileName):
                 # http://stackoverflow.com/questions/2709800/how-to-pickle-yourself
                 # We can't pickle nested classes. So we have to pickle the variables of the class
                 # If we ever change this, we need to ensure it doesn't break FileID
-                cPickle.dump(self.__dict__, fH, cPickle.HIGHEST_PROTOCOL)
+                dill.dump(self.__dict__, fH)
             os.rename(fileName + '.tmp', fileName)
 
         def isBalanced(self):
@@ -1001,6 +1018,79 @@ def purgeRequired(self, jobReqs):
             # totalFree = totalStats.f_bavail * totalStats.f_frsize
             # return totalFree < jobReqs
 
+    # Functions related to the deferred function logic
+    @classmethod
+    def findAndHandleDeadJobs(cls, cacheInfo, batchSystemShutdown=False):
+        """
+        This function looks at the state of all jobs registered in the cache state file and will
+        handle them (clean up the cache, and run any registered defer functions)
+
+        :param FileStore._CacheState cacheInfo: The state of the cache a _CacheState object
+        :param bool batchSystemShutdown: Is the batch system in the process of shutting down?
+        :return:
+        """
+        # A list of tuples of (hashed job id, pid or process running job)
+        registeredJobs = [(jid, state['pid']) for jid, state in cacheInfo.jobState.items()]
+        for jobID, jobPID in registeredJobs:
+            if not cls.HarbingerFile._pidExists(jobPID):
+                jobState = FileStore._JobState(cacheInfo.jobState[jobID])
+                logger.warning('Detected that job (%s) prematurely terminated.  Fixing the state '
+                               'of the cache.', jobState.jobName)
+                if not batchSystemShutdown:
+                    logger.debug("Returning dead job's used disk to cache.")
+                    # Delete the old work directory if it still exists, to remove unwanted nlinks.
+                    # Do this only during the life of the program and dont' do it during the
+                    # batch system cleanup.  Leave that to the batch system cleanup code.
+                    if os.path.exists(jobState.jobDir):
+                        shutil.rmtree(jobState.jobDir)  # Ignore_errors?
+                    cacheInfo.sigmaJob -= jobState.jobReqs
+                # Run any deferred functions associated with the job
+                logger.debug('Running user-defined deferred functions.')
+                deferredFunctions = jobState.deferredFunctions
+                cls._runDeferredFunctions(deferredFunctions)
+                # Remove it from the cache state file
+                cacheInfo.jobState.pop(jobID)
+
+    def _registerDeferredFunction(self, callable, *args, **kwargs):
+        """
+        This is the backend function to register a function that will be run after the completion of
+        the current job on the same node
+        :param function callable: The function to be run after this job.
+        :param list args: The arguments to the function
+        :param dict kwargs: The keyword arguments to the function
+        """
+        with self._CacheState.open(self) as cacheInfo:
+            deferredFunction = DeferredFunction(callable=callable,
+                                                args=args,
+                                                kwargs=kwargs)
+            cacheInfo.jobState[self.jobID]['deferredFunctions'].append(deferredFunction)
+            logger.info('Registered the deferred function "%s" to job "%s".', callable.__name__,
+                        self.jobName)
+
+    @staticmethod
+    def _runDeferredFunctions(deferredFunctions):
+        """
+        This function will run all the registered deferred functions for the job.
+
+        :param list deferredFunctions: A list of DeferredFunctions to run
+        the cache
+        :returns: list of failed functions
+        :rtype: list
+        """
+        failures = []
+        for deferredFunction in deferredFunctions:
+            #TODO: better representation of the callable in the output string
+            try:
+                logger.debug('Running deferred function "%s".', deferredFunction.callable.__name__)
+                deferredFunction.callable(*deferredFunction.args, **deferredFunction.kwargs)
+            except:
+                # This has to be generic because we don't know what is in the function
+                failures.append(deferredFunction.callable.__name__)
+                logger.exception('Deferred function "%s" failed.',
+                                 deferredFunction.callable.__name__)
+        return failures
+
+
     def _accountForNlinkEquals2(self, localFilePath):
         '''
         This is a utility function that accounts for the fact that if nlinkThreshold == 2, the
@@ -1013,7 +1103,7 @@ def _accountForNlinkEquals2(self, localFilePath):
         assert fileStats.st_nlink >= self.nlinkThreshold
         with self._CacheState.open(self) as cacheInfo:
             cacheInfo.sigmaJob -= fileStats.st_size
-            jobState = self._JobState(cacheInfo.jobState[self.hashedJobCommand])
+            jobState = self._JobState(cacheInfo.jobState[self.jobID])
             jobState.updateJobReqs(fileStats.st_size, 'remove')
 
     class _JobState(object):
@@ -1039,9 +1129,9 @@ def updateJobSpecificFiles(cls, outer, jobStoreFileID, filePath, fileSize, cache
             :return: None
             '''
             with outer._CacheState.open(outer) as cacheInfo:
-                jobState = cls(cacheInfo.jobState[outer.hashedJobCommand])
+                jobState = cls(cacheInfo.jobState[outer.jobID])
                 jobState.addToJobSpecFiles(jobStoreFileID, filePath, fileSize, cached)
-                cacheInfo.jobState[outer.hashedJobCommand] = jobState.__dict__
+                cacheInfo.jobState[outer.jobID] = jobState.__dict__
 
         def addToJobSpecFiles(self, jobStoreFileID, filePath, fileSize, cached):
             '''
@@ -1441,6 +1531,29 @@ def forPath(cls, fileStoreID, filePath):
         return cls(fileStoreID, os.stat(filePath).st_size)
 
 
+def shutdownCache(cacheDir):
+    """
+    Run the deferred functions from any prematurely terminated jobs still lingering on the system
+    and delete the cache directory.
+
+    This is a destructive operation and it is important to ensure that there are no other running
+    processes on the system that are modifying or using the cache state file.
+
+    This is the intended to be the last call to the file store in a Toil run, called by the
+    batch system cleanup function upon batch system shutdown.
+
+    :param cacheDir: The path to the cache directory
+    :return: None
+    """
+    if os.path.exists(cacheDir):
+        # The presence of the cacheDir suggests this was a cached run. We don't need the cache lock
+        # for any of this since this is the final cleanup of a job and there should be  no other
+        # conflicting processes using the cache.
+        cacheInfo = FileStore._CacheState._load(os.path.join(cacheDir, '_cacheState'))
+        FileStore.findAndHandleDeadJobs(cacheInfo, batchSystemShutdown=True)
+        shutil.rmtree(cacheDir)
+
+
 class CacheError(Exception):
     '''
     Error Raised if the user attempts to add a non-local file to cache
diff --git a/src/toil/job.py b/src/toil/job.py
--- a/src/toil/job.py
+++ b/src/toil/job.py
@@ -477,6 +477,35 @@ def checkNewCheckpointsAreLeafVertices(self):
                 if len(y._children) != 0 and len(y._followOns) != 0 and len(y._services) != 0:
                     raise JobGraphDeadlockException("New checkpoint job %s is not a leaf in the job graph" % y)
 
+    def defer(self, callable, *args, **kwargs):
+        """
+        Register a deferred function, i.e. a callable that will be invoked after the current attempt
+        at running this job concludes. A job attempt is said to conclude when the job function (or
+        the Job.run method for class-based jobs) returns, raises an exception or after the process
+        running it terminates abnormally. A deferred function will be called on the node that
+        attempted to run the job, even if a subsequent attempt is made on another node. A deferred
+        function should be idempotent because it may be called multiple times on the same node or
+        even in the same process. More than one deferred function may be registered per job attempt
+        by calling this method repeatedly with different arguments. If the same callable is
+        registered twice, it will be called twice per job attempt.
+
+        The functions one would typically provide here are cleanup functions that handle
+        Toil-external events upon a failure within Toil (killing Docker containers, etc).
+
+        :param function callable: The function to be run after this job.
+        :param list args: The arguments to the function
+        :param dict kwargs: The keyword arguments to the function
+        :return: None
+        """
+        try:
+            getattr(self, 'fileStore')
+        except AttributeError:
+            raise RuntimeError('A deferred function may only be registered from within the job it '
+                               'is being registered with. "%s" was illegally registered.',
+                               callable.__name__)
+        self.fileStore._registerDeferredFunction(callable, *args, **kwargs)
+
+
     ####################################################
     #The following nested classes are used for
     #creating jobtrees (Job.Runner),


Test patch:
diff --git a/src/toil/test/batchSystems/batchSystemTest.py b/src/toil/test/batchSystems/batchSystemTest.py
--- a/src/toil/test/batchSystems/batchSystemTest.py
+++ b/src/toil/test/batchSystems/batchSystemTest.py
@@ -79,8 +79,18 @@ def createBatchSystem(self):
         def supportsWallTime(self):
             return False
 
-        def _createDummyConfig(self):
-            return Config()
+        @staticmethod
+        def _createDummyConfig():
+            """
+            Returns a dummy config for the batch system tests.  We need a workflowID to be set up
+            since we are running tests without setting up a jobstore.
+
+            :rtype: toil.common.Config
+            """
+            config = Config()
+            from uuid import uuid4
+            config.workflowID = str(uuid4())
+            return config
 
         @classmethod
         def setUpClass(cls):
@@ -374,12 +384,12 @@ def test(self):
                     jobs = int(maxCores / coresPerJob * load)
                     if jobs >= 1 and minCores <= coresPerJob < maxCores:
                         self.assertEquals(maxCores, float(maxCores))
-                        bs = SingleMachineBatchSystem(config=Config(),
-                                                      maxCores=float(maxCores),
-                                                      # Ensure that memory or disk requirements
-                                                      # don't get in the way.
-                                                      maxMemory=jobs * 10,
-                                                      maxDisk=jobs * 10)
+                        bs = SingleMachineBatchSystem(
+                            config=hidden.AbstractBatchSystemTest._createDummyConfig(),
+                            maxCores=float(maxCores),
+                            # Ensure that memory or disk requirements don't get in the way.
+                            maxMemory=jobs * 10,
+                            maxDisk=jobs * 10)
                         try:
                             jobIds = set()
                             for i in range(0, int(jobs)):
diff --git a/src/toil/test/src/jobCacheTest.py b/src/toil/test/src/jobCacheTest.py
--- a/src/toil/test/src/jobCacheTest.py
+++ b/src/toil/test/src/jobCacheTest.py
@@ -28,7 +28,7 @@
 from uuid import uuid4
 
 from toil.job import Job
-from toil.fileStore import IllegalDeletionCacheError
+from toil.fileStore import IllegalDeletionCacheError, FileStore
 from toil.test import ToilTest, needs_aws, needs_azure, needs_google, experimental
 from toil.leader import FailedJobsException
 from toil.jobStores.abstractJobStore import NoSuchFileException
@@ -717,7 +717,7 @@ def _requirementsConcur(job, jobDisk, cached):
             state file is equal to the values we expect.
             """
             with job.fileStore._CacheState.open(job.fileStore) as cacheInfo:
-                jobState = cacheInfo.jobState[job.fileStore.hashedJobCommand]
+                jobState = cacheInfo.jobState[job.fileStore.jobID]
                 # cached should have a value only if the job store is on a different file system
                 # than the cache
                 if cacheInfo.nlink != 2:
@@ -871,6 +871,283 @@ def _deleteLocalFileFn(job, nonLocalDir):
             except NoSuchFileException:
                 pass
 
+        # These tests don't test the cache functions but they test code that heavily inherits cache
+        # code
+        def testDeferredFunctionRunsWithMethod(self):
+            """
+            Refer docstring in _testDeferredFunctionRuns.
+            Test with Method
+            """
+            self._testDeferredFunctionRuns(self._writeNonLocalFilesMethod)
+
+        def testDeferredFunctionRunsWithClassMethod(self):
+            """
+            Refer docstring in _testDeferredFunctionRuns.
+            Test with Class Method
+            """
+            self._testDeferredFunctionRuns(self._writeNonLocalFilesClassMethod)
+
+        def testDeferredFunctionRunsWithLambda(self):
+            """
+            Refer docstring in _testDeferredFunctionRuns.
+            Test with Lambda
+            """
+            self._testDeferredFunctionRuns(self._writeNonLocalFilesLambda)
+
+        def _testDeferredFunctionRuns(self, callableFn):
+            """
+            Create 2 files. Make a job that writes data to them. Register a deferred function that
+            deletes the two files (one passed as an arg, adn one as a kwarg) and later assert that
+            the files have been deleted.
+
+            :param function callableFn: The function to use in the test.
+            :return: None
+            """
+            workdir = self._createTempDir(purpose='nonLocalDir')
+            nonLocalFile1 = os.path.join(workdir, str(uuid4()))
+            nonLocalFile2 = os.path.join(workdir, str(uuid4()))
+            open(nonLocalFile1, 'w').close()
+            open(nonLocalFile2, 'w').close()
+            assert os.path.exists(nonLocalFile1)
+            assert os.path.exists(nonLocalFile2)
+            A = Job.wrapJobFn(callableFn, files=(nonLocalFile1, nonLocalFile2))
+            Job.Runner.startToil(A, self.options)
+            assert not os.path.exists(nonLocalFile1)
+            assert not os.path.exists(nonLocalFile2)
+
+        @staticmethod
+        def _writeNonLocalFilesMethod(job, files):
+            """
+            Write some data to 2 files.  Pass them to a registered deferred method.
+
+            :param tuple files: the tuple of the two files to work with
+            :return: None
+            """
+            for nlf in files:
+                with open(nlf, 'w') as nonLocalFileHandle:
+                    nonLocalFileHandle.write(os.urandom(1 * 1024 * 1024))
+            job.defer(_deleteMethods._deleteFileMethod, files[0], nlf=files[1])
+            return None
+
+        @staticmethod
+        def _writeNonLocalFilesClassMethod(job, files):
+            """
+            Write some data to 2 files.  Pass them to a registered deferred class method.
+
+            :param tuple files: the tuple of the two files to work with
+            :return: None
+            """
+            for nlf in files:
+                with open(nlf, 'w') as nonLocalFileHandle:
+                    nonLocalFileHandle.write(os.urandom(1 * 1024 * 1024))
+            job.defer(_deleteMethods._deleteFileClassMethod, files[0], nlf=files[1])
+            return None
+
+        @staticmethod
+        def _writeNonLocalFilesLambda(job, files):
+            """
+            Write some data to 2 files.  Pass them to a registered deferred Lambda.
+
+            :param tuple files: the tuple of the two files to work with
+            :return: None
+            """
+            lmd = lambda x, nlf: [os.remove(x), os.remove(nlf)]
+            for nlf in files:
+                with open(nlf, 'w') as nonLocalFileHandle:
+                    nonLocalFileHandle.write(os.urandom(1 * 1024 * 1024))
+            job.defer(lmd, files[0], nlf=files[1])
+            return None
+
+        def testDeferredFunctionRunsWithFailures(self):
+            """
+            Create 2 non local filesto use as flags.  Create a job that registers a function that
+            deletes one non local file.  If that file exists, the job SIGKILLs itself. If it doesn't
+            exist, the job registers a second deferred function to delete the second non local file
+            and exits normally.
+
+            Initially the first file exists, so the job should SIGKILL itself and neither deferred
+            function will run (in fact, the second should not even be registered). On the restart,
+            the first deferred function should run and the first file should not exist, but the
+            second one should.  We assert the presence of the second, then register the second
+            deferred function and exit normally.  At the end of the test, neither file should exist.
+
+            Incidentally, this also tests for multiple registered deferred functions, and the case
+            where a deferred function fails (since the first file doesn't exist on the retry).
+            """
+            self.options.retryCount = 1
+            workdir = self._createTempDir(purpose='nonLocalDir')
+            nonLocalFile1 = os.path.join(workdir, str(uuid4()))
+            nonLocalFile2 = os.path.join(workdir, str(uuid4()))
+            open(nonLocalFile1, 'w').close()
+            open(nonLocalFile2, 'w').close()
+            assert os.path.exists(nonLocalFile1)
+            assert os.path.exists(nonLocalFile2)
+            A = Job.wrapJobFn(self._deferredFunctionRunsWithFailuresFn,
+                              files=(nonLocalFile1, nonLocalFile2))
+            Job.Runner.startToil(A, self.options)
+            assert not os.path.exists(nonLocalFile1)
+            assert not os.path.exists(nonLocalFile2)
+
+        @staticmethod
+        def _deferredFunctionRunsWithFailuresFn(job, files):
+            """
+            Refer testDeferredFunctionRunsWithFailures
+
+            :param tuple files: the tuple of the two files to work with
+            :return: None
+            """
+            job.defer(hidden.AbstractCacheTest._deleteFile, files[0])
+            if os.path.exists(files[0]):
+                os.kill(os.getpid(), signal.SIGKILL)
+            else:
+                assert os.path.exists(files[1])
+                job.defer(hidden.AbstractCacheTest._deleteFile, files[1])
+
+        @staticmethod
+        def _deleteFile(nonLocalFile, nlf=None):
+            """
+            Delete nonLocalFile and nlf
+            :param str nonLocalFile:
+            :param str nlf:
+            :return: None
+            """
+            os.remove(nonLocalFile)
+            if nlf is not None:
+                os.remove(nlf)
+
+        def testNewJobsCanHandleOtherJobDeaths(self):
+            """
+            Create 2 non-local files and then create 2 jobs. The first job registers a deferred job
+            to delete the second non-local file, deletes the first non-local file and then kills
+            itself.  The second job waits for the first file to be deleted, then sleeps for a few
+            seconds and then spawns a child. the child of the second does nothing. However starting
+            it should handle the untimely demise of the first job and run the registered deferred
+            function that deletes the first file.  We assert the absence of the two files at the
+            end of the run.
+            """
+            # There can be no retries
+            self.options.retryCount = 0
+            workdir = self._createTempDir(purpose='nonLocalDir')
+            nonLocalFile1 = os.path.join(workdir, str(uuid4()))
+            nonLocalFile2 = os.path.join(workdir, str(uuid4()))
+            open(nonLocalFile1, 'w').close()
+            open(nonLocalFile2, 'w').close()
+            assert os.path.exists(nonLocalFile1)
+            assert os.path.exists(nonLocalFile2)
+            files = [nonLocalFile1, nonLocalFile2]
+            root = Job()
+            A = Job.wrapJobFn(self._testNewJobsCanHandleOtherJobDeaths_A, files=files)
+            B = Job.wrapJobFn(self._testNewJobsCanHandleOtherJobDeaths_B, files=files)
+            C = Job.wrapJobFn(self._testNewJobsCanHandleOtherJobDeaths_C, files=files,
+                              expectedResult=False)
+            root.addChild(A)
+            root.addChild(B)
+            B.addChild(C)
+            try:
+                Job.Runner.startToil(root, self.options)
+            except FailedJobsException as e:
+                pass
+
+        @staticmethod
+        def _testNewJobsCanHandleOtherJobDeaths_A(job, files):
+            """
+            Defer deletion of files[1], then wait for _testNewJobsCanHandleOtherJobDeaths_B to
+            start up, and finally delete files[0] before sigkilling self.
+
+            :param tuple files: the tuple of the two files to work with
+            :return: None
+            """
+            # Write the pid to files[1] such that we can be sure that this process has died before
+            # we spawn the next job that will do the cleanup.
+            with open(files[1], 'w') as fileHandle:
+                fileHandle.write(str(os.getpid()))
+            job.defer(hidden.AbstractCacheTest._deleteFile, files[1])
+            while os.stat(files[0]).st_size == 0:
+                time.sleep(0.5)
+            os.remove(files[0])
+            os.kill(os.getpid(), signal.SIGKILL)
+
+        @staticmethod
+        def _testNewJobsCanHandleOtherJobDeaths_B(job, files):
+            # Write something to files[0] such that we can be sure that this process has started
+            # before _testNewJobsCanHandleOtherJobDeaths_A kills itself.
+            with open(files[0], 'w') as fileHandle:
+                fileHandle.write(str(os.getpid()))
+            while os.path.exists(files[0]):
+                time.sleep(0.5)
+            # Get the pid of _testNewJobsCanHandleOtherJobDeaths_A and wait for it to truly be dead.
+            with open(files[1], 'r') as fileHandle:
+                meeseeksPID = int(fileHandle.read())
+            while FileStore.HarbingerFile._pidExists(meeseeksPID):
+                time.sleep(0.5)
+            # Now that we are convinced that_testNewJobsCanHandleOtherJobDeaths_A has died, we can
+            # spawn the next job
+            return None
+
+        @staticmethod
+        def _testNewJobsCanHandleOtherJobDeaths_C(job, files, expectedResult):
+            """
+            Asserts whether the files exist or not.
+
+            :param Job job: Job
+            :param list files: list of files to test
+            :param bool expectedResult: Are we expecting the files to exist or not?
+            """
+            for testFile in files:
+                assert os.path.exists(testFile) is expectedResult
+
+        def testBatchSystemCleanupCanHandleWorkerDeaths(self):
+            """
+            Create a non-local files. Create a job that registers a deferred job to delete the file
+            and then kills itself.
+
+            Assert that the file is missing after the pipeline fails.
+            """
+            # There can be no retries
+            self.options.retryCount = 0
+            workdir = self._createTempDir(purpose='nonLocalDir')
+            nonLocalFile1 = os.path.join(workdir, str(uuid4()))
+            nonLocalFile2 = os.path.join(workdir, str(uuid4()))
+            # The first file has to be non zero or meseeks will go into an infinite sleep
+            file1 = open(nonLocalFile1, 'w')
+            file1.write('test')
+            file1.close()
+            open(nonLocalFile2, 'w').close()
+            assert os.path.exists(nonLocalFile1)
+            assert os.path.exists(nonLocalFile2)
+            A = Job.wrapJobFn(self._testNewJobsCanHandleOtherJobDeaths_A,
+                              files=(nonLocalFile1, nonLocalFile2))
+            try:
+                Job.Runner.startToil(A, self.options)
+            except FailedJobsException:
+                pass
+            assert not os.path.exists(nonLocalFile1)
+            assert not os.path.exists(nonLocalFile2)
+
+
+class _deleteMethods(object):
+    @staticmethod
+    def _deleteFileMethod(nonLocalFile, nlf=None):
+        """
+        Delete nonLocalFile and nlf
+
+        :return: None
+        """
+        os.remove(nonLocalFile)
+        if nlf is not None:
+            os.remove(nlf)
+
+    @classmethod
+    def _deleteFileClassMethod(cls, nonLocalFile, nlf=None):
+        """
+        Delete nonLocalFile and nlf
+
+        :return: None
+        """
+        os.remove(nonLocalFile)
+        if nlf is not None:
+            os.remove(nlf)
+
 
 class FileJobStoreCacheTest(hidden.AbstractCacheTest):
     def _getTestJobStore(self):