queue.py

Go to the documentation of this file.

from __future__ import division, absolute_import
    
__all__ = ["queue"]

import glob
from opscore.actor import ActorDispatcher
import os.path
from RO.Comm.TCPConnection import TCPConnection
from twistedActor import CommandError, ScriptRunner, BaseWrapper
from tcc.base import getJobsDir

_CurrJob = None

def queue(tccActor, userCmd):
    """Implement the Queue command

    @param[in] tccActor  tcc actor
    @param[in,out] userCmd  a twistedActor BaseCommand with parseCmd attribute
    """
    global _CurrJob

    cmdParam = userCmd.parsedCmd.paramDict["command"].valueList[0]
    cmdName = cmdParam.keyword.lower()
    if cmdName == "run":
        if _CurrJob and _CurrJob.isRunning:
            raise CommandError("%s is already running" % (_CurrJob.scriptRunner.name))

        # start running a new job
        name = cmdParam.valueList[0]
        _CurrJob = BatchJob(
            tccActor = tccActor,
            name = name,
            userCmd = userCmd,
        )

    elif cmdName == "list":
        # list available jobs
        jobDir = getJobsDir()
        jobPaths = glob.glob(os.path.join(jobDir, "*.py"))
        names = [os.path.basename(jp)[0:-3] for jp in jobPaths]
        hubMsg = "JobNames=%s" % (", ".join(names))
        userCmd.setState(userCmd.Done, hubMsg=hubMsg)

    elif cmdName == "stop":
        # stop current job
        if not _CurrJob:
            userCmd.setState(userCmd.Done, textMsg="No batch job running; nothing to do")
            return

        _CurrJob.stop(userCmd)

    elif cmdName == "status":
        # show job status
        if _CurrJob:
            name = _CurrJob.name
            jobState = _CurrJob.state
        else:
            name = '""'
            jobState = BatchJob.Done
        hubMsg = "JobStatus=%s, %s" % (name, jobState)
        userCmd.setState(userCmd.Done, hubMsg=hubMsg)

    else:
        raise CommandError("Unrecognized queues sub-command %r" % (cmdName,))


class BatchJob(object):
    """A TCC batch job

    Attributes:
    - scriptRunner: the current batch job script (a twistedActor.ScriptRunner)
    - dispatcher: a dispatcher that talks to the TCC
    """
    Connecting = "Connecting"
    ConnectionCancelled = "ConnectionCancelled"
    ConnectionFailed = "ConnectionFailed"
    Starting = "Starting"
    Running = "Running"
    Done = "Done"
    Failed = "Failed"
    Cancelled = "Cancelled"
    CancelStates = set((ConnectionCancelled, Cancelled))
    ErrorStates = CancelStates | set((ConnectionFailed, Failed))
    DoneStates = ErrorStates | set((Done,))

    def __init__(self, tccActor, name, userCmd):
        """Create and start a batch job

        @param[in] tccActor  tcc actor
        @param[in] name  job name (name of file in jobs dir, without the .py suffix
        @param[in,out] userCmd  user command; the command is reported done when the command starts running,
            or failed if it cannot be started.
        """
        self.name = name
        self._runCmd = userCmd # command used to start job
        self._stopCmd = None
        self.state = self.Connecting
        self.writeToUsers = tccActor.writeToUsers

        # parse batch job file
        self.scriptRunner = None
        filePath = os.path.join(getJobsDir(), self.name + ".py")
        self._srArgs = self._getScript(filePath)

        # connect to TCC
        self.dw = TCCDispatcherWrapper(
            userPort = tccActor.server.port,
            stateCallback = self._wrapperStateChanged,
            debug = False,
        )

    def stop(self, userCmd):
        """Abort the batch job

        @param[in,out] userCmd  user command; the command is reported done when the command starts running,
            or failed if it cannot be started.
        """
        self._stopCmd = userCmd
        if self.scriptRunner and self.scriptRunner.isExecuting:
            self.scriptRunner.cancel()
        elif not self.dw.isDone:
            self.dw.close()
        else:
            userCmd.setState(userCmd.Done, textMsg="No batch job running; nothing to do")

    @property
    def isDone(self):
        return self.state in self.DoneStates

    @property
    def didFail(self):
        return self.state in self.ErrorStates

    @property
    def wasCancelled(self):
        return self.state in self.CancelStates

    @property
    def isRunning(self):
        return self.state == self.Running

    def setState(self, state):
        if self.isDone:
            return

        self.state = state

        if self.didFail:
            msgCode = "f"
        elif self.isDone:
            msgCode = ":"
        else:
            msgCode = "i"
        self.writeToUsers(msgCode, "JobStatus=%s, %s" % (self.name, state))

        if self.isDone:
            if not self._runCmd.isDone:
                if self.didFail:
                    self._runCmd.setState(self._runCmd.Failed, textMsg=state)
                else:
                    self._runCmd.setState(self._runCmd.Done)
            if self._stopCmd and not self._stopCmd.isDone:
                if self.wasCancelled:
                    self._stopCmd.setState(self._stopCmd.Done)
                elif self.didFail:
                    self._stopCmd.setState(self._stopCmd.Failed, textMsg=state)
                else:
                    self._stopCmd.setState(self._stopCmd.Done)
        elif self.isRunning and not self._runCmd.isDone:
            self._runCmd.setState(self._runCmd.Done)

    def _wrapperStateChanged(self, dw):
        """The TCC device wrapper's state changed

        @param[in] dw  tcc device wrapper
        """
        if self.isDone:
            return

        if self._runCmd.isDone:
            # the user command finished before the wrapper finished connecting or failing;
            # close the wrapper and give up
            self.setState(self.ConnectionCancelled)
            dw.close()

        elif dw.didFail:
            # could not connect
            self.setState(self.ConnectionFailed)

        elif dw.isReady:
            self.setState(self.Starting)
            self.scriptRunner = ScriptRunner(
                name = self.name,
                dispatcher = self.dw.dispatcher,
                stateFunc = self._scriptStateChanged,
            **self._srArgs)
            self.scriptRunner.start()
        
    def _scriptStateChanged(self, sr):
        """ScriptRunner state changed
        """
        if sr.isDone:
            if sr.state == sr.Cancelled:
                self.setState(self.Cancelled)
            elif sr.didFail:
                self.setState(self.Failed)
            else:
                self.setState(self.Done)
            self.dw.close()
        elif sr.isExecuting:
            self.setState(self.Running)

    def _getScript(self, filePath):
        """Load a script file, returning the functions as a dict of twistedActor.ScriptRunner args

        @param[in] fileName  name of script file (path is relative to the standard jobs dir)
        """
        if not os.path.isfile(filePath):
            raise CommandError("Job file %r not found" % (filePath,))

        scriptLocals = {"__file__": filePath}
        execfile(filePath, scriptLocals)

        scriptArgDict = dict()

        scriptClass = scriptLocals.get("ScriptClass")
        if scriptClass:
            scriptArgDict["scriptClass"] = scriptClass
        else:
            for attrName in ("run", "init", "end"):
                attr = scriptLocals.get(attrName)
                if attr:
                    scriptArgDict["%sFunc" % attrName] = attr
                elif attrName == "run":
                    raise RuntimeError("%r has no %s function" % (filePath, attrName))
        return scriptArgDict

    def __repr__(self):
        return "%s(%s)" % (type(self).__name__, self.name)


