Repository: databiosphere/toil
Issue: AzureJobStore.clean() is very slow

I'm attempting to restart a run with somewhere between several tens of thousands and hundreds of jobs in an AzureJobStore.

The first step of the restart process appears to be to "clean" the job store, which involves walking the entire job graph, in serial, before scheduling any jobs. With 100,000 jobs, and one (or more, with chunking) round-trip HTTP request to the strongly-consistent-but-slow Azure Tables system per job, all happening in serial, this could take a long time. It's already taken 25 minutes so far for me, which is time my cluster is just sitting there not doing useful work.

Ideally the cleaning process (and, if possible, all the scheduling) should be able to be done on-line, without deserializing all the jobs. It could just be implicit: if the modification timestamp on a job is earlier than the start time of the Toil run, we just fix up its retry count when we load it. At the very least the traversal should be able to be done in parallel, so we don't need to wait for the round-trip latency to the table service for every single job.


Changed files:
src/toil/jobStores/azureJobStore.py, src/toil/jobStores/abstractJobStore.py, src/toil/job.py, src/toil/utils/toilRestart.py, src/toil/test/jobStores/jobStoreTest.py, src/toil/leader.py

diff --git a/src/toil/job.py b/src/toil/job.py
index dc879445..fe75ff40 100644
--- a/src/toil/job.py
+++ b/src/toil/job.py
@@ -419,8 +419,13 @@ class Job(object):
             """
             setLoggingFromOptions(options)
             with setupToil(options, userScript=job.getUserScript()) as (config, batchSystem, jobStore):
+                logger.info("Downloading entire JobStore")
+                jobCache = {jobWrapper.jobStoreID: jobWrapper
+                    for jobWrapper in jobStore.jobs()}
+                logger.info("{} jobs downloaded.".format(len(jobCache)))
                 if options.restart:
-                    jobStore.clean(job._loadRootJob(jobStore)) #This cleans up any half written jobs after a restart
+                    #This cleans up any half written jobs after a restart
+                    jobStore.clean(job._loadRootJob(jobStore), jobCache=jobCache) 
                     rootJob = job._loadRootJob(jobStore)
                 else:
                     #Make a file to store the root jobs return value in
@@ -434,7 +439,9 @@ class Job(object):
                         fH.write(jobStoreFileID)
                     #Setup the first wrapper.
                     rootJob = job._serialiseFirstJob(jobStore)
-                return mainLoop(config, batchSystem, jobStore, rootJob)
+                    #Make sure it's cached
+                    jobCache[rootJob.jobStoreID] = rootJob
+                return mainLoop(config, batchSystem, jobStore, rootJob, jobCache=jobCache)
 
     class FileStore( object ):
         """
diff --git a/src/toil/jobStores/abstractJobStore.py b/src/toil/jobStores/abstractJobStore.py
index 30736fd1..380671ba 100644
--- a/src/toil/jobStores/abstractJobStore.py
+++ b/src/toil/jobStores/abstractJobStore.py
@@ -126,31 +126,57 @@ class AbstractJobStore(object):
 
     ##Cleanup functions
 
-    def clean(self, rootJobWrapper):
+    def clean(self, rootJobWrapper, jobCache=None):
         """
         Function to cleanup the state of a jobStore after a restart.
         Fixes jobs that might have been partially updated.
         Resets the try counts.
-        Removes jobs that are not successors of the rootJobWrapper. 
+        Removes jobs that are not successors of the rootJobWrapper.
+        
+        If jobCache is passed, it must be a dict from job ID to JobWrapper
+        object. Jobs will be loaded from the cache (which can be downloaded from
+        the jobStore in a batch) instead of piecemeal when recursed into.
         """
         # Iterate from the root jobWrapper and collate all jobs that are reachable from it
         # All other jobs returned by self.jobs() are orphaned and can be removed
         reachableFromRoot = set()
 
+        if jobCache is None:
+            logger.warning("Cleaning jobStore recursively. This may be slow.")
+
+        def getJob(jobId):
+            if jobCache is not None:
+                return jobCache[jobId]
+            else:
+                return self.load(jobId)
+                
+        def haveJob(jobId):
+            if jobCache is not None:
+                return jobCache.has_key(jobId)
+            else:
+                return self.exists(jobId)
+                
+        def getJobs():
+            if jobCache is not None:
+                return jobCache.itervalues()
+            else:
+                return self.jobs()
+
         def getConnectedJobs(jobWrapper):
             if jobWrapper.jobStoreID in reachableFromRoot:
                 return
             reachableFromRoot.add(jobWrapper.jobStoreID)
             for jobs in jobWrapper.stack:
                 for successorJobStoreID in map(lambda x: x[0], jobs):
