Created
March 18, 2024 05:12
-
-
Save mtreviso/cad9d43067e36fe590b00c80af3ecf0f to your computer and use it in GitHub Desktop.
Pretty Print for Slurm `sinfo`. It uses [rich](https://github.com/Textualize/rich). You can install it globally via `sudo pip install rich`.
This file contains 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 -*- | |
"""psinfo.py | |
This script displays sinfo in a pretty table. | |
""" | |
import platform | |
import subprocess | |
import sys | |
try: | |
from rich.console import Console | |
from rich.table import Table | |
except ImportError: | |
sys.exit("Please install rich: pip install rich") | |
def run_command(args): | |
# Command construction | |
command = ['sinfo', '-N'] + args | |
try: | |
result = subprocess.run(command, check=True, stdout=subprocess.PIPE, text=True) | |
return result.stdout | |
except subprocess.CalledProcessError as e: | |
print(f"Error running sinfo: {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() | |
# Parsing each row | |
rows = [line.split() for line in lines[1:]] | |
return headers, rows | |
def display_table(headers, rows, title="sinfo"): | |
console = Console() | |
table = Table(title=title, show_header=True, highlight=True) | |
# Renaming the headers | |
rename_map = { | |
"NODELIST": "NODE", | |
"GRES": "GPUS", | |
"GRES_USED": "GPUS_USED", | |
"FREE_MEM": "MEM_USED", | |
"MEMORY": "MEMORY", | |
"CPUS": "CPUS", | |
"CPU_LOAD": "CPU_LOAD", | |
"STATE": "STATE", | |
"REASON": "REASON", | |
} | |
headers = [rename_map.get(h, h) for h in headers] | |
# Adding columns to the table | |
for header in headers: | |
table.add_column(header, no_wrap=True) | |
# Get the current user's name to highlight it | |
current_node = platform.node() | |
# Get the index of the "USER" column | |
node_index = headers.index("NODE") if "NODE" in headers else None | |
gpus_index = headers.index("GPUS") if "GPUS" in headers else None | |
gpusused_index = headers.index("GPUS_USED") if "GPUS_USED" in headers else None | |
reason_index = headers.index("REASON") if "REASON" in headers else None | |
cpuload_index = headers.index("CPU_LOAD") if "CPU_LOAD" in headers else None | |
memused_index = headers.index("MEM_USED") if "MEM_USED" in headers else None | |
mem_index = headers.index("MEMORY") if "MEMORY" in headers else None | |
# Adding rows to the table | |
for row in rows: | |
# Replace "gpu:" with "" | |
if gpus_index is not None: | |
row[gpus_index] = row[gpus_index].replace("gpu:", "") | |
if gpusused_index is not None: | |
row[gpusused_index] = row[gpusused_index].replace("gpu:", "") | |
# Replace "None" with empty string | |
if reason_index is not None and row[reason_index].lower() == "none": | |
row[reason_index] = "" | |
# Add % symbol with 2 decimal points to CPU_LOAD | |
if cpuload_index is not None: | |
row[cpuload_index] = f"{row[cpuload_index]}%" | |
# Convert FREE_MEM and MEMORY to GB and display GB, and convert FREE_MEM to MEM_USED | |
if memused_index is not None and mem_index is not None: | |
y = int(row[mem_index]) / 1024 | |
x = (int(row[mem_index]) - int(row[memused_index])) / 1024 | |
row[memused_index] = f"{x:.2f} GB" | |
if mem_index is not None: | |
y = int(row[mem_index]) / 1024 | |
row[mem_index] = f"{y:.2f} GB" | |
# Highlight the current node's name | |
style = None | |
if node_index is not None and row[node_index] == current_node: | |
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=NodeList,GresUsed,Gres,FreeMem,Memory,CPUsLoad,CPUs,StateLong,Reason"] # noqa | |
output = run_command(additional_args) | |
headers, rows = parse_output(output) | |
display_table(headers, rows, title="sinfo "+args_str) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Example:
