#!/usr/bin/env python

### A specific account can be given as first argument (If you have access to partitions from more than one account)
### example for getting info from the chemistry account:
### ./sfree chemistry

import subprocess
import sys
import os

DEVNULL = open(os.devnull, "wb")

# Run bash command and parse output
def run(command):
    output = [line.split() for line in subprocess.Popen(command, shell=True, stdout=subprocess.PIPE,stderr=DEVNULL).communicate()[0].split("\n")]
    return output

def get_queues(lines):
    queues = []
    for line in lines[:-1]:
        queues.append(line[0])
    return queues

def parse_snodes(snodes):
    # catches an error where a node with down state is in the list:
    for i, node in enumerate(snodes[::-1]):
        if node[-1] == "0":
            snodes.pop(-i-1)

    cpu_list = []
    memory_list = []
    free_dict = {}
    node_list = [x[1] for x in snodes]
    memory_dict = get_memory_info(node_list)
    for line in snodes:
        partition = line[5]
        node_name = line[1]
        cpu_load = float(line[4])
        cpu_tot = int(line[2])
        cpu_alloc = int(line[3].split("/")[0])
        real_memory = int(line[0])
        threads_per_core = int(line[-1])
        alloc_memory = memory_dict[node_name]

        free_cpus = cpu_tot - cpu_alloc
        free_memory = real_memory - alloc_memory

        # checks if cpu load is within 30% of allocated cpus.
        cpu_load_state = check_cpu_load(cpu_load, cpu_alloc,threads_per_core)
        job_list = []

        if not cpu_load_state:
            job_list = get_job_info(node_name)
            cpu_list.append([node_name, cpu_alloc, cpu_load, partition, threads_per_core, job_list])

        # Check if requested memory results in cpus being unused on the nodes.
        # An arbitrary 2 gb of memory is assumed required for future jobs
        # If the script is expanded to take needed memory as user input, this should be changed.
        if free_cpus > 0:

            # check if allocated memory in percent is larger than allocated cpus in percent.
            if cpu_alloc > 0 and float(free_memory)/real_memory < float(free_cpus) / cpu_tot and float(free_memory * threads_per_core) / free_cpus < 2000:
                if len(job_list) == 0:
                    job_list = get_job_info(node_name)
                # get list of jobs where allocated memory in percent is larger than allocated cpus in percent.
                memory_misuse_list = check_memory_usage(job_list,real_memory,cpu_tot)
                memory_list.append([free_cpus, threads_per_core, node_name, partition, free_memory, memory_misuse_list])

            if free_memory > 0:
                # Only count cpu as available if 2gb ram is unallocated
                free_cpus = min(free_cpus, free_memory/2000)
                if partition not in free_dict.keys(): free_dict[partition] = []
                free_dict[partition].append([free_cpus, free_memory, threads_per_core])

    return free_dict, cpu_list, memory_list

# get list of jobs where allocated memory in percent is larger than allocated cpus in percent
def check_memory_usage(full_job_list,real_memory, cpu_tot):
    job_list = []
    for job in full_job_list:
        job_id, user, job_name, num_nodes, num_cpus, mem, time = job
        if float(mem) / num_cpus > float(real_memory) / cpu_tot:
            job_list.append(job)
    return job_list

# print number of jobs possible to submit on all queues, for 1,2,4,8 etc. cpus.
def print_free_dict(free_dict):
    print_list = []
    print_list.append("\n\nPossible submissions:\n")
    l = max([len(x) for x in free_dict.keys()])

    # get window size
    rows, columns = os.popen('stty size', 'r').read().split()
    columns = int(columns)

    size_dict = {}
    for partition in free_dict.keys():
        remaining_columns = columns
        ones = sum([x[0]/x[2] for x in free_dict[partition]])
        size_dict[2] = sum([x[0]/(2*x[2]) for x in free_dict[partition]])
        size_dict[4] = sum([x[0]/(4*x[2]) for x in free_dict[partition]])
        size_dict[8] = sum([x[0]/(8*x[2]) for x in free_dict[partition]])
        size_dict[12] = sum([x[0]/(12*x[2]) for x in free_dict[partition]])
        size_dict[16] = sum([x[0]/(16*x[2]) for x in free_dict[partition]])
        size_dict[24] = sum([x[0]/(24*x[2]) for x in free_dict[partition]])

        print_list.append("{0:{1:0d}s}: {2:3d} jobs on {3:1d} thread".format(partition, l, ones, 1))
        remaining_columns -= 22+l
        for size in sorted(size_dict.keys()):
            if remaining_columns > 22:
                print_list.append(", {0:3d} jobs on {1:1d} threads".format(size_dict[size], size))
                remaining_columns -= 23
        print_list.append("\n")
    print "".join(print_list)


