schedulers.py

from datetime import datetime

from lofar.common.cache import cache

from lofar.sas.resourceassignment.database import radb

from lofar.mom.momqueryservice.momqueryrpc import MoMQueryRPC
from lofar.mom.momqueryservice.config import DEFAULT_MOMQUERY_BUSNAME, DEFAULT_MOMQUERY_SERVICENAME

from lofar.sas.resourceassignment.resourceassigner.resource_availability_checker import CouldNotFindClaimException

from lofar.mac.config import DEFAULT_OBSERVATION_CONTROL_BUS_NAME, DEFAULT_OBSERVATION_CONTROL_SERVICE_NAME
from lofar.mac.observation_control_rpc import ObservationControl


""" Scheduling is the process of looking for a suitable resource slot for a given task.

    The following code allows three levels of scheduling:

    	* The /BasicScheduler/ locates resources for a task between its specified start- and endtime.
    	* The /PriorityScheduler/ enhances the BasicScheduler, and kills lower-priority tasks to make
    	  room for the given task.
    	* The /DwellScheduler/ enhances the PriorityScheduler, by increasing the starttime
    	  until the task can be scheduled.

    Each level contains hooks to support the next.

    All schedulers modify the database directly, and commit() if a solution if found. If not, a rollback()
    is performed. To use:

        sched = <schedulerclass>(...)
        success = sched.allocate_resources(estimates)
"""


class ScheduleException(Exception):
    def __init__(self, message):
        super(ScheduleException, self).__init__(message)


class BasicScheduler:
    """ A Scheduler that allocates resources at a fixed time. Resources are searched for. """

    def __init__(self, task_id, resource_availability_checker, radbcreds=None):
        self.task_id = task_id
        self.resource_availability_checker = resource_availability_checker

        # We need a DIRECT connection to the database in order to do client-side (that's us) transactions
        self.radb = radb(radbcreds)

        # Cache our task info
        self.task = self.radb.getTasks(task_id=self.task_id)[0]
        self.starttime = self.task["starttime"]
        self.endtime = self.task["endtime"]

        # Any resources that we cannot schedule on for some reason
        self.unusable_resources = []

    def allocate_resources(self, estimates):
        # any resources we cannot schedule on, either because we've already
        # tried, or because there's already a claim of us there

        try:
            # pre process
            self._pre_process_allocation()

            # keep iterating until all estimates are turned into tentative claims
            while estimates:
                estimates = self._try_schedule(estimates)

            # post process
            self._post_process_allocation()
        except ScheduleException:
            self.radb.rollback()
            return False

        self.radb.commit()
        return True

    def _pre_process_allocation(self):
        pass

    def _post_process_allocation(self):
        # move all claims from tentative -> claimed
        if not self.radb.updateResourceClaims(task_id=self.task_id, claim_status="claimed", commit=False):
            raise ScheduleException("Failed to put tentative claims to claimed")

    def _get_resource_availability(self):
        """ Get the available resources on which we can schedule. """

        # available resources
        resources = self.radb.getResources(include_availability=True, claimable_capacity_lower_bound=self.starttime,
                                           claimable_capacity_upper_bound=self.endtime)

        # resources WE claimed (note that their size is already accounted for in 'resources')
        tentative_claims = self.radb.getResourceClaims(task_ids=self.task_id, status="tentative")
        tentative_resources = [c["resource_id"] for c in tentative_claims]

        # disable resources we cannot schedule on (we cannot schedule on resources with tentative
        # claims until we implement merging tentative claims)
        disabled_resources = self.unusable_resources + tentative_resources

        # disable resources for various reasons
        return [r for r in resources if r["id"] not in disabled_resources]

    def _finalise_claims(self, claims):
        for c in claims:
            c["status"] = "tentative"  # we MUST insert as "tentative", as radb demands so
            c["starttime"] = self.starttime
            c["endtime"] = self.endtime

    def _try_schedule(self, requested_resources):
        """ Schedule the estimates within the available resources.

            Returns the set of estimates that could not be scheduled.
        """

        available_resources = self._get_resource_availability()
        try:
            claims = self.resource_availability_checker.get_is_claimable(requested_resources, available_resources)
        except CouldNotFindClaimException:
            raise ScheduleException("Could not schedule")

        # add static info to all claims
        self._finalise_claims(claims)

        # insert all claims to reserve the resources in the next call to findfit and to find the conflicts according to
        # the DB
        claim_ids = self.radb.insertResourceClaims(self.task_id, claims, _, _, commit=False)

        # tie the claim ids to the estimates
        claim_to_estimates = {cid: claims[cid]["requested_resources"] for cid in claim_ids}

        # handle any conflicts. We need NOT resolve ALL conflicts: removing one conflict can free up more resources as
        # a by-product, in which case other conflicts can simply be shifted to those newly freed resources.
        conflict_claims = self.radb.getResourceClaims(task_ids=self.task_id, status="conflict")
        if not any([self._resolve_conflict(c) for c in conflict_claims]):
            # Could not resolve any conflict
            raise ScheduleException("Could not resolve any conflict")

        # remove conflicting claims (allowing the next iteration to propose alternatives). Note that _handle_conflicts
        # could have reduced the number of conflicting claims.
        conflict_claims = self.radb.getResourceClaims(task_ids=self.task_id, status="conflict")
        conflict_claim_ids = [c["id"] for c in conflict_claims]
        self.radb.deleteResourceClaims(conflict_claim_ids, commit=False)

        # return any estimates that we could not fullfil
        remaining_estimates = set().union([claim_to_estimates[cid] for cid in conflict_claim_ids])
        return remaining_estimates

    def _get_conflicting_claims_and_tasks(self, conflict_claim):
        """ Return all claims that are conflicting with ours, and return the tasks associated with those claims. """

        conflicting_claims = self.radb.get_conflicting_overlapping_claims(c)
        conflicting_task_ids = set([c["task_id"] for c in conflicting_claims])
        conflicting_tasks = self.radb.getTasks(task_ids=conflicting_task_ids)

        return conflicting_claims, conflicting_tasks


    def _resolve_conflict(self, conflict_claim):
        """ Resolve one conflict, making it is useful to try to schedule again.

            :returns True if the conflict might have been resolved.
            :returns False if nothing can be done.
        """

        return False

