richqueue 0.7__tar.gz → 0.9__tar.gz

{richqueue-0.7 → richqueue-0.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: richqueue
-Version: 0.7
+Version: 0.9
 Summary: RichQueue
 Project-URL: Homepage, https://github.com/mwinokan/RichQueue/issues
 Project-URL: Bug Tracker, https://github.com/mwinokan/RichQueue/issues

{richqueue-0.7 → richqueue-0.9}/pyproject.toml

@@ -3,7 +3,7 @@ requires = ["hatchling"]
 build-backend = "hatchling.build"
 [project]
 name = "richqueue"
-version = "0.7"
+version = "0.9"
 authors = [
     { name = "Max Winokan", email = "max@winokan.com" },
 ]
@@ -30,4 +30,4 @@ include = [
 ]
 
 [project.scripts]
-rq = "richqueue.slurm:main"
+rq = "richqueue.main:main"

{richqueue-0.7 → richqueue-0.9}/richqueue/main.py

@@ -4,7 +4,7 @@ from rich.live import Live
 import time
 from .console import console
 from .tools import curry
-from .slurm import get_layout_pair, get_node_layout, get_hist_layout
+from .slurm import get_layout_pair, get_node_layout, get_hist_layout, get_job_layout
 
 # set up singletons
 app = Typer()
@@ -20,6 +20,7 @@ def show(
     loop: bool = True,
     hist: int | None = None,
     hist_unit: str = "weeks",
+    job: int | None = None,
     # screen: bool = True,
     # disappear: bool = True,
 ):
@@ -36,18 +37,21 @@ def show(
         "hist_unit": hist_unit,
         "screen": screen,
         "disappear": disappear,
+        "job": job,
     }
 
     # console.print(kwargs)
 
-    match (bool(idle), bool(hist)):
-        case (True, False):
+    match (bool(idle), bool(hist), bool(job)):
+        case (True, False, False):
             loop = False
             layout_func = get_node_layout
-        case (False, True):
+        case (False, True, False):
             loop = False
             layout_func = get_hist_layout
-        case (False, False):
+        case (False, False, True):
+            layout_func = get_job_layout
+        case (False, False, False):
             layout_func = curry(dual_layout, get_layout_pair)
         case _:
             raise Exception("Unsupported CLI options")

{richqueue-0.7 → richqueue-0.9}/richqueue/slurm.py

@@ -6,7 +6,7 @@ import json
 from .console import console
 from pathlib import Path
 import datetime
-from .table import running_job_table, history_job_table, node_table
+from .table import running_job_table, history_job_table, node_table, job_table
 from .tools import human_timedelta
 
 # from numpy import isnat
@@ -119,12 +119,32 @@ def get_node_layout(idle: bool = True, **kwargs):
     return Panel(table, expand=False)
 
 
+def get_job_layout(job: int, **kwargs):
+
+    df = get_squeue(job=job, **kwargs)
+
+    if "end_time" not in df:
+        df["end_time"] = None
+    df["run_time"] = add_run_time(df)
+
+    assert len(df) == 1
+
+    table = job_table(df.iloc[0], job=job, **kwargs)
+
+    return Panel(table, expand=False)
+
+
 ### GET QUEUE DFs
 
 
-def get_squeue(user: str | None = None, **kwargs) -> "pandas.DataFrame":
+def get_squeue(
+    user: str | None = None, job: int | None = None, **kwargs
+) -> "pandas.DataFrame":
 
-    if user:
+    if job:
+        assert isinstance(job, int)
+        command = f"squeue --job={job} --json"
+    elif user:
         command = f"squeue -u {user} --json"
     else:
         command = f"squeue --json"
@@ -134,13 +154,22 @@ def get_squeue(user: str | None = None, **kwargs) -> "pandas.DataFrame":
             [command], shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE
         )
         output = process.communicate()
+
+        if job and "slurm_load_jobs error: Invalid job id specified" in str(output[1]):
+            return get_sacct(user=user, job=job, **kwargs)
+
         payload = json.loads(output[0])
