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