diff --git a/slurm-eff-tool.py b/slurm-eff-tool.py index 729fd14..ec69a75 100755 --- a/slurm-eff-tool.py +++ b/slurm-eff-tool.py @@ -59,6 +59,7 @@ SACCT_FIELDS = [ DEFAULT_COLUMNS = [ "username", + "JobID", "CPUs", "Nodes", "ReqMem", @@ -74,6 +75,7 @@ DEFAULT_COLUMNS = [ # One-character aliases for sorting and output format specifications. ALIASES = { "u": "username", + "i": "JobID", "c": "CPUs", "N": "Nodes", "m": "ReqMem", @@ -122,6 +124,7 @@ class JobRecord: @dataclass class OutputRow: username: str + JobID: str CPUs: int Nodes: int ReqMem: float | None @@ -139,6 +142,7 @@ class OutputRow: def as_dict(self, sdev: bool = False) -> dict[str, Any]: d: dict[str, Any] = { "username": self.username, + "JobID": self.JobID, "CPUs": self.CPUs, "Nodes": self.Nodes, "ReqMem": self.ReqMem, @@ -455,6 +459,7 @@ def aggregate_records(records: list[JobRecord], args: argparse.Namespace) -> lis def make_single_row(rec: JobRecord) -> OutputRow: return OutputRow( username=rec.username, + JobID=rec.raw.get("JobIDRaw", rec.raw.get("JobID", "")), CPUs=rec.cpus, Nodes=rec.nodes, ReqMem=rec.reqmem_gb, @@ -478,6 +483,7 @@ def make_aggregate_row(records: list[JobRecord], username: str, jobname: str) -> time_vals = [r.time_eff for r in records if r.time_eff is not None] return OutputRow( username=username, + JobID="", CPUs=first.cpus, Nodes=first.nodes, ReqMem=first.reqmem_gb, @@ -654,7 +660,7 @@ def parse_args(argv: list[str]) -> argparse.Namespace: "--format", help=( "Slurm-like output format using aliases: " - "u=username,c=CPUs,N=Nodes,m=ReqMem,p=MemPerCPU,l=ReqWalltime,C=Count," + "u=username,i=JobID,c=CPUs,N=Nodes,m=ReqMem,p=MemPerCPU,l=ReqWalltime,C=Count," "e=CPU_Efficiency,M=Memory_Efficiency,t=Time_Efficiency,j=jobname" ), ) @@ -682,12 +688,17 @@ def main(argv: list[str] | None = None) -> int: out_rows = aggregate_records(records, args) out_rows = sort_rows(out_rows, args.sort) + if args.aggr_user or args.aggr_regexp: + output_columns = [c for c in DEFAULT_COLUMNS if c != "JobID"] + else: + output_columns = DEFAULT_COLUMNS + if args.json: print(json.dumps([r.as_dict(sdev=args.sdev) for r in out_rows], indent=2)) elif args.format: print_custom_format(out_rows, args.format, args.sdev) else: - print_table(out_rows, DEFAULT_COLUMNS, args.sdev) + print_table(out_rows, output_columns, args.sdev) return 0