Move to xautodl

xautodl/utils/gpu_manager.py

@@ -0,0 +1,86 @@
+import os
+
+
+class GPUManager:
+    queries = (
+        "index",
+        "gpu_name",
+        "memory.free",
+        "memory.used",
+        "memory.total",
+        "power.draw",
+        "power.limit",
+    )
+
+    def __init__(self):
+        all_gpus = self.query_gpu(False)
+
+    def get_info(self, ctype):
+        cmd = "nvidia-smi --query-gpu={} --format=csv,noheader".format(ctype)
+        lines = os.popen(cmd).readlines()
+        lines = [line.strip("\n") for line in lines]
+        return lines
+
+    def query_gpu(self, show=True):
+        num_gpus = len(self.get_info("index"))
+        all_gpus = [{} for i in range(num_gpus)]
+        for query in self.queries:
+            infos = self.get_info(query)
+            for idx, info in enumerate(infos):
+                all_gpus[idx][query] = info
+
+        if "CUDA_VISIBLE_DEVICES" in os.environ:
+            CUDA_VISIBLE_DEVICES = os.environ["CUDA_VISIBLE_DEVICES"].split(",")
+            selected_gpus = []
+            for idx, CUDA_VISIBLE_DEVICE in enumerate(CUDA_VISIBLE_DEVICES):
+                find = False
+                for gpu in all_gpus:
+                    if gpu["index"] == CUDA_VISIBLE_DEVICE:
+                        assert not find, "Duplicate cuda device index : {}".format(
+                            CUDA_VISIBLE_DEVICE
+                        )
+                        find = True
+                        selected_gpus.append(gpu.copy())
+                        selected_gpus[-1]["index"] = "{}".format(idx)
+                assert find, "Does not find the device : {}".format(CUDA_VISIBLE_DEVICE)
+            all_gpus = selected_gpus
+
+        if show:
+            allstrings = ""
+            for gpu in all_gpus:
+                string = "| "
+                for query in self.queries:
+                    if query.find("memory") == 0:
+                        xinfo = "{:>9}".format(gpu[query])
+                    else:
+                        xinfo = gpu[query]
+                    string = string + query + " : " + xinfo + " | "
+                allstrings = allstrings + string + "\n"
+            return allstrings
+        else:
+            return all_gpus
+
+    def select_by_memory(self, numbers=1):
+        all_gpus = self.query_gpu(False)
+        assert numbers <= len(all_gpus), "Require {} gpus more than you have".format(
+            numbers
+        )
+        alls = []
+        for idx, gpu in enumerate(all_gpus):
+            free_memory = gpu["memory.free"]
+            free_memory = free_memory.split(" ")[0]
+            free_memory = int(free_memory)
+            index = gpu["index"]
+            alls.append((free_memory, index))
+        alls.sort(reverse=True)
+        alls = [int(alls[i][1]) for i in range(numbers)]
+        return sorted(alls)
+
+
+"""
+if __name__ == '__main__':
+  manager = GPUManager()
+  manager.query_gpu(True)
+  indexes = manager.select_by_memory(3)
+  print (indexes)
+"""