diff --git a/slurm-eff-tool.py b/slurm-eff-tool.py index e48a16a..cf44f53 100755 --- a/slurm-eff-tool.py +++ b/slurm-eff-tool.py @@ -645,6 +645,8 @@ def stdev_or_none(values: list[float]) -> float | None: def mean_or_none(values: list[float]) -> float | None: return statistics.mean(values) if values else None +def median_or_none(values: list[float]) -> float | None: + return statistics.median(values) if values else None def pct(numerator: float | None, denominator: float | None) -> float | None: if numerator is None or denominator is None or denominator <= 0: @@ -926,6 +928,17 @@ def aggregate_records(records: list[JobRecord], args: argparse.Namespace) -> lis if args.dflt_mpcpu: default_mempercpu_gb = float(args.dflt_mpcpu) + avg_function = median_or_none + if args.avg_function: + avgfns = {"mean": mean_or_none, + "median" : median_or_none} + try: + avg_function = avgfns[args.avg_function] + except KeyError: + sys.stderr.write(f"ERROR: No such averaging function: {args.avg_function}\n") + sys.exit(1) + + storvecs = [] if args.histo: if not (args.aggr_regexp or args.aggr_user): @@ -960,7 +973,8 @@ def aggregate_records(records: list[JobRecord], args: argparse.Namespace) -> lis out = [make_aggregate_row(v, username=k[1], jobname=k[0], \ partition=k[2], dflt_mpcpu=default_mempercpu_gb, - storvecs=storvecs) \ + storvecs=storvecs, + avg_function=avg_function) \ for k, v in buckets.items()] out.extend(make_single_row(r, dflt_mpcpu=default_mempercpu_gb, storvecs=storvecs) \ @@ -981,7 +995,8 @@ def aggregate_records(records: list[JobRecord], args: argparse.Namespace) -> lis buckets[key].append(rec) return [make_aggregate_row(v, username=k[0],partition=k[1], \ jobname=f"{common_prefix([r.jobname for r in v])}", - dflt_mpcpu=default_mempercpu_gb, storvecs=storvecs) + dflt_mpcpu=default_mempercpu_gb, storvecs=storvecs, + avg_function=avg_function) for k, v in buckets.items()] return [make_single_row(r, dflt_mpcpu=default_mempercpu_gb, @@ -1069,26 +1084,27 @@ def make_single_row(rec: JobRecord, dflt_mpcpu: float, def make_aggregate_row(records: list[JobRecord], username: str, partition: str, jobname: str, dflt_mpcpu: float, - storvecs: list[str] | None = None) -> OutputRow: + storvecs: list[str] | None = None, + avg_function=median_or_none) -> OutputRow: """Returns an OutputRow based on the given list of job records.""" first = records[0] cpu_eff_vals = [r.cpu_eff for r in records if r.cpu_eff is not None] mem_eff_vals = [r.mem_eff for r in records if r.mem_eff is not None] time_eff_vals = [r.time_eff for r in records if r.time_eff is not None] - walltime=mean_or_none([r.elapsed_sec for r in records if r.elapsed_sec is not None]) + walltime=avg_function([r.elapsed_sec for r in records if r.elapsed_sec is not None]) if walltime is not None: walltime /= 3600 walltime_max=max([r.elapsed_sec for r in records if r.elapsed_sec is not None]) / 3600 - planned_sec = mean_or_none([r.planned_sec for r in records + planned_sec = avg_function([r.planned_sec for r in records if r.planned_sec is not None]) planned_time_h = None if planned_sec is not None: planned_time_h = planned_sec/3600 - elig_qtime_sec = mean_or_none([r.elig_qtime_sec for r in records + elig_qtime_sec = avg_function([r.elig_qtime_sec for r in records if r.elig_qtime_sec is not None]) elig_qtime_h = None if elig_qtime_sec is not None: @@ -1098,12 +1114,12 @@ def make_aggregate_row(records: list[JobRecord], username: str, partition: str, # different allocations, even though all of them had the same user required # Memory. Maybe should generate a warning. Most of the time, what the user # requested should match what the scheduler gave - alloc_mem = mean_or_none([r.mem_alloc_tres for r in records if r.mem_alloc_tres is not None]) - used_mem = mean_or_none([r.mem_used_tres for r in records if r.mem_used_tres is not None]) + alloc_mem = avg_function([r.mem_alloc_tres for r in records if r.mem_alloc_tres is not None]) + used_mem = avg_function([r.mem_used_tres for r in records if r.mem_used_tres is not None]) maxrss_max = max([r.maxrss_bytes for r in records if r.maxrss_bytes is not None]) / (1024**3) or None - memory_efficiency=mean_or_none(mem_eff_vals) + memory_efficiency=avg_function(mem_eff_vals) count=len(records) waste_mem = 0 @@ -1112,7 +1128,7 @@ def make_aggregate_row(records: list[JobRecord], username: str, partition: str, waste_mem = max(0,count * walltime * (100-memory_efficiency)/100 \ * (first.reqmem_gb - first.cpus * dflt_mpcpu)) - cpu_efficiency = mean_or_none(cpu_eff_vals) + cpu_efficiency = avg_function(cpu_eff_vals) waste_cpu=0 if cpu_efficiency is not None and walltime is not None: @@ -1167,7 +1183,7 @@ def make_aggregate_row(records: list[JobRecord], username: str, partition: str, Count=count, CPU_Eff=cpu_efficiency, Mem_Eff=memory_efficiency, - Time_Eff=mean_or_none(time_eff_vals), + Time_Eff=avg_function(time_eff_vals), jobname=jobname, maxrss_max=maxrss_max, walltime=walltime, @@ -1539,6 +1555,8 @@ def parse_args(argv: list[str]) -> argparse.Namespace: """ ) + p.add_argument("--avg-function", choices=["median","mean"], default="median", + help="averaging function to use for reporting aggregate rows",) p.add_argument("-S", "--start", help="sacct start time, passed to sacct -S", default="now - 24 hours") p.add_argument("-E", "--end", help="sacct end time, passed to sacct -E",