HPCE/slurm-efficiency-tools
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add binary cache writing by gzipped msgpack format

Browse Source 
adds the dependency on msgpack-python, but it's
available as a standard packages on RedHat and SUSE
This commit is contained in:
feichtinger
2026-06-26 17:36:03 +02:00
parent 652acc0b9a
commit fd66073623
1 changed files with 59 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
slurm-eff-tool.py
+59 -18
View File
@@ -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)