slurm_users.py

slurm_stats.py

import json
import pandas as pd
import subprocess
import sys

def get_msg(
    backticks=True   # whether to add backticks for Discord formatting or not
    ):
    "gets a list of cluster usage from squeue and creates a text message from it"
    a = json.loads(subprocess.check_output(['squeue','--json']).decode("utf8"))
    df = pd.DataFrame(a["jobs"])

    # tally up number of idle nodes
    tally_choice = 'rom1504'   
    if tally_choice == 'rom1504': # json method by rom1504
        a = json.loads(subprocess.check_output(['sinfo','--json']).decode("utf8"))
        num_idles = sum([1 for a in  a['nodes'] if a['state'] == 'idle' and 'gpu' in a['name'] and 'POWERED_DOWN' not in a['state_flags']])
    else:   # shell-based method by drscotthawley
        cmd = "sinfo | grep ' idle ' | awk '{print $4}' | head -n 1"
        ps = subprocess.Popen(cmd,shell=True,stdout=subprocess.PIPE,stderr=subprocess.STDOUT)
        num_idles = ps.communicate()[0].decode("utf-8").rstrip()

    preemptible_accounts = [e[0] for e in [l.split("|") for l in subprocess.check_output(['sacctmgr', 'list', '--parsable', 'Account']).decode("utf8").split("\n")] if len(e) >= 3 and e[2] == "root"]

    def group_per_user_name(df):
        return str(df.groupby(["account", "user_name"]).sum("node_count")[["node_count"]].sort_values("node_count"))

    preemptible = df[df["account"].isin(preemptible_accounts)]
    non_preemptible = df[~df["account"].isin(preemptible_accounts)]

    preemptible_count = sum(preemptible["node_count"].values)
    non_preemptible_count = sum(preemptible["node_count"].values)

    msg = "```\n" if backticks else ""
    msg += "Preemptible:\n"+group_per_user_name(preemptible)+"\n\n"
    msg += "Non-preemptible:\n"+group_per_user_name(non_preemptible)+"\n\n"
    msg += f"Idle: {num_idles} nodes\n"
    msg += f"Preemptible count (these jobs will be killed if needed by non preemtible): {preemptible_count} nodes\n"
    msg += f"Non pre emptible count: {non_preemptible_count} nodes\n"
    msg = msg+"\n```" if backticks else msg
    return msg

print(get_msg())

slurm_users.py

import json
import pandas as pd
import subprocess
a = json.loads(subprocess.check_output(['squeue','--json']).decode("utf8"))
df = pd.DataFrame(a["jobs"])
print(df.groupby("user_name").sum("node_count")[["node_count"]].sort_values("node_count"))

moved to https://github.com/rom1504/slurm-tracking-bot