PopulateDatabase

#!/usr/bin/env python3


"""
Desc: Script to populate the database for the last N days of use.

Notes:
    - This script is built to be run on all users, en-masse

Authors:
    - Chance Nelson <chance-nelson@nau.edu>
"""


import subprocess
import sys
import time
import re
from datetime import datetime
from datetime import timedelta
import mysql.connector as mariadb
import configparser
import os
import json

def timeStringToSeconds(timeString):
    """
    Desc: Convert a time string to seconds total elapsed.

    Args:
        timeString (string): Timestamp of the form 'D-HH:MM:SS.SSS'

    Returns:
        Integer or floating point representation of seconds elapsed.
    """
    rightDecimal = 0
    timeInSeconds = 0
    timeIndex = 0

    timeParsed = timeString.split('.')

    if len(timeParsed) > 1:
        timeInSeconds += float('0.' + timeParsed[1])

    timeParsed = timeParsed[0]
    
    timeParsed = re.split('-|:', timeParsed)

    # Does elapsed contain days? (*1-*00:00:00)
    if len(timeParsed) == 4:
        timeInSeconds += int(timeParsed[0]) * 24 * 60 * 60

    # Does time contain hours?
    if len(timeParsed) >= 3:
        timeInSeconds += int(timeParsed[-3]) * 60 * 60  # Hours

    # Does time contain minutes?
    if len(timeParsed) >= 2:
        timeInSeconds += int(timeParsed[-2]) * 60       # Minutes
    
    # Does time contain seconds?
    if len(timeParsed) >= 1:
        timeInSeconds += int(timeParsed[-1])            # Seconds

    return timeInSeconds

def getSlurmAccountUsers(acc):
    """
    Desc: Get a list of slurm users who belong to a given account.

    Args:
        acc (string): name of slurm account to search

    Returns:
        List of account user names. None of account doesn't exist/no users
        belong to it.
    """
    out = subprocess.getstatusoutput('sshare -a -P')[1].split('\n')
    users = []

    for i in out:
        users.append(i.split('|'))
        
    users = [i for i in users if acc in i[0]]
    users = [i[1] for i in users if not i[1] == 'parent' and not i[1] == '']

    return users


def getSlurmAccounts():
    """
    Desc: Retrieve a list of all available slurm accounts in the cluster
    """
    out = subprocess.getstatusoutput('sshare -l -P')[1].split('\n')[1::]
    accounts = []

    for i in out:
        accounts.append(i.split('|'))
        
    account_names = []
    for i in accounts :
        if i[0].strip() not in account_names:
            account_names.append(i[0].strip())

    return account_names


def timeStringToSeconds(time_string):
    """
    Desc: Convert a time string to seconds total elapsed.

    Args:
        time_string (string): Timestamp of the form 'D-HH:MM:SS.SSS'

    Returns:
        Integer or floating point representation of seconds elapsed.
    """
    right_decimal = 0
    time_in_seconds = 0
    time_index = 0

    time_parsed = time_string.split('.')

    if len(time_parsed) > 1:
        time_in_seconds += float('0.' + time_parsed[1])

    time_parsed = time_parsed[0]
    
    time_parsed = re.split('-|:', time_parsed)

    # Does elapsed contain days? (*1-*00:00:00)
    if len(time_parsed) == 4:
        time_in_seconds += int(time_parsed[0]) * 24 * 60 * 60

    # Does time contain hours?
    if len(time_parsed) >= 3:
        time_in_seconds += int(time_parsed[-3]) * 60 * 60  # Hours

    # Does time contain minutes?
    if len(time_parsed) >= 2:
        time_in_seconds += int(time_parsed[-2]) * 60       # Minutes
    
    # Does time contain seconds?
    if len(time_parsed) >= 1:
        time_in_seconds += float(time_parsed[-1])          # Seconds

    return time_in_seconds


