richqueue 0.1__tar.gz

richqueue-0.1/.gitignore

@@ -0,0 +1,162 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/

richqueue-0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Max Winokan
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

richqueue-0.1/PKG-INFO

@@ -0,0 +1,19 @@
+Metadata-Version: 2.3
+Name: richqueue
+Version: 0.1
+Summary: RichQueue
+Project-URL: Homepage, https://github.com/mwinokan/RichQueue/issues
+Project-URL: Bug Tracker, https://github.com/mwinokan/RichQueue/issues
+Author-email: Max Winokan <max@winokan.com>
+License-File: LICENSE
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Requires-Python: >=3.10
+Requires-Dist: pandas
+Requires-Dist: rich
+Requires-Dist: typer
+Description-Content-Type: text/markdown
+
+# RichQueue
+💰 RichQueue: A colourful and pythonic SLURM queue viewer

richqueue-0.1/README.md

@@ -0,0 +1,2 @@
+# RichQueue
+💰 RichQueue: A colourful and pythonic SLURM queue viewer

richqueue-0.1/pyproject.toml

@@ -0,0 +1,30 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+[project]
+name = "richqueue"
+version = "0.1"
+authors = [
+    { name = "Max Winokan", email = "max@winokan.com" },
+]
+description = "RichQueue"
+readme = "README.md"
+requires-python = ">=3.10"
+requires = []
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+]
+dependencies = [
+    "typer",
+    "rich",
+    "pandas",
+]
+[project.urls]
+"Homepage" = "https://github.com/mwinokan/RichQueue/issues"
+"Bug Tracker" = "https://github.com/mwinokan/RichQueue/issues"
+[tool.hatch.build]
+include = [
+    "richqueue/*.py",
+]

richqueue-0.1/richqueue/slurm.py