class PriorityScheduler(BasicScheduler):
    """ A Scheduler that searches for an allocation with a fixed start time, but flexible resources.
        Conflict resolution is done by killing jobs with lower priority. """

    def __init__(self, task_id, resource_availability_checker,
                 radbcreds=None,
                 mom_busname=DEFAULT_MOMQUERY_BUSNAME,
                 mom_servicename=DEFAULT_MOMQUERY_SERVICENAME,
                 observation_control_busname=DEFAULT_OBSERVATION_CONTROL_BUS_NAME,
                 observation_control_servicename=DEFAULT_OBSERVATION_CONTROL_SERVICE_NAME,
                 broker=None):

        super(PriorityScheduler, self).__init__(task_id, resource_availability_checker, radbcreds)

        self.momqueryservice = MoMQueryRPC(servicename=mom_servicename, busname=mom_busname, broker=broker, timeout=180)
        self.obscontrol = ObservationControl(sevicename=observation_control_servicename,
                                             busname=observation_control_busname,
                                             broker=broker,
                                             timeout=120)

    def open(self):
        self.momqueryservice.open()
        self.obscontrol.open()

    def close(self):
        self.obscontrol.close()
        self.momqueryservice.close()

    def _pre_process_allocation(self):
        super(PriorityScheduler, self)._pre_process_allocation()

        self.earliest_potential_starttime = datetime.max
        self.tasks_to_kill = []

    def _post_process_allocation(self):
        super(PriorityScheduler, self)._post_process_allocation()

        # kill all jobs in self.tasks_to_kill
        for t in set(self.tasks_to_kill):
            if t["type"] == "observation":
                self.obscontrol.abort_observation(t["otdb_id"])
            else:
                # Killing scheduled pipelines only makes sense when they use resources other than storage, which is not
                # the case for current pipelines
                raise ScheduleException("Cannot kill jobs of type %s yet" % t["type"])

    @cache
    def _my_task_priority(self):
        my_momid = self.task["mom_id"]
        priority_dict = self.momqueryservice.get_project_priorities_for_objects([my_momid])
        return priority_dict[my_momid]

    def _add_task_to_kill(self, task_id):
        # add it to the list to actually kill later
        self.tasks_to_kill.append(task_id)

        # emulate by setting task endtime to now()
        self.radb.updateTaskAndResourceClaims(task_id=task_id, endtime=now(), commit=False)

    def _propose_potential_starttime(self, newstarttime):
        """ Propose a new start time at which allocation could succeed. We only collect the earliest one proposed. """
        self.earliest_potential_starttime = min(self.earliest_potential_starttime, newstarttime)

    def _resolve_conflict(self, conflict_claim):
        # try to resolve the conflict, and mark any resources unavailable if resolution fails
        tasks_to_kill = self._kill_conflict_tasks(conflict_claim)

        if not tasks_to_kill:
            # record which resources cannot be used anymore, because we can't kill anything on it
            self.unusable_resources.append(conflict_claim["resource"])

        # kill all tasks as requested
        for k in tasks_to_kill:
            self._kill_task(k)

        # Return True if we killed anything
        return tasks_to_kill != []

    def _kill_conflict_tasks(self, conflict_claim):
        """
            Modify the RA database such that one conflicting claim is resolved.

            Returns (kill_task_list, unfreeable_resources)
        """

        if conflict_claim["resource_type_id"] == self.resource_availability_checker.resource_types['storage']:
            raise ScheduleException("Could not resolve conflict on storage resource")

        # find all conflicting claims & which tasks they belong to
        conflicting_claims, conflicting_tasks = self._get_conflicting_claims_and_tasks(conflict_claim)
        conflicting_task_momids = [t["mom_id"] for t in conflicting_tasks]

        # check which tasks we can kill
        task_priorities = self.momqueryservice.get_project_priorities_for_objects(conflicting_task_momids)
        kill_task_list = [t for t in conflicting_tasks if task_priorities[t["mom_id"]] < self._my_task_priority()]

        # update if we're blocked by an earlier task than we know so far
        earliest_potential_starttime = min([t["endtime"] for t in conflicting_tasks if t not in kill_task_list])
        self._propose_potential_starttime(earliest_potential_starttime)

        return kill_task_list


class DwellScheduler(PriorityScheduler):
    """ A Scheduler that searches for an allocation with a flexible start time.

        Example:

        sched = DwellScheduler(task_id, min_starttime, max_starttime, duration)
        success = sched.allocate_resources(estimates)
        starttime = sched.starttime
    """

    def __init__(self, task_id, min_starttime, max_starttime, duration, **kwargs):
        super(DwellScheduler, self).__init__(task_id, **kwargs)

        self.task_id = task_id
        self.min_starttime = min_starttime
        self.max_starttime = max_starttime
        self.duration = duration

    def _new_starttime(self, starttime):
        self.starttime = starttime
        self.endtime = starttime + self.duration

    def _post_process_allocation(self):
        super(DwellScheduler, self)._post_process_allocation()

        # Update the task specification and resources to the new starttime and endtime
        self.radb.updateTaskAndResourceClaims(task_id=self.task_id,
                                              starttime=self.starttime,
                                              endtime=self.endtime,
                                              commit=False)

    def allocate_resources(self, estimates):
        """ Scan between [min_starttime, max_starttime) to find the first
            possible slot where the task can be scheduled. """

        self._new_starttime(self.min_starttime)
        success = False

        while not success and self.starttime <= self.max_starttime:
            # Find a solution
            success = super(DwellScheduler, self).allocate_resources(estimates)
            if success:
                break

            # Try the next slot
            new_starttime = self.earliest_potential_starttime
            assert new_starttime > self.starttime, "DwellScheduler cannot advance"

            self._new_starttime(new_starttime)

        return success