Repository: databiosphere/toil
Issue: Make killBatchJobs consistent between batch systems

Currently, the leader treats killed batch jobs separately from jobs that ended by themselves. JobBatcher.killJobs kills jobs and then sends them to JobBatcher.processFinishedJob, with the exit status manually set to 1. Also, in each iteration of the main loop, processFinishedJobs is called on the latest updated job. This means that if batchSystem.killBatchJobs adds the killed jobs to the updated queue, they will be processed twice. (single machine currently does this)
I think JobBatcher.killJobs shouldn't call processFinishedJobs, but all batch systems should add the killed jobs to the updated jobs queue with exit status 1. This means that all finished jobs are processed in the same place.


Changed files:
src/toil/batchSystems/gridengine.py, src/toil/batchSystems/mesos/batchSystem.py, src/toil/batchSystems/singleMachine.py, src/toil/test/batchSystems/batchSystemTest.py, src/toil/batchSystems/abstractBatchSystem.py, src/toil/batchSystems/parasol.py

diff --git a/src/toil/batchSystems/abstractBatchSystem.py b/src/toil/batchSystems/abstractBatchSystem.py
index 0002e46a..25a1f5b7 100644
--- a/src/toil/batchSystems/abstractBatchSystem.py
+++ b/src/toil/batchSystems/abstractBatchSystem.py
@@ -90,7 +90,7 @@ class AbstractBatchSystem:
         """Gets a job that has updated its status,
         according to the job manager. Max wait gives the number of seconds to pause
         waiting for a result. If a result is available returns (jobID, exitValue)
-        else it returns None.
+        else it returns None. Does not return anything for jobs that were killed.
         """
         raise NotImplementedError('Abstract method: getUpdatedBatchJob')
 
diff --git a/src/toil/batchSystems/gridengine.py b/src/toil/batchSystems/gridengine.py
index 8c455f30..f2edc87e 100644
--- a/src/toil/batchSystems/gridengine.py
+++ b/src/toil/batchSystems/gridengine.py
@@ -302,7 +302,7 @@ class GridengineBatchSystem(AbstractBatchSystem):
         return self.worker.getRunningJobIDs()
 
     def getUpdatedBatchJob(self, maxWait):
-        i = self.updatedJobsQueue.get()
+        i = self.getFromQueueSafely(self.updatedJobsQueue, maxWait)
         logger.debug('UpdatedJobsQueue Item: %s', i)
         if i is None:
             return None
diff --git a/src/toil/batchSystems/mesos/batchSystem.py b/src/toil/batchSystems/mesos/batchSystem.py
index ae8fe889..8d5eb6fd 100644
--- a/src/toil/batchSystems/mesos/batchSystem.py
+++ b/src/toil/batchSystems/mesos/batchSystem.py
@@ -68,6 +68,11 @@ class MesosBatchSystem(AbstractBatchSystem, mesos.interface.Scheduler):
         # queue of jobs to kill, by jobID.
         self.killSet = set()
 
+        # contains jobs on which killBatchJobs were called,
+        #regardless of whether or not they actually were killed or
+        #ended by themselves.
+        self.intendedKill = set()
+
         # Dict of launched jobIDs to TaskData named tuple. Contains start time, executorID, and slaveID.
         self.runningJobMap = {}
 
@@ -127,6 +132,7 @@ class MesosBatchSystem(AbstractBatchSystem, mesos.interface.Scheduler):
             log.debug("passing tasks to kill to Mesos driver")
             self.killSet.add(jobID)
             localSet.add(jobID)
+            self.intendedKill.add(jobID)
 
             if jobID not in self.getIssuedBatchJobIDs():
                 self.killSet.remove(jobID)
@@ -180,6 +186,9 @@ class MesosBatchSystem(AbstractBatchSystem, mesos.interface.Scheduler):
             return None
         jobID, retcode = i
         self.updatedJobsQueue.task_done()
+        if jobID in self.intendedKill:
+            self.intendedKill.remove(jobID)
+            return self.getUpdatedBatchJob(maxWait)
         log.debug("Job updated with code {}".format(retcode))
         return i
 
