richqueue 1.4__tar.gz → 1.6__tar.gz

{richqueue-1.4 → richqueue-1.6}/PKG-INFO

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

{richqueue-1.4 → richqueue-1.6}/pyproject.toml

@@ -3,7 +3,7 @@ requires = ["hatchling"]
 build-backend = "hatchling.build"
 [project]
 name = "richqueue"
-version = "1.4"
+version = "1.6"
 authors = [
     { name = "Max Winokan", email = "max@winokan.com" },
 ]
@@ -31,4 +31,5 @@ include = [
 
 [project.scripts]
 rq = "richqueue.main:main"
+res = "richqueue.log:main"
 confucius = "richqueue.wisdom:main"

richqueue-1.6/richqueue/log.py

@@ -0,0 +1,157 @@
+import time
+
+from rich.live import Live
+from rich.panel import Panel
+from rich.layout import Layout
+
+from typing_extensions import Annotated
+from typer import Typer, Argument, Option
+
+from .console import console
+from .layout import dual_layout
+from .slurm import combined_df, get_user, PANEL_PADDING
+from .table import job_table, log_table
+from .tools import curry
+
+# access logs
+
+JOB_DF = None
+
+app = Typer()
+
+
+@app.callback(invoke_without_command=True)
+def show_log(
+    job: Annotated[int, Argument(help="Show logs for this job")] = None,
+    user: Annotated[
+        str,
+        Option(
+            "-u",
+            "--user",
+            help="Query jobs for another user",
+        ),
+    ] = None,
+    long: Annotated[bool, Option("-v", "--long", help="More detailed output")] = False,
+    wisdom: Annotated[
+        bool, Option("-w", "--wisdom", help="Show a confucius quote on exit")
+    ] = False,
+):
+
+    if user is None:
+        user = get_user()
+
+    if job is None:
+
+        def job_getter():
+            return guess_current_job(user=user)
+
+    else:
+
+        def job_getter():
+            return get_specific_job(job=job, user=None)
+
+    job = job_getter()
+
+    if job is None:
+        return None
+
+    layout_func = curry(dual_layout, log_layout_pair)
+
+    layout = layout_func(job_getter=job_getter, long=long)
+
+    with Live(
+        layout,
+        refresh_per_second=1,
+        screen=True,
+        transient=True,
+        vertical_overflow="visible",
+    ) as live:
+
+        try:
+            while True:
+                layout = layout_func(job_getter=job_getter, long=long)
+                live.update(layout)
+                time.sleep(1)
+        except KeyboardInterrupt:
+            live.stop()
+            if wisdom:
+                from .wisdom import print_random_quote
+
+                print_random_quote()
+
+
+def log_layout_pair(job_getter, **kwargs):
+    job = job_getter()
+    # console.print(job)
+
+    upper = job_table(row=job, job=job.job_id, **kwargs)
+    upper = Panel(upper, expand=False)
+
+    limit = console.size.height - 2 * PANEL_PADDING - upper.renderable.row_count
+
+    lower = log_table(
+        stdout=job.standard_output, stderr=job.standard_error, limit=limit
+    )
+    lower = Panel(lower, expand=False)
+
+    return upper, lower
+
+
+def guess_current_job(user: str, **kwargs):
+
+    global JOB_DF
+    JOB_DF = combined_df(user=user, **kwargs)
+
+    assert len(JOB_DF), "No jobs found"
+
+    pending = JOB_DF[JOB_DF["job_state"] == "PENDING"]
+    running = JOB_DF[JOB_DF["job_state"] == "RUNNING"]
+    history = JOB_DF[~JOB_DF["job_state"].isin(["PENDING", "RUNNING"])]
+
+    # pick the single active job
+    if len(running) == 1:
+        return running.iloc[0]
+
+    # pick the most recently submitted active job
+    elif len(running) > 1:
+        return running.sort_values(by="submit_time", ascending=False).iloc[0]
+
+    # pick the single pending job
+    elif len(pending) == 1:
+        return pending.iloc[0]
+
+    # pick the most recently submitted pending job
+    elif len(pending) > 1:
+        return pending.sort_values(by="submit_time", ascending=False).iloc[0]
+
+    # pick the most recently submitted pending job
+    else:
+        return pending.sort_values(by="submit_time", ascending=False).iloc[0]
+
+    return job
+
+
+def get_specific_job(job: int, user: str, **kwargs):
+
+    global JOB_DF
+    JOB_DF = combined_df(user=user, **kwargs)
+
+    assert len(JOB_DF), "No jobs found"
+
+    matches = JOB_DF[JOB_DF["job_id"] == job]
+
+    if not len(matches):
+        console.print("[red bold]Could not find job with ID {job_id=}")
+        return None
+
+    assert len(matches) == 1
+
+    return matches.iloc[0]
+
+
+def main():
+    app()
+
+
+if __name__ == "__main__":
+    main()

{richqueue-1.4 → richqueue-1.6}/richqueue/slurm.py

@@ -18,14 +18,18 @@ PANEL_PADDING = 4
 ### CONSTRUCT LAYOUT
 
 
+def get_user() -> str:
+    x = subprocess.Popen(["whoami"], shell=True, stdout=subprocess.PIPE)
+    output = x.communicate()
+    return output[0].strip().decode("utf-8")
+
+
 def get_layout_pair(user: str | None, **kwargs):
 
     if user == "all":
         user = None
     elif user is None:
-        x = subprocess.Popen(["whoami"], shell=True, stdout=subprocess.PIPE)
-        output = x.communicate()
-        user = output[0].strip().decode("utf-8")
+        user = get_user()
 
     df = combined_df(user=user, **kwargs)
 
@@ -186,10 +190,12 @@ def get_squeue(
 
     global METADATA
 
+    meta = payload["meta"]
+
     METADATA = {
-        "cluster_name": payload["meta"]["slurm"]["cluster"],
-        "user": payload["meta"]["client"]["user"],
-        "group": payload["meta"]["client"]["group"],
+        "cluster_name": meta.get("slurm", {}).get("cluster", "unknown HPC"),
+        "user": meta.get("client", {}).get("user", "unknown user"),
+        "group": meta.get("client", {}).get("group", "unknown group"),
     }
 
     # parse payload
@@ -223,7 +229,7 @@ def get_squeue(
 
     extract_list(df, "job_state")
     extract_list(df, "exclusive")
-    extract_json(df, "exclusive")
+    # extract_json(df, "exclusive")
 
     return df
 
@@ -266,10 +272,12 @@ def get_sacct(
 
     global METADATA
 
+    meta = payload["meta"]
+
     METADATA = {
-        "cluster_name": payload["meta"]["slurm"]["cluster"],
-        "user": payload["meta"]["client"]["user"],
-        "group": payload["meta"]["client"]["group"],
+        "cluster_name": meta.get("slurm", {}).get("cluster", "unknown HPC"),
+        "user": meta.get("client", {}).get("user", "unknown user"),
+        "group": meta.get("client", {}).get("group", "unknown group"),
         "hist": hist,
         "hist_unit": hist_unit,
     }
@@ -316,10 +324,12 @@ def get_sinfo(**kwargs) -> "pandas.DataFrame":
 
     global METADATA
 
+    meta = payload["meta"]
+
     METADATA = {
-        "cluster_name": payload["meta"]["slurm"]["cluster"],
-        "user": payload["meta"]["client"]["user"],
-        "group": payload["meta"]["client"]["group"],
+        "cluster_name": meta.get("slurm", {}).get("cluster", "unknown HPC"),
+        "user": meta.get("client", {}).get("user", "unknown user"),
+        "group": meta.get("client", {}).get("group", "unknown group"),
     }
 
     df = DataFrame(payload["sinfo"])

{richqueue-1.4 → richqueue-1.6}/richqueue/table.py

@@ -2,6 +2,7 @@ from rich.table import Table
 from .tools import human_datetime
 from pandas import isnull
 from .console import console
+import subprocess
 
 ### TABLES
 
@@ -200,6 +201,40 @@ def job_table(row, job: int, long: bool = True, **kwargs):
     return table
 
 
+def log_table(stdout, stderr, limit):
+
+    from rich.text import Text
+
+    dual = stdout != stderr
+
+    if dual:
+        raise NotImplementedError("separate log files not yet supported")
+
+    else:
+
+        table = Table(title="stdout", box=None, header_style="")
+
+        table.add_column("[bold underline]#", style="bold")
+        table.add_column("[bold underline]Log", no_wrap=True)
+
+        lines = tail(stdout, limit)
+
+        for i, line in enumerate(lines):
+            line = line.decode("utf-8")
+
+            line = line.split("\r")[-1]
+
+            table.add_row(str(i), Text.from_ansi(line))
+
+    return table
+
+
+def tail(path, n):
+    proc = subprocess.Popen(["tail", "-n", str(n), path], stdout=subprocess.PIPE)
+    lines = proc.stdout.readlines()
+    return lines
+
+
 ### FORMATTERS