-                    if successorJobStoreID not in reachableFromRoot and self.exists(
-                            successorJobStoreID):
-                        getConnectedJobs(self.load(successorJobStoreID))
+                    if successorJobStoreID not in reachableFromRoot and haveJob(successorJobStoreID):
+                        getConnectedJobs(getJob(successorJobStoreID))
 
+        logger.info("Checking job graph connectivity...")
         getConnectedJobs(rootJobWrapper)
+        logger.info("%d jobs reachable from root." % len(reachableFromRoot))
 
         # Cleanup the state of each jobWrapper
-        for jobWrapper in self.jobs():
+        for jobWrapper in getJobs():
             changed = False  # Flag to indicate if we need to update the jobWrapper
             # on disk
 
@@ -174,7 +200,7 @@ class AbstractJobStore(object):
             # those jobs from the stack (this cleans up the case that the jobWrapper
             # had successors to run, but had not been updated to reflect this)
             while len(jobWrapper.stack) > 0:
-                jobs = [command for command in jobWrapper.stack[-1] if self.exists(command[0])]
+                jobs = [command for command in jobWrapper.stack[-1] if haveJob(command[0])]
                 if len(jobs) < len(jobWrapper.stack[-1]):
                     changed = True
                     if len(jobs) > 0:
@@ -198,10 +224,14 @@ class AbstractJobStore(object):
                 changed = True
 
             if changed:  # Update, but only if a change has occurred
+                logger.critical("Repairing job: %s" % jobWrapper.jobStoreID)
                 self.update(jobWrapper)
 
         # Remove any crufty stats/logging files from the previous run
+        logger.info("Discarding old statistics and logs...")
         self.readStatsAndLogging(lambda x: None)
+        
+        logger.info("Job store is clean")
 
     ##########################################
     # The following methods deal with creating/loading/updating/writing/checking for the
diff --git a/src/toil/jobStores/azureJobStore.py b/src/toil/jobStores/azureJobStore.py
index 577d6889..97cd926c 100644
--- a/src/toil/jobStores/azureJobStore.py
+++ b/src/toil/jobStores/azureJobStore.py
@@ -39,7 +39,7 @@ from toil.jobStores.abstractJobStore import (AbstractJobStore, NoSuchJobExceptio
                                              NoSuchFileException)
 import toil.lib.encryption as encryption
 
-log = logging.getLogger(__name__)
+logger = logging.getLogger(__name__)
 
 credential_file_path = '~/.toilAzureCredentials'
 
@@ -63,8 +63,8 @@ def _fetchAzureAccountKey(accountName):
     try:
         return configParser.get('AzureStorageCredentials', accountName)
     except NoOptionError:
-        raise RuntimeError("No account key found for %s, please provide it in " +
-                           credential_file_path % accountName)
+        raise RuntimeError("No account key found for %s, please provide it in %s" %
+                           (accountName, credential_file_path))
 
 
 maxAzureTablePropertySize = 64 * 1024
@@ -82,7 +82,7 @@ class AzureJobStore(AbstractJobStore):
 
         # Table names have strict requirements in Azure
         self.namePrefix = self._sanitizeTableName(namePrefix)
-        log.debug("Creating job store with name prefix '%s'" % self.namePrefix)
+        logger.debug("Creating job store with name prefix '%s'" % self.namePrefix)
 
         # These are the main API entrypoints.
         self.tableService = TableService(account_key=self.account_key, account_name=accountName)
@@ -123,8 +123,21 @@ class AzureJobStore(AbstractJobStore):
         return self.namePrefix + self.nameSeparator + name
 
     def jobs(self):
-        for jobEntity in self.jobItems.query_entities():
+        
+        # How many jobs have we done?
+        total_processed = 0
+            
+        for jobEntity in self.jobItems.query_entities_auto():
+            # Process the items in the page
             yield AzureJob.fromEntity(jobEntity)
+            total_processed += 1
+            
+            if total_processed % 1000 == 0:
+                # Produce some feedback for the user, because this can take
+                # a long time on, for example, Azure
+                logger.info("Processed %d total jobs" % total_processed)
+            
+        logger.info("Processed %d total jobs" % total_processed)
 
     def create(self, command, memory, cores, disk,
                predecessorNumber=0):
@@ -532,7 +545,47 @@ class AzureTable(object):
             return self.__getattr__('get_entity')(**kwargs)
         except WindowsAzureMissingResourceError:
             return None
