text: # Copyright (C) 2015-2021 Regents of the University of California
  #
  # 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.
  import logging
  
  from toil.job import CheckpointJobDescription, JobDescription
  
  logger = logging.getLogger(__name__)
  
  
  class ToilState:
      """
      Holds the leader's scheduling information that does not need to be
      persisted back to the JobStore (such as information on completed and
      outstanding predecessors).
      
      Only holds JobDescription objects, not Job objects, and those
      JobDescription objects only exist in single copies.
      
      Everything else in the leader should reference JobDescriptions by ID
      instead of moving them around between lists.
      """
      def __init__(self, jobStore, rootJob, jobCache=None):
          """
          Loads the state from the jobStore, using the rootJob 
          as the root of the job DAG.
  
          The jobCache is a map from jobStoreID to JobDescription or None. Is
          used to speed up the building of the state when loading initially from
          the JobStore, and is not preserved.
  
          :param toil.jobStores.abstractJobStore.AbstractJobStore jobStore 
          :param toil.job.JobDescription rootJob
          """
          
          # Maps from successor (child or follow-on) jobStoreID to predecessor jobStoreID
          self.successorJobStoreIDToPredecessorJobs = {}
          
          # Hash of jobStoreIDs to counts of numbers of successors issued.
          # There are no entries for jobs without successors in this map.
          self.successorCounts = {}
  
          # This is a hash of service jobs, referenced by jobStoreID, to their predecessor job
          self.serviceJobStoreIDToPredecessorJob = {}
  
          # Hash of jobStoreIDs mapping to dict from service ID to service host JobDescription for issued services
          self.servicesIssued = {}
          
          # Jobs that are ready to be processed.
          # Stored as a dict from job store ID to a pair of (job, result status),
          # where a status other than 0 indicates that an error occurred when
          # running the job.
          self.updatedJobs = {} 
          
          # The set of totally failed jobs - this needs to be filtered at the
          # end to remove jobs that were removed by checkpoints
          self.totalFailedJobs = set()
          
          # Jobs (as jobStoreIDs) with successors that have totally failed
          self.hasFailedSuccessors = set()
          
          # The set of successors of failed jobs as a set of jobStoreIds
          self.failedSuccessors = set()
          
          # Set of jobs that have multiple predecessors that have one or more predecessors
          # finished, but not all of them. This acts as a cache for these jobs.
          # Stored as hash from jobStoreIDs to JobDescriptions
          self.jobsToBeScheduledWithMultiplePredecessors = {}
          
          ##Algorithm to build this information
          self._buildToilState(rootJob, jobStore, jobCache)
          
          
      def allJobDescriptions(self):
          """
          Returns an iterator over all JobDescription objects referenced by the
          ToilState, with some possibly being visited multiple times.
          """
          
          for item in self.serviceJobStoreIDToPredecessorJob.values():
              assert isinstance(item, JobDescription)
              yield item
              
          for item in (desc for mapping in self.servicesIssued.values() for desc in mapping.values()):
              assert isinstance(item, JobDescription)
              yield item
              
          for item in (pair[0] for pair in self.updatedJobs.values()):
              assert isinstance(item, JobDescription)
              yield item
              
          for item in self.totalFailedJobs:
              assert isinstance(item, JobDescription)
              yield item
              
          for item in self.jobsToBeScheduledWithMultiplePredecessors.values():
              assert isinstance(item, JobDescription)
              yield item
          
      def _buildToilState(self, jobDesc, jobStore, jobCache=None):
          """
          Traverses tree of jobs from the root JobDescription (rootJob) building the
          ToilState class.
  
          If jobCache is passed, it must be a dict from job ID to JobDescription
          object. Jobs will be loaded from the cache (which can be downloaded from
          the jobStore in a batch) instead of piecemeal when recursed into.
  
          :param jobDesc: Object representing a job.
          :param jobStore: Object inheriting toil.jobStores.abstractJobStore.AbstractJobStore.
          :param jobCache:
          :return:
          """
  
          def getJob(jobId):
              if jobCache is not None:
                  if jobId in jobCache:
                      return jobCache[jobId]
              return jobStore.load(jobId)
          
          # If the job description has a command, is a checkpoint, has services
          # or is ready to be deleted it is ready to be processed
          if jobDesc.command is not None or (isinstance(jobDesc, CheckpointJobDescription) and jobDesc.checkpoint is not None) or len(jobDesc.services) > 0 or jobDesc.nextSuccessors() is None:
              logger.debug('Found job to run: %s, with command: %s, with checkpoint: %s, '
                           'with  services: %s, with no next successors: %s', jobDesc.jobStoreID,
                           jobDesc.command is not None, isinstance(jobDesc, CheckpointJobDescription) and jobDesc.checkpoint is not None,
                           len(jobDesc.services) > 0, jobDesc.nextSuccessors() is None)
              self.updatedJobs[jobDesc.jobStoreID] = (jobDesc, 0)
  
              if isinstance(jobDesc, CheckpointJobDescription) and jobDesc.checkpoint is not None:
                  jobDesc.command = jobDesc.checkpoint
  
          else: # There exist successors
              logger.debug("Adding job: %s to the state with %s successors" % (jobDesc.jobStoreID, len(jobDesc.nextSuccessors())))
              
              # Record the number of successors
              self.successorCounts[jobDesc.jobStoreID] = len(jobDesc.nextSuccessors())
  
              def processSuccessorWithMultiplePredecessors(successor):
                  # If jobDesc is not reported as complete by the successor
                  if jobDesc.jobStoreID not in successor.predecessorsFinished:
  
                      # Update the successor's status to mark the predecessor complete
                      successor.predecessorsFinished.add(jobDesc.jobStoreID)
  
                  # If the successor has no predecessors to finish
                  assert len(successor.predecessorsFinished) <= successor.predecessorNumber
                  if len(successor.predecessorsFinished) == successor.predecessorNumber:
                      
                      # It is ready to be run, so remove it from the cache
                      self.jobsToBeScheduledWithMultiplePredecessors.pop(successorJobStoreID)
                      
                      # Recursively consider the successor
                      self._buildToilState(successor, jobStore, jobCache=jobCache)
              
              # For each successor
              for successorJobStoreID in jobDesc.nextSuccessors():
                  
                  # If the successor does not yet point back at a
                  # predecessor we have not yet considered it
                  if successorJobStoreID not in self.successorJobStoreIDToPredecessorJobs:
  
                      # Add the job as a predecessor
                      self.successorJobStoreIDToPredecessorJobs[successorJobStoreID] = [jobDesc]
                      
                      # We load the successor job
                      successor = getJob(successorJobStoreID)
                      
                      # If predecessor number > 1 then the successor has multiple predecessors
                      if successor.predecessorNumber > 1:
                          
                          # We put the successor job in the cache of successor jobs with multiple predecessors
                          assert successorJobStoreID not in self.jobsToBeScheduledWithMultiplePredecessors
                          self.jobsToBeScheduledWithMultiplePredecessors[successorJobStoreID] = successor
                          
                          # Process successor
                          processSuccessorWithMultiplePredecessors(successor)
                              
                      else:
                          # The successor has only this job as a predecessor so
                          # recursively consider the successor
                          self._buildToilState(successor, jobStore, jobCache=jobCache)
                  
                  else:
                      # We've already seen the successor
                      
                      # Add the job as a predecessor
                      assert jobDesc not in self.successorJobStoreIDToPredecessorJobs[successorJobStoreID]
                      self.successorJobStoreIDToPredecessorJobs[successorJobStoreID].append(jobDesc)
                      
                      # If the successor has multiple predecessors
                      if successorJobStoreID in self.jobsToBeScheduledWithMultiplePredecessors:
                          
                          # Get the successor from cache
                          successor = self.jobsToBeScheduledWithMultiplePredecessors[successorJobStoreID]
                          
                          # Process successor
                          processSuccessorWithMultiplePredecessors(successor)
  