renderer.py 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114
  1. from math import ceil
  2. class Renderer:
  3. def __init__(self, progress_bar_width=50, columns=1, use_space_lines=True, node_names=None, display_power=True) -> None:
  4. self.progress_bar_width = progress_bar_width
  5. self.columns = columns
  6. self.use_space_lines = use_space_lines
  7. self.node_names = node_names
  8. self.display_power = display_power
  9. def render_info_dict(self, info_dict):
  10. line_blocks = []
  11. for node_dict in info_dict:
  12. if self.node_names is None or node_dict["name"] in self.node_names:
  13. lines = self.render_node(node_dict)
  14. line_blocks.append(lines)
  15. first_line = "|" + "=" * (len(lines[-1]) - 2) + "|"
  16. final_lines = []
  17. n_rows = ceil(len(line_blocks) / self.columns)
  18. # Format rows and columns
  19. for row in range(n_rows):
  20. lines = []
  21. for col in range(self.columns):
  22. if col * n_rows + row < len(line_blocks):
  23. if len(lines) == 0:
  24. lines.extend(line_blocks[col * n_rows + row])
  25. else:
  26. for i, line in enumerate(line_blocks[col * n_rows + row]):
  27. lines[i] += line
  28. final_lines.extend(lines)
  29. final_lines.insert(0, first_line * self.columns)
  30. #lines.append("=" * len(lines[-1]))
  31. return "\n".join(final_lines)
  32. def render_node(self, node_dict):
  33. name = node_dict["name"]
  34. mem_used = node_dict["latest_info"]["used_memory_mb"]
  35. mem_total = node_dict["total_memory_mb"]
  36. utilization = node_dict["latest_info"]["cpu_utilization"]
  37. temp = node_dict["latest_info"]["temperature"]
  38. head_line = "|- Node: " + name + " "
  39. info_line = f"| CPU: {utilization:>4.1f}% Memory: {mem_used:>6}/{mem_total:<6} MB Temp: {temp:>3}°C"
  40. lines = []
  41. for i, gpu_dict in enumerate(node_dict["gpus"]):
  42. lines.extend(self.get_rendered_gpu_lines(gpu_dict))
  43. if i != len(node_dict["gpus"]) - 1:
  44. if self.use_space_lines:
  45. lines.append("|" + "-" * (len(lines[-1]) - 2) + "|")
  46. else:
  47. lines.append("|" + " " * (len(lines[-1]) - 2) + "|")
  48. else:
  49. lines.append("|" + "=" * (len(lines[-1]) - 2) + "|")
  50. head_line = head_line + "-" * (len(lines[-1]) - len(head_line) - 1) + "|"
  51. info_line = info_line + " " * (len(lines[-1]) - len(info_line) - 1) + "|"
  52. pad_line = "|" + "-" * (len(lines[-1]) - 2) + "|"
  53. pad_line_empty = "|" + " " * (len(lines[-1]) - 2) + "|"
  54. lines.insert(0, pad_line)
  55. if self.use_space_lines:
  56. lines.insert(0, pad_line_empty)
  57. lines.insert(0, info_line)
  58. if self.use_space_lines:
  59. lines.insert(0, pad_line_empty)
  60. lines.insert(0, head_line)
  61. return lines
  62. def get_rendered_gpu_lines(self, gpu_dict):
  63. gpu_type = gpu_dict["type"]
  64. index = gpu_dict["index"]
  65. mem_used = gpu_dict["latest_info"]["used_memory_mb"]
  66. mem_total = gpu_dict["total_memory_mb"]
  67. utilization = gpu_dict["latest_info"]["utilization"]
  68. temp = gpu_dict["latest_info"]["temperature"]
  69. n_processes = len(gpu_dict["running_processes"])
  70. power = gpu_dict["latest_info"]["power_draw"]
  71. mem_used_percent = int(self.progress_bar_width * mem_used / mem_total)
  72. rest_mem = self.progress_bar_width - mem_used_percent
  73. line_util = "| [" + \
  74. "=" * mem_used_percent + \
  75. " " * rest_mem + "]" + \
  76. f"{mem_used:>6}/{mem_total:<6} MB, Util: {int(utilization):>3}% |"
  77. line_meta = f"| GPU #{index} ({gpu_type}): #Proc.: {n_processes} Temp: {temp:>3}°C " + (f"Power: {int(power):>3} W" if self.display_power else "")
  78. line_meta = line_meta + " " * (len(line_util) - len(line_meta) - 1) + "|"
  79. empty_line = "|" + " " * (len(line_meta) - 2) + "|"
  80. if self.use_space_lines:
  81. return [empty_line, line_meta, line_util, empty_line]
  82. else:
  83. return [line_meta, line_util]