class TCCDispatcherWrapper(BaseWrapper):
    """A wrapper for an opscore.ActorDispatcher talking to the TCC

    This wrapper is responsible for starting and stopping everything:
    - It builds a connection
    - It builds a dispatcher when the connection is ready
    - It stops the connection and dispatcher on close

    Public attributes include:
    - connection: the connection (RO.Conn.TwistedSocket.TCPConnection)
    - dispatcher: the actor dispatcher (twistedActor.ActorDispatcher); None until ready
    - readyDeferred: called when the dispatcher is ready
      (for tracking closure use the Deferred returned by the close method, or stateCallback).
    """
    def __init__(self,
        userPort,
        name = "tcc",
        readCallback = None,
        stateCallback = None,
        debug = False,
    ):
        """Construct a TCCDispatcherWrapper that manages everything

        @param[in] userPort  TCC port
        @param[in] readCallback  function to call when the actor dispatcher has data to read
        @param[in] stateCallback  function to call when connection state of of any socket changes;
            receives one argument: this actor wrapper
        @param[in] debug  print debug messages to stdout?
        """
        BaseWrapper.__init__(self,
            name = name,
            stateCallback = stateCallback,
            callNow = False,
            debug = debug,
        )
        self.actor = "tcc"
        self._readCallback = readCallback
        self.dispatcher = None # the ActorDispatcher, once it's built

        connection = TCPConnection(
            host = 'localhost',
            port = userPort,
            readLines = True,
            name = "tccjob",
        )
        self._makeDispatcher(connection)
        connection.addStateCallback(self._stateChanged)
        if self._readCallback:
            connection.addReadCallback(self._readCallback)
        connection.connect()
        self._stateChanged()

    def _makeDispatcher(self, connection):
        self.debugMsg("_makeDispatcher()")
        self.dispatcher = ActorDispatcher(
            connection = connection,
            name = self.actor, # name of keyword dictionary
        )

    @property
    def isReady(self):
        """Return True if the dispatcher is connected to the TCC
        """
        return self.dispatcher is not None and self.dispatcher.connection.isConnected

    @property
    def isDone(self):
        """Return True if the dispatcher is disconnected from the TCC
        """
        return self.dispatcher is not None and self.dispatcher.connection.isDisconnected

    @property
    def isFailing(self):
        """Return True if there is a failure
        """
        return self.dispatcher is not None and self.dispatcher.connection.didFail

    def _basicClose(self):
        """Close dispatcher and actor
        """
        if self.dispatcher:
            self.dispatcher.disconnect()