diff --git a/src/toil/batchSystems/parasol.py b/src/toil/batchSystems/parasol.py
index 25a7530d..9f7fee5f 100644
--- a/src/toil/batchSystems/parasol.py
+++ b/src/toil/batchSystems/parasol.py
@@ -127,6 +127,12 @@ class ParasolBatchSystem(AbstractBatchSystem):
         worker.start()
         self.usedCpus = 0
         self.jobIDsToCpu = {}
+
+        #Set of jobs that have been issued but aren't known to
+        #have finished or been killed yet. Jobs that end
+        #by themselves are removed in getUpdatedJob, and jobs
+        #that are killed are removed in killBatchJobs.
+        self.runningJobs = set()
          
     def issueBatchJob(self, command, memory, cores, disk):
         """Issues parasol with job commands.
@@ -161,6 +167,7 @@ class ParasolBatchSystem(AbstractBatchSystem):
                 time.sleep(5)
         jobID = int(match.group(1))
         self.jobIDsToCpu[jobID] = cores
+        self.runningJobs.add(jobID)
         logger.debug("Got the parasol job id: %s from line: %s" % (jobID, line))
         logger.debug("Issued the job command: %s with (parasol) job id: %i " % (parasolCommand, jobID))
         return jobID
@@ -171,6 +178,8 @@ class ParasolBatchSystem(AbstractBatchSystem):
         """
         while True:
             for jobID in jobIDs:
+                if jobID in self.runningJobs:
+                    self.runningJobs.remove(jobID)
                 exitValue = popenParasolCommand("%s remove job %i" % (self.parasolCommand, jobID), runUntilSuccessful=False)[0]
                 logger.info("Tried to remove jobID: %i, with exit value: %i" % (jobID, exitValue))
             runningJobs = self.getIssuedBatchJobIDs()
@@ -215,6 +224,12 @@ class ParasolBatchSystem(AbstractBatchSystem):
         jobID = self.getFromQueueSafely(self.outputQueue2, maxWait)
         if jobID != None:
             self.outputQueue2.task_done()
+            if jobID[0] not in self.runningJobs:
+                #we tried to kill this job, but it ended by itself
+                #instead, so skip it.
+                return self.getUpdatedJob(maxWait)
+            else:
+                self.runningJobs.remove(jobID[0])
         return jobID
     
     @classmethod
diff --git a/src/toil/batchSystems/singleMachine.py b/src/toil/batchSystems/singleMachine.py
index fd7801a7..5dabff40 100644
--- a/src/toil/batchSystems/singleMachine.py
+++ b/src/toil/batchSystems/singleMachine.py
@@ -131,7 +131,8 @@ class SingleMachineBatchSystem(AbstractBatchSystem):
             finally:
                 log.debug('Finished job. self.coreFractions ~ %s and self.memory ~ %s',
                           self.coreFractions.value, self.memory.value)
-                self.outputQueue.put((jobID, statusCode))
+                if not info.killIntended:
+                    self.outputQueue.put((jobID, statusCode))
         log.debug('Exiting worker thread normally.')
 
     def issueBatchJob(self, command, memory, cores, disk):
diff --git a/src/toil/test/batchSystems/batchSystemTest.py b/src/toil/test/batchSystems/batchSystemTest.py
index 2810a5be..ffad2046 100644
--- a/src/toil/test/batchSystems/batchSystemTest.py
+++ b/src/toil/test/batchSystems/batchSystemTest.py
@@ -137,6 +137,12 @@ class hidden:
             # self.assertEqual([0], self.batchSystem.getRunningJobIDs().keys())
             self.batchSystem.killBatchJobs([jobID])
             self.assertEqual({}, self.batchSystem.getRunningBatchJobIDs())
+
+            #Make sure the batch system doesn't add killed jobs
+            #to the updated jobs queue if they were killed.
+            updatedJobID = self.batchSystem.getUpdatedBatchJob(5);
+            self.assertEqual(updatedJobID, None)
+
             # Make sure that killJob doesn't hang / raise KeyError on unknown job IDs
             self.batchSystem.killBatchJobs([0])