nodes-busy

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

'''
Author : Sara Willis, University of Arizona
Date   : April 11, 2023
'''

import subprocess, sys, os, random, datetime, getopt, math, re, yaml

from config.visualization import *
from input.input_flags import *

from display.usage import usage
from display.system_view import system_display
from display.single_node import single_node_display
from display.text import display_text_summary
from display.summary import display_system_summary

from slurm.squeue import get_user_jobids
from slurm.scontrol import get_scontrol_job_data, get_scontrol_node_data
from slurm.sinfo import get_partitions

from formatting.sorting import sort_node_types
from formatting.slurm_output_reformatting import merge

    
    
'''
------------------------------------------------------------------------------------------
##########################################################################################
                                    Program Executes Below                                
##########################################################################################
------------------------------------------------------------------------------------------
'''

if __name__ == '__main__':
    
    script_dir = os.path.dirname(os.path.abspath(__file__))
    system_custom_file = os.path.join(script_dir,"config/system_customization.yml")
    with open(system_custom_file, 'r') as file:
        system_specifications = yaml.safe_load(file)

    output_formatting_file = os.path.join(script_dir,"config/output_formatting.yml")
    with open(output_formatting_file, 'r') as file:
        output_formatting=yaml.safe_load(file)
        
    # Made job options a class instead of a function so it wouldn't have to be batted around 
    # like a beach ball
    job_options = Args(sys.argv[1:],output_formatting)

    user      = job_options.user
    group     = job_options.group
    partition = job_options.partition
    
    # Check if valid user
    if user != None:
        p = subprocess.Popen(['getent passwd '+user],stdout=subprocess.PIPE, shell=True)
        out,err = p.communicate()
        if out.decode('utf-8','ignore') == '':
            print("\nOops! User not found. Check your options and try again.\n")
            usage(1)
            
    # Check if valid group
    if group != None:
        p = subprocess.Popen(['getent group '+group],stdout=subprocess.PIPE, shell=True)
        out,err = p.communicate()
        if out.decode('utf-8','ignore') == '':
            print("\nOops! Group not found. Check your options and try again.\n")
            usage(1)
          
    # Check if valid partition.         
    valid_partitions = get_partitions()
    if partition != None and partition not in valid_partitions:
        print("\nOops! Partition not recognized.\n")
        usage(1)
        
    # If user wants to get random colors, a random seed is set
    if job_options.color_random != False:
        seed = random.randint(1,1000)
    else:
        seed = None

    # Pull job data using "scontrol show jobs --all"
    job_data = get_scontrol_job_data(system_specifications)
    
    # Pull job data using "scontrol show nodes --all"
    node_data = get_scontrol_node_data(system_specifications)
    
    # Merge the two dictionaries above to create one 
    merged = merge(job_data, node_data)

    # Set vis options
    vis = VisOptions(seed,job_options.use_ascii)
    vis.width = job_options.scaled_width
    
    # We won't need node_data or job_data again since everything's in merged
    del node_data, job_data

    if job_options.display_all == True:
        system_display(merged,job_options,vis,system_specifications)
        
    elif job_options.single_node == True:
        single_node_display(job_options.node,merged,job_options.use_ascii,vis,system_specifications)
        
    elif job_options.text == True:
        display_text_summary(merged,job_options,vis,system_specifications)
        
    elif job_options.summary == True:
        display_system_summary(merged,job_options,vis,system_specifications)
    else:
        print("It's unclear how we got here")