diff --git a/slurm-eff-tool.py b/slurm-eff-tool.py index 8b001cc..d62b88c 100755 --- a/slurm-eff-tool.py +++ b/slurm-eff-tool.py @@ -30,11 +30,15 @@ import subprocess import sys import signal from collections import defaultdict -from dataclasses import dataclass, field, fields +from dataclasses import dataclass, field, fields, asdict from pathlib import Path import ast import operator as op from typing import Any, Iterable, Mapping +import msgpack +import gzip + +VERSION=0.2 # Default GB per CPUs for the cluster default_mempercpu_gb = 2.0 @@ -275,7 +279,8 @@ class SafeExpression: @dataclass class JobRecord: - raw: dict[str, str] + jobid: str + jobidraw: str username: str jobname: str reqtasks: int @@ -571,14 +576,20 @@ def parse_pipe_rows(lines: Iterable[str]) -> list[dict[str, str]]: return rows -def write_cache(path: str, rows: list[dict[str, str]]) -> None: +def write_cache_raw(path: str, rows: list[dict[str, str]]) -> None: + """Write a raw sacct output file""" + if os.path.exists(path): + sys.stderr.write(f"WARNING, file exists: {path}. Balking out...\n") + sys.exit(1) + with open(path, "w", newline="", encoding="utf-8") as f: writer = csv.DictWriter(f, fieldnames=SACCT_FIELDS, delimiter="|", lineterminator="\n") writer.writeheader() writer.writerows(rows) -def read_cache(path: str) -> list[dict[str, str]]: +def read_cache_raw(path: str) -> list[dict[str, str]]: + """Read a raw sacct output file.""" with open(path, newline="", encoding="utf-8") as f: reader = csv.DictReader(f, delimiter="|") missing = [x for x in SACCT_FIELDS if x not in (reader.fieldnames or [])] @@ -590,6 +601,18 @@ def read_cache(path: str) -> list[dict[str, str]]: return [{k: (row.get(k) or "").strip() for k in SACCT_FIELDS} for row in reader] +def write_binary_cache(records: list[JobRecord], filename): + if os.path.exists(filename): + sys.stderr.write(f"WARNING, file exists: {filename}. Balking out...\n") + sys.exit(1) + cache = [asdict(r) for r in records] + cache_blob = { + "slurm_eff_version": VERSION, + "records": cache +} + with gzip.open(filename, "wb") as f: + msgpack.pack(cache_blob, f) + def base_job_id(jobid: str) -> str: return re.split(r"[._]", jobid, maxsplit=1)[0] @@ -717,7 +740,8 @@ def build_job_records(rows: list[dict[str, str]], records.append( JobRecord( - raw=top, + jobid=top.get("JobID", ""), + jobidraw=top.get("JobIDRaw", ""), username=top.get("User", ""), jobname=top.get("JobName", ""), reqtasks=req_ntasks, @@ -809,9 +833,13 @@ def make_single_row(rec: JobRecord, dflt_mpcpu: float) -> OutputRow: if rec.cpu_eff is not None: waste_cpu = walltime * (100-rec.cpu_eff) * rec.cpus + jobid = rec.jobidraw + if jobid == "": + jobid = rec.jobid + return OutputRow( username=rec.username, - JobID=rec.raw.get("JobIDRaw", rec.raw.get("JobID", "")), + JobID=jobid, state=state_mappings_inv[rec.state], NTasks=rec.reqtasks, CPUs=rec.cpus, @@ -1037,9 +1065,10 @@ def parse_args(argv: list[str]) -> argparse.Namespace: default="COMPLETED", help="sacct state filter, e.g. COMPLETED,FAILED,TIMEOUT") - p.add_argument("-O", "--output-cache", help="write raw sacct output cache to this file") - p.add_argument("-F", "--from-cache", help="read raw sacct output cache from this file instead of running sacct") - + p.add_argument("-O", "--output-raw", help="write raw sacct output cache to this file") + p.add_argument("-F", "--from-raw", help="read raw sacct output cache from this file instead of running sacct") + p.add_argument("-B", "--write-binary-cache", help="write a binary cache file in msgpack format") + p.add_argument("-L", "--load-binary-cache", help="load a binary cache file in msgpack format") p.add_argument("-U", "--aggr-user", action="store_true", help="aggregate jobs by user, CPUs, nodes, ReqMem, and timelimit") p.add_argument( "-R", @@ -1065,8 +1094,6 @@ def parse_args(argv: list[str]) -> argparse.Namespace: args = p.parse_args(argv) - # if args.from_cache and any([args.start, args.end, args.state]): - # print("warning: -S/-E/--state are ignored when using -F/--from-cache", file=sys.stderr) if args.aggr_user and args.aggr_regexp: die("choose only one aggregation mode: --aggr-user or --aggr-regexp") return args @@ -1090,15 +1117,29 @@ def main(argv: list[str] | None = None) -> int: else: output_columns = [c for c in output_columns if c != "Count"] - if args.from_cache: - rows = read_cache(args.from_cache) + # CREATE RECORDS FROM SACCT QUERY OR FILE, OR LOAD PROCESSED RECORDS FROM BINARY CACHE + if args.load_binary_cache: + with gzip.open(args.load_binary_cache, "rb") as f: + blob = msgpack.unpack(f, raw=False) + version_str = blob["slurm_eff_version"] + if version_str != VERSION: + sys.stderr.write(f"WARNING: cache was written by version f{version_str}, but this is version {VERSION}\n") + # print(f"DEBUG: Loading from a binary cache created by version {version_str}") + records = [JobRecord(**d) for d in blob["records"]] else: - rows = run_sacct(args) - if args.output_cache: - write_cache(args.output_cache, rows) + if args.from_raw: + rows = read_cache_raw(args.from_raw) + else: + rows = run_sacct(args) + if args.output_raw: + write_cache_raw(args.output_raw, rows) + + records = build_job_records(rows, filter_user=args.user, + filter_state=args.state) + + if args.write_binary_cache: + write_binary_cache(records, args.write_binary_cache) - records = build_job_records(rows, filter_user=args.user, - filter_state=args.state) out_rows = aggregate_records(records, args) # , dflt_mpcpu=default_mempercpu_gb if args.expr: expr = SafeExpression(args.expr)