# print which jobs that blocks cpu's from being used.
def print_memory_list(memory_list):
    print_lines = ["\nMisuse of memory:"]
    for node in memory_list:
        real_cpus = node[0] / node[1]
        if node[4] > 2000:
            while real_cpus > 0 and float(node[4]) / real_cpus < 2000:
                real_cpus -= 1
        print_lines.append("{0:5s} {1:10s}   {2:2d} CPUs blocked by: ".format(node[3], node[2], real_cpus))

        for job in node[-1]:
            print_lines.append("      User: {0:10s}   job-id: {1:7s}   Allocated memory: {2:6d}   Allocated CPUs: {3:2d}".format(job[1], job[0], job[5], job[4]/node[1]))
    if len(print_lines) > 1:
        for line in print_lines:
            print line

# prints the node where cpu load is more than 30% off compared to the number of allocated cpus
# prints info on all jobs running on that node.
def print_cpu_list(cpu_list):
    if len(cpu_list) == 0:
        return
    print_lines = ["\nPossible bad CPU utilization:"]
    for node in cpu_list:
        # flag to register if there's jobs that has been running for less than 10 minutes.
        # Flag to register if any jobs has been running for less than 10 minutes on the node, since cpu load takes a while to update.
        flag = False
        for job in node[5]:
            time_split = job[6].split(":")
            if not ("-" in time_split[0] or len(time_split) > 2 or int(time_split[-2]) >= 10):
                flag = True
                continue
        if flag:
            continue
        real_cpus = node[1]/node[4]
        print_lines.append("{0:5s} {1:10s} {2:12s} {3:<5d} {4:9s} {5:4.1f}".format(node[3], node[0], "Allocated CPUs:", real_cpus, "CPU load:", node[2]))
        for job in node[5]:
            print_lines.append("      User: {0:10s}   job-id: {1:6s}   job-name: {2:30s}".format(job[1], job[0], job[2]))
    if len(print_lines) > 1:
        for line in print_lines:
            print line


def get_job_info(node_name):
    job_list = []
    output = run("squeue | grep " + node_name)
    for job in output[:-1]:
        job_id = job[0]
        user = job[3]
        time = job[5]

        job_output = run("scontrol show jobid -d " + job_id)

        job_name = job_output[0][-1].split("=")[1]
        num_nodes = int(job_output[14][0].split("=")[1])
        num_cpus = int(job_output[14][1].split("=")[1])
        mem = int(job_output[16][2].split("=")[1])

        job_list.append((job_id, user, job_name, num_nodes, num_cpus, mem, time))
    return job_list

# checks if cpu load is within 30% of number of allocated cpus on node.
def check_cpu_load(cpu_load,cpu_alloc,threads_per_core):
    threshold = 0.3
    if cpu_alloc == 0:
        return True
    if cpu_load < 1-threshold:
        return False
    if ((cpu_alloc / cpu_load < 1+threshold) and (cpu_alloc / cpu_load > 1-threshold)) \
           or ((cpu_alloc / (threads_per_core*cpu_load) < 1+threshold) \
           and (cpu_alloc / (threads_per_core*cpu_load) > 1-threshold)):
       return True
    else:
        return False


def get_memory_info(node_list):
    memory_dict = {}
    command = "scontrol show nodes "
    for node in node_list[:-1]:
        command += node + ","
    command += node_list[-1]

    output = run(command)

    counter = 0
    for line in output:
        if len(line) == 5 and "AllocMem" in line[2]:
            memory_dict[node_list[counter]] = int(line[2].split("=")[1])
            counter +=1
    return memory_dict


# Get the partitions the user have access to.
command = "sacctmgr show associations format=Partition%-30 -n user=$USER"

# Only get partitions from the user given account
account = None
if len(sys.argv) > 1:
    account = sys.argv[1]
    command += " Account=" + account

slimits = run(command)
queues = get_queues(slimits)

# Analyze all nodes on the partitions
arg_string = "sinfo -h -o '%m %n %c %C %O %P %Z' -r -t mix,idle,alloc"
if len(queues) > 0:
    arg_string += " -p "
    for item in queues[:-1]:
        arg_string += item + ","
    arg_string += queues[-1]

# print possible bad cpu utilization and memory misuse.
snodes = run(arg_string)[:-1]

# get number of free cpus
free_dict, cpu_list, memory_list = parse_snodes(snodes)

print_cpu_list(cpu_list)
print_memory_list(memory_list)
print_free_dict(free_dict)