@@ -0,0 +1,407 @@
+from typer import Typer
+import json
+import datetime
+from rich.console import Console
+from rich.panel import Panel
+from rich.layout import Layout
+from pandas import DataFrame
+
+console = Console()
+
+METADATA = {}
+
+app = Typer()
+
+def color_by_state(state):
+
+    if state == "RUNNING":
+        return "[bold bright_green]Running"
+    elif state == "PENDING":
+        return "[bright_yellow]Pending"
+    elif state == "CANCELLED":
+        return "[orange3]Cancelled"
+    elif state == "FAILED":
+        return "[bold bright_red]Failed"
+    elif state == "COMPLETED":
+        return "[bold bright_green]Completed"
+    else:
+        return state
+
+def extract_inner(df, key, inner):
+
+    def _inner(x):
+        d = x[key] 
+        if "set" in d:
+            if d["set"]:
+                return d[inner]
+            else:
+                return None
+        else:
+            return d[inner]
+
+    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_list(df, key):
+    def inner(x):
+        if len(x[key]) == 1:
+            return x[key][0]
+        else:
+            return x[key]
+
+    df[key] = df.apply(inner, axis=1)
+
+
+def job_table(
+    df,
+    title: str = "jobs",
+    long: bool = False,
+    return_table: bool = False,
+    box: bool = False,
+    panel: bool = True,
+):
+
+    from rich.table import Table
+
+    if box:
+        from rich.box import SIMPLE
+
+        box = SIMPLE
+
+    from .tools import human_timedelta, human_datetime
+
+    num_nodes = "node_count" in df.columns
+    num_cores = "cpus" in df.columns
+    start_time = "start_time" in df.columns
+
+    table = Table(title=title, box=box, header_style="")
+
+    table.add_column(
+        "[bold underline]Job Id", justify="right", style="bold", no_wrap=True
+    )
+    table.add_column(
+        "[underline cyan]Job Name", justify="left", style="cyan", no_wrap=False
+    )
+    if num_nodes:
+        table.add_column(
+            "[underline magenta]#N", justify="right", style="magenta", no_wrap=True
+        )
+    if num_cores:
+        table.add_column(
+            "[underline magenta]#C", justify="right", style="magenta", no_wrap=True
+        )
+
+    if long and start_time:
+        table.add_column(
+            "[underline dodger_blue2]Start Time",
+            justify="right",
+            style="dodger_blue2",
+            no_wrap=True,
+        )
+
+    if start_time:
+        table.add_column(
+            "[underline dodger_blue2]Run Time",
+            justify="right",
+            style="dodger_blue2",
+            no_wrap=True,
+        )
+
+    if long:
+        table.add_column(
+            "[underline green_yellow]Partition",
+            justify="right",
+            style="green_yellow",
+            no_wrap=True,
+        )
+        table.add_column(
+            "[underline green_yellow]Nodes",
+            justify="left",
+            style="green_yellow",
+            no_wrap=False,
+        )
+
+    table.add_column(
+        "[underline bold]State", justify="left", style=None, no_wrap=True
+    )
+
+    for i, row in df.iterrows():
+
+        if start_time:
+            start_time = human_datetime(row.start_time)
+            run_time = human_timedelta(datetime.datetime.now() - row.start_time)
+
+        values = [row.job_id, row["name"]]
+
+        if num_nodes:
+            values.append(row.node_count)
+        if num_cores:
+            values.append(row.cpus)
+
+        if long and start_time:
+            values.append(start_time)
+
+        if start_time:
+            values.append(run_time)
+
+        if long:
+            values += [row.partition, row.nodes]
+
+        values.append(color_by_state(row.job_state))
+
+        table.add_row(*[str(v) for v in values])
+
+    if panel:
+        table = Panel(table, expand=False)
+
+    if return_table:
+        return table
+
+    console.print(table)
+
+
+def parse_squeue_json(payload: dict) -> "DataFrame":
+
+    command = "squeue"
+
+    global METADATA
+
+    METADATA = {
+        "cluster_name": payload["meta"]["slurm"]["cluster"],
+        "user": payload["meta"]["client"]["user"],
+        "group": payload["meta"]["client"]["group"],
+    }
+
+    # parse payload
+    df = DataFrame(payload["jobs"])
+
+    # filter columns
+
+    columns = COLUMNS[command]
+
+    try:
+        df = df[columns]
+
+    except KeyError:
+
+        for key in columns:
+            if key not in df.columns:
+                raise KeyError(key)
+
+    extract_inner(df, "cpus", "number")
+    extract_inner(df, "node_count", "number")
+    extract_inner(df, "cpus_per_task", "number")
+    extract_inner(df, "threads_per_core", "number")
+
+    extract_time(df, "end_time")
+    extract_time(df, "start_time")
+    extract_time(df, "submit_time")
+    extract_time(df, "time_limit")
+
+    extract_list(df, "job_state")
+
+    return df
+
+def parse_sacct_json(payload: dict) -> "DataFrame":
+
+    command = "sacct"
+
+    global METADATA
+
+    METADATA.update({
+        "cluster_name": payload["meta"]["slurm"]["cluster"],
+        "user": payload["meta"]["client"]["user"],
+        "group": payload["meta"]["client"]["group"],
+    })
+
+    # parse payload
+    df = DataFrame(payload["jobs"])
+
+    # filter columns
+
+    columns = COLUMNS[command]
+
+    try:
+        df = df[columns]
+
+    except KeyError:
+
+        for key in columns:
+            if key not in df.columns:
+                raise KeyError(key)
+
+    df = df.rename(columns={"user":"user_name", "state":"job_state"})
+
+    extract_inner(df, "job_state", "current")
+
+    # extract_number(df, "cpus")
+    # extract_number(df, "node_count")
+    # extract_number(df, "cpus_per_task")
+    # extract_number(df, "threads_per_core")
+
+    # extract_time(df, "end_time")
+    # extract_time(df, "start_time")
+    # extract_time(df, "submit_time")
+    # extract_time(df, "time_limit")
+
+    extract_list(df, "job_state")
+
+    df = df[df["job_state"] != "RUNNING"]
+
+    return df
+
+def show_queue(
+    command: str,
+    user: None | str = None,
+    long: bool = False,
+    return_table: bool = False,
+    box: bool = False,
+):
+
+    import subprocess
+
+    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")
+
+    if user:
+        s_command = f"{command} -u {user} --json"
+    else:
+        s_command = f"{command} --json"
+
+    x = subprocess.Popen([s_command], shell=True, stdout=subprocess.PIPE)
+    output = x.communicate()
+
+    payload = json.loads(output[0])
+
+    parser = PARSERS[command]
+
+    df = parser(payload)
+
+    cluster = METADATA["cluster_name"]
+
+    if command == "sacct":
+        title = f"[bold]{user}'s history on {cluster}"
+
+    elif user:
+        df = df[df["user_name"] == user]
+        title = f"[bold]{user}'s jobs on {cluster}"
+    else:
+        title = f"[bold]jobs on {cluster}"
+
+    return job_table(df, title=title, long=long, return_table=return_table, box=box)
+
+def dual_layout(
+    user: None | str = None,
+    long: bool = False,
+    return_table: bool = True,
+    box: bool = False,
+):
+
+    layout = Layout()
+
+    panel1 = show_queue(user=user, command="squeue", long=long, return_table=True)
+    panel2 = show_queue(user=user, command="sacct", long=long, return_table=True)
+
+    upper = Layout(renderable=panel1, name="upper")
+    lower = Layout(renderable=panel2, name="lower")
+
+    upper.size = panel1.renderable.row_count+4
+    lower.size = panel2.renderable.row_count+4
+
+    layout.split_column(
+        upper,
+        lower,
+    )
+
+    if not return_table:
+        console.print(layout)
+    else:
+        return layout
+
+@app.command()
+def show(user: None | str = None, long: bool = False):
+
+    loop = True
+
+    if loop:
+
+        from rich.live import Live
+        import time
+
+        layout = dual_layout(user=user, long=long)
+
+        with Live(layout, refresh_per_second=4, screen=True) as live:
+            while True:
+
+                layout = dual_layout(user=user, long=long)
+                live.update(layout)
+
+                time.sleep(1)
+
+    else:
+        layout = dual_layout(user=user, long=long)
+        console.print(layout)
+        # show_queue(command="squeue", user=user, long=long)
+        # show_queue(command="sacct", user=user, long=long)
+
+
+def main():
+
+    squeue = json.load(open("../example_data/squeue.json"))
+
+    parse_squeue_json(squeue)
+
+PARSERS = {"squeue":parse_squeue_json, "sacct":parse_sacct_json}
+
+COLUMNS = {
+"sacct":[
+    "job_id",
+    "state",
+    "name",
+    "nodes",
+    "partition",
+    "user",
+],
+
+"squeue":[
+    "command",
+    "cpus_per_task",
+    "dependency",
+    "derived_exit_code",
+    "end_time",
+    "group_name",
+    "job_id",
+    "job_state",
+    "name",
+    "nodes",
+    "node_count",
+    "cpus",
+    "tasks",
+    "partition",
+    "memory_per_cpu",
+    "memory_per_node",
+    "qos",
+    "restart_cnt",
+    "requeue",
+    "exclusive",
+    "start_time",
+    "standard_error",
+    "standard_output",
+    "submit_time",
+    "time_limit",
+    "threads_per_core",
+    "user_name",
+    "current_working_directory",
+]}
+
+
+if __name__ == "__main__":
+    app()

richqueue-0.1/richqueue/tools.py

@@ -0,0 +1,58 @@
+import datetime
+
+
+def human_timedelta(delta):
+
+    bits = []
+
+    days = delta.days
+
+    s = delta.seconds
+
+    m, s = divmod(s, 60)
+    h, m = divmod(m, 60)
+
+    if days:
+        bits.append(f"{days}d")
+
+    if h:
+        bits.append(f"{h}h")
+
+    if m:
+        bits.append(f"{m}m")
+
+    if s:
+        bits.append(f"{s}s")
+
+    bits = (str(b) for b in bits)
+
+    return " ".join(bits)
+
+
+def human_datetime(dt):
+
+    now = datetime.datetime.now()
+
+    bits = []
+
+    bits.append(dt.strftime("%b"))
+
+    day = dt.day
+
+    match day % 10:
+        case 1:
+            suffix = "st"
+        case 2:
+            suffix = "nd"
+        case 3:
+            suffix = "rd"
+        case _:
+            suffix = "th"
+
+    bits.append(f"{day}{suffix}")
+
+    bits.append(dt.strftime("%H:%M"))
+
+    bits = (str(b) for b in bits)
+
+    return " ".join(bits)