Created
March 18, 2024 05:13
-
-
Save mtreviso/56ff754e0b0b60f45992c29bc715efe5 to your computer and use it in GitHub Desktop.
Pretty Print for Slurm `squeue`. It uses [rich](https://github.com/Textualize/rich). You can install it globally via `sudo pip install rich`.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/usr/bin/env python3 | |
# -*- coding: utf-8 -*- | |
"""psqueue.py | |
This script displays squeue in a pretty table. | |
""" | |
import os | |
import subprocess | |
import sys | |
try: | |
from rich.console import Console | |
from rich.table import Table | |
except ImportError: | |
sys.exit("Please install rich as sudo: sudo pip install rich") | |
try: | |
command = ['sinfo', '--format=%N', '--noheader'] | |
result = subprocess.run(command, check=True, stdout=subprocess.PIPE, text=True) | |
NODE_NAMES = [x.strip() for x in result.stdout.split(',')] | |
except subprocess.CalledProcessError as e: | |
NODE_NAMES = [] | |
def run_squeue(args): | |
# Command construction | |
command = ['squeue'] + args | |
try: | |
# Running squeue command | |
result = subprocess.run(command, check=True, stdout=subprocess.PIPE, text=True) | |
return result.stdout | |
except subprocess.CalledProcessError as e: | |
print(f"Error running squeue: {e}") | |
sys.exit(1) | |
def parse_output(output): | |
# Splitting the output into lines | |
lines = output.strip().split('\n') | |
# Assuming first line is header | |
headers = lines[0].split() | |
# Get the index of the "NODELIST" column | |
nodelist_index = headers.index("NODELIST") if "NODELIST" in headers else None | |
# Parsing each row | |
rows = [] | |
for line in lines[1:]: | |
parts = line.split() | |
# in case the "NodeList" column is not present, add it | |
if nodelist_index is not None: | |
nodelist_str = parts[nodelist_index] | |
# dirty hack to add "None" to the nodelist column | |
if nodelist_str.strip() not in NODE_NAMES: | |
parts.insert(nodelist_index, "None") | |
rows.append(parts) | |
return headers, rows | |
def get_state_color(state): | |
if state in ["COMPLETED", "CD"]: | |
return "white" | |
elif state in ["RUNNING", "COMPLETING", "R", "CG"]: | |
return "green" | |
elif state in ["PENDING", "PD"]: | |
return "cyan" | |
return "bright_red" | |
def get_time_left_color(time_left): | |
def split_time(t): | |
if ':' in t: | |
p = t.split(':') | |
if len(p) == 3: | |
return int(p[0]), int(p[1]), int(p[2]) | |
elif len(p) == 2: | |
return 0, int(p[0]), int(p[1]) | |
return 0, 0, int(t) | |
if '-' in time_left: | |
days, time = time_left.split('-') | |
days = int(days) | |
hours, minutes, seconds = split_time(time) | |
else: | |
days = 0 | |
hours, minutes, seconds = split_time(time_left) | |
if days >= 1: | |
return "white" | |
elif hours >= 12: | |
return "navajo_white1" | |
elif hours >= 1: | |
return "light_salmon1" | |
elif minutes >= 30: | |
return "red1" | |
return "bright_red" | |
def display_table(headers, rows, title="squeue"): | |
console = Console() | |
table = Table(title=title, show_header=True, highlight=True) | |
# merge TRES_PER_JOB and TRES_PER_NODE | |
if "TRES_PER_JOB" in headers and "TRES_PER_NODE" in headers: | |
index_job = headers.index("TRES_PER_JOB") | |
index_node = headers.index("TRES_PER_NODE") | |
for row in rows: | |
if row[index_job] == "N/A": | |
row[index_job] = row[index_node] | |
row.pop(index_node) | |
headers.pop(index_node) | |
# if the number of columns of a row is greater than the number of headers, concat the last columns | |
for i, row in enumerate(rows): | |
if len(row) > len(headers): | |
rows[i] = row[:len(headers)-1] + [' '.join(row[len(headers)-1:])] | |
# Adding columns to the table | |
for header in headers: | |
if header == "TRES_PER_JOB": | |
table.add_column("GPUS", no_wrap=True) | |
else: | |
table.add_column(header, no_wrap=True) | |
# Get the current user's name to highlight it | |
current_user = os.getenv("USER") | |
# Get the index of the "USER" column | |
user_index = headers.index("USER") if "USER" in headers else None | |
reason_index = headers.index("REASON") if "REASON" in headers else None | |
state_index = headers.index("ST") if "ST" in headers else None | |
state_index = headers.index("STATE") if "STATE" in headers else state_index | |
tres_index = headers.index("TRES_PER_JOB") if "TRES_PER_JOB" in headers else None | |
# timeleft_index = headers.index("TIME_LEFT") if "TIME_LEFT" in headers else None | |
# Adding rows to the table | |
for row in rows: | |
# Replace "gres:gpu:" with "" | |
if tres_index is not None: | |
row[tres_index] = row[tres_index].replace("gres:gpu:", "") | |
# Replace "None" with empty string | |
if reason_index is not None and row[reason_index].lower() == "none": | |
row[reason_index] = "" | |
# Highlight the state | |
if state_index is not None: | |
state = row[state_index] | |
color = get_state_color(state) | |
row[state_index] = f"[{color}]{state}[/]" | |
# Highlight the time left | |
# if timeleft_index is not None: | |
# time_left = row[timeleft_index] | |
# color = get_time_left_color(time_left) | |
# row[timeleft_index] = f"[{color}]{time_left}[/]" | |
# Highlight the current user's name | |
style = None | |
if user_index is not None and row[user_index] == current_user: | |
style = "bold" | |
table.add_row(*row, style=style) | |
# Displaying the table | |
console.print(table) | |
if __name__ == "__main__": | |
# Getting additional arguments passed to the script | |
additional_args = sys.argv[1:] | |
# If no specific format is provided, use the default format | |
args_str = ' '.join(additional_args) | |
if '--Format' not in args_str and '-O' not in args_str: | |
additional_args += ["--Format=JobID,Partition,Name,UserName,QOS,TimeUsed,TimeLeft,NumCPUs,tres-per-job,tres-per-node,MinMemory,NodeList,State,Reason:64"] # noqa | |
# if '--format' not in args_str and '-o' not in args_str: | |
# additional_args += ["--format=%.7i %9P %35j %.8u %.12q %.12N %.12M %.12L %.5C %.7m %.4D %.12T %r"] | |
output = run_squeue(additional_args) | |
headers, rows = parse_output(output) | |
display_table(headers, rows, title="squeue "+args_str) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Example:
