implement expr comparisons for string variables

This commit is contained in:
2026-07-03 21:28:32 +02:00
parent 5b952870e5
commit 8f32991734
+48 -27
View File
@@ -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)