mtreviso · March 18, 2024 05:12 · Mar 18, 2024
diff --git a/psinfo.py b/psinfo.py
@@ -0,0 +1,118 @@
+#!/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)