implement expr comparisons for string variables
This commit is contained in:
+48
-27
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user