cvgpu_console/python/renderer.py

from math import ceil

__all__ = ["Renderer"]

STANDARD_DICT = {
    "Memory": "Memory",
    "#Proc": "#Proc",
    "Temp": "Temp",
    "Util": "Util",
}

SHORTENED_DICT = {
    "Memory": "Mem",
    "#Proc": "#Pr",
    "Temp": "T",
    "Util": "Ut",
}


def build_line(fill_char, length):
    return "|" + fill_char * (length - 2) + "|"

def fill_line(line, length, fill_char=" "):
    return line + fill_char * (length - len(line) - 1) + "|"


class Renderer:
    def __init__(
        self, 
        progress_bar_width=50, 
        columns=1, 
        gpu_inner_spacings=True,
        gpu_outer_spacings=True,
        node_names=None, 
        display_power=True,
        display_users=False,
        dict_type="default",
        available_gpus_only=False,
    ) -> None:
        assert dict_type in ("default", "shortened")

        self.progress_bar_width = progress_bar_width
        self.columns = columns
        self.gpu_inner_spacings = gpu_inner_spacings
        self.gpu_outer_spacings = gpu_outer_spacings
        self.node_names = node_names
        self.display_power = display_power
        self.display_users = display_users
        self.dict_type = dict_type
        self.available_gpus_only=available_gpus_only

        if dict_type == "default":
            self.act_dict = STANDARD_DICT
        else:
            self.act_dict = SHORTENED_DICT

    def render_info_dict(self, info_dict):
        line_blocks = []
        line_len = -1

        for node_dict in sorted(info_dict, key=lambda n: n["name"]):
            if self.node_names is None or node_dict["name"] in self.node_names:
                lines = self.render_node(node_dict)

                if lines:
                    line_blocks.append(lines)
                    
                    if len(lines[-1]) > line_len:
                        line_len = len(lines[-1])

        if line_len > -1:
            first_line = build_line("=", line_len)
        else: # No gpus found with current settings.
            return "--- No available GPUs found ---"
        
        final_lines = []

        n_rows = ceil(len(line_blocks) / self.columns)
        
        calc_index = lambda row, col, n_rows, n_cols: row * n_cols + col

        # Format rows and columns
        for row in range(n_rows):
            lines = []

            max_size = -1

            # Find max size of the blocks in the same row:
            for col in range(self.columns):
                idx = calc_index(row, col, n_rows, self.columns)
                if idx < len(line_blocks):
                    if len(line_blocks[idx]) > max_size:
                        max_size = len(line_blocks[idx])

            for col in range(self.columns):
                idx = calc_index(row, col, n_rows, self.columns)
                if idx < len(line_blocks):
                    if len(lines) == 0:
                        lines.extend(
                            self.expand_rendered_block_to_size(
                                line_blocks[idx], 
                                max_size
                            )
                        )
                    else:
                        for i, line in enumerate(self.expand_rendered_block_to_size(line_blocks[idx], max_size)):
                            lines[i] += line
            
            final_lines.extend(lines)            

        final_lines.insert(0, first_line * min(self.columns, len(line_blocks)))

        #lines.append("=" * len(lines[-1]))

        return "\n".join(final_lines)

    def render_node(self, node_dict):
        name = node_dict["name"]
        mem_used = node_dict["latest_info"]["used_memory_mb"]
        mem_total = node_dict["total_memory_mb"]
        utilization = node_dict["latest_info"]["cpu_utilization"]
        temp = node_dict["latest_info"]["temperature"]

        head_line = "|- Node: " + name + " "
        info_line = f"| CPU: {utilization:>4.1f}%   {self.act_dict['Memory']}: {mem_used:>6}/{mem_total:<6} MB   {self.act_dict['Temp']}: {temp:>3}°C"

        lines = []
        
        line_len = -1

        for i, gpu_dict in enumerate(node_dict["gpus"]):
            if "detached" in gpu_dict and gpu_dict["detached"]:
                continue
            
            new_lines = self.get_rendered_gpu_lines(gpu_dict)
            
            if line_len == -1:
                line_len = len(new_lines[-1])
            
            if not self.available_gpus_only or \
                (len(gpu_dict["running_processes"]) == 0 and \
                    gpu_dict["latest_info"]["used_memory_mb"] < 100):
                lines.extend(new_lines)
        
        # If no GPUs are available for this node, skip the GPU display
        if self.available_gpus_only and len(lines) == 0:
            return None

        head_line = fill_line(head_line, line_len, fill_char="-")
        info_line = fill_line(info_line, line_len)
        pad_line = build_line("-", line_len)
        pad_line_empty = build_line(" ", line_len)

        lines.append(build_line("=", line_len))

        if self.gpu_outer_spacings:
            lines.insert(0, pad_line_empty)
        lines.insert(0, pad_line)
        if self.gpu_outer_spacings:
            lines.insert(0, pad_line_empty)
        lines.insert(0, info_line)
        if self.gpu_outer_spacings:
            lines.insert(0, pad_line_empty)
        lines.insert(0, head_line)

        return lines

    def get_rendered_gpu_lines(self, gpu_dict):
        gpu_type = self.maybe_shorten_gpu_name(gpu_dict["type"].replace("NVIDIA", "").strip())
        index = gpu_dict["index"]
        mem_used = gpu_dict["latest_info"]["used_memory_mb"]
        mem_total = gpu_dict["total_memory_mb"]
        utilization = gpu_dict["latest_info"]["utilization"]
        temp = gpu_dict["latest_info"]["temperature"]
        n_processes = len(gpu_dict["running_processes"])
        power = gpu_dict["latest_info"]["power_draw"]

        mem_used_percent = int(self.progress_bar_width * mem_used / mem_total)
        rest_mem = self.progress_bar_width - mem_used_percent

        line_util = "| [" + \
            "=" * mem_used_percent + \
            " " * rest_mem + "]" + \
            f"{mem_used:>6}/{mem_total:<6} MB,  {self.act_dict['Util']}: {int(utilization):>3}% |"

        line_meta = f"| " \
            f"GPU #{index} ({gpu_type}):  " + \
            f"{self.act_dict['#Proc']}: {n_processes}   " + \
            f"{self.act_dict['Temp']}: {temp:>3}°C   " + \
            (f"Pow: {int(power):>3} W" if self.display_power else "")

        line_meta = fill_line(line_meta, len(line_util))
        empty_line = build_line(" ", len(line_meta))

        lines = [line_meta, line_util]

        if self.display_users:
            lines.append(self.get_rendered_users_line(gpu_dict["running_processes"], len(line_meta)))
            
            if self.gpu_outer_spacings:
                lines.append(empty_line)
        else:
            if self.gpu_inner_spacings:
                lines.append(empty_line)

        return lines

    def get_rendered_users_line(self, processes_list, line_len):
        user_names = [p_dict["user_name"] for p_dict in processes_list]
        
        line = "| Users: " + ", ".join(user_names)
        line = fill_line(line, line_len)

        return line

    def expand_rendered_block_to_size(self, block_lines, max_size):
        buffer_line = "|" + " " * (len(block_lines[-1]) - 2) + "|"

        block_lines_new = block_lines[:-1] + [buffer_line] * (max_size - len(block_lines)) + block_lines[-1:]

        return block_lines_new

    def maybe_shorten_gpu_name(self, name):
        if self.dict_type == "shortened":
            name = name.replace("GeForce", "").strip()

            if len(name) > 13:
                name = name[:10] + "..."

        return name