diff --git a/slurm-eff-tool.py b/slurm-eff-tool.py index eea1bdd..2ebc38d 100755 --- a/slurm-eff-tool.py +++ b/slurm-eff-tool.py @@ -38,7 +38,7 @@ from dataclasses import dataclass, field, fields, asdict from pathlib import Path import ast import operator as op -from typing import Any, Iterable, Mapping +from typing import Any, Iterable, Mapping, Union, List, Tuple, cast import msgpack import gzip @@ -275,6 +275,8 @@ _ALLOWED_CMPOPS = { ast.LtE: op.le, ast.Gt: op.gt, ast.GtE: op.ge, + ast.In: lambda a, b: a in b, + ast.NotIn: lambda a, b: a not in b, } ########################################## @@ -282,30 +284,39 @@ _ALLOWED_CMPOPS = { #################################### # Class for expression evaluations # #################################### +# TODO: Note that we currently cannot filter on state strings, because +# we map state to numeric ID values +SEValue = Union[int, float, str, bool, None] +SEEvalValue = Union[SEValue, List[SEValue], Tuple[SEValue, ...]] class SafeExpression: def __init__(self, expression: str): self.expression = expression.lower() self.tree = ast.parse(self.expression, mode="eval") - def evaluate(self, variables: Mapping[str, int | float]) -> int | float | bool: + def evaluate(self, variables: Mapping[str, SEValue]) -> bool: # normalize to downcased variable names variables = {k.lower(): v for k,v in variables.items()} - try: - return self._eval(self.tree.body, variables) - except TypeError: - return False - except: - raise + return bool(self._eval(self.tree.body, variables)) - def _eval(self, node: ast.AST, variables: Mapping[str, int | float]): + def _eval(self, node: ast.AST, variables: Mapping[str, SEValue]) -> SEEvalValue: if isinstance(node, ast.Constant): - if isinstance(node.value, (int, float, bool)): + if isinstance(node.value, (int, float, str, bool)): return node.value raise ValueError(f"Unsupported constant: {node.value!r}") + if isinstance(node, ast.List): + return [cast(SEValue, self._eval(element, variables)) + for element in node.elts] + + if isinstance(node, ast.Tuple): + return tuple(cast(SEValue, self._eval(element, variables)) + for element in node.elts) + + # Variable names if isinstance(node, ast.Name): + name = node.id.lower() try: - return variables[node.id] + return variables[name] except KeyError: raise ValueError(f"Unknown variable: {node.id}") from None @@ -315,6 +326,9 @@ class SafeExpression: raise ValueError(f"Unsupported operator: {op_type.__name__}") left = self._eval(node.left, variables) right = self._eval(node.right, variables) + # deal with columns containing None value + if left is None or right is None: + return False return _ALLOWED_BINOPS[op_type](left, right) if isinstance(node, ast.UnaryOp): @@ -347,6 +361,9 @@ class SafeExpression: raise ValueError(f"Unsupported comparison: {op_type.__name__}") right = self._eval(comparator, variables) + # deal with columns containing None value + if left is None or right is None: + return False if not _ALLOWED_CMPOPS[op_type](left, right): return False @@ -1199,17 +1216,21 @@ def parse_args(argv: list[str]) -> argparse.Namespace: # list that user's single jobs slurm-eff-tool.py -L sacct.cache -u dfeich - # supports multiple sort keys - slurm-eff-tool.py -L sacct.cache --aggr-user --sdev -s cpu,-mem,time + # supports sorting, also by multiple keys + slurm-eff-tool.py -L sacct.cache --aggr-user -s-waste_total + slurm-eff-tool.py -L sacct.cache --aggr-user -s cpu,-mem,time - # you can cluster jobs by Regexps applying to the job names + # you can cluster jobs by multiple regexps applying to the job names slurm-eff-tool.py -L sacct.cache -u dfeich -R '^vasp','^gromacs' - # supports flexibel output formatting + # supports flexibel output formatting by defining columns to print slurm-eff-tool.py -L sacct.cache -o username,Y # only print rows that evaluate to true based on arithmetic expressions slurm-eff-tool.py -L sacct.cache -U --expr "(waste_Mem > 2000 and Mem_Eff < 20) and MaxRSS_max/AllocMem < 0.5" + slurm-eff-tool.py -L sacct.cache -U --expr 'username=="feichtinger"' + slurm-eff-tool.py -L sacct.cache -U --expr 'partition in ["standard","short"]' + """ ) @@ -1217,16 +1238,7 @@ def parse_args(argv: list[str]) -> argparse.Namespace: default="now - 24 hours") p.add_argument("-E", "--end", help="sacct end time, passed to sacct -E", default="now") - p.add_argument("-u", "--user", help="restrict to one user; passed as sacct -u unless reading from cache") - p.add_argument("--state", "--job-state", dest="state", - default=None, - help="sacct state filter, e.g. COMPLETED,FAILED,TIMEOUT") - - p.add_argument("-O", "--output-raw", help="write raw sacct output cache to this file (text format, large).") - p.add_argument("-F", "--from-raw", help="read raw sacct output cache from this file.") - p.add_argument("-B", "--write-binary-cache", help="write a binary cache file in gzipped msgpack format") - p.add_argument("-L", "--load-binary-cache", help="load a binary cache file in gzipped msgpack format") - p.add_argument("-i", "--info", help="show metadata information for the given binary cache file") + p.add_argument("-u", "--user", help="restrict to a single user at sacct DB search") p.add_argument("-U", "--aggr-user", action="store_true", help="aggregate jobs by user, CPUs, nodes, ReqMem, and timelimit") p.add_argument( "-R", @@ -1234,18 +1246,27 @@ def parse_args(argv: list[str]) -> argparse.Namespace: action="append", help="aggregate jobs based on a regexp applied to the jobnames. Option may be repeated multiple times.", ) + p.add_argument("--state", "--job-state", dest="state", + default=None, + help="state filter passed to sacct, e.g. COMPLETED,FAILED,TIMEOUT") + + p.add_argument("-B", "--write-binary-cache", help="write a binary cache file in gzipped msgpack format") + p.add_argument("-L", "--load-binary-cache", help="load a binary cache file in gzipped msgpack format") + p.add_argument("-O", "--output-raw", help="write raw sacct output cache to this file (text format, large).") + p.add_argument("-F", "--from-raw", help="read raw sacct output cache from this file.") + p.add_argument("-i", "--info", help="show metadata information for the given binary cache file") p.add_argument("--dflt-mpcpu", help=f"Default memory/CPU ratio of cluster [{default_mempercpu_gb} GB/cpu]. Used in memory waste calculation.") p.add_argument("--sdev", action="store_true", help="after each efficiency average, add sdev, max, and min columns") p.add_argument("--json", action="store_true", help="emit JSON instead of an ASCII table") - p.add_argument("-s", "--sort", help="comma-separated numeric sort columns or short aliases; prefix with - for descending") + p.add_argument("-s", "--sort", help="comma-separated numeric sort columns or short aliases; prefix with - for descending.") fmthelp = ", ".join([f"{k}:{ALIASES[k]}" for k in ALIASES]) p.add_argument( "-o", "--format", help= f"String of comma separated column names or short aliases defining the output format [{fmthelp}]", ) - p.add_argument("--expr", help="filter using an arithmetic expression using column names as variables (case insensitive)", + p.add_argument("--expr", help="filter using an arithmetic expression using column names as variables (case insensitive). Don't use for state (use --state flag instead).", default=None) p.add_argument("-p", "--preset", help=f"use one of several preset output column formats [{','.join(PRESET_COLUMNS.keys())}]", default=None)