Config

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

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)

gpustat/core.py

@@ -7,7 +7,7 @@
 @url https://github.com/wookayin/gpustat
 """
 
-from typing import Sequence
+from typing import Sequence, TypeVar
 import json
 import locale
 import os.path
@@ -21,6 +21,7 @@
 from blessed import Terminal
 
 import gpustat.util as util
+from gpustat.config import PrintConfig
 from gpustat.nvml import pynvml as N
 
 
@@ -31,6 +32,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 +748,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 = 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:
     '''