diff --git a/bin/computecanada/diskusage_report.py b/bin/computecanada/diskusage_report.py index e45b0b4..f549ba7 100755 --- a/bin/computecanada/diskusage_report.py +++ b/bin/computecanada/diskusage_report.py @@ -7,6 +7,7 @@ import sys from pathlib import Path from datetime import datetime +from colorama import Fore, Style import pwd SUPPORTED_FS_TYPES = {'lustre', 'nfs', 'gpfs', 'nfs4'} @@ -28,6 +29,7 @@ ('links/nearlines', '*'), ], 'gpfs_diskusage_location': None, + 'warning_threshold': 0.95, } cfg = DEFAULT_CONFIG @@ -38,6 +40,7 @@ parser.add_argument("--nearline", default=False, action='store_true', help="Display information for the nearline filesystem") parser.add_argument("--per_user", default=False, action='store_true', help="Display per-user breakdown if available") parser.add_argument("--all_users", default=False, action='store_true', help="Display information for all users of the project") +parser.add_argument('--color', default=True, action=argparse.BooleanOptionalAction, help="Display any usage over quota in red") args = parser.parse_args() def get_network_filesystems(): @@ -218,7 +221,8 @@ def sizeof_fmt(num, suffix="B", scale=1024, units=None): num /= scale return f"{num:.0f}{units[-1]}{suffix}" -def report_quotas(paths_info): +def report_quotas(paths_info, with_color): + any_overquota = False header = ["Description", "Space", "# of files"] has_explorer = False print(f"{header[0]:>39} {header[1]:>20} {header[2]:>18}") @@ -236,9 +240,22 @@ def report_quotas(paths_info): continue quota_type = 'user' if path_info['filesystem'] in ('/home', '/scratch') else quota_info['quota_type'] description = f"{path_info['filesystem']} ({quota_type} {quota_info['identity_name']})" - space = f"{sizeof_fmt(quota_info['space_used_bytes'], scale=space_display_scale)}/{sizeof_fmt(quota_info['space_quota_bytes'], scale=space_display_scale)}" - files = f"{sizeof_fmt(quota_info['file_used'], suffix='', scale=1000)}/{sizeof_fmt(quota_info['file_quota'], suffix='', scale=1000)}" - print(f"{description:>39} {space:>20} {files:>18}") + space_usage, space_quota = quota_info['space_used_bytes'], quota_info['space_quota_bytes'] + files_usage, files_quota = quota_info['file_used'], quota_info['file_quota'] + space_isoverquota = space_usage / (space_quota + 1) > cfg['warning_threshold'] + files_isoverquota = files_usage / (files_quota + 1) > cfg['warning_threshold'] + space_style = Fore.RED + Style.BRIGHT if with_color and space_isoverquota else '' + files_style = Fore.RED + Style.BRIGHT if with_color and files_isoverquota else '' + space_marker = '-> ' if space_isoverquota else '' + files_marker = '-> ' if files_isoverquota else '' + space_ratio = f"{space_marker}{sizeof_fmt(space_usage, scale=space_display_scale)}/{sizeof_fmt(space_quota, scale=space_display_scale)}" + files_ratio = f"{files_marker}{sizeof_fmt(files_usage, suffix='', scale=1000)}/{sizeof_fmt(files_quota, suffix='', scale=1000)}" + print( + f"{description:>39}", + space_style + f"{space_ratio:>20}" + Style.RESET_ALL, + files_style + f"{files_ratio:>18}" + Style.RESET_ALL + ) + any_overquota |= space_isoverquota | files_isoverquota # display breakdowns per user if requested for path, path_info in paths_info.items(): @@ -253,6 +270,12 @@ def report_quotas(paths_info): subprocess.run(["diskusage_rbh", fs[1:], path_info['group']] + diskusage_rbh_arg) print("\n") + if any_overquota: + note_style = Fore.RED + Style.BRIGHT if with_color else '' + print( + f"\n{note_style}-> Usage near or over quota", + f"(warning threshold at {cfg['warning_threshold']*100:.0f}% of the quota)" + Style.RESET_ALL + ) # report breakdown commands if has_explorer: @@ -313,7 +336,7 @@ def deep_update_dict(dict1, dict2): relevant_paths = get_relevant_paths() network_filesystems = get_network_filesystems() paths_info = get_paths_info(relevant_paths, network_filesystems) - report_quotas(paths_info) + report_quotas(paths_info, args.color) if not args.per_user and any([x in cfg['filesystems'].keys() for x in ['/project', '/nearline']]): print("--")