def getJobStats(cursor, user, account, start_date_time, track_single_core_jobs):
    """
    Desc: Retrieve the job stats for a list of users, and populate the gpuinfo
          table

    Args:
        user (string): User to search.
        account (string): Slurm account of user
        hoursBack (int) (optional): How many hours in the past to search back

    Returns:
        [float, float, float, float]
        List containing the stats for memory, cores, timelimit, and total,
        respectively.
    """
    cmd = 'jobstats -o Comment,End,AllocGRES -p'  # Add in username and make parsable
    cmd += " -u {0} -A {1}".format(user, account)

    if start_date_time:
        cmd += ' -S {}'.format(start_date_time.strftime('%m/%d/%y'))
        cmd += ' -E {}'.format((start_date_time + timedelta(1)).strftime('%m/%d/%y'))

    cmd += ' --state=COMPLETED'  # filter for only completed jobs on that day 
                                 # to prevent duplicates

    res       = subprocess.getstatusoutput(cmd)[1].split('\n')
    job_count = len(res) 

    user_stats = {
        'memory-req': None,
        'ideal-cpu-time':  None,
        'tlimit-req': None,
        'memory-use': None,
        'cpu-time':   None,
        'tlimit-use': None,
        'job-count':  job_count
    }

    # Get the resources required/used per job
    for job in res:
        job = job.split('|')
        # Skip failed jobs
        if len(job) != 13 or job[0] == 'JobID':
            user_stats['job-count'] -= 1
            continue

        mem_req       = None
        mem_use       = None
        core_req      = None 
        cpu_time      = None
        tlimit_req    = None
        tlimit_use    = None
        comment       = None # contains GPU usage info

        try:
            mem_req = int(job[2].replace('M', ''))
            mem_use = int(float(job[3].replace('M', '')))

            if mem_req > 0:
                if not user_stats['memory-req']:
                    user_stats['memory-req'] = 0   

                if not user_stats['memory-use']:
                    user_stats['memory-use'] = 0

                user_stats['memory-req'] += mem_req
                user_stats['memory-use'] += mem_use

        except:
            pass

        try:
            tlimit_req = int(timeStringToSeconds(job[6]))
            tlimit_use = int(timeStringToSeconds(job[7]))


            if tlimit_req > 0:
                if not user_stats['tlimit-req']:
                    user_stats['tlimit-req'] = 0   

                if not user_stats['tlimit-use']:
                    user_stats['tlimit-use'] = 0

                user_stats['tlimit-req'] += tlimit_req
                user_stats['tlimit-use'] += tlimit_use

        except:
            pass

        try:
            core_req = int(job[4])
            cpu_time = float(timeStringToSeconds(job[5]))

            if core_req > 1 or track_single_core_jobs:
                if not user_stats['ideal-cpu-time']:
                    user_stats['ideal-cpu-time'] = 0   

                if not user_stats['cpu-time']:
                    user_stats['cpu-time'] = 0

                user_stats['ideal-cpu-time'] += (core_req * tlimit_use)
                user_stats['cpu-time']  += cpu_time            
        
        except:
            pass

        # add row to gpuinfo (1 row per job)
        gpustats = dict()
        gpustats['username'] = user
        gpustats['account'] = account
        gpustats['date'] = job[10][0:job[10].index('T')]
        gpustats['jobid'] = job[0]
        comment = job[9]
        allocGRES = job[11]
        gpuinfo = None
        try:
            gpuinfo = json.loads(comment)
        except:
            pass
        if gpuinfo is not None and allocGRES.find('gpu') != -1:
            gpustats['ngpu'] = gpuinfo['ngpu']
            gpustats['gputime'] = 60.0 * gpuinfo['step'] * gpuinfo['ngpu'] * gpuinfo['gpu_util'] / 100.0
            try:
                gpustats['idealgpu'] = gpuinfo['ngpu'] * tlimit_use
            except:
                # tlimit_use must be None
                gpustats['idealgpu'] = gpuinfo['ngpu'] * 60.0 * gpuinfo['step']

            # insert query
            statement = 'INSERT IGNORE INTO gpuinfo ' + \
                '(username,account,date,jobid,ngpu,gputime,idealgpu) ' + \
                'VALUES (%s, %s, %s, %s, %s, %s, %s)'
            args = (gpustats['username'], gpustats['account'], gpustats['date'],
                    gpustats['jobid'], gpustats['ngpu'], gpustats['gputime'],
                    gpustats['idealgpu'])
            cursor.execute(statement, args)

    return user_stats


def updateDatabase(user, account, stats, date, cursor):
    """
    Desc: Insert some data into the database

    Args:
        user (string): Username associated with account
        account (string): Slurm account of user
        stats (dict): Dictionary of stats values
        date (datetime): Datetime object for stats
        cursor (cursor): SQL cursor
    """
    statement = 'INSERT IGNORE INTO jobs (username, account, date, ' + \
                'idealcpu, memoryreq, tlimitreq, cputime, tlimituse, ' + \
                'memoryuse, jobsum) VALUES (%s, %s, %s, %s, %s, ' + \
                '%s, %s, %s, %s, %s)'

    args = (user, account, date, stats['ideal-cpu-time'], stats['memory-req'],
            stats['tlimit-req'], stats['cpu-time'], stats['tlimit-use'],
            stats['memory-use'], stats['job-count'])
    cursor.execute(statement, args)


if __name__ == '__main__':
    config = configparser.ConfigParser()
    config.read(os.path.dirname(__file__) + str('/config.ini'))

    username = None
    password = None
    host     = None
    TRACK_SINGLE_CORE_JOBS = False
    

    try:
        username = config['jobstats-db']['username']
        host     = config['jobstats-db']['host']
        password = config['jobstats-db']['password']
        TRACK_SINGLE_CORE_JOBS = config['jobstats-db']['TRACK_SINGLE_CORE_JOBS'] == 'True'

    except:
        print("Error: configuration values not found")


    try:
        from_date = datetime.strptime(sys.argv[1], '%m/%d/%y')
   
    except:
        print("{0} date")
        exit()

    print("Processing jobs for", from_date)

    accounts = getSlurmAccounts()
    conn = mariadb.connect(user=username, password=password, host=host, autocommit=True)
    cursor = conn.cursor()

    cursor.execute('USE jobstats')
    account_total = len(accounts)

    for account in accounts:
        if account == 'orphan':
            continue
        user_total = len(getSlurmAccountUsers(account))
        for user in getSlurmAccountUsers(account):
            stats = getJobStats(cursor,
                                user,
                                account,
                                from_date,
                                TRACK_SINGLE_CORE_JOBS)

            if stats['job-count'] != 0:
                updateDatabase(user, account, stats, from_date, cursor)

    conn.commit()    
    cursor.close()
    conn.close()