+
     except json.JSONDecodeError:
         # console.print('[orange1 bold]Warning: using example data')
+
+        if job:
+            example = "squeue_job.json"
+        else:
+            example = "squeue_long.json"
+
         payload = json.load(
-            open(
-                Path(__file__).parent.parent / "example_data" / "squeue_long.json", "rt"
-            )
+            open(Path(__file__).parent.parent / "example_data" / example, "rt")
         )
 
     global METADATA
@@ -168,23 +197,35 @@ def get_squeue(user: str | None = None, **kwargs) -> "pandas.DataFrame":
     extract_inner(df, "node_count", "number")
     extract_inner(df, "cpus_per_task", "number")
     extract_inner(df, "threads_per_core", "number")
+    extract_inner(df, "memory_per_node", "number")
+    extract_inner(df, "memory_per_cpu", "number")
+    extract_inner(df, "tasks", "number")
 
     extract_time(df, "start_time")
     extract_time(df, "submit_time")
-    extract_time(df, "time_limit")
+    # extract_inner(df, "time_limit", "number")
+    extract_time_limit(df, "time_limit")
 
     extract_list(df, "job_state")
+    extract_list(df, "exclusive")
+    extract_json(df, "exclusive")
 
     return df
 
 
 def get_sacct(
-    user: str | None = None, hist: int | None = 4, hist_unit: str = "weeks", **kwargs
+    user: str | None = None,
+    hist: int | None = 4,
+    job: int | None = None,
+    hist_unit: str = "weeks",
+    **kwargs,
 ) -> "pandas.DataFrame":
 
     hist = hist or 4
 
-    if user:
+    if job:
+        command = f"sacct --job={job} --json"
+    elif user:
         command = f"sacct -u {user} --json -S now-{hist}{hist_unit}"
     else:
         command = f"sacct --json -S now-{hist}{hist_unit}"
