diff --git a/slurm-eff-tool.py b/slurm-eff-tool.py index ecb9154..87cfaf6 100755 --- a/slurm-eff-tool.py +++ b/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)