add filtering by arithmetic expression

slurm-eff-tool.py

@@ -32,7 +32,9 @@ import signal
 from collections import defaultdict
 from dataclasses import dataclass, field, fields
 from pathlib import Path
-from typing import Any, Iterable
+import ast
+import operator as op
+from typing import Any, Iterable, Mapping
 
 # Default GB per CPUs for the cluster
 default_mempercpu_gb = 2.0
@@ -117,6 +119,108 @@ NUMERIC_COLUMNS = {
     "waste_Mem",
 }
 
+##########################################
+# Definitions for expression evaluations #
+##########################################
+
+_ALLOWED_BINOPS = {
+    ast.Add: op.add,
+    ast.Sub: op.sub,
+    ast.Mult: op.mul,
+    ast.Div: op.truediv,
+    ast.FloorDiv: op.floordiv,
+    ast.Mod: op.mod,
+    ast.Pow: op.pow,
+}
+
+_ALLOWED_UNARYOPS = {
+    ast.UAdd: op.pos,
+    ast.USub: op.neg,
+    ast.Not: op.not_,
+}
+
+_ALLOWED_CMPOPS = {
+    ast.Eq: op.eq,
+    ast.NotEq: op.ne,
+    ast.Lt: op.lt,
+    ast.LtE: op.le,
+    ast.Gt: op.gt,
+    ast.GtE: op.ge,
+}
+
+##########################################
+
+####################################
+# Class for expression evaluations #
+####################################
+class SafeExpression:
+    def __init__(self, expression: str):
+        self.expression = expression
+        self.tree = ast.parse(expression, mode="eval")
+
+    def evaluate(self, variables: Mapping[str, int | float]) -> int | float | bool:
+        return self._eval(self.tree.body, variables)
+
+    def _eval(self, node: ast.AST, variables: Mapping[str, int | float]):
+        if isinstance(node, ast.Constant):
+            if isinstance(node.value, (int, float, bool)):
+                return node.value
+            raise ValueError(f"Unsupported constant: {node.value!r}")
+
+        if isinstance(node, ast.Name):
+            try:
+                return variables[node.id]
+            except KeyError:
+                raise ValueError(f"Unknown variable: {node.id}") from None
+
+        if isinstance(node, ast.BinOp):
+            op_type = type(node.op)
+            if op_type not in _ALLOWED_BINOPS:
+                raise ValueError(f"Unsupported operator: {op_type.__name__}")
+            left = self._eval(node.left, variables)
+            right = self._eval(node.right, variables)
+            return _ALLOWED_BINOPS[op_type](left, right)
+
+        if isinstance(node, ast.UnaryOp):
+            op_type = type(node.op)
+            if op_type not in _ALLOWED_UNARYOPS:
+                raise ValueError(f"Unsupported unary operator: {op_type.__name__}")
+            return _ALLOWED_UNARYOPS[op_type](self._eval(node.operand, variables))
+
+        if isinstance(node, ast.BoolOp):
+            if isinstance(node.op, ast.And):
+                for value in node.values:
+                    if not self._eval(value, variables):
+                        return False
+                return True
+
+            if isinstance(node.op, ast.Or):
+                for value in node.values:
+                    if self._eval(value, variables):
+                        return True
+                return False
+
+            raise ValueError(f"Unsupported boolean operator: {type(node.op).__name__}")
+
+        if isinstance(node, ast.Compare):
+            left = self._eval(node.left, variables)
+
+            for operator_node, comparator in zip(node.ops, node.comparators):
+                op_type = type(operator_node)
+                if op_type not in _ALLOWED_CMPOPS:
+                    raise ValueError(f"Unsupported comparison: {op_type.__name__}")
+
+                right = self._eval(comparator, variables)
+
+                if not _ALLOWED_CMPOPS[op_type](left, right):
+                    return False
+
+                left = right
+
+            return True
+
+        raise ValueError(f"Unsupported expression element: {type(node).__name__}")
+
 
 @dataclass
 class JobRecord:
@@ -588,8 +692,11 @@ def aggregate_records(records: list[JobRecord], args: argparse.Namespace) -> lis
             if not matched:
                 unmatched.append(rec)
 
-        out = [make_aggregate_row(v, username=k[1], jobname=k[0], ) for k, v in buckets.items()]
-        out.extend(make_single_row(r, dflt_mpcpu=default_mempercpu_gb) for r in unmatched)
+        out = [make_aggregate_row(v, username=k[1], jobname=k[0], \
+                                  dflt_mpcpu=default_mempercpu_gb) \
+                                  for k, v in buckets.items()]
+        out.extend(make_single_row(r, dflt_mpcpu=default_mempercpu_gb) \
+                   for r in unmatched)
         return out
 
     if args.aggr_user:
@@ -828,6 +935,8 @@ def parse_args(argv: list[str]) -> argparse.Namespace:
   slurm-eff-tool.py -F sacct.cache -u dfeich -R '^vasp','^gromacs' --json
   # supports flexibel output formatting
   slurm-eff-tool.py -F sacct.cache -o username,Y
+  # only print rows that evaluate to true based on expression
+  slurm-eff-tool.py -F sacct.cache -U --expr "(waste_Mem > 2000 and Mem_Eff < 20) and MaxRSS_max/AllocMem < 0.5"
         """
     )
 
@@ -861,6 +970,8 @@ def parse_args(argv: list[str]) -> argparse.Namespace:
         "--format",
         help= f"String of comma separated column names or short names defining the output format:\n{fmthelp}",
     )
+    p.add_argument("--expr", help="arithmetic expression using column names",
+                   default=None)
 
     args = p.parse_args(argv)
 
@@ -896,6 +1007,10 @@ def main(argv: list[str] | None = None) -> int:
 
     records = build_job_records(rows, only_user=args.user)
     out_rows = aggregate_records(records, args) # , dflt_mpcpu=default_mempercpu_gb
+    if args.expr:
+        expr = SafeExpression(args.expr)
+        out_rows = [row for row in out_rows \
+                    if expr.evaluate(row.as_dict())]
     out_rows = sort_rows(out_rows, args.sort)