Browse Source

Addition of availability flag

koerschens 3 years ago
parent
commit
c43ee834c4
2 changed files with 33 additions and 7 deletions
  1. 9 0
      python/gpu_status.py
  2. 24 7
      python/renderer.py

+ 9 - 0
python/gpu_status.py

@@ -73,6 +73,12 @@ parser.add_argument(
     default=None,
     help="The node names of nodes to filter for."
 )
+parser.add_argument(
+    "-a",
+    "--available",
+    action="store_true",
+    help="Flag; if set, displays only available GPUs."
+)
 
 args = parser.parse_args()
 
@@ -90,6 +96,7 @@ if args.ultracompact:
         display_power=False,
         display_users=args.users,
         dict_type="shortened",
+        available_gpus_only=args.available,
     )
 elif args.compact:
     renderer = Renderer(
@@ -100,6 +107,7 @@ elif args.compact:
         node_names=args.filter,
         display_power=True,
         display_users=args.users,
+        available_gpus_only=args.available,
     )
 else:
     renderer = Renderer(
@@ -109,6 +117,7 @@ else:
         gpu_outer_spacings=True,
         node_names=args.filter,
         display_users=args.users,
+        available_gpus_only=args.available,
     )
 
 print(renderer.render_info_dict(response.json()))

+ 24 - 7
python/renderer.py

@@ -35,6 +35,7 @@ class Renderer:
         display_power=True,
         display_users=False,
         dict_type="default",
+        available_gpus_only=False,
     ) -> None:
         assert dict_type in ("default", "shortened")
 
@@ -46,6 +47,7 @@ class Renderer:
         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
@@ -59,7 +61,8 @@ class Renderer:
             if self.node_names is None or node_dict["name"] in self.node_names:
                 lines = self.render_node(node_dict)
 
-                line_blocks.append(lines)
+                if lines:
+                    line_blocks.append(lines)
 
         first_line = build_line("=", len(lines[-1]))
         
@@ -111,18 +114,32 @@ class Renderer:
         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"]):
-            lines.extend(self.get_rendered_gpu_lines(gpu_dict))
+            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, len(lines[-1]), fill_char="-")
-        info_line = fill_line(info_line, len(lines[-1]))
-        pad_line = build_line("-", len(lines[-1]))
-        pad_line_empty = build_line(" ", len(lines[-1]))
+        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)
 
         if self.gpu_outer_spacings:
             lines.append(pad_line_empty)
-        lines.append(build_line("=", len(lines[-1])))
+        lines.append(build_line("=", line_len))
 
         if self.gpu_outer_spacings:
             lines.insert(0, pad_line_empty)