#!/usr/bin/python
# /* vim: set expandtab tabstop=4 shiftwidth=4 softtabstop=4: */

###########################################################################
# slurm_meter_running
#
# Python-based Gratia probe for SLURM accounting database
# This probe reports ComputeElementRecords for running and waiting jobs
# 
# John Thiltges, 2012-Jun-19
# Based on condor_meter by Brian Bockelman
# 
# Copyright 2012 University of Nebraska-Lincoln. Released under GPL v2.
###########################################################################

from SlurmProbe import SlurmProbe

import optparse
import datetime

import gratia.common.Gratia as Gratia
import gratia.services.ComputeElement as ComputeElement
import gratia.services.ComputeElementRecord as ComputeElementRecord


def get_ce(server_id, slurm_version, cluster, time_now):
    r = ComputeElement.ComputeElement() 

    ce_name = Gratia.Config.getConfigAttribute('SlurmCEName')

    r.UniqueID('slurm-running:%s/%s' % (server_id, cluster))
    r.CEName(ce_name)
    r.Cluster(ce_name)
    r.HostName(ce_name)
    r.Timestamp(str(time_now) + "Z")
    r.LrmsType('SLURM')
    r.LrmsVersion(slurm_version)

    return r

def user_to_cer(user, server_id, time_now):
    r = ComputeElementRecord.ComputeElementRecord() 

    r.UniqueID('slurm-running:%s/%s' % (server_id, user['cluster']))
    r.VO(user['user'])
    r.Timestamp(str(time_now) + "Z")
    r.RunningJobs(user['cpus_running'])
    r.WaitingJobs(user['cpus_pending'])
    r.TotalJobs(user['cpus_running'] + user['cpus_pending'])

    return r

class SlurmMeterRunning(SlurmProbe):
    def parse_opts(self):
        parser = optparse.OptionParser(usage="%prog [options]")
        parser.add_option("-f", "--gratia_config", help="Location of the Gratia"
            " config [default: %default].",
            dest="gratia_config", default="/etc/gratia/slurm-running/ProbeConfig")
        parser.add_option("-s", "--sleep", help="Do a random amount of sleep, "
            "up to the specified number of seconds before running.",
            dest="sleep", default=0, type="int")
        parser.add_option("-v", "--verbose", help="Enable verbose logging to "
            "stdout.",
            default=False, action="store_true", dest="verbose")
        parser.add_option("-c", "--checkpoint", help="Only reports records past"
            " checkpoint; default is to report all records.",
            default=False, action="store_true", dest="checkpoint")

        # Options are stored into opts/args class variables
        return parser.parse_args()

    def main(self):
        time_now = datetime.datetime.utcnow()
        time_end = None
        server_id = self.get_db_server_id()
        slurm_version = self.get_slurm_version()

        # Send ComputeElement
        Gratia.Send(get_ce(server_id, slurm_version, self.cluster, time_now))

        # Loop over running jobs summarized by user, sending CERs
        # Also include users with jobs completed since the checkpoint
        for user in self.sacct.running_users(self.checkpoint.val if self.opts.checkpoint else None):
            r = user_to_cer(user, server_id, time_now)
            Gratia.Send(r)

            # The query sorted the results by time_end, so our last value will
            # be the greatest
            time_end = user['time_end']

        # If we found at least one record, but the time_end has not increased since
        # the previous run, increase the checkpoint by one so we avoid continually
        # reprocessing the last records.
        # (This assumes the probe won't be run more than once per second.)
        if self.checkpoint.val == time_end:
            self.checkpoint.val = time_end + 1

if __name__ == "__main__":
    SlurmMeterRunning().main()
