diff --git a/gpustat/cli.py b/gpustat/cli.py index be3b296..bed7c2a 100644 --- a/gpustat/cli.py +++ b/gpustat/cli.py @@ -9,6 +9,7 @@ from blessed import Terminal from gpustat import __version__ +from gpustat.custom_color_action import CustomColorAction from .core import GPUStatCollection @@ -78,6 +79,8 @@ def main(*argv): help='Force to output with colors') parser_color.add_argument('--no-color', action='store_true', help='Suppress colored output') + parser_color.add_argument('--custom-colors', action=CustomColorAction, + help='Specify output colors') parser.add_argument('-a', '--show-all', action='store_true', help='Display all gpu properties above') diff --git a/gpustat/core.py b/gpustat/core.py index a6c4709..3cd937a 100644 --- a/gpustat/core.py +++ b/gpustat/core.py @@ -36,6 +36,7 @@ class GPUStat(object): + CUSTOM_COLORS = {} def __init__(self, entry): if not isinstance(entry, dict): @@ -186,6 +187,7 @@ def print_to(self, fp, show_power=None, gpuname_width=16, term=None, + custom_colors=None, ): if term is None: term = Terminal(stream=sys.stdout) @@ -193,6 +195,9 @@ def print_to(self, fp, # color settings colors = {} + if custom_colors is None: + custom_colors = {} + def _conditional(cond_fn, true_value, false_value, error_value=term.bold_black): try: @@ -202,19 +207,30 @@ def _conditional(cond_fn, true_value, false_value, _ENC_THRESHOLD = 50 - colors['C0'] = term.normal - colors['C1'] = term.cyan - colors['CBold'] = term.bold - colors['CName'] = term.blue + def _set_color(name: str, default_color: str): + color = custom_colors.get(name, '') + if name in GPUStat.CUSTOM_COLORS: + return GPUStat.CUSTOM_COLORS[name] + elif color and getattr(term, color): + GPUStat.CUSTOM_COLORS[name] = getattr(term, color) + return GPUStat.CUSTOM_COLORS[name] + else: + GPUStat.CUSTOM_COLORS[name] = getattr(term, default_color) + return GPUStat.CUSTOM_COLORS[name] + + colors['C0'] = _set_color('C0', 'normal') + colors['C1'] = _set_color('C1', 'cyan') + colors['CBold'] = _set_color('CBold', 'bold') + colors['CName'] = _set_color('CName', 'blue') colors['CTemp'] = _conditional(lambda: self.temperature < 50, term.red, term.bold_red) colors['FSpeed'] = _conditional(lambda: self.fan_speed < 30, term.cyan, term.bold_cyan) - colors['CMemU'] = term.bold_yellow - colors['CMemT'] = term.yellow - colors['CMemP'] = term.yellow - colors['CCPUMemU'] = term.yellow - colors['CUser'] = term.bold_black # gray + colors['CMemU'] = _set_color('CMemU', 'bold_yellow') + colors['CMemT'] = _set_color('CMemT', 'yellow') + colors['CMemP'] = _set_color('CMemP', 'yellow') + colors['CCPUMemU'] = _set_color('CCPUMemU', 'yellow') + colors['CUser'] = _set_color('CUser', 'bold_black') colors['CUtil'] = _conditional(lambda: self.utilization < 30, term.green, term.bold_green) colors['CUtilEnc'] = _conditional( @@ -223,13 +239,13 @@ def _conditional(cond_fn, true_value, false_value, colors['CUtilDec'] = _conditional( lambda: self.utilization_dec < _ENC_THRESHOLD, term.green, term.bold_green) - colors['CCPUUtil'] = term.green + colors['CCPUUtil'] = _set_color('CCPUUtil', 'green') colors['CPowU'] = _conditional( lambda: (self.power_limit is not None and float(self.power_draw) / self.power_limit < 0.4), term.magenta, term.bold_magenta ) - colors['CPowL'] = term.magenta + colors['CPowL'] = _set_color('CPowL', 'magenta') colors['CCmd'] = term.color(24) # a bit dark if not with_colors: @@ -574,7 +590,7 @@ def print_formatted(self, fp=sys.stdout, force_color=False, no_color=False, show_pid=False, show_fan_speed=None, show_codec="", show_power=None, gpuname_width=16, show_header=True, - eol_char=os.linesep, + eol_char=os.linesep, custom_colors=None ): # ANSI color configuration if force_color and no_color: @@ -631,7 +647,8 @@ def print_formatted(self, fp=sys.stdout, force_color=False, no_color=False, show_codec=show_codec, show_power=show_power, gpuname_width=gpuname_width, - term=t_color) + term=t_color, + custom_colors=custom_colors) fp.write(eol_char) fp.flush() diff --git a/gpustat/custom_color_action.py b/gpustat/custom_color_action.py new file mode 100644 index 0000000..b72e98f --- /dev/null +++ b/gpustat/custom_color_action.py @@ -0,0 +1,15 @@ +import argparse + + +class CustomColorAction(argparse.Action): + def __call__(self, _parser, namespace, values, _option_strings=None): + param_dict = getattr(namespace, self.dest) + if param_dict is None: + param_dict = {} + + try: + k, v = values.split("=") + param_dict[k] = v + except ValueError: + pass + setattr(namespace, self.dest, param_dict)