-
+    
+    def query_entities_auto(self, **kwargs):
+        """
+        An automatically-paged version of query_entities. The iterator just
+        yields all entities matching the query, occasionally going back to Azure
+        for the next page.
+        """
+        
+        # We need to page through the results, since we only get some of them at
+        # a time. Just like in the BlobService. See the only documentation
+        # available: the API bindings source code, at:
+        # https://github.com/Azure/azure-storage-python/blob/09e9f186740407672777d6cb6646c33a2273e1a8/azure/storage/table/tableservice.py#L385
+        
+        # These two together constitute the primary key for an item. 
+        next_partition_key = None
+        next_row_key = None
+    
+        while True:
+            # Get a page (up to 1000 items)
+            kwargs['next_partition_key'] = next_partition_key
+            kwargs['next_row_key'] = next_row_key
+            page = self.query_entities(**kwargs)
+            
+            for result in page:
+                # Yield each item one at a time
+                yield result
+                
+            if hasattr(page, 'x_ms_continuation'):
+                # Next time ask for the next page. If you use .get() you need
+                # the lower-case versions, but this is some kind of fancy case-
+                # insensitive dictionary.
+                next_partition_key = page.x_ms_continuation['NextPartitionKey']
+                next_row_key = page.x_ms_continuation['NextRowKey']
+            else:
+                # No continuation to check
+                next_partition_key = None
+                next_row_key = None
+            
+            if not next_partition_key and not next_row_key:
+                # If we run out of pages, stop
+                break
 
 class AzureBlobContainer(object):
     """
@@ -690,11 +743,11 @@ def retry_on_error(num_tries=5, retriable_exceptions=(socket.error, socket.gaier
             if last:
                 raise
             else:
-                log.info("Got a retriable exception %s, trying again" % e.__class__.__name__)
+                logger.info("Got a retriable exception %s, trying again" % e.__class__.__name__)
         except Exception as e:
             # For other exceptions, the retriable_check function determines whether to retry
             if retriable_check(e):
-                log.info("Exception %s passed predicate, trying again" % e.__class__.__name__)
+                logger.info("Exception %s passed predicate, trying again" % e.__class__.__name__)
             else:
                 raise
         else:
diff --git a/src/toil/leader.py b/src/toil/leader.py
index bf633b04..0fd8682b 100644
--- a/src/toil/leader.py
+++ b/src/toil/leader.py
@@ -269,7 +269,8 @@ class ToilState( object ):
     """
     Represents a snapshot of the jobs in the jobStore.
     """
-    def __init__( self, jobStore, rootJob ):
+    def __init__( self, jobStore, rootJob, jobCache=None):
+        
         # This is a hash of jobs, referenced by jobStoreID, to their predecessor jobs.
         self.successorJobStoreIDToPredecessorJobs = { }
         # Hash of jobs to counts of numbers of successors issued.
@@ -279,13 +280,25 @@ class ToilState( object ):
         # Jobs that are ready to be processed
         self.updatedJobs = set( )
         ##Algorithm to build this information
-        self._buildToilState(rootJob, jobStore)
+        logger.info("(Re)building internal scheduler state")
+        self._buildToilState(rootJob, jobStore, jobCache)
 
-    def _buildToilState(self, jobWrapper, jobStore):
+    def _buildToilState(self, jobWrapper, jobStore, jobCache=None):
         """
         Traverses tree of jobs from the root jobWrapper (rootJob) building the
         ToilState class.
+        
+        If jobCache is passed, it must be a dict from job ID to JobWrapper
+        object. Jobs will be loaded from the cache (which can be downloaded from
+        the jobStore in a batch) instead of piecemeal when recursed into.
         """
+        
+        def getJob(jobId):
+            if jobCache is not None:
+                return jobCache[jobId]
+            else:
+                return jobStore.load(jobId)
+                
         if jobWrapper.command != None or len(jobWrapper.stack) == 0: #If the jobWrapper has a command
             #or is ready to be deleted it is ready to be processed
             self.updatedJobs.add((jobWrapper, 0))
@@ -298,7 +311,7 @@ class ToilState( object ):
                     #predecessor we have not yet considered it, so we call the function
                     #on the successor
                     self.successorJobStoreIDToPredecessorJobs[successorJobStoreID] = [jobWrapper]
-                    self._buildToilState(jobStore.load(successorJobStoreID), jobStore)
+                    self._buildToilState(getJob(successorJobStoreID), jobStore, jobCache=jobCache)
                 else:
                     #We have already looked at the successor, so we don't recurse,
                     #but we add back a predecessor link
@@ -310,10 +323,14 @@ class FailedJobsException( Exception ):
         self.jobStoreString = jobStoreString
         self.numberOfFailedJobs = numberOfFailedJobs
 
-def mainLoop(config, batchSystem, jobStore, rootJobWrapper):
+def mainLoop(config, batchSystem, jobStore, rootJobWrapper, jobCache=None):
     """
     This is the main loop from which jobs are issued and processed.
     