@@ -197,8 +238,13 @@ def get_sacct(
         payload = json.loads(output[0])
     except json.JSONDecodeError:
         # console.print('[orange1 bold]Warning: using example data')
+        if job:
+            example = "sacct_job.json"
+        else:
+            example = "sacct.json"
+
         payload = json.load(
-            open(Path(__file__).parent.parent / "example_data" / "sacct.json", "rt")
+            open(Path(__file__).parent.parent / "example_data" / example, "rt")
         )
 
     global METADATA
@@ -348,12 +394,26 @@ def extract_inner(df, key, inner):
     df[key] = df.apply(_inner, axis=1)
 
 
+def extract_json(df, key):
+
+    def _inner(x):
+        return json.loads(x[key])
+
+    df[key] = df.apply(_inner, axis=1)
+
+
 def extract_time(df, key):
     df[key] = df.apply(
         lambda x: datetime.datetime.fromtimestamp(x[key]["number"]), axis=1
     )
 
 
+def extract_time_limit(df, key):
+    df[key] = df.apply(
+        lambda x: human_timedelta(datetime.timedelta(minutes=x[key]["number"])), axis=1
+    )
+
+
 def extract_sacct_times(df):
     df["start_time"] = df.apply(
         lambda x: datetime.datetime.fromtimestamp(x["time"]["start"]), axis=1

{richqueue-0.7 → richqueue-0.9}/richqueue/table.py

@@ -145,6 +145,61 @@ def node_table(df, long: bool = False, **kwargs):
     return table
 
 
+def job_table(row, job: int, long: bool = True, **kwargs):
+
+    from .slurm import METADATA
+
+    title = f"[bold]{job=} on {METADATA['cluster_name']}"
+
+    table = Table(title=title, box=None, header_style="")
+
+    columns = JOB_COLUMNS[long]
+
+    table.add_column("[bold underline]Key", style="bold")
+    table.add_column("[bold underline]Value")
+
+    for col in columns:
+        if not col:
+            table.add_row()
+            continue
+
+        col_data = COLUMNS[col]
+
+        key = col_data.get("header") or col
+
+        if (
+            "standard_output" in row
+            and col == "standard_error"
+            and row["standard_output"] == row[col]
+        ):
+            continue
+
+        if "]" in key:
+            key = key.split("]")[1]
+
+        try:
+            value = row[col]
+
+            if long or value:
+                formatter = FORMATTERS.get(col, str)
+                value = formatter(value)
+                style = col_data.get("style")
+
+                if style:
+                    value = f"[{style}]{value}"
+
+        except KeyError:
+            if long:
+                value = "[red]Unknown"
+            else:
+                value = None
+
+        if value or long:
+            table.add_row(key, value)
+
+    return table
+
+
 ### FORMATTERS
 
 
@@ -181,6 +236,78 @@ def mem_string(mb):
 
 ### SPECIFICATION
 
+JOB_COLUMNS = {
+    True: [
+        None,
+        "job_id",
+        "name",
+        "job_state",
+        "nodes",
+        None,
+        "user_name",
+        "group_name",
+        "partition",
+        "qos",
+        None,
+        "node_count",
+        "cpus",
+        "tasks",
+        "cpus_per_task",
+        "memory_per_cpu",
+        "memory_per_node",
+        "threads_per_core",
+        None,
+        "submit_time",
+        "start_time",
+        "run_time",
+        "time_limit",
+        None,
+        "command",
+        "current_working_directory",
+        "standard_output",
+        "standard_error",
+        None,
+        "exclusive",
+        "requeue",
+        "dependency",
+        "restart_cnt",
+        # "derived_exit_code",
+        # # "end_time",
+    ],
+    False: [
+        "job_id",
+        "name",
+        "job_state",
+        # "nodes",
+        "user_name",
+        # "group_name",
+        "partition",
+        # "qos",
+        "node_count",
+        "cpus",
+        # "tasks",
+        # "cpus_per_task",
+        # "memory_per_cpu",
+        # "memory_per_node",
+        # "threads_per_core",
+        # "submit_time",
+        "start_time",
+        "run_time",
+        "time_limit",
+        "command",
+        "current_working_directory",
+        "standard_output",
+        "standard_error",
+        "exclusive",
+        "requeue",
+        "dependency",
+        "restart_cnt",
+        # None,
+        # "derived_exit_code",
+        # # "end_time",
+    ],
+}
+
 RUNNING_JOB_COLUMNS = {
     True: [
         "job_id",
@@ -355,6 +482,29 @@ COLUMNS = {
         "style": "magenta",
         "no_wrap": True,
     },
+    "command": {"header": "Script", "style": "bright_yellow"},
+    "current_working_directory": {"header": "Directory", "style": "bright_yellow"},
+    "standard_output": {"header": "Log", "style": "bright_yellow"},
+    "standard_error": {"header": "Error Log", "style": "bright_yellow"},
+    "cpus_per_task": {"header": "#CPUs/Task", "style": "magenta"},
+    "dependency": {"header": "Dependencies", "style": "green_yellow"},
+    "derived_exit_code": {},
+    "group_name": {"header": "User Group", "style": "green_yellow"},
+    "tasks": {"header": "#Tasks", "style": "magenta"},
+    "memory_per_cpu": {"header": "RAM/CPU", "style": "magenta"},
+    "memory_per_node": {"header": "RAM/Node", "style": "magenta"},
+    "qos": {"header": "QoS", "style": "green_yellow"},
+    "restart_cnt": {"header": "#Restarts", "style": "green_yellow"},
+    "requeue": {"header": "Requeue?", "style": "green_yellow"},
+    "exclusive": {"header": "Exclusive?", "style": "green_yellow"},
+    "time_limit": {
+        "header": "[underline dodger_blue2]Time Limit",
+        "justify": "right",
+        "style": "dodger_blue2",
+        "no_wrap": True,
+    },
+    "threads_per_core": {"header": "#Threads/core", "style": "magenta"},
+    "user_name": {"header": "User Group", "style": "green_yellow"},
 }
 
 FORMATTERS = {