From c4acac8dce6d3846296f53facd2c79017f3db766 Mon Sep 17 00:00:00 2001 From: jenuk Date: Fri, 7 Apr 2023 13:28:30 +0200 Subject: [PATCH 1/3] Make output configurable --- gpustat/config.py | 157 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 157 insertions(+) create mode 100644 gpustat/config.py diff --git a/gpustat/config.py b/gpustat/config.py new file mode 100644 index 0000000..ea5dd5e --- /dev/null +++ b/gpustat/config.py @@ -0,0 +1,157 @@ +""" +Defines a dataclass `PrintConfig` to handle options for printing gpustat to the +terminal, such as colors, other modifiers and overall arangement +""" + +import re +from copy import deepcopy +from dataclasses import dataclass, field, fields +from typing import Dict, Literal, Optional, TypeVar, Union + +from blessed import Terminal +from omegaconf import OmegaConf + + +@dataclass +class ConditionalFormat: + # TODO: Allow refernce to something else? + mode: Literal["Larger", "Smaller", "Equal"] + val: Union[int, float, str] + eval_true: str + eval_false: str + eval_error: str = "normal" + +StrPlus = TypeVar("StrPlus", str, ConditionalFormat) + +def str_to_term(s: StrPlus, terminal: Terminal) -> StrPlus: + """Converts a color entry below from human readable form to a ANSI-escape + code of the specific terminal""" + if isinstance(s, ConditionalFormat): + s.eval_true = str_to_term(s.eval_true, terminal) + s.eval_false = str_to_term(s.eval_false, terminal) + s.eval_error = str_to_term(s.eval_error, terminal) + return s + + code = terminal.normal # reset previous color/style + for part in s.split(";"): + if hasattr(terminal, part): + code += getattr(terminal, part) + else: + try: + # TODO: background modifier ? + r, g, b = map(int, part.split(",")) + code += terminal.color_rgb(r, g, b) + except: + raise ValueError(f"Unknown color/style modifier: {part}") + return code + +@dataclass +class ProcecessFontModifiers: + """Set colors for each kind of metric of a process running on a gpu""" + username: str = "bold;black" # change here for readable dark terminals + command: str = "blue" + # TODO: full_command colors command uses two colors + full_command: str = "cyan" + gpu_memory_usage: str = "yellow" + cpu_percent: str = "green" + cpu_memory_usage: str = "yellow" + pid: str = "normal" + + def to_term(self, terminal: Terminal): + for var in fields(self): + value = getattr(self, var.name) + if not isinstance(value, (str, ConditionalFormat)): + continue + modifier = str_to_term(value, terminal) + setattr(self, var.name, modifier) + + +@dataclass +class GPUFontModifiers: + """Set colors for each kind of metric of a gpu""" + # NOTE: dots in json are replaced with underscores here + index: str = "cyan" + uuid: str = "normal" # Not a default option + name: str = "blue" + temperature_gpu: str = "red" # TODO: is Conditional + fan_speed: str = "cyan" # TODO: is Conditional + utilization_gpu: str = "green" # TODO: is Conditional + utilization_enc: str = "green" # TODO: is Conditional + utilization_dec: str = "green" # TODO: is Conditional + power_draw: str = "magenta" # TODO: is Conditional + enforced_power_limit: str = "magenta" + memory_used: str = "yellow" + memory_total: str = "yellow" + processes_font_modifiers: ProcecessFontModifiers = ProcecessFontModifiers() + + def to_term(self, terminal: Terminal): + self.processes_font_modifiers.to_term(terminal) + for var in fields(self): + value = getattr(self, var.name) + if not isinstance(value, (str, ConditionalFormat)): + continue + modifier = str_to_term(value, terminal) + setattr(self, var.name, modifier) + +@dataclass +class FontModifiers: + """Set colors for each kind of all metric indepentendly""" + # NOTE: Should this be overwritable in the string? Seems useless tbh + gpu_font_modifiers: GPUFontModifiers = GPUFontModifiers() + hostname: str = "bold;gray" + driver_version: str = "bold;black" + query_time: str = "normal" + + def to_term(self, terminal: Terminal): + self.gpu_font_modifiers.to_term(terminal) + for var in fields(self): + value = getattr(self, var.name) + if not isinstance(value, (str, ConditionalFormat)): + continue + modifier = str_to_term(value, terminal) + setattr(self, var.name, modifier) + +@dataclass +class PrintConfig: + font_modifiers: FontModifiers = FontModifiers() + # NOTE: Introduce color reset shortcut? + + # # Config 1: Roughly default `gpustat` + # header: Optional[str] = "$hostname:{width}$ $query_time$ $driver_version$\n{gpus}" + # gpus: Optional[str] = "[$index$] $name:{width}$ | $temperature_gpu:2${t.red}°C{t.normal}, $utilization_gpu:3$ {t.green}%{t.normal} | $memory_used:5$ / $memory_total:5$ MB | {processes}" + # process: Optional[str] = "$username$ ($gpu_memory_usage$ MB)" + # gpu_sep: str = "\n" + # processes_sep: str = ", " + + # Config 2: Boxed config + header: Optional[str] = ( + "┌{empty:─>{width}}{empty:─>39}┐\n" + "│ $hostname:{width}$ $query_time$ $driver_version:>11$ │\n" + "{gpus}\n" + "└{empty:─>{width}}{empty:─>39}┘" + ) + gpus: Optional[str] = "│ [$index$] $name:{width}$ │ $temperature_gpu:2${t.red}°C{t.normal}, $utilization_gpu:3$ {t.green}%{t.normal} │ $memory_used:5$ / $memory_total:5$ MB │{processes}" + process: Optional[str] = "\n│ $username:>{width}$ │ $gpu_memory_usage$ MB {empty: >22}│" + gpu_sep: str = "\n" + processes_sep: str = "" + + gpuname_width: Optional[int] = None + use_color: Optional[bool] = None + extra_colors: Dict[str, str] = field(default_factory=dict) # dict() + + def __post_init__(self): + def inner_prepare(s): + return re.sub(r"\$([^:$]*?)(:[^:$]*?)?\$", r"{mods.\1}{\1\2}{t.normal}", s) + self.header = inner_prepare(self.header) + self.gpus = inner_prepare(self.gpus) + self.process = inner_prepare(self.process) + + def to_term(self, terminal: Terminal): + cp = deepcopy(self) + for key, color in cp.extra_colors.values(): + cp.extra_colors[key] = str_to_term(color, terminal) + cp.font_modifiers.to_term(terminal) + return cp + + def to_yaml(self) -> str: + return OmegaConf.to_yaml(self) From 4640a04a6aefdd13bdef4e9a3faa7cab57afadee Mon Sep 17 00:00:00 2001 From: jenuk Date: Fri, 7 Apr 2023 13:29:33 +0200 Subject: [PATCH 2/3] integration with rest --- gpustat/cli.py | 12 ++++++-- gpustat/core.py | 82 ++++++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 91 insertions(+), 3 deletions(-) diff --git a/gpustat/cli.py b/gpustat/cli.py index dbc4e61..d1262e0 100644 --- a/gpustat/cli.py +++ b/gpustat/cli.py @@ -6,6 +6,7 @@ from blessed import Terminal from gpustat import __version__ +from gpustat.config import PrintConfig from gpustat.core import GPUStatCollection @@ -52,7 +53,7 @@ def get_complete_for_one_or_zero(input): return output -def print_gpustat(*, id=None, json=False, debug=False, **kwargs): +def print_gpustat(*, id=None, json=False, debug=False, config=None, **kwargs): '''Display the GPU query results into standard output.''' try: gpu_stats = GPUStatCollection.new_query(debug=debug, id=id) @@ -76,7 +77,9 @@ def print_gpustat(*, id=None, json=False, debug=False, **kwargs): sys.stderr.flush() sys.exit(1) - if json: + if config is not None: + gpu_stats.print_from_config(config, sys.stdout) + elif json: gpu_stats.print_json(sys.stdout) else: gpu_stats.print_formatted(sys.stdout, **kwargs) @@ -189,7 +192,12 @@ def nonnegative_int(value): ) parser.add_argument('-v', '--version', action='version', version=('gpustat %s' % __version__)) + parser.add_argument('--config', default=None, const=True, nargs="?") args = parser.parse_args(argv[1:]) + + + if args.config is not None: + args.config = PrintConfig() # TypeError: GPUStatCollection.print_formatted() got an unexpected keyword argument 'print_completion' with suppress(AttributeError): del args.print_completion # type: ignore diff --git a/gpustat/core.py b/gpustat/core.py index dffda79..4a373dc 100644 --- a/gpustat/core.py +++ b/gpustat/core.py @@ -7,7 +7,7 @@ @url https://github.com/wookayin/gpustat """ -from typing import Sequence +from typing import Any, Dict, Sequence, TypeVar import json import locale import os.path @@ -19,10 +19,13 @@ import psutil from blessed import Terminal +from gpustat.config import PrintConfig import gpustat.util as util from gpustat.nvml import pynvml as N +from pprint import pprint + NOT_SUPPORTED = 'Not Supported' MB = 1024 * 1024 @@ -31,6 +34,17 @@ IS_WINDOWS = 'windows' in platform.platform().lower() +T = TypeVar("T") +def recursive_replace(input: T) -> T: + """Replaces . (dots) in all keys with _ (underscores). Expects json-type + input and returns json-type output""" + if isinstance(input, list): + return [recursive_replace(x) for x in input] + elif isinstance(input, dict): + return {key.replace(".", "_"): recursive_replace(val) for key, val in input.items()} + else: + return input + class GPUStat(object): @@ -736,6 +750,72 @@ def date_handler(obj): fp.write(os.linesep) fp.flush() + def print_from_config(self, config: PrintConfig, fp=sys.stdout, + *, eol_char: str=os.linesep): + # ANSI color configuration + if config.use_color: + TERM = os.getenv('TERM') or 'xterm-256color' + t_color = Terminal(kind=TERM, force_styling=True) + + # workaround of issue #32 (watch doesn't recognize sgr0 characters) + t_color._normal = u'\x1b[0;10m' + elif config.use_color is not None: # <=> is set to false + t_color = Terminal(force_styling=None) + else: + t_color = Terminal() # auto, depending on isatty + config = config.to_term(t_color) + assert config.header is not None + assert config.gpus is not None + assert config.process is not None + + # appearance settings + if config.gpuname_width is None: + if len(self) > 0: + config.gpuname_width = max(len(g.entry['name']) for g in self) + else: + config.gpuname_width = 0 + extra = { + "t": t_color, + "c": config.extra_colors, + "width": config.gpuname_width, + "empty": "", + } + + # Gathering the info + vars = self.jsonify() + vars: dict = recursive_replace(vars) + if IS_WINDOWS: + # no localization is available; just use a reasonable default + # same as str(timestr) but without ms + timestr = self.query_time.strftime('%Y-%m-%d %H:%M:%S') + else: + time_format = locale.nl_langinfo(locale.D_T_FMT) + timestr = self.query_time.strftime(time_format) + vars["query_time"] = timestr + + # Actually building the output + gpu_info: list[str] = [] + for gpu in vars["gpus"]: + processes = [] + for p in gpu["processes"]: + processes.append(config.process.format( + mods=config.font_modifiers.gpu_font_modifiers.processes_font_modifiers, + **p, **extra, + )) + gpu["processes"] = config.processes_sep.join(processes) + gpu_info.append(config.gpus.format( + mods=config.font_modifiers.gpu_font_modifiers, + **gpu, **extra, + )) + vars["gpus"] = config.gpu_sep.join(gpu_info) + complete = config.header.format( + mods=config.font_modifiers, + **vars, **extra, + ) + + fp.write(complete.strip()) + fp.write(eol_char) + fp.flush() def new_query() -> GPUStatCollection: ''' From 72208dbd35d075d16165a86d7f6c3812b3581e0b Mon Sep 17 00:00:00 2001 From: jenuk Date: Fri, 7 Apr 2023 13:36:58 +0200 Subject: [PATCH 3/3] clean up imports --- gpustat/config.py | 6 +++--- gpustat/core.py | 8 +++----- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/gpustat/config.py b/gpustat/config.py index ea5dd5e..fe613dc 100644 --- a/gpustat/config.py +++ b/gpustat/config.py @@ -9,7 +9,7 @@ from typing import Dict, Literal, Optional, TypeVar, Union from blessed import Terminal -from omegaconf import OmegaConf +# from omegaconf import OmegaConf @dataclass @@ -153,5 +153,5 @@ def to_term(self, terminal: Terminal): cp.font_modifiers.to_term(terminal) return cp - def to_yaml(self) -> str: - return OmegaConf.to_yaml(self) + # def to_yaml(self) -> str: + # return OmegaConf.to_yaml(self) diff --git a/gpustat/core.py b/gpustat/core.py index 4a373dc..ce45eea 100644 --- a/gpustat/core.py +++ b/gpustat/core.py @@ -7,7 +7,7 @@ @url https://github.com/wookayin/gpustat """ -from typing import Any, Dict, Sequence, TypeVar +from typing import Sequence, TypeVar import json import locale import os.path @@ -19,13 +19,11 @@ import psutil from blessed import Terminal -from gpustat.config import PrintConfig import gpustat.util as util +from gpustat.config import PrintConfig from gpustat.nvml import pynvml as N -from pprint import pprint - NOT_SUPPORTED = 'Not Supported' MB = 1024 * 1024 @@ -783,7 +781,7 @@ def print_from_config(self, config: PrintConfig, fp=sys.stdout, # Gathering the info vars = self.jsonify() - vars: dict = recursive_replace(vars) + vars = recursive_replace(vars) if IS_WINDOWS: # no localization is available; just use a reasonable default # same as str(timestr) but without ms