+    If jobCache is passed, it must be a dict from job ID to pre-existing
+    JobWrapper objects. Jobs will be loaded from the cache (which can be
+    downloaded from the jobStore in a batch).
+    
     :raises: toil.leader.FailedJobsException if at the end of function their remain
     failed jobs
     
@@ -324,7 +341,7 @@ def mainLoop(config, batchSystem, jobStore, rootJobWrapper):
     #Get a snap shot of the current state of the jobs in the jobStore
     ##########################################
 
-    toilState = ToilState(jobStore, rootJobWrapper)
+    toilState = ToilState(jobStore, rootJobWrapper, jobCache=jobCache)
 
     ##########################################
     #Load the jobBatcher class - used to track jobs submitted to the batch-system
diff --git a/src/toil/test/jobStores/jobStoreTest.py b/src/toil/test/jobStores/jobStoreTest.py
index b06a8329..b002cabe 100644
--- a/src/toil/test/jobStores/jobStoreTest.py
+++ b/src/toil/test/jobStores/jobStoreTest.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from __future__ import absolute_import
+from __future__ import print_function
 from Queue import Queue
 from abc import abstractmethod, ABCMeta
 import hashlib
@@ -24,6 +25,8 @@ from threading import Thread
 import tempfile
 import uuid
 import shutil
+import time
+from unittest import skip
 from toil.common import Config
 from toil.jobStores.abstractJobStore import (AbstractJobStore, NoSuchJobException,
                                              NoSuchFileException)
@@ -423,6 +426,62 @@ class hidden:
                     urllib2.urlopen(urllib2.Request(url))
                 except:
                     self.fail()
+                    
+        def testCleanCache(self):
+            # Make a bunch of jobs
+            master = self.master
+
+            # Create parent job
+            rootJob = master.create('rootjob', 12, 34, 35)
+            
+            # Create a bunch of child jobs
+            for i in range(100):
+                child = master.create("child%s" % i, 23, 45, 46, 1)
+                rootJob.stack.append(((child.jobStoreID, 23, 45, 46, 1),))
+            master.update(rootJob)
+            
+            # See how long it takes to clean with no cache
+            noCacheStart = time.time()
+            master.clean(rootJob)
+            noCacheEnd = time.time()
+            
+            noCacheTime = noCacheEnd - noCacheStart
+            
+            # See how long it takes to clean with cache
+            jobCache = {jobWrapper.jobStoreID: jobWrapper
+                        for jobWrapper in master.jobs()}
+            cacheStart = time.time()
+            master.clean(rootJob, jobCache=jobCache)
+            cacheEnd = time.time()
+            
+            cacheTime = cacheEnd - cacheStart
+            
+            print("Without cache: %s With cache: %s" % (str(noCacheTime), str(cacheTime)))
+            
+            # Running with the cache should be faster.
+            self.assertTrue(cacheTime <= noCacheTime)
+            
+        @skip("too slow") # This takes a long time on the remote JobStores
+        def testManyJobs(self):
+            # Make sure we can store large numbers of jobs
+            
+            # Make a bunch of jobs
+            master = self.master
+
+            # Create parent job
+            rootJob = master.create('rootjob', 12, 34, 35)
+            
+            # Create a bunch of child jobs
+            for i in range(3000):
+                child = master.create("child%s" % i, 23, 45, 46, 1)
+                rootJob.stack.append(((child.jobStoreID, 23, 45, 46, 1),))
+            master.update(rootJob)
+            
+            # Pull them all back out again
+            allJobs = list(master.jobs())
+            
+            # Make sure we have the right number of jobs
+            self.assertEquals(len(allJobs), 3001)
 
         # Sub-classes may want to override these in order to maximize test coverage
 
diff --git a/src/toil/utils/toilRestart.py b/src/toil/utils/toilRestart.py
index 612d50b4..beaac122 100644
--- a/src/toil/utils/toilRestart.py
+++ b/src/toil/utils/toilRestart.py
@@ -57,8 +57,13 @@ def main():
     setLoggingFromOptions(options)
     options.restart = True
     with setupToil(options) as (config, batchSystem, jobStore):
-        jobStore.clean(Job._loadRootJob(jobStore))
-        mainLoop(config, batchSystem, jobStore, Job._loadRootJob(jobStore))
+        # Load the whole jobstore into memory in a batch
+        logger.info("Downloading entire JobStore")
+        jobCache = {jobWrapper.jobStoreID: jobWrapper
+            for jobWrapper in jobStore.jobs()}
+        logger.info("{} jobs downloaded.".format(len(jobCache)))
+        jobStore.clean(Job._loadRootJob(jobStore), jobCache=jobCache)
+        mainLoop(config, batchSystem, jobStore, Job._loadRootJob(jobStore), jobCache=jobCache)
     
 def _test():
     import doctest