Source code for smif.controller.job.serial_job_scheduler

"""Job Schedulers are used to run job graphs.

Runs a job graph by calling execute_model_step for each operation in order
"""
import itertools
import logging
import traceback
from collections import defaultdict

import networkx
from smif.controller.execute_step import (execute_model_before_step,
                                          execute_model_step)
from smif.model import ModelOperation


class SerialJobScheduler(object):
    """Run JobGraphs produced by a :class:`~smif.controller.modelrun.ModelRun`
    """
    def __init__(self, store=None):
        self._status = defaultdict(lambda: 'unstarted')
        self._id_counter = itertools.count()
        self.logger = logging.getLogger(__name__)
        self.store = store

    def add(self, job_graph, dry_run=False):
        """Add a JobGraph to the SerialJobScheduler and run directly

        Arguments
        ---------
        job_graph: :class:`networkx.graph`
        dry_run: boolean, optional
            If True, print job steps without running
        """
        job_graph_id = self._next_id()
        try:
            self._run(job_graph, job_graph_id, dry_run)
        except Exception as ex:
            self._status[job_graph_id] = 'failed'
            traceback.print_exc()
            return job_graph_id, ex

        return job_graph_id, None

    def kill(self, job_graph_id):
        """Kill a job_graph that is already running - not implemented

        Parameters
        ----------
        job_graph_id: int
        """
        raise NotImplementedError

    def get_status(self, job_graph_id):
        """Get job graph status

        Parameters
        ----------
        job_graph_id: int

        Returns
        -------
        dict: A message containing the status

        Notes
        -----
        Possible statuses:

        unstarted:
            Job graph has not yet started
        running:
            Job graph is running
        done:
            Job graph was completed succesfully
        failed:
            Job graph completed running with an exit code
        """
        return {'status': self._status[job_graph_id]}

    def _run(self, job_graph, job_graph_id, dry_run=False):
        """Run a job graph
        - sort the jobs into a single list
        - unpack model, data_handle and operation from each node
        """
        try:
            self.logger.profiling_start(
                'SerialJobScheduler._run()', 'graph_' + str(job_graph_id))
        except AttributeError:
            self.logger.info('START SerialJobScheduler._run():graph_%s', job_graph_id)

        self._status[job_graph_id] = 'running'

        for job_node_id, job in self._get_run_order(job_graph):
            self._run_job(job_node_id, job, dry_run)

        self._status[job_graph_id] = 'done'
        try:
            self.logger.profiling_stop(
                'SerialJobScheduler._run()', 'graph_' + str(job_graph_id))
        except AttributeError:
            self.logger.info(
                'STOP SerialJobScheduler._run():graph_%s', job_graph_id)

    def _run_job(self, job_node_id, job, dry_run=False):
        self.logger.info("Job %s", job_node_id)  # Call root logger to satisfy CLI test
        try:
            self.logger.profiling_start('SerialJobScheduler._run()', 'job_' + job_node_id)
        except AttributeError:
            self.logger.info('START SerialJobScheduler._run():job_%s', job_node_id)

        if job['operation'] == ModelOperation.SIMULATE:
            execute_model_step(
                job['modelrun_name'],
                job['model'].name,
                job['current_timestep'],
                job['decision_iteration'],
                self.store,
                dry_run
            )
        elif job['operation'] == ModelOperation.BEFORE_MODEL_RUN:
            execute_model_before_step(
                job['modelrun_name'],
                job['model'].name,
                self.store,
                dry_run
            )
        else:
            raise ValueError("Model operation not recognised", job)

        try:
            self.logger.profiling_stop('SerialJobScheduler._run()', 'job_' + job_node_id)
        except AttributeError:
            self.logger.info('STOP SerialJobScheduler._run():job_%s', job_node_id)

    def _next_id(self):
        return next(self._id_counter)

    @staticmethod
    def _get_run_order(graph):
        """Returns a list of jobs in a runnable order.

        Returns
        -------
        list
            A list of job nodes
        """
        try:
            # topological sort gives a single list from directed graph,
            # ignoring opportunities to run independent models in parallel
            run_order = networkx.topological_sort(graph)

            # list of Models (typically ScenarioModel and SectorModel)
            ordered_jobs = [(run, graph.nodes[run]) for run in run_order]
        except networkx.NetworkXUnfeasible:
            raise NotImplementedError("Job graphs must not contain cycles")

        return ordered_jobs