generate_qos_cmds.py

#!/usr/bin/python

from __future__ import print_function

#-----------------------------------------------------------------------
#  Imports
#-----------------------------------------------------------------------
#  Standard modules
import sys, time, getopt
#  Custom module
import mysqldump_to_csv

#-----------------------------------------------------------------------
#  Constants
#-----------------------------------------------------------------------
verbose = False

#  Column labels for table 'qos_table"
QOS_TABLE_COLS = ("creation_time","mod_time","deleted","id","name","description","flags","grace_time","max_jobs_per_user","max_submit_jobs_per_user","max_tres_pj","max_tres_pn","max_tres_pu","max_tres_mins_pj","max_tres_run_mins_pu","min_tres_pj","max_wall_duration_per_job","grp_jobs","grp_submit_jobs","grp_tres","grp_tres_mins","grp_tres_run_mins","grp_wall","preempt","preempt_mode","priority","usage_factor","usage_thres")

#  Make the following column labels to sacctmgr options
COL_LABELS_TO_OPTS = {
   "grace_time": "GraceTime",
   "grp_jobs": "GrpJobs",
   "grp_submit_jobs": "GrpSubmitJobs",
   "grp_tres": "GrpTRES",
   "grp_tres_mins": "GrpTRESMins",
   "grp_tres_run_mins": "GrpTRESRunMins",
   "grp_wall": "GrpWall",
   "max_jobs_per_user": "MaxJobs",
   "max_submit_jobs_per_user": "MaxSubmitJobs",
   "max_tres_mins_pj": "MaxTRESMins",
   "max_tres_pj": "MaxTRESPerJob",
   "max_tres_pn": "MaxTRESPerNode",
   "max_tres_pu": "MaxTRESPerUser",
   "max_wall_duration_per_job": "MaxWall",
   "min_tres_pj": "MinTRESPerJob",
   "preempt": "Preempt",
   "preempt_mode": "PreemptMode",
   "priority": "Priority",
   "usage_factor": "UsageFactor",
   "usage_thres": "UsageThreshold",
}

SKIP_COLUMNS = set(("max_tres_run_mins_pu",))

#-----------------------------------------------------------------------
#  Functions
#-----------------------------------------------------------------------
def Usage():
		print("""
   Usage:  generate_qos_cmds.py -v  FILE|-

      Read plain or gzipped msyql database dump from FILE or on standard input
      ('-') and print out sacctmgr command to set GrpTres on QOS's
   
      You can find the daily archives of the Slurm MySQL databases here:
   
         /gpfs/gpfs2/backup/slurmdb/archive/*
   
      -v  Verbose

   EXAMPLE:

     generate_qos_cmds.py /gpfs/gpfs2/backup/slurmdb/archive/20200921-2300.slurmdb.sql.gz 
   
		""")
		sys.exit()

def format_key(k):
	return ''.join( s.capitalize() for s in k.split("_") )

def format_val(v):
	return v.replace("1=","cpu=").replace("2=","mem=").replace("4=","node=").replace("5=","gres/gpu=")

#  For each non-deleted table row, generate corresponding 'sacctmgr modify qos' command
def make_config_line(row):
	#  If this record was marked as deleted, return None
	if not row[2]==0: return None
	name = row[4]
	config_line = ''
	for key,val in zip(QOS_TABLE_COLS,row)[7:-2]:  #  Omit first 7 columns, and last 2 columns
		val = str(val).strip()
		if key in SKIP_COLUMNS: continue
		#test
		if val in ('','0'): continue  # Skip default values
		config_line += " %s=%s" % ( COL_LABELS_TO_OPTS[key], format_val(val) )
	if config_line:
		return "sacctmgr -i modify qos %s set %s flags=denyonlimit" % (name, config_line)
	else:
		return None
		

#-----------------------------------------------------------------------
#  Main
#-----------------------------------------------------------------------

def main():
	global verbose

	def print_verbose(*msgs):
		global verbose
		if verbose: print("VERBOSE %s:" % time.strftime("%Y-%m-%d %H:%M:%S"), *msgs)

	opts, args = getopt.getopt(sys.argv[1:],'v')
	verbose = '-v' in dict(opts)
	print_verbose("Starting...")
	if not args: Usage()

	print_verbose("Reading dump file ...")
	if verbose: mysqldump_to_csv.VERBOSE = True
	data = mysqldump_to_csv.parse_dump_file(args[0],{'include':['qos_table']})
	print_verbose("Finished reading dump file")
	headers, rows = data.get_table("qos_table")
	print_verbose("Have headers and rows")
	for row in rows:
		cmd = make_config_line(row)
		if cmd: print